from pyFAI.units import to_unit
from ewoksid02.id02_format import Path
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,
)
from ewoksid02.utils.mathutils import (
    is_dataset,
    get_variance_from_sigma,
)
from ewoksid02.headers import HS32

DEFAULT_NPT_RAD = 1600
DEFAULT_NPT_AZIM = 360


class AzimuthalTask(
    ID02ProcessingTask,
    optional_input_names=[
        "filename_mask_azimuthal",
        "filename_dark_azimuthal",
        "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_azimuthal (str): Path to the mask file for masking invalid pixels.
        - filename_dark_azimuthal (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 get_processing_parameters(
        self,
    ) -> dict:
        if not self.processing_params:
            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_azimuthal": self.get_input_value(
                    "filename_mask_azimuthal",
                    HS32.get_mask_beamstop_filename(headers=self.headers),
                ),
                "filename_dark_azimuthal": self.get_input_value(
                    "filename_dark_azimuthal",
                    HS32.get_dark_filename(headers=self.headers, missing_ok=True),
                ),
                "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", default="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()
            self.processing_params = processing_params
        return self.processing_params

    def process(self) -> None:
        processing_params = self.get_processing_parameters()
        dataset_variance = None
        if is_dataset(self.dataset_variance):
            dataset_variance = self.dataset_variance
        elif is_dataset(self.dataset_sigma):
            dataset_variance = get_variance_from_sigma(
                dataset_sigma=self.dataset_sigma,
                Dummy=processing_params.get("Dummy", 0.0),
            )

        (
            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

    def save(self) -> None:
        super().save()
        with self.open_id02_file(
            filename=self.processing_filename,
            mode="a",
        ) as processed_file:
            processed_file.update_dataset(
                added_dataset=self.outputs.dataset_sum_signal,
                h5_dataset_path=str(Path(processed_file.nxdata_path) / "sum_signal"),
                index_read=self.index_range_last,
                datatype=self.get_input_value("datatype", "float32"),
            )
            processed_file.update_dataset(
                added_dataset=self.outputs.dataset_sum_normalization,
                h5_dataset_path=str(
                    Path(processed_file.nxdata_path) / "sum_normalization"
                ),
                index_read=self.index_range_last,
                datatype=self.get_input_value("datatype", "float32"),
            )
            processed_file.update_dataset(
                added_dataset=self.outputs.dataset_sum_variance,
                h5_dataset_path=str(Path(processed_file.nxdata_path) / "sum_variance"),
                index_read=self.index_range_last,
                datatype=self.get_input_value("datatype", "float32"),
            )

            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:
                nexus_data_grp = processed_file[processed_file.nxdata_path]
                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.update(
                        {
                            "axis": "3",
                            "interpretation": "scalar",
                            "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.update(
                        {
                            "axis": "2",
                            "interpretation": "scalar",
                            "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_parameters())
