import nibabel as nib

import numpy as np

import matplotlib.pyplot as plt

from skimage.transform import resize

from MSKMulticlass import CreateMasks4MulticlassMSK

def superimpose_images(image1, image2):

    image1 = image1 / np.max(image1)

    image2 = image2 / np.max(image2)

    alpha = 0.5

    superimposed_image = alpha * image1 + (1 - alpha) * image2

    return superimposed_image

def uniform_resizing(images, new_size):

    """

    Applies uniform resizing to an array of images. ### bicubic_interpolation

    Parameters:

        images (numpy.ndarray): Input array of images with shape (x, 512, 512, 1).

        new_size (tuple or int): The desired size of the resized images. If it's an int,

                                 the new size will be (new_size, new_size).

    Returns:

        numpy.ndarray: Array of resized images with shape (x, new_height, new_width, 1).

    """

    x, height, width, _ = images.shape

    if isinstance(new_size, int):

        new_height, new_width = new_size, new_size

    elif isinstance(new_size, tuple) and len(new_size) == 2:

        new_height, new_width = new_size

    else:

        raise ValueError("Invalid new_size argument. It should be an int or a tuple of two ints.")

    resized_images = np.zeros((x, new_height, new_width, 1))

    for i in range(x):

        resized_images[i, ..., 0] = resize(images[i, ..., 0], (new_height, new_width), mode='reflect', order=1)

        # resized_images[i, ..., 0] = resize(images[i, ..., 0], (new_height, new_width), mode='constant', preserve_range=True)

    return resized_images

# Pranav Change!

runVisualisationONLY = True

# Base Directory (Make a new folder!)

base_dir = 'D:/MRI - Friedlander'

# The mask index is the msk_00X number

# mask_index = [1] # Does one at a time

mask_index = [2,3,4,5] # Does multiple

# Make the following folders from your base directory

# 1)    First make a folder called nnUNet Data 

# 2)    From nnUNet folder make: scans, masks, multiclass_masks

# 3)    For unprocessed_scans folder load tbe scans as nii.gz using the msk_00X naming convention (ALWAYS back up scans outside of this file it )

# 4)    For masks folder look on discord for a reference pic

# 5)    For scans folder leave empty for processed out files

individual_mask_directory = ('{}/nnUNet Data/masks').format(base_dir)

multiclass_mask_output_dir = ('{}/nnUNet Data/multiclass_masks').format(base_dir)

scan_dir = ('{}/nnUNet Data/scans').format(base_dir) #output scans directory not unprocessed scans

input_scan_dir = ('{}/nnUNet Data/unprocessed_scans').format(base_dir)

# DONT EDIT!

TIBIA_encoding = 1

FEMUR_encoding = 2

FIBULA_encoding = 3

PELVIS_encoding = 4

AOIThresholding = True

FriedLanderDataset = True

if (runVisualisationONLY == False):

    print('Multi-Class Segmentation Task Data Preparation!')

    for i in range (len(mask_index)):

        CreateMasks4MulticlassMSK(input_scan_dir, scan_dir, individual_mask_directory, mask_index[i], TIBIA_encoding, FEMUR_encoding, FIBULA_encoding, PELVIS_encoding, multiclass_mask_output_dir, AOIThresholding, FriedLanderDataset)

# Visualisations

# Pranav Change!

slice = 650

msk_index = 3

img = nib.load(('{}/msk_00{}.nii.gz').format(scan_dir, msk_index))

img_data = img.get_fdata()

print('Training Scan Shape: ', img_data.shape)

mask = nib.load(('{}/msk_00{}.nii.gz').format(multiclass_mask_output_dir, msk_index))

mask_data = mask.get_fdata()

print('Training Mask Shape: ', mask_data.shape)

superimposed_data = superimpose_images(img_data[slice, :, :, 0], mask_data[slice, :, :, 0])

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

plt.axis('off')

plt.show()

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

plt.axis('off')

plt.show()

plt.imshow(superimposed_data, cmap='gray')

plt.axis('off')

plt.show()