from contextlib import ExitStack

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.average import calculate_average, get_array_limit


class AverageTask(
    ID02ProcessingTask,
    optional_input_names=[
        "dataset_sum_signal",
        "dataset_sum_normalization",
        "dataset_sum_variance",
        "radial_array",
        "azimuth_array",
        "Dummy",
        "unit",
        "azimuth_range",
        "pca_parameters",
    ],
    output_names=[
        "dataset_average_signal_norm",
        "radial_array",
    ],
):
    """The `AverageTask` class is responsible for calculating the average of datasets in the ID02 SAXS pipeline.
    It extends the `ID02ProcessingTask` class and provides additional functionality for handling averaging-specific
    inputs and processing logic. If azimuth_range is not provided, a full average will be performed.

    Optional Inputs:
        - dataset_sum_signal (numpy.ndarray): Sum of signal, non-normalized from an ai.integrate2d result.
        - dataset_sum_normalization (numpy.ndarray): Sum of normalized pixels, from an ai.integrate2d result.
        - dataset_sum_variance (numpy.ndarray): Sum of variance, from an ai.integrate2d result.
        - radial_array (numpy.ndarray): Radial axis array for the dataset.
        - azimuth_array (numpy.ndarray): Azimuthal axis array for the dataset.
        - Dummy (float): Value to replace invalid pixels in the dataset.
        - unit (str): Unit for the radial axis (e.g., "q_nm^-1").
        - azimuth_range (list of tuples): Azimuthal ranges for averaging (e.g., `[(0, 90), (270, 360)]`).
        - pca_parameters (dict): Parameters for Principal Component Analysis (PCA) if applicable.
    Outputs Parameters:
        - dataset_average_signal_norm (numpy.ndarray): Normalized average signal dataset.
        - radial_array (numpy.ndarray): Radial axis array for the averaged data.
    """

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

    def get_processing_inputs(self) -> dict:
        azimuth_range = self.get_input_value("azimuth_range", None)
        array_ranges = None

        if azimuth_range is None:
            self.log_info("There is no azimuth_range. Full average will be done.")
            array_ranges = [(0, -1)]
        else:
            azimuth_array = self.get_input_value("azimuth_array", None)
            if azimuth_array is None:
                self.log_warning(
                    "There is no azimuth_array , azimuth ranges cannot be transformed into array limits. Full average will be done."
                )
                array_ranges = [(0, -1)]
            else:
                array_ranges = [
                    get_array_limit(azimuth_array=azimuth_array, azimuth_range=az_range)
                    for az_range in azimuth_range
                ]

        params_ave = {
            "array_ranges": array_ranges,
            "Dummy": self.get_parameter("Dummy"),
            "azimuth_range": azimuth_range,
        }

        return params_ave

    def process(self):
        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

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

            (
                dataset_average_intensity,
                dataset_average_signal_norm,
                dataset_average_variance,
                dataset_average_sigma,
            ) = calculate_average(
                dataset_intensity=self.dataset_signal,
                dataset_sum_signal=self.get_input_value("dataset_sum_signal", None),
                dataset_sum_norm=self.get_input_value(
                    "dataset_sum_normalization", None
                ),
                dataset_sum_variance=self.get_input_value("dataset_sum_variance", None),
                calculate_variance=self.get_input_value("save_variance", False),
                **processing_params,
            )

            self.outputs.dataset_signal = dataset_average_intensity
            self.outputs.dataset_average_signal_norm = dataset_average_signal_norm
            self.outputs.dataset_variance = dataset_average_variance
            self.outputs.dataset_sigma = dataset_average_sigma
            self.outputs.radial_array = self.get_input_value("radial_array", None)

        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_average_signal_norm,
                h5_group=nexus_data_grp,
                h5_dataset_name="data_signal_norm",
                index_read=self.index_range_last,
            )

            # Add nexus data information (only once)
            unit = self.get_parameter("unit")
            radial_unit = to_unit(unit)

            if radial_unit is not None and radial_unit.short_name not in nexus_data_grp:
                if self.get_input_value("radial_array", None) is not None:
                    radial_dset = nexus_data_grp.create_dataset(
                        name=radial_unit.short_name,
                        data=self.get_input_value("radial_array", None),
                    )
                    radial_dset.attrs["unit"] = str(radial_unit)
                    radial_dset.attrs["interpretation"] = "scalar"
                    radial_dset.attrs["axis"] = "2"
                    nexus_data_grp.attrs["axes"] = [".", radial_unit.short_name]
                    nexus_data_grp.attrs["signal"] = "data"

    def processing_info(self) -> list:
        azimuth_range = self.get_input_value("azimuth_range", [0, 360])
        return [{"h5path": "entry_0000", "name": "ave_limits", "value": azimuth_range}]

    def _save_in_gallery(self) -> str:
        if self.loop_nb != 1:
            return

        dataset_signal = self.outputs.dataset_signal
        if dataset_signal is None or dataset_signal.size == 0:
            return

        import matplotlib.pyplot as plt

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

        signal = dataset_signal.mean(axis=0)
        q = self.outputs.radial_array
        fig, ax = plt.subplots(figsize=(15, 10))
        ax.plot(q, signal)
        ax.set_xlabel(f"{radial_unit.short_name} ({radial_unit.unit_symbol})")
        ax.set_ylabel("Intensity (arb. units)")
        ax.set_title(f"scan {self.scan_nb} - average")
        filename_png = self._get_filename_gallery()
        fig.savefig(filename_png)
