import numpy
from ewoksid02.tasks.id02processingtask import ID02ProcessingTask
from ewoksid02.utils.caving import process_data_caving
from ewoksid02.headers import HS32

DEFAULT_ALGORITHM_CAVE = "cupy"


class CavingTask(
    ID02ProcessingTask,
    optional_input_names=[
        "Dummy",
        "Center_1",
        "Center_2",
        "filename_mask_static",
        "filename_mask_reference",
        "flip_caving",
        "algorithm_cave",
    ],
):
    """The `CavingTask` class is responsible for applying a "caving" operation to datasets in the ID02 SAXS pipeline.
    The caving is applied to those pixels whose intensity matches with a dummy value + those pixels provided by a mask file.

    Optional Inputs:
        - Dummy (float): Value to replace invalid pixels in the dataset.
        - Center_1 (float): Beam center in the first dimension.
        - Center_2 (float): Beam center in the second dimension.
        - filename_mask_static (str): Path to the mask file used for the caving operation.
        - filename_mask_reference (str): Path to the reference mask file (kind of a negative mask).
        - flip_caving (bool): Cave the image with its flipped projection, both horizontal and vertical. WARNING: it is physically not correct!
        - algorithm_cave (str): Implementation to perform the caving (numpy or cupy).
    """

    def get_processing_parameters(self) -> dict:
        if not self.processing_params:
            self.processing_params = {
                "Dummy": self.get_parameter("Dummy"),
                "Center_1": self.get_parameter("Center_1"),
                "Center_2": self.get_parameter("Center_2"),
                "filename_mask_static": self.get_input_value(
                    "filename_mask_static", None
                ),
                "filename_mask_reference": self.get_input_value(
                    "filename_mask_reference", None
                ),
                "flip_caving": self.get_input_value(
                    "flip_caving",
                    bool(self.headers.get("cave_flip")),
                ),
                "algorithm_cave": self.get_input_value(
                    "algorithm_cave", DEFAULT_ALGORITHM_CAVE
                ),
            }
        return self.processing_params

    def process(self) -> None:
        processing_params = self.get_processing_parameters()
        caved_dataset = process_data_caving(
            data=self.dataset_signal,
            **processing_params,
        )

        if not self.get_input_value("save_variance", False):
            caved_dataset_variance = None
        elif self.dataset_variance is not None:
            caved_dataset_variance = process_data_caving(
                data=self.dataset_variance,
                **processing_params,
            )
        else:
            caved_dataset_variance = numpy.zeros_like(caved_dataset)

        if self.dataset_sigma is not None:
            caved_dataset_sigma = process_data_caving(
                data=self.dataset_sigma,
                **processing_params,
            )
        else:
            caved_dataset_sigma = numpy.zeros_like(caved_dataset)
        self.outputs.dataset_signal = caved_dataset
        self.outputs.dataset_variance = caved_dataset_variance
        self.outputs.dataset_sigma = caved_dataset_sigma


class CavingGapsTask(
    CavingTask,
):
    """The `CavingGapsTask` inherits 'CavingTask'.
    Here, it will take the parameters from the headers to perform the caving operation
    and cover the gaps in the dataset and avoid the beamstop. Hence, no explicit mask filenames are needed.
    """

    def get_processing_parameters(self) -> dict:
        if not self.processing_params:
            self.processing_params = {
                "Dummy": self.get_parameter("Dummy"),
                "Center_1": self.get_parameter("Center_1"),
                "Center_2": self.get_parameter(
                    "Center_2",
                ),
                "filename_mask_static": self.get_input_value(
                    "filename_mask_static",
                    HS32.get_mask_gaps_filename(headers=self.headers),
                ),
                "filename_mask_reference": self.get_input_value(
                    "filename_mask_reference",
                    HS32.get_mask_beamstop_filename(headers=self.headers),
                ),
                "flip_caving": self.get_input_value(
                    "flip_caving",
                    bool(self.headers.get("nw_cave_flip")),
                ),
                "algorithm_cave": self.get_input_value(
                    "algorithm_cave", DEFAULT_ALGORITHM_CAVE
                ),
            }
        return self.processing_params


class CavingBeamstopTask(
    CavingTask,
):
    """The `CavingBeamstopTask` inherits 'CavingTask'.
    Here, it will take the parameters from the headers to perform the caving operation
    and cover the beamstop. Hence, no explicit mask filename is needed.
    """

    def get_processing_parameters(self) -> dict:
        if not self.processing_params:
            self.processing_params = {
                "Dummy": self.get_parameter("Dummy"),
                "Center_1": self.get_parameter("Center_1"),
                "Center_2": self.get_parameter("Center_2"),
                "filename_mask_static": self.get_input_value(
                    "filename_mask_static",
                    HS32.get_mask_beamstop_filename(headers=self.headers),
                ),
                "filename_mask_reference": None,
                "flip_caving": self.get_input_value(
                    "flip_caving",
                    bool(self.headers.get("nw_cave_flip")),
                ),
                "algorithm_cave": self.get_input_value(
                    "algorithm_cave", DEFAULT_ALGORITHM_CAVE
                ),
            }
        return self.processing_params
