""" build_continuous_historical_cache.py ==================================== High-fidelity continuous historical cache for Indonesian rice zones. """ from __future__ import annotations import argparse import json import logging import pickle import time from datetime import datetime, timedelta, timezone from pathlib import Path from typing import Any, Dict, List, Optional import zone_observation as _zo assert _zo.SCHEMA_VERSION == 3 from zone_observation import ForecastConfig, DataSource, CropStage from indonesia_zones import ( register_indonesia_zones, INDONESIA_ZONES, crop_stage_for_date, ) 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") # Priority rice zones PRIORITY_ZONES = [ "karawang_rice", "indramayu_rice", "central_java_rice", "east_java_rice", "lampung_rice", "south_sumatra_rice", "banten_rice", "south_sulawesi_rice", ] # Continuous paradigmatic seasons PARADIGMATIC_SEASONS = [ ("elnino_2015_16_vstrong", "2015-05-01", "2016-04-30", "el_nino_very_strong"), ("elnino_2018_19", "2018-06-01", "2019-05-31", "el_nino_moderate"), ("elnino_2023_24_strong", "2023-05-01", "2024-04-30", "el_nino_strong"), ("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"), ("neutral_2017_18", "2017-05-01", "2018-04-30", "neutral"), ] 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 = 5): cur = start while cur <= end: yield cur cur += timedelta(days=step_days) 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 "extras" not in obs_dict or obs_dict["extras"] is None: obs_dict["extras"] = {} if season_name: obs_dict["extras"]["season_name"] = season_name except Exception: pass return obs_dict def _safe_to_dict(obj) -> Optional[Dict]: 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 build_continuous_cache( output_path: str = "historical_continuous_indonesia_v1.pkl", step_days: int = 5, sleep_s: float = 0.7, max_days_per_zone_season: int = 75, resume: bool = True, ) -> None: register_indonesia_zones() available = {z.zone_id for z in INDONESIA_ZONES} zones = [z for z in PRIORITY_ZONES if z in available] logger.info("Priority zones (%d): %s", len(zones), zones) cfg = ForecastConfig( forecast_backend="baseline", use_climatology_anomalies=True, include_basin_context=True, 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) dmi_warned = False # Resume if resume and out.exists(): try: with open(out, "rb") as f: existing = pickle.load(f) 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) 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: break window_end = day + timedelta(days=30) try: ctx = fetch_episode_context(zone_id, (day, window_end), cfg) 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": str(ctx.data_source), } traj_points.append(point) days_fetched += 1 time.sleep(sleep_s) except Exception as e: msg = str(e) if "dmi.data" in msg.lower() or "DMI" in msg or "404" in msg: if not dmi_warned: logger.warning( "DMI/IOD source unavailable (404) — using synthetic IOD. " "All other real data (Open-Meteo weather, climatology anomalies, ENSO, crop stage) remains intact." ) dmi_warned = True else: logger.warning(" Fail %s @ %s: %s", zone_id, day.date(), msg[:120]) 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, }, "trajectory": traj_points, }) logger.info(" %s: %d ordered points saved", zone_id, len(traj_points)) if len(trajectories) % 3 == 0: _save(trajectories, out, failures, zones, cfg) _save(trajectories, out, failures, zones, cfg) elapsed = (time.time() - t0) / 60 total_points = sum(t["meta"]["n_points"] for t in trajectories) logger.info("=" * 70) logger.info("CONTINUOUS HISTORICAL CACHE COMPLETE") logger.info(" Trajectories : %d", len(trajectories)) logger.info(" Total points : %d", total_points) logger.info(" Failures : %d", failures) logger.info(" Elapsed : %.1f min", elapsed) logger.info(" Output : %s", out) logger.info("=" * 70) def _save(trajectories, out: Path, failures: int, zones, cfg): payload = { "version": "indonesia_continuous_v2", "created_utc": datetime.now(timezone.utc).isoformat(), "design": "continuous_paradigmatic_seasons", "n_trajectories": len(trajectories), "total_points": sum(t["meta"]["n_points"] for t in trajectories), "priority_zones": zones, "config_snapshot": { "forecast_backend": cfg.forecast_backend, "use_climatology_anomalies": cfg.use_climatology_anomalies, "include_basin_context": cfg.include_basin_context, }, "trajectories": trajectories, } with open(out, "wb") as f: pickle.dump(payload, f, protocol=pickle.HIGHEST_PROTOCOL) summary = { "version": payload["version"], "n_trajectories": payload["n_trajectories"], "total_points": payload["total_points"], "failures": failures, "zones": zones, "seasons": [s[0] for s in PARADIGMATIC_SEASONS], "output": str(out), } with open(out.with_suffix(".summary.json"), "w") as f: json.dump(summary, f, indent=2) def main(): p = argparse.ArgumentParser() p.add_argument("--output", default="historical_continuous_indonesia_v1.pkl") p.add_argument("--step-days", type=int, default=5) p.add_argument("--sleep", type=float, default=0.7) p.add_argument("--max-days", type=int, default=75) p.add_argument("--no-resume", action="store_true") 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, ) if __name__ == "__main__": main()