"""

Developed by: Daniel Crovo

"""

import numpy as np

import csv

from hs_dataset import HSDataset

from torch.utils.data import DataLoader

import torch

import os

import torchvision

from torchmetrics.classification import BinaryJaccardIndex

import torchvision.transforms as transforms

import torchvision.utils as vutils

import torch

from PIL import Image

def get_loaders(train_dir, train_maskdir, val_dir, val_maskdir, batch_size, train_transform, 

                val_transform, num_workers, w_level, w_width, pin_memory = True, normalize=False):

    """_summary_

    Args:

        train_dir (_type_): _description_

        train_maskdir (_type_): _description_

        val_dir (_type_): _description_

        val_maskdir (_type_): _description_

        batch_size (_type_): _description_

        train_transform (_type_): _description_

        val_transform (_type_): _description_

        num_workers (_type_): _description_

        pin_memory (bool, optional): _description_. Defaults to True.

    Returns:

        _type_: _description_

    """

    train_img_mask = HSDataset(image_dir = train_dir, 

                               mask_dir = train_maskdir, transform = train_transform, 

                               normalized=normalize, w_level=w_level, w_width=w_width)

    train_loader = DataLoader(train_img_mask, batch_size = batch_size, 

                              num_workers = num_workers, pin_memory = pin_memory, shuffle = True)

    val_img_mask = HSDataset(image_dir = val_dir, mask_dir = val_maskdir, 

                             transform = val_transform, normalized=normalize,

                             w_level=w_level, w_width=w_width,)

    val_loader = DataLoader(val_img_mask, batch_size = batch_size, 

                            num_workers = num_workers, pin_memory = pin_memory, shuffle = False)

    return train_loader, val_loader

def load_checkpoint(checkpoint, model): 

    try: 

        model.load_state_dict(checkpoint['state_dict'])

        print('\nCheckpoint importado exitosamente.')

    except: 

        print('Error en la importación del Checkpoint.')

def save_checkpoint(state, filename = 'my_checkpoint.pth.tar'): 

    try: 

        torch.save(state, filename)

        print('\nCheckpoint almacenado exitosamente.')

    except: 

        print('Error en la importación del Checkpoint.')

def compute_jaccard_index(preds, targets):

    intersection = torch.logical_and(preds, targets).sum()

    union = torch.logical_or(preds, targets).sum()

    jaccard = intersection.item() / (union.item()+ 1e-10)

    return jaccard

def diceCoef(y_true, y_pred, smooth=1.):

  y_true_f = y_true.flatten()

  y_pred_f = y_pred.flatten()

  intersection = np.sum(y_true_f * y_pred_f)

  dice = (2. * intersection + smooth) / (np.sum(y_true_f) + np.sum(y_pred_f) + smooth)

  return round(float(dice), 3)

def perf(loader, model, device): 

    dice_score = 0.0

    jaccard = 0.0

    model.eval()

    with torch.no_grad(): # deshabilitar el cálculo y almacenamiento de gradientes en el grafo computacional de PyTorch

        for x, y in loader: 

            x = x.to(device = device)

            y = y.to(device).unsqueeze(1)

            preds = torch.sigmoid(model(x))

            preds = (preds > 0.5).float()

            if((preds.sum() == 0) and (y.sum() == 0)):

                jaccard +=1

                dice_score +=1

            else:

                jaccard += compute_jaccard_index(preds, y)

                dice_score += (2*(preds*y).sum())/((preds + y).sum() + 1e-10)

    dice_s = dice_score/len(loader) 

    jaccard_idx = jaccard/len(loader)

    print('\nDice score: {}'.format(dice_s))

    print('Jaccard index: {}\n'.format(jaccard_idx))

    model.train()

    return dice_s, jaccard_idx

def save_preds_as_imgs(loader, model, device, folder = '/home/danielcrovo/Documents/01.Study/01.MSc/02.MSc AI/Deep Learning/Proyecto/saved_images'): 

    model.eval()

    for idx, (x, y) in enumerate(loader):

        x = x.to(device = device)

        with torch.no_grad(): # deshabilitar el cálculo y almacenamiento de gradientes en el grafo computacional de PyTorch

            preds = torch.sigmoid(model(x))

            preds = (preds > 0.5).float()

        torchvision.utils.save_image(preds, f'{folder}/y_hat_{idx}.png') # almacenamiento de máscaras predichas

        y = torch.unsqueeze(y, 1).to(torch.float32)

        torchvision.utils.save_image(y, f'{folder}/y_{idx}.png') # almacenamiento de máscaras reales

        torchvision.utils.save_image(x, f'{folder}/x_{idx}.png') # almacenamiento de máscaras reales

        #masked_img= add_mask_to_rgb_image(x,preds)

        #torchvision.utils.save_image(masked_img, f'{folder}/masked_{idx}.png') # almacenamiento de máscaras reales

    model.train()

def add_mask_to_rgb_image(rgb_image_tensor, mask_tensor):

    # Apply the mask to the RGB image tensor

    masked_image_tensor = torch.where(mask_tensor > 0, torch.tensor([32,255, 0, 0]), rgb_image_tensor)

    return masked_image_tensor    #print(dice_score)

    #train_loss = [tensor.cpu() for tensor in train_loss]

    #train_dice = [tensor.cpu() for tensor in train_dice]

def save_metrics_to_csv(epoch, train_loss, train_dice, train_jaccard, dice_score, jaccard_score, filename):

    metrics = {

        'epoch': epoch,

        'train_loss': train_loss,

        'train_dice': train_dice,

        'train_jaccard': train_jaccard,

        'dice_score': dice_score,

        'jaccard_score': jaccard_score

    }

    file_exists = os.path.isfile(filename)

    with open(filename, mode='a', newline='') as file:

        writer = csv.DictWriter(file, fieldnames=metrics.keys())

        if not file_exists:

            writer.writeheader()

        writer.writerow(metrics)

def save_preds_as_imgs2(loader, model, device, folder='/home/danielcrovo/Documents/01.Study/01.MSc/02.MSc AI/Deep Learning/Proyecto/saved_images'):

    model.eval()

    for idx, (x, y) in enumerate(loader):

        x = x.to(device=device)

        with torch.no_grad():

            preds = torch.sigmoid(model(x))

            preds = (preds > 0.5).float()

        # Convert single-channel mask to RGB format

        y_rgb = y.repeat(1, 3, 1, 1)

        # Merge predicted mask with input image

        merged_img = x.clone()

        merged_img[:, :3] = torch.where(preds > 0, torch.tensor([1.0, 0.0, 0.0]), merged_img[:, :3])

        # Save the merged image

        vutils.save_image(merged_img, f'{folder}/merged_{idx}.png')