#Libraries

import numpy as np

import os

import pydicom

import SimpleITK as sitk

import pandas as pd

import trimesh

from pyntcloud import PyntCloud

from skimage.transform import resize

import sys

import matplotlib.pyplot as plt

from tqdm import tqdm

import nibabel as nib

import re

import cv2

from PIL import Image

# Helper Function

def flatten_3d_to_2d(array_3d):

    # Get the dimensions of the 3D array

    depth, height, width = array_3d.shape

    # Reshape the 3D array to a 2D array

    array_2d = np.reshape(array_3d, (depth, height * width))

    return array_2d

def transform_string(input_string):

    parts = input_string.split('_')

    if len(parts) < 2:

        return None

    prefix = parts[-2].upper()

    try:

        number_part = int(parts[0])

    except ValueError:

        return None

    new_string = f'{prefix}_{number_part:03d}'

    return new_string

# Helper Function

def flatten_2d_array(arr):

    flattened = []

    for row in arr:

        flattened.extend(row)

    return flattened

# Helper Function

# Read in entire scan of single patient

# folders = [f for f in os.listdir('MRI Scans - Tairawhiti') if os.path.isdir(os.path.join('MRI Scans - Tairawhiti', f))]

def ListFolders(directory):

    folder_names = []

    for root, dirs, files in os.walk(directory):

        for folder in dirs:

            folder_names.append(folder)

    return folder_names

# Helper Function

def read_dicom_files(directory):

    dicom_files = []

    files = os.listdir(directory)

    sorted_files = sorted(files)

    # print(sorted_files)

    for filename in sorted_files:

        filepath = os.path.join(directory, filename)

        if os.path.isfile(filepath) and filename.endswith('.dcm'):

            try:

                dicom_file = pydicom.dcmread(filepath)

                dicom_files.append(dicom_file)

            except pydicom.errors.InvalidDicomError:

                print(f"Skipping file: {filename}. It is not a valid DICOM file.")

    return dicom_files

# Helper Function

def get_ram_usage(variable, variable_name):

    size_in_bytes = sys.getsizeof(variable)

    size_in_kb = size_in_bytes / 1024

    size_in_mb = size_in_kb / 1024

    size_in_gb = size_in_mb / 1024

    message = "Memory usage of %s: %d %s." % (variable_name, size_in_mb, 'MB')

    print(message)

#Helper Function

def convert_4d_to_3d(array_4d, axis):

    array_3d = np.squeeze(array_4d, axis=axis)

    return array_3d

def plot_3d_data(x, y, z):

    # Define the size of the figure (width, height) in inches

    fig = plt.figure(figsize=(4, 4))

    # Plot the 3D scatter plot

    ax = fig.add_subplot(111, projection='3d')

    ax.scatter(x, y, z, c='blue')

    ax.set_xlabel('X')

    ax.set_ylabel('Y')

    ax.set_zlabel('Z')

    plt.show()

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 ReadIn_MRIScans_Masks(scans_path, folders):

    print('Patient Scan Data: ', folders)

    scan_pixel_data = []

    scan_coordinate_data = []

    single_scan_pixel_data = []

    scan_coordinate_data = []

    scan_orientation_data = []

    scan_pixelspacing_data = []

    single_paitent_scans_path =  scans_path + '/Raw DICOM MRI Scans/{}'.format(folders)

    dicom_files = read_dicom_files(single_paitent_scans_path)

    # Extracting pixel data

    for i in range (len(dicom_files)):

        normalized_data = (dicom_files[i].pixel_array)/(np.max(dicom_files[i].pixel_array))

        single_scan_pixel_data.append(normalized_data)

        print("Max pixel value in image stack is: ", np.max(normalized_data))

    scan_pixel_data.append(single_scan_pixel_data)

    training_scans = flatten_2d_array(scan_pixel_data)

    training_scans = np.array(training_scans)

    print("Max pixel value in image stack is: ", training_scans.max())

    # Coordinate Data

    single_paitent_scans_path =  scans_path + '/{}'.format(folders)

    for i in range (len(dicom_files)):

        scan_coordinate_data.append(dicom_files[i].ImagePositionPatient)

        scan_orientation_data.append(dicom_files[i].ImageOrientationPatient)

        scan_pixelspacing_data.append(dicom_files[i].PixelSpacing)

    coord_data = pd.DataFrame(scan_coordinate_data, columns=["x", "y", "z"])

    return training_scans

# Mapping coordinate data from groundtruth mask/label to mri training data

def MappingCoordinateData(filename_label, coord_data):

    # Load in mesh of label data

    mesh = trimesh.load_mesh(('C:/Users/GGPC/OneDrive/Desktop/Part 4 Project/Part4Project/SegmentationMasks/{}.ply').format(filename_label))

    # Convert the mesh vertices to a DataFrame

    vertices = pd.DataFrame(mesh.vertices, columns=["x", "y", "z"])

    # Pranav Ordering

    # coord_data = coord_data.sort_values('z', ascending = False)

    # coord_data = coord_data.reset_index(drop = True)

    vertices = vertices.sort_values('z', ascending = False)

    vertices = vertices.reset_index(drop = True)

    print('Height of Paitent in mm: ', np.abs(coord_data.iloc[-1][2] - coord_data.iloc[0][2]))

    print('Length of Paitent AOI (tibia) in mm: ', np.abs(vertices.iloc[-1][2] - vertices.iloc[0][2]))

    # vertices['z'] = vertices['z'].apply(lambda x: round(x, 1))

    coord_data['z'] = coord_data['z'].apply(lambda x: round(x, 1))

    # plot_3d_data(np.array(vertices['x']), np.array(vertices['y']), np.array(vertices['z']))

    if True:

        vertices.to_csv('ExactMaskCoordinateData.csv')

        coord_data.to_csv('ExactScanCoordinateData.csv')

    # vertices['z'] = np.round(vertices['z'] * 2) / 2

    # coord_data['z'] = np.round(coord_data['z'] * 2) / 2

    vertices['z'] = np.round(vertices['z'] * 2) / 2

    # coord_data['z'] = np.round(coord_data['z'] * 10) / 10

    # vertices['z'] = vertices['z'].apply(lambda x: round(x, 1))

    # coord_data['z'] = coord_data['z'].apply(lambda x: round(x, 1))

    if True:

        vertices.to_csv('RoundedMaskCoordinateData.csv')

    merged_df = pd.merge(coord_data, vertices, on='z')

    # condensed_df = merged_df.groupby('z').mean().reset_index()

    condensed_df = merged_df.groupby('z').median().reset_index()

    mapping_dict = dict(zip(condensed_df['z'], ['AOI']*len(condensed_df)))

    coord_data['SegmentationRegionSlice'] = coord_data['z'].map(mapping_dict).fillna('Outside of AOI')

    slices_aoi_start = (coord_data.loc[coord_data['SegmentationRegionSlice'] == 'AOI'].index)[0]

    slices_aoi_end = (coord_data.loc[coord_data['SegmentationRegionSlice'] == 'AOI'].index)[-1]

    slice_aoi_range = (slices_aoi_end - slices_aoi_start + 1)

    print('AOI Slice Start: ', slices_aoi_start)

    print('AOI Slice End: ', slices_aoi_end)

    print('AOI Slice Range: ', slice_aoi_range)

    # CSV Format 

    if True:

        coord_data.to_csv('tibia_mri_coord.csv')

        merged_df.to_csv('mergedcoordsystems.csv')

        condensed_df.to_csv('condensedmergedcoordsystems.csv')

    # print(coord_data)

    return slices_aoi_start, slices_aoi_end, slice_aoi_range, coord_data

def VoxelisationMask(filename_label, slice_aoi_range):

    # Load in mesh of label data

    mesh = trimesh.load_mesh(('C:/Users/GGPC/OneDrive/Desktop/Part 4 Project/Part4Project/SegmentationMasks/{}.ply').format(filename_label))

    # # Convert the mesh vertices to a DataFrame

    # vertices = pd.DataFrame(mesh.vertices, columns=["x", "y", "z"])

    # # Convert the mesh to a PyntCloud object

    # cloud = PyntCloud(vertices)

    vertices = pd.DataFrame(mesh.vertices, columns=["x", "y", "z"])

    faces = pd.DataFrame(mesh.faces, columns=['v1', 'v2', 'v3'])

    cloud = PyntCloud(points=vertices, mesh=faces)

    # Set the desired resolution

    desired_resolution = [slice_aoi_range, 512, 512]

    # Voxelize the mesh using the PyntCloud voxelization module

    voxelgrid_id = cloud.add_structure("voxelgrid", n_x=desired_resolution[0], n_y=desired_resolution[1], n_z=desired_resolution[2])

    voxel_grid = cloud.structures[voxelgrid_id].get_feature_vector().reshape(desired_resolution)

    # Transpose and swap axes to change the voxel grid orientation

    voxel_grid = np.transpose(voxel_grid, axes=(2, 0, 1))

    # Resize the voxel grid to match the desired dimensions

    voxel_grid = resize(voxel_grid, desired_resolution, anti_aliasing=False)

    voxel_grid = np.where(voxel_grid > 0, 1, 0)

    # print('Mask Slices Normalized to MRI Scans Shape (Purely AOI): ', voxel_grid.shape)

    return voxel_grid

def MaskCreation(basedir, filename_label):

    seg_masks_15A_tibia = sitk.ReadImage(('{}/Raw NIFITI Segmentation Masks (3D Slicer Output)/{}').format(basedir, filename_label))

    seg_masks_15A_tibia_data = sitk.GetArrayFromImage(seg_masks_15A_tibia)

    voxel_dimensions = seg_masks_15A_tibia.GetSpacing()

    voxel_dimensions = pd.DataFrame(data = voxel_dimensions)

    seg_masks_15A_tibia_data = seg_masks_15A_tibia_data[::-1, :, :]

    seg_masks_15A_tibia_data = np.where(seg_masks_15A_tibia_data != 0, 1, 0)

    seg_masks_15A_tibia_data = np.array(seg_masks_15A_tibia_data)

    indices = np.transpose(np.nonzero(seg_masks_15A_tibia_data != 0))

    if indices.size > 0:

        first_non_zero_index_2d = tuple(indices[0])

    else:

        first_non_zero_index_2d = None

    if indices.size > 0:

        last_non_zero_index_2d = tuple(indices[-1])

    else:

        last_non_zero_index_2d = None

    print("AOI Slice Start: ", first_non_zero_index_2d[0])

    print("AOI Slice End: ", last_non_zero_index_2d[0])

    print("AOI Slice Range: ", (last_non_zero_index_2d[0] - first_non_zero_index_2d[0] + 1))

    return seg_masks_15A_tibia_data, first_non_zero_index_2d[0], last_non_zero_index_2d[0]

def VisualValidationMSK(basedir, colab_fname, slice_idx_bin, slice_idx_multi, multiclass_dir, mask_index):

    nii_img_scan = nib.load(('{}/nnUNet Data/scans/msk_00{}.nii.gz').format(basedir, int(((colab_fname[0]).split('_'))[1])))

    nii_img_mask = nib.load(('{}/nnUNet Data/masks/{}/{}.nii.gz').format(basedir, ((colab_fname[0]).split('_'))[0],colab_fname[0]))

    mask_data = nii_img_mask.get_fdata()

    scan_data = nii_img_scan.get_fdata()

    image1 = scan_data[int(slice_idx_bin), :, :, 0]

    image2 = mask_data[int(slice_idx_bin), :, :, 0]

    superimposed_image = superimpose_images(image1, image2)

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

    plt.title(('Validating Scan & Mask (Binary) on Slice {} of {}').format(slice_idx_bin, colab_fname[0]))

    plt.axis('off')

    plt.show()

    if (multiclass_dir != False):

        nii_img_mask_multiclass = nib.load(('{}/msk_00{}.nii.gz').format(multiclass_dir, mask_index))

        multiclass_mask_data = nii_img_mask_multiclass.get_fdata()

        image1 = scan_data[slice_idx_multi, :, :, 0]

        image2 = multiclass_mask_data[slice_idx_multi, :, :, 0]

        superimposed_image = superimpose_images(image1, image2)

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

        plt.title(('Validating Scan & Mask (Multiclass) on Slice {}/{} of msk_00{}').format(slice_idx_multi, multiclass_mask_data.shape[0], mask_index))

        plt.axis('off')

        plt.show()

def preprocessing(scans_path, filename_labels, folders, total_slices_raw_data, DataOnlyAOI, Cropping):

    train_mask_tibia_labels, training_scans, start_slices_aoi, end_slices_aoi, slice_aoi_ranges  = [], [], [], [], []

    filename_labels = [filename_labels]

    print('\n')

    print('Patient Scan Data Folders Included in Run: ', folders)

    for index, filename_label in enumerate(filename_labels):

        print('\n')

        print('Segmentation Mask: ',('{}'.format(filename_label)))

        training_scan = ReadIn_MRIScans_Masks(scans_path, folders[index])

        seg_masks_15A_tibia_data, slices_aoi_start, slices_aoi_end = MaskCreation(scans_path, filename_label)

        median_aoi_index = int(np.abs(slices_aoi_end - slices_aoi_start) / 2) + slices_aoi_start

        if(DataOnlyAOI == True):

            if (index == 0):

                train_mask_tibia_labels = seg_masks_15A_tibia_data[(slices_aoi_start):(slices_aoi_end+1)]

                training_scans = training_scan[(slices_aoi_start):(slices_aoi_end+1)]

            else:

                train_mask_tibia_labels = np.concatenate((train_mask_tibia_labels, seg_masks_15A_tibia_data[(slices_aoi_start):(slices_aoi_end+1)]), axis=0)

                training_scans = np.concatenate((training_scans, training_scan[(slices_aoi_start):(slices_aoi_end+1)]), axis=0)

        if (DataOnlyAOI == False):

            train_mask_tibia_labels.append(seg_masks_15A_tibia_data)

            training_scans.append(training_scan)

        start_slices_aoi.append(slices_aoi_start)

        end_slices_aoi.append(slices_aoi_end)

        print('\n')

    train_mask_tibia_labels = np.array(train_mask_tibia_labels)

    training_scans = np.array(training_scans)

    # Normalization and Binarization

    training_scans = training_scans.astype('float32')

    training_scans /= 255.  # scale scans to [0, 1]

    train_mask_tibia_labels = np.where(train_mask_tibia_labels != 0, 1, 0)

    print("Scans Normalized! [0-1]")

    print("Max pixel value in image stack is: ", training_scans.max())

    print("Masks Binarised! [0,1]")

    print("Labels in the mask are : ", np.unique(train_mask_tibia_labels))

    print("\n")

    if (DataOnlyAOI == False):

        training_scans = np.reshape(training_scans, (len(folders) * training_scans.shape[1], 512, 512))

        train_mask_tibia_labels = np.reshape(train_mask_tibia_labels, (len(filename_labels) * train_mask_tibia_labels.shape[1], 512, 512))

        training_scans = np.expand_dims(training_scans, axis=-1)

        train_mask_tibia_labels = np.expand_dims(train_mask_tibia_labels, axis=-1)

    if (DataOnlyAOI == True):

        training_scans = np.expand_dims(training_scans, axis=-1)

        train_mask_tibia_labels = np.expand_dims(train_mask_tibia_labels, axis=-1)

    if (Cropping == True):

    # Resize 'training_scans'

        training_scans_resized = np.empty((training_scans.shape[0], 256, 256, 1), dtype=np.float32)

        for i in range(training_scans.shape[0]):

            input_image = training_scans[i, :, :, 0]  # Extract the 2D image from the 4D array

            pil_image = Image.fromarray(input_image)  # Convert to Pillow Image

            resized_image = pil_image.resize((256, 256), Image.BILINEAR)  # Resize using Pillow

            training_scans_resized[i, :, :, 0] = np.array(resized_image)  # Store the resized image in the output array

        # Resize 'train_mask_tibia_labels'

        train_mask_tibia_labels_resized = np.empty((train_mask_tibia_labels.shape[0], 256, 256, 1), dtype=np.uint8)

        for i in range(train_mask_tibia_labels.shape[0]):

            input_image = train_mask_tibia_labels[i, :, :, 0]  # Extract the 2D image from the 4D array

            pil_image = Image.fromarray(input_image)  # Convert to Pillow Image

            resized_image = pil_image.resize((256, 256), Image.BILINEAR)  # Resize using Pillow

            train_mask_tibia_labels_resized[i, :, :, 0] = np.array(resized_image)  # Store the resized image in the output array

        training_scans = training_scans_resized

        train_mask_tibia_labels = train_mask_tibia_labels_resized

    print('Training Scans Input Shape (Full 3D Stack): ', training_scans.shape)

    print('Training Masks Input Shape (Full 3D Stack): ', train_mask_tibia_labels.shape)

    return training_scans, train_mask_tibia_labels, median_aoi_index