import SimpleITK as sitk

import os

import tempfile

import numpy as np

import matplotlib.pyplot as plt

from scipy.ndimage import binary_closing

from skimage import measure

def read_raw(

    binary_file_name,

    image_size,

    sitk_pixel_type,

    image_spacing=None,

    image_origin=None,

    big_endian=False,

):

    """

    Read a raw binary scalar image.

    Source: https://simpleitk.readthedocs.io/en/master/link_RawImageReading_docs.html

    Parameters

    ----------

    binary_file_name (str): Raw, binary image file content.

    image_size (tuple like): Size of image (e.g. [2048,2048])

    sitk_pixel_type (SimpleITK pixel type: Pixel type of data (e.g.

        sitk.sitkUInt16).

    image_spacing (tuple like): Optional image spacing, if none given assumed

        to be [1]*dim.

    image_origin (tuple like): Optional image origin, if none given assumed to

        be [0]*dim.

    big_endian (bool): Optional byte order indicator, if True big endian, else

        little endian.

    Returns

    -------

    SimpleITK image or None if fails.

    """

    pixel_dict = {

        sitk.sitkUInt8: "MET_UCHAR",

        sitk.sitkInt8: "MET_CHAR",

        sitk.sitkUInt16: "MET_USHORT",

        sitk.sitkInt16: "MET_SHORT",

        sitk.sitkUInt32: "MET_UINT",

        sitk.sitkInt32: "MET_INT",

        sitk.sitkUInt64: "MET_ULONG_LONG",

        sitk.sitkInt64: "MET_LONG_LONG",

        sitk.sitkFloat32: "MET_FLOAT",

        sitk.sitkFloat64: "MET_DOUBLE",

    }

    direction_cosine = [

        "1 0 0 1",

        "1 0 0 0 1 0 0 0 1",

        "1 0 0 0 0 1 0 0 0 0 1 0 0 0 0 1",

    ]

    dim = len(image_size)

    header = [

        "ObjectType = Image\n".encode(),

        (f"NDims = {dim}\n").encode(),

        (

            "DimSize = " + " ".join([str(v) for v in image_size]) + "\n"

        ).encode(),

        (

            "ElementSpacing = "

            + (

                " ".join([str(v) for v in image_spacing])

                if image_spacing

                else " ".join(["1"] * dim)

            )

            + "\n"

        ).encode(),

        (

            "Offset = "

            + (

                " ".join([str(v) for v in image_origin])

                if image_origin

                else " ".join(["0"] * dim) + "\n"

            )

        ).encode(),

        ("TransformMatrix = " + direction_cosine[dim - 2] + "\n").encode(),

        ("ElementType = " + pixel_dict[sitk_pixel_type] + "\n").encode(),

        "BinaryData = True\n".encode(),

        ("BinaryDataByteOrderMSB = " + str(big_endian) + "\n").encode(),

        # ElementDataFile must be the last entry in the header

        (

            "ElementDataFile = " + os.path.abspath(binary_file_name) + "\n"

        ).encode(),

    ]

    fp = tempfile.NamedTemporaryFile(suffix=".mhd", delete=False)

    # print(header)

    # Not using the tempfile with a context manager and auto-delete

    # because on windows we can't open the file a second time for ReadImage.

    fp.writelines(header)

    fp.close()

    img = sitk.ReadImage(fp.name)

    os.remove(fp.name)

    return img

def convert_raw_to_nifti(

        dataset_dir: list, 

        output_dir: str,

        metadata: dict):

    '''

    Convert the dataset images from raw (.raw) to nifti (nii.gz) format and export it to a 

    given output directory. 

    Args:

        dataset_dir (list): list of paths to the dataset directories

        output_dir (str): path to the output directory

        metadata (dict): dictionary containing the metadata of the dataset

    Returns:

        None

    '''

    pass

def create_mask(volume, threshold = 700):

    '''

    Create a mask from a given volume using a given threshold.

    Args:

        volume (numpy array): volume to be masked

        threshold (int): threshold to be used for the masking

    Returns:

        numpy array: masked volume

    '''

    return np.where(volume <= threshold, 1, 0)

def label_regions(mask):

    '''

    Label connected components in a binary mask.

    Args:

        mask (numpy array): Binary mask.

    Returns:

        tuple: A tuple containing labeled mask and the number of labels.

    '''

    labeled_mask, num_labels = measure.label(mask, connectivity=2, return_num=True, background=0)

    return labeled_mask, num_labels

def get_largest_regions(labeled_mask, num_regions=2):

    '''

    Get the largest connected regions in a labeled mask.

    Args:

        labeled_mask (numpy array): Labeled mask.

        num_regions (int): Number of largest regions to retrieve.

    Returns:

        list: List of region properties for the largest regions.

    '''

    regions = measure.regionprops(labeled_mask)

    regions.sort(key=lambda x: x.area, reverse=True)

    regions = regions[:min(num_regions, len(regions))]

    # print([regions[i].axis_major_length for i in range(len(regions))])

    # print([regions[i].axis_minor_length for i in range(len(regions))])

    return regions

def create_masks(labeled_mask, regions):

    '''

    Create masks for specific regions in a labeled mask.

    Args:

        labeled_mask (numpy array): Labeled mask.

        regions (list): List of region properties for which masks need to be created.

    Returns:

        list: List of masks corresponding to the specified regions.

    '''

    masks = [labeled_mask == region.label for region in regions]

    return masks

def fill_holes_and_erode(mask, structure=(7, 7, 5)):

    '''

    Fill holes in a binary mask and perform erosion.

    Args:

        mask (numpy array): Binary mask.

        dilation_structure (tuple): Dilation structure for binary dilation.

        erosion_structure (tuple): Erosion structure for binary erosion.

    Returns:

        numpy array: Processed mask after filling holes and erosion.

    '''

    processed_mask = binary_closing(mask, structure=np.ones(structure))

    return processed_mask

def remove_trachea(largest_masks, get_largest_regions, create_masks):

    '''

    Remove the trachea from a set of largest masks with a shape (Slice, H, W).

    Args:

        largest_masks (numpy array): 3D array of largest masks.

        get_largest_regions (function): Function to get largest regions.

        create_masks (function): Function to create masks.

    Returns:

        numpy array: 3D array of masks with trachea removed.

    '''

    # Find bounding boxes for each region in the 3D mask

    labeled_mask_slices = np.array([label_regions(largest_masks[idx, :, :])[0] for idx in range(largest_masks.shape[0])])

    # labeled_mask_slices = np.transpose(labeled_mask_slices, (1, 2, 0))

    largest_regions_slices = [

        get_largest_regions(labeled_mask_slices[idx, :, :], num_regions=3)

        for idx in range(labeled_mask_slices.shape[0])

    ]

    largest_regions_masks = [

        # we filter the trachea by checking the difference between the major and minor axis length when there is only 1 region

        create_masks(labeled_mask_slices[idx, :, :], region)[0] if (len(region) == 1 and (abs(region[0].axis_major_length - region[0].axis_minor_length) > 30))

        # this handles the very first few slices with trachea that has a very small difference between the major and minor axis length

        else np.zeros_like(labeled_mask_slices[idx, :, :]) if (len(region) == 1 and (abs(region[0].axis_major_length - region[0].axis_minor_length) < 30)) 

        # remove the trachea if there are 3 regions, it will be the 3rd region as we sort by area (highest to lowest)

        else create_masks(labeled_mask_slices[idx, :, :], region)[0] + create_masks(labeled_mask_slices[idx, :, :], region)[1] if len(region) == 3 

        # when there are only 2 regions, we check the difference in the area (area of the first region has to be atleast 50 more than the second region) to indicate that it is a lung not a trachea

        # also check if the minor axis of the second region (trachea) is less than 100

        # this condition happens when both lungs are touching each other as a region, and trachea as another region

        else create_masks(labeled_mask_slices[idx, :, :], region)[0] if len(region) == 2 and (getattr(region[0], 'area') - getattr(region[1], 'area') > 50) and (region[1].axis_minor_length < 100) 

        # when there are only 2 regions, we combine them. This is after the previous condition is met (when only 2 lungs are detected)

        else create_masks(labeled_mask_slices[idx, :, :], region)[0] + create_masks(labeled_mask_slices[idx, :, :], region)[1] if len(region) == 2 

        else np.zeros_like(labeled_mask_slices[idx, :, :])

        for idx, region in enumerate(largest_regions_slices)

    ]

    # largest_regions_masks = np.transpose(largest_regions_masks, (1, 2, 0))

    return largest_regions_masks

def segment_lungs_and_remove_trachea(volume, threshold=700, structure=(7, 7, 5), fill_holes_before_trachea_removal=False):

    '''

    Segment lungs and remove trachea from a given 3D volume with shape (Slice, H, W). Note that this shape is a must for 

    the internal functions to compute as expected.

    Args:

        volume (numpy array): 3D volume shape (slice, H, W).

        threshold (int): Threshold for creating the initial mask.

        dilation_structure (tuple): Dilation structure for binary dilation.

        erosion_structure (tuple): Erosion structure for binary erosion.

    Returns:

        initial_mask (numpy array): Initial mask created from the volume.

        labeled_mask (numpy array): Labeled mask.

        largest_masks (numpy array): 3D array of largest masks.

        processed_mask_without_trachea (numpy array): 3D binary array of masks with trachea removed.

    '''

    # create a mask

    initial_mask = create_mask(volume, threshold=threshold)

    # Label connected components

    labeled_mask, _ = label_regions(initial_mask)

    # Get the largest three regions (two lungs and trachea)

    largest_regions = get_largest_regions(labeled_mask, num_regions=3)

    # Create masks for the largest three regions

    largest_masks = create_masks(labeled_mask, largest_regions)[1]

    # fill holes of the largest mask

    if fill_holes_before_trachea_removal:

        largest_masks = fill_holes_and_erode(largest_masks, structure=tuple([2*x for x in structure]))

    # remove the trachea

    largest_masks_without_trachea = remove_trachea(largest_masks, get_largest_regions, create_masks)

    # Exclude the trachea by subtracting it from the processed mask

    processed_mask_without_trachea = fill_holes_and_erode(largest_masks_without_trachea, structure=structure)

    return initial_mask, labeled_mask, largest_masks, processed_mask_without_trachea.astype(np.uint8)

def segment_body(image, threshold=700):

    '''

    Segment the body from a given 3D volume with shape (Slice, H, W). Note that this shape is a must for

    the internal functions to compute as expected.

    Args:

        image (numpy array): 3D volume shape (slice, H, W).

        threshold (int): Threshold for creating the initial mask.

    Returns:

        mask (numpy array): Initial mask created from the volume.

        labeled_mask (numpy array): Labeled mask.

        largest_masks (numpy array): 3D array of largest masks.

        body_segmented (numpy array): 3D binary array of masks with body segmented.

    '''

    mask = create_mask(image, threshold=threshold)

    labeled_mask, _ = label_regions(mask)

    largest_regions = get_largest_regions(labeled_mask, num_regions=3)

    largest_masks = create_masks(labeled_mask, largest_regions)[0]

    # to have zeros and ones instead of binary false and true

    largest_masks = largest_masks.astype(np.int8)

    body_segmented = np.zeros_like(image)

    body_segmented[largest_masks == 0] = image[largest_masks == 0]

    return mask, labeled_mask, largest_masks, body_segmented

def display_two_volumes(volume1, volume2, title1, title2, slice=70):

    '''

    Display two volumes side by side.

    Args:

        volume1 (numpy array): first volume to be displayed

        volume2 (numpy array): second volume to be displayed

        title1 (str): title of the first volume

        title2 (str): title of the second volume

        slice (int): slice to be displayed

    Returns:

        None

    '''

    plt.figure(figsize=(9, 6))

    plt.subplot(1, 2, 1)

    plt.imshow(volume1[slice, :, :], cmap='gray') 

    plt.title(title1)

    plt.axis('off')

    plt.subplot(1, 2, 2)

    plt.imshow(volume2[slice, :, :], cmap='gray') 

    plt.title(title2)

    plt.axis('off')

    plt.show()

def display_volumes(*volumes, **titles_and_slices):

    '''

    Display multiple volumes side by side.

    Args:

        volumes (tuple of numpy arrays): volumes to be displayed

        titles_and_slices (dict): titles and slices for each volume

    Returns:

        None

    '''

    num_volumes = len(volumes)

    plt.figure(figsize=(6 * num_volumes, 6))

    for i, volume in enumerate(volumes, start=1):

        title = titles_and_slices.get(f'title{i}', f'Title {i}')

        slice_val = titles_and_slices.get(f'slice{i}', 70)

        plt.subplot(1, num_volumes, i)

        plt.imshow(volume[slice_val, :, :], cmap='gray') #gray

        plt.title(title)

        plt.axis('off')

    plt.show()

def min_max_normalization(image, mask = None, max_value=None):

    '''

    Perform min-max normalization on a given image.

    Args:

        image ('np.array'): Input image to normalize.

        mask ('np.array'): Mask to be applied to the image.

        max_value ('float'): Maximum value for normalization.

    Returns:

        normalized_image ('np.array'): Min-max normalized image.

    '''

    if max_value is None:

        max_value = np.iinfo(image.dtype).max

        print(f"The maximum value for this volume {image.dtype} is: {max_value}")

    print("Using mask for normalization" if mask is not None else "Not using mask for normalization")

    # Ensure the image is a NumPy array for efficient calculations

    image = np.array(image)

    # Calculate the minimum and maximum pixel values

    min_value = np.min(image[mask == 1]) if mask is not None else np.min(image)

    max_actual = np.max(image[mask == 1]) if mask is not None else np.max(image)

    # Perform min-max normalization

    normalized_image = (image - min_value) / (max_actual - min_value) * max_value

    normalized_image = np.clip(normalized_image, 0, max_value)

    return normalized_image.astype(image.dtype)