'''

Copyright (c) Microsoft Corporation. All rights reserved.

Licensed under the MIT License.

'''

from monai.transforms import (

    AsDiscrete,

    Compose,

)

import argparse

from monai.inferers import sliding_window_inference

from monai.data import CacheDataset, DataLoader, decollate_batch

import torch

import matplotlib.pyplot as plt

import os

import pandas as pd

import time

from torch.utils.data.distributed import DistributedSampler

from torch.nn.parallel import DistributedDataParallel as DDP

import torch.distributed as dist 

import os

from initialize_train import (

    create_data_split_files,

    get_train_valid_data_in_dict_format, 

    get_train_transforms, 

    get_valid_transforms, 

    get_model, 

    get_loss_function,

    get_optimizer, 

    get_scheduler,

    get_metric,

    get_validation_sliding_window_size

)

import sys

config_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), "..")

sys.path.append(config_dir)

from config import RESULTS_FOLDER

torch.backends.cudnn.benchmark = True

#%%

def ddp_setup():

    dist.init_process_group(backend='nccl', init_method="env://")

def convert_to_4digits(str_num):

    if len(str_num) == 1:

        new_num = '000' + str_num

    elif len(str_num) == 2:

        new_num = '00' + str_num

    elif len(str_num) == 3:

        new_num = '0' + str_num

    else:

        new_num = str_num

    return new_num

#%%

def load_train_objects(args):

    train_data, valid_data = get_train_valid_data_in_dict_format(args.fold) 

    train_transforms = get_train_transforms(args.input_patch_size)

    valid_transforms = get_valid_transforms()

    model = get_model(args.network_name, args.input_patch_size) 

    optimizer = get_optimizer(model, learning_rate=args.lr, weight_decay=args.wd)

    loss_function = get_loss_function()

    scheduler = get_scheduler(optimizer, args.epochs)

    metric = get_metric()

    return (

        train_data,

        valid_data,

        train_transforms,

        valid_transforms,

        model,

        loss_function,

        optimizer,

        scheduler,

        metric

    )

def prepare_dataset(data, transforms, args):

    dataset = CacheDataset(data=data, transform=transforms, cache_rate=args.cache_rate, num_workers=args.num_workers)

    return dataset

def main_worker(save_models_dir, save_logs_dir, args):

    # init_process_group

    ddp_setup() 

    # get local rank on the GPU

    local_rank = int(dist.get_rank())

    if local_rank == 0:

        print(f"Training {args.network_name} on fold {args.fold}")

        print(f"The models will be saved in {save_models_dir}")

        print(f"The training/validation logs will be saved in {save_logs_dir}")

    # get all training and validation objects

    train_data, valid_data, train_transforms, valid_transforms, model, loss_function, optimizer, scheduler, metric = load_train_objects(args)

    # get dataset of object-type CacheDataset 

    train_dataset = prepare_dataset(train_data, train_transforms, args)

    valid_dataset = prepare_dataset(valid_data, valid_transforms, args)

    # get DistributedSampler instances for both training and validation dataloader

    # this will be used to split data into different GPUs

    train_sampler = DistributedSampler(dataset=train_dataset, shuffle=True)

    valid_sampler = DistributedSampler(dataset=valid_dataset, shuffle=False)

    # initializing train and valid dataloaders

    train_dataloader = DataLoader(

        train_dataset,

        batch_size=args.train_bs,

        pin_memory=True,

        shuffle=False,

        sampler=train_sampler,

        num_workers=args.num_workers

    )

    valid_dataloader = DataLoader(

        valid_dataset,

        batch_size=1,

        pin_memory=True,

        shuffle=False,

        sampler=valid_sampler,

        num_workers=args.num_workers

    )

    post_pred = Compose([AsDiscrete(argmax=True, to_onehot=2)])

    post_label = Compose([AsDiscrete(to_onehot=2)])

    # filepaths for storing training and validation logs from different GPUs

    trainlog_fpath = os.path.join(save_logs_dir, f'trainlog_gpu{local_rank}.csv')

    validlog_fpath = os.path.join(save_logs_dir, f'validlog_gpu{local_rank}.csv')

    # initialize the GPU device    

    device = torch.device(f"cuda:{local_rank}")

    torch.cuda.set_device(device)

    # number of epochs and epoch interval for running validation

    max_epochs = args.epochs

    val_interval = args.val_interval

    # push models to device

    model = model.to(device)

    epoch_loss_values = []

    metric_values = []

    # wrap the model with DDP

    model = DDP(model, device_ids=[device])

    experiment_start_time = time.time()

    for epoch in range(max_epochs):

        epoch_start_time = time.time()

        print(f"[GPU{local_rank}]: Running training: epoch = {epoch + 1}")

        model.train()

        epoch_loss = 0

        step = 0

        train_sampler.set_epoch(epoch)

        for batch_data in train_dataloader:

            step += 1

            inputs, labels = (

                batch_data['CTPT'].to(device),

                batch_data['GT'].to(device),

            )

            optimizer.zero_grad()

            outputs = model(inputs)

            loss = loss_function(outputs, labels)

            loss.backward()

            optimizer.step()

            epoch_loss += loss.item()

        epoch_loss /= step

        print(f"[GPU:{local_rank}]: epoch {epoch + 1}/{max_epochs}: average loss: {epoch_loss:.4f}")

        epoch_loss_values.append(epoch_loss)

        # steps forward the CosineAnnealingLR scheduler

        scheduler.step()

        # update the training log file

        epoch_loss_values_df = pd.DataFrame(data=epoch_loss_values, columns=['Loss'])

        epoch_loss_values_df.to_csv(trainlog_fpath, index=False)

        if (epoch + 1) % val_interval == 0:

            print(f"[GPU{local_rank}]: Running validation")

            model.eval()

            with torch.no_grad():

                for val_data in valid_dataloader:

                    val_inputs, val_labels = (

                        val_data['CTPT'].to(device),

                        val_data['GT'].to(device),

                    )

                    roi_size = get_validation_sliding_window_size(args.input_patch_size) 

                    sw_batch_size = args.sw_bs

                    val_outputs = sliding_window_inference(

                        val_inputs, roi_size, sw_batch_size, model)

                    val_outputs = [post_pred(i) for i in decollate_batch(val_outputs)]

                    val_labels = [post_label(i) for i in decollate_batch(val_labels)]

                    # compute metric for current iteration

                    metric(y_pred=val_outputs, y=val_labels)

                # aggregate the final mean dice result

                metric_val = metric.aggregate().item()

                metric.reset()

                metric_values.append(metric_val)

                metric_values_df = pd.DataFrame(data=metric_values, columns=['Metric'])

                metric_values_df.to_csv(validlog_fpath, index=False)

                print(f"[GPU:{local_rank}] SAVING MODEL at epoch: {epoch + 1}; Mean DSC: {metric_val:.4f}")

                savepath = os.path.join(save_models_dir, "model_ep="+convert_to_4digits(str(int(epoch + 1)))+".pth")

                torch.save(model.module.state_dict(), savepath)

        epoch_end_time = (time.time() - epoch_start_time)/60

        print(f"[GPU:{local_rank}]: Epoch {epoch + 1} time: {round(epoch_end_time,2)} min")

    experiment_end_time = (time.time() - experiment_start_time)/(60*60)

    print(f"[GPU:{local_rank}]: Total time: {round(experiment_end_time,2)} hr")

    dist.destroy_process_group()

def main(args):

    os.environ['OMP_NUM_THREADS'] = '6'

    fold = args.fold

    network = args.network_name

    inputsize = f'randcrop{args.input_patch_size}'

    experiment_code = f"{network}_fold{fold}_{inputsize}"

    #save models folder

    save_models_dir = os.path.join(RESULTS_FOLDER,'models')

    save_models_dir = os.path.join(save_models_dir, 'fold'+str(fold), network, experiment_code)

    os.makedirs(save_models_dir, exist_ok=True)

    # save train and valid logs folder

    save_logs_dir = os.path.join(RESULTS_FOLDER,'logs')

    save_logs_dir = os.path.join(save_logs_dir, 'fold'+str(fold), network, experiment_code)

    os.makedirs(save_logs_dir, exist_ok=True)

    main_worker(save_models_dir, save_logs_dir, args)

if __name__ == "__main__": 

    # create datasplit files for train and test images

    # follow all the instructions for dataset directory creation and images/labels file names as given in: LINK

    create_data_split_files() 

    parser = argparse.ArgumentParser(description='Lymphoma PET/CT lesion segmentation using MONAI-PyTorch')

    parser.add_argument('--fold', type=int, default=0, metavar='fold',

                        help='validation fold (default: 0), remaining folds will be used for training')

    parser.add_argument('--network-name', type=str, default='unet', metavar='netname',

                        help='network name for training (default: unet)')

    parser.add_argument('--epochs', type=int, default=500, metavar='epochs',

                        help='number of epochs to train (default: 10)')

    parser.add_argument('--input-patch-size', type=int, default=192, metavar='inputsize',

                        help='size of cropped input patch for training (default: 192)')

    parser.add_argument('--train-bs', type=int, default=1, metavar='train-bs',

                        help='mini-batchsize for training (default: 1)')

    parser.add_argument('--num_workers', type=int, default=2, metavar='nw',

                        help='num_workers for train and validation dataloaders (default: 2)')

    parser.add_argument('--cache-rate', type=float, default=0.1, metavar='cr',

                        help='cache_rate for CacheDataset from MONAI (default=0.1)')

    parser.add_argument('--lr', type=float, default=2e-4, metavar='lr',

                        help='initial learning rate for AdamW optimizer (default=2e-4); Cosine scheduler will decrease this to 0 in args.epochs epochs')

    parser.add_argument('--wd', type=float, default=1e-5, metavar='wd',

                        help='weight-decay for AdamW optimizer (default=1e-5)')

    parser.add_argument('--val-interval', type=int, default=2, metavar='val-interval',

                        help='epochs interval for which validation will be performed (default=2)')

    parser.add_argument('--sw-bs', type=int, default=2, metavar='sw-bs',

                        help='batchsize for sliding window inference (default=2)')

    args = parser.parse_args()

    main(args)