import copy

import nibabel as nib

import numpy as np

import os

import tarfile

import json

from sklearn.utils import shuffle

from torch.utils.data import Dataset, DataLoader

import torch

from torch.utils.data import random_split

from config import (

    DATASET_PATH, TASK_ID, TRAIN_VAL_TEST_SPLIT,

    TRAIN_BATCH_SIZE, VAL_BATCH_SIZE, TEST_BATCH_SIZE

)

#Utility function to extract .tar file formats into ./Datasets directory

def ExtractTar(Directory):

        try:

            print("Extracting tar file ...")

            tarfile.open(Directory).extractall('./Datasets')

        except:

            raise "File extraction failed!"

        print("Extraction completed!")

        return 

#The dict representing segmentation tasks along with their IDs

task_names = {

    "01": "BrainTumour",

    "02": "Heart",

    "03": "Liver",

    "04": "Hippocampus",

    "05": "Prostate",

    "06": "Lung",

    "07": "Pancreas",

    "08": "HepaticVessel",

    "09": "Spleen",

    "10": "Colon"

}

class MedicalSegmentationDecathlon(Dataset):

    """

    The base dataset class for Decathlon segmentation tasks

    -- __init__()

    :param task_number -> represent the organ dataset ID (see task_names above for hints)

    :param dir_path -> the dataset directory path to .tar files

    :param transform -> optional - transforms to be applied on each instance

    """

    def __init__(self, task_number, dir_path, split_ratios = [0.8, 0.1, 0.1], transforms = None, mode = None) -> None:

        super(MedicalSegmentationDecathlon, self).__init__()

        #Rectify the task ID representaion

        self.task_number = str(task_number)

        if len(self.task_number) == 1:

            self.task_number = "0" + self.task_number

        #Building the file name according to task ID

        self.file_name = f"Task{self.task_number}_{task_names[self.task_number]}"

        #Extracting .tar file

        if not os.path.exists(os.path.join(os.getcwd(), "Datasets", self.file_name)):

            ExtractTar(os.path.join(dir_path, f"{self.file_name}.tar"))

        #Path to extracted dataset

        self.dir = os.path.join(os.getcwd(), "Datasets", self.file_name)

        #Meta data about the dataset

        self.meta = json.load(open(os.path.join(self.dir, "dataset.json")))

        self.splits = split_ratios

        self.transform = transforms

        #Calculating split number of images

        num_training_imgs =  self.meta["numTraining"]

        train_val_test = [int(x * num_training_imgs) for x in split_ratios]

        if(sum(train_val_test) != num_training_imgs): train_val_test[0] += (num_training_imgs - sum(train_val_test))

        train_val_test = [x for x in train_val_test if x!=0]

        # train_val_test = [(x-1) for x in train_val_test]

        self.mode = mode

        #Spliting dataset

        samples = self.meta["training"]

        shuffle(samples)

        self.train = samples[0:train_val_test[0]]

        self.val = samples[train_val_test[0]:train_val_test[0] + train_val_test[1]]

        self.test = samples[train_val_test[1]:train_val_test[1] + train_val_test[2]]

    def set_mode(self, mode):

        self.mode = mode

    def __len__(self):

        if self.mode == "train":

            return len(self.train)

        elif self.mode == "val":

            return len(self.val)

        elif self.mode == "test":

            return len(self.test)

        return self.meta["numTraining"]

    def __getitem__(self, idx):

        if torch.is_tensor(idx):

            idx = idx.tolist()

        #Obtaining image name by given index and the mode using meta data

        if self.mode == "train":

            name = self.train[idx]['image'].split('/')[-1]

        elif self.mode == "val":

            name = self.val[idx]['image'].split('/')[-1]

        elif self.mode == "test":

            name = self.test[idx]['image'].split('/')[-1]

        else:

            name = self.meta["training"][idx]['image'].split('/')[-1]

        img_path = os.path.join(self.dir, "imagesTr", name)

        label_path = os.path.join(self.dir, "labelsTr", name)

        img_object = nib.load(img_path)

        label_object = nib.load(label_path)

        img_array = img_object.get_fdata()

        #Converting to channel-first numpy array

        img_array = np.moveaxis(img_array, -1, 0)

        label_array = label_object.get_fdata()

        label_array = np.moveaxis(label_array, -1, 0)

        proccessed_out = {'name': name, 'image': img_array, 'label': label_array} 

        if self.transform:

            if self.mode == "train":

                proccessed_out = self.transform[0](proccessed_out)

            elif self.mode == "val":

                proccessed_out = self.transform[1](proccessed_out)

            elif self.mode == "test":

                proccessed_out = self.transform[2](proccessed_out)

            else:

                proccessed_out = self.transform(proccessed_out)

        #The output numpy array is in channel-first format

        return proccessed_out

def get_train_val_test_Dataloaders(train_transforms, val_transforms, test_transforms):

    """

    The utility function to generate splitted train, validation and test dataloaders

    Note: all the configs to generate dataloaders in included in "config.py"

    """

    dataset = MedicalSegmentationDecathlon(task_number=TASK_ID, dir_path=DATASET_PATH, split_ratios=TRAIN_VAL_TEST_SPLIT, transforms=[train_transforms, val_transforms, test_transforms])

    #Spliting dataset and building their respective DataLoaders

    train_set, val_set, test_set = copy.deepcopy(dataset), copy.deepcopy(dataset), copy.deepcopy(dataset)

    train_set.set_mode('train')

    val_set.set_mode('val')

    test_set.set_mode('test')

    train_dataloader = DataLoader(dataset= train_set, batch_size= TRAIN_BATCH_SIZE, shuffle= False)

    val_dataloader = DataLoader(dataset= val_set, batch_size= VAL_BATCH_SIZE, shuffle= False)

    test_dataloader = DataLoader(dataset= test_set, batch_size= TEST_BATCH_SIZE, shuffle= False)

    return train_dataloader, val_dataloader, test_dataloader