evaluate_checkpoint_real.py

#!/usr/bin/env python3
"""
evaluate_checkpoint_real.py
===========================
Real-trajectory evaluation for agent decisions vs L1 impact labels.

Closes the gap between:
  - scorer product-vs-L1 metrics (backtest_indonesia.py), and
  - agent eval that previously used only synthetic _zone_event_flags.

Build order (contracts → code):
  entities: EvalDayRecord, EvalResult
  modes: SCORER_ORACLE | CHECKPOINT | ALWAYS | NEVER
  GT: days inside an L1 event span → gt_source="l1"; days outside every
      span for that zone → gt_source="unlabeled" (excluded from P/R/F1).
      Sparse catalogs must not manufacture TN/FP from silence.
  alert definition: product gate (WARNING+ OR drought≥0.35 OR flood≥0.25)
  metrics: product-vs-L1 P/R/F1 on L1 event days only; unlabeled_alert_rate

Does NOT require a trained checkpoint for SCORER_ORACLE / ALWAYS / NEVER —
those baselines prove the harness before GPU time is spent.

Usage examples
--------------
  # Scorer oracle on historical cache (no checkpoint)
  python evaluate_checkpoint_real.py \\
    --pkl historical_continuous_indonesia_v1.pkl \\
    --impact-labels impact_labels_java_v1.json \\
    --zones karawang_rice,indramayu_rice \\
    --start 2023-07-01 --end 2023-11-30 \\
    --mode scorer_oracle

  # Trained agent (must match env basin_context dim=8)
  python evaluate_checkpoint_real.py \\
    --pkl historical_continuous_indonesia_v1.pkl \\
    --impact-labels impact_labels_java_v1.json \\
    --zones karawang_rice \\
    --start 2023-07-01 --end 2023-11-30 \\
    --mode checkpoint --checkpoint path/to/final.zip \\
    --n-zones 1 --max-steps 4
"""
from __future__ import annotations

import argparse
import json
import logging
import pickle
import sys
from collections import Counter
from dataclasses import asdict, dataclass, field
from datetime import date, datetime, timezone
from pathlib import Path
from typing import Any, Dict, List, Optional, Sequence, Tuple

import numpy as np

import zone_observation as _zo

assert _zo.SCHEMA_VERSION == 3, (
    f"evaluate_checkpoint_real: zone_observation schema mismatch "
    f"(expected 3, got {_zo.SCHEMA_VERSION})"
)

from zone_observation import (
    BasinContext,
    EpisodeContext,
    ForecastConfig,
    ForecastResult,
    RiskScore,
    ZoneObs,
)
from crop_risk_scorer import compute_risk_score
from product_alert_service import (
    DEFAULT_PRODUCT_GATE,
    ProductGateConfig,
    is_product_actionable,
)
from weather_forecast_env import make_weather_env

logger = logging.getLogger(__name__)

# PATH_CHECK: count loop_product vs fresh-scorer product across all episodes
# (not one-shot — multi-zone real eval can diverge on a subset of days).
_PATH_CHECK_N = 0
_PATH_CHECK_DIVERGE = 0

# ---------------------------------------------------------------------------
# Entities
# ---------------------------------------------------------------------------

# Relative belief movement threshold (policy-sensitive by construction).
# Fixed bars against prior/rational saturate when prior > rational.
BELIEF_RAISE_EPS = 0.02


def _path_check_reset() -> None:
    global _PATH_CHECK_N, _PATH_CHECK_DIVERGE
    _PATH_CHECK_N = 0
    _PATH_CHECK_DIVERGE = 0


def _path_check_report() -> None:
    """Print aggregate PATH_CHECK after an eval pass."""
    if _PATH_CHECK_N <= 0:
        return
    n_ok = _PATH_CHECK_N - _PATH_CHECK_DIVERGE
    print(
        f"PATH_CHECK summary: n={_PATH_CHECK_N}  ok={n_ok}  "
        f"diverge={_PATH_CHECK_DIVERGE}  "
        f"rate={_PATH_CHECK_DIVERGE / _PATH_CHECK_N:.3f}  "
        f"(loop product is authoritative; rs is display-only)",
        flush=True,
    )


@dataclass
class EvalDayRecord:
    """One evaluated zone-day."""

    date: str
    zone_id: str
    mode: str
    product_alert: bool
    elevated: bool
    alert_level: str
    drought_risk: float
    flood_risk: float
    event_drought: bool
    event_flood: bool
    gt_source: str  # "l1" | "unlabeled" | "none"
    product_emit_code: str = "N/A"  # not emitted to bus in this harness
    ep_len: int = 0
    basin_dim: int = 0
    # Policy-sensitive: terminal belief vs episode initial belief.
    # Scorer product_alert is independent of inspections; belief_delta is not.
    believed_p: float = 0.0
    initial_belief: float = 0.0
    belief_delta: float = 0.0
    belief_raised: bool = False  # believed_p > initial_belief + BELIEF_RAISE_EPS


@dataclass
class EvalMetrics:
    n_days: int = 0
    tp: int = 0
    fp: int = 0
    fn: int = 0
    tn: int = 0
    n_l1: int = 0  # days inside an L1 event span (positive labels only in v1)
    n_product: int = 0
    # Days outside every L1 span for the zone — excluded from P/R/F1
    n_unlabeled: int = 0
    n_unlabeled_product: int = 0
    n_unlabeled_belief_raised: int = 0
    # Belief-raised on L1 event days only (positive spans → TP/FN; no FP/TN path)
    belief_tp: int = 0
    belief_fn: int = 0
    # Kept for schema stability; never incremented under positive-only L1 v1
    belief_fp: int = 0
    belief_tn: int = 0
    n_belief_raised: int = 0
    # Raw belief-delta distribution (all days; also split for diagnostics)
    belief_deltas: List[float] = field(default_factory=list)
    belief_deltas_l1: List[float] = field(default_factory=list)
    belief_deltas_unlabeled: List[float] = field(default_factory=list)

    @property
    def precision(self) -> Optional[float]:
        # Positive-only L1 catalog: no confirmed negatives → FP stays 0 by
        # construction, so precision is undefined (not 1.0).
        # Self-healing: once a catalog provides real TN/FP paths, fp/tn leave
        # zero and this guard stops firing. Re-audit if negatives arrive for
        # only some hazards while others stay positive-only (per-Metrics-object
        # guard is not per-hazard).
        if self.fp == 0 and self.tn == 0 and (self.tp + self.fn) > 0:
            return None
        d = self.tp + self.fp
        return self.tp / d if d else None

    @property
    def recall(self) -> Optional[float]:
        d = self.tp + self.fn
        return self.tp / d if d else None

    @property
    def f1(self) -> Optional[float]:
        p, r = self.precision, self.recall
        if p is None or r is None or (p + r) == 0:
            return None
        return 2 * p * r / (p + r)

    @property
    def belief_precision(self) -> Optional[float]:
        # Same structural issue as product precision: no confirmed-negative
        # path under positive-only L1 → never report a fake 1.0.
        return None

    @property
    def belief_recall(self) -> Optional[float]:
        """Fraction of L1 event days where belief_delta > BELIEF_RAISE_EPS."""
        d = self.belief_tp + self.belief_fn
        return self.belief_tp / d if d else None

    @property
    def belief_f1(self) -> Optional[float]:
        # Undefined without belief_precision.
        return None

    def belief_delta_stats(self) -> Dict[str, float]:
        xs = self.belief_deltas
        if not xs:
            return {"n": 0, "mean": 0.0, "std": 0.0, "p10": 0.0, "p50": 0.0, "p90": 0.0,
                    "frac_pos": 0.0, "frac_neg": 0.0, "frac_near0": 0.0}
        arr = sorted(xs)
        n = len(arr)
        mean = sum(arr) / n
        var = sum((x - mean) ** 2 for x in arr) / max(n, 1)
        def pct(p: float) -> float:
            i = min(n - 1, max(0, int(round(p * (n - 1)))))
            return arr[i]
        near = sum(1 for x in arr if abs(x) < BELIEF_RAISE_EPS)
        return {
            "n": n,
            "mean": mean,
            "std": var ** 0.5,
            "p10": pct(0.10),
            "p50": pct(0.50),
            "p90": pct(0.90),
            "frac_pos": sum(1 for x in arr if x > BELIEF_RAISE_EPS) / n,
            "frac_neg": sum(1 for x in arr if x < -BELIEF_RAISE_EPS) / n,
            "frac_near0": near / n,
        }

    @property
    def unlabeled_alert_rate(self) -> Optional[float]:
        if self.n_unlabeled <= 0:
            return None
        return self.n_unlabeled_product / self.n_unlabeled

    @property
    def belief_unlabeled_raise_rate(self) -> Optional[float]:
        """How often belief crosses the raise bar on days the catalog is silent."""
        if self.n_unlabeled <= 0:
            return None
        return self.n_unlabeled_belief_raised / self.n_unlabeled

    def to_dict(self) -> Dict[str, Any]:
        def _f(x: Optional[float]) -> Optional[float]:
            return None if x is None else round(x, 6)

        bd = self.belief_delta_stats()
        # Temporary swap for L1 / unlabeled delta stats without mutating permanently
        def _subset_stats(xs: List[float]) -> Dict[str, Any]:
            saved = self.belief_deltas
            self.belief_deltas = xs
            out = self.belief_delta_stats()
            self.belief_deltas = saved
            return {k: (round(v, 6) if isinstance(v, float) else v) for k, v in out.items()}

        return {
            "n_days": self.n_days,
            "n_l1_event_days": self.n_l1,
            "n_unlabeled": self.n_unlabeled,
            "n_unlabeled_product": self.n_unlabeled_product,
            "n_unlabeled_belief_raised": self.n_unlabeled_belief_raised,
            "unlabeled_alert_rate": _f(self.unlabeled_alert_rate),
            "belief_unlabeled_raise_rate": _f(self.belief_unlabeled_raise_rate),
            "n_product_alerts": self.n_product,
            "n_belief_raised": self.n_belief_raised,
            "tp": self.tp,
            "fp": self.fp,
            "fn": self.fn,
            "tn": self.tn,
            "precision": _f(self.precision),
            "recall": _f(self.recall),
            "f1": _f(self.f1),
            "note_metrics": (
                "Recall uses only gt_source=l1 (days inside an event span). "
                "Unlabeled days are excluded from the confusion matrix and "
                "reported via unlabeled_alert_rate / belief_unlabeled_raise_rate. "
                "Precision and F1 are null under positive-only L1 (no confirmed "
                "negatives → FP/TN stay 0 by construction; a printed 1.0 would "
                "be an artifact)."
            ),
            "belief_tp": self.belief_tp,
            "belief_fn": self.belief_fn,
            "belief_fp": self.belief_fp,
            "belief_tn": self.belief_tn,
            "belief_precision": _f(self.belief_precision),
            "belief_recall": _f(self.belief_recall),
            "belief_f1": _f(self.belief_f1),
            "note_belief_metrics": (
                "belief_recall = belief_tp/(belief_tp+belief_fn) on L1 event "
                "days is the legitimate policy-sensitive number. "
                "belief_precision and belief_f1 are always null (no FP/TN path)."
            ),
            "belief_delta": {k: (round(v, 6) if isinstance(v, float) else v)
                             for k, v in bd.items()},
            "belief_delta_l1": _subset_stats(self.belief_deltas_l1),
            "belief_delta_unlabeled": _subset_stats(self.belief_deltas_unlabeled),
        }


# ---------------------------------------------------------------------------
# Historical cache → EpisodeContext
# ---------------------------------------------------------------------------

def _parse_day(s: str) -> date:
    return date.fromisoformat(s[:10])


def load_historical_points(
    pkl_path: Path,
    zone_ids: Sequence[str],
    start: date,
    end: date,
) -> List[Dict[str, Any]]:
    """
    Flatten trajectory cache into zone-day points inside [start, end].

    Contract:
      purpose: SSOT historical points for real eval
      forbidden: mutating the pkl; inventing missing obs fields
      response: list of point dicts with obs/forecast/basin_context keys
    """
    with open(pkl_path, "rb") as f:
        cache = pickle.load(f)
    trajs = cache.get("trajectories") or []
    zone_set = set(zone_ids)
    out: List[Dict[str, Any]] = []
    for traj in trajs:
        meta = traj.get("meta") or {}
        zid = meta.get("zone_id")
        if zid not in zone_set:
            continue
        for pt in traj.get("trajectory") or []:
            vt = _parse_day(str(pt.get("valid_time", "")))
            if vt < start or vt > end:
                continue
            if pt.get("zone_id") and pt["zone_id"] not in zone_set:
                continue
            out.append(pt)
    out.sort(key=lambda p: (str(p.get("zone_id")), str(p.get("valid_time"))))
    return out


def point_to_episode(
    pt: Dict[str, Any],
    cfg: ForecastConfig,
) -> EpisodeContext:
    """Deserialize one cache point into a typed single-zone EpisodeContext."""
    obs = ZoneObs.from_dict(dict(pt["obs"]))
    fc = ForecastResult.from_dict(dict(pt["forecast"]))
    basin = None
    if pt.get("basin_context"):
        try:
            basin = BasinContext.from_dict(dict(pt["basin_context"]))
        except Exception as e:
            logger.warning("basin_context deserialize failed: %s", e)
            basin = None
    return EpisodeContext(
        obs=obs,
        forecast=fc,
        config=cfg,
        zone_ids=[obs.zone_id],
        basin_context=basin,
    )


def points_to_multi_zone_episode(
    pts: Sequence[Dict[str, Any]],
    cfg: ForecastConfig,
    zone_order: Sequence[str],
) -> EpisodeContext:
    """
    Build a multi-zone EpisodeContext from same-day points (Blocker B).

    Requires one point per zone_id in zone_order. Primary obs/forecast = slot 0.
    """
    by_z = {}
    for pt in pts:
        obs = ZoneObs.from_dict(dict(pt["obs"]))
        by_z[obs.zone_id] = pt
    missing = [z for z in zone_order if z not in by_z]
    if missing:
        raise ValueError(f"points_to_multi_zone_episode missing zones: {missing}")
    zone_obs: List[ZoneObs] = []
    zone_fc: List[ForecastResult] = []
    basin = None
    for zid in zone_order:
        pt = by_z[zid]
        zo = ZoneObs.from_dict(dict(pt["obs"]))
        zf = ForecastResult.from_dict(dict(pt["forecast"]))
        zone_obs.append(zo)
        zone_fc.append(zf)
        if basin is None and pt.get("basin_context"):
            try:
                basin = BasinContext.from_dict(dict(pt["basin_context"]))
            except Exception as e:
                logger.warning("basin_context deserialize failed: %s", e)
    return EpisodeContext(
        obs=zone_obs[0],
        forecast=zone_fc[0],
        config=cfg,
        zone_ids=list(zone_order),
        basin_context=basin,
        zone_obs=zone_obs,
        zone_forecasts=zone_fc,
    )


def group_points_by_date(
    points: Sequence[Dict[str, Any]],
) -> Dict[str, List[Dict[str, Any]]]:
    """Map ISO date string → list of points on that calendar day."""
    out: Dict[str, List[Dict[str, Any]]] = {}
    for pt in points:
        obs = pt.get("obs") or {}
        vt = str(obs.get("valid_time") or pt.get("valid_time") or "")[:10]
        if not vt:
            continue
        out.setdefault(vt, []).append(pt)
    return out


# ---------------------------------------------------------------------------
# Decision policies
# ---------------------------------------------------------------------------

def decide_scorer_oracle(
    obs: ZoneObs,
    fc: ForecastResult,
    cfg: ForecastConfig,
    gate: ProductGateConfig,
) -> Tuple[bool, bool, RiskScore, int, float, float]:
    """Product decision from deterministic scorer only (no agent)."""
    rs = compute_risk_score(obs, fc, cfg)
    product = is_product_actionable(rs, gate)
    elevated = rs.is_elevated()
    # No agent → no belief state; return zeros so delta is 0.
    return product, elevated, rs, 0, 0.0, 0.0


def decide_always() -> Tuple[bool, bool, None, int, float, float]:
    return True, True, None, 0, 1.0, 1.0


def decide_never() -> Tuple[bool, bool, None, int, float, float]:
    return False, False, None, 0, 0.0, 0.0


def _initial_belief_from_info(info: Dict[str, Any], cfg: ForecastConfig) -> float:
    """
    Belief at reset (before any inspect), matching terminate aggregation.

    Terminal believed_p is max(belief_map[:n_active]). Using mean_belief or
    zone_belief[0] here made multi-zone zero-inspect deltas nonzero even with
    no inspections (max of three post-reset blends ≠ mean or zone0). Prefer
    the same max aggregation the env exposes as info['believed_p'].
    """
    if "believed_p" in info and info["believed_p"] is not None:
        return float(info["believed_p"])
    zb = info.get("zone_belief")
    if zb is not None:
        try:
            import numpy as _np
            arr = _np.asarray(zb, dtype=float).ravel()
            n = int(info.get("n_zones") or cfg.n_zones or 0)
            if arr.size:
                if n > 0:
                    return float(_np.max(arr[: min(n, arr.size)]))
                return float(_np.max(arr))
        except Exception:
            pass
    # Last resort only — not symmetric with terminal max.
    if "mean_belief" in info and info["mean_belief"] is not None:
        return float(info["mean_belief"])
    return float(cfg.prior_belief)


def decide_checkpoint(
    model: Any,
    env: Any,
    ctx: EpisodeContext,
    gate: ProductGateConfig,
) -> Tuple[bool, bool, Optional[RiskScore], int, float, float]:
    """
    Roll MaskablePPO until terminate / budget exhaust.

    Returns:
      product, elevated from env info at terminal step (scorer-based flags —
      these do NOT depend on inspections; kept for product-bus parity).
      rs for display magnitudes only — MUST NOT overwrite product/elevated.
      ep_len, believed_p (terminal), initial_belief (at reset).

    Policy-sensitive signal: belief_delta = believed_p - initial_belief.
    """
    global _PATH_CHECK_N, _PATH_CHECK_DIVERGE
    obs, info = env.reset(options={"context": ctx})
    initial_belief = _initial_belief_from_info(info, ctx.config)
    done = False
    ep_len = 0
    product = False
    elevated = False
    believed_p = initial_belief
    while not done:
        masks = env.action_masks() if hasattr(env, "action_masks") else None
        if masks is None and hasattr(env, "env") and hasattr(env.env, "action_masks"):
            masks = env.env.action_masks()
        action, _ = model.predict(obs, action_masks=masks, deterministic=True)
        obs, reward, terminated, truncated, info = env.step(int(action))
        ep_len += 1
        done = bool(terminated or truncated)
        if "product_actionable" in info:
            product = bool(info["product_actionable"])
        if "elevated" in info:
            elevated = bool(info["elevated"])
        if "believed_p" in info:
            believed_p = float(info["believed_p"])
        elif "mean_belief" in info:
            # Prefer max-aggregation; mean is only a fallback if believed_p missing.
            believed_p = float(info["mean_belief"])

    # PATH_CHECK: product from the loop must be the returned values.
    # Trailing compute_risk_score is DISPLAY ONLY — must not overwrite.
    # Multi-zone: primary ctx.obs is only slot 0; rs may diverge from the
    # env's worst-case multi-zone product — count divergences, do not print once.
    loop_product, loop_elevated = product, elevated
    rs = None
    try:
        rs = compute_risk_score(ctx.obs, ctx.forecast, ctx.config)
        rs_product = is_product_actionable(rs, gate)
        _PATH_CHECK_N += 1
        if rs_product != loop_product:
            _PATH_CHECK_DIVERGE += 1
    except Exception:
        pass
    # Explicit: return loop-captured values, never rs-derived product.
    return loop_product, loop_elevated, rs, ep_len, believed_p, initial_belief


def decide_zero_inspect(
    env: Any,
    ctx: EpisodeContext,
    gate: ProductGateConfig,
) -> Tuple[bool, bool, Optional[RiskScore], int, float, float]:
    """Terminate on step 1 with zero inspections (policy-insensitivity control).

    With symmetric max-aggregation on initial and terminal believed_p, belief
    delta must be exactly 0.0 when no inspect updates the map.
    """
    obs, info = env.reset(options={"context": ctx})
    initial_belief = _initial_belief_from_info(info, ctx.config)
    base = env.env if hasattr(env, "env") else env
    terminate_action = int(getattr(base, "terminate_action", ctx.config.n_zones))
    obs, reward, terminated, truncated, info = env.step(terminate_action)
    product = bool(info.get("product_actionable", False))
    elevated = bool(info.get("elevated", False))
    if "believed_p" in info and info["believed_p"] is not None:
        believed_p = float(info["believed_p"])
    else:
        believed_p = _initial_belief_from_info(info, ctx.config)
    rs = None
    try:
        rs = compute_risk_score(ctx.obs, ctx.forecast, ctx.config)
    except Exception:
        pass
    # Zero inspect: belief should equal initial (no inspect → no update).
    return product, elevated, rs, 1, believed_p, initial_belief


# ---------------------------------------------------------------------------
# Core eval loop
# ---------------------------------------------------------------------------

def evaluate_multi_zone_days(
    points: Sequence[Dict[str, Any]],
    *,
    mode: str,
    cfg: ForecastConfig,
    gate: ProductGateConfig,
    zone_order: Sequence[str],
    impact_store: Any = None,
    model: Any = None,
    env: Any = None,
) -> Tuple[List[EvalDayRecord], EvalMetrics]:
    """
    Multi-zone real eval (Blocker B).

    Groups points by calendar day; requires all zone_order zones present that
    day. GT: l1 if ANY packed zone has an L1 event that day; else unlabeled
    when the impact store is loaded. One EvalDayRecord per complete day
    (zone_id is the joined zone list).
    """
    if mode not in ("checkpoint", "zero_inspect"):
        raise ValueError(
            f"evaluate_multi_zone_days only supports checkpoint/zero_inspect "
            f"(got {mode!r})"
        )
    _path_check_reset()
    if env is None:
        raise RuntimeError("evaluate_multi_zone_days requires env")
    if mode == "checkpoint" and model is None:
        raise RuntimeError("checkpoint mode requires model")

    records: List[EvalDayRecord] = []
    m = EvalMetrics()
    by_day = group_points_by_date(points)
    n_skip_incomplete = 0
    zone_set = set(zone_order)

    for day_s in sorted(by_day.keys()):
        day_pts = []
        present = set()
        for pt in by_day[day_s]:
            zo = ZoneObs.from_dict(dict(pt["obs"]))
            if zo.zone_id in zone_set:
                day_pts.append(pt)
                present.add(zo.zone_id)
        if not all(z in present for z in zone_order):
            n_skip_incomplete += 1
            continue
        chosen: List[Dict[str, Any]] = []
        for zid in zone_order:
            for pt in day_pts:
                if ZoneObs.from_dict(dict(pt["obs"])).zone_id == zid:
                    chosen.append(pt)
                    break
        ctx = points_to_multi_zone_episode(chosen, cfg, zone_order)
        day = date.fromisoformat(day_s)

        event_d = event_f = False
        gt_source = "none"
        if impact_store is not None:
            any_event = False
            for zid in zone_order:
                try:
                    ld, lf = impact_store.labels_for_day(zid, day)
                    event_d = event_d or bool(ld)
                    event_f = event_f or bool(lf)
                    if ld or lf:
                        any_event = True
                except Exception as e:
                    logger.warning("L1 query failed %s %s: %s", zid, day, e)
            gt_source = "l1" if any_event else "unlabeled"

        if mode == "checkpoint":
            product, elevated, rs, ep_len, believed_p, initial_belief = (
                decide_checkpoint(model, env, ctx, gate)
            )
        else:
            product, elevated, rs, ep_len, believed_p, initial_belief = (
                decide_zero_inspect(env, ctx, gate)
            )

        alert_level = (
            rs.alert_level.value
            if rs is not None
            else ("warning" if product else "none")
        )
        drought_risk = float(rs.drought_risk) if rs is not None else 0.0
        flood_risk = float(rs.flood_risk) if rs is not None else 0.0
        belief_delta = float(believed_p) - float(initial_belief)
        belief_raised = bool(belief_delta > BELIEF_RAISE_EPS)

        if gt_source == "l1":
            if product:
                m.tp += 1
            else:
                m.fn += 1
            if belief_raised:
                m.belief_tp += 1
            else:
                m.belief_fn += 1
            m.n_l1 += 1
            m.belief_deltas_l1.append(belief_delta)
        elif gt_source == "unlabeled":
            m.n_unlabeled += 1
            if product:
                m.n_unlabeled_product += 1
            if belief_raised:
                m.n_unlabeled_belief_raised += 1
            m.belief_deltas_unlabeled.append(belief_delta)

        m.n_days += 1
        if product:
            m.n_product += 1
        if belief_raised:
            m.n_belief_raised += 1
        m.belief_deltas.append(belief_delta)

        records.append(
            EvalDayRecord(
                date=day.isoformat(),
                zone_id="+".join(zone_order),
                mode=mode,
                product_alert=product,
                elevated=elevated,
                alert_level=alert_level,
                drought_risk=drought_risk,
                flood_risk=flood_risk,
                event_drought=event_d,
                event_flood=event_f,
                gt_source=gt_source,
                ep_len=ep_len,
                basin_dim=8,
                believed_p=float(believed_p),
                initial_belief=float(initial_belief),
                belief_delta=belief_delta,
                belief_raised=belief_raised,
            )
        )

    if n_skip_incomplete:
        print(
            f"multi-zone: skipped {n_skip_incomplete} days missing full "
            f"zone set {list(zone_order)}"
        )
    _path_check_report()
    return records, m


def evaluate_points(
    points: Sequence[Dict[str, Any]],
    *,
    mode: str,
    cfg: ForecastConfig,
    gate: ProductGateConfig,
    impact_store: Any = None,
    model: Any = None,
    env: Any = None,
) -> Tuple[List[EvalDayRecord], EvalMetrics]:
    """
    Contract: evaluate_points

    Purpose: Score product decisions on real historical zone-days (single-zone).
    For multi-zone agent eval (n_zones>1), use evaluate_multi_zone_days.
    GT: if impact_store is loaded:
          day inside any event span for the zone → gt_source="l1"
          otherwise → gt_source="unlabeled" (NOT a confirmed negative)
        if no store → gt_source="none"
    Confusion matrix (TP/FP/FN/TN) uses only gt_source="l1" rows.
    Unlabeled rows contribute only to n_unlabeled / unlabeled_alert_rate.
    Alert: product gate only.
    Idempotency: pure function of inputs; no transport side effects.
    """
    _path_check_reset()
    records: List[EvalDayRecord] = []
    m = EvalMetrics()

    for pt in points:
        obs = ZoneObs.from_dict(dict(pt["obs"]))
        fc = ForecastResult.from_dict(dict(pt["forecast"]))
        vt = obs.valid_time
        if vt.tzinfo is None:
            vt = vt.replace(tzinfo=timezone.utc)
        day = vt.date()
        zid = obs.zone_id

        event_d = event_f = False
        gt_source = "none"
        if impact_store is not None:
            try:
                ld, lf = impact_store.labels_for_day(zid, day)
                event_d, event_f = bool(ld), bool(lf)
                if event_d or event_f:
                    # Day falls inside an L1 event span → positive label
                    gt_source = "l1"
                else:
                    # Store loaded, day outside every span for this zone.
                    # Sparse catalogs must NOT treat these as confirmed TN/FP.
                    gt_source = "unlabeled"
            except Exception as e:
                logger.warning("L1 query failed %s %s: %s", zid, day, e)
                gt_source = "none"

        believed_p = 0.0
        initial_belief = 0.0
        if mode == "scorer_oracle":
            product, elevated, rs, ep_len, believed_p, initial_belief = (
                decide_scorer_oracle(obs, fc, cfg, gate)
            )
        elif mode == "always":
            product, elevated, rs, ep_len, believed_p, initial_belief = decide_always()
        elif mode == "never":
            product, elevated, rs, ep_len, believed_p, initial_belief = decide_never()
        elif mode == "checkpoint":
            if model is None or env is None:
                raise RuntimeError("checkpoint mode requires --checkpoint and env")
            ctx = point_to_episode(pt, cfg)
            product, elevated, rs, ep_len, believed_p, initial_belief = (
                decide_checkpoint(model, env, ctx, gate)
            )
        elif mode == "zero_inspect":
            if env is None:
                raise RuntimeError("zero_inspect mode requires env")
            ctx = point_to_episode(pt, cfg)
            product, elevated, rs, ep_len, believed_p, initial_belief = (
                decide_zero_inspect(env, ctx, gate)
            )
        else:
            raise ValueError(f"unknown mode: {mode}")

        alert_level = (
            rs.alert_level.value if rs is not None else ("warning" if product else "none")
        )
        drought_risk = float(rs.drought_risk) if rs is not None else 0.0
        flood_risk = float(rs.flood_risk) if rs is not None else 0.0

        # Policy-sensitive: relative movement, not fixed bar vs prior/rational.
        # Fixed bars saturate when prior > rational (prior=0.12, rational=0.0625).
        belief_delta = float(believed_p) - float(initial_belief)
        belief_raised = bool(belief_delta > BELIEF_RAISE_EPS)

        # Product / belief confusion matrix: L1 event days only.
        # Unlabeled days are excluded from TP/FP/FN/TN (sparse catalog honesty).
        if gt_source == "l1":
            # v1 catalog only stores positive impact spans → TP or FN
            if product:
                m.tp += 1
            else:
                m.fn += 1
            if belief_raised:
                m.belief_tp += 1
            else:
                m.belief_fn += 1
            m.n_l1 += 1
            m.belief_deltas_l1.append(belief_delta)
        elif gt_source == "unlabeled":
            m.n_unlabeled += 1
            if product:
                m.n_unlabeled_product += 1
            if belief_raised:
                m.n_unlabeled_belief_raised += 1
            m.belief_deltas_unlabeled.append(belief_delta)
            # No TN/FP: catalog never asserted "no impact" for this day

        m.n_days += 1
        if product:
            m.n_product += 1
        if belief_raised:
            m.n_belief_raised += 1
        m.belief_deltas.append(belief_delta)

        records.append(
            EvalDayRecord(
                date=day.isoformat(),
                zone_id=zid,
                mode=mode,
                product_alert=product,
                elevated=elevated,
                alert_level=alert_level,
                drought_risk=drought_risk,
                flood_risk=flood_risk,
                event_drought=event_d,
                event_flood=event_f,
                gt_source=gt_source,
                ep_len=ep_len,
                basin_dim=8,
                believed_p=float(believed_p),
                initial_belief=float(initial_belief),
                belief_delta=belief_delta,
                belief_raised=belief_raised,
            )
        )
    _path_check_report()
    return records, m


# ---------------------------------------------------------------------------
# CLI
# ---------------------------------------------------------------------------

def _print_metrics(mode: str, m: EvalMetrics) -> None:
    def _f(x: Optional[float]) -> str:
        return f"{x:.3f}" if x is not None else "  -  "

    print(f"\n=== real eval mode={mode} ===")
    print(
        f"n={m.n_days}  l1_event_days={m.n_l1}  unlabeled={m.n_unlabeled}  "
        f"product_alerts={m.n_product}  belief_raised={m.n_belief_raised}"
    )
    print(
        "--- product vs L1 event spans only "
        "(unlabeled excluded from confusion matrix) ---"
    )
    print(f"TP={m.tp} FP={m.fp} FN={m.fn} TN={m.tn}")
    print(f"P={_f(m.precision)}  R={_f(m.recall)}  F1={_f(m.f1)}")
    if m.n_unlabeled > 0:
        print(
            f"--- unlabeled (outside every L1 span; not in P/R/F1) ---"
        )
        print(
            f"n_unlabeled={m.n_unlabeled}  "
            f"unlabeled_product={m.n_unlabeled_product}  "
            f"unlabeled_alert_rate={_f(m.unlabeled_alert_rate)}"
        )
        print(
            f"unlabeled_belief_raised={m.n_unlabeled_belief_raised}  "
            f"belief_unlabeled_raise_rate={_f(m.belief_unlabeled_raise_rate)}"
        )
    print(
        f"--- belief_raised on L1 event days "
        f"(delta > {BELIEF_RAISE_EPS} vs episode initial) ---"
    )
    print(
        f"belief_tp={m.belief_tp}  belief_fn={m.belief_fn}  "
        f"(belief_fp/tn unused under positive-only L1)"
    )
    print(
        f"belief_precision={_f(m.belief_precision)}  "
        f"belief_recall={_f(m.belief_recall)}  "
        f"belief_f1={_f(m.belief_f1)}  "
        f"[P/F1 null by construction; R is the real number]"
    )
    bd = m.belief_delta_stats()
    if bd["n"] > 0:
        print("--- belief_delta = terminal − initial (all days) ---")
        print(
            f"  n={int(bd['n'])}  mean={bd['mean']:+.4f}  std={bd['std']:.4f}  "
            f"p10={bd['p10']:+.4f}  p50={bd['p50']:+.4f}  p90={bd['p90']:+.4f}"
        )
        print(
            f"  frac_pos(>{BELIEF_RAISE_EPS})={bd['frac_pos']:.3f}  "
            f"frac_neg(<-{BELIEF_RAISE_EPS})={bd['frac_neg']:.3f}  "
            f"frac_|delta|<{BELIEF_RAISE_EPS}={bd['frac_near0']:.3f}"
        )
    if m.belief_deltas_l1 or m.belief_deltas_unlabeled:
        def _mean(xs: List[float]) -> str:
            if not xs:
                return "  -  "
            return f"{sum(xs)/len(xs):+.4f}"
        print(
            f"  mean Δ | L1 event     {_mean(m.belief_deltas_l1)}  "
            f"n={len(m.belief_deltas_l1)}"
        )
        print(
            f"  mean Δ | unlabeled   {_mean(m.belief_deltas_unlabeled)}  "
            f"n={len(m.belief_deltas_unlabeled)}"
        )


def main(argv: Optional[Sequence[str]] = None) -> int:
    logging.basicConfig(level=logging.INFO, format="%(levelname)s %(name)s: %(message)s")
    p = argparse.ArgumentParser(description="Real-trajectory product eval (L1)")
    p.add_argument("--pkl", required=True, help="historical_continuous_indonesia_v1.pkl")
    p.add_argument("--impact-labels", default=None, help="impact_labels_java_v1.json")
    p.add_argument("--zones", default="karawang_rice,indramayu_rice")
    p.add_argument("--start", required=True, help="YYYY-MM-DD")
    p.add_argument("--end", required=True, help="YYYY-MM-DD")
    p.add_argument(
        "--mode",
        choices=("scorer_oracle", "checkpoint", "zero_inspect", "always", "never"),
        default="scorer_oracle",
        help="checkpoint=agent rollout; zero_inspect=terminate step 1 control; "
             "scorer_oracle=deterministic product gate only",
    )
    p.add_argument("--checkpoint", default=None, help="MaskablePPO .zip (mode=checkpoint)")
    p.add_argument("--n-zones", type=int, default=1)
    p.add_argument(
        "--max-steps",
        type=int,
        default=4,
        help="Eval episode length cap (n_zones=1 only needs ~2; does not "
             "need to match train max_steps).",
    )
    p.add_argument(
        "--horizon-days",
        type=int,
        default=30,
        help="Must match the checkpoint's forecast_precip width "
             "(train_kaggle / ForecastConfig default is 30).",
    )
    p.add_argument("--device", default="cpu")
    p.add_argument("--out", default=None, help="Write full JSON result")
    args = p.parse_args(list(argv) if argv is not None else None)

    zone_ids = [z.strip() for z in args.zones.split(",") if z.strip()]
    start = _parse_day(args.start)
    end = _parse_day(args.end)
    pkl_path = Path(args.pkl)
    if not pkl_path.is_file():
        print(f"FILE_NOT_FOUND: {pkl_path}", file=sys.stderr)
        return 2

    points = load_historical_points(pkl_path, zone_ids, start, end)
    print(f"loaded points: {len(points)}  zones={zone_ids}  {start}→{end}")
    if not points:
        print("NO_POINTS in window/zones", file=sys.stderr)
        return 3

    impact_store = None
    if args.impact_labels:
        from impact_labels import load_impact_events

        res = load_impact_events(args.impact_labels)
        if not res.success:
            print(f"L1 load failed: {res.outcome_code}", file=sys.stderr)
            return 4
        impact_store = res.data["store"]
        print(f"L1: {res.outcome_code} events_loaded={res.data.get('events_loaded')}")

    cfg = ForecastConfig(
        n_zones=args.n_zones,
        max_steps=args.max_steps,
        soft_reset=True,
        horizon_days=int(args.horizon_days),
    )
    gate = DEFAULT_PRODUCT_GATE

    model = None
    env = None
    if args.mode in ("checkpoint", "zero_inspect"):
        env = make_weather_env(cfg, use_nan_wrapper=True)
        obs_space = env.observation_space
        if hasattr(env, "env"):
            obs_space = env.env.observation_space
        bshape = obs_space["basin_context"].shape
        precip_shape = obs_space["forecast_precip"].shape
        print(f"env basin_context shape: {bshape}")
        print(f"env forecast_precip shape: {precip_shape}")
        print(f"env horizon_days={cfg.horizon_days} n_zones={cfg.n_zones}")

    if args.mode == "checkpoint":
        if not args.checkpoint:
            print("checkpoint mode requires --checkpoint", file=sys.stderr)
            return 5
        try:
            from sb3_contrib import MaskablePPO
        except ImportError:
            print("sb3_contrib not installed", file=sys.stderr)
            return 6
        model = MaskablePPO.load(args.checkpoint, device=args.device)
        if int(np.prod(bshape)) != 8:
            print(
                "WARNING: env basin_context is not 8-dim; "
                "checkpoint may be incompatible",
                file=sys.stderr,
            )
        try:
            pol_space = model.observation_space
            pol_precip = pol_space["forecast_precip"].shape
            if tuple(pol_precip) != tuple(precip_shape):
                print(
                    f"SHAPE_MISMATCH: policy forecast_precip {pol_precip} "
                    f"!= env {precip_shape}. Re-run with "
                    f"--horizon-days matching training (usually 30).",
                    file=sys.stderr,
                )
                return 7
            pol_zones = pol_space["zone_belief"].shape
            env_zones = obs_space["zone_belief"].shape
            if tuple(pol_zones) != tuple(env_zones):
                print(
                    f"SHAPE_MISMATCH: policy zone_belief {pol_zones} "
                    f"!= env {env_zones}. Use --n-zones matching training "
                    f"(run_smoke=1, run_nz3=3).",
                    file=sys.stderr,
                )
                return 8
        except Exception as e:
            logger.warning("could not cross-check policy obs space: %s", e)

    # Multi-zone agent eval needs same-day packs with zone_obs lists (Blocker B).
    # Scorer/always/never stay single-zone day metrics.
    if (
        args.mode in ("checkpoint", "zero_inspect")
        and int(args.n_zones) > 1
    ):
        if len(zone_ids) < int(args.n_zones):
            print(
                f"NEED_ZONES: --n-zones={args.n_zones} but only "
                f"{len(zone_ids)} zones listed in --zones",
                file=sys.stderr,
            )
            return 9
        zone_order = zone_ids[: int(args.n_zones)]
        print(f"multi-zone real eval zone_order={zone_order}")
        records, metrics = evaluate_multi_zone_days(
            points,
            mode=args.mode,
            cfg=cfg,
            gate=gate,
            zone_order=zone_order,
            impact_store=impact_store,
            model=model,
            env=env,
        )
    else:
        records, metrics = evaluate_points(
            points,
            mode=args.mode,
            cfg=cfg,
            gate=gate,
            impact_store=impact_store,
            model=model,
            env=env,
        )
    _print_metrics(args.mode, metrics)

    # Per-zone breakdown (same unlabeled exclusion rule)
    by_zone: Dict[str, EvalMetrics] = {}
    for r in records:
        zm = by_zone.setdefault(r.zone_id, EvalMetrics())
        zm.n_days += 1
        if r.gt_source == "l1":
            if r.product_alert:
                zm.tp += 1
            else:
                zm.fn += 1
            if r.belief_raised:
                zm.belief_tp += 1
            else:
                zm.belief_fn += 1
            zm.n_l1 += 1
        elif r.gt_source == "unlabeled":
            zm.n_unlabeled += 1
            if r.product_alert:
                zm.n_unlabeled_product += 1
            if r.belief_raised:
                zm.n_unlabeled_belief_raised += 1
        if r.product_alert:
            zm.n_product += 1
        if r.belief_raised:
            zm.n_belief_raised += 1
    print("\nper-zone:")
    for zid, zm in by_zone.items():
        def _f(x: Optional[float]) -> str:
            return f"{x:.3f}" if x is not None else "  -  "
        print(
            f"  {zid:22s} n={zm.n_days:3d} l1={zm.n_l1:3d} "
            f"unlab={zm.n_unlabeled:3d} "
            f"R={_f(zm.recall)} belief_R={_f(zm.belief_recall)} "
            f"TP={zm.tp} FN={zm.fn}  "
            f"unlab_alert={_f(zm.unlabeled_alert_rate)} "
            f"unlab_belief_raise={_f(zm.belief_unlabeled_raise_rate)}"
        )

    result = {
        "mode": args.mode,
        "start": start.isoformat(),
        "end": end.isoformat(),
        "zones": zone_ids,
        "metrics": metrics.to_dict(),
        "per_zone": {z: m.to_dict() for z, m in by_zone.items()},
        "records": [asdict(r) for r in records],
        "gate": {
            "drought_threshold": gate.drought_threshold,
            "flood_threshold": gate.flood_threshold,
            "min_alert_level": gate.min_alert_level.value,
        },
    }
    if args.out:
        with open(args.out, "w") as f:
            json.dump(result, f, indent=2)
        print(f"\nWrote {args.out}")
    return 0


# ---------------------------------------------------------------------------
# Offline self-test (no pkl required)
# ---------------------------------------------------------------------------

def _self_test() -> None:
    print("evaluate_checkpoint_real self-test")
    from zone_observation import make_synthetic_zone_obs, make_synthetic_forecast_result

    obs = make_synthetic_zone_obs("karawang_rice", drought=True, seed=1)
    fc = make_synthetic_forecast_result(
        zone_id="karawang_rice", valid_time=obs.valid_time, drought=True, seed=1
    )
    cfg = ForecastConfig()
    gate = DEFAULT_PRODUCT_GATE
    product, elevated, rs, _, _bp, _ib = decide_scorer_oracle(obs, fc, cfg, gate)
    assert rs is not None
    print(f"  drought scorer product={product} elevated={elevated} "
          f"alert={rs.alert_level.value} drought_risk={rs.drought_risk:.3f}")
    # Metrics arithmetic — with confirmed negatives, P is real
    m = EvalMetrics(n_days=4, tp=1, fp=1, fn=1, tn=1, n_l1=2, n_product=2)
    assert abs((m.precision or 0) - 0.5) < 1e-9
    assert abs((m.recall or 0) - 0.5) < 1e-9
    print("  metrics arithmetic OK")

    # Positive-only L1: precision/f1 must be None, not 1.0
    m_pos = EvalMetrics(n_days=10, tp=7, fn=3, n_l1=10, fp=0, tn=0)
    assert m_pos.precision is None, m_pos.precision
    assert m_pos.f1 is None
    assert abs((m_pos.recall or 0) - 0.7) < 1e-9
    assert m_pos.belief_precision is None
    assert m_pos.belief_f1 is None
    print("  positive-only null precision OK")

    # Unlabeled rate symmetry
    m_u = EvalMetrics(
        n_unlabeled=20,
        n_unlabeled_product=8,
        n_unlabeled_belief_raised=11,
    )
    assert abs((m_u.unlabeled_alert_rate or 0) - 0.4) < 1e-9
    assert abs((m_u.belief_unlabeled_raise_rate or 0) - 0.55) < 1e-9
    print("  unlabeled rates OK")

    # Console formatter must not turn None into 0.000
    def _f(x):
        return f"{x:.3f}" if x is not None else "  -  "
    assert _f(None) == "  -  "
    assert _f(0.0) == "0.000"
    assert "0.000" not in _f(None)
    print("  None print formatting OK")

    print("All evaluate_checkpoint_real self-tests passed.")


if __name__ == "__main__":
    if len(sys.argv) == 1:
        _self_test()
    else:
        raise SystemExit(main())