import numpy as np
import pandas as pd
import SimpleITK as sitk
import os
import itertools
from pandarallel import pandarallel
from tqdm import tqdm
import matplotlib.pyplot as plt

class RobustnessPhase:
    """
    Pipeline para evaluar la robustez de las características radiómicas.

    Atributos:
        output_file_path: Ruta para guardar el archivo CSV de las puntuaciones de robustez.
        dataset_path: Ruta al conjunto de datos NIFTI procesados.
        features_path: Ruta a los resultados de la extracción de características.
        features: Diccionario que contiene dataframes de características de pyradiomics para cada tipo de segmentación.
        annotation_boxes: DataFrame que contiene boxes de anotación.
        segmentation_types: Lista de tipos de segmentación considerados.
        robustness_percent: Porcentaje de robustez a calcular.
        R: Matriz utilizada para cálculos de robustez.
    """


    def __init__(self, global_path, annotation_boxes_path, segmentation_types, robustness_percent=100):
        """
        Inicializa RobustnessPipeline con las rutas requeridas.
        """
        
        self.output_file_path = os.path.join(global_path, 'Robustness Scores/robustness_scores_'+segmentation_types[1]+'.csv')
        self.dataset_path = os.path.join(global_path, 'Processed NIFTI Dataset/')
        self.features1_path = os.path.join(global_path, 'Feature Extraction Results/', f"pyradiomics_features_{segmentation_types[0]}.csv")
        self.features2_path = os.path.join(global_path, 'Feature Extraction Results/', f"pyradiomics_features_{segmentation_types[1]}.csv")
        self.segmentation_types = segmentation_types
        self.annotation_boxes = pd.read_csv(annotation_boxes_path)
        self.segmentation_types = segmentation_types
        self.robustness_percent = robustness_percent
        self.R = self.create_robustness_matrix()

    def calculate_dice_coefficient(self, image1, image2):
        import SimpleITK as sitk
        import numpy as np

        array1 = sitk.GetArrayFromImage(image1) > 0
        array2 = sitk.GetArrayFromImage(image2) > 0

        # Asegurarse de que ambas imágenes tienen el mismo tamaño
        if array1.shape != array2.shape:
            # Determinar la forma máxima en cada dimensión
            max_shape = np.maximum(array1.shape, array2.shape)

            # Crear nuevas matrices de la forma máxima
            new_array1 = np.zeros(max_shape, dtype=bool)
            new_array2 = np.zeros(max_shape, dtype=bool)

            # Rellenar las matrices originales en las nuevas matrices
            new_array1[:array1.shape[0], :array1.shape[1], :array1.shape[2]] = array1
            new_array2[:array2.shape[0], :array2.shape[1], :array2.shape[2]] = array2

            array1, array2 = new_array1, new_array2

        intersection = np.logical_and(array1, array2)
        return 2.0 * intersection.sum() / (array1.sum() + array2.sum())

    def calculate_dice(self, image_path1, image_path2):
        """
        Calcula el coeficiente de Dice para un par de trazados de imagen.

        Args:
            image_path1: Ruta a la primera imagen.
            image_path2: Ruta a la segunda imagen.

        Returns:
            Dice coefficient as a float.
        """
        import SimpleITK as sitk
        image1 = sitk.ReadImage(image_path1)
        image2 = sitk.ReadImage(image_path2)
        return self.calculate_dice_coefficient(image1, image2)
    
    def calculate_relative_error_matrix(self, df1, df2, columns_to_exclude=['patient', 'sequence']):
        """
        Calcular la matriz de errores relativos entre dos DataFrames.

        Args:
            df1 (pd.DataFrame): Primer Dataframe.
            df2 (pd.DataFrame): Segundo Dataframe.
            columns_to_exclude (list): Lista de columnas para excluir.

        Returns:
            pd.DataFrame: DataFrame de los errores relativos.
        """
        def convert_to_real(x):
            if isinstance(x, str) and 'j' in x:
                return complex(x).real
            return x

        # Convierte cualquier número complejo (almacenado como cadena) a su parte real
        df1_real = df1.applymap(convert_to_real)
        df2_real = df2.applymap(convert_to_real)

        df1_numeric = df1_real.drop(columns=columns_to_exclude).select_dtypes(include=[np.number])
        df2_numeric = df2_real.drop(columns=columns_to_exclude).select_dtypes(include=[np.number])
        
        # Calcula la matriz de error relativo
        relative_error_matrix = abs(df1_numeric - df2_numeric) / (df1_numeric + 1e-8)

        return pd.DataFrame(relative_error_matrix, columns=df1_numeric.columns)
    
    def create_robustness_matrix(self):
        """
        Crea la matriz R 10x10 R para el cálculo de la robustez.
        """
        R = np.zeros((10, 10))

        # La última fila con todos unos
        R[9, :] = 1

        # La penúltima fila con todos unos, excepto el último elemento que es 0
        R[8, :-1] = 1

        # La antepenúltima fila con todos unos, excepto los últimos dos elementos que son 0
        R[7, :-2] = 1

        return R

    def calculate_robustness_score(self, dice_coefficients, relative_errors_matrix):
        """
        Calculate robustness scores for a pair of segmentations.
        """
        num_columns = relative_errors_matrix.shape[1]
        robustness_scores = []
        print(relative_errors_matrix)
        for i in range(num_columns):
            errors = relative_errors_matrix.iloc[:, i].values
            dice_coeffs_normalized = np.clip(np.array(dice_coefficients) * 10, 0, 9)
            relative_errors_normalized = np.clip(errors * 10, 0, 9)
            H = np.zeros((10, 10))
            for dice, error in zip(dice_coeffs_normalized, relative_errors_normalized):
                H[9-int(error), int(dice)] += 1
            r = np.sum(self.R * H)/relative_errors_matrix.shape[0]
            robustness_scores.append(r)
            print(H)
        return robustness_scores
   
    def plot_robustness_scores(self, df):
        """
        Genera y muestra un boxplot de las puntuaciones de robustez.

        Args:
            robustness_df (pd.DataFrame): DataFrame que contiene las puntuaciones de robustez.
        """
        feature_groups = {
            'Shape-based 3D': 'original_shape',
            'First Order Statistics': 'original_firstorder',
            'GLCM': 'original_glcm',
            'GLDM': 'original_gldm',
            'GLRM': 'original_glrlm',
            'GLSZM': 'original_glszm',
            'NGTDM': 'original_ngtdm'
        }

        group_colors = {
            'Shape-based 3D': 'red',
            'First Order Statistics': 'blue',
            'GLCM': 'green',
            'GLDM': 'yellow',
            'GLRM': 'orange',
            'GLSZM': 'purple',
            'NGTDM': 'pink',
            'Other': 'grey'  # Color para el grupo 'Other'
        }

        def extract_group_name(feature_name):
            """
            Extrae el nombre del grupo de una característica radiómica.

            Esta función recorre los grupos de características definidos y devuelve el nombre 
            del grupo al que pertenece una característica dada. Si la característica no coincide 
            con ninguno de los grupos predefinidos, se asigna al grupo 'Other'.

            Args:
                feature_name (str): Nombre de la característica radiómica cuyo grupo se busca identificar.

            Returns:
                str: El nombre del grupo al que pertenece la característica. Si no coincide con 
                     ningún grupo, devuelve 'Other'.
            """ 
            for group_name, prefix in feature_groups.items():
                if prefix in feature_name:
                    return group_name
            return 'Other'

        def simplify_feature_name(feature_name):
            """
            Simplifica el nombre de una característica radiómica eliminando ciertos prefijos.

            Esta función itera sobre los valores de prefijos de los grupos de características. 
            Si encuentra un prefijo dentro del nombre de la característica, lo elimina, así como 
            el prefijo 'original_', para simplificar el nombre. Si el prefijo del grupo no se 
            encuentra en el nombre de la característica, solo elimina el prefijo 'original_'.

            Args:
                feature_name (str): El nombre de la característica radiómica que se va a simplificar.

            Returns:
                str: El nombre simplificado de la característica.
            """
            for prefix in feature_groups.values():
                if prefix in feature_name:
                    return feature_name.replace(prefix + '_', '').replace('original_', '')
            return feature_name.replace('original_', '')

        df['Feature'] = df['Feature'].astype(str)

        df['Group'] = df['Feature'].apply(extract_group_name)
        custom_order = list(feature_groups.keys()) + ['Other']
        df['Group'] = pd.Categorical(df['Group'], categories=custom_order, ordered=True)
        df = df.sort_values('Group')
        df['Simplified Feature'] = df['Feature'].apply(simplify_feature_name)
        threshold = 0.6
        df['Color'] = df['Robustness Score'].apply(lambda x: 'green' if x >= threshold else 'grey')

        groups_first_half = ['Shape-based 3D', 'First Order Statistics', 'GLCM']
        groups_second_half = [g for g in custom_order if g not in groups_first_half]

        df_first_half = df[df['Group'].isin(groups_first_half)]
        df_second_half = df[df['Group'].isin(groups_second_half)]

        def plot_half_with_centered_labels(df_half, ax):
            """
            Genera un gráfico de barras con etiquetas centradas para una mitad del dataframe.

            Esta función crea un gráfico de barras horizontal donde cada barra representa una 
            característica radiómica. Las características se agrupan y se separan con un espacio 
            definido por 'group_gap'. Además, se asigna una posición a cada característica para 
            mantener el agrupamiento y el orden en el gráfico.

            Args:
                df_half (DataFrame): La mitad del dataframe que contiene las características, 
                                     sus grupos, colores y puntuaciones de robustez.
                ax (Axes): El objeto Axes de Matplotlib en el que se dibujará el gráfico.

            """ 
            group_gap = 2.5
            current_position = 0
            positions = []
            prev_group = None
            for index, row in df_half.iterrows():
                if row['Group'] != prev_group:
                    if prev_group is not None:
                        current_position += group_gap
                    prev_group = row['Group']
                positions.append(current_position)
                current_position += 1

            df_half['Position'] = positions
            for color in ['green', 'grey']:
                subset = df_half[df_half['Color'] == color]
                ax.barh(subset['Position'], subset['Robustness Score'], color=color, align='center')

            for group, group_data in df_half.groupby('Group'):
                min_pos = group_data['Position'].min()
                max_pos = group_data['Position'].max()
                middle_pos = (min_pos + max_pos) / 2
                ax.text(0, min_pos - 1, group, fontsize=12, verticalalignment='center', horizontalalignment='right')

            ax.set_yticks(df_half['Position'])
            ax.set_yticklabels(df_half['Simplified Feature'], rotation=0, fontsize=8)
            ax.invert_yaxis()
            for spine in ax.spines.values():
                spine.set_visible(False)
            ax.grid(axis='x')

        fig, axs = plt.subplots(1, 2, figsize=(11, 15))
        plot_half_with_centered_labels(df_first_half, axs[0])
        plot_half_with_centered_labels(df_second_half, axs[1])
        fig.suptitle('Feature Robustness Scores', fontsize=16)
        axs[0].set_xlabel('PDR', fontsize=14)
        axs[1].set_xlabel('PDR', fontsize=14)

        plt.tight_layout(rect=[0, 0.03, 1, 0.95])
        plt.savefig('../graficos/individuales/grafico_' + self.segmentation_types[1] + '.png')

    def get_image_path(self, patient_id, segmentation_type):
        """
        Construye y retorna la ruta al archivo de imagen correspondiente al ID del paciente y al tipo de segmentación.

        Args:
            patient_id (str): El identificador único del paciente.
            segmentation_type (str): El tipo de segmentación, puede ser 'manual' u otro tipo especificado como 'otsu', 'kmeans', etc.

        Returns:
            str: Ruta completa al archivo de imagen. Si el tipo de segmentación es 'manual', devuelve la ruta a la imagen con ejes                       invertidos. Para otros tipos de segmentación, devuelve la ruta a la imagen de segmentación correspondiente.

        """
        if segmentation_type=='manual':
            return os.path.join(self.dataset_path, patient_id, f"{patient_id}_inverted_axes.nii.gz")
        else:
            return os.path.join(self.dataset_path, patient_id, f'{segmentation_type}_segmentation.img.gz')
        
    def run_pipeline(self, optimal_workers=3):
        """
        Run the robustness analysis pipeline for two segmentation types.
        """
        import os
        tqdm.pandas()
        pandarallel.initialize(progress_bar=True, nb_workers=optimal_workers)
        
        assert len(self.segmentation_types) == 2, "There must be exactly two segmentation types."
        self.features1 = pd.read_csv(self.features1_path)
        self.features2 = pd.read_csv(self.features2_path)
        seg1, seg2 = self.segmentation_types
        feature_df1, feature_df2 = self.features1, self.features2

        image_paths1 = [self.get_image_path(patient_id, seg1) for patient_id in self.annotation_boxes['Patient_ID']]
        image_paths2 = [self.get_image_path(patient_id, seg2) for patient_id in self.annotation_boxes['Patient_ID']]
        
        # Crear un DataFrame temporal para el cálculo paralelo
        temp_df = pd.DataFrame({'path1': image_paths1, 'path2': image_paths2})
        
        # Calcular los coeficientes de Dice en paralelo
        dice_coefficients = temp_df.parallel_apply(lambda x: self.calculate_dice(x['path1'], x['path2']), axis=1)

        relative_errors = self.calculate_relative_error_matrix(feature_df1, feature_df2)
        robustness_scores = self.calculate_robustness_score(dice_coefficients, relative_errors)

        feature_names = list(self.features1.columns)

        columns_to_exclude=['patient', 'sequence', 'segmentation_type']
        feature_names_numeric = [name for name in feature_names if name not in columns_to_exclude]
        
        # Crear DataFrame para guardar en CSV
        robustness_df = pd.DataFrame({'Feature': feature_names_numeric, 'Robustness Score': robustness_scores})

        # Ordenar el DataFrame por 'Feature' en orden alfabético
        robustness_df = robustness_df.sort_values('Feature')

        # Guardar en CSV
        robustness_df.to_csv(self.output_file_path, index=False)
        
        # Generación del gráfico
        self.plot_robustness_scores(robustness_df)

        return robustness_df
