import os

import h5py
import nibabel as nb
import numpy as np
import torch
import torch.utils.data as data
from torchvision import transforms
import utils.preprocessor as preprocessor


# transform_train = transforms.Compose([
#     transforms.RandomCrop(200, padding=56),
#     transforms.ToTensor(),
# ])


class ImdbData(data.Dataset):
    def __init__(self, X, y, w, transforms=None):
        self.X = X if len(X.shape) == 4 else X[:, np.newaxis, :, :]
        self.y = y
        self.w = w
        self.transforms = transforms

    def __getitem__(self, index):
        img = torch.from_numpy(self.X[index])
        label = torch.from_numpy(self.y[index])
        weight = torch.from_numpy(self.w[index])
        return img, label, weight

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


def get_imdb_dataset(data_params):
    data_train = h5py.File(os.path.join(data_params['data_dir'], data_params['train_data_file']), 'r')
    label_train = h5py.File(os.path.join(data_params['data_dir'], data_params['train_label_file']), 'r')
    class_weight_train = h5py.File(os.path.join(data_params['data_dir'], data_params['train_class_weights_file']), 'r')
    weight_train = h5py.File(os.path.join(data_params['data_dir'], data_params['train_weights_file']), 'r')

    data_test = h5py.File(os.path.join(data_params['data_dir'], data_params['test_data_file']), 'r')
    label_test = h5py.File(os.path.join(data_params['data_dir'], data_params['test_label_file']), 'r')
    class_weight_test = h5py.File(os.path.join(data_params['data_dir'], data_params['test_class_weights_file']), 'r')
    weight_test = h5py.File(os.path.join(data_params['data_dir'], data_params['test_weights_file']), 'r')

    return (ImdbData(data_train['data'][()], label_train['label'][()], class_weight_train['class_weights'][()]),
            ImdbData(data_test['data'][()], label_test['label'][()], class_weight_test['class_weights'][()]))


def load_dataset(file_paths,
                 orientation,
                 remap_config,
                 return_weights=False,
                 reduce_slices=False,
                 remove_black=False):
    print("Loading and preprocessing data...")
    volume_list, labelmap_list, headers, class_weights_list, weights_list = [], [], [], [], []

    for file_path in file_paths:
        volume, labelmap, class_weights, weights, header = load_and_preprocess(file_path, orientation,
                                                                               remap_config=remap_config,
                                                                               reduce_slices=reduce_slices,
                                                                               remove_black=remove_black,
                                                                               return_weights=return_weights)

        volume_list.append(volume)
        labelmap_list.append(labelmap)

        if return_weights:
            class_weights_list.append(class_weights)
            weights_list.append(weights)

        headers.append(header)

        print("#", end='', flush=True)
    print("100%", flush=True)
    if return_weights:
        return volume_list, labelmap_list, class_weights_list, weights_list, headers
    else:
        return volume_list, labelmap_list, headers


def load_and_preprocess(file_path, orientation, remap_config, reduce_slices=False,
                        remove_black=False,
                        return_weights=False):
    volume, labelmap, header = load_data(file_path, orientation)

    volume, labelmap, class_weights, weights = preprocess(volume, labelmap, remap_config=remap_config,
                                                          reduce_slices=reduce_slices,
                                                          remove_black=remove_black,
                                                          return_weights=return_weights)
    return volume, labelmap, class_weights, weights, header


def load_and_preprocess_eval(file_path, orientation, notlabel=True):
    volume_nifty = nb.load(file_path[0])
    header = volume_nifty.header
    volume = volume_nifty.get_fdata()
    if notlabel:
        volume = (volume - np.min(volume)) / (np.max(volume) - np.min(volume))
    else:
        volume = np.round(volume)
    if orientation == "COR":
        volume = volume.transpose((2, 0, 1))
    elif orientation == "AXI":
        volume = volume.transpose((1, 2, 0))
    return volume, header


def load_data(file_path, orientation):
    volume_nifty, labelmap_nifty = nb.load(file_path[0]), nb.load(file_path[1])
    volume, labelmap = volume_nifty.get_fdata(), labelmap_nifty.get_fdata()
    volume = (volume - np.min(volume)) / (np.max(volume) - np.min(volume))
    volume, labelmap = preprocessor.rotate_orientation(volume, labelmap, orientation)
    return volume, labelmap, volume_nifty.header


def preprocess(volume, labelmap, remap_config, reduce_slices=False, remove_black=False, return_weights=False):
    if reduce_slices:
        volume, labelmap = preprocessor.reduce_slices(volume, labelmap)

    if remap_config:
        labelmap = preprocessor.remap_labels(labelmap, remap_config)

    if remove_black:
        volume, labelmap = preprocessor.remove_black(volume, labelmap)

    if return_weights:
        class_weights, weights = preprocessor.estimate_weights_mfb(labelmap)
        return volume, labelmap, class_weights, weights
    else:
        return volume, labelmap, None, None


# def load_file_paths(data_dir, label_dir, volumes_txt_file=None):
#     """
#     This function returns the file paths combined as a list where each element is a 2 element tuple, 0th being data and 1st being label.
#     It should be modified to suit the need of the project
#     :param data_dir: Directory which contains the data files
#     :param label_dir: Directory which contains the label files
#     :param volumes_txt_file: (Optional) Path to the a csv file, when provided only these data points will be read
#     :return: list of file paths as string
#     """
#
#     volume_exclude_list = ['IXI290', 'IXI423']
#     if volumes_txt_file:
#         with open(volumes_txt_file) as file_handle:
#             volumes_to_use = file_handle.read().splitlines()
#     else:
#         volumes_to_use = [name for name in os.listdir(data_dir) if
#                           name.startswith('IXI') and name not in volume_exclude_list]
#
#     file_paths = [
#         [os.path.join(data_dir, vol, 'mri/orig.mgz'), os.path.join(label_dir, vol, 'mri/aseg.auto_noCCseg.mgz')]
#         for
#         vol in volumes_to_use]
#     return file_paths


def load_file_paths(data_dir, label_dir, data_id, volumes_txt_file=None):
    """
    This function returns the file paths combined as a list where each element is a 2 element tuple, 0th being data and 1st being label.
    It should be modified to suit the need of the project
    :param data_dir: Directory which contains the data files
    :param label_dir: Directory which contains the label files
    :param data_id: A flag indicates the name of Dataset for proper file reading
    :param volumes_txt_file: (Optional) Path to the a csv file, when provided only these data points will be read
    :return: list of file paths as string
    """

    if volumes_txt_file:
        with open(volumes_txt_file) as file_handle:
            volumes_to_use = file_handle.read().splitlines()
    else:
        volumes_to_use = [name for name in os.listdir(data_dir)]

    if data_id == "MALC":
        file_paths = [
            [os.path.join(data_dir, vol, 'mri/orig.mgz'), os.path.join(label_dir, vol + '_glm.mgz')]
            for
            vol in volumes_to_use]
    elif data_id == "ADNI":
        file_paths = [
            [os.path.join(data_dir, vol, 'orig.mgz'), os.path.join(label_dir, vol, 'Lab_con.mgz')]
            for
            vol in volumes_to_use]
    elif data_id == "CANDI":
        file_paths = [
            [os.path.join(data_dir, vol + '/' + vol + '_1.mgz'),
             os.path.join(label_dir, vol + '/' + vol + '_1_seg.mgz')]
            for
            vol in volumes_to_use]
    elif data_id == "IBSR":
        file_paths = [
            [os.path.join(data_dir, vol, 'mri/orig.mgz'), os.path.join(label_dir, vol + '_map.nii.gz')]
            for
            vol in volumes_to_use]
    elif data_id == "BORIS": #BORIS
        file_paths = [
            [os.path.join(data_dir, vol), os.path.join(label_dir, vol.replace('.nii', '_seg.nii'))]
            for
            vol in volumes_to_use]
    else:
        raise ValueError("Invalid entry, valid options are MALC, ADNI, CANDI and IBSR")

    return file_paths


def load_file_paths_eval(data_dir, volumes_txt_file, dir_struct):
    """
    This function returns the file paths combined as a list where each element is a 2 element tuple, 0th being data and 1st being label.
    It should be modified to suit the need of the project
    :param data_dir: Directory which contains the data files
    :param volumes_txt_file:  Path to the a csv file, when provided only these data points will be read
    :param dir_struct: If the id_list is in FreeSurfer style or normal
    :return: list of file paths as string
    """

    with open(volumes_txt_file) as file_handle:
        volumes_to_use = file_handle.read().splitlines()
    if dir_struct == "FS":
        file_paths = [
            [os.path.join(data_dir, vol, 'mri/orig.mgz')]
            for
            vol in volumes_to_use]
    elif dir_struct == "Linear":
        file_paths = [
            [os.path.join(data_dir, vol)]
            for
            vol in volumes_to_use]
    elif dir_struct == "part_FS":
        file_paths = [
            [os.path.join(data_dir, vol, 'orig.mgz')]
            for
            vol in volumes_to_use]
    else:
        raise ValueError("Invalid entry, valid options are FS and Linear")
    return file_paths

[((96, 19, 96, 40), 'nibabel.load', 'nb.load', ({(96, 27, 96, 39): 'file_path[0]'}, {}), '(file_path[0])', True, 'import nibabel as nb\n'), ((114, 23, 114, 85), 'utils.preprocessor.rotate_orientation', 'preprocessor.rotate_orientation', ({(114, 55, 114, 61): 'volume', (114, 63, 114, 71): 'labelmap', (114, 73, 114, 84): 'orientation'}, {}), '(volume, labelmap, orientation)', True, 'import utils.preprocessor as preprocessor\n'), ((26, 14, 26, 45), 'torch.from_numpy', 'torch.from_numpy', ({(26, 31, 26, 44): 'self.X[index]'}, {}), '(self.X[index])', False, 'import torch\n'), ((27, 16, 27, 47), 'torch.from_numpy', 'torch.from_numpy', ({(27, 33, 27, 46): 'self.y[index]'}, {}), '(self.y[index])', False, 'import torch\n'), ((28, 17, 28, 48), 'torch.from_numpy', 'torch.from_numpy', ({(28, 34, 28, 47): 'self.w[index]'}, {}), '(self.w[index])', False, 'import torch\n'), ((36, 27, 36, 96), 'os.path.join', 'os.path.join', ({(36, 40, 36, 63): "data_params['data_dir']", (36, 65, 36, 95): "data_params['train_data_file']"}, {}), "(data_params['data_dir'], data_params['train_data_file'])", False, 'import os\n'), ((37, 28, 37, 98), 'os.path.join', 'os.path.join', ({(37, 41, 37, 64): "data_params['data_dir']", (37, 66, 37, 97): "data_params['train_label_file']"}, {}), "(data_params['data_dir'], data_params['train_label_file'])", False, 'import os\n'), ((38, 35, 38, 113), 'os.path.join', 'os.path.join', ({(38, 48, 38, 71): "data_params['data_dir']", (38, 73, 38, 112): "data_params['train_class_weights_file']"}, {}), "(data_params['data_dir'], data_params['train_class_weights_file'])", False, 'import os\n'), ((39, 29, 39, 101), 'os.path.join', 'os.path.join', ({(39, 42, 39, 65): "data_params['data_dir']", (39, 67, 39, 100): "data_params['train_weights_file']"}, {}), "(data_params['data_dir'], data_params['train_weights_file'])", False, 'import os\n'), ((41, 26, 41, 94), 'os.path.join', 'os.path.join', ({(41, 39, 41, 62): "data_params['data_dir']", (41, 64, 41, 93): "data_params['test_data_file']"}, {}), "(data_params['data_dir'], data_params['test_data_file'])", False, 'import os\n'), ((42, 27, 42, 96), 'os.path.join', 'os.path.join', ({(42, 40, 42, 63): "data_params['data_dir']", (42, 65, 42, 95): "data_params['test_label_file']"}, {}), "(data_params['data_dir'], data_params['test_label_file'])", False, 'import os\n'), ((43, 34, 43, 111), 'os.path.join', 'os.path.join', ({(43, 47, 43, 70): "data_params['data_dir']", (43, 72, 43, 110): "data_params['test_class_weights_file']"}, {}), "(data_params['data_dir'], data_params['test_class_weights_file'])", False, 'import os\n'), ((44, 28, 44, 99), 'os.path.join', 'os.path.join', ({(44, 41, 44, 64): "data_params['data_dir']", (44, 66, 44, 98): "data_params['test_weights_file']"}, {}), "(data_params['data_dir'], data_params['test_weights_file'])", False, 'import os\n'), ((102, 17, 102, 33), 'numpy.round', 'np.round', ({(102, 26, 102, 32): 'volume'}, {}), '(volume)', True, 'import numpy as np\n'), ((111, 35, 111, 56), 'nibabel.load', 'nb.load', ({(111, 43, 111, 55): 'file_path[0]'}, {}), '(file_path[0])', True, 'import nibabel as nb\n'), ((111, 58, 111, 79), 'nibabel.load', 'nb.load', ({(111, 66, 111, 78): 'file_path[1]'}, {}), '(file_path[1])', True, 'import nibabel as nb\n'), ((120, 27, 120, 71), 'utils.preprocessor.reduce_slices', 'preprocessor.reduce_slices', ({(120, 54, 120, 60): 'volume', (120, 62, 120, 70): 'labelmap'}, {}), '(volume, labelmap)', True, 'import utils.preprocessor as preprocessor\n'), ((123, 19, 123, 68), 'utils.preprocessor.remap_labels', 'preprocessor.remap_labels', ({(123, 45, 123, 53): 'labelmap', (123, 55, 123, 67): 'remap_config'}, {}), '(labelmap, remap_config)', True, 'import utils.preprocessor as preprocessor\n'), ((126, 27, 126, 70), 'utils.preprocessor.remove_black', 'preprocessor.remove_black', ({(126, 53, 126, 59): 'volume', (126, 61, 126, 69): 'labelmap'}, {}), '(volume, labelmap)', True, 'import utils.preprocessor as preprocessor\n'), ((129, 33, 129, 76), 'utils.preprocessor.estimate_weights_mfb', 'preprocessor.estimate_weights_mfb', ({(129, 67, 129, 75): 'labelmap'}, {}), '(labelmap)', True, 'import utils.preprocessor as preprocessor\n'), ((113, 23, 113, 37), 'numpy.min', 'np.min', ({(113, 30, 113, 36): 'volume'}, {}), '(volume)', True, 'import numpy as np\n'), ((113, 42, 113, 56), 'numpy.max', 'np.max', ({(113, 49, 113, 55): 'volume'}, {}), '(volume)', True, 'import numpy as np\n'), ((113, 59, 113, 73), 'numpy.min', 'np.min', ({(113, 66, 113, 72): 'volume'}, {}), '(volume)', True, 'import numpy as np\n'), ((100, 27, 100, 41), 'numpy.min', 'np.min', ({(100, 34, 100, 40): 'volume'}, {}), '(volume)', True, 'import numpy as np\n'), ((100, 46, 100, 60), 'numpy.max', 'np.max', ({(100, 53, 100, 59): 'volume'}, {}), '(volume)', True, 'import numpy as np\n'), ((100, 63, 100, 77), 'numpy.min', 'np.min', ({(100, 70, 100, 76): 'volume'}, {}), '(volume)', True, 'import numpy as np\n'), ((175, 43, 175, 63), 'os.listdir', 'os.listdir', ({(175, 54, 175, 62): 'data_dir'}, {}), '(data_dir)', False, 'import os\n'), ((179, 13, 179, 56), 'os.path.join', 'os.path.join', ({(179, 26, 179, 34): 'data_dir', (179, 36, 179, 39): 'vol', (179, 41, 179, 55): '"""mri/orig.mgz"""'}, {}), "(data_dir, vol, 'mri/orig.mgz')", False, 'import os\n'), ((179, 58, 179, 99), 'os.path.join', 'os.path.join', ({(179, 71, 179, 80): 'label_dir', (179, 82, 179, 98): "(vol + '_glm.mgz')"}, {}), "(label_dir, vol + '_glm.mgz')", False, 'import os\n'), ((223, 13, 223, 56), 'os.path.join', 'os.path.join', ({(223, 26, 223, 34): 'data_dir', (223, 36, 223, 39): 'vol', (223, 41, 223, 55): '"""mri/orig.mgz"""'}, {}), "(data_dir, vol, 'mri/orig.mgz')", False, 'import os\n'), ((184, 13, 184, 52), 'os.path.join', 'os.path.join', ({(184, 26, 184, 34): 'data_dir', (184, 36, 184, 39): 'vol', (184, 41, 184, 51): '"""orig.mgz"""'}, {}), "(data_dir, vol, 'orig.mgz')", False, 'import os\n'), ((184, 54, 184, 97), 'os.path.join', 'os.path.join', ({(184, 67, 184, 76): 'label_dir', (184, 78, 184, 81): 'vol', (184, 83, 184, 96): '"""Lab_con.mgz"""'}, {}), "(label_dir, vol, 'Lab_con.mgz')", False, 'import os\n'), ((228, 13, 228, 40), 'os.path.join', 'os.path.join', ({(228, 26, 228, 34): 'data_dir', (228, 36, 228, 39): 'vol'}, {}), '(data_dir, vol)', False, 'import os\n'), ((189, 13, 189, 63), 'os.path.join', 'os.path.join', ({(189, 26, 189, 34): 'data_dir', (189, 36, 189, 62): "(vol + '/' + vol + '_1.mgz')"}, {}), "(data_dir, vol + '/' + vol + '_1.mgz')", False, 'import os\n'), ((190, 13, 190, 68), 'os.path.join', 'os.path.join', ({(190, 26, 190, 35): 'label_dir', (190, 37, 190, 67): "(vol + '/' + vol + '_1_seg.mgz')"}, {}), "(label_dir, vol + '/' + vol + '_1_seg.mgz')", False, 'import os\n'), ((233, 13, 233, 52), 'os.path.join', 'os.path.join', ({(233, 26, 233, 34): 'data_dir', (233, 36, 233, 39): 'vol', (233, 41, 233, 51): '"""orig.mgz"""'}, {}), "(data_dir, vol, 'orig.mgz')", False, 'import os\n'), ((195, 13, 195, 56), 'os.path.join', 'os.path.join', ({(195, 26, 195, 34): 'data_dir', (195, 36, 195, 39): 'vol', (195, 41, 195, 55): '"""mri/orig.mgz"""'}, {}), "(data_dir, vol, 'mri/orig.mgz')", False, 'import os\n'), ((195, 58, 195, 102), 'os.path.join', 'os.path.join', ({(195, 71, 195, 80): 'label_dir', (195, 82, 195, 101): "(vol + '_map.nii.gz')"}, {}), "(label_dir, vol + '_map.nii.gz')", False, 'import os\n'), ((200, 13, 200, 40), 'os.path.join', 'os.path.join', ({(200, 26, 200, 34): 'data_dir', (200, 36, 200, 39): 'vol'}, {}), '(data_dir, vol)', False, 'import os\n')]