"""Hysteretic OBSERVE policy contract for LKAlert-BD.

Implements the policy described in the plan:

  H_0 = 0
  H_t = max(decay·H_{t-1} + risk_t − β·clear_t, risk_t)
  state = SILENT   if H_t < τ_observe AND clear-streak ≥ K
           OBSERVE  if τ_observe ≤ H_t < τ_alert(TTA, U)
           ALERT   if H_t ≥ τ_alert(TTA, U)
  τ_alert(TTA, U) = τ_alert_base − γ·max(0, 1.5 − TTA_seconds) + δ·U

OBSERVE is intentionally absorbing:
  * Entry: low threshold τ_observe.
  * Release to SILENT requires K consecutive frames with `clear_t` high.
  * Promotion to ALERT happens whenever H_t crosses τ_alert(TTA, U).

This module is pure NumPy — no PyTorch dependency — so it can be applied
to per-step series.json files produced by `tools/rolling_inference_*.py`
without loading any model.
"""
from __future__ import annotations

from dataclasses import dataclass, field
from enum import IntEnum
from typing import Dict, List, Optional, Sequence

import numpy as np


class State(IntEnum):
    SILENT  = 0
    OBSERVE = 1
    ALERT   = 2


@dataclass
class HysteresisConfig:
    decay:           float = 0.85       # H_t memory decay per step
    beta:            float = 0.50       # weight on clearance evidence
    tau_observe:     float = 0.40       # entry to OBSERVE
    tau_alert_base:  float = 0.65       # base ALERT threshold
    gamma:           float = 0.10       # ALERT threshold drop per missing TTA-second
    delta:           float = 0.05       # ALERT threshold rise per uncertainty unit
    clear_threshold: float = 0.60       # what counts as "clear evidence" per step
    clear_streak_K:  int   = 3          # consecutive clear frames to release SILENT
    tta_default_s:   float = 1.5        # if no TTA available, treat as 1.5 s
    # ── Risk-C logit-additive form (default; robust under near-zero factors)
    # observe_score_form ∈ {"logit_additive", "multiplicative"}
    observe_score_form: str = "logit_additive"
    # logit-additive coefficients (a, b, c, d) for
    # observe_logit = a·logit(risk) + b·logit(ego) − c·logit(imm) − d·logit(clear)
    obs_a: float = 1.0   # risk_exists
    obs_b: float = 1.0   # ego_risk
    obs_c: float = 1.0   # 1 - risk_imminent
    obs_d: float = 0.5   # 1 - clear


def _logit(p: float, eps: float = 1e-4) -> float:
    p = max(eps, min(1.0 - eps, float(p)))
    import math
    return math.log(p / (1.0 - p))


def _sigmoid(x: float) -> float:
    import math
    return 1.0 / (1.0 + math.exp(-x))


def observe_score(cfg: HysteresisConfig,
                   risk_exists: float, ego_risk: float,
                   risk_imminent: float, clear: float) -> float:
    """Compute OBSERVE entry score in either form (Risk-C dual)."""
    if cfg.observe_score_form == "logit_additive":
        z = (cfg.obs_a * _logit(risk_exists)
             + cfg.obs_b * _logit(ego_risk)
             - cfg.obs_c * _logit(risk_imminent)
             - cfg.obs_d * _logit(clear))
        return _sigmoid(z)
    # else multiplicative (intuitive form)
    return float(risk_exists) * float(ego_risk) \
            * (1.0 - float(risk_imminent)) * (1.0 - float(clear))


@dataclass
class TraceStep:
    t_seconds:   float
    prob:        float
    H:           float
    state:       int
    tau_alert:   float
    clear_streak: int


def step_threshold(cfg: HysteresisConfig, tta_seconds: float,
                    uncertainty: float) -> float:
    """ALERT threshold drops as TTA shrinks; rises with uncertainty."""
    return (cfg.tau_alert_base
            - cfg.gamma * max(0.0, 1.5 - tta_seconds)
            + cfg.delta * uncertainty)


def simulate_clip(probs: Sequence[float],
                   t_seconds: Sequence[float],
                   tta_seq: Optional[Sequence[float]] = None,
                   uncertainty_seq: Optional[Sequence[float]] = None,
                   p_ego: Optional[Sequence[float]] = None,
                   p_resolution: Optional[Sequence[float]] = None,
                   cfg: HysteresisConfig = HysteresisConfig()
                   ) -> List[TraceStep]:
    """Run the hysteresis simulator over one clip.

    Inputs are per-step sequences; only `probs` and `t_seconds` are required.
    `p_ego`, `p_resolution` default to {prob, 1-prob} when not supplied.
    `tta_seq` defaults to `cfg.tta_default_s` everywhere.
    """
    n = len(probs)
    assert len(t_seconds) == n, (n, len(t_seconds))
    p_ego = list(p_ego) if p_ego is not None else list(probs)
    p_clear = list(p_resolution) if p_resolution is not None \
                else [1.0 - p for p in probs]
    tta = list(tta_seq) if tta_seq is not None \
            else [cfg.tta_default_s] * n
    U = list(uncertainty_seq) if uncertainty_seq is not None else [0.0] * n

    H = 0.0
    state = State.SILENT
    streak = 0
    trace: List[TraceStep] = []
    for i in range(n):
        risk = float(probs[i]) * float(p_ego[i])
        clear = float(p_clear[i])
        H = max(cfg.decay * H + risk - cfg.beta * clear, risk)
        H = float(max(0.0, min(1.0, H)))

        if clear >= cfg.clear_threshold:
            streak += 1
        else:
            streak = 0

        tau_alert = step_threshold(cfg, float(tta[i]), float(U[i]))

        # state transitions (asymmetric release)
        if state == State.ALERT:
            if H < cfg.tau_observe and streak >= cfg.clear_streak_K:
                state = State.SILENT
            elif H < tau_alert:
                state = State.OBSERVE
            else:
                state = State.ALERT
        elif state == State.OBSERVE:
            if H >= tau_alert:
                state = State.ALERT
            elif H < cfg.tau_observe and streak >= cfg.clear_streak_K:
                state = State.SILENT
            else:
                state = State.OBSERVE
        else:  # SILENT
            if H >= tau_alert:
                state = State.ALERT
            elif H >= cfg.tau_observe:
                state = State.OBSERVE
            else:
                state = State.SILENT

        trace.append(TraceStep(
            t_seconds=float(t_seconds[i]),
            prob=float(probs[i]),
            H=float(H),
            state=int(state),
            tau_alert=float(tau_alert),
            clear_streak=int(streak),
        ))
    return trace


# ─── single-clip summary helpers ──────────────────────────────────────────────

def first_alert_lead(trace: List[TraceStep]) -> Optional[float]:
    """Time-before-collision (positive seconds) of first ALERT, or None."""
    for s in trace:
        if s.state == State.ALERT:
            return -s.t_seconds if s.t_seconds < 0 else 0.0
    return None


def observe_duration_seconds(trace: List[TraceStep]) -> float:
    """Total OBSERVE-state duration assuming uniform step spacing."""
    if len(trace) < 2:
        return 0.0
    dt = trace[1].t_seconds - trace[0].t_seconds
    return float(sum(1 for s in trace if s.state == State.OBSERVE) * dt)


def n_release_to_silent(trace: List[TraceStep]) -> int:
    """Count of SILENT releases (transitions to SILENT after non-SILENT)."""
    n = 0
    for prev, cur in zip(trace[:-1], trace[1:]):
        if prev.state != State.SILENT and cur.state == State.SILENT:
            n += 1
    return n


def summarize(trace: List[TraceStep]) -> Dict:
    return {
        "n_steps":       len(trace),
        "first_alert_lead_s":  first_alert_lead(trace),
        "observe_duration_s":   observe_duration_seconds(trace),
        "n_silent_releases":    n_release_to_silent(trace),
        "max_H":         float(max(s.H for s in trace)) if trace else 0.0,
        "mean_H":        float(np.mean([s.H for s in trace])) if trace else 0.0,
        "any_alert":     int(any(s.state == State.ALERT for s in trace)),
        "any_observe":   int(any(s.state == State.OBSERVE for s in trace)),
    }