##########################

# Nicola Altini (2020)

# V-Net for Hippocampus Segmentation from MRI with PyTorch

##########################

# python run/validate.py

# python run/validate.py --dir=logs/no_augm_torchio

# python run/validate.py --dir=logs/no_augm_torchio --write=0

# python run/validate.py --dir=path/to/logs/dir --write=WRITE --verbose=VERBOSE

##########################

# Imports

##########################

import torch

import numpy as np

import os

from sklearn.model_selection import KFold

import argparse

import sys

##########################

# Local Imports

##########################

current_path_abs = os.path.abspath('.')

sys.path.append(current_path_abs)

print('{} appended to sys!'.format(current_path_abs))

from config.paths import ( train_images_folder, train_labels_folder, train_prediction_folder,

                           train_images, train_labels,

                           test_images_folder, test_images, test_prediction_folder)

from run.utils import (train_val_split, print_folder, nii_load, sitk_load, nii_write, print_config,

                       sitk_write, print_test, np3d_to_torch5d, torch5d_to_np3d, print_metrics, plot_confusion_matrix)

from config.config import SemSegMRIConfig

from semseg.utils import multi_dice_coeff

from sklearn.metrics import confusion_matrix, f1_score

def run(logs_dir="logs", write_out=False, plot_conf=False):

    ##########################

    # Config

    ##########################

    config = SemSegMRIConfig()

    print_config(config)

    ###########################

    # Load Net

    ###########################

    cuda_dev = torch.device("cuda")

    # Load From State Dict

    # path_net = "logs/model_epoch_0080.pht"

    # net = VNet3D(num_outs=config.num_outs, channels=config.num_channels)

    # net.load_state_dict(torch.load(path_net))

    path_net = os.path.join(logs_dir,"model.pt")

    path_nets_crossval = [os.path.join(logs_dir,"model_folder_{:d}.pt".format(idx))

                          for idx in range(config.num_folders)]

    ###########################

    # Eval Loop

    ###########################

    use_nib = True

    pad_ref = (48,64,48)

    multi_dices = list()

    f1_scores = list()

    os.makedirs(train_prediction_folder, exist_ok=True)

    os.makedirs(test_prediction_folder, exist_ok=True)

    train_and_test = [True, False]

    train_and_test_images = [train_images, test_images]

    train_and_test_images_folder = [train_images_folder, test_images_folder]

    train_and_test_prediction_folder = [train_prediction_folder, test_prediction_folder]

    os.makedirs(train_prediction_folder,exist_ok=True)

    os.makedirs(test_prediction_folder,exist_ok=True)

    train_confusion_matrix = np.zeros((config.num_outs, config.num_outs))

    for train_or_test_images, train_or_test_images_folder, train_or_test_prediction_folder, is_training in \

            zip(train_and_test_images, train_and_test_images_folder, train_and_test_prediction_folder, train_and_test):

        print("Images Folder: {}".format(train_or_test_images_folder))

        print("IsTraining: {}".format(is_training))

        kf = KFold(n_splits=config.num_folders)

        for idx_crossval, (train_index, val_index) in enumerate(kf.split(train_images)):

            if is_training:

                print_folder(idx_crossval, train_index, val_index)

                model_path = path_nets_crossval[idx_crossval]

                print("Model: {}".format(model_path))

                net = torch.load(model_path)

                _, train_or_test_images, _, train_labels_crossval = \

                    train_val_split(train_images, train_labels, train_index, val_index)

            else:

                print_test()

                net = torch.load(path_net)

            net = net.cuda(cuda_dev)

            net.eval()

            for idx, train_image in enumerate(train_or_test_images):

                print("Iter {} on {}".format(idx,len(train_or_test_images)))

                print("Image: {}".format(train_image))

                train_image_path = os.path.join(train_or_test_images_folder, train_image)

                if use_nib:

                    train_image_np, affine = nii_load(train_image_path)

                else:

                    train_image_np, meta_sitk = sitk_load(train_image_path)

                with torch.no_grad():

                    inputs = np3d_to_torch5d(train_image_np, pad_ref, cuda_dev)

                    outputs = net(inputs)

                    outputs_np = torch5d_to_np3d(outputs, train_image_np.shape)

                if write_out:

                    filename_out = os.path.join(train_or_test_prediction_folder, train_image)

                    if use_nib:

                        nii_write(outputs_np, affine, filename_out)

                    else:

                        sitk_write(outputs_np, meta_sitk, filename_out)

                if is_training:

                    train_label = train_labels_crossval[idx]

                    train_label_path = os.path.join(train_labels_folder, train_label)

                    if use_nib:

                        train_label_np, _ = nii_load(train_label_path)

                    else:

                        train_label_np, _ = sitk_load(train_label_path)

                    multi_dice = multi_dice_coeff(np.expand_dims(train_label_np,axis=0),

                                                  np.expand_dims(outputs_np,axis=0),

                                                  config.num_outs)

                    print("Multi Class Dice Coeff = {:.4f}".format(multi_dice))

                    multi_dices.append(multi_dice)

                    f1_score_idx = f1_score(train_label_np.flatten(), outputs_np.flatten(), average=None)

                    cm_idx = confusion_matrix(train_label_np.flatten(), outputs_np.flatten())

                    train_confusion_matrix += cm_idx

                    f1_scores.append(f1_score_idx)

            if not is_training:

                break

    print_metrics(multi_dices, f1_scores, train_confusion_matrix)

    if plot_conf:

        plot_confusion_matrix(train_confusion_matrix,

                              target_names=None, title='Cross-Validation Confusion matrix',

                              cmap=None, normalize=False, already_normalized=False,

                              path_out="images/conf_matrix_no_norm_no_augm_torchio.png")

        plot_confusion_matrix(train_confusion_matrix,

                              target_names=None, title='Cross-Validation Confusion matrix (row-normalized)',

                              cmap=None, normalize=True, already_normalized=False,

                              path_out="images/conf_matrix_normalized_row_no_augm_torchio.png")

############################

# MAIN

############################

if __name__ == "__main__":

    parser = argparse.ArgumentParser(description="Run Validation for Hippocampus Segmentation")

    parser.add_argument(

        "-V",

        "--verbose",

        default=False, type=bool,

        help="Boolean flag. Set to true for VERBOSE mode; false otherwise."

    )

    parser.add_argument(

        "-D",

        "--dir",

        default="logs", type=str,

        help="Local path to logs dir"

    )

    parser.add_argument(

        "-W",

        "--write",

        default=False, type=bool,

        help="Boolean flag. Set to true for WRITE mode; false otherwise."

    )

    parser.add_argument(

        "--net",

        default='vnet',

        help="Specify the network to use [unet | vnet] ** FOR FUTURE RELEASES **"

    )

    args = parser.parse_args()

    run(logs_dir=args.dir, write_out=args.write, plot_conf=args.verbose)