##########################
# Nicola Altini (2020)
# V-Net for Hippocampus Segmentation from MRI with PyTorch
##########################
# python run/train.py
# python run/train.py --epochs=NUM_EPOCHS --batch=BATCH_SIZE --workers=NUM_WORKERS --lr=LR
# python run/train.py --epochs=5 --batch=1 --net=unet

##########################
# Imports
##########################
import argparse
import os
import sys
import numpy as np
import torch
import torch.optim as optim
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 print_config, check_train_set, check_torch_loader, print_folder, train_val_split_config
from config.config import SemSegMRIConfig
from config.paths import logs_folder
from semseg.train import train_model, val_model
from semseg.data_loader import TorchIODataLoader3DTraining, TorchIODataLoader3DValidation
from models.vnet3d import VNet3D
from models.unet3d import UNet3D


def get_net(config):
    name = config.net
    assert name in ['unet', 'vnet'], "Network Name not valid or not supported! Use one of ['unet', 'vnet']"
    if name == 'vnet':
        return VNet3D(num_outs=config.num_outs, channels=config.num_channels)
    elif name == 'unet':
        return UNet3D(num_out_classes=config.num_outs, input_channels=1, init_feat_channels=32)


def run(config):
    ##########################
    # Check training set
    ##########################
    check_train_set(config)

    ##########################
    # Config
    ##########################
    print_config(config)

    ##########################
    # Check Torch DataLoader and Net
    ##########################
    check_torch_loader(config, check_net=False)

    ##########################
    # Training loop
    ##########################
    cuda_dev = torch.device('cuda')

    if config.do_crossval:
        ##########################
        # Training (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(config.train_images)):
            print_folder(idx, train_index, val_index)
            config_crossval = train_val_split_config(config, train_index, val_index)

            ##########################
            # Training (cross-validation)
            ##########################
            net = get_net(config_crossval)
            config_crossval.lr = 0.01
            optimizer = optim.Adam(net.parameters(), lr=config_crossval.lr)
            train_data_loader_3D = TorchIODataLoader3DTraining(config_crossval)
            net = train_model(net, optimizer, train_data_loader_3D,
                              config_crossval, device=cuda_dev, logs_folder=logs_folder)

            ##########################
            # 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.8728 +/- 0.0227

    ##########################
    # Training (full training set)
    ##########################
    net = get_net(config)
    config.lr = 0.01
    optimizer = optim.Adam(net.parameters(), lr=config.lr)
    train_data_loader_3D = TorchIODataLoader3DTraining(config)
    net = train_model(net, optimizer, train_data_loader_3D,
                      config, device=cuda_dev, logs_folder=logs_folder)

    torch.save(net,os.path.join(logs_folder,"model.pt"))


############################
# MAIN
############################
if __name__ == "__main__":
    config = SemSegMRIConfig()

    parser = argparse.ArgumentParser(description="Run Training on Hippocampus Segmentation")
    parser.add_argument(
        "-e",
        "--epochs",
        default=config.epochs, type=int,
        help="Specify the number of epochs required for training"
    )
    parser.add_argument(
        "-b",
        "--batch",
        default=config.batch_size, type=int,
        help="Specify the batch size"
    )
    parser.add_argument(
        "-v",
        "--val_epochs",
        default=config.val_epochs, type=int,
        help="Specify the number of validation epochs during training ** FOR FUTURE RELEASES **"
    )
    parser.add_argument(
        "-w",
        "--workers",
        default=config.num_workers, type=int,
        help="Specify the number of workers"
    )
    parser.add_argument(
        "--net",
        default='vnet',
        help="Specify the network to use [unet | vnet] ** FOR FUTURE RELEASES **"
    )
    parser.add_argument(
        "--lr",
        default=config.lr, type=float,
        help="Learning Rate"
    )

    args = parser.parse_args()
    config.net = args.net
    config.epochs = args.epochs
    config.batch_size = args.batch
    config.val_epochs = args.val_epochs
    config.num_workers = args.workers
    config.lr = args.lr

    run(config)