"""Event-gated alert policy with refractory + Bayesian belief update.

Turns a continuous score time-series into a sparse stream of *event*
timestamps. Three properties distinguish it from a state-based policy:

  1. Refractory period — at most one alert event per `refractory_sec`.
  2. Score-reset requirement — score must dip below `tau_silent_floor`
     before another alert is allowed (transition, not state).
  3. Bayesian belief update — sustained-high score WITHOUT a confirmed
     real event decays `prior_alert_real`; the *effective* threshold for
     the next alert rises accordingly.

Designed as a pure post-processor on already-saved score series; no
PyTorch dependency.

Self-test: ``python -m training.Policy.event_gated_policy``
"""
from __future__ import annotations

from collections import deque
from dataclasses import asdict, dataclass, field
from typing import Deque, Dict, List, Optional, Sequence, Tuple


@dataclass
class EventGatedConfig:
    # Threshold gating
    tau_alert:        float = 0.70   # base alert threshold (score in [0,1])
    tau_silent_floor: float = 0.30   # score must drop below this to re-alert
    delta_min:        float = 0.20   # minimum upward Δ over baseline

    # Refractory
    refractory_sec:   float = 3.0    # absolute lockout between alerts

    # Bayesian belief update
    belief_decay:     float = 0.15   # /sec prior decay if sustained-high w/o event
    belief_min:       float = 0.30   # never decay below this floor
    belief_restore:   float = 1.0    # value on confirmed real event

    # Baseline window
    baseline_window:  float = 5.0    # seconds of recent low-score history
                                      # used to estimate baseline


@dataclass
class AlertEvent:
    t:              float          # timestamp (seconds from video start)
    score:          float          # raw score at fire moment
    prior_at_fire:  float          # belief prior at fire (1.0=fresh, low=demoted)
    last_baseline:  float          # baseline against which Δ was measured
    last_alert_dt:  Optional[float]  # seconds since previous alert; None if first

    def to_dict(self) -> Dict:
        return asdict(self)


class EventGatedPolicy:
    """Stateful sequential decision module.

    Usage::

        policy = EventGatedPolicy()
        policy.reset()
        for t, score in zip(times, scores):
            ev = policy.step(t, score)
            if ev is not None:
                events.append(ev)
    """

    def __init__(self, cfg: Optional[EventGatedConfig] = None) -> None:
        self.cfg = cfg or EventGatedConfig()
        self.reset()

    # ── lifecycle ────────────────────────────────────────────────────────
    def reset(self) -> None:
        cfg = self.cfg
        self._history: Deque[Tuple[float, float]] = deque()
        self._last_alert_t: float = -float("inf")
        self._last_step_t: Optional[float] = None
        self._prev_alert_t: Optional[float] = None
        self._prior: float = 1.0
        # require a reset BEFORE the very first alert too — start as already-reset
        self._seen_reset: bool = True

    # ── core step ───────────────────────────────────────────────────────
    def step(self, t: float, score: float,
             event_observed: bool = False) -> Optional[AlertEvent]:
        """Process one tick. Returns AlertEvent on transition, else None."""
        cfg = self.cfg
        score = float(score)
        t = float(t)

        # confirmed real event → restore prior
        if event_observed:
            self._prior = cfg.belief_restore

        dt = (t - self._last_step_t) if self._last_step_t is not None else 0.0
        self._last_step_t = t

        # maintain rolling history (used for baseline estimation)
        self._history.append((t, score))
        cutoff = t - cfg.baseline_window
        while self._history and self._history[0][0] < cutoff:
            self._history.popleft()

        # 1. refractory lockout
        elapsed_since_alert = t - self._last_alert_t
        if elapsed_since_alert < cfg.refractory_sec:
            # observe-only: decay belief if sustained-high without confirmed event
            if score > cfg.tau_alert and not event_observed and dt > 0:
                self._prior = max(cfg.belief_min,
                                   self._prior - cfg.belief_decay * dt)
            return None

        # 2. score-reset gate (transition, not state)
        if not self._seen_reset:
            if score < cfg.tau_silent_floor:
                self._seen_reset = True
            return None

        # 3. transition gate — Δ over recent low-score baseline
        baseline = self._baseline()
        if (score - baseline) < cfg.delta_min:
            return None

        # 4. belief-modulated effective threshold
        effective_tau = cfg.tau_alert / max(self._prior, cfg.belief_min)
        if score < effective_tau:
            return None

        # 5. fire
        last_dt = (t - self._prev_alert_t) if self._prev_alert_t is not None \
                    else None
        ev = AlertEvent(
            t=t, score=score,
            prior_at_fire=float(self._prior),
            last_baseline=float(baseline),
            last_alert_dt=(float(last_dt) if last_dt is not None else None),
        )
        self._prev_alert_t = self._last_alert_t \
            if self._last_alert_t != -float("inf") else None
        self._last_alert_t = t
        self._seen_reset = False
        return ev

    # ── helpers ──────────────────────────────────────────────────────────
    def _baseline(self) -> float:
        """Mean of recent low-score history; falls back to tau_silent_floor."""
        cfg = self.cfg
        lows = [s for (_, s) in self._history if s < cfg.tau_silent_floor]
        if not lows:
            return cfg.tau_silent_floor
        return sum(lows) / len(lows)

    @property
    def state(self) -> Dict:
        return {
            "last_alert_t":        (None if self._last_alert_t == -float("inf")
                                     else self._last_alert_t),
            "prev_alert_t":        self._prev_alert_t,
            "prior":               self._prior,
            "seen_reset":          self._seen_reset,
            "history_len":         len(self._history),
            "current_baseline":    self._baseline(),
        }


# ─── convenience: apply to a full series in one call ─────────────────────

def apply_policy_to_series(scores: Sequence[float],
                            times: Optional[Sequence[float]] = None,
                            dt: Optional[float] = None,
                            cfg: Optional[EventGatedConfig] = None,
                            event_observed_at: Optional[Sequence[bool]] = None
                            ) -> Tuple[List[AlertEvent], List[Dict]]:
    """Run the policy over a precomputed (t, score) series.

    Either `times` (per-tick timestamps) or `dt` (uniform tick spacing) must
    be supplied. Returns (events, traces) where traces[i] is the policy
    .state snapshot AFTER step i (used for visualization / belief plots).
    """
    n = len(scores)
    if times is None:
        if dt is None:
            raise ValueError("either times or dt must be supplied")
        times = [i * dt for i in range(n)]
    else:
        if len(times) != n:
            raise ValueError(f"times/scores length mismatch: {len(times)} vs {n}")
    if event_observed_at is None:
        event_observed_at = [False] * n
    elif len(event_observed_at) != n:
        raise ValueError("event_observed_at length must match scores length")

    policy = EventGatedPolicy(cfg=cfg)
    policy.reset()
    events: List[AlertEvent] = []
    traces: List[Dict] = []
    for i in range(n):
        ev = policy.step(times[i], scores[i],
                          event_observed=bool(event_observed_at[i]))
        if ev is not None:
            events.append(ev)
        snap = dict(policy.state)
        snap["t"] = float(times[i])
        snap["score"] = float(scores[i])
        snap["fired"] = ev is not None
        traces.append(snap)
    return events, traces


# ─── self-test ───────────────────────────────────────────────────────────

def _self_test() -> int:
    """Synthetic series: brief spike, sustained high, dip, second spike."""
    cfg = EventGatedConfig()
    dt = 0.5  # 2 Hz
    series: List[Tuple[float, float, str]] = []
    # phase A: silent baseline (10 s)
    for i in range(20):
        series.append((i * dt, 0.10, "silent"))
    # phase B: SPIKE at t=10s — should fire
    for i in range(20, 24):
        series.append((i * dt, 0.85, "spike1"))
    # phase C: sustained high (8 s) — should NOT fire (refractory + reset gate +
    # belief decay)
    for i in range(24, 40):
        series.append((i * dt, 0.80, "sustained"))
    # phase D: dip below silent floor (4 s)
    for i in range(40, 48):
        series.append((i * dt, 0.10, "dip"))
    # phase E: SPIKE again — should fire (belief partially recovered? actually
    # belief-decay only happens during refractory, and we're well past that;
    # but belief-restore only on event_observed=True — so prior may be low.
    # we use a strong spike to exceed effective threshold even at min prior)
    for i in range(48, 56):
        series.append((i * dt, 0.95, "spike2"))

    times = [s[0] for s in series]
    scores = [s[1] for s in series]
    events, traces = apply_policy_to_series(scores, times=times, cfg=cfg)

    print(f"[self-test] n_events = {len(events)}")
    for ev in events:
        print(f"   fire at t={ev.t:.2f}  score={ev.score:.2f}  "
              f"prior={ev.prior_at_fire:.2f}  baseline={ev.last_baseline:.2f}  "
              f"last_dt={ev.last_alert_dt}")
    fail = 0
    if len(events) < 1:
        print("FAIL: expected at least 1 event from initial spike")
        fail += 1
    if len(events) > 3:
        print(f"FAIL: too many events ({len(events)}), refractory broken")
        fail += 1
    if events:
        first = events[0]
        if not (9.5 < first.t < 11.0):
            print(f"FAIL: first event at unexpected t={first.t:.2f}")
            fail += 1
    # second spike should fire (after dip)
    if len(events) >= 2:
        second = events[1]
        if not (24.0 <= second.t < 28.0):
            print(f"FAIL: second event at unexpected t={second.t:.2f}")
            fail += 1
    print(f"[self-test] {'PASS' if fail == 0 else 'FAIL'} ({fail} fails)")
    return fail


if __name__ == "__main__":
    import sys
    sys.exit(_self_test())