from contextlib import ExitStack

import numexpr
from pyFAI.units import to_unit
from silx.io.h5py_utils import open_item as open_item_silx

from ewoksid02.tasks.id02processingtask import ID02ProcessingTask
from ewoksid02.utils.pyfai import (
    get_gpu_method,
    get_persistent_azimuthal_integrator,
    _get_persistent_pyfai_worker,
    process_dataset_azim,
    guess_npt2_rad,
)

DEFAULT_NPT_RAD = 1600
DEFAULT_NPT_AZIM = 360


class AzimuthalTask(
    ID02ProcessingTask,
    optional_input_names=[
        "filename_mask",
        "filename_dark",
        "npt2_rad",
        "npt2_azim",
        "unit",
        "Dummy",
        "DDummy",
        "Center_1",
        "Center_2",
        "PSize_1",
        "PSize_2",
        "BSize_1",
        "BSize_2",
        "SampleDistance",
        "WaveLength",
        "DetectorRotation_1",
        "DetectorRotation_2",
        "DetectorRotation_3",
        "method",
        "integration_options",
        "do_variance_formula",
        "variance_formula",
        "save_sum",
    ],
    output_names=[
        "dataset_sum_signal",
        "dataset_sum_normalization",
        "dataset_sum_variance",
        "radial_array",
        "azimuth_array",
    ],
):
    """The `AzimuthalTask` class is responsible for performing azimuthal integration on datasets in the ID02 SAXS pipeline.
    It extends the `ID02ProcessingTask` class and provides additional functionality for handling azimuthal integration-specific
    inputs and processing logic using the pyFAI library.

    Optional Inputs:
        - filename_mask (str): Path to the mask file for masking invalid pixels.
        - filename_dark (str): Path to the file with a dark-current correction.
        - npt2_rad (int): Number of radial bins for the integration.
        - npt2_azim (int): Number of azimuthal bins for the integration.
        - unit (str): Unit for the radial axis (e.g., "q_nm^-1").
        - Dummy (float): Value to replace invalid pixels in the dataset.
        - DDummy (float): Tolerance for dummy pixel replacement.
        - 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.
        - method (str): Integration method to be used by pyFAI (e.g., "bbox", "csr").
        - integration_options (dict): Additional options for pyFAI integration.
        - do_variance_formula (bool): Flag to enable variance calculation using a formula. Default is `False`.
        - variance_formula (str): Formula for calculating variance in the dataset.
        - save_sum (bool): To save or not the arrays sum_signal, sum_normalization and sum_variance.
    Outputs:
        - dataset_sum_signal (numpy.ndarray): Summed signal dataset after azimuthal integration.
        - dataset_sum_normalization (numpy.ndarray): Summed normalization dataset after azimuthal integration.
        - dataset_sum_variance (numpy.ndarray): Summed variance dataset after azimuthal integration.
        - radial_array (numpy.ndarray): Radial axis array for the integrated data.
        - azimuth_array (numpy.ndarray): Azimuthal axis array for the integrated data.
    """

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

    def get_processing_inputs(
        self,
    ) -> dict:
        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,
        )

        processing_params = {
            "filename_mask": self.get_input_value(
                "filename_mask",
                self.get_mask_beamstop_filename(),
            ),
            "filename_dark": self.get_input_value(
                "filename_dark",
                self.get_dark_filename(),
            ),
            "Dummy": self.get_parameter("Dummy"),
            "DDummy": self.get_parameter("DDummy"),
            "npt2_rad": self.get_parameter("npt2_rad"),
            "npt2_azim": self.get_parameter("npt2_azim"),
            "unit": self.get_parameter("unit", defaut="q_nm^-1"),
            "method": self.get_input_value("method", get_gpu_method()),
            "integration_options": self.get_input_value("integration_options", {}),
            "do_variance_formula": self.get_input_value("do_variance_formula", False),
            "variance_formula": self.get_parameter("variance_formula"),
            "datatype": self.get_input_value("datatype", "float32"),
            "binning": (
                self.get_parameter(
                    "BSize_1",
                ),
                self.get_parameter(
                    "BSize_2",
                ),
            ),
            "save_variance": self.get_input_value("save_variance", False),
            "save_sum": self.get_input_value("save_sum", False),
            "azimuthal_integrator": azimuthal_integrator,
            **params_azimuthalintegrator,
        }
        if not processing_params.get("npt2_rad"):
            processing_params["npt2_rad"] = guess_npt2_rad(
                azimuthal_integrator=azimuthal_integrator
            )
        if not processing_params.get("npt2_azim"):
            processing_params["npt2_azim"] = DEFAULT_NPT_AZIM
        if not processing_params.get("method"):
            processing_params["method"] = get_gpu_method()
        return processing_params

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

        with ExitStack() as stack:
            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

            if self.dataset_variance is not None:
                dataset_variance = self.dataset_variance
            elif self.dataset_sigma is not None:
                dataset_sigma = self.dataset_sigma  # noqa
                Dummy = processing_params.get("Dummy", 0.0)  # noqa
                dataset_variance = numexpr.evaluate(
                    "where(dataset_sigma <= 0.0, Dummy, dataset_sigma ** 2)"
                )
            else:
                dataset_variance = None

            self.log_info("Performing 2D integration...")

            (
                dataset_signal_azim,
                dataset_variance_azim,
                dataset_sigma_azim,
                dataset_sumsignal_azim,
                dataset_sumnorm_azim,
                dataset_sumvariance_azim,
                array_radial,
                array_azim,
            ) = process_dataset_azim(
                dataset_signal=self.dataset_signal,
                dataset_variance=dataset_variance,
                **processing_params,
            )

            self.outputs.dataset_signal = dataset_signal_azim
            self.outputs.dataset_variance = dataset_variance_azim
            self.outputs.dataset_sigma = dataset_sigma_azim
            self.outputs.dataset_sum_signal = dataset_sumsignal_azim
            self.outputs.dataset_sum_normalization = dataset_sumnorm_azim
            self.outputs.dataset_sum_variance = dataset_sumvariance_azim
            self.outputs.radial_array = array_radial
            self.outputs.azimuth_array = array_azim

        self._log_benchmark(self.bench_process)

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

        filename_output, h5path_nxdata = self.links["destination"]["nxdata"].split("::")
        params = {
            "filename": filename_output,
            "name": h5path_nxdata,
            "retry_timeout": 0.1,
            "mode": "a",
        }
        with open_item_silx(**params) as nexus_data_grp:
            self._update_dataset(
                added_dataset=self.outputs.dataset_sum_signal,
                h5_group=nexus_data_grp,
                h5_dataset_name="sum_signal",
                index_read=self.index_range_last,
            )

            self._update_dataset(
                added_dataset=self.outputs.dataset_sum_normalization,
                h5_group=nexus_data_grp,
                h5_dataset_name="sum_normalization",
                index_read=self.index_range_last,
            )

            self._update_dataset(
                added_dataset=self.outputs.dataset_sum_variance,
                h5_group=nexus_data_grp,
                h5_dataset_name="sum_variance",
                index_read=self.index_range_last,
            )

            unit = self.get_parameter("unit")
            radial_unit = to_unit(unit)

            # Update radial and azimuthal arrays only once
            if radial_unit and self.outputs.radial_array is not None:
                if radial_unit.short_name not in nexus_data_grp:
                    radial_dset = nexus_data_grp.create_dataset(
                        name=radial_unit.short_name,
                        data=self.outputs.radial_array,
                    )
                    radial_dset.attrs["axis"] = "3"
                    radial_dset.attrs["interpretation"] = "scalar"
                    radial_dset.attrs["unit"] = str(radial_unit)

            if self.outputs.azimuth_array is not None:
                if "chi" not in nexus_data_grp:
                    chi_dset = nexus_data_grp.create_dataset(
                        name="chi",
                        data=self.outputs.azimuth_array,
                    )
                    chi_dset.attrs["axis"] = "2"
                    chi_dset.attrs["interpretation"] = "scalar"
                    chi_dset.attrs["unit"] = "deg"

                    nexus_data_grp.attrs["axes"] = [".", "chi", radial_unit.short_name]

    def _get_last_worker(self):
        return _get_persistent_pyfai_worker(**self.get_processing_inputs())
