from utils import ( # .py previamente desarrollado

                       get_loaders,

                       load_checkpoint,

                       save_checkpoint,

                       perf,

                       save_preds_as_imgs,

                       save_metrics_to_csv,

                       compute_jaccard_index,

                      )

import torch

from segmentation_models_pytorch.encoders import get_preprocessing_fn

import multiprocessing # para conocer el número de workers (núcleos)

import albumentations as A # para emplear data augmentation

from albumentations.pytorch import ToTensorV2 # para convertir imágenes y etiquetas en tensores de Pytorch

import torch.optim as optim # para optimización de parámetros

import segmentation_models_pytorch as smp

CHECKPOINT_PATH = '/home/danielcrovo/Documents/01.Study/01.MSc/02.MSc AI/Deep Learning/Heart_Segmentation/checkpoints/MANET_my_checkpoint.pth.tar'

TEST_IMAGES = '/home/danielcrovo/Documents/01.Study/01.MSc/02.MSc AI/Deep Learning/Heart_Segmentation/2d_data/test_images'

TEST_MASK = '/home/danielcrovo/Documents/01.Study/01.MSc/02.MSc AI/Deep Learning/Heart_Segmentation/2d_data/test_masks'

TRAIN_IMG_DIR = '/home/danielcrovo/Documents/01.Study/01.MSc/02.MSc AI/Deep Learning/Heart_Segmentation/2d_data/images'

TRAIN_MAKS_DIR = '/home/danielcrovo/Documents/01.Study/01.MSc/02.MSc AI/Deep Learning/Heart_Segmentation/2d_data/masks'

IMAGE_HEIGHT = 192 # 1280 px.

IMAGE_WIDTH = 192 # 1918 px.

PREDS_PATH = '/home/danielcrovo/Documents/01.Study/01.MSc/02.MSc AI/Deep Learning/Heart_Segmentation/saved_images_test'

DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'

def evaluate_training_set(model, loader, device):

    model.eval()

    dice_score = 0.0

    jaccard = 0.0

    num_samples = 0

    with torch.no_grad():

        for data, y in loader:

            data = data.to(device)

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

            preds = model(data)

            preds = torch.sigmoid(model(data))

            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)

    return dice_s, jaccard_idx

def main():

    x_start = 50

    x_end = 462

    y_start = 50

    y_end = 462 

    val_transforms = A.Compose(

                               [A.Crop(x_start, y_start, x_end, y_end, always_apply= False),

                                A.Resize(height = IMAGE_HEIGHT, width = IMAGE_WIDTH), 

                                A.Normalize(

                                            mean = [0.0, 0.0, 0.0],

                                            std = [1.0, 1.0, 1.0], 

                                            max_pixel_value = 255.0

                                         ),

                              ToTensorV2(), # conversión a tensor de Pytorch

                             ],

                            )

    train_loader, test_loader = get_loaders( # cargar y preprocesar los datos de entrenamiento y validación con base en los formatos tensoriales

                                           TRAIN_IMG_DIR, 

                                           TRAIN_MAKS_DIR, 

                                           TEST_IMAGES, 

                                           TEST_MASK,

                                           16,

                                           None,

                                           val_transforms, 

                                           16,

                                           w_level=40, 

                                           w_width=350,

                                           pin_memory= True, 

                                           normalize=False, 

                                          )

    model = smp.MAnet(   

                encoder_name="mit_b3",        # choose encoder, e.g. mobilenet_v2 or efficientnet-b7

                encoder_weights="imagenet",     # use `imagenet` pre-trained weights for encoder initialization

                in_channels=3,                  # model input channels (1 for gray-scale images, 3 for RGB, etc.)

                classes=1,

                decoder_use_batchnorm=True,

                ).to(DEVICE)

    load_checkpoint(torch.load(CHECKPOINT_PATH), model)

    dice_score, jaccard_index = evaluate_training_set(model, test_loader, DEVICE)

    print(f"Average Dice Score (Training Set): {dice_score:.4f}")

    print(f"Average Jaccard Index (Training Set): {jaccard_index:.4f}")

    save_preds_as_imgs(test_loader, model, DEVICE,PREDS_PATH)

if __name__ == '__main__': 

    main()