from contextlib import ExitStack
from ewoksid02.tasks.id02processingtask import ID02ProcessingTask
from ewoksid02.utils.normalization import (
    normalize_dataset,
)
from ewoksid02.utils.pyfai import get_persistent_azimuthal_integrator


class NormalizationTask(
    ID02ProcessingTask,
    optional_input_names=[
        "filename_mask",
        "filename_dark",
        "filename_flat",
        "Dummy",
        "DDummy",
        "NormalizationFactor",
        "polarization_factor",
        "polarization_axis_offset",
        "Center_1",
        "Center_2",
        "PSize_1",
        "PSize_2",
        "BSize_1",
        "BSize_2",
        "SampleDistance",
        "WaveLength",
        "DetectorRotation_1",
        "DetectorRotation_2",
        "DetectorRotation_3",
        "pin_monitor",
        "header_pin_monitor",
        "variance_formula",
        "algorithm",
        "dark_filter",
        "dark_filter_quantil_lower",
        "dark_filter_quantil_upper",
    ],
):
    """The `NormalizationTask` class is responsible for normalizing datasets in the ID02 SAXS pipeline.
    It extends the `ID02ProcessingTask` class and provides additional functionality to apply a standard pyFAI normalization:
    - Methods to read monitor values from the metadata file or from blissdata
    - Methods to read normalization parameters from the metadata file or from the headers
    - Methods to cache pyFAI azimuthal integrator and apply normalization
    - Applies corrections such as masking, dark frame subtraction, flat field correction, and polarization adjustments.

    Optional Inputs:
        - filename_mask (str): Path to the mask file for correcting detector gaps or bad pixels.
        - filename_dark (str): Path to the file for dark current correction.
        - filename_flat (str): Path to the file for flat field correction.
        - Dummy (float): Value to replace invalid pixels in the dataset.
        - DDummy (float): Tolerance for dummy pixel replacement.
        - NormalizationFactor (float): Factor for normalizing the dataset.
        - polarization_factor (float): Factor for polarization correction.
        - polarization_axis_offset (float): Axis for polarization correction.
        - Center_1 (float): Beam center in the first dimension.
        - Center_2 (float): Beam center in the second dimension.
        - PSize_1 (float): Pixel size 1.
        - PSize_2 (float): Pixel size 2.
        - BSize_1 (float): Pixel binning 1.
        - BSize_2 (float): Pixel binning 2.
        - SampleDistance (float): Sample to detector distance in meters.
        - WaveLength (float): Wavelength of beam in meters.
        - DetectorRotation_1 (float): rot2 of pyFAI.
        - DetectorRotation_2 (float): rot1 of pyFAI.
        - DetectorRotation_3 (float): rot3 of pyFAI.
        - pin_monitor (str): Pin to the monitor stream.
        - header_pin_monitor (str): Header key used to monitor values.
        - variance_formula (str): Formula for calculating variance in the dataset.
        - algorithm (str): Implementation to perform the normalization (cython or cupy).
        - dark_filter (str): Filter to use in case the dark-current file is a multi-frame scan file.
        - dark_filter_quantil_lower (str): In case the dark_filter is quantil, lower limit.
        - dark_filter_quantil_upper (str): In case the dark_filter is quantil, upper limit.
    """

    def run(self):
        super().run(processing_type="norm")

    def get_processing_inputs(self) -> dict:
        monitor_values = self.get_monitor_1_values()

        params_azimuthalintegrator = {
            "Center_1": self.get_parameter("Center_1"),
            "Center_2": self.get_parameter("Center_2"),
            "PSize_1": self.get_parameter("PSize_1"),
            "PSize_2": self.get_parameter("PSize_2"),
            "SampleDistance": self.get_parameter("SampleDistance"),
            "WaveLength": self.get_parameter("WaveLength"),
            "BSize_1": self.get_parameter("BSize_1"),
            "BSize_2": self.get_parameter("BSize_2"),
            "DetectorRotation_1": self.get_parameter("DetectorRotation_1"),
            "DetectorRotation_2": self.get_parameter("DetectorRotation_2"),
            "DetectorRotation_3": self.get_parameter("DetectorRotation_3"),
        }

        azimuthal_integrator = get_persistent_azimuthal_integrator(
            data_signal_shape=self.dataset_signal[0].shape,
            **params_azimuthalintegrator,
        )

        params_normalization = {
            "filename_mask": self.get_input_value(
                "filename_mask",
                self.get_mask_gaps_filename(),
            ),
            "filename_dark": self.get_input_value(
                "filename_dark",
                self.get_dark_filename(),
            ),
            "filename_flat": self.get_input_value(
                "filename_flat",
                self.get_flat_filename(),
            ),
            "Dummy": self.get_parameter("Dummy"),
            "DDummy": self.get_parameter("DDummy"),
            "variance_formula": self.get_parameter("variance_formula"),
            "binning": (
                self.get_parameter("BSize_1", to_integer=True),
                self.get_parameter("BSize_2", to_integer=True),
            ),
            "algorithm": self.get_input_value("algorithm", "cython"),
            "datatype": self.get_input_value("datatype", "float32"),
            "monitor_values": monitor_values,
            "azimuthal_integrator": azimuthal_integrator,
            "NormalizationFactor": self.get_parameter("NormalizationFactor"),
            **params_azimuthalintegrator,
        }

        do_polarization = self.get_parameter(
            "polarization_factor"
        ) or self.get_parameter("do_polarization")
        if do_polarization:
            params_normalization["polarization_factor"] = self.get_parameter(
                "polarization_factor"
            )
            params_normalization["polarization_axis_offset"] = self.get_parameter(
                "polarization_axis_offset", to_integer=True
            )

        if params_normalization.get("filename_dark"):
            params_normalization["dark_filter"] = self.get_parameter(
                "dark_filter", defaut="quantil"
            )
            params_normalization["dark_filter_quantil_lower"] = self.get_parameter(
                "dark_filter_quantil_lower", defaut=0.1
            )
            params_normalization["dark_filter_quantil_upper"] = self.get_parameter(
                "dark_filter_quantil_upper", defaut=0.9
            )

        return params_normalization

    def process(self) -> None:
        do_process = super().process()
        if do_process is False:
            return

        with ExitStack() as stack:
            if len(self.dataset_signal) == 0:
                return

            self.bench_process = self.Benchmark(
                nb_frames=len(self.dataset_signal), benchmark_name="processing"
            )
            stack.enter_context(self.bench_process)

            processing_params = self.get_processing_inputs()
            self.processing_params = processing_params

            self.log_info("Performing normalization...")

            (
                dataset_signal_normalized,
                dataset_variance_normalized,
                dataset_sigma_normalized,
            ) = normalize_dataset(
                dataset_signal=self.dataset_signal,
                **processing_params,
            )
            self.log_debug("Processing done")
            self.outputs.dataset_signal = dataset_signal_normalized
            self.outputs.dataset_variance = dataset_variance_normalized
            self.outputs.dataset_sigma = dataset_sigma_normalized

        self._log_benchmark(self.bench_process)

    def save(self) -> None:
        do_save = super().save()
        if not do_save:
            return

    def get_monitor_1_values(
        self,
    ):
        """Generic method to read monitor_values online (from blissdata streams)
        or offline (from a group in the scalers file)
        """
        stream_monitor_info = self.get_stream_monitor_1()
        monitor_values, _, __ = self._read_from_stream(
            stream_object=stream_monitor_info.get("stream"),
            slice_init=self.index_range_last[0],
            slice_end=self.index_range_last[-1],
        )
        return monitor_values
