""" Script to preprocess a given FDG PET image in .nii.
"""
# Import libraries
import glob
import os

import numpy as np
from numpy import ndarray
from numpy.random import seed
from tqdm import tqdm
from typing import List, Tuple
import matplotlib.pyplot as plt
import warnings

from skimage.transform import resize
import scipy.ndimage
import nibabel as nib

# seed random number generator
seed(1)


def directory_exist(dir_check: str = None):
    """ Check if the given path exists.

    Args:
        dir_check: directory path to check
    Raises:
        Raises exception if the given directory path doesn't exist.
    """
    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: %s not found \n" % (dir_check))


def generate_mip_show(pet: int = None, gt: int = None, identifier: str = None):
    """ Display the MIP images of a given 3D FDG PET image , and corresponding ground truth gt.

    Args:
        pet: 3D array of PET images
        gt: 3D array of reference or ground truth lymphoma segmentation.
        identifier: Name (identifier) of the patient.
    """
    # consider axis 0 for sagittal and axis 1 for coronal views
    pet, gt = np.array(pet), np.array(gt)
    pet = [np.amax(pet, axis=0), np.amax(pet, axis=1)]
    gt = [np.amax(gt, axis=0), np.amax(gt, axis=1)]
    try:
        pet = np.squeeze(pet)
        gt = np.squeeze(gt)
    except:
        pass

    img_gt = np.concatenate((pet, gt, pet), axis=0)
    display_image(img_gt, identifier=identifier)


def display_image(im_display: ndarray, identifier: str = None):
    """ display given array of images.

    Args:
        im_display: array of images to show
        identifier: patient name to display as title
    """
    plt.figure(figsize=(10, 1))
    plt.subplots_adjust(hspace=0.015)
    plt.suptitle("Showing image: " + str(identifier), fontsize=12, y=0.95)
    # loop through the length of tickers and keep track of index
    for n, im in enumerate(im_display):
        # add a new subplot iteratively
        plt.subplot(int(len(im_display) // 2), 2, n + 1)
        plt.imshow(np.log(im + 1))
    plt.show()


def get_pet_gt(current_dir: str = None):
    """ Read pet and corresponding reference images.

    Args:
        current_dir: directory path to the PET and corresponding ground truth image.
    Returns:
        Returns array of pet and ground truth images.

    """

    def get_nii_files(nii_path):
        """ Read .nii data in the given nii_path.
        if there are more .nii or .nii.gz file inside the nii_path directory, the function will return the
        first read image.

        Args:
            nii_path: directory path to the .nii or .nii.gz file.

        Returns:
            Returns loaded nibble format if it exists, else it returns empty array.

        """
        # more than one .nii or .nii.gz is found in the folder the first will be returned
        full_path = [path_i for path_i in glob.glob(str(nii_path) + "/*.nii.gz")]

        if not len(full_path):  # if no file exists that ends wtih .nii.gz
            full_path = [path_i for path_i in glob.glob(str(nii_path) + "/*.nii")]

        if not len(full_path):
            get_image = None  # raise Exception("No .nii or .nii.gz found in %s dirctory" % nii_path)

        else:
            print("%d files found in %s \t the first will be read " % (len(full_path), nii_path))
            print("reading ... %s" % full_path[0])
            get_image = nib.load(full_path[0])

        return get_image

    # all folders files
    # all_ = [path_i for path_i in glob.glob(str(nii_path) + "/*.nii")]

    # Get pet image
    # if "pet" in str()
    pet_path = str(current_dir) + "/pet/"
    pet = get_nii_files(pet_path)

    # Get gt image
    try:
        # if the ground truth exists
        gt_path = str(current_dir) + "/gt/"
        gt = get_nii_files(gt_path)
    except:
        #  if the ground truth does not exist
        gt = None

    return [pet, gt]


def resize_nii_to_desired_spacing(
        data: int = None, data_spacing: Tuple[float] = None, desired_spacing: ndarray = None,
        interpolation_order_value: int = None
        ):
    """ resizes a given input data into the desired spacing using the specified interpolation order.

    Args:
        data: array of input data to resize.
        data_spacing:  original input data spacing.
        desired_spacing: required output data spacing.
        interpolation_order_value: interpolation order. E.g., 0 for binary images, and 3 for cubic interpolation.

    Returns:
        resized data.

    """
    if desired_spacing is None:
        desired_spacing_x, desired_spacing_y, desired_spacing_z = [4.0, 4.0, 4.0]
    else:
        desired_spacing_x, desired_spacing_y, desired_spacing_z = desired_spacing

    print("Given data spacing \t Desired spacing \n")
    print(data_spacing, "\t", end="")
    print(desired_spacing, "\n")

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

    # New resolution, consider at least the input image is two-dimensional
    new_x_resolution = np.ceil(data.shape[0] * (data_spacing[0] / desired_spacing_x))
    new_y_resolution = np.ceil(data.shape[1] * (data_spacing[1] / desired_spacing_y))

    print('resizing')
    print(np.array(data).shape)

    if len(data_spacing) == 3 and len(np.squeeze(data).shape) == 3:  # 3D input image
        new_z_resolution = np.ceil(data.shape[2] * (data_spacing[2] / desired_spacing_z))

        # resize to new image resolution
        image_resized = resize(
            data, (new_x_resolution, new_y_resolution, new_z_resolution), order=interpolation_order_value,
            preserve_range=True, anti_aliasing=False
            )

    else:  # if the given input image is 2D
        image_resized = resize(
            data, (new_x_resolution, new_y_resolution), order=interpolation_order_value, preserve_range=True,
            anti_aliasing=False
            )

    return image_resized


def z_score(image):
    """ z-score operation on the input data.

    Args:
        image: input data to apply the z-score.

    Returns:
        Standardized value of the input using z-score. Z-score = (input - mean(input))/(std(input))

    """
    image = image.copy()
    if not isinstance(image, ndarray):
        image = np.array(image)

    image_nan = np.where(image > 0, image, np.nan)
    means = np.nanmean(image_nan, axis=(image.shape))
    stds = np.nanstd(image_nan, axis=(image.shape))
    image_norm = (image - means) / (stds + 1e-8)

    return image_norm


# cropping from the center based:
def crop_nii_to_desired_resolution(data: ndarray = None, cropped_resolution: List[int] = None):
    """ Crops the input data in the given cropped resolution.

    Args:
        data: input data to be cropped.
        cropped_resolution: desired output resolution.

    Returns:
        Returns cropped data of size cropped resolution.

    """
    #
    if not isinstance(data, list):  # if data is not array change it to array
        data = np.array(data)

    try:
        data = np.squeeze(data)
    except:
        pass

    if cropped_resolution is not None:
        cropped_resolution = [128, 128, 256]

    print("\n Initial data size \t Cropped data size ")
    print(data.shape, "\t", end=" ")
    print(cropped_resolution, "\n")

    if len(cropped_resolution) == 3 and len(data.shape) == 3:  # 3D data
        x, y, z = data.shape
    else:
        raise Exception("Input image not !")

    # start cropping : get middle x, y, z values by dividing by 2 and subtract the desired center of image resolution
    start_cropping_at_x = (x // 2 - (cropped_resolution[0] // 2))
    start_cropping_at_y = (y // 2 - (cropped_resolution[1] // 2))
    start_cropping_at_z = (z // 2 - (cropped_resolution[2] // 2))

    # check for off sets: mean the new cropping resolution is bigger than the input image's resolution
    off_set_x, off_set_y, off_set_z = 0, 0, 0
    if start_cropping_at_x < 0:
        off_set_x = np.abs(start_cropping_at_x)
        # set the first pixel at zero
        start_cropping_at_x = 0
    if start_cropping_at_y < 0:
        off_set_y = np.abs(start_cropping_at_y)
        # set the first pixel at zero
        start_cropping_at_y = 0
    if start_cropping_at_z < 0:
        off_set_z = np.abs(start_cropping_at_z)
        # set the first pixel at zero
        start_cropping_at_z = 0
    else:
        # take [0:cropped_resolution[0]]
        # Patients are given [x, y, z] when we say z it starts at zero (leg) to z (head).
        start_cropping_at_z = 2 * (z // 2 - (cropped_resolution[2] // 2))

    npad = ((off_set_x, off_set_x), (off_set_y, off_set_y), (2 * off_set_z, 0))

    # set zero value to the off set pixels mode='constant',
    data = np.pad(data, pad_width=npad, constant_values=0)

    # cropping to the given or set cropping resolution
    data = data[start_cropping_at_x:start_cropping_at_x + cropped_resolution[0],
           start_cropping_at_y:start_cropping_at_y + cropped_resolution[1],
           start_cropping_at_z:start_cropping_at_z + cropped_resolution[2]]

    return data


def save_nii_images(
        image: List[ndarray] = None, affine: ndarray = None, path_save: str = None, identifier: str = None,
        name: List[str] = None
        ):
    """ Save given images into the given directory. If no saving directory is given it will save into
    ./data/predicted/' directory.

    Args:
        image: data to save.
        affine: affine value.
        path_save: saving directory path.
        identifier: unique name of the case.
        name: name of the case to read.
    Returns:
        Saved image.
    """
    try:
        directory_exist(path_save)
    except:
        try:
            os.mkdir(path_save)
        except:
            if not os.path.exists("../data"):
                os.mkdir("../data")

            if not os.path.exists('../data/predicted/'):
                os.mkdir('../data/predicted/')

            path_save = '../data/predicted/'
    if identifier is not None:  # associate file name e.g. patient name
        dir = os.path.join(path_save, str(identifier))
    else:
        dir = path_save

    if not os.path.exists(dir):
        os.mkdir(dir)  # os.mkdir("./" + str(identifier))

    # print('saving it to %s\n' % dir)
    # .nii name, e.g. pet, gt, but if not given the base name of given directory
    if name is None:
        name = ['image_' + str(os.path.basename(dir)).split('.')[0]]

    if affine is None:
        affine = np.diag([4, 4, 4, 1])

    # image = np.flip(image, axis=-1)
    for select_image in range(len(image)):
        im_ = nib.Nifti1Image(image[select_image], affine)
        save_to = str(dir) + "/" + str(name[select_image])
        im_.to_filename(save_to)


def generate_mip_from_3D(given_3d_nii: ndarray = None, mip_axis: int = 0):
    """ Projects a 3D data into MIP along the selected MIP axis.

    Args:
        given_3d_nii: 3D array to project into MIP.
        mip_axis: Projecting axis.

    Returns:
        Returns MIP image.

    """
    if given_3d_nii is None:
        raise Exception("3D image not given")

    if not isinstance(given_3d_nii, list):
        given_3d_nii = np.array(given_3d_nii)

    mip = np.amax(given_3d_nii, axis=mip_axis)

    mip = np.asarray(mip)

    return mip


def transform_coordinate_space(modality_1, modality_2, mode='nearest'):
    """ Apply affine transformation on given data using affine from the second data.

    Adapted from: https://gist.github.com/zivy/79d7ee0490faee1156c1277a78e4a4c4
   Transfers coordinate space from modality_2 to modality_1
   Input images are in nifty/nibabel format (.nii or .nii.gz)


    Args:
        modality_1: reference modality.
        modality_2:  image modality to apply affine transformation.

    Returns:
        Returns affine transformed form of modality_2 using modality_1 as reference.

    """
    aff_t1 = modality_1.affine
    try:
        aff_t2 = modality_2.affine
    except:
        aff_t2 = aff_t1

    inv_af_2 = np.linalg.inv(aff_t2)
    out_shape = modality_1.get_fdata().shape

    # desired transformation
    T = inv_af_2.dot(aff_t1)

    # apply transformation
    transformed_img = scipy.ndimage.affine_transform(modality_2.get_fdata(), T, output_shape=out_shape, mode=mode)
    return transformed_img


# read PET and GT images
def read_pet_gt_resize_crop_save_as_3d_andor_mip(
        data_path: str = None, data_name: str = None, saving_dir: str = None, save_3D: bool = False, crop: bool = True,
        output_resolution: List[int] = None, desired_spacing: List[float] = None, generate_mip: bool = False
        ):
    """ Read pet and ground truth images from teh input data path. It also apply resize, and cropping operations.

    Args:
        data_path: directory to the raw 3D pet and ground truth .nii fies.
        data_name: unique the whole data identifier.
        saving_dir: directory path to save the processed images.
        save_3D: Save the processed 3D data or not.
        crop: apply cropping operations.
        output_resolution: desired output resolution.
        desired_spacing: required to be output spacing.
        generate_mip: project the processed 3D data into MIPs.

    Returns:
        Returns processed data.

    """
    if output_resolution is not None:
        # output resized and cropped image resolution
        rows, columns, depths = output_resolution
    else:  # default values
        # output resized and cropped image resolution=
        output_resolution = [128, 128, 256]

    if data_name is None:
        data_name = "unknown_data"

    '''
    get directory name  or patient id of the PET and gt images
    Assuming the file structure is :
   ---patient_id
         -- PET
            -- *.nii or .nii.gz
        -- gt
            -- *.nii or .nii.gz
    '''

    # check if the directory exist
    directory_exist(data_path)

    # by default the processed 3d and 2D MIP will be saved into the 'data' subdirectory, respectively with name tags
    # as '_default_3d_dir' and '_default_MIP_dir'

    def create_directory(directory_to_create: list):
        """

        :param directory_to_create:
        """
        for directory in directory_to_create:
            if not os.path.exists(directory):
                os.mkdir(directory)  # os.mkdir("./" + str(identifier))

    # if saving_dir is None:
    #     if not os.path.exists("../data"):
    #         os.mkdir("../data")
    #
    #     saving_dir_3d = "../data/" + str(data_name) + "_default_3d_dir"
    #
    #     # create directory in the data with name resized_cropped_MIP_default
    #     saving_dir_mip = "../data/" + str(data_name) + "_default_MIP_dir"
    #
    #     create_directory([saving_dir_3d, saving_dir_mip])
    #
    # # If the saving directory is given, a sub folder for the processed 3d and 2D MIP will be created
    # else:
    data_3d = str(data_name) + "_default_3d_dir_"
    data_mip = str(data_name) + "_default_MIP_dir"
    saving_dir_3d = os.path.join(saving_dir, data_3d)
    saving_dir_mip = os.path.join(saving_dir, data_mip)

    create_directory([saving_dir_3d, saving_dir_mip])

    # all patient ids
    case_ids = os.listdir(data_path)

    if not len(case_ids):  # reise exception if the directory is empty
        raise Exception("Directory %s is empty" % data_path)
    else:  # read the pet and gt images
        print("Continuing to read %d  cases" % len(case_ids))

    #  resize, to  given resolution
    if desired_spacing is None:
        desired_spacing = [4.0, 4.0, 4.0]

    # print where will be saved the 3d and mip
    print("\n")
    print(40 * "**==**")
    print("3D resized and cropped images will be saved to %s, if set save" % saving_dir_3d)
    print("Generated MIPs will be saved to %s, if set to save" % saving_dir_mip)
    print(40 * "**==**")
    print("\n")

    for image_name in tqdm(list(case_ids)):
        # if image_name in images:
        current_id = os.path.join(data_path, image_name)
        # read, resize, crop and save as 3D
        pet, gt = get_pet_gt(current_id)

        # resolution
        res_pet = pet.header.get_zooms()
        res_pet = tuple([float(round(res, 2)) for res in list(res_pet[:3])])
        affine = pet.affine
        print(f'Size of the PET input image: {np.asanyarray(pet.dataobj).shape}')

        # if there is a ground truth:
        if gt is not None:
            res_gt = gt.header.get_zooms()
            res_gt = tuple([float(round(res, 2)) for res in list(res_gt[:3])])

            # check if pet and gt are on the same spacing, otherwise gt could be generated from CT images
            if not res_pet[:3] == res_gt[:3]:  # assert (res_pet == res_gt)
                print("pet, and gt resolutions, respectively \n")
                print(res_pet, "\t", res_gt)
                warnings.warn("Pet and gt have different spacing, Alignment continue...")
                # apply affine transformation to move the gt to the PET space,
                # probably the gt was generated from the CT images
                gt = transform_coordinate_space(pet, gt, mode='constant')
                # raise Exception('Ground truth and pet images are not on the same spacing')
                gt[gt >= 1] = 1
                gt[gt < 1] = 0
            else:
                gt = np.asanyarray(gt.dataobj)

            """
               For remarc data the ground truth is inverted along the z-axis
               """
            if data_name == "remarc":
                gt = np.flip(gt, axis=-1)

            gt = resize_nii_to_desired_spacing(
                gt, data_spacing=res_pet, desired_spacing=desired_spacing, interpolation_order_value=0
                )

        pet = np.asanyarray(pet.dataobj)
        # if the given image has stacked as channel example one image in remarc : 175x175x274x2
        if pet.shape[-1] == 2:
            pet = pet[..., 0]

        # generate_mip_show(pet, gt, identifier=str(image_name))

        pet = resize_nii_to_desired_spacing(
            pet, data_spacing=res_pet, desired_spacing=desired_spacing, interpolation_order_value=3
            )

        '''
        if most data are with brain images at the very top, avoid cropping the brain images,instead crop from bottom
        or the leg part
        '''
        if gt is not None:
            if str(data_name).lower() == 'lnh':
                # flip left to right to mach lnh data to remarc
                gt = np.flip(gt, axis=-1)

        crop_zero_above_brain = True
        if crop_zero_above_brain:
            # remove all zero pixels just before the brian image
            xs, ys, zs = np.where(pet != 0)
            if len(zs):  # use zs for cropping the ground truth data also
                pet = pet[:, :, min(zs):]

        # if gt is None:
        #     generate_mip_show(pet, pet, identifier=str(image_name))
        # else:
        #     generate_mip_show(pet, gt, identifier=str(image_name))

        if crop:
            pet = crop_nii_to_desired_resolution(pet.copy(), cropped_resolution=output_resolution.copy())

            if gt is not None:
                # remove the zero values of pet image above the brain
                if len(zs):  # use zs for cropping the ground truth data also
                    gt = gt[:, :, min(zs):]

                gt = crop_nii_to_desired_resolution(gt.copy(), cropped_resolution=output_resolution.copy())

        # if gt is None:
        #     generate_mip_show(pet, pet, identifier=str(image_name))
        # else:
        #     generate_mip_show(pet, gt, identifier=str(image_name))

        if save_3D:
            # output image affine
            affine = np.diag([desired_spacing[0], desired_spacing[1], desired_spacing[2], 1])
            if gt is not None:
                save_nii_images(
                    [pet, gt], affine=affine, path_save=saving_dir_3d, identifier=str(image_name),
                    name=['pet', 'ground_truth']
                    )
            else:
                save_nii_images([pet], affine=affine, path_save=saving_dir_3d, identifier=str(image_name), name=['pet'])

        # generate Sagittal and coronal MIPs
        if generate_mip:
            for sagittal_coronal in range(2):
                pet_mip = generate_mip_from_3D(pet, mip_axis=int(sagittal_coronal))  # pet mip

                # assuming sagittal is on axis 0, and coronal is on axis 1
                if sagittal_coronal == 0:  # sagittal
                    naming_ = "sagittal"
                elif sagittal_coronal == 1:
                    naming_ = "coronal"

                if gt is not None:
                    gt_mip = generate_mip_from_3D(gt, mip_axis=int(sagittal_coronal))  # gt mip
                    # save the generated MIP
                    save_nii_images(
                        [pet_mip, gt_mip], affine, path_save=saving_dir_mip, identifier=str(image_name),
                        name=['pet_' + str(naming_), 'ground_truth_' + str(naming_)]
                        )
                else:
                    # save the generated MIP
                    save_nii_images(
                        [pet_mip], affine, path_save=saving_dir_mip, identifier=str(image_name),
                        name=['pet_' + str(naming_)]
                        )
    return saving_dir_mip


# Read .nii files using itk
if __name__ == '__main__':
    # how to run examples.
    # input_path = r"F:\Data\Remarc\REMARC/"
    # data_ = "remarc"
    #
    input_path = r"F:\Data\Vienna\No_ground_truth/"
    data_ = "LNH"
    saving_dir_mip = read_pet_gt_resize_crop_save_as_3d_andor_mip(
        data_path=input_path, data_name=data_, saving_dir=None, save_3D=True, crop=True,
        output_resolution=[128, 128, 256], desired_spacing=None, generate_mip=True
        )