import pandas as pd

import numpy as np

import matplotlib.pyplot as plt

import nibabel as nib

import os

from matplotlib.colors import LinearSegmentedColormap

def write_landmarks_to_list(landmarks, file_path):

    '''

    Write the landmarks to a text file.

    Args:

        landmarks ('list'): List of landmarks.

        file_path ('str'): Path to the text file. It includes the file name and extension.

    Returns:

        None

    '''

    # Write the landmarks to a text file

    with open(file_path, 'w') as file:

        for row in landmarks:

            file.write('\t'.join(map(str, row)) + '\n')

def get_landmarks_from_txt(transformed_file_path, search_key='OutputIndexFixed'):

    '''

    Get the transformed landmarks from the text file using a given search_key column index.

    The searcu_key column index is the column name in the text file, where the landmarks (input or transformed)

    are stored. The default value is 'OutputIndexFixed', which is the transformed landmarks. Only two values are

    possible: 'OutputIndexFixed' or 'InputIndex'.

    Args:

        transformed_file_path ('str'): Path to the transformed text file.

        search_key ('str'): Column name in the text file where the landmarks are stored.

    Returns:

        landmarks_list ('list'): List of transformed landmarks.

    '''

    # validate the search key

    assert search_key in ['OutputIndexFixed', 'InputIndex'], "The search_key must be either 'OutputIndexFixed' or 'InputIndex'."

    # Define the column names based on the data in the text file

    columns = [

        'Point', 'InputIndex', 'InputPoint',

        'OutputIndexFixed', 'OutputPoint', 'Deformation'

    ]

    # Read the text file into a pandas DataFrame

    df = pd.read_csv(transformed_file_path, sep='\t;', comment=';', header=None, names=columns, engine='python')

    # select the required column

    df_col = df[search_key]

    # convert the column values to a list of lists

    landmarks_list = [list(map(int, df_col[idx].split(' ')[-4:-1])) for idx in range(len(df_col))]

    return landmarks_list

def visualize_landmarks(slice_index=70, subject='copd1', split='train'):

    '''

    Visualize the landmarks on the reference image or a mask.

    '''

    # Define the paths

    landmarks_path = os.path.join(os.getcwd(),f'../dataset/{split}/{subject}/{subject}_300_iBH_xyz_r1.txt')

    reference_image_path = os.path.join(os.getcwd(),f'../dataset/{split}/{subject}/{subject}_iBHCT.nii.gz') # _lung

    reference_mask_path = os.path.join(os.getcwd(),f'../dataset/{split}/{subject}/{subject}_iBHCT_lung.nii.gz') # _lung

    # reference_mask_path = os.path.join(os.getcwd(),f'../dataset/segmentation_trails/mask1/train/{subject}_eBHCT_lung.nii.gz') # _lung

    # -1 to match the MATLAB visualizer indexing result

    slice_index = slice_index - 1

    # Load the reference image

    nii_image = nib.load(reference_image_path)

    reference_image = nii_image.get_fdata()

    # Load the reference mask

    nii_mask = nib.load(reference_mask_path)

    reference_mask = nii_mask.get_fdata()

    # transpose the axis to rotate for visualization

    reference_image = reference_image.transpose(2, 1, 0)

    reference_mask = reference_mask.transpose(2, 1, 0)

    # Load 3D landmarks from the file

    landmarks_data = np.loadtxt(landmarks_path, skiprows=2)

    slice_landmarsk = np.array([inner_list for inner_list in landmarks_data if inner_list[2] == slice_index+1])

    # Create a red-green colormap with opacity

    cmap = LinearSegmentedColormap.from_list('red_green', ['red', 'green'], N=256)

    # Visualize a specific slice

    plt.figure(figsize=(8, 8))

    plt.imshow(reference_image[slice_index, :, :], cmap='gray')

    # Overlay the mask with color on the landmarks

    mask_overlay = np.ma.masked_where(reference_mask[slice_index, :, :] == 0, reference_mask[slice_index, :, :])

    plt.imshow(mask_overlay, cmap=cmap, alpha=0.5)

    if len(slice_landmarsk) > 0:

        # Extract x, y, z coordinates

        x_coords = slice_landmarsk[:, 0].astype(int)

        y_coords = slice_landmarsk[:, 1].astype(int)

        z_coords = slice_landmarsk[:, 2].astype(int)

        # Plot landmarks on the current slice

        plt.scatter(x_coords, y_coords, c='b', marker='+', label=f'{landmarks_path.split("/")[-1].split(".txt")[0]} landmarks')

    plt.legend()

    plt.axis('off')

    # plt.title(f"Slice {slice_index+1}")

    plt.show()