from imgtools.io import read_dicom_series
from itertools import chain
from joblib import Parallel, delayed

# from readii.utils.logging_config import logger  # Updated import
from radiomics import featureextractor, imageoperations
import os
import pandas as pd
import numpy as np
import SimpleITK as sitk

from readii.image_processing import (
    flattenImage, 
    alignImages, 
    getROIVoxelLabel,    
)

from readii.loaders import (
    loadSegmentation,
) 

from readii.metadata import (
    saveDataframeCSV,
)
from readii.negative_controls import (
    applyNegativeControl,
)

from typing import Optional, Any
from collections import OrderedDict

import logging

from readii.utils.logging_config import logger

# disable radiomics logging
logging.getLogger("radiomics").setLevel(logging.ERROR)

def singleRadiomicFeatureExtraction(
    ctImage: sitk.Image,
    roiImage: sitk.Image,
    pyradiomicsParamFilePath: Optional[str] = "./src/readii/data/default_pyradiomics.yaml",
    negativeControl: Optional[str] = None,
    randomSeed: Optional[int] = None,
) -> OrderedDict[Any, Any]:
    """Function to perform radiomic feature extraction for a single CT image and its corresponding segmentation.
       CT and segmentation will be aligned and cropped prior to extraction.

    Parameters
    ----------
    ctImage : sitk.Image
        CT image to perform feature extraction on. Will be cropped and potentially generate a negative control (see negativeControl arg)
    roiImage : sitk.Image
        Region of interest (ROI) to extract radiomic features from within the CT.
    pyradiomicsParamFilePath : str
        Path to file containing configuration settings for pyradiomics feature extraction. Will use the provided config file in 'data/' by default if no file passed in.
    negativeControl : str
        Name of negative control to generate from the CT to perform feature extraction on. If set to None, will extract features from original CT image.
    randomSeed : int
        Value to set random seed with for negative control creation to be reproducible.
        
    Returns
    -------
    OrderedDict[Any, Any]
        Dictionary containing image metadata, versions for key packages used for extraction, and radiomic features
    """
    # If no pyradiomics paramater file passed, use default
    if pyradiomicsParamFilePath == None:
        pyradiomicsParamFilePath = "./src/readii/data/default_pyradiomics.yaml"

    # In case segmentation contains extra axis, flatten to 3D by removing it
    roiImage = flattenImage(roiImage)
    # Segmentation has different origin, align it to the CT for proper feature extraction
    alignedROIImage = alignImages(ctImage, roiImage)

    # Get pixel value for the segmentation
    segmentationLabel: int = getROIVoxelLabel(alignedROIImage)

    # Check that CT and segmentation correspond, segmentationLabel is present, and dimensions match
    segBoundingBox, correctedROIImage = imageoperations.checkMask(ctImage, alignedROIImage, label=segmentationLabel)
    
    # Update the ROI image if a correction was generated by checkMask
    if correctedROIImage is not None:
        alignedROIImage = correctedROIImage

    try:
        if negativeControl != None:
            logger.info(f"Generating {negativeControl} negative control for CT.")
            # Split negative control type into negative control and region of interest
            try:
                if "non_roi" in negativeControl:
                    negativeControlType =  negativeControl.rsplit("_", 2)[0]
                    negativeControlRegion = "non_roi"
                else:
                    negativeControlComponents = negativeControl.rsplit("_", 1)
                    negativeControlType = negativeControlComponents[0]
                    negativeControlRegion = negativeControlComponents[1]
                logger.debug(f"Negative control region: {negativeControlRegion}")
                logger.debug(f"Negative control type: {negativeControlType}")
                # Make negative control version of ctImage
                ctImage_nc: sitk.Image | np.ndarray = applyNegativeControl(
                    baseImage=ctImage,
                    negativeControlType=negativeControlType,
                    negativeControlRegion=negativeControlRegion,
                    roiMask=alignedROIImage,
                    randomSeed=randomSeed
                )
            except Exception as e:
                logger.error(f"Error generating {negativeControl} negative control for CT: {e}")
                raise

            croppedCT, croppedROI = imageoperations.cropToTumorMask(ctImage_nc, alignedROIImage, segBoundingBox)
        else:
            # Crop the image and mask to a bounding box around the mask to reduce volume size to process
            croppedCT, croppedROI = imageoperations.cropToTumorMask(ctImage, alignedROIImage, segBoundingBox)
    except Exception as e:
        logger.error(f"Error cropping image and mask: {e}")
        raise

    # Load PyRadiomics feature extraction parameters to use
    # Initialize feature extractor with parameters
    try:
        logger.info("Setting up Pyradiomics feature extractor...")
        featureExtractor = featureextractor.RadiomicsFeatureExtractor(pyradiomicsParamFilePath)
    except OSError as e:
        logger.error(f"Supplied pyradiomics parameter file {pyradiomicsParamFilePath} does not exist or is not at that location: {e}")
        raise

    try:
        logger.info("Starting radiomic feature extraction...")
        # Extract radiomic features from CT with segmentation as mask
        idFeatureVector = featureExtractor.execute(croppedCT, croppedROI, label=segmentationLabel)
    except Exception as e:
        logger.error(f"An error occurred while extracting radiomic features: {e}")
        raise

    return idFeatureVector


def radiomicFeatureExtraction(
    imageMetadataPath: str,
    imageDirPath: str,
    roiNames: Optional[str] = None,
    pyradiomicsParamFilePath: Optional[str] = "src/readii/data/default_pyradiomics.yaml",
    outputDirPath: Optional[str] = None,
    negativeControl: Optional[str] = None,
    randomSeed: Optional[int] = None,
    parallel: bool = False,
    keep_running: bool = False
) -> pd.DataFrame:
    """Perform radiomic feature extraction using PyRadiomics on CT images with a corresponding segmentation.
       Utilizes outputs from med-imagetools (https://github.com/bhklab/med-imagetools) run on the image dataset.

    Parameters
    ----------
    imageMetadataPath : str
        Path to csv file created by matchCTtoSegmentation function that contains a CT and matching segmentation in each row.
    imageDirPath : str
        Path to the directory containing the directory of CT and segmentation images. This directory should contain the .imgtools directory from the med-imagetools run
        and be the same as the input path used in med-imagetools
    roiNames : str
        Name pattern for the ROIs to load for the RTSTRUCTs. Can be None for DICOM SEG segmentations.
    pyradiomicsParamFilePath : str
        Path to file containing configuration settings for pyradiomics feature extraction. Will use the provided config file in 'data/' by default if no file passed in.
    outputDirPath : str
        Path to directory save the dataframe of extracted features to as a csv
    negativeControl : str
        Name of negative control to generate from the CT to perform feature extraction on. If set to None, will extract features from original CT image.
    randomSeed : int
        Value to set random seed with for negative control creation to be reproducible.
    parallel : bool
        Flag to decide whether to run extraction in parallel.
    keep_running : bool
        Flag to keep pipeline running even when feature extraction for a patient fails.
    Returns
    -------
    pd.DataFrame
        Dataframe containing the image metadata and extracted radiomic features.
    """

    # If no pyradiomics paramater file passed, use default
    if pyradiomicsParamFilePath == None:
        pyradiomicsParamFilePath = "./src/readii/data/default_pyradiomics.yaml"

    # Load in summary file generated by radiogenomic_pipeline
    pdImageInfo = pd.read_csv(imageMetadataPath, header=0)

    # Get array of unique CT series' IDs to iterate over
    ctSeriesIDList = pdImageInfo["series_CT"].unique()

    def featureExtraction(ctSeriesID):
        """Function to extract PyRadiomics features for all ROIs present in a CT. Inner function so it can be run in parallel with joblib."""
        # Get all info rows for this ctSeries
        ctSeriesInfo = pdImageInfo.loc[pdImageInfo["series_CT"] == ctSeriesID]
        patID = ctSeriesInfo.iloc[0]["patient_ID"]

        logger.info(f"Processing {patID}")
        plogger = logger.bind(patient=patID)

        try:
            # Get absolute path to CT image files
            ctDirPath = os.path.join(imageDirPath, ctSeriesInfo.iloc[0]["folder_CT"])
            # Load CT by passing in specific series to find in a directory
            ctImage = read_dicom_series(path=ctDirPath, series_id=ctSeriesID)

            # Get list of segmentations to iterate over
            segSeriesIDList = ctSeriesInfo["series_seg"].unique()

            # Initialize dictionary to store radiomics data for each segmentation (image metadata + features)
            ctAllData = []

            # Loop over every segmentation associated with this CT - only loading CT once
            for segCount, segSeriesID in enumerate(segSeriesIDList):
                segSeriesInfo = ctSeriesInfo.loc[ctSeriesInfo["series_seg"] == segSeriesID]

                # Check that a single segmentation file is being processed
                if len(segSeriesInfo) > 1:
                    # Check that if there are multiple rows that it's not due to a CT with subseries (this is fine, the whole series is loaded)
                    if not segSeriesInfo.duplicated(subset=["series_CT"], keep=False).all():
                        raise RuntimeError(
                            "Some kind of duplication of segmentation and CT matches not being caught. Check seg_and_ct_dicom_list in radiogenomic_output."
                        )

                # Get absolute path to segmentation image file
                segFilePath = os.path.join(
                    imageDirPath, segSeriesInfo.iloc[0]["file_path_seg"]
                )
                # Get dictionary of ROI sitk Images for this segmentation file
                segImages = loadSegmentation(
                    segFilePath,
                    modality=segSeriesInfo.iloc[0]["modality_seg"],
                    baseImageDirPath=ctDirPath,
                    roiNames=roiNames,
                )

                # Check that this series has ROIs to extract from (dictionary isn't empty)
                if not segImages:
                    log_msg = f"CT {ctSeriesID} and segmentation {segSeriesID} has no ROIs or no ROIs with the label {roiNames}. Moving to next segmentation."
                    plogger.warning(log_msg)

                else:
                    # Loop over each ROI contained in the segmentation to perform radiomic feature extraction
                    for roiCount, roiImageName in enumerate(segImages):
                        # ROI counter for image metadata output
                        roiNum = roiCount + 1

                        # Extract features listed in the parameter file
                        plogger.info(
                            f"Calculating radiomic features for segmentation: {roiImageName}"
                        )

                        # Get sitk Image object for this ROI
                        roiImage = segImages[roiImageName]

                        # Check if segmentation just has an extra axis with a size of 1 and remove it
                        if roiImage.GetDimension() > 3 and roiImage.GetSize()[3] == 1:
                            roiImage = flattenImage(roiImage)

                        # Check that image and segmentation mask have the same dimensions
                        if ctImage.GetSize() != roiImage.GetSize():
                            plogger.error(
                                "Segmentation mask is smaller than the CT.",
                                ct_size=ctImage.GetSize(),
                                roi_size=roiImage.GetSize(),
                                ctSeriesID=ctSeriesID,
                                segSeriesID=segSeriesID,
                            )
                            raise RuntimeError(
                                "CT and ROI dimensions do not match. "
                                "This is likely due to a segmentation mask that is smaller than the CT. "
                                "The padding workaround has been disabled. Please raise an issue on GitHub"
                                " to discuss migration options."
                            )

                        # Extract radiomic features from this CT/segmentation pair
                        idFeatureVector = singleRadiomicFeatureExtraction(
                            ctImage,
                            roiImage=roiImage,
                            pyradiomicsParamFilePath=pyradiomicsParamFilePath,
                            negativeControl=negativeControl,
                            randomSeed=randomSeed,
                        )

                        # Create dictionary of image metadata to append to front of output table
                        sampleROIData = {
                            "patient_ID": patID,
                            "study_description": segSeriesInfo.iloc[0][
                                "study_description_CT"
                            ],
                            "series_UID": segSeriesInfo.iloc[0]["series_CT"],
                            "series_description": segSeriesInfo.iloc[0][
                                "series_description_CT"
                            ],
                            "image_modality": segSeriesInfo.iloc[0]["modality_CT"],
                            "instances": segSeriesInfo.iloc[0]["instances_CT"],
                            "seg_series_UID": segSeriesInfo.iloc[0]["series_seg"],
                            "seg_modality": segSeriesInfo.iloc[0]["modality_seg"],
                            "seg_ref_image": segSeriesInfo.iloc[0]["reference_ct_seg"],
                            "roi": roiImageName,
                            "roi_number": roiNum,
                            "negative_control": negativeControl,
                        }

                        # Concatenate image metadata with PyRadiomics features
                        sampleROIData.update(idFeatureVector)
                        # Store this ROI's info in the segmentation level list
                        ctAllData.append(sampleROIData)

            return ctAllData
            ###### END featureExtraction #######
        except RuntimeError as re:
            raise re
        except Exception as e:
            plogger.error(
                f"Error processing patient", ctSeriesID=ctSeriesID
            )  #  {patID}, series {ctSeriesID}: {e}
            if keep_running:
                # Log the error and continue without raising the exception
                return
            raise e

    # Extract radiomic features for each CT, get a list of dictionaries
    # Each dictioary contains features for each ROI in a single CT
    if not parallel:
        # Run feature extraction over samples in sequence - will be slower
        features = [featureExtraction(ctSeriesID) for ctSeriesID in ctSeriesIDList]
    else:
        # Run feature extraction in parallel
        features = Parallel(n_jobs=-1, require="sharedmem")(
            delayed(featureExtraction)(ctSeriesID) for ctSeriesID in ctSeriesIDList
        )

    logger.info("Finished feature extraction.")

    # Filter out None and ensure each result is a list (even if it's empty)
    features = [f if isinstance(f, list) else [f] for f in features if f is not None]

    # Flatten the list of dictionaries (happens when there are multiple ROIs or SEGs associated with a single CT)
    flatFeatures = list(chain.from_iterable(features))
    # Convert list of feature sets into a pandas dataframe to save out
    featuresTable = pd.DataFrame(flatFeatures)

    if outputDirPath != None:
        if not os.path.exists(outputDirPath):
            logger.info(f"Directory {outputDirPath} does not exist. Creating...")
            os.makedirs(outputDirPath)
        else:
            logger.warning(f"Directory {outputDirPath} already exists. Will overwrite contents.")

        logger.info("Saving output to file...")
        datasetName = imageMetadataPath.partition("match_list_")[2]
        # Setup output file name with the dataset name as a suffix
        if negativeControl == None:
            outFileName = "radiomicfeatures_original_" + datasetName
        else:
            # Add negative control identifier to output file name
            outFileName = "radiomicfeatures_" + negativeControl + "_" + datasetName

        # Join outputDirPath, a features directory, and the output file name
        outputFilePath = os.path.join(outputDirPath, "features/", outFileName)
        # Save out the features
        saveDataframeCSV(featuresTable, outputFilePath)

    return featuresTable
