from torch.utils.data import Dataset
from torchvision import transforms
import os

import json
import cv2
import numpy as np

import dataloader.pt_data_loader.mytransforms as mytransforms
import dataloader.pt_data_loader.dataset_parameterset as dps
import dataloader.file_io.get_path as gp
import dataloader.file_io.dir_lister as dl


class BaseDataset(Dataset):
    """Image Dataset which can be used to load several images and their corresponding additional information"""

    def __init__(self,
                 dataset,
                 trainvaltest_split,
                 video_mode='mono',
                 stereo_mode='mono',
                 cluster_mode=None,
                 simple_mode=False,
                 labels=None,
                 labels_mode=None,
                 data_transforms=None,
                 scales=None,
                 keys_to_load=None,
                 keys_to_video=None,
                 keys_to_stereo=None,
                 split=None,
                 video_frames=None,
                 disable_const_items=True,
                 folders_to_load=None,
                 files_to_load=None,
                 n_files=None,
                 output_filenames=False,
                 flow_validation_mode=True
                 ):
        """Initialises the dataset by loading the desired data from the json file

        :param dataset: name of the dataset
        :param trainvaltest_split: can be train, validation or test
        :param video_mode: can be mono or video and defines if only the images or image sequences are to be loaded
        :param stereo_mode: can be mono or stereo and defines if the stereo images are to be loaded
        :param simple_mode: if True, the Data is read directly from a folder without using a .json file
        :param labels: gives the labels as defined in the named tuples style in Cityscapes. Get the labels from
            defintions folder
        :param labels_mode: can be fromid or fromrgb and defines if the segmentation masks are given as id or color
        :param data_transforms: takes the transforms.compose list
        :param scales: list of all scales at which the images should be loaded (list of exponents for powers of 2)
        :param keys_to_load: defines all keys which should be loaded
        :param keys_to_video: defines for which keys the sequences are to be loaded
        :param keys_to_stereo: defines for which keys the stereo images are supposed to be loaded
        :param split: dataset split that is supposed to be loaded. default is the complete dataset itself
        :param video_frames: all frames of the sequence that are supposed to be loaded (list of frame numbers relative
            to the main frame, e.g. [0, -2, -1, 1, 2])
        :param disable_const_items: removes the constant items like camera calibration from loading procedure
        :param folders_to_load: list of folders from which data should be loaded; folders not mentioned are skipped in
            the respective set. Only the last folder in a path is considered; filter is case insensitive.
            Default: None -> all folders are loaded from dataset
        :param files_to_load: list of files that should be loaded; files not mentioned are skipped in the respective
            set. File names need not be complete; filter is case insensitive.
            Default: None -> all files are loaded from dataset
        :param n_files: How many files shall be loaded. Files are selected randomly if there are more files than n_files
            Seeded by numpy.random.seed()
        """
        super(BaseDataset, self).__init__()
        assert isinstance(dataset, str)
        assert trainvaltest_split in ('train', 'validation', 'test'), '''trainvaltest_split must be train,
        validation or test'''
        assert video_mode in ('mono', 'video'), 'video_mode must be mono or video'
        assert stereo_mode in ('mono', 'stereo'), 'stereo_mode must be mono or stereo'
        assert isinstance(simple_mode, bool)
        if data_transforms is None:
            data_transforms = [mytransforms.CreateScaledImage(),
                               mytransforms.CreateColoraug(),
                               mytransforms.ToTensor()]
        if scales is None:
            scales = [0]
        if keys_to_load is None:
            keys_to_load = ['color']
        if keys_to_stereo is None and stereo_mode == 'stereo':
            keys_to_stereo = ['color']
        if keys_to_video is None and video_mode == 'video':
            keys_to_video = ['color']
        if video_frames is None:
            video_frames = [0, -1, 1]

        self.dataset = dataset
        self.video_mode = video_mode
        self.stereo_mode = stereo_mode
        self.scales = scales
        self.disable_const_items = disable_const_items
        self.output_filenames = output_filenames
        self.parameters = dps.DatasetParameterset(dataset)
        if labels is not None:
            self.parameters.labels = labels
        if labels_mode is not None:
            self.parameters.labels_mode = labels_mode
        path_getter = gp.GetPath()
        dataset_folder = path_getter.get_data_path()
        datasetpath = os.path.join(dataset_folder, self.dataset)
        self.datasetpath = datasetpath
        if split is None:
            splitpath = None
        else:
            splitpath = os.path.join(dataset_folder, self.dataset + '_' + split)

        if simple_mode is False:
            self.data = self.read_json_file(datasetpath, splitpath, trainvaltest_split,
                                            keys_to_load, keys_to_stereo, keys_to_video,
                                            video_frames, folders_to_load, files_to_load, n_files)
        else:
            self.data = self.read_from_folder(datasetpath, keys_to_load, video_mode, video_frames)

        self.load_transforms = transforms.Compose(
            [mytransforms.LoadRGB(),
             mytransforms.LoadSegmentation(),
             mytransforms.LoadDepth(),
             mytransforms.LoadFlow(validation_mode=flow_validation_mode),
             mytransforms.LoadNumerics()
             ])

        # IMPORTANT to create a new list if the same list is passed to multiple datasets. Otherwise, due to the
        # mutability of lists, ConvertSegmentation will only be added once. Hence, the labels may be wrong for the 2nd,
        # 3rd, ... dataset!
        self.data_transforms = list(data_transforms)

        # Error if CreateColorAug and CreateScaledImage not in transforms.
        if mytransforms.CreateScaledImage not in data_transforms:
            raise Exception('The transform CreateScaledImage() has to be part of the data_transforms list')
        if mytransforms.CreateColoraug not in data_transforms:
            raise Exception('The transform CreateColoraug() has to be part of the data_transforms list')

        # Error if depth, segmentation or flow keys are given but not the corresponding Convert-Transform
        if any([key.startswith('segmentation') for key in keys_to_load]) and \
                mytransforms.ConvertSegmentation not in self.data_transforms:
            raise Exception('When loading segmentation images, please add mytransforms.ConvertSegmentation() to '
                            'the data_transforms')
        if any([key.startswith('depth') for key in keys_to_load]) and \
                mytransforms.ConvertDepth not in self.data_transforms:
            raise Exception('When loading depth images, please add mytransforms.ConvertDepth() to the data_transforms')
        if any([key.startswith('flow') for key in keys_to_load]) and \
                mytransforms.ConvertFlow not in self.data_transforms:
            raise Exception('When loading flow images, please add mytransforms.ConvertFlow() to the data_transforms')

        # In the flow validation mode, it is not allowed to use data-altering transforms
        if any([key.startswith('flow') for key in keys_to_load]) and flow_validation_mode:
            allowed_transforms = [mytransforms.CreateScaledImage,
                                  mytransforms.CreateColoraug,
                                  mytransforms.ConvertSegmentation,
                                  mytransforms.ConvertDepth,
                                  mytransforms.ConvertFlow,
                                  mytransforms.RemoveOriginals,
                                  mytransforms.ToTensor,
                                  mytransforms.Relabel,
                                  mytransforms.OneHotEncoding,
                                  mytransforms.NormalizeZeroMean,
                                  mytransforms.AdjustKeys,
                                  mytransforms.RemapKeys,
                                  mytransforms.AddKeyValue
                                  ]
            for transform in self.data_transforms:
                if transform not in allowed_transforms:
                    raise Exception('In flow validation mode, it is not allowed to use data-altering transforms')

        # Set the correct parameters to the ConvertDepth and ConvertSegmentation transforms
        for i, transform in zip(range(len(self.data_transforms)), self.data_transforms):
            if isinstance(transform, mytransforms.ConvertDepth):
                transform.set_mode(self.parameters.depth_mode)
            elif isinstance(transform, mytransforms.ConvertSegmentation):
                transform.set_mode(self.parameters.labels, self.parameters.labels_mode)
            elif isinstance(transform, mytransforms.ConvertFlow):
                transform.set_mode(self.parameters.flow_mode, flow_validation_mode)

        self.data_transforms = transforms.Compose(self.data_transforms)

    def __len__(self):
        """Return the number of elements inside the dataset"""
        dict_keys = list(self.data.keys())
        return len(self.data[dict_keys[0]])

    def __getitem__(self, number):
        """Dataset element with index number 'number' is loaded"""
        sample = {}
        for item in list(self.data.keys()):
            if isinstance(self.data[item][number], str):
                element = self.read_image_file(self.data[item][number])
            else:
                element = self.data[item][number]
            sample.update({item: element})
        if not self.disable_const_items:
            sample = self.add_const_dataset_items(sample)
        sample = self.load_transforms(sample)
        sample = self.data_transforms(sample)
        if self.output_filenames:
            sample['filename'] = {}
            for item in list(self.data.keys()):
                if isinstance(self.data[item][number], str):
                    sample['filename'][item] = (self.data[item][number])
        return sample

    def add_const_dataset_items(self, sample):
        """Add dataset specific constants or items"""
        raise NotImplementedError

    def read_image_file(self, filepath):
        """Returns an image as a numpy array"""
        filepath = os.path.join(self.datasetpath, filepath)
        filepath = filepath.replace('/', os.sep)
        filepath = filepath.replace('\\', os.sep)
        image = cv2.imread(filepath, -1)
        return image

    def read_json_file(self, datasetpath, splitpath, trainvaltest_split, keys_to_load,
                       keys_to_stereo, keys_to_video, video_frames, folders_to_load, files_to_load, n_files):
        """Reads a json file from a dataset and outputs its data for the data loader
        here one might include filtering by folders for video data"""

        assert self.video_mode in ['mono', 'video'], 'video mode is not supported'
        assert self.stereo_mode in ['mono', 'stereo'], 'stereo mode is not supported'

        if splitpath is None:
            splitpath = datasetpath

        splitpath = os.path.join(splitpath, trainvaltest_split + '.json')

        datasetpath = os.path.join(datasetpath, 'basic_files' + '.json')
        assert os.path.isfile(datasetpath), 'Path to basic files is not valid'
        assert os.path.isfile(splitpath), 'Path to the split is not valid. Please use another argument for split.'

        # Load the basic filenames file to get the video data
        if self.video_mode == 'video':
            with open(datasetpath) as file:
                basic_json_data = json.load(file)
            basic_names = basic_json_data['names']
            basic_files = basic_json_data['files']
            basic_numerics = basic_json_data['numerical_values']

        # Load the split file
        with open(splitpath) as file:
            split_json_data = json.load(file)
        split_names = split_json_data['names']
        split_types = split_json_data['types']
        split_folders = split_json_data['folders']
        split_files = split_json_data['files']
        split_positions = split_json_data['positions']

        # Load all data necessary. Create dictionaries data_files and data_positions of the form
        # data_files = {('color', 0, -1): ['rgb_images/file_001_001', 'rgb_images/file_001_002', ...],
        #               ('depth', 0, -1): ['depth_images/file_001_001', 'depth_images/file_001_001', ...],
        #               ...}
        # data_positions = {('color', 0, -1): [[0, 0, 0, 0], [1, 0, 0, 1], ...],
        #                   ('depth', 0, -1): [[0, 0, 0, 0], [1, 0, 0, 1], ...],
        #                   ...}
        # The keys have three entries: the data category, a frame index (0 at the moment, can take different values
        # for video data) and a resolution parameter which has -1 as a default value und will be different after certain
        # data transforms have been performed.
        data_files = {}
        data_positions = {}
        frame_index = 0
        resolution = -1
        existing_positions_all = []

        keys_to_load = list(keys_to_load)
        # If the stereo mode is switched on, then all stereo keys are added if they exist in the split_names.
        # If this is not the case then nothing is added
        if self.stereo_mode == 'stereo':
            for name in split_names:
                if name in keys_to_stereo:
                    if name[-6:] == '_right':
                        stereo_name = name[:-6]
                    else:
                        stereo_name = name + '_right'
                    if stereo_name in split_names and stereo_name not in keys_to_load:
                        keys_to_load.append(stereo_name)
        keys_to_load = tuple(keys_to_load)

        for name, filetype, folder, file, position in \
                zip(split_names, split_types, split_folders, split_files, split_positions):
            if name in keys_to_load:
                data_files.update({(name, frame_index, resolution): file})
                data_positions.update({(name, frame_index, resolution): position})
        # Keep only files where all data is available. existing_positions_all contains all global positions that are
        # present in every data category. Then, remove all entries in the value lists of data_files and data_positions
        # that do not occur in every data category.
        for i, name in zip(range(len(data_files.keys())), data_files.keys()):
            if i == 0:
                existing_positions_all = np.array(data_positions[name])[:, 0]   # [:, 0] extracts the global position
            else:
                existing_positions_one = np.array(data_positions[name])[:, 0]
                existing_positions_all = list(set(existing_positions_all).intersection(existing_positions_one))
        existing_positions_all = sorted(existing_positions_all)

        for name in data_files.keys():
            existing_positions_one = np.array(data_positions[name])[:, 0]
            indices_to_keep = []
            index_one = 0
            index_all = 0
            while index_all < len(existing_positions_all):
                if existing_positions_one[index_one] == existing_positions_all[index_all]:
                    indices_to_keep.append(index_one)
                    index_all += 1
                index_one += 1
            data_files[name] = [data_files[name][j] for j in indices_to_keep]
            data_positions[name] = [data_positions[name][j] for j in indices_to_keep]

        # Keep only files in certain folders. Lists in data_files and data_positions must have the same size and order!
        if folders_to_load is not None:
            assert isinstance(folders_to_load, list), "please provide a list for folders_to_load"
            assert len(folders_to_load) > 0, "please provide a non-empty list for folders_to_load"

            # Get a key from data_files; which key does not matter as lists (dict values) should be of same size and
            # order.
            key = next(iter(data_files.keys()))
            list_to_work_on = data_files[key]
            indices_to_keep = []

            # Iterate over list_to_work_on as well as folders_to_load. If the last folder of a path "i" in
            # list_to_work_on is equal to any folder in folders_to_load, its index "i" is appended to indices_to_keep.
            folders_to_load = set([folder.lower() for folder in folders_to_load])
            for i in range(len(list_to_work_on)):
                dir_to_compare = list_to_work_on[i].split(os.sep)
                if not set(dir_to_compare).isdisjoint(folders_to_load):
                    indices_to_keep.append(i)

            # Filter all lists based on indices_to_keep.
            assert len(indices_to_keep) > 0, "given folders_to_load is/are not existing in dataset"
            for name in data_files.keys():
                data_files[name] = [data_files[name][j] for j in indices_to_keep]
                data_positions[name] = [data_positions[name][j] for j in indices_to_keep]

        # Keep only certain files. Lists in data_files and data_positions must have the same size and order!
        if files_to_load is not None:
            assert isinstance(files_to_load, list), "please provide a list for files_to_load"
            assert len(files_to_load) > 0, "please provide a non-empty list for files_to_load"

            # Get a key from data_files; which key does not matter as lists (dict values) should be of same size and
            # order.
            key = next(iter(data_files.keys()))
            list_to_work_on = data_files[key]
            indices_to_keep = []

            # Iterate over list_to_work_on as well as files_to_load. If the filename of a file "i" in
            # list_to_work_on contains any file from files_to_load, its index "i" is appended to indices_to_keep.
            files_to_load = [file.lower() for file in files_to_load]
            for i in range(len(list_to_work_on)):
                file_to_compare = list_to_work_on[i].split(os.sep)[-1].lower()
                for file in files_to_load:
                    if file_to_compare[:len(file)] == file:
                        indices_to_keep.append(i)
                        break

            # Filter all lists based on indices_to_keep.
            assert len(indices_to_keep) > 0, 'given files_to_load is/are not existing in dataset'
            for name in data_files.keys():
                data_files[name] = [data_files[name][j] for j in indices_to_keep]
                data_positions[name] = [data_positions[name][j] for j in indices_to_keep]

        # Keep only files with compatible video modes
        if self.video_mode == 'video':
            existing_positions_all = []
            min_frame = np.min(np.array(video_frames))
            max_frame = np.max(np.array(video_frames))

            # Create a list existing_positions_all containing all global positions for which a sufficient amount of
            # preceding and succeeding frames is available in every data category
            for name in data_files.keys():
                if name[0] in keys_to_video:
                    # Create boolean arrays indicating whether each frame has enough preceding/succeeding frames
                    existing_min_frames = np.array(data_positions[name])[:, 1] >= -min_frame
                    existing_max_frames = np.array(data_positions[name])[:, 2] >= max_frame
                    # Get the position entries of all frames that have enough preceding and succeeding frames
                    existing_positions_one = np.array(data_positions[name])[:, 0]
                    existing_positions_one = existing_positions_one[np.logical_and(existing_min_frames,
                                                                                   existing_max_frames)]
                    # existing_positions_all is the list of global positions that have occured in every list
                    # existing_positions_one so far
                    if len(existing_positions_all) == 0:
                        existing_positions_all = existing_positions_one
                    else:
                        existing_positions_all = list(set(existing_positions_all).intersection(existing_positions_one))
            existing_positions_all = sorted(existing_positions_all)

            # Remove all entries in the value lists of every data category which do not have enough preceding and
            # succeeding frames
            for name in data_files.keys():
                existing_positions_one = np.array(data_positions[name])[:, 0]
                indices_to_keep = []
                index_one = 0
                index_all = 0
                while index_all < len(existing_positions_all):
                    if existing_positions_one[index_one] == existing_positions_all[index_all]:
                        indices_to_keep.append(index_one)
                        index_all += 1
                    index_one += 1
                data_files[name] = [data_files[name][j] for j in indices_to_keep]
                data_positions[name] = [data_positions[name][j] for j in indices_to_keep]

        # Add video files to data
        if self.video_mode == 'video':
            original_keys = list(data_files.keys())
            for name in original_keys:
                if isinstance(name, tuple) and name[0] in keys_to_video:
                    basic_name_index = basic_names.index(name[0])
                    indices = np.array(data_positions[name])[:, 3]
                    # For every video frame index,
                    for frame_index in video_frames:
                        if frame_index == 0:
                            continue
                        else:
                            # If a numerical entry exists, it will be loaded instead of a file (e.g. for camera
                            # intrinsic parameters).
                            if basic_numerics[basic_name_index] is not None:
                                frame_file = [basic_numerics[basic_name_index][j + frame_index] for j in indices]
                            else:
                                frame_file = [basic_files[basic_name_index][j + frame_index] for j in indices]
                            data_files.update({(name[0], frame_index, resolution): frame_file})

        # Select n_files to load, based on a uniform distribution
        if n_files is not None:
            # Get a key from data_files; which key does not matter as lists (dict values) should be of same size and
            # order.
            key = next(iter(data_files.keys()))
            list_to_work_on = data_files[key]

            # If n_files is greater or equal to len(list_to_work_on), we don't have to do anything.
            if n_files < len(list_to_work_on):
                indices_to_keep = np.random.choice(len(list_to_work_on), size=n_files, replace=False)
                for name in data_files.keys():
                    data_files[name] = [data_files[name][j] for j in indices_to_keep]

        return data_files

    def read_from_folder(self, path, keys_to_load, video_mode, video_frames):
        """
        Creates the data dictionary directly from the folder without a .json-File. Only suitable for simple datasets.

        Folders should have the same name as keys. Folder structure is assumed to be as follows:
        <path>
          color
            <image_01>
            ...
            <image_n>
          sgementation
            <image_01>
            ...
            <image_n>
          ...

        :param path: path of the dataset/dataset split to use
        :param keys_to_load:
        :param video_mode:
        :return: a dictionary with all files for each key, sorted alphabetically by filename
        """

        assert keys_to_load is not None, 'in simple mode, the keys must be specified'
        root_stringlength = len(path) + 1
        folders = dl.DirLister.get_directories(path)
        folders = sorted(folders, key=str.lower)
        keys_to_load = sorted(keys_to_load, key=str.lower)
        data_files = {}

        if video_mode == 'mono':
            frame_index = 0
            resolution = -1
            for key in keys_to_load:
                files = []
                # include all folders containing the key
                key_folders = dl.DirLister.include_dirs_by_name(folders, key)
                for folder in key_folders:
                    new_files = dl.DirLister.list_files_in_directory(folder)
                    # remove root path from file
                    for i in range(len(new_files)):
                        new_files[i] = new_files[i][root_stringlength:]
                    files.extend(new_files)
                    # sort alphabetically by file name (not by the full path)
                    files = sorted(files, key=lambda file: os.path.split(file)[1].lower())
                data_files.update({(key, frame_index, resolution): files})

        elif video_mode == 'video':
            # determine how many frames before and after each frame should be loaded. Based on that, determine the first
            # frame to load with index 0
            min_frame = np.min(np.array(video_frames))
            max_frame = np.max(np.array(video_frames))
            resolution = -1

            for key in keys_to_load:
                # include all folders containing the key
                key_folders = dl.DirLister.include_dirs_by_name(folders, key)
                all_files = {}
                # Create a dictionary that contains all file names for every folder
                for folder in key_folders:
                    new_files = dl.DirLister.list_files_in_directory(folder)
                    # remove root path from file
                    for i in range(len(new_files)):
                        new_files[i] = new_files[i][root_stringlength:]
                    all_files.update({folder: new_files})

                # Update the data files dictionary for every frame index
                for frame_index in video_frames:
                    files = []
                    for folder in key_folders:
                        first_frame = -min_frame + frame_index
                        last_frame = len(all_files[folder]) - max_frame + frame_index
                        new_files = all_files[folder][first_frame:last_frame]
                        files.extend(new_files)
                    data_files.update({(key, frame_index, resolution): files})
        for key in data_files.keys():
            print(str(key) + ': ' + str(data_files[key]))
            print()
        return data_files
