"""Test ewoks task for the demo stop_scan_xrpd_id31 workflow.

Integrates detector frames via PyFAI, accumulates the 1D profile and sends
a ``{'type': 'stop'}`` event on a Redis stream once the accumulated mean relative error
drops below the input ``acc_profile_threshold``.
"""

from __future__ import annotations

import logging
from contextlib import ExitStack
from pathlib import Path

import h5py
from blissdata.h5api.dynamic_hdf5 import File
from ewokscore import Task
from ewoksxrpd.tasks.integrate import _BaseIntegrateSingleGeometry

from .stop_scan import _EventStream

logger = logging.getLogger(__name__)


class IntegrateAndSendMsg(
    _BaseIntegrateSingleGeometry,
    input_names=[
        "filename",
        "scan_number",
        "detector_name",
        "acc_profile_threshold",
        "redis_url",
        "response_stream_name",
    ],
    optional_input_names=[
        "scan_key",
        "retry_timeout",
        "retry_period",
    ],
):
    def run(self):
        with ExitStack() as stack:
            self._run_in_stack(stack)

    def _run_in_stack(self, stack: ExitStack):
        import numpy as np
        from pyFAI.containers import Integrate1dResult

        with self._worker() as (worker, _):
            event_stream = _EventStream(
                self.inputs.redis_url, self.inputs.response_stream_name
            )

            detector_name = str(self.inputs.detector_name)

            # --- data iterator ------------------------------------------------
            if self.inputs.scan_key:
                logger.info("PyFAI integrate data from %r", self.inputs.scan_key)
                data_iterator = self.iter_bliss_data_from_memory(
                    self.inputs.scan_key,
                    lima_names=[detector_name],
                    retry_timeout=self.get_input_value("retry_timeout", None),
                    retry_period=self.get_input_value("retry_period", None),
                )
            else:
                logger.info(
                    "PyFAI integrate data from '%s::%d.%s'",
                    self.inputs.filename,
                    self.inputs.scan_number,
                )
                data_iterator_context = self.iter_bliss_data(
                    self.inputs.filename,
                    self.inputs.scan_number,
                    lima_names=[detector_name],
                    retry_timeout=self.get_input_value("retry_timeout", None),
                    retry_period=self.get_input_value("retry_period", None),
                )
                data_iterator = stack.enter_context(data_iterator_context)

            # --- integrate + accumulate ---------------------------------------
            acc_intensity = None
            acc_variance = None  # sum of sigma**2 across frames
            nframes = 0
            for i, scan_point_data_dict in enumerate(data_iterator):
                image = scan_point_data_dict[detector_name]
                result: Integrate1dResult = worker.process(image)

                frame_intensity = np.array(result.intensity)
                frame_sigma = (
                    np.array(result.sigma) if result.sigma is not None else None
                )

                if acc_intensity is None:
                    acc_intensity = frame_intensity.copy()
                    acc_variance = frame_sigma**2 if frame_sigma is not None else None
                else:
                    acc_intensity += frame_intensity
                    if frame_sigma is not None and acc_variance is not None:
                        acc_variance += frame_sigma**2
                nframes += 1

                if acc_variance is None:
                    logger.info(
                        "frames: %d - no sigma from worker, skipping stop criterion",
                        nframes,
                    )
                    continue

                # mean relative error over bins with positive accumulated intensity:
                #   rel_err = mean( sqrt(Σσ²) / ΣI )
                mask = acc_intensity > 0
                if not mask.any():
                    continue
                rel_err = float(
                    np.mean(np.sqrt(acc_variance[mask]) / acc_intensity[mask])
                )
                logger.info("rel_err: %.4f / frames: %d", rel_err, nframes)

                if rel_err < self.inputs.acc_profile_threshold:
                    event_stream.send_event({"type": "stop"})
                    logger.warning(
                        "Stopping scan: rel_err=%.4f < threshold=%.4f (frames: %d)",
                        rel_err,
                        self.inputs.acc_profile_threshold,
                        nframes,
                    )
                    break


class WriteSumFrame(
    Task,
    input_names=[
        "filename",
        "scan_number",
        "detector_name",
    ],
):
    """Sum all raw detector frames of a completed loopscan and write ``*_sum.h5``.

    Reads raw 2D frames from the scan HDF5 file, sums them across the frame
    axis, and writes the sum alongside the frame count to a ``*_sum.h5`` file
    next to the original scan file.
    """

    def run(self):

        filename = self.inputs.filename
        scan_nb = self.inputs.scan_number
        detector_name = self.inputs.detector_name

        ds_path = f"/{scan_nb}.1/measurement/{detector_name}"

        with File(filename, retry_timeout=10, retry_period=0.1) as root:
            _ = root[f"/{scan_nb}.1/end_time"]  # ensure the scan is finished
            frames = root[ds_path][()]  # type: ignore

        nframes = int(frames.shape[0]) if frames.ndim > 2 else 1
        sum_frame = frames.sum(axis=0) if frames.ndim > 2 else frames

        stem = Path(filename).stem
        sum_filename = str(Path(filename).parent / f"{stem}_{scan_nb}_sum.h5")
        with h5py.File(sum_filename, "w") as f:
            f.create_dataset("sum_frame", data=sum_frame)
            f.create_dataset("nframes", data=nframes)
        logger.info("Auto-stop sum written: %s (%d frames)", sum_filename, nframes)
