#%%

'''

Copyright (c) Microsoft Corporation. All rights reserved.

Licensed under the MIT License.

'''

import numpy as np 

import pandas as pd 

import SimpleITK as sitk 

import os 

from glob import glob 

import sys 

import argparse

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

sys.path.append(config_dir)

from config import RESULTS_FOLDER

from metrics.metrics import (

    get_3darray_from_niftipath,

    calculate_patient_level_dice_score,

    calculate_patient_level_false_positive_volume,

    calculate_patient_level_false_negative_volume,

    calculate_patient_level_tp_fp_fn

)

def get_spacing_from_niftipath(path):

    image = sitk.ReadImage(path)

    return image.GetSpacing()

def get_column_statistics(col):

    mean = col.mean()

    std = col.std()

    median = col.median()

    quantile25 = col.quantile(q=0.25)

    quantile75 = col.quantile(q=0.75)

    return (mean, std, median, quantile25, quantile75)

def get_prediction_statistics(data_df):

    dsc_stats = get_column_statistics(data_df['DSC'].astype(float))

    fpv_stats = get_column_statistics(data_df['FPV'].astype(float))

    fnv_stats = get_column_statistics(data_df['FNV'].astype(float))

    c1_sensitivity = data_df[f'TP_C1']/(data_df[f'TP_C1'] + data_df[f'FN_C1'])

    c2_sensitivity = data_df[f'TP_C2']/(data_df[f'TP_C2'] + data_df[f'FN_C2'])

    c3_sensitivity = data_df[f'TP_C3']/(data_df[f'TP_C3'] + data_df[f'FN_C3'])

    sens_c1_stats = get_column_statistics(c1_sensitivity)

    sens_c2_stats = get_column_statistics(c2_sensitivity)

    sens_c3_stats = get_column_statistics(c3_sensitivity)

    fp_c1_stats = get_column_statistics(data_df['FP_M1'].astype(float))

    fp_c2_stats = get_column_statistics(data_df['FP_M2'].astype(float))

    fp_c3_stats = get_column_statistics(data_df['FP_M3'].astype(float))

    dsc_stats = [round(d, 2) for d in dsc_stats]

    fpv_stats = [round(d, 2) for d in fpv_stats]

    fnv_stats = [round(d, 2) for d in fnv_stats]

    sens_c1_stats = [round(d, 2) for d in sens_c1_stats]

    sens_c2_stats = [round(d, 2) for d in sens_c2_stats]

    sens_c3_stats = [round(d, 2) for d in sens_c3_stats]

    fp_c1_stats = [round(d, 0) for d in fp_c1_stats]

    fp_c2_stats = [round(d, 0) for d in fp_c2_stats]

    fp_c3_stats = [round(d, 0) for d in fp_c3_stats]

    print(f"DSC (Mean): {dsc_stats[0]} +/- {dsc_stats[1]}")

    print(f"DSC (Median): {dsc_stats[2]} [{dsc_stats[3]}, {dsc_stats[4]}]")

    print(f"FPV (Median): {fpv_stats[2]} [{fpv_stats[3]}, {fpv_stats[4]}]")

    print(f"FNV (Median): {fnv_stats[2]} [{fnv_stats[3]}, {fnv_stats[4]}]")

    print(f"Sensitivity - Criterion1 (Median): {sens_c1_stats[2]} [{sens_c1_stats[3]}, {sens_c1_stats[4]}]")

    print(f"FP - Criterion1 (Median): {fp_c1_stats[2]} [{fp_c1_stats[3]}, {fp_c1_stats[4]}]")

    print(f"Sensitivity - Criterion2 (Median): {sens_c2_stats[2]} [{sens_c2_stats[3]}, {sens_c2_stats[4]}]")

    print(f"FP - Criterion1 (Median): {fp_c2_stats[2]} [{fp_c2_stats[3]}, {fp_c2_stats[4]}]")

    print(f"Sensitivity - Criterion3 (Median): {sens_c3_stats[2]} [{sens_c3_stats[3]}, {sens_c3_stats[4]}]")

    print(f"FP - Criterion3 (Median): {fp_c3_stats[2]} [{fp_c3_stats[3]}, {fp_c3_stats[4]}]")

    print('\n')

#%%

def main(args):

    fold = args.fold

    network = args.network_name

    inputsize = args.input_patch_size

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

    preddir = os.path.join(RESULTS_FOLDER, 'predictions', f'fold{fold}', network, experiment_code)

    predpaths = sorted(glob(os.path.join(preddir, '*.nii.gz')))

    gtpaths = sorted(list(pd.read_csv('./../data_split/test_filepaths.csv')['GTPATH']))

    ptpaths = sorted(list(pd.read_csv('./../data_split/test_filepaths.csv')['PTPATH'])) # PET image paths (ptpaths) for calculating the detection metrics using criterion3 

    imageids = [os.path.basename(path)[:-7] for path in gtpaths]

    TEST_DSCs, TEST_FPVs, TEST_FNVs = [], [], []

    TEST_TP_criterion1, TEST_FP_criterion1, TEST_FN_criterion1 = [], [], []

    TEST_TP_criterion2, TEST_FP_criterion2, TEST_FN_criterion2 = [], [], []

    TEST_TP_criterion3, TEST_FP_criterion3, TEST_FN_criterion3 = [], [], []

    for i in range(len(gtpaths)):

        gtpath = gtpaths[i]

        ptpath = ptpaths[i]

        predpath = predpaths[i]

        gtarray = get_3darray_from_niftipath(gtpath)

        ptarray = get_3darray_from_niftipath(ptpath)

        predarray = get_3darray_from_niftipath(predpath)

        spacing = get_spacing_from_niftipath(gtpath)

        dsc = calculate_patient_level_dice_score(gtarray, predarray)

        fpv = calculate_patient_level_false_positive_volume(gtarray, predarray, spacing)

        fnv = calculate_patient_level_false_negative_volume(gtarray, predarray, spacing)

        tp_c1, fp_c1, fn_c1 = calculate_patient_level_tp_fp_fn(gtarray, predarray, criterion='criterion1')

        tp_c2, fp_c2, fn_c2 = calculate_patient_level_tp_fp_fn(gtarray, predarray, criterion='criterion2', threshold=0.5)

        tp_c3, fp_c3, fn_c3 = calculate_patient_level_tp_fp_fn(gtarray, predarray, criterion='criterion3', ptarray=ptarray)

        TEST_DSCs.append(dsc)

        TEST_FPVs.append(fpv)

        TEST_FNVs.append(fnv)

        TEST_TP_criterion1.append(tp_c1)

        TEST_FP_criterion1.append(fp_c1)

        TEST_FN_criterion1.append(fn_c1)

        TEST_TP_criterion2.append(tp_c2)

        TEST_FP_criterion2.append(fp_c2)

        TEST_FN_criterion2.append(fn_c2)

        TEST_TP_criterion3.append(tp_c3)

        TEST_FP_criterion3.append(fp_c3)

        TEST_FN_criterion3.append(fn_c3)

        print(f"{imageids[i]}: DSC = {round(dsc, 4)}\nFPV = {round(fpv, 4)} ml\nFNV = {round(fnv, 4)} ml")

    save_testmetrics_dir = os.path.join(RESULTS_FOLDER, 'test_metrics', 'fold'+str(fold), network, experiment_code)

    os.makedirs(save_testmetrics_dir, exist_ok=True)

    save_testmetrics_fpath = os.path.join(save_testmetrics_dir, 'testmetrics.csv')

    data = np.column_stack(

        (

            imageids, TEST_DSCs, TEST_FPVs, TEST_FNVs,

            TEST_TP_criterion1, TEST_FP_criterion1, TEST_FN_criterion1,

            TEST_TP_criterion2, TEST_FP_criterion2, TEST_FN_criterion2,

            TEST_TP_criterion3, TEST_FP_criterion3, TEST_FN_criterion3

        )

    )

    column_names = [

        'PatientID', 'DSC', 'FPV', 'FNV',

        'TP_C1', 'FP_C1', 'FN_C1',

        'TP_C2', 'FP_C2', 'FN_C2',

        'TP_C3', 'FP_C3', 'FN_C3',

    ]

    data_df = pd.DataFrame(data=data, columns=column_names)

    data_df.to_csv(save_testmetrics_fpath, index=False)

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)')

    args = parser.parse_args()

    main(args)

# %%