from token import OP
from venv import logger
from imgtools.io import read_dicom_series
from itertools import chain
from joblib import Parallel, delayed
from radiomics import featureextractor, imageoperations, logging

import os
import pandas as pd
import numpy as np
import SimpleITK as sitk

from readii.image_processing import (
    flattenImage, 
    alignImages, 
    padSegToMatchCT, 
    getROIVoxelLabel, 
    displayImageSlice, 
    displayCTSegOverlay, 
    getROICenterCoords, 
    getCroppedImages,    
)

from readii.loaders import (
    loadDicomSITK, 
    loadRTSTRUCTSITK, 
    loadSegmentation,
) 

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

from readii.utils import get_logger

from typing import Optional, Any
from collections import OrderedDict

logger = get_logger()

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

    if negativeControl != None:
        logger.info(f"Generating {negativeControl} negative control for CT.")
        # Split negative control type into negative control and region of interest
        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
        )
        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)

    # 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,
) -> 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.

    Returns
    -------
    pd.DataFrame
        Dataframe containing the image metadata and extracted radiomic features.
    """
    # Setting pyradiomics verbosity lower
    radiomics_logger: logging.Logger = logging.getLogger("radiomics")
    radiomics_logger.setLevel(logging.ERROR)

    # 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}")

        # 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."
                logger.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
                    logger.info(f"Calculating radiomic features for segmentation: {roiImageName}")

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

                    # Exception catch for if the segmentation dimensions do not match that original image
                    try:
                        # 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():
                            # Checking if number of segmentation slices is less than CT
                            if ctImage.GetSize()[2] > roiImage.GetSize()[2]:
                                logger.warning(
                                    f"Slice number mismatch between CT and segmentation for {patID}."
                                    f"ctImage.GetSize(): {ctImage.GetSize()}"
                                    f"roiImage.GetSize(): {roiImage.GetSize()}"
                                    "Padding segmentation to match."
                                )
                                roiImage = padSegToMatchCT(
                                    ctDirPath, segFilePath, ctImage, roiImage
                                )
                                logger.warning(
                                    f"Padded segmentation to match CT for {patID}."
                                    "roiImage.GetSize() after padding: {roiImage.GetSize()}"
                                )
                            else:
                                raise RuntimeError(
                                    "CT and ROI dimensions do not match."
                                )

                    # Catching CT and segmentation size mismatch error
                    except RuntimeError as e:
                        logger.error(str(e))

                    # 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 #######

    # 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:
        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_" + 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
