from ewoksid02.tasks.id02processingtask import ID02ProcessingTask
from ewoksid02.utils.normalization import (
    normalize_dataset,
)
from ewoksid02.utils.pyfai import get_persistent_azimuthal_integrator
from ewoksid02.headers import HS32

DEFAULT_ALGORITHM_NORMALIZATION = "cupy"


class NormalizationTask(
    ID02ProcessingTask,
    optional_input_names=[
        "filename_mask_normalization",
        "filename_dark_normalization",
        "filename_flat_normalization",
        "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_normalization",
        "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_normalization (str): Path to the mask file for correcting detector gaps or bad pixels.
        - filename_dark_normalization (str): Path to the file for dark current correction.
        - filename_flat_normalization (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_normalization (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 get_processing_parameters(self) -> dict:
        if not self.processing_params:
            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_normalization": self.get_input_value(
                    "filename_mask_normalization",
                    HS32.get_mask_gaps_filename(headers=self.headers),
                ),
                "filename_dark_normalization": self.get_input_value(
                    "filename_dark_normalization",
                    HS32.get_dark_filename(headers=self.headers, missing_ok=True),
                ),
                "filename_flat_normalization": self.get_input_value(
                    "filename_flat_normalization",
                    HS32.get_flat_filename(headers=self.headers),
                ),
                "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_normalization": self.get_input_value(
                    "algorithm_normalization", DEFAULT_ALGORITHM_NORMALIZATION
                ),
                "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_normalization"):
                params_normalization["dark_filter"] = self.get_parameter(
                    "dark_filter", default="quantil"
                )
                params_normalization["dark_filter_quantil_lower"] = self.get_parameter(
                    "dark_filter_quantil_lower", default=0.1
                )
                params_normalization["dark_filter_quantil_upper"] = self.get_parameter(
                    "dark_filter_quantil_upper", default=0.9
                )
            self.processing_params = params_normalization
        return self.processing_params

    def process(self) -> None:
        (
            dataset_signal_normalized,
            dataset_variance_normalized,
            dataset_sigma_normalized,
        ) = normalize_dataset(
            dataset_signal=self.dataset_signal,
            **self.get_processing_parameters(),
        )
        self.outputs.dataset_signal = dataset_signal_normalized
        self.outputs.dataset_variance = dataset_variance_normalized
        self.outputs.dataset_sigma = dataset_sigma_normalized

    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_name, stream_array = self.get_stream_monitor_1()
        monitor_values, slice_init, slice_end = self._read_from_stream(
            stream_sliceable=stream_array,
            stream_name=stream_name,
            slice_init=self.index_range_last[0],
            slice_end=self.index_range_last[-1],
            datatype="float64",
            timeout=1,
        )
        return monitor_values
