""" 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

        )