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

# Nicola Altini (2020)

# V-Net for Hippocampus Segmentation from MRI with PyTorch

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

# python run/validate_torchio.py

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

# python run/validate_torchio.py --dir=path/to/logs/dir --verbose=VERBOSE

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

# Imports

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

import os

import sys

import argparse

import numpy as np

import torch

from sklearn.model_selection import KFold

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

# Local Imports

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

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

sys.path.append(current_path_abs)

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

from run.utils import (train_val_split_config, print_folder, print_config, check_train_set)

from config.config import *

from config.paths import logs_folder, train_images, train_labels

from semseg.train import val_model

from semseg.data_loader import TorchIODataLoader3DValidation

def run(logs_dir="logs"):

    config = SemSegMRIConfig()

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

    # Check training set

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

    check_train_set(config)

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

    # Config

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

    config.batch_size = 1

    print_config(config)

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

                          for idx in range(config.num_folders)]

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

    # Val loop

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

    cuda_dev = torch.device('cuda')

    if config.do_crossval:

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

        # cross-validation

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

        multi_dices_crossval = list()

        mean_multi_dice_crossval = list()

        std_multi_dice_crossval = list()

        kf = KFold(n_splits=config.num_folders)

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

            print_folder(idx, train_index, val_index)

            config_crossval = train_val_split_config(config, train_index, val_index)

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

            # Training (cross-validation)

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

            model_path = path_nets_crossval[idx]

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

            net = torch.load(model_path)

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

            # Validation (cross-validation)

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

            val_data_loader_3D = TorchIODataLoader3DValidation(config_crossval)

            multi_dices, mean_multi_dice, std_multi_dice = val_model(net, val_data_loader_3D,

                                                                     config_crossval, device=cuda_dev)

            multi_dices_crossval.append(multi_dices)

            mean_multi_dice_crossval.append(mean_multi_dice)

            std_multi_dice_crossval.append(std_multi_dice)

            torch.save(net, os.path.join(logs_folder, "model_folder_{:d}.pt".format(idx)))

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

        # Saving Validation Results

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

        multi_dices_crossval_flatten = [item for sublist in multi_dices_crossval for item in sublist]

        mean_multi_dice_crossval_flatten = np.mean(multi_dices_crossval_flatten)

        std_multi_dice_crossval_flatten = np.std(multi_dices_crossval_flatten)

        print("Multi-Dice: {:.4f} +/- {:.4f}".format(mean_multi_dice_crossval_flatten, std_multi_dice_crossval_flatten))

        # Multi-Dice: 0.8668 +/- 0.0337

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

# MAIN

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

if __name__ == "__main__":

    parser = argparse.ArgumentParser(description="Run Validation (With torchio based Data Loader) "

                                                 "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(

        "--net",

        default='vnet',

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

    )

    args = parser.parse_args()

    run(logs_dir=args.dir)