import os
from typing import Optional

import pyradise.data as ps_data
import pyradise.fileio as ps_io

class ExampleModalityExtractor(ps_io.ModalityExtractor):

    def extract_from_dicom(self,
                           path: str
                           ) -> Optional[ps_data.Modality]:
        # Extract the necessary attributes from the DICOM file
        tags = (ps_io.Tag((0x0008, 0x0060)),  # Modality
                ps_io.Tag((0x0008, 0x103e)))  # Series Description
        dataset_dict = self._load_dicom_attributes(tags, path)

        # Identify the modality rule-based
        modality = dataset_dict.get('Modality', {}).get('value', None)
        series_desc = dataset_dict.get('Series Description', {}).get('value', '')
        if modality == 'MR':
            if 't1' in series_desc.lower():
                return ps_data.Modality('T1')
            elif 't2' in series_desc.lower():
                return ps_data.Modality('T2')
            else:
                return None
        else:
            return None

    def extract_from_path(self,
                          path: str
                          ) -> Optional[ps_data.Modality]:
        # Identify the discrete image file's modality rule-based
        filename = os.path.basename(path)

        # Check if the image contains an img prefix
        # (i.e., it is a intensity image)
        if not filename.startswith('img'):
            return None

        # Check if the image contains a modality search string
        if 'T1' in filename:
            return ps_data.Modality('T1')
        elif 'T2' in filename:
            return ps_data.Modality('T2')
        else:
            return None


class ExampleOrganExtractor(ps_io.OrganExtractor):

    def extract(self,
                path: str
                ) -> Optional[ps_data.Organ]:
        # Identify the discrete image file's organ rule-based
        filename = os.path.basename(path)

        # Check if the image contains a seg prefix
        # (i.e., it is a segmentation)
        if not filename.startswith('seg'):
            return None

        # Split the filename for extracting the organ name
        organ_name = filename.split('_')[-1].split('.')[0]
        return ps_data.Organ(organ_name)


class ExampleAnnotatorExtractor(ps_io.AnnotatorExtractor):

    def extract(self,
                path: str
                ) -> Optional[ps_data.Annotator]:
        # Identify the discrete image file's annotator rule-based
        filename = os.path.basename(path)

        # Check if the image contains a seg prefix
        # (i.e., it is a segmentation)
        if not filename.startswith('seg'):
            return None

        # Split the filename for extracting the annotator name
        annotator_name = filename.split('_')[2]
        return ps_data.Annotator(annotator_name)
    
    
def convert_subject_to_dicom_rtss(input_dir_path: str,
                                  output_dir_path: str,
                                  dicom_image_dir_path: str,
                                  use_3d_conversion: bool = True
                                  ) -> None:
    # Specify a reference modalities
    # This is the modality of the DICOM image series that will be
    # referenced in the DICOM-RTSS.
    reference_modality = 'T1'

    # Create the loader
    loader = ps_io.SubjectLoader()

    # Create the writer and specify the output file name of the
    # DICOM-RTSS files
    writer = ps_io.DicomSeriesSubjectWriter()
    rtss_filename = 'rtss.dcm'

    # (optional)
    # Instantiate a new selection to exclude the original DICOM-RTSS SeriesInfo
    # Note: If this is omitted the original DICOM-RTSS will be copied to the
    # corresponding output directory.
    selection = ps_io.NoRTSSInfoSelector()

    # Create the file crawler for the discrete image files and
    # loop through the subjects
    crawler = ps_io.DatasetFileCrawler(input_dir_path,
                                       extension='.nii.gz',
                                       modality_extractor=ExampleModalityExtractor(),
                                       organ_extractor=ExampleOrganExtractor(),
                                       annotator_extractor=ExampleAnnotatorExtractor())
    for series_info in crawler:
        # Load the subject
        subject = loader.load(series_info)

        # Print the progress
        print(f'Converting subject {subject.get_name()}...')

        # Construct the path to the subject's DICOM images
        dicom_subject_path = os.path.join(dicom_image_dir_path, subject.get_name())

        # Construct a DICOM crawler to retrieve the reference
        # DICOM image series info
        dcm_crawler = ps_io.SubjectDicomCrawler(dicom_subject_path,
                                                modality_extractor=ExampleModalityExtractor())
        dicom_series_info = dcm_crawler.execute()

        # (optional)
        # Keep all SeriesInfo entries that do not describe a DICOM-RTSS for loading
        dicom_series_info = selection.execute(dicom_series_info)

        # (optional)
        # Define the metadata for the DICOM-RTSS
        # Note: For some attributes, the value must follow the value
        # representation of the DICOM standard.
        meta_data = ps_io.RTSSMetaData(patient_size='180',
                                       patient_weight='80',
                                       patient_age='050Y',
                                       series_description='Converted from NIfTI')

        # Convert the segmentations to a DICOM-RTSS with standard smoothing settings.
        # For the conversion we can either use a 2D or a 3D algorithm (see API reference
        # for details).
        # Note: Inappropriate smoothing leads to corrupted structures if their size
        # is too small
        if use_3d_conversion:
            conv_conf = ps_io.RTSSConverter3DConfiguration()
        else:
            conv_conf = ps_io.RTSSConverter2DConfiguration()

        converter = ps_io.SubjectToRTSSConverter(subject,
                                                 dicom_series_info,
                                                 reference_modality,
                                                 conv_conf,
                                                 meta_data)
        rtss = converter.convert()

        # Combine the DICOM-RTSS with its output file name
        rtss_combination = ((rtss_filename, rtss),)

        # Write the DICOM-RTSS to a separate subject directory
        # and include the DICOM files crawled before
        # Note: If you want to output just a subset of the
        # original DICOM files you may use additional selectors
        writer.write(rtss_combination, output_dir_path,
                     subject.get_name(), dicom_series_info)

        
# The indicator if the 2D or the 3D conversion algorithm should
# be used.
use_3d_algorithm = True

# The input path pointing to the top-level directory containing the
# NIfTI subject directories
input_dataset_path = '//YOUR/PATH/TO/THE/EXAMPLE/DATA/nifti_data'
input_dataset_path = '/nn/3378417/20241112/output/7_AX_3D_fl3d_mt_FS_+_c_20241112142646.nii.gz'
input_dataset_path = '/nn/3378417/20241112/output'

# The input path pointing to the top-level directory containing the
# DICOM subject directories that will get referenced in the output
# DICOM-RTSS files
dicom_dataset_path = '//YOUR/PATH/TO/THE/EXAMPLE/DATA/dicom_data'
dicom_dataset_path = '/nn/3378417/20241112/CT/8'

# The output path pointing to an empty directory where the output
# will be saved
output_dataset_path = '//YOUR/PATH/TO/THE/OUTPUT/DIRECTORY/'
output_dataset_path = '0'

# Execution of the conversion procedure
convert_subject_to_dicom_rtss(input_dataset_path,
                              output_dataset_path,
                              dicom_dataset_path,
                              use_3d_algorithm)
