build_continuous_historical.py

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