build_continuous_historical.py

"""
build_continuous_historical.py
==============================
Causal historical cache for Indonesian rice zones.

Pairs with:
  - era5_data_pipeline.fetch_episode_context (valid_time-anchored, backward
    30-day obs windows, archive URL for dates older than today-5)
  - zone_observation.DataSource.OPENMETEO_ARCHIVE
  - real_episode_sampler.DEFAULT_HOLDOUT_RANGES

This builder does NOT issue NWP. forecast_backend="baseline" is persistence
of causal precip_30d. Do not report forecast skill against this file.

Holdout seasons are stored in the pickle so evaluate_checkpoint_real.py can
see them. RealEpisodeIndex excludes DEFAULT_HOLDOUT_RANGES at train time.
"""
from __future__ import annotations

import argparse
import json
import logging
import pickle
import time
from collections import Counter
from datetime import datetime, timedelta, timezone
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple

import zone_observation as _zo

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

from zone_observation import CropStage, DataSource, ForecastConfig
from indonesia_zones import (
    INDONESIA_ZONES,
    crop_stage_for_date,
    register_indonesia_zones,
)
from era5_data_pipeline import fetch_episode_context

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s | %(levelname)s | %(message)s",
)
logger = logging.getLogger("continuous_cache")

CACHE_VERSION = "indonesia_continuous_v3_causal_obs"

PRIORITY_ZONES = [
    "karawang_rice",
    "indramayu_rice",
    "central_java_rice",
    "east_java_rice",
    "lampung_rice",
    "south_sumatra_rice",
    "banten_rice",
    "south_sulawesi_rice",
]

# Seasons stored in the pickle. Train vs eval is decided by
# real_episode_sampler.DEFAULT_HOLDOUT_RANGES, not by omitting rows here.
#
# DEFAULT_HOLDOUT_RANGES currently:
#   2023-05-01 -> 2024-04-30   el_nino_2023_24_strong
#   2017-05-01 -> 2018-04-30   neutral_2017_18
#
# neutral_2016_17 is the Neutral year that remains eligible for training.
PARADIGMATIC_SEASONS: List[Tuple[str, str, str, str]] = [
    ("elnino_2015_16_vstrong", "2015-05-01", "2016-04-30", "el_nino_very_strong"),
    ("neutral_2016_17",        "2016-05-01", "2017-04-30", "neutral"),
    ("neutral_2017_18",        "2017-05-01", "2018-04-30", "neutral"),
    ("elnino_2018_19",         "2018-06-01", "2019-05-31", "el_nino_moderate"),
    ("lanina_2020_21",         "2020-09-01", "2021-05-31", "la_nina_moderate"),
    ("lanina_2021_22",         "2021-09-01", "2022-05-31", "la_nina_moderate"),
    ("lanina_2022_23",         "2022-09-01", "2023-04-30", "la_nina_weak_moderate"),
    ("elnino_2023_24_strong",  "2023-05-01", "2024-04-30", "el_nino_strong"),
]

# Mirrors real_episode_sampler.DEFAULT_HOLDOUT_RANGES for the QA sidecar.
# If you change the sampler holdout, change this list too (or pass --holdout).
DEFAULT_HOLDOUT_RANGES: Tuple[Tuple[str, str], ...] = (
    ("2023-05-01", "2024-04-30"),
    ("2017-05-01", "2018-04-30"),
)

_ARCHIVE_LAG_DAYS = 5


def _parse(s: str) -> datetime:
    return datetime.strptime(s, "%Y-%m-%d").replace(tzinfo=timezone.utc)


def _daterange(start: datetime, end: datetime, step_days: int = 3):
    if step_days < 1:
        raise ValueError(f"step_days must be >= 1, got {step_days}")
    cur = start
    while cur <= end:
        yield cur
        cur += timedelta(days=step_days)


def _in_holdout(day: datetime, holdout: Tuple[Tuple[str, str], ...]) -> bool:
    d = day.date().isoformat()
    for start_s, end_s in holdout:
        if start_s <= d <= end_s:
            return True
    return False


def _enrich_with_crop_stage(
    obs_dict: Dict[str, Any], zone_id: str, valid_time: datetime
) -> Dict[str, Any]:
    try:
        stage, days_to_harvest, season_name = crop_stage_for_date(zone_id, valid_time)
        obs_dict["crop_stage"] = stage.value if isinstance(stage, CropStage) else str(stage)
        obs_dict["days_to_harvest"] = days_to_harvest
        if obs_dict.get("extras") is None:
            obs_dict["extras"] = {}
        if season_name:
            obs_dict["extras"]["season_name"] = season_name
    except Exception as e:
        logger.debug("crop_stage_for_date failed for %s @ %s: %s", zone_id, valid_time.date(), e)
    return obs_dict


def _safe_to_dict(obj) -> Optional[Dict[str, Any]]:
    if obj is None:
        return None
    if hasattr(obj, "to_dict"):
        return obj.to_dict()
    try:
        return dict(obj.__dict__)
    except Exception:
        return None


def _source_value(obs) -> str:
    src = getattr(obs, "source", None)
    if src is None:
        return "unknown"
    return src.value if hasattr(src, "value") else str(src)


def _assert_causal_obs(ctx: Any, day: datetime) -> None:
    """Fail the point rather than write a leaked or synthetic row."""
    obs = ctx.obs
    if obs is None:
        raise RuntimeError("fetch_episode_context returned ctx.obs=None")

    obs_day = obs.valid_time
    if getattr(obs_day, "tzinfo", None) is None:
        obs_day = obs_day.replace(tzinfo=timezone.utc)
    if obs_day.date() != day.date():
        raise RuntimeError(
            f"obs.valid_time {obs_day.date()} != anchor {day.date()}"
        )

    src = getattr(obs, "source", None)
    if src == DataSource.SYNTHETIC:
        raise RuntimeError(
            f"synthetic obs refused at {day.date()} zone={obs.zone_id}"
        )

    forecast = getattr(ctx, "forecast", None)
    if forecast is not None:
        ft = getattr(forecast, "forecast_time", None)
        if ft is not None:
            if getattr(ft, "tzinfo", None) is None:
                ft = ft.replace(tzinfo=timezone.utc)
            if ft.date() > day.date():
                raise RuntimeError(
                    f"forecast_time {ft.date()} is after anchor {day.date()}"
                )

    today = datetime.now(timezone.utc).date()
    if day.date() < today - timedelta(days=_ARCHIVE_LAG_DAYS):
        archive = getattr(DataSource, "OPENMETEO_ARCHIVE", None)
        if archive is not None and src not in (
            archive,
            DataSource.ERA5_REANALYSIS,
            DataSource.SATELLITE_PRECIP,
            DataSource.SATELLITE_SOIL,
        ):
            raise RuntimeError(
                f"historical day {day.date()} stamped {src} "
                f"(expected OPENMETEO_ARCHIVE / ERA5 / satellite). "
                f"zone_observation.DataSource.OPENMETEO_ARCHIVE is missing "
                f"or the pipeline on this machine is stale."
            )


def _qa_payload(
    trajectories: List[Dict[str, Any]],
    failures: int,
    zones: List[str],
    cfg: ForecastConfig,
    step_days: int,
    holdout: Tuple[Tuple[str, str], ...],
) -> Dict[str, Any]:
    source_counts: Counter = Counter()
    model_counts: Counter = Counter()
    n_train = 0
    n_holdout = 0
    n_missing_basin = 0
    n_synthetic = 0
    n_points = 0

    for traj in trajectories:
        for pt in traj.get("trajectory", []):
            n_points += 1
            obs = pt.get("obs") or {}
            src = str(obs.get("source") or pt.get("data_source") or "unknown")
            source_counts[src] += 1
            if src.endswith("synthetic") or src == "DataSource.SYNTHETIC":
                n_synthetic += 1
            fcast = pt.get("forecast") or {}
            model_counts[str(fcast.get("model_id") or "missing")] += 1
            if pt.get("basin_context") is None:
                n_missing_basin += 1
            vt = pt.get("valid_time") or ""
            day = _parse(vt[:10]) if vt else None
            if day is None:
                continue
            if _in_holdout(day, holdout):
                n_holdout += 1
            else:
                n_train += 1

    return {
        "version": CACHE_VERSION,
        "n_trajectories": len(trajectories),
        "total_points": n_points,
        "failures_dropped": failures,
        "zones": zones,
        "seasons": [s[0] for s in PARADIGMATIC_SEASONS],
        "step_days": step_days,
        "holdout_ranges": [list(h) for h in holdout],
        "n_train_eligible": n_train,
        "n_holdout_stored": n_holdout,
        "n_missing_basin": n_missing_basin,
        "n_synthetic_obs": n_synthetic,
        "source_histogram": dict(source_counts),
        "forecast_model_histogram": dict(model_counts),
        "config_snapshot": {
            "forecast_backend": cfg.forecast_backend,
            "use_climatology_anomalies": cfg.use_climatology_anomalies,
            "include_basin_context": cfg.include_basin_context,
            "require_real_basin_context": cfg.require_real_basin_context,
            "force_data_source": (
                cfg.force_data_source.value if cfg.force_data_source else None
            ),
            "climatology_years": cfg.climatology_years,
        },
        "notes": [
            "forecast_backend=baseline is persistence of causal precip_30d, not NWP",
            "holdout rows are stored for eval; RealEpisodeIndex excludes them at train",
            "n_synthetic_obs must be 0 before this pickle is used for training",
        ],
    }


def build_continuous_cache(
    output_path: str = "historical_continuous_indonesia_v3_causal_obs.pkl",
    step_days: int = 3,
    sleep_s: float = 0.7,
    max_days_per_zone_season: int = 400,
    resume: bool = True,
    require_real_basin_context: bool = False,
    holdout_ranges: Tuple[Tuple[str, str], ...] = DEFAULT_HOLDOUT_RANGES,
) -> None:
    register_indonesia_zones()
    available = {z.zone_id for z in INDONESIA_ZONES}
    zones = [z for z in PRIORITY_ZONES if z in available]
    missing = [z for z in PRIORITY_ZONES if z not in available]
    if missing:
        logger.warning("Priority zones not in registry (skipped): %s", missing)
    if not zones:
        raise RuntimeError("No priority zones registered — aborting")
    logger.info("Priority zones (%d): %s", len(zones), zones)

    if not hasattr(DataSource, "OPENMETEO_ARCHIVE"):
        logger.warning(
            "DataSource.OPENMETEO_ARCHIVE is missing. Historical points "
            "will not pass the archive-stamp assertion. Upload the updated "
            "zone_observation.py before rebuilding."
        )

    cfg = ForecastConfig(
        forecast_backend="baseline",
        use_climatology_anomalies=True,
        include_basin_context=True,
        require_real_basin_context=require_real_basin_context,
        force_data_source=DataSource.OPENMETEO_LIVE,
        real_data_ratio=1.0,
        climatology_years=10,
    )

    trajectories: List[Dict[str, Any]] = []
    failures = 0
    t0 = time.time()
    out = Path(output_path)

    if resume and out.exists():
        try:
            with open(out, "rb") as f:
                existing = pickle.load(f)
            if existing.get("version") != CACHE_VERSION:
                logger.warning(
                    "Resume file version=%s != %s — starting fresh so a "
                    "leaky v1/v2 cache cannot be extended.",
                    existing.get("version"), CACHE_VERSION,
                )
            else:
                trajectories = existing.get("trajectories", [])
                logger.info("Resuming from %d existing trajectories", len(trajectories))
        except Exception as e:
            logger.warning("Resume failed (%s) — starting fresh", e)
            trajectories = []

    already_done = {
        (t["meta"]["label"], t["meta"]["zone_id"]) for t in trajectories
    }

    for label, start_s, end_s, regime in PARADIGMATIC_SEASONS:
        start = _parse(start_s)
        end = _parse(end_s)
        logger.info("=== %s  (%s → %s)  [%s] ===", label, start_s, end_s, regime)

        for zone_id in zones:
            key = (label, zone_id)
            if key in already_done:
                logger.info("  %s already present — skipping", zone_id)
                continue

            traj_points: List[Dict[str, Any]] = []
            days_fetched = 0

            for day in _daterange(start, end, step_days=step_days):
                if days_fetched >= max_days_per_zone_season:
                    logger.warning(
                        "  %s hit max_days=%d before %s — raise --max-days",
                        zone_id, max_days_per_zone_season, end_s,
                    )
                    break

                window_end = day + timedelta(days=30)
                try:
                    ctx = fetch_episode_context(zone_id, (day, window_end), cfg)
                    _assert_causal_obs(ctx, day)

                    obs_dict = _safe_to_dict(ctx.obs) or {}
                    obs_dict = _enrich_with_crop_stage(obs_dict, zone_id, day)

                    point = {
                        "valid_time": day.isoformat(),
                        "zone_id": zone_id,
                        "obs": obs_dict,
                        "forecast": _safe_to_dict(ctx.forecast),
                        "basin_context": _safe_to_dict(
                            getattr(ctx, "basin_context", None)
                        ),
                        "data_source": _source_value(ctx.obs),
                        "split": (
                            "holdout" if _in_holdout(day, holdout_ranges) else "train"
                        ),
                    }
                    traj_points.append(point)
                    days_fetched += 1
                    time.sleep(sleep_s)

                except Exception as e:
                    msg = str(e)
                    logger.warning("  Drop %s @ %s: %s", zone_id, day.date(), msg[:200])
                    failures += 1
                    time.sleep(sleep_s * 1.3)
                    continue

            if traj_points:
                trajectories.append({
                    "meta": {
                        "label": label,
                        "regime": regime,
                        "zone_id": zone_id,
                        "start": start_s,
                        "end": end_s,
                        "n_points": len(traj_points),
                        "step_days": step_days,
                        "n_holdout": sum(
                            1 for p in traj_points if p.get("split") == "holdout"
                        ),
                    },
                    "trajectory": traj_points,
                })
                logger.info("  %s: %d causal points saved", zone_id, len(traj_points))
                if len(trajectories) % 3 == 0:
                    _save(
                        trajectories, out, failures, zones, cfg,
                        step_days, holdout_ranges,
                    )

    _save(trajectories, out, failures, zones, cfg, step_days, holdout_ranges)

    elapsed = (time.time() - t0) / 60
    qa = _qa_payload(trajectories, failures, zones, cfg, step_days, holdout_ranges)
    logger.info("=" * 70)
    logger.info("CAUSAL HISTORICAL CACHE COMPLETE")
    logger.info("  Trajectories     : %d", qa["n_trajectories"])
    logger.info("  Total points     : %d", qa["total_points"])
    logger.info("  Train-eligible   : %d", qa["n_train_eligible"])
    logger.info("  Holdout stored   : %d", qa["n_holdout_stored"])
    logger.info("  Synthetic obs    : %d", qa["n_synthetic_obs"])
    logger.info("  Missing basin    : %d", qa["n_missing_basin"])
    logger.info("  Dropped          : %d", failures)
    logger.info("  Sources          : %s", qa["source_histogram"])
    logger.info("  Forecast models  : %s", qa["forecast_model_histogram"])
    logger.info("  Elapsed          : %.1f min", elapsed)
    logger.info("  Output           : %s", out)
    logger.info("=" * 70)
    if qa["n_synthetic_obs"] != 0:
        raise RuntimeError(
            f"cache contains {qa['n_synthetic_obs']} synthetic obs rows — "
            "do not train on this file"
        )
    if qa["n_train_eligible"] == 0:
        raise RuntimeError("cache has 0 train-eligible points after holdout")


def _save(
    trajectories,
    out: Path,
    failures: int,
    zones,
    cfg: ForecastConfig,
    step_days: int,
    holdout: Tuple[Tuple[str, str], ...],
) -> None:
    qa = _qa_payload(trajectories, failures, zones, cfg, step_days, holdout)
    payload = {
        "version": CACHE_VERSION,
        "created_utc": datetime.now(timezone.utc).isoformat(),
        "design": "continuous_paradigmatic_seasons_causal_obs",
        "n_trajectories": qa["n_trajectories"],
        "total_points": qa["total_points"],
        "priority_zones": zones,
        "holdout_ranges": [list(h) for h in holdout],
        "config_snapshot": qa["config_snapshot"],
        "qa": qa,
        "trajectories": trajectories,
    }
    tmp = out.with_suffix(out.suffix + ".tmp")
    with open(tmp, "wb") as f:
        pickle.dump(payload, f, protocol=pickle.HIGHEST_PROTOCOL)
    tmp.replace(out)

    with open(out.with_suffix(".summary.json"), "w") as f:
        json.dump(qa, f, indent=2)


def main() -> None:
    p = argparse.ArgumentParser(
        description="Build a causal-obs Indonesia historical cache (v3)."
    )
    p.add_argument(
        "--output",
        default="historical_continuous_indonesia_v3_causal_obs.pkl",
    )
    p.add_argument("--step-days", type=int, default=3)
    p.add_argument("--sleep", type=float, default=0.7)
    p.add_argument("--max-days", type=int, default=400)
    p.add_argument("--no-resume", action="store_true")
    p.add_argument(
        "--require-real-basin",
        action="store_true",
        help="Drop the point if RONI/DMI/SWPC cannot be fetched. "
             "Off by default so a DMI 404 does not empty the cache; "
             "ENSO/obs still have to be real.",
    )
    args = p.parse_args()

    build_continuous_cache(
        output_path=args.output,
        step_days=args.step_days,
        sleep_s=args.sleep,
        max_days_per_zone_season=args.max_days,
        resume=not args.no_resume,
        require_real_basin_context=args.require_real_basin,
    )


if __name__ == "__main__":
    main()