import os
import pandas as pd
import SimpleITK as sitk
from radiomics import featureextractor
from tqdm import tqdm
from pandarallel import pandarallel

class FeatureExtractionPhase:
    """
    Un pipeline para extraer características radiómicas de imágenes médicas.

    Atributos:
        DATASET_PATH (str): Ruta al conjunto de datos.
        annotation_boxes (pd.DataFrame): DataFrame que contiene las anotaciones.
        segmentation_types (list): Lista de tipos de segmentación a procesar.
        output_path (str): Ruta para guardar las características extraídas.
        feat_ext (radiomics.featureextractor.RadiomicsFeatureExtractor): Objeto extractor de características.
    """

    def __init__(self, features_output_path, global_path, annotation_boxes_path, segmentation_types):
        """
        Inicializa FeatureExtractionPipeline con ruta del conjunto de datos, anotaciones, tipos de segmentación y ruta de salida.

        Args:
            features_output_path (str): Ruta para guardar las características extraídas.
            global_path (str): Ruta global al conjunto de datos.
            annotation_boxes_path (str): Ruta al archivo de anotaciones.
            segmentation_types (list): Tipos de segmentación a procesar.
        """
        self.DATASET_PATH = os.path.join(global_path, 'Processed NIFTI Dataset/')
        self.annotation_boxes = pd.read_csv(annotation_boxes_path)
        self.segmentation_types = segmentation_types
        self.output_path = features_output_path
        self.feat_ext = featureextractor.RadiomicsFeatureExtractor()
        self.feat_ext.settings['minimumROIDimensions'] = 2
        self.feat_ext.settings['normalize'] = True
        self.feat_ext.settings['additionalInfo'] = False
        self.feat_ext.settings['correctMask'] = True

    def extract_features(self, x, seg_type):
        """
        Extrae características para una fila dada del DataFrame de anotaciones y un tipo de segmentación.

        Args:
            x (pd.Series): Una fila del DataFrame de anotaciones.
            seg_type (str): El tipo de segmentación.

        Return:
            dict: Características extraídas.
        """
        import os
        import SimpleITK as sitk
        patient_id = x['Patient_ID']
        file_path = os.path.join(self.DATASET_PATH, patient_id, 'post_1.img.gz')
        if seg_type == "manual":
            segmentation_path = os.path.join(self.DATASET_PATH, patient_id, f"{patient_id}_inverted_axes.nii.gz")
        else:
            segmentation_path = os.path.join(self.DATASET_PATH, patient_id, f'{seg_type}_segmentation.img.gz')
        
        if not os.path.exists(segmentation_path):
            return None

        img = sitk.ReadImage(file_path)
        smt = sitk.ReadImage(segmentation_path)

        smt.SetOrigin(img.GetOrigin())
        smt.SetSpacing(img.GetSpacing())
        smt.SetDirection(img.GetDirection())
        
        print(x)

        feats = self.feat_ext.execute(img, smt)
        feats['patient'] = patient_id
        feats['sequence'] = 'post_1'
        feats['segmentation_type'] = seg_type

        return feats

    def run_pipeline(self, optimal_workers=3):
        """
        Ejecuta la fase de extracción de características para cada tipo de segmentación y guarda los resultados.
        
        Args:
            optimal_workers (int): Número óptimo de trabajadores para paralelización.
        """
        import os
        tqdm.pandas()
        pandarallel.initialize(progress_bar=True, nb_workers=optimal_workers)

        # Asegurarse de que el directorio de salida existe
        if not os.path.exists(self.output_path):
            os.makedirs(self.output_path)

        for seg_type in self.segmentation_types:
            extracted_features = self.annotation_boxes.parallel_apply(self.extract_features, args=(seg_type,), axis=1)
            feature_df = pd.DataFrame([feat for feat in extracted_features if feat is not None])
            output_file = os.path.join(self.output_path, f'pyradiomics_features_{seg_type}.csv')
            feature_df.to_csv(output_file, index=False)
            print(f"Características para {seg_type} guardadas en {output_file}")