"""
Pipeline processing for SATOP31-ATT01 arrival time tool.

Produces per-shot AND N-shot averaged outputs:
  - Single-shot: {device}:signal, edge_pos, xcorr, xcorr_ampl, arrival_time, ...
  - Averaged:    {device}:avg_signal, avg_edge_pos, avg_xcorr, avg_xcorr_ampl, avg_arrival_time

The averaged channels are always present once the first valid FEL-on shot
arrives (average of 1 = the shot itself), so output types and sizes never
change mid-stream.  The averaging depth grows until it reaches
``params["avg_length"]`` (default 50), after which it is a sliding window.
"""

from collections import deque
from logging import getLogger
import numpy as np
from scipy import signal

_logger = getLogger(__name__)

# ---------------------------------------------------------------------------
# Module-level state (persists across calls within one pipeline instance)
# ---------------------------------------------------------------------------
buffer_dark    = deque()
buffer_savgol  = deque()
buffer_signal  = None          # deque for N-shot averaged signal
initialized    = False


def initialize(params):
    global buffer_dark, buffer_savgol, buffer_signal, initialized

    buf_len = params["buffer_length"]
    avg_len = params.get("avg_length", 50)

    buffer_dark   = deque(maxlen=buf_len)
    buffer_savgol = deque(maxlen=buf_len)
    buffer_signal = deque(maxlen=avg_len)

    initialized = True
    _logger.info(
        f"ATT01 pipeline initialised: buffer_length={buf_len}, avg_length={avg_len}"
    )


# ---------------------------------------------------------------------------
# Edge-finding (unchanged logic, extracted for reuse)
# ---------------------------------------------------------------------------
def find_edge(data, step_length=50, edge_type="falling", roi=None):
    """
    Find an edge via cross-correlation with a step waveform.

    If *roi* is ``[start, end]`` the search is restricted to that slice
    and the returned edge position is offset back to full-frame coordinates.
    """
    if roi is not None:
        data_roi = data[roi[0]:roi[1]]
    else:
        data_roi = data

    step_waveform = np.ones(step_length)
    if edge_type == "rising":
        step_waveform[: step_length // 2] = -1
    elif edge_type == "falling":
        step_waveform[step_length // 2 :] = -1

    xcorr = signal.correlate(data_roi, step_waveform, mode="valid")
    edge_position = np.argmax(xcorr) + np.floor(step_length / 2)

    if roi is not None:
        edge_position += roi[0]

    return {
        "edge_pos":    edge_position,
        "xcorr":       xcorr,
        "xcorr_ampl":  float(np.amax(xcorr)),
        "signal":      data,
    }


# ---------------------------------------------------------------------------
# Main entry point
# ---------------------------------------------------------------------------
def process(data, pulse_id, timestamp, params):
    global initialized

    device       = params["device"]
    step_length  = params["step_length"]
    edge_type    = params["edge_type"]
    dark_event   = params["dark_event"]
    fel_on_event = params["fel_on_event"]
    calib        = params["calib"]
    filter_window = params["filter_window"]
    prof_sig     = data[params["prof_sig"]]
    events       = data[params["events"]]

    if not initialized:
        initialize(params)

    # ------------------------------------------------------------------
    # Pre-processing (unchanged)
    # ------------------------------------------------------------------
    prof_sig_savgol = signal.savgol_filter(prof_sig, filter_window, 3)

    if events[dark_event]:
        buffer_dark.append(prof_sig)
        buffer_savgol.append(prof_sig_savgol)

    if buffer_savgol:
        prof_sig_norm = prof_sig_savgol / np.mean(np.array(buffer_savgol), axis=0)
    else:
        prof_sig_norm = prof_sig_savgol

    # ------------------------------------------------------------------
    # Single-shot edge finding
    # ------------------------------------------------------------------
    if events[fel_on_event] and not events[dark_event]:
        edge_results = find_edge(
            prof_sig_norm, step_length, edge_type, roi=params.get("roi")
        )
        edge_results["arrival_time"] = np.polyval(calib, edge_results["edge_pos"])
    else:
        edge_results = {
            "edge_pos":      None,
            "xcorr":         None,
            "xcorr_ampl":    None,
            "signal":        prof_sig_norm,
            "arrival_time":  None,
        }

    # ------------------------------------------------------------------
    # Build single-shot output (unchanged keys)
    # ------------------------------------------------------------------
    output = {}
    for key, value in edge_results.items():
        output[f"{device}:{key}"] = value

    output[f"{device}:raw_wf"]        = prof_sig
    output[f"{device}:raw_wf_savgol"] = prof_sig_savgol

    if buffer_dark:
        output[f"{device}:avg_dark_wf"] = np.mean(buffer_dark, axis=0)
    else:
        output[f"{device}:avg_dark_wf"] = None

    # ------------------------------------------------------------------
    # N-shot signal averaging + averaged edge finding
    # ------------------------------------------------------------------
    # Only accumulate valid FEL-on normalised signals
    if events[fel_on_event] and not events[dark_event]:
        buffer_signal.append(prof_sig_norm.copy())

    if len(buffer_signal) > 0:
        # Compute average over available shots (1..avg_length)
        avg_signal = np.mean(np.array(buffer_signal), axis=0)

        # Run edge finding on the averaged signal
        avg_edge_results = find_edge(
            avg_signal, step_length, edge_type, roi=params.get("roi")
        )
        avg_edge_results["arrival_time"] = np.polyval(
            calib, avg_edge_results["edge_pos"]
        )

        output[f"{device}:avg_signal"]       = avg_signal
        output[f"{device}:avg_edge_pos"]     = avg_edge_results["edge_pos"]
        output[f"{device}:avg_xcorr"]        = avg_edge_results["xcorr"]
        output[f"{device}:avg_xcorr_ampl"]   = avg_edge_results["xcorr_ampl"]
        output[f"{device}:avg_arrival_time"] = avg_edge_results["arrival_time"]
        output[f"{device}:avg_nshots"]       = len(buffer_signal)
    else:
        # No valid shots yet — don't include avg keys at all.
        # (Once the first FEL-on shot arrives, these keys appear and
        #  stay with consistent types/sizes for the rest of the run.)
        pass

    return output