import os

import torch

from trainer import train_model

import utils as ut

from loss.diceloss import diceloss

from metrics import m

from models.UnetAttention import UnnetAttention

def run_nn():

    """

        Version requirements:

            PyTorch Version:  >1.2.0

            Torchvision Version:  >0.4.0a0+6b959ee

    """

    """

        Parameters to execute the method

    """

    root_dir = r'./data'

    epochs = 100

    batch_size = 8

    # Filename of the final model weigths

    weight_filename = "weights_final.pt"

    data_aug = 'online'

    log_path = './weights/'

    """

        Main 

    """

    ut.create_nested_dir(log_path)

    # Loads the distribution of the cases between train and val

    cases = ut.load_dataset_dist()

    # Create the dataloader

    dataloaders = ut.get_data_loaders(

        data_aug, cases, root_dir, batch_size)

    model = UnnetAttention()

    model.train()

    # Load the loss object by name

    criterion = diceloss()

    # Specify the optimizer with a lower learning rate

    optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)

    exp_lr_scheduler = torch.optim.lr_scheduler.StepLR(

        optimizer, step_size=7, gamma=0.9

    )

    # Specify the evalutation metrics

    metrics = {'dice': m.mean_dice_coef,

               'dice_target': m.mean_dice_coef_remove_empty}

    train_model(model, criterion, dataloaders,

                optimizer, exp_lr_scheduler, bpath=log_path, metrics=metrics, num_epochs=epochs)

    # Save the trained model

    torch.save(model, os.path.join(log_path, weight_filename))

    print('\n\n ### ===> Training finished sucessfully!\n\n')

if __name__ == '__main__':

    run_nn()