backtest_indonesia.py

"""
backtest_indonesia.py
=====================
Historical replay / evaluation harness for the Indonesia weather-risk stack.

PIPELINE PER REPLAY STEP (per zone, per date)
---------------------------------------------
  1. obs: live mode -> era5_data_pipeline.fetch_zone_obs (Open-Meteo archive
     for historical dates -- no credentials needed; ERA5/CDS and GEE are
     deliberately bypassed by forcing OPENMETEO_LIVE so the replay does not
     depend on paid/configured services).
     synthetic mode -> deterministic make_synthetic_zone_obs with planted
     event blocks (offline, CI-friendly).
  2. anomalies: climatology.apply_climatology_anomalies with a PINNED
     climatology whose period ends BEFORE the replay window starts
     (end_year = replay_start.year - 1).
  3. forecast: timesfm_wrapper 'baseline' backend (persistence /
     climatology-reverting, derived from the obs only -- no look-ahead).
  4. score: crop_risk_scorer.compute_risk_score -> alert level per day.
"""

from __future__ import annotations

import argparse
import json
import logging
import sys
from dataclasses import dataclass, field, asdict
from datetime import datetime, timedelta, timezone
from typing import Any, Dict, List, Optional, Sequence, Tuple

import zone_observation as _zo

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

from zone_observation import (
    AlertLevel,
    DataSource,
    ForecastConfig,
    ZoneObs,
    make_synthetic_zone_obs,
)

from climatology import (
    ZoneClimatology,
    apply_climatology_anomalies,
    get_zone_climatology,
)
from crop_risk_scorer import compute_risk_score
from indonesia_zones import INDONESIA_ZONES, get_zone, register_indonesia_zones

logger = logging.getLogger(__name__)

_ALERT_POSITIVE = (AlertLevel.ADVISORY, AlertLevel.WARNING, AlertLevel.CRITICAL)

# Ground-truth percentile cut-points (Gaussian approx, see module docstring).
_DROUGHT_Z = 0.84   # ~20th percentile of the 30-day window aggregate
_FLOOD_Z = 1.28     # ~90th percentile of the 7-day window aggregate

# Event runs separated by a single quiet day are merged (monsoon hazards
# are persistent; a 1-day lull is not a new event).
_RUN_MERGE_GAP_DAYS = 1
# An alert this many days before an event run's start counts as early warning.
_LEAD_WINDOW_DAYS = 21


# ---------------------------------------------------------------------------
# Per-day record + metrics
# ---------------------------------------------------------------------------

@dataclass
class DayRecord:
    date: str                 # ISO date
    zone_id: str
    alert: bool
    alert_level: str
    drought_risk: float
    flood_risk: float
    event_drought: bool
    event_flood: bool
    precip_30d_mm: float
    precip_anomaly_idx: float
    source: str
    gt_source: str = "proxy"          # "proxy" | "l1" | "proxy+l1_miss"
    product_emit: str = "skipped"     # outcome_code or skipped/disabled

    @property
    def event(self) -> bool:
        return self.event_drought or self.event_flood


@dataclass
class BacktestMetrics:
    n_days: int = 0
    tp: int = 0
    fp: int = 0
    fn: int = 0
    tn: int = 0
    n_event_runs: int = 0
    n_runs_detected: int = 0
    mean_lead_days: Optional[float] = None
    drought_recall: Optional[float] = None
    flood_recall: Optional[float] = None

    @property
    def precision(self) -> Optional[float]:
        return self.tp / (self.tp + self.fp) if (self.tp + self.fp) else None

    @property
    def recall(self) -> Optional[float]:
        return self.tp / (self.tp + self.fn) if (self.tp + self.fn) 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 event_detection_rate(self) -> Optional[float]:
        return (self.n_runs_detected / self.n_event_runs
                if self.n_event_runs else None)

    def to_dict(self) -> Dict[str, Any]:
        d = asdict(self)
        d["precision"] = self.precision
        d["recall"] = self.recall
        d["f1"] = self.f1
        d["event_detection_rate"] = self.event_detection_rate
        return d


def compute_metrics(records: Sequence[DayRecord]) -> BacktestMetrics:
    m = BacktestMetrics(n_days=len(records))
    for r in records:
        if r.event and r.alert:
            m.tp += 1
        elif r.event and not r.alert:
            m.fn += 1
        elif not r.event and r.alert:
            m.fp += 1
        else:
            m.tn += 1

    drought_days = [r for r in records if r.event_drought]
    flood_days = [r for r in records if r.event_flood]
    if drought_days:
        m.drought_recall = sum(1 for r in drought_days if r.alert) / len(drought_days)
    if flood_days:
        m.flood_recall = sum(1 for r in flood_days if r.alert) / len(flood_days)

    # --- Build event runs ---
    runs: List[Tuple[int, int]] = []  # (start_idx, end_idx) inclusive
    i = 0
    n = len(records)
    while i < n:
        if records[i].event:
            j = i
            while j + 1 < n and (
                records[j + 1].event
                or (j + 2 < n and records[j + 2].event)  # peek over 1 gap day
            ):
                if records[j + 1].event:
                    j += 1
                elif j + 2 < n and records[j + 2].event and (j + 2) - (j + 1) <= _RUN_MERGE_GAP_DAYS:
                    j += 2
                else:
                    break
            runs.append((i, j))
            i = j + 1
        else:
            i += 1

    m.n_event_runs = len(runs)
    leads: List[int] = []
    dates = [datetime.fromisoformat(r.date) for r in records]
    for (s, e) in runs:
        first_alert_idx: Optional[int] = None
        for k in range(max(0, s - _LEAD_WINDOW_DAYS), e + 1):
            if records[k].alert:
                first_alert_idx = k
                break
        if first_alert_idx is not None:
            m.n_runs_detected += 1
            leads.append(max(0, (dates[s] - dates[first_alert_idx]).days))
    if leads:
        m.mean_lead_days = sum(leads) / len(leads)
    return m


def compute_product_l1_metrics(records: Sequence[Dict[str, Any]]) -> Dict[str, Any]:
    n = len(records)
    tp = fp = fn = tn = 0
    n_l1 = 0
    n_emitted = 0
    for r in records:
        emitted = r.get("product_emit") == "EMITTED"
        l1 = r.get("gt_source") == "l1" and (
            r.get("event_drought") or r.get("event_flood")
        )
        if l1:
            n_l1 += 1
        if emitted:
            n_emitted += 1
        if l1 and emitted:
            tp += 1
        elif l1 and not emitted:
            fn += 1
        elif not l1 and emitted:
            fp += 1
        else:
            tn += 1

    def _div(a: int, b: int) -> Optional[float]:
        return a / b if b else None

    p = _div(tp, tp + fp)
    r = _div(tp, tp + fn)
    f1 = (2 * p * r / (p + r)) if (p and r and (p + r)) else None
    return {
        "n_days": n,
        "n_l1_event_days": n_l1,
        "n_emitted": n_emitted,
        "tp": tp,
        "fp": fp,
        "fn": fn,
        "tn": tn,
        "precision": p,
        "recall": r,
        "f1": f1,
        "emit_rate": _div(n_emitted, n),
        "l1_coverage": _div(n_l1, n),
    }


# ---------------------------------------------------------------------------
# Replay engine
# ---------------------------------------------------------------------------

def _classify_events(obs: ZoneObs, clim: ZoneClimatology) -> Tuple[bool, bool]:
    event_drought = False
    event_flood = False
    if obs.precip_30d_mm > 0.0:
        mean30, std30 = clim.window_precip_stats(obs.valid_time, 30)
        event_drought = obs.precip_30d_mm < (mean30 - _DROUGHT_Z * std30)
    if obs.precip_7d_mm > 0.0:
        mean7, std7 = clim.window_precip_stats(obs.valid_time, 7)
        event_flood = obs.precip_7d_mm > (mean7 + _FLOOD_Z * std7)
    return event_drought, event_flood


def replay_zone(
    zone_id: str,
    start: datetime,
    end: datetime,
    step_days: int = 3,
    mode: str = "synthetic",
    climatology_years: int = 10,
    planted_events: Optional[Dict[str, Tuple[datetime, datetime]]] = None,
    impact_store: Any = None,
    emit_product: bool = False,
    emission_ledger: Any = None,
    transport: Any = None,
) -> List[DayRecord]:
    if start.tzinfo is None:
        start = start.replace(tzinfo=timezone.utc)
    if end.tzinfo is None:
        end = end.replace(tzinfo=timezone.utc)

    z = get_zone(zone_id)

    # Pinned climatology: strictly before the replay window (no look-ahead).
    clim = get_zone_climatology(
        zone_id, z.lat, z.lon,
        years=climatology_years,
        end_year=start.year - 1,
        prefer_real=(mode == "live"),
    )
    logger.info(
        "replay %s: climatology source=%s period=%d-%d",
        zone_id, clim.source, clim.period_start_year, clim.period_end_year,
    )

    if mode == "live":
        import era5_data_pipeline as edp
        from timesfm_wrapper import create_forecast_backend
        forecast_backend = create_forecast_backend(mode="baseline", horizon_days=14)
        cfg = ForecastConfig(force_data_source=DataSource.OPENMETEO_LIVE)

    records: List[DayRecord] = []
    vt = start
    while vt <= end:
        if mode == "live":
            dr = (vt - timedelta(days=35), vt)
            obs = edp.fetch_zone_obs(zone_id, dr, cfg)
            forecast = forecast_backend.forecast(obs)
        else:
            flag: Dict[str, bool] = {}
            for hazard, blk in (planted_events or {}).items():
                if blk[0] <= vt <= blk[1]:
                    flag[hazard] = True
            obs = make_synthetic_zone_obs(
                zone_id,
                seed=_zo._stable_seed(f"{zone_id}|{vt.date().isoformat()}"),
                **flag,
            )
            _d = obs.to_dict()
            _d.pop("_schema_version", None)
            _d["valid_time"] = vt.isoformat()
            obs = ZoneObs.from_dict(_d)
            from zone_observation import make_synthetic_forecast_result
            forecast = make_synthetic_forecast_result(
                zone_id, valid_time=vt,
                seed=_zo._stable_seed(f"f|{zone_id}|{vt.date().isoformat()}"),
                **flag,
            )

        obs = apply_climatology_anomalies(obs, clim)
        risk = compute_risk_score(obs, forecast, ForecastConfig())
        ev_drought, ev_flood = _classify_events(obs, clim)
        gt_source = "proxy"

        if impact_store is not None:
            try:
                l1_d, l1_f = impact_store.labels_for_day(zone_id, vt.date())
                if l1_d or l1_f:
                    ev_drought, ev_flood = l1_d, l1_f
                    gt_source = "l1"
                else:
                    gt_source = "proxy+l1_miss"
            except Exception as e:
                logger.warning("impact_store lookup failed: %s", e)

        product_emit = "disabled"
        if emit_product and emission_ledger is not None and transport is not None:
            try:
                import asyncio
                from product_alert_service import emit_product_alert

                async def _one():
                    return await emit_product_alert(
                        transport, risk,
                        ledger=emission_ledger,
                        valid_time=vt,
                    )

                try:
                    loop = asyncio.get_running_loop()
                except RuntimeError:
                    loop = None
                if loop and loop.is_running():
                    import concurrent.futures
                    with concurrent.futures.ThreadPoolExecutor(max_workers=1) as pool:
                        product_emit = pool.submit(lambda: asyncio.run(_one())).result().outcome_code
                else:
                    product_emit = asyncio.run(_one()).outcome_code
            except Exception as e:
                logger.warning("product emit failed: %s", e)
                product_emit = "TRANSPORT_FAILED"

        records.append(DayRecord(
            date=vt.date().isoformat(),
            zone_id=zone_id,
            alert=risk.alert_level in _ALERT_POSITIVE,
            alert_level=risk.alert_level.value,
            drought_risk=round(risk.drought_risk, 4),
            flood_risk=round(risk.flood_risk, 4),
            event_drought=ev_drought,
            event_flood=ev_flood,
            precip_30d_mm=round(obs.precip_30d_mm, 1),
            precip_anomaly_idx=round(obs.precip_anomaly_idx, 3),
            source=obs.source.value,
            gt_source=gt_source,
            product_emit=product_emit,
        ))
        vt += timedelta(days=step_days)
    return records


def run_backtest(
    zone_ids: Sequence[str],
    start: datetime,
    end: datetime,
    step_days: int = 3,
    mode: str = "synthetic",
    climatology_years: int = 10,
    planted_events: Optional[Dict[str, Tuple[datetime, datetime]]] = None,
    impact_labels_path: Optional[str] = None,
    emit_product_alerts: bool = False,
) -> Dict[str, Any]:
    if mode == "live":
        register_indonesia_zones()

    impact_store = None
    if impact_labels_path:
        from impact_labels import load_impact_events
        loaded = load_impact_events(impact_labels_path)
        if not loaded.success:
            raise RuntimeError(
                f"impact labels load failed: {loaded.outcome_code} {loaded.data}"
            )
        impact_store = loaded.data["store"]
        logger.info(
            "L1 impact labels: %s events_loaded=%s",
            loaded.outcome_code, loaded.data.get("events_loaded"),
        )

    emission_ledger = None
    transport = None
    if emit_product_alerts:
        from product_alert_service import EmissionLedger
        from node_transport import LocalTransport
        emission_ledger = EmissionLedger()
        transport = LocalTransport()

    all_records: List[DayRecord] = []
    per_zone: Dict[str, Any] = {}
    for zid in zone_ids:
        recs = replay_zone(
            zid, start, end, step_days=step_days, mode=mode,
            climatology_years=climatology_years,
            planted_events=planted_events,
            impact_store=impact_store,
            emit_product=emit_product_alerts,
            emission_ledger=emission_ledger,
            transport=transport,
        )
        all_records.extend(recs)
        per_zone[zid] = {
            "metrics": compute_metrics(recs).to_dict(),
            "n_records": len(recs),
            "sources": sorted({r.source for r in recs}),
        }
        logger.info("zone %s: %s", zid, per_zone[zid]["metrics"])

    overall = compute_metrics(all_records)
    return {
        "mode": mode,
        "window": [start.date().isoformat(), end.date().isoformat()],
        "step_days": step_days,
        "climatology_years": climatology_years,
        "overall": overall.to_dict(),
        "per_zone": per_zone,
        "records": [asdict(r) for r in all_records],
    }


def _print_report(result: Dict[str, Any]) -> None:
    o = result["overall"]

    def _f(x: Optional[float]) -> str:
        return f"{x:.3f}" if isinstance(x, float) else "  -  "

    print(f"\n=== backtest [{result['mode']}] "
          f"{result['window'][0]} -> {result['window'][1]} "
          f"(step {result['step_days']}d) ===")
    print(f"days={o['n_days']}  TP={o['tp']} FP={o['fp']} FN={o['fn']} TN={o['tn']}")
    print(f"precision={_f(o['precision'])}  recall={_f(o['recall'])}  f1={_f(o['f1'])}")
    print(f"event runs: {o['n_runs_detected']}/{o['n_event_runs']} detected "
          f"(rate={_f(o['event_detection_rate'])})  "
          f"mean lead={_f(o['mean_lead_days'])}d")
    print(f"drought recall={_f(o['drought_recall'])}  "
          f"flood recall={_f(o['flood_recall'])}")
    for zid, zr in result["per_zone"].items():
        zm = zr["metrics"]
        print(f"  {zid:24s} n={zr['n_records']:3d} src={','.join(zr['sources'])} "
              f"P={_f(zm['precision'])} R={_f(zm['recall'])} "
              f"runs={zm['n_runs_detected']}/{zm['n_event_runs']}")


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

def _parse_args(argv: Optional[List[str]] = None) -> argparse.Namespace:
    p = argparse.ArgumentParser(
        description="Backtest the Indonesia weather-risk stack over a "
                    "historical window (default invocation with NO arguments "
                    "runs the offline self-test instead).",
    )
    p.add_argument("--mode", choices=["synthetic", "live"], default="live",
                   help="'live' = real Open-Meteo archive data (network); "
                        "'synthetic' = deterministic offline replay.")
    p.add_argument("--zones", default="karawang_rice",
                   help="Comma-separated zone ids from indonesia_zones.")
    p.add_argument("--start", required=True, help="YYYY-MM-DD")
    p.add_argument("--end", required=True, help="YYYY-MM-DD")
    p.add_argument("--step-days", type=int, default=3)
    p.add_argument("--climatology-years", type=int, default=10)
    p.add_argument("--out", default=None, help="Optional JSON output path.")
    p.add_argument(
        "--impact-labels", default=None,
        help="Path to impact_labels JSON (L1). Overrides proxy GT when "
             "zone-day is labeled.",
    )
    p.add_argument(
        "--emit-product-alerts", action="store_true",
        help="After each score, idempotently emit product alerts via "
             "product_alert_service (LocalTransport).",
    )
    return p.parse_args(argv)


def _main(argv: Optional[List[str]] = None) -> int:
    args = _parse_args(argv)
    logging.basicConfig(level=logging.INFO,
                        format="%(levelname)s %(name)s: %(message)s")
    start = datetime.fromisoformat(args.start).replace(tzinfo=timezone.utc)
    end = datetime.fromisoformat(args.end).replace(tzinfo=timezone.utc)
    result = run_backtest(
        zone_ids=[z.strip() for z in args.zones.split(",") if z.strip()],
        start=start, end=end,
        step_days=args.step_days,
        mode=args.mode,
        climatology_years=args.climatology_years,
        impact_labels_path=args.impact_labels,
        emit_product_alerts=args.emit_product_alerts,
    )
    _print_report(result)
    # Product emit + L1 primary metrics
    recs = result.get("records") or []
    if args.emit_product_alerts and recs:
        from collections import Counter
        c = Counter(r.get("product_emit", "skipped") for r in recs)
        print("\nproduct_emit counts:", dict(c))
    if args.impact_labels and recs:
        from collections import Counter
        c = Counter(r.get("gt_source", "proxy") for r in recs)
        print("gt_source counts:", dict(c))
    if args.emit_product_alerts and args.impact_labels and recs:
        pm = compute_product_l1_metrics(recs)
        result["product_l1_metrics"] = pm

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

        print("\n--- product vs L1 (primary product skill) ---")
        print(
            f"n={pm['n_days']}  l1_days={pm['n_l1_event_days']}  "
            f"emitted={pm['n_emitted']}  "
            f"emit_rate={_f(pm['emit_rate'])}  l1_coverage={_f(pm['l1_coverage'])}"
        )
        print(
            f"TP={pm['tp']} FP={pm['fp']} FN={pm['fn']} TN={pm['tn']}  "
            f"P={_f(pm['precision'])} R={_f(pm['recall'])} F1={_f(pm['f1'])}"
        )
    if args.out:
        with open(args.out, "w") as f:
            json.dump(result, f, indent=2)
        print(f"\nWrote {args.out}")
    return 0


# ---------------------------------------------------------------------------
# Self-test (python backtest_indonesia.py) -- fully offline
# ---------------------------------------------------------------------------

def _self_test() -> int:
    logging.basicConfig(level=logging.WARNING)
    print("backtest_indonesia.py self-test (offline, synthetic)\n")
    failures: List[str] = []

    def _assert(cond: bool, msg: str) -> None:
        if not cond:
            failures.append(msg)
            print(f"  FAIL: {msg}")

    def _rec(day: int, alert: bool = False, event: bool = False,
             drought: bool = False, flood: bool = False) -> DayRecord:
        d = (datetime(2024, 1, 1, tzinfo=timezone.utc) + timedelta(days=day))
        return DayRecord(
            date=d.date().isoformat(), zone_id="t",
            alert=alert, alert_level="advisory" if alert else "none",
            drought_risk=0.0, flood_risk=0.0,
            event_drought=(drought or (event and not flood)),
            event_flood=flood,
            precip_30d_mm=0.0, precip_anomaly_idx=0.0, source="synthetic",
        )

    # 1. Metrics on a fabricated series with known expected values.
    #    alerts: 5,6,7 (early), 25 (isolated), 40,41 (inside run 2)
    #    events: 10-14 (run 1), 40-42 (run 2)
    recs: List[DayRecord] = []
    for day in range(50):
        alert = day in (5, 6, 7, 25, 40, 41)
        event = (10 <= day <= 14) or (40 <= day <= 42)
        recs.append(_rec(day, alert=alert, event=event, drought=event))
    m = compute_metrics(recs)
    _assert(m.tp == 2 and m.fn == 6 and m.fp == 4 and m.tn == 38,
            f"confusion wrong: tp={m.tp} fn={m.fn} fp={m.fp} tn={m.tn}")
    _assert(abs((m.precision or 0) - 2 / 6) < 1e-9, f"precision {m.precision}")
    _assert(abs((m.recall or 0) - 2 / 8) < 1e-9, f"recall {m.recall}")
    _assert(m.n_event_runs == 2, f"runs={m.n_event_runs}")
    _assert(m.n_runs_detected == 2, f"detected={m.n_runs_detected}")
    _assert(abs((m.mean_lead_days or 0) - 10.0) < 1e-9,
            f"mean lead {m.mean_lead_days} (expected 10: 5d + 15d)")
    print(f"  Metrics OK: P={m.precision:.3f} R={m.recall:.3f} "
          f"runs {m.n_runs_detected}/{m.n_event_runs} lead={m.mean_lead_days}d")

    # 2. Run merging across a 1-day lull: events 10,11,13 = ONE run.
    recs2 = [_rec(day, event=day in (10, 11, 13), drought=True) for day in range(30)]
    m2 = compute_metrics(recs2)
    _assert(m2.n_event_runs == 1, f"1-day lull should merge: runs={m2.n_event_runs}")
    print(f"  Run-merge OK (runs={m2.n_event_runs})")

    # 3. Event classifier: planted drought reads as drought event vs climatology
    clim = get_zone_climatology("karawang_rice", -6.30, 107.30, years=5,
                                end_year=2020, prefer_real=False)
    vt = datetime(2021, 8, 15, tzinfo=timezone.utc)

    def _obs_at(flag: str, seed: int) -> ZoneObs:
        o = make_synthetic_zone_obs("karawang_rice", seed=seed, **{flag: True})
        _d = o.to_dict()
        _d.pop("_schema_version", None)
        _d["valid_time"] = vt.isoformat()
        return ZoneObs.from_dict(_d)

    dry_obs = _obs_at("drought", 11)
    ev_d, ev_f = _classify_events(dry_obs, clim)
    _assert(ev_d and not ev_f, f"planted drought misclassified: d={ev_d} f={ev_f}")
    wet_obs = _obs_at("flood", 12)
    ev_d2, ev_f2 = _classify_events(wet_obs, clim)
    _assert(ev_f2 and not ev_d2, f"planted flood misclassified: d={ev_d2} f={ev_f2}")
    print("  Event classifier OK (drought/flood classified correctly)")

    # 4. Synthetic end-to-end: 5-month replay with a planted drought block.
    start = datetime(2021, 6, 1, tzinfo=timezone.utc)
    end = datetime(2021, 10, 31, tzinfo=timezone.utc)
    planted = {"drought": (datetime(2021, 8, 1, tzinfo=timezone.utc),
                           datetime(2021, 9, 10, tzinfo=timezone.utc))}
    result = run_backtest(["karawang_rice"], start, end, step_days=5,
                          mode="synthetic", climatology_years=5,
                          planted_events=planted)
    o = result["overall"]
    _assert(o["n_days"] > 25, f"too few replay days: {o['n_days']}")
    _assert((o["drought_recall"] or 0) >= 0.9,
            f"planted drought should be caught: drought_recall={o['drought_recall']}")
    _assert(o["n_event_runs"] >= 1, "no event runs found")
    _assert(o["n_runs_detected"] >= 1, "planted drought run not detected")
    json.dumps(result)  # whole result must be JSON-serialisable
    print(f"  End-to-end OK: drought_recall={o['drought_recall']:.2f} "
          f"runs={o['n_runs_detected']}/{o['n_event_runs']} "
          f"P={o['precision'] if o['precision'] is not None else float('nan'):.2f}")
    _print_report(result)

    print()
    if failures:
        print(f"FAILED  {len(failures)} test(s):")
        for f in failures:
            print(f"  - {f}")
        return 1
    print("All 4 test groups passed.")
    return 0


if __name__ == "__main__":
    if len(sys.argv) > 1:
        sys.exit(_main())
    else:
        sys.exit(_self_test())