import pandas as pd
from tqdm import tqdm
from pandarallel import pandarallel
import os
import SimpleITK as sitk
import numpy as np
from sklearn.cluster import KMeans
from skimage.morphology import ball, dilation, erosion

class SegmentationPhase:
    """
    Fase del pipeline para segmentar imágenes médicas utilizando varios métodos.

    Atributos:
        dataset_path (str): Ruta al conjunto de datos.
        segmentation_types (list[str]): Tipos de segmentación a aplicar.
        annotation_boxes (DataFrame): DataFrame que contiene las anotaciones.
    """

    def __init__(self, global_path, annotation_boxes_path, segmentation_types):
        """
        Inicializa la fase de segmentación.

        Args:
            global_path (str): Ruta global al conjunto de datos.
            annotation_boxes_path (str): Ruta al archivo de anotaciones.
            segmentation_types (list[str]): Tipos de segmentación a aplicar.
        """
        self.DATASET_PATH = os.path.join(global_path, 'Processed NIFTI Dataset/')
        self.segmentation_types = segmentation_types
        self.annotation_boxes = pd.read_csv(annotation_boxes_path)

    def segment_row(self, x):
        """
        Segmenta una fila de imagen basada en anotaciones y tipos de segmentación.

        Args:
            x (Series): Una fila del DataFrame de anotaciones.
        """
        patient_id = x['Patient_ID']
        patient_path = os.path.join(self.DATASET_PATH, patient_id)

        # Obtener índices de anotación
        row1, row2 = x['Start Row'], x['End Row']
        col1, col2 = x['Start Column'], x['End Column']
        slice1, slice2 = x['Start Slice'], x['End Slice']

        if os.path.isdir(patient_path):
            img_path = os.path.join(patient_path, 'post_1.img.gz')
            img = sitk.ReadImage(img_path)
            img_array = sitk.GetArrayFromImage(img)
            for seg_type in self.segmentation_types:
                self.apply_segmentation(img, img_array, patient_path, seg_type, slice1, slice2, row1, row2, col1, col2, patient_id)
                print(f"Segmentaciones para {patient_id} guardadas")
        
    def apply_segmentation(self, img, img_array, patient_path, seg_type, slice1, slice2, row1, row2, col1, col2, patient_id):
        """
        Aplica un tipo específico de segmentación a una imagen.

        Args:
            img (Image): Objeto de imagen SimpleITK.
            img_array (ndarray): Arreglo Numpy de la imagen.
            patient_path (str): Ruta a los datos del paciente.
            seg_type (str): Tipo de segmentación a aplicar.
            slice1 (int): Índice de inicio de corte.
            slice2 (int): Índice final de corte.
            row1 (int): Índice de inicio de fila.
            row2 (int): Índice final de fila.
            col1 (int): Índice de inicio de columna.
            col2 (int): Índice final de columna.
            patient_id (str): ID del paciente.
        """
        import SimpleITK as sitk 
        import numpy as np
        import os
        from sklearn.cluster import KMeans
        from skimage.morphology import ball, dilation, erosion
        
        if seg_type == 'otsu':
            # Otsu's thresholding
            otsu_filter = sitk.OtsuThresholdImageFilter()
            otsu_filter.SetInsideValue(0)
            otsu_filter.SetOutsideValue(1)
            seg_otsu = otsu_filter.Execute(img)
            seg_otsu_array = sitk.GetArrayFromImage(seg_otsu)
            mask_otsu = np.zeros_like(seg_otsu_array)
            mask_otsu[slice1:slice2, row1:row2, col1:col2] = seg_otsu_array[slice1:slice2, row1:row2, col1:col2]
            segment_otsu = sitk.GetImageFromArray(mask_otsu)
            segment_otsu.CopyInformation(img)
            sitk.WriteImage(segment_otsu, os.path.join(patient_path, 'otsu_segmentation.img.gz'))

        elif seg_type == 'kmeans':
            # K-means clustering
            bbox_img_array = img_array[slice1:slice2, row1:row2, col1:col2].reshape((-1, 1))
            kmeans = KMeans(n_clusters=2, random_state=0).fit(bbox_img_array)
            seg_kmeans_labels = kmeans.labels_.reshape(slice2-slice1, row2-row1, col2-col1)
            mask_kmeans = np.zeros_like(img_array)
            mask_kmeans[slice1:slice2, row1:row2, col1:col2] = seg_kmeans_labels
            segment_kmeans = sitk.GetImageFromArray(mask_kmeans)
            segment_kmeans.CopyInformation(img)
            sitk.WriteImage(segment_kmeans, os.path.join(patient_path, 'kmeans_segmentation.img.gz'))

        elif seg_type == 'box':
            # Segmentación ortoédrica
            mask_box = np.zeros_like(img_array)
            mask_box[slice1:slice2, row1:row2, col1:col2] = 1
            segment_box = sitk.GetImageFromArray(mask_box)
            segment_box.CopyInformation(img)
            sitk.WriteImage(segment_box, os.path.join(patient_path, 'box_segmentation.img.gz'))
        
        elif seg_type == 'ellipse':
            # Segmentación elipsoidal
            center_z, center_y, center_x = (slice1 + slice2) // 2, (row1 + row2) // 2, (col1 + col2) // 2
            semiaxis_z, semiaxis_y, semiaxis_x = (slice2 - slice1) // 2, (row2 - row1) // 2, (col2 - col1) // 2

            # Los semiejes no tienen que ser muy pequeños
            semiaxis_z = max(semiaxis_z, 1)
            semiaxis_y = max(semiaxis_y, 1)
            semiaxis_x = max(semiaxis_x, 1)

            mask_ellipse = np.zeros_like(img_array)
            for z in range(slice1, slice2):
                for y in range(row1, row2):
                    for x in range(col1, col2):
                        # Normalizar las distáncias
                        nz = (z - center_z) / semiaxis_z
                        ny = (y - center_y) / semiaxis_y
                        nx = (x - center_x) / semiaxis_x
                        if nz**2 + ny**2 + nx**2 < 1:
                            mask_ellipse[z, y, x] = 1
                            
            segment_ellipse = sitk.GetImageFromArray(mask_ellipse)
            segment_ellipse.CopyInformation(img)
            sitk.WriteImage(segment_ellipse, os.path.join(patient_path, 'ellipse_segmentation.img.gz'))
        
        elif seg_type == 'closing':
            segmentation_path = os.path.join(patient_path, f"{patient_id}_inverted_axes.nii.gz")
            if os.path.exists(segmentation_path):
                existing_segmentation = sitk.ReadImage(segmentation_path)
                existing_segmentation_array = sitk.GetArrayFromImage(existing_segmentation)

                # Aplicar dilatación y después erosion para realizar 'closing'
                mask_dilated = dilation(existing_segmentation_array, ball(2))
                mask_closing = erosion(mask_dilated, ball(2))

                segment_closing = sitk.GetImageFromArray(mask_closing)
                segment_closing.CopyInformation(existing_segmentation)
                sitk.WriteImage(segment_closing, os.path.join(patient_path, 'closing_segmentation.img.gz'))
                print('Closing aplicado y guardado en ', patient_path)
            else:
                print(f"Ruta a la segmentación existente no proporcionada o no encontrada para el paciente {patient_id}.")

        elif seg_type == 'opening':
            segmentation_path = os.path.join(patient_path, f"{patient_id}_inverted_axes.nii.gz")
            if os.path.exists(segmentation_path):
                existing_segmentation = sitk.ReadImage(segmentation_path)
                existing_segmentation_array = sitk.GetArrayFromImage(existing_segmentation)

                # Aplicar erosión y después erosion para realizar 'opening'
                mask_eroded = erosion(existing_segmentation_array, ball(2))
                mask_opening = dilation(mask_eroded, ball(2))

                segment_opening = sitk.GetImageFromArray(mask_opening)
                segment_opening.CopyInformation(existing_segmentation)
                sitk.WriteImage(segment_opening, os.path.join(patient_path, 'opening_segmentation.img.gz'))
                print('Opening aplicado y guardado en ', patient_path)
            else:
                print(f"Ruta a la segmentación existente no proporcionada o no encontrada para el paciente {patient_id}.")

    def run_pipeline(self, optimal_workers=3):
        """
        Ejecuta la fase de segmentación.

        Argumentos:
            optimal_workers (int): Número óptimo de trabajadores para procesamiento paralelo.
        """
        tqdm.pandas()
        pandarallel.initialize(progress_bar=True, nb_workers=optimal_workers)
        self.annotation_boxes.parallel_apply(self.segment_row, axis=1)
