"""

"""
import os
import glob
import numpy as np
from numpy import ndarray
import matplotlib.pyplot as plt
import nibabel as nib
from typing import List, Tuple
from numpy.random import seed

# seed random number generator
seed(1)


class DataLoader:
    """
    read preprocessed pet and gt MIP data for training
    """

    def __init__(self, data_dir: str, ids_to_read: ndarray = None, shuffle=True, training: bool = True):
        self.data_dir = data_dir
        self.ids_to_read = ids_to_read
        self.shuffle = shuffle
        self.training = training

    def get_batch_of_data(self):
        """
        data structure:
        -- main directory
        ------case Name:
                -- pet.nii.gz
                -- gt.nii.gz
        --Given list of training and testing on .text files
            -- train.text
            -- valid.text
        """

        # check directory
        self.directory_exist(self.data_dir)

        # get all names of the directories under data_dir
        case_ids = os.listdir(self.data_dir)

        # store batch data
        image_batch, ground_truth_batch = [], []

        # if there are file in data dir
        if not len(case_ids):
            raise Exception("No files found in %s" % self.data_dir)

        # else continue getting.reading the files
        for get_id in list(case_ids):
            if str(get_id) in list(self.ids_to_read):
                try:
                    # consider there four images in each folder name get_id:
                    # e.g. : coronal (gt_1, pet_1) and sagittal  (gt_0, pet_0)
                    current_dir = os.path.join(self.data_dir, str(get_id))
                    # read sagittal and coronal as independent images
                    pet_sagittla_coronal, gt_sagittal_coronal = self.get_nii_files_path(current_dir)

                    # pet, normalization, standardization
                    if len(pet_sagittla_coronal):  # if image is read
                        pet_sagittla_coronal = self.data_normalization_standardization(pet_sagittla_coronal,
                                                                                       z_score=True,
                                                                                       z_score_include_zeros=False)

                        gt_sagittal_coronal = self.data_normalization_standardization(gt_sagittal_coronal, threshold=True)

                        # display or save samples
                        # self.mip_show(pet=pet_sagittla_coronal, gt=gt_sagittal_coronal, identifier=str(get_id))

                        # collect all images with case_id
                        if not bool(len(image_batch)):  # if it is empty; first time
                            image_batch = pet_sagittla_coronal
                            ground_truth_batch = gt_sagittal_coronal
                        else:
                            image_batch = np.concatenate((image_batch, pet_sagittla_coronal), axis=0)
                            ground_truth_batch = np.concatenate((ground_truth_batch, gt_sagittal_coronal), axis=0)
                except:
                    print('Not read %s' %(str(get_id)))

        return [image_batch, ground_truth_batch]

    @staticmethod
    def directory_exist(dir_check: str = None) -> None:
        """
        :param dir_check:
        """
        if os.path.exists(dir_check):
            #  print("The directory %s does exist \n" % dir_check)
            pass
        else:
            raise Exception(
                "Please provide the correct path to the processed data ! \n %s not found \n" % (dir_check))

    @staticmethod
    def mip_show(pet: ndarray = None, gt: ndarray = None, identifier: str = None) -> None:
        """

        :param pet:
        :param gt:
        :param identifier:
        :return:
        """
        # consider axis 0 for sagittal and axis 1 for coronal views
        fig, axs = plt.subplots(1, 4, figsize=(15, 15))
        plt.title(str(identifier))
        try:
            pet = np.squeeze(pet)
            gt = np.squeeze(gt)
        except:
            pass

        axs[0].imshow(np.rot90(np.log(pet[0] + 1)))
        axs[0].set_title('pet_project_on_axis_0')
        axs[1].imshow(np.rot90(np.log(gt[0] + 1)))
        axs[1].set_title('gt_project_on_axis_0')
        axs[2].imshow(np.rot90(np.log(pet[1] + 1)))
        axs[2].set_title('project_on_axis_1')
        axs[3].imshow(np.rot90(np.log(gt[1] + 1)))
        axs[3].set_title('gt_project_on_axis_1')
        plt.show()

    @staticmethod
    def get_nii_files_path(data_directory: str) -> List[ndarray]:
        """
        read .nii or .nii.gz files from a given folder of path data_directory
        :param data_directory:
        :return:
        """
        # more than one .nii or .nii.gz is found in the folder the first will be returned
        types = ('/*.nii', '/*.nii.gz')  # the tuple of file types
        nii_paths = []
        for files in types:
            nii_paths.extend([i for i in glob.glob(str(data_directory) + files)])

        pet, gt = [], []
        if not len(nii_paths):  # if no file exists that ends wtih .nii.gz or .nii
            # raise Exception("No .nii or .nii.gz found in %s dirctory" % data_directory)
            pass
        else:
            # assuming the folder contains coronal mips: pet_1, gt_1, and sagittal mips: pet_0, gt_0,
            pet_saggital, pet_coronal, gt_saggital, gt_coronal = [], [], [], []
            for path in list(nii_paths):
                # get the base name: means the file name
                identifier_base_name = str(os.path.basename(path)).split('.')[0]
                if "pet_sagittal" == str(identifier_base_name):
                    pet_saggital = np.asanyarray(nib.load(path).dataobj)
                    pet_saggital = np.expand_dims(pet_saggital, axis=0)

                elif "pet_coronal" == str(identifier_base_name):
                    pet_coronal = np.asanyarray(nib.load(path).dataobj)
                    pet_coronal = np.expand_dims(pet_coronal, axis=0)

                if "ground_truth_sagittal" == str(identifier_base_name):
                    gt_saggital = np.asanyarray(nib.load(path).dataobj)
                    gt_saggital = np.expand_dims(gt_saggital, axis=0)

                elif "ground_truth_coronal" == str(identifier_base_name):
                    gt_coronal = np.asanyarray(nib.load(path).dataobj)
                    gt_coronal = np.expand_dims(gt_coronal, axis=0)

            # concatenate coronal and sagita images
            # show
            pet = np.concatenate((pet_saggital, pet_coronal), axis=0)
            gt = np.concatenate((gt_saggital, gt_coronal), axis=0)
        return [pet, gt]

    @staticmethod
    def z_score(image: ndarray, include_zeros: bool = False):
        """

        :param image:
        :param include_zeros:
        :return:
        """
        # include zeros
        if include_zeros:
            image = (image - np.mean(image)) / (np.std(image) + 1e-8)
        else:
            # Don't include zeros
            means = np.true_divide(image.sum(), (image != 0).sum())
            stds = np.nanstd(np.where(np.isclose(image, 0), np.nan, image))
            image = (image - means) / (stds + 1e-8)
        return image

    def data_normalization_standardization(self, data: ndarray, threshold: bool = False, z_score: bool = False,
                                           z_score_include_zeros: bool = False,
                                           min_max_scale: bool = False, log_transform: bool = False) -> List[ndarray]:
        """
        Data normalization and standardization function
        :param data:
        :param threshold:
        :param z_score:
        :param z_score_include_zeros:
        :param min_max_scale:
        :param log_transform:
        :return:
        """

        if not isinstance(data, List):
            data = np.array(data)

        # groundtruh > 0 is 1 and <=0 is 0
        if threshold:
            data[data > 0] = 1

        if z_score:
            data = self.z_score(data, include_zeros=z_score_include_zeros)

        if min_max_scale:
            data = (data - min(data)) / (max(data) - min(data))

        if log_transform:
            data = np.log(data + 1)

        return data


if __name__ == '__main__':
    # for Example
    print("data_loader for preprocessed coronal and sagittal MIPs, pet, and gt")
    data_dir = "../data/vienna_default_MIP_dir/"
    ids_to_read = os.listdir(data_dir)

    data_loader = DataLoader(data_dir=data_dir, ids_to_read=ids_to_read)
    loaded_data = data_loader.get_batch_of_data()
    print(np.array(loaded_data).shape)