from typing import (
    Optional,
    Tuple,
    Union)
from copy import deepcopy

from pyradise.data import (
    Subject,
    IntensityImage,
    SegmentationImage,
    TransformInfo,
    Modality,
    OrganAnnotatorCombination,
    seq_to_modalities,
    seq_to_organ_annotator_combinations)
from pyradise.process import (
    Filter,
    FilterParams)


__all__ = ['PlaybackTransformTapeFilterParams', 'PlaybackTransformTapeFilter']


class PlaybackTransformTapeFilterParams(FilterParams):
    """A filter parameter class for the :class:`~pyradise.process.invertibility.PlaybackTransformTapeFilter` class.

    Args:
        modalities (Optional[Tuple[Union[str, Modality], ...]]): A tuple of modalities for which the transform tape
         should be played back. If None, the transform tape will be played back for all modalities (default: None).
        organ_annotator_combinations (Optional[Tuple[Union[Tuple[str, str], OrganRaterCombination], ...]]): A tuple of
         organ-annotator combinations for which the transform tape should be played back. If None, the transform tape
         will be played back for all organ-annotator combinations (default: None).
    """

    def __init__(self,
                 modalities: Optional[Tuple[Union[str, Modality], ...]] = None,
                 organ_annotator_combinations: Optional[Tuple[Union[Tuple[str, str],
                                                                    OrganAnnotatorCombination], ...]] = None
                 ) -> None:
        super().__init__()

        if modalities is not None:
            self.modalities = seq_to_modalities(modalities)
        else:
            self.modalities = None

        if organ_annotator_combinations is not None:
            self.organ_annotator_combinations = seq_to_organ_annotator_combinations(organ_annotator_combinations)
        else:
            self.organ_annotator_combinations = organ_annotator_combinations


class PlaybackTransformTapeFilter(Filter):
    """A filter class for playing back the transform tape of specific or all :class:`~pyradise.data.image.Image`
    instances of the provided :class:`~pyradise.data.subject.Subject` instance.

    This filter is helpful for restoring the spatial properties of the loaded data such that the output data of the
    processing pipeline has identical spatial properties as the input data.
    """

    @staticmethod
    def is_invertible() -> bool:
        """Return False because the filter is not invertible.

        Returns:
            bool: False.
        """
        return False

    def execute(self,
                subject: Subject,
                params: Optional[PlaybackTransformTapeFilterParams] = None
                ) -> Subject:
        """Execute the filter on the provided :class:`~pyradise.data.subject.Subject` instance.

        Args:
            subject (Subject): The :class:`~pyradise.data.subject.Subject` instance.
            params (Optional[FilterParams]): The filter parameters.

        Returns:
            Subject: The :class:`~pyradise.data.subject.Subject` instance with played back images.
        """
        original_images = []
        changed_images = []

        images = subject.get_images()
        for image in images:

            # exclude images not matching the provided criteria
            if isinstance(image, IntensityImage) and params.modalities is not None:
                if not image.get_modality() in params.modalities:
                    continue

            if isinstance(image, SegmentationImage) and params.organ_annotator_combinations is not None:
                if not image.get_organ_annotator_combination() in params.organ_annotator_combinations:
                    continue

            # copy the original subject
            temp_subject = deepcopy(subject)
            transform_tape = image.get_transform_tape()
            transform_infos = transform_tape.get_recorded_elements(reverse=True)

            # play back the transform tape
            for transform_info in transform_infos:
                filter_ = transform_info.get_filter()
                if not filter_.is_invertible():
                    continue

                temp_subject = filter_.execute_inverse(temp_subject, transform_info)

            # collect the modified images
            changed_image_candidates = [img for img in temp_subject.get_images() if img == image]
            changed_images.append(changed_image_candidates[0])
            original_images.append(image)

        # replace the original images with the modified images and reset the transform tape
        for original_image, changed_image in zip(original_images, changed_images):

            original_image.set_image_data(changed_image.get_image_data())
            original_image.get_transform_tape().reset()

        return subject

    def execute_inverse(self,
                        subject: Subject,
                        transform_info: TransformInfo
                        ) -> Subject:
        """Return the provided :class:`~pyradise.data.subject.Subject` instance without any processing because
        :class:`~pyradise.data.taping.TransformTape` playback is not invertible.

        Args:
            subject (Subject): The :class:`~pyradise.data.subject.Subject` instance.
            transform_info (TransformInfo): The :class:`~pyradise.data.taping.TransformInfo` instance.

        Returns:
            Subject: The provided :class:`~pyradise.data.subject.Subject` instance.
        """
        return subject
