# NEEDS TO BE FINISHED
# TODO CHECK ABOUT IMAGE DIMENSIONS
# TODO ENSURE ITERATION WORKS
import glob
import nibabel as nib
import numpy as np
import random
import torch
from torch.utils.data import Dataset


'''
Prepares CustomDatasets for training, validating, and testing CNN
'''
def prepare_datasets(mri_dir, val_split=0.2, seed=50):

    rndm = random.Random(seed)

    raw_data = glob.glob(mri_dir + "*")
    AD_list = []
    NL_list = []

    print("--- DATA INFO ---")
    print("Amount of images: " + str(len(raw_data)))

    # TODO Check that image is in CSV?
    for image in raw_data:
        if "NL" in image:
            NL_list.append(image)
        elif "AD" in image:
            AD_list.append(image)

    print("Total AD: " + str(len(AD_list)))
    print("Total NL: " + str(len(NL_list)))

    rndm.shuffle(AD_list)
    rndm.shuffle(NL_list)

    train_list, val_list, test_list = get_train_val_test(AD_list, NL_list, val_split)

    rndm.shuffle(train_list)
    rndm.shuffle(val_list)
    rndm.shuffle(test_list)

    train_dataset = CustomDataset(train_list)
    val_dataset = CustomDataset(val_list)
    test_dataset = CustomDataset(test_list)

    return train_dataset, val_dataset, test_dataset

    # TODO  Normalize data? Later add / Exctract clinical data? Which data?



'''
Returns train_list, val_list and test_list in format [(image, id), ...] each
'''
def get_train_val_test(AD_list, NL_list, val_split):

    train_list, val_list, test_list = [], [], []

    num_test_ad = int(len(AD_list) * val_split)
    num_test_nl = int(len(NL_list) * val_split)

    num_val_ad = int((len(AD_list) - num_test_ad) * val_split)
    num_val_nl = int((len(NL_list) - num_test_nl) * val_split)

    # Sets up ADs
    for image in AD_list[0:num_val_ad]:
        val_list.append((image, 1))

    for image in AD_list[num_val_ad:num_test_ad]:
        test_list.append((image, 1))

    for image in AD_list[num_test_ad:]:
        train_list.append((image, 1))

    # Sets up NLs
    for image in NL_list[0:num_val_nl]:
        val_list.append((image, 0))

    for image in NL_list[num_val_nl:num_test_nl]:
        test_list.append((image, 0))

    for image in NL_list[num_test_nl:]:
        train_list.append((image, 0))

    return train_list, val_list, test_list


class CustomDataset(Dataset):
    def __init__(self, list):
        self.data = list        # DATA IS A LIST WITH TUPLES (image_dir, class_id)

    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):     # RETURNS TUPLE WITH IMAGE AND CLASS_ID, BASED ON INDEX IDX
        mri_path, class_id = self.data[idx]
        mri = nib.load(mri_path)
        mri_data = mri.get_fdata()
        # mri_array = np.array(mri)
        # mri_tensor = torch.from_numpy(mri_array)
        # class_id = torch.tensor([class_id]) TODO return tensor or just id (0, 1)??
        return mri_data, class_id