import numpy as np

import os

import pandas as pd

import matplotlib.pyplot as plt

from glob import glob

def plot_boxplot(experiment_name, output_dir, exclude=[], title="Boxplot"):

    '''

    Plot boxplot for the given data.

    Args:

        experiment_name (str): Name of the experiment.

        output_dir (str): Path to the output directory.

        exclude (list): List of columns to exclude from the boxplot.

        title (str): Title of the plot.

    Note:

        The dataframe holds the data in columns. Each column represents an experiment (single box plot) that we want to plot.

        We add the data to specific column of the experiment in the datafraame.

        Each row in the dataframe represents the result obtained from each subject in the experiment.

        >> df.head()

        >>          Par0003.affine  Par0003.bs-R1-fg  Par0003.bs-R6-ug  experiment_name

        >> copd1    10.62             26.25              1.34            ..

        >> copd2    10.07             21.45              2.68            ..

        >> copd3    03.57             12.04              1.27            ..

        >> copd4    07.48             29.45              1.53            ..

        If we describe the dataframe, we get the following:

        >> stats = df.describe()

        >> stats

        >>       Par0003.affine  Par0003.bs-R1-fg  Par0003.bs-R6-ug

        >> count  4.000000        4.000000          4.000000

        >> mean   7.935000        22.795000         1.705000

        >> std    3.417692        7.221071          0.700713

        >> min    3.570000        12.040000         1.270000

        >> 25%    6.345000        19.522500         1.330000

        >> 50%    8.775000        23.850000         1.435000

        >> 75%    10.365000       27.122500         2.060000

        >> max    10.620000       29.450000         2.680000

    Returns:

        None. The function generates and displays the box plot.

    '''

    # Get the data

    columns = os.listdir(f'../output/{experiment_name}/')

    TRE_sample_results = [path.replace('\\', '/') for path in glob(os.path.join(output_dir, experiment_name, "***", "points", "TRE_sample_results.csv"), recursive=True)]

    # Remove the excluded columns

    for column in exclude:

        if column in columns:

            columns.remove(column)

        TRE_sample_results = [item for item in TRE_sample_results if column not in item]

    # debugging

    # columns = columns[:5]

    # TRE_sample_results = TRE_sample_results[:5]

    # print(columns)

    # print(TRE_sample_results)

    # assert len(columns) == len(TRE_sample_results)

    assert  len(columns) == len(TRE_sample_results), f"Number of columns ({len(columns)}) does not match number of results ({len(TRE_sample_results)})"

    # Create a dataframe

    df = pd.DataFrame(columns=columns)

    for i, path in enumerate(TRE_sample_results):

        # Read the csv file

        data = pd.read_csv(path, index_col=0)

        # Add the data to the dataframe

        columns[i] = columns[i] + f" ({data['TRE_mean'].mean():.3f})"

        df[columns[i]] = data['TRE_mean']

    # Plot the boxplot

    boxplot = df.boxplot(column=columns, rot=90)

    # Get the lowest values for each column

    lowest_values = df.mean()

    # Get the column with the overall lowest minimum value

    lowest_column = lowest_values.idxmin()

    # Highlight the entire boxplot for the column with the lowest minimum value in red

    position = columns.index(lowest_column) + 1

    # 7 is the number of data that represents a single boxplot (divide len(boxplot.get_lines())//len(columns) to get the number of data per boxplot)

    boxplot.get_lines()[position * 7 - 7].set(color='red', linewidth=3) 

    # Set plot title

    plt.title(title)

    # Show the plot

    plt.show()

def compute_TRE(pts_exhale_file, pts_inhale_file, voxel_size):

    """

    Computes the Target Registration Error (TRE) to quantify the accuracy of the registration process. The TRE is calculated using 3D Euclidean 

    distance between the keypoints in the reference image (File 1) and the transformed keypoints in the registered image (File 2).

    Args:

        pts_exhale_file (str): path to the file containing the coordinates of the moving points

        pts_inhale_file (str): path to the file containing the coordinates of the fixed points

        voxel_size (tuple): voxel size in mm

    Returns:

        mean_TRE (float): mean TRE in mm

        std_TRE (float): standard deviation of the TRE in mm

    """

    # load the files if paths are provided

    pts_inhale = np.loadtxt(pts_inhale_file)

    pts_exhale = np.loadtxt(pts_exhale_file)

    # Check if the number of points in both files is the same

    if len(pts_inhale) != len(pts_exhale):

        raise ValueError("The number of points in the fixed and moving files must be the same.")

    # Compute the TRE - square root of the sum of the squared elements

    TRE = np.linalg.norm((pts_inhale - pts_exhale) * voxel_size, axis=1)

    return np.round(np.mean(TRE),2), np.round(np.std(TRE),2)