"""Brain Mapper Run File

Description:
-------------
This file contains all the relevant functions for running BrainMapper.
The network can be ran in one of these modes:
    - train
    - evaluate path
    - evaluate whole

TODO: Might be worth adding some information on uncertaintiy estimation, later down the line


Usage
-------------
In order to run the network, in the terminal, the user needs to pass it relevant arguments:
    - (TODO: ADD ARGUMENTS)

"""

import torch
from utils.data_utils import get_datasets
from BrainMapperUNet import BrainMapperUNet
import torch.utils.data as data
from solver import Solver
import os
from utils.data_logging_utils import LogWriter
import utils.data_evaluation_utils as evaluations
import shutil

# Set the default floating point tensor type to FloatTensor

torch.set_default_tensor_type(torch.FloatTensor)

def load_data(data_parameters):
    """Dataset Loader

    This function loads the training and testing datasets.
    TODO: Will need to define if all the training data is loaded as bulk or individually!

    Args:
        data_parameters (dict): Dictionary containing relevant information for the datafiles.
        data_parameters = {
            'data_directory': 'path/to/directory'
            'train_data_file': 'training_data'
            'train_output_targets': 'training_targets'
            'test_data_file': 'testing_data'
            'test_target_file': 'testing_targets'
        }

    Returns:
        train_data (dataset object): Pytorch map-style dataset object, mapping indices to training data samples.
        test_data (dataset object): Pytorch map-style dataset object, mapping indices to testing data samples.

    Raises:
        None

    """
    print("Data is loading...")
    train_data, test_data = get_datasets(data_parameters)
    print("Data has loaded!")
    print("Training dataset size is {}".format(len(train_data)))
    print("Testing dataset size is {}".format(len(test_data)))

    return train_data, test_data

def train(data_parameters, training_parameters, network_parameters, misc_parameters):
    """Training Function
    
    This function trains a given model using the provided training data.
    Currently, the data loaded is set to have multiple sub-processes. 
    A high enough number of workers assures that CPU computations are efficiently managed, i.e. that the bottleneck is indeed the neural network's forward and backward operations on the GPU (and not data generation)
    Loader memory is also pinned, to speed up data transfer from CPU to GPU  by using the page-locked memory.
    Train data is also re-shuffled at each training epoch. 

    Args:
        data_parameters (dict): Dictionary containing relevant information for the datafiles.
        data_parameters = {
            'data_directory': 'path/to/directory'
            'train_data_file': 'training_data'
            'train_output_targets': 'training_targets'
            'test_data_file': 'testing_data'
            'test_target_file': 'testing_targets'
        }

        training_parameters(dict): Dictionary containing relevant hyperparameters for training the network.
        training_parameters = {
            'training_batch_size': 5
            'test_batch_size: 5
            'use_pre_trained': False
            'pre_trained_path': 'pre_trained/path'
            'experiment_name': 'experiment_name'
            'learning_rate': 1e-4
            'optimizer_beta': (0.9, 0.999)
            'optimizer_epsilon': 1e-8
            'optimizer_weigth_decay': 1e-5
            'number_of_epochs': 10
            'loss_log_period': 50
            'learning_rate_scheduler_step_size': 3
            'learning_rate_scheduler_gamma': 1e-1
            'use_last_checkpoint': True
            'final_model_output_file': 'path/to/model'
        }

        network_parameters (dict): Contains information relevant parameters 
        network_parameters= {
            'kernel_heigth': 5
            'kernel_width': 5
            'kernel_classification': 1
            'input_channels': 1
            'output_channels': 64
            'convolution_stride': 1
            'dropout': 0.2
            'pool_kernel_size': 2
            'pool_stride': 2
            'up_mode': 'upconv'
            'number_of_classes': 1
        }

        misc_parameters (dict): Dictionary of aditional hyperparameters
        misc_parameters = {
            'save_model_directory': 'directory_name'
            'model_name': 'BrainMapper'
            'logs_directory': 'log-directory'
            'device': 1
            'experiments_directory': 'experiments-directory'
        }

    Returns:
        None

    Raises:
        None
    """

    train_data, test_data = load_data(data_parameters)

    train_loader = data.DataLoader(
        dataset= train_data,
        batch_size= training_parameters['training_batch_size'],
        shuffle= True,
        num_workers= 4,
        pin_memory= True
    )

    test_loader = data.DataLoader(
        dataset= test_data,
        batch_size= training_parameters['test_batch_size'],
        shuffle= False,
        num_workers= 4,
        pin_memory= True 
    )

    if training_parameters['use_pre_trained']:
        BrainMapperModel = torch.load(training_parameters['pre_trained_path'])
    else:
        BrainMapperModel = BrainMapperUNet(network_parameters)

    solver = Solver(model= BrainMapperModel,
                    device= misc_parameters['device'],
                    number_of_classes= network_parameters['number_of_classes'],
                    experiment_name= training_parameters['experiment_name'],
                    optimizer_arguments = {'lr': training_parameters['learning_rate'],
                                            'betas': training_parameters['optimizer_beta'],
                                            'eps': training_parameters['optimizer_epsilon'],
                                            'weight_decay': training_parameters['optimizer_weigth_decay']
                                            },
                    model_name = misc_parameters['model_name'],
                    number_epochs = training_parameters['number_of_epochs'],
                    loss_log_period = training_parameters['loss_log_period'],
                    learning_rate_scheduler_step_size = training_parameters['learning_rate_scheduler_step_size'],
                    learning_rate_scheduler_gamma = training_parameters['learning_rate_scheduler_gamma'],
                    use_last_checkpoint = training_parameters['use_last_checkpoint'],
                    experiment_directory = misc_parameters['experiments_directory'],
                    logs_directory = misc_parameters['logs_directory']
                    )

    solver.train(train_loader, test_loader)

    model_output_path = os.path.join(misc_parameters['save_model_directory'], training_parameters['final_model_output_file'])
    BrainMapperModel.save(model_output_path)

    print("Final Model Saved in: {}".format(model_output_path))

def evaluate_score(data_parameters, training_parameters, network_parameters, misc_parameters, evaluation_parameters):
    """Mapping Score Evaluator

    This function evaluates a given trained model by calculating the it's dice score prediction.

    Args:
        data_parameters (dict): Dictionary containing relevant information for the datafiles.
        data_parameters = {
            'data_directory': 'path/to/directory'
            'train_data_file': 'training_data'
            'train_output_targets': 'training_targets'
            'test_data_file': 'testing_data'
            'test_target_file': 'testing_targets'
        }

        training_parameters(dict): Dictionary containing relevant hyperparameters for training the network.
        training_parameters = {
            'training_batch_size': 5
            'test_batch_size: 5
            'use_pre_trained': False
            'pre_trained_path': 'pre_trained/path'
            'experiment_name': 'experiment_name'
            'learning_rate': 1e-4
            'optimizer_beta': (0.9, 0.999)
            'optimizer_epsilon': 1e-8
            'optimizer_weigth_decay': 1e-5
            'number_of_epochs': 10
            'loss_log_period': 50
            'learning_rate_scheduler_step_size': 3
            'learning_rate_scheduler_gamma': 1e-1
            'use_last_checkpoint': True
            'final_model_output_file': 'path/to/model'
        }

        network_parameters (dict): Contains information relevant parameters 
        network_parameters= {
            'kernel_heigth': 5
            'kernel_width': 5
            'kernel_classification': 1
            'input_channels': 1
            'output_channels': 64
            'convolution_stride': 1
            'dropout': 0.2
            'pool_kernel_size': 2
            'pool_stride': 2
            'up_mode': 'upconv'
            'number_of_classes': 1
        }

        misc_parameters (dict): Dictionary of aditional hyperparameters
        misc_parameters = {
            'save_model_directory': 'directory_name'
            'model_name': 'BrainMapper'
            'logs_directory': 'log-directory'
            'device': 1
            'experiments_directory': 'experiments-directory'
        }

        evaluation_parameters (dict): Dictionary of parameters useful during evaluation.
        evaluation_parameters = {
            'trained_model_path': 'path/to/model'
            'data_directory': 'path/to/data'
            'targets_directory': 'path/to/targets'
            'data_list': 'path/to/datalist.txt/
            'orientation': 'coronal'
            'saved_predictions_directory': 'directory-of-saved-predictions'
        }

    Returns:
        None

    Raises:
        None
    """

    logWriter = LogWriter(number_of_classes= network_parameters['number_of_classes'],
                        logs_directory= misc_parameters['logs_directory'], 
                        experiment_name= training_parameters['experiment_name']
                        )

    prediction_output_path = os.path.join(misc_parameters['experiments_directory'],
                                            training_parameters['experiment_name'],
                                            evaluation_parameters['saved_predictions_directory']
                                            )

    average_dice_score = evaluations.evaluate_dice_score(trained_model_path= evaluation_parameters['trained_model_path'],
                                                        number_of_classes= network_parameters['number_of_classes'],
                                                        data_directory= evaluation_parameters['data_directory'],
                                                        targets_directory= evaluation_parameters['targets_directory'],
                                                        data_list= evaluation_parameters['data_list'],
                                                        orientation= evaluation_parameters['orientation'],
                                                        prediction_output_path= prediction_output_path,
                                                        device= misc_parameters['device'],
                                                        LogWriter= logWriter
                                                        )

    logWriter.close()

def evaluate_mapping(mapping_evaluation_parameters):
    """Mapping Evaluator

    This function passes through the network an input and generates the rsfMRI outputs.
    This function allows the user to either use one or two or three paths.

    The convention for the different model paths is as follows:
    - model1: coronal
    - model2: axial
    - model3: saggital

    However, this convention can be changed either bellow or the settings file.

    Args:
        mapping_evaluation_parameters (dict): Dictionary of parameters useful during mapping evaluation.
        mapping_evaluation_parameters = {
            'trained_model1_path': 'path/to/model1'
            'trained_model2_path': 'path/to/model2'
            'trained_model3_path': 'path/to/model3'
            'data_directory': 'path/to/data'
            'data_list': 'path/to/datalist.txt/
            'orientation1': 'coronal'
            'orientation2': 'axial'
            'orientation3': 'sagittal'
            'prediction_output_path': 'directory-of-saved-predictions'
            'batch_size': 2
            'device': 0
            'exit_on_error': True
            'number_of_paths': 3
        }

    Returns:
        None

    Raises:
        None
    """
    trained_model1_path = mapping_evaluation_parameters['trained_model1_path']
    trained_model2_path = mapping_evaluation_parameters['trained_model2_path']
    trained_model3_path = mapping_evaluation_parameters['trained_model3_path']
    data_directory = mapping_evaluation_parameters['data_directory']
    data_list = mapping_evaluation_parameters['data_list']
    orientation1 = mapping_evaluation_parameters['orientation1']
    orientation2 = mapping_evaluation_parameters['orientation2']
    orientation3 =mapping_evaluation_parameters['orientation3']
    prediction_output_path = mapping_evaluation_parameters['prediction_output_path']
    batch_size = mapping_evaluation_parameters['batch_size']
    device= mapping_evaluation_parameters['device']
    exit_on_error = mapping_evaluation_parameters['exit_on_error']
    

    if mapping_evaluation_parameters['number_of_paths'] == 1:
        evaluations.evaluate_single_path(trained_model_path= trained_model1_path,
                        data_directory,
                        data_list,
                        orientation= orientation1,
                        prediction_output_path,
                        batch_size,
                        device= device,
                        exit_on_error= exit_on_error):
    elif mapping_evaluation_parameters['number_of_paths'] == 2:
        evaluations.evaluate_two_paths(trained_model1_path,
                        trained_model2_path,
                        data_directory,
                        data_list,
                        orientation1,
                        orientation2,
                        prediction_output_path,
                        batch_size,
                        device= device,
                        exit_on_error= exit_on_error)
    elif mapping_evaluation_parameters['number_of_paths'] == 3:
        evaluations.evaluate_all_paths(trained_model1_path,
                        trained_model2_path,
                        trained_model3_path,
                        data_directory,
                        data_list,
                        orientation1,
                        orientation2,
                        orientation3,
                        prediction_output_path,
                        batch_size,
                        device= device
                        exit_on_error= exit_on_error)

def delete_files():
    """ Clear Folder Contents
    
    Function which clears contents (like experiments or logs)

    Args:
        folder (str): Name of folders whose conents is to be deleted
    
    Returns:
        None
    
    Raises:
        Exception: Any error
    """
    
    for object_name in os.listdir(folder):
        file_path = os.path.join(folder, object_name)
        try:
            if os.path.isfile(file_path):
                os.unlink(file_path)
            elif os.path.isdir(file_path):
                shutil.rmtree(file_path)
        except Exception as exception:
            print(exception)


if __name__ == '__main__':
    pass