import pandas as pd
import nibabel as nib
from random import sample
from tqdm import tqdm
from Excel_Processing import ProcessSpreadsheets

DATASET_INFO_PATH = "C:\\Users\\leotu\\OneDrive\\Documents\\ImageCHD_dataset\\imageCHD_dataset_info.xlsx"
SCAN_INFO_PATH = "C:\\Users\\leotu\\OneDrive\\Documents\\ImageCHD_dataset\\imagechd_dataset_image_info.xlsx"

dataset_info = ProcessSpreadsheets(DATASET_INFO_PATH, SCAN_INFO_PATH)

try:
    dataset_info.to_csv(path_or_buf = "C:\\Users\\leotu\\OneDrive\\Documents\\ImageCHD_dataset\\patient_info.csv",
                        index = False, mode = 'x')
    print('Patient CSV saved.')
except:
    print('Patient CSV already present. Moving on.')

dataset_info = dataset_info.drop(['ASD', 'VSD', 'AVSD', 'ToF', 'TGA', 'CA',
                                  'PA', 'PDA', 'COUNT', 'PatientSex',
                                  'PatientBirthDate', 'AcquisitionDate'],
                                  axis = 1)

axial_count = [nib.load("C:\\Users\\leotu\\Downloads\\ImageCHD_dataset\\ImageCHD_dataset\\ct_" \
                         + str(x) + "_image.nii.gz") \
               .header['dim'][3] for x in dataset_info['index'].sort_values()]
dataset_info['Axial_count'] = axial_count
train = list()
evaluation = list()
test = list()

for index, row in tqdm(dataset_info.iterrows()):
    temp = row.copy(deep = True)
    ax_c = temp['Axial_count']
    temp = temp.rename({'Axial_count': 'Adjacency_count'})

    temp_sagittal = list()
    temp_coronal = list()
    temp_axial = list()

    for i in range(0, 512):
        temp['Type'] = 'S'
        temp['Indice'] = i
        temp_sagittal.append(temp.copy(deep = True))
    
    for i in range(0, 512):
        temp['Type'] = 'C'
        temp['Indice'] = i
        temp_coronal.append(temp.copy(deep = True))
    
    for i in range(0, ax_c):
        temp['Type'] = 'A'
        temp['Indice'] = i
        temp['Adjacency_count'] = 512
        temp_axial.append(temp.copy(deep = True))
    
    new_train = sample(temp_sagittal, k = round(512*0.7))
    tmp = list()

    good = False
    for sag in temp_sagittal:
        good = True
        for tr in new_train:
            if sag.equals(tr):
                good = False
                break
        if good:
            tmp.append(sag)
    
    temp_sagittal = tmp.copy()
    for tr in new_train:
        train.append(tr)

    new_train = sample(temp_coronal, k = round(512*0.7))
    tmp = list()

    good = False
    for cor in temp_coronal:
        good = True
        for tr in new_train:
            if cor.equals(tr):
                good = False
                break
        if good:
            tmp.append(cor)
    
    temp_coronal = tmp.copy()
    for tr in new_train:
        train.append(tr)

    new_train = sample(temp_axial, k = round(ax_c*0.7))
    tmp = list()

    good = False
    for ax in temp_axial:
        good = True
        for tr in new_train:
            if ax.equals(tr):
                good = False
                break
        if good:
            tmp.append(ax)
    
    temp_axial = tmp.copy()
    for tr in new_train:
        train.append(tr)

    new_eval = sample(temp_sagittal, k = round(512*0.2))
    tmp = list()

    good = False
    for sag in temp_sagittal:
        good = True
        for ev in new_eval:
            if sag.equals(ev):
                good = False
                break
        if good:
            tmp.append(sag)
    
    temp_sagittal = tmp.copy()
    for ev in new_eval:
        evaluation.append(ev)

    new_eval = sample(temp_coronal, k = round(512*0.2))
    tmp = list()

    good = False
    for cor in temp_coronal:
        good = True
        for ev in new_eval:
            if cor.equals(ev):
                good = False
                break
        if good:
            tmp.append(cor)
    
    temp_coronal = tmp.copy()
    for ev in new_eval:
        evaluation.append(ev)

    new_eval = sample(temp_axial, k = round(ax_c*0.2))
    tmp = list()

    good = False
    for ax in temp_axial:
        good = True
        for ev in new_eval:
            if ax.equals(ev):
                good = False
                break
        if good:
            tmp.append(ax)
    
    temp_axial = tmp.copy()
    for ev in new_eval:
        evaluation.append(ev)

    for te in temp_sagittal:
        test.append(te)
    for te in temp_coronal:
        test.append(te)
    for te in temp_axial:
        test.append(te)

train_dataset = pd.DataFrame(train).reset_index().drop('level_0', axis = 1)
eval_dataset = pd.DataFrame(evaluation).reset_index().drop('level_0', axis = 1)
test_dataset = pd.DataFrame(test).reset_index().drop('level_0', axis = 1)

print("Train: ", str(train_dataset.__len__()))
print("Eval: ", str(eval_dataset.__len__()))
print("Test: ", str(test_dataset.__len__()))

try:
    train_dataset.to_csv(path_or_buf = "C:\\Users\\leotu\\OneDrive\\Documents\\ImageCHD_dataset\\train_dataset_info.csv",
                         index = False, mode = 'x')
    print('Train dataset CSV saved.')
except:
    print('Train dataset CSV already present.')

try:
    eval_dataset.to_csv(path_or_buf = "C:\\Users\\leotu\\OneDrive\\Documents\\ImageCHD_dataset\\eval_dataset_info.csv",
                        index = False, mode = 'x')
    print('Eval dataset CSV saved.')
except:
    print('Eval dataset CSV already present.')

try:
    test_dataset.to_csv(path_or_buf = "C:\\Users\\leotu\\OneDrive\\Documents\\ImageCHD_dataset\\test_dataset_info.csv",
                        index = False, mode = 'x')
    print('Test dataset CSV saved.')
except:
    print('Test dataset CSV already present.')