#%%

'''

Copyright (c) Microsoft Corporation. All rights reserved.

Licensed under the MIT License.

'''

import numpy as np 

import glob

import os 

import pandas as pd 

import SimpleITK as sitk

import sys

import argparse

from monai.inferers import sliding_window_inference

from monai.data import DataLoader, Dataset, decollate_batch

import torch

import os

import glob

import pandas as pd

import numpy as np

import torch.nn as nn

import time

from initialize_train import (

    get_validation_sliding_window_size,

    get_model,

    get_test_data_in_dict_format,

    get_valid_transforms,

    get_post_transforms

)

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

sys.path.append(config_dir)

from config import RESULTS_FOLDER

#%%

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 create_dictionary_ctptgt(ctpaths, ptpaths, gtpaths):

    data = []

    for i in range(len(gtpaths)):

        ctpath = ctpaths[i]

        ptpath = ptpaths[i]

        gtpath = gtpaths[i]

        data.append({'CT':ctpath, 'PT':ptpath, 'GT':gtpath})

    return data

def read_image_array(path):

    img =  sitk.ReadImage(path)

    array = np.transpose(sitk.GetArrayFromImage(img), (2,1,0))

    return array

#%%

def main(args):

    # initialize inference

    fold = args.fold

    network = args.network_name

    inputsize = args.input_patch_size

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

    sw_roi_size = get_validation_sliding_window_size(inputsize) # get sliding_window inference size for given input patch size

    # find the best model for this experiment from the training/validation logs

    # best model is the model with the best validation `Metric` (DSC)

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

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

    validlog = pd.read_csv(validlog_fname)

    best_epoch = 2*(np.argmax(validlog['Metric']) + 1)

    best_metric = np.max(validlog['Metric'])

    print(f"Using the {network} model at epoch={best_epoch} with mean valid DSC = {round(best_metric, 4)}")

    # get the best model and push it to device=cuda:0

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

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

    best_model_fname = 'model_ep=' + convert_to_4digits(str(best_epoch)) +'.pth'

    model_path = os.path.join(save_models_dir, best_model_fname)

    device = torch.device(f"cuda:0")

    model = get_model(network, input_patch_size=inputsize)

    model.load_state_dict(torch.load(model_path, map_location=device))

    model.to(device)

    # initialize the location to save predicted masks

    save_preds_dir = os.path.join(RESULTS_FOLDER, f'predictions')

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

    os.makedirs(save_preds_dir, exist_ok=True)

    # get test data (in dictionary format for MONAI dataloader), test_transforms and post_transforms

    test_data = get_test_data_in_dict_format()

    test_transforms = get_valid_transforms()

    post_transforms = get_post_transforms(test_transforms, save_preds_dir)

    # initalize PyTorch dataset and Dataloader

    dataset_test = Dataset(data=test_data, transform=test_transforms)

    dataloader_test = DataLoader(dataset_test, batch_size=1, shuffle=False, num_workers=args.num_workers)

    model.eval()

    with torch.no_grad():

        for data in dataloader_test:

            inputs = data['CTPT'].to(device)

            sw_batch_size = args.sw_bs

            print(sw_batch_size)

            data['Pred'] = sliding_window_inference(inputs, sw_roi_size, sw_batch_size, model)

            data = [post_transforms(i) for i in decollate_batch(data)]

if __name__ == "__main__": 

    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('--input-patch-size', type=int, default=192, metavar='inputsize',

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

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

                        help='num_workers for train and validation dataloaders (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)