#%%
'''
Copyright (c) Microsoft Corporation. All rights reserved.
Licensed under the MIT License.
'''
import pandas as pd 
import numpy as np
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 *

def get_spacing_from_niftipath(path):
    spacing = sitk.ReadImage(path).GetSpacing()
    return spacing


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]
    DSC = [] 
    SUVmean_orig, SUVmean_pred = [], []
    SUVmax_orig, SUVmax_pred = [], [] 
    LesionCount_orig, LesionCount_pred = [], [] 
    TMTV_orig, TMTV_pred = [], []
    TLG_orig, TLG_pred = [], []
    Dmax_orig, Dmax_pred = [], []
    
    for i in range(len(gtpaths)):
        ptpath = ptpaths[i]
        gtpath = gtpaths[i]
        predpath = predpaths[i]
        
        ptarray = get_3darray_from_niftipath(ptpath)
        gtarray = get_3darray_from_niftipath(gtpath)
        predarray = get_3darray_from_niftipath(predpath)
        spacing = get_spacing_from_niftipath(gtpath)

        # Dice score between mask gt and pred
        dsc = calculate_patient_level_dice_score(gtarray, predarray)
        # Lesion SUVmean
        suvmean_orig = calculate_patient_level_lesion_suvmean_suvmax(ptarray, gtarray, marker='SUVmean')
        suvmean_pred = calculate_patient_level_lesion_suvmean_suvmax(ptarray, predarray, marker='SUVmean')
        # Lesion SUVmax
        suvmax_orig = calculate_patient_level_lesion_suvmean_suvmax(ptarray, gtarray, marker='SUVmax')
        suvmax_pred = calculate_patient_level_lesion_suvmean_suvmax(ptarray, predarray, marker='SUVmax')
        # Lesion Count 
        lesioncount_orig = calculate_patient_level_lesion_count(gtarray)
        lesioncount_pred = calculate_patient_level_lesion_count(predarray)
        # TMTV
        tmtv_orig = calculate_patient_level_tmtv(gtarray, spacing)
        tmtv_pred = calculate_patient_level_tmtv(predarray, spacing)
        # TLG
        tlg_orig = calculate_patient_level_tlg(ptarray, gtarray, spacing)
        tlg_pred = calculate_patient_level_tlg(ptarray, predarray, spacing)
        # Dmax
        dmax_orig = calculate_patient_level_dissemination(gtarray, spacing)
        dmax_pred = calculate_patient_level_dissemination(predarray, spacing)
        
        DSC.append(dsc)
        SUVmean_orig.append(suvmean_orig)
        SUVmean_pred.append(suvmean_pred)
        SUVmax_orig.append(suvmax_orig)
        SUVmax_pred.append(suvmax_pred)
        LesionCount_orig.append(lesioncount_orig)
        LesionCount_pred.append(lesioncount_pred)
        TMTV_orig.append(tmtv_orig)
        TMTV_pred.append(tmtv_pred)
        TLG_orig.append(tlg_orig)
        TLG_pred.append(tlg_pred)
        Dmax_orig.append(dmax_orig)
        Dmax_pred.append(dmax_pred)
        
        
        print(f"{i}: {imageids[i]}")
        print(f"Dice Score: {round(dsc,4)}")
        print(f"SUVmean: GT: {suvmean_orig}, Pred: {suvmean_pred}")
        print(f"SUVmax: GT: {suvmax_orig}, Pred: {suvmax_pred}")
        print(f"LesionCount: GT: {lesioncount_orig}, Pred: {lesioncount_pred}")
        print(f"TMTV: GT: {tmtv_orig} ml, Pred: {tmtv_pred} ml")
        print(f"TLG: GT: {tlg_orig} ml, Pred: {tlg_pred} ml")
        print(f"Dmax: GT: {dmax_orig} cm, Pred: {dmax_pred} cm")
        print("\n")

    save_lesionmeasures_dir = os.path.join(RESULTS_FOLDER, f'test_lesion_measures', 'fold'+str(fold), network, experiment_code)
    os.makedirs(save_lesionmeasures_dir, exist_ok=True)
    filepath = os.path.join(save_lesionmeasures_dir, f'testlesionmeasures.csv')
    
    data = np.column_stack(
            [
                imageids,
                DSC,
                SUVmean_orig,
                SUVmean_pred,
                SUVmax_orig,
                SUVmax_pred,
                LesionCount_orig,
                LesionCount_pred,
                TMTV_orig,
                TMTV_pred,
                TLG_orig,
                TLG_pred,
                Dmax_orig,
                Dmax_pred
            ]
        )

    data_df = pd.DataFrame(
        data=data,
        columns=[
            'PatientID',
            'DSC',
            'SUVmean_orig',
            'SUVmean_pred',
            'SUVmax_orig',
            'SUVmax_pred',
            'LesionCount_orig',
            'LesionCount_pred',
            'TMTV_orig',
            'TMTV_pred',
            'TLG_orig',
            'TLG_pred',
            'Dmax_orig',
            'Dmax_pred'
        ]
    )
    data_df.to_csv(filepath, 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)