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_N = 0
_PATH_CHECK_DIVERGE = 0

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

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:
    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:

    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
    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
    n_unlabeled: int = 0
    n_unlabeled_product: int = 0
    n_unlabeled_belief_raised: int = 0
    belief_tp: int = 0
    belief_fn: int = 0
    belief_fp: int = 0
    belief_tn: int = 0
    n_belief_raised: int = 0
    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]:
        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]:
        return None

    @property
    def belief_recall(self) -> Optional[float]:
        d = self.belief_tp + self.belief_fn
        return self.belief_tp / d if d else None

    @property
    def belief_f1(self) -> Optional[float]:
        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]:
        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()
        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]]:
    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:
    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:
    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]]]:
    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]:
    rs = compute_risk_score(obs, fc, cfg)
    product = is_product_actionable(rs, gate)
    elevated = rs.is_elevated()
    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:
    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
    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]:
    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:
            believed_p = float(info["mean_belief"])

    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
    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]:
    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
    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]:
    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]:
    _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:
                    gt_source = "l1"
                else:
                    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

        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=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)

    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)

    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}")
    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")

    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")

    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")

    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())