real_episode_sampler.py

"""
real_episode_sampler.py
========================
Loads a historical trajectory cache (the same format
build_continuous_historical.py produces) and exposes
a sampler that WeatherForecastEnv.reset() can draw
real multi-zone EpisodeContexts from during training.
Import-light by design: weather_forecast_env.py must
not become dependent on the era5_data_pipeline stack
just to read a cache dict that was already built.

--------------------------------------------------------------
DEFAULT_HOLDOUT_RANGES below encodes the exclusion by
whole paradigmatic season (matching
build_continuous_historical.py's own PARADIGMATIC_SEASONS
boundaries), so it unambiguously covers every window already
quoted in technical_details.md and this project's
evaluate_checkpoint_real.py runs:

  - 2023-07-01 -> 2023-11-30 (the "dry product freeze" window) falls
    entirely inside el_nino_2023_24_strong (2023-05-01 -> 2024-04-30).
  - 2017-06-10 (the PATH_CHECK "argmax miss / OR hit" reference example)
    falls inside neutral_2017_18 (2017-05-01 -> 2018-04-30).

Both whole seasons are held out, not just the specific cited days, so
that any future real-eval run inside those seasons remains comparable
to what's already on record.
"""

from __future__ import annotations

import logging
import pickle
import random
from datetime import date
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple

from zone_observation import (
    BasinContext,
    EpisodeContext,
    ForecastConfig,
    ForecastResult,
    ZoneObs,
)

logger = logging.getLogger(__name__)


DEFAULT_HOLDOUT_RANGES: Tuple[Tuple[str, str], ...] = (
    ("2023-05-01", "2024-04-30"),  # el_nino_2023_24_strong
    ("2017-05-01", "2018-04-30"),  # neutral_2017_18
)


def _parse_day(s: str) -> date:
    return date.fromisoformat(str(s)[:10])


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


# ---------------------------------------------------------------------------
# Precip window perturbation
# ---------------------------------------------------------------------------

_PRECIP_WINDOW_FIELDS = (
    "precip_24h_mm", "precip_7d_mm", "precip_14d_mm", "precip_30d_mm",
)


def _perturb_precip_windows(d: Dict[str, Any], scale: float, rng: random.Random) -> None:
    h24 = float(d.get("precip_24h_mm") or 0.0)
    d7  = float(d.get("precip_7d_mm") or 0.0)
    d14 = float(d.get("precip_14d_mm") or 0.0)
    d30 = float(d.get("precip_30d_mm") or 0.0)

    bucket_1d     = h24
    bucket_2_7d   = max(0.0, d7 - h24)
    bucket_8_14d  = max(0.0, d14 - d7)
    bucket_15_30d = max(0.0, d30 - d14)

    def _jitter(x: float) -> float:
        return max(0.0, x * (1.0 + rng.gauss(0.0, scale)))

    bucket_1d     = _jitter(bucket_1d)
    bucket_2_7d   = _jitter(bucket_2_7d)
    bucket_8_14d  = _jitter(bucket_8_14d)
    bucket_15_30d = _jitter(bucket_15_30d)

    d["precip_24h_mm"] = bucket_1d
    d["precip_7d_mm"]  = d["precip_24h_mm"] + bucket_2_7d
    d["precip_14d_mm"] = d["precip_7d_mm"]  + bucket_8_14d
    d["precip_30d_mm"] = d["precip_14d_mm"] + bucket_15_30d


def _perturb_zone_obs(obs: ZoneObs, scale: float, rng: random.Random) -> ZoneObs:
    d = obs.to_dict()
    d.pop("_schema_version", None)

    _perturb_precip_windows(d, scale, rng)

    for field_name in (
        "temp_mean_c", "temp_max_c", "temp_min_c",
        "soil_moisture_pct", "rh_mean_pct", "rh_max_pct",
        "wind_speed_max_ms", "wind_speed_mean_ms", "evapotranspiration_mm",
    ):
        val = d.get(field_name)
        if val is None:
            continue
        jitter = 1.0 + rng.gauss(0.0, scale)
        d[field_name] = max(0.0, float(val) * jitter)

    for field_name in (
        "precip_anomaly_idx", "temp_anomaly_idx",
        "soil_moisture_anom", "rh_anomaly_idx",
    ):
        val = d.get(field_name)
        if val is None:
            continue
        d[field_name] = float(val) + rng.gauss(0.0, scale * 2.0)

    return ZoneObs.from_dict(d)


class RealEpisodeIndex:

    def __init__(
        self,
        pkl_path: str,
        holdout_ranges: Tuple[Tuple[str, str], ...] = DEFAULT_HOLDOUT_RANGES,
    ) -> None:
        self.pkl_path = pkl_path
        self.holdout_ranges = holdout_ranges
        self._by_date: Dict[str, Dict[str, Dict[str, Any]]] = {}
        self._eligible_dates: List[str] = []
        self._all_zone_ids: List[str] = []
        self._load()

    def _load(self) -> None:
        p = Path(self.pkl_path)
        if not p.is_file():
            raise FileNotFoundError(
                f"RealEpisodeIndex: pkl not found at {self.pkl_path}"
            )
        with open(p, "rb") as f:
            cache = pickle.load(f)

        trajs = cache.get("trajectories") or []
        n_points = 0
        n_excluded = 0
        zone_set = set()

        for traj in trajs:
            for pt in traj.get("trajectory") or []:
                obs = pt.get("obs") or {}
                vt = str(obs.get("valid_time") or pt.get("valid_time") or "")[:10]
                zid = obs.get("zone_id") or pt.get("zone_id")
                if not vt or not zid:
                    continue
                try:
                    d = _parse_day(vt)
                except Exception:
                    continue
                if _in_holdout(d, self.holdout_ranges):
                    n_excluded += 1
                    continue
                self._by_date.setdefault(vt, {})[zid] = pt
                zone_set.add(zid)
                n_points += 1

        self._all_zone_ids = sorted(zone_set)
        self._eligible_dates = sorted(self._by_date.keys())

        logger.info(
            "RealEpisodeIndex: %d training-eligible points (%d excluded by "
            "holdout=%s) across %d zones, %d dates, from %s",
            n_points, n_excluded, self.holdout_ranges,
            len(self._all_zone_ids), len(self._eligible_dates), self.pkl_path,
        )
        if not self._eligible_dates:
            raise ValueError(
                f"RealEpisodeIndex: 0 eligible points after holdout filter "
                f"-- check holdout_ranges={self.holdout_ranges} against "
                f"the actual date span of {self.pkl_path}"
            )

    @property
    def n_eligible_dates(self) -> int:
        return len(self._eligible_dates)

    @property
    def all_zone_ids(self) -> List[str]:
        return list(self._all_zone_ids)

    def is_holdout(self, day: date) -> bool:
        return _in_holdout(day, self.holdout_ranges)

    def sample(
        self,
        rng: random.Random,
        n_zones: int,
        config: ForecastConfig,
        inject_noise: bool = False,
        noise_scale: float = 0.05,
        max_attempts: int = 20,
    ) -> Optional[EpisodeContext]:
        for _ in range(max_attempts):
            day_s = rng.choice(self._eligible_dates)
            day_points = self._by_date[day_s]
            if len(day_points) < n_zones:
                continue
            zone_ids = rng.sample(sorted(day_points.keys()), n_zones)

            try:
                zone_obs: List[ZoneObs] = []
                zone_fc: List[ForecastResult] = []
                basin: Optional[BasinContext] = None
                for zid in zone_ids:
                    pt = day_points[zid]
                    zo = ZoneObs.from_dict(dict(pt["obs"]))
                    zf = ForecastResult.from_dict(dict(pt["forecast"]))
                    if inject_noise:
                        zo = _perturb_zone_obs(zo, noise_scale, rng)
                    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.debug(
                                "RealEpisodeIndex: basin_context skipped "
                                "for %s/%s (%s)", day_s, zid, e,
                            )

                return EpisodeContext(
                    obs=zone_obs[0],
                    forecast=zone_fc[0],
                    config=config,
                    zone_ids=list(zone_ids),
                    basin_context=basin,
                    zone_obs=zone_obs,
                    zone_forecasts=zone_fc,
                    data_source=zone_obs[0].source,
                )
            except Exception as e:
                logger.debug(
                    "RealEpisodeIndex: sample build failed for %s (%s), "
                    "retrying", day_s, e,
                )
                continue
        return None


# ---------------------------------------------------------------------------
# Self-test (python real_episode_sampler.py) -- builds its own tiny
# synthetic-shaped pkl on disk so it needs no external data file.
# ---------------------------------------------------------------------------

def _self_test() -> int:
    import sys
    import tempfile
    from datetime import timezone

    from zone_observation import make_synthetic_zone_obs, make_synthetic_forecast_result

    logging.basicConfig(level=logging.WARNING)
    print("real_episode_sampler.py self-test\n")
    failures: List[str] = []

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

    zone_ids = ["zone_a", "zone_b", "zone_c"]
    dates_in_holdout = ["2023-08-01", "2023-08-06"]
    dates_eligible = ["2019-08-01", "2019-08-06", "2019-08-11"]

    trajectories = []
    for zid in zone_ids:
        points = []
        for day_s in dates_in_holdout + dates_eligible:
            seed = hash((zid, day_s)) % 100000
            obs = make_synthetic_zone_obs(zid, seed=seed)
            d = obs.to_dict()
            d["valid_time"] = day_s + "T00:00:00+00:00"
            fc = make_synthetic_forecast_result(zid, valid_time=obs.valid_time, seed=seed)
            points.append({
                "valid_time": day_s,
                "zone_id": zid,
                "obs": d,
                "forecast": fc.to_dict(),
                "basin_context": None,
            })
        trajectories.append({"meta": {"zone_id": zid}, "trajectory": points})

    cache = {"trajectories": trajectories}

    with tempfile.TemporaryDirectory() as td:
        pkl_path = str(Path(td) / "test_cache.pkl")
        with open(pkl_path, "wb") as f:
            pickle.dump(cache, f)

        idx = RealEpisodeIndex(pkl_path)

        # 1. Holdout dates must never be indexed.
        for day_s in dates_in_holdout:
            _assert(day_s not in idx._by_date,
                    f"holdout date {day_s} was indexed (must be excluded)")
        for day_s in dates_eligible:
            _assert(day_s in idx._by_date,
                    f"eligible date {day_s} was NOT indexed")
        _assert(idx.n_eligible_dates == len(dates_eligible),
                f"expected {len(dates_eligible)} eligible dates, "
                f"got {idx.n_eligible_dates}")
        print(f"  Holdout filtering OK "
              f"({idx.n_eligible_dates} eligible / "
              f"{len(dates_in_holdout)} excluded)")

        # 2. is_holdout() agrees with what actually got indexed.
        _assert(idx.is_holdout(_parse_day("2023-08-01")),
                "is_holdout() should be True for a holdout-season date")
        _assert(not idx.is_holdout(_parse_day("2019-08-01")),
                "is_holdout() should be False for an eligible date")

        # 3. Sampling only ever returns eligible-date zones, never holdout.
        cfg = ForecastConfig(n_zones=3)
        rng = random.Random(7)
        seen_dates = set()
        for _ in range(50):
            ctx = idx.sample(rng, n_zones=3, config=cfg)
            _assert(ctx is not None, "sample() unexpectedly returned None")
            if ctx is not None:
                d = ctx.obs.valid_time.date()
                _assert(not idx.is_holdout(d),
                        f"sample() returned a holdout date: {d}")
                seen_dates.add(d.isoformat())
                _assert(len(ctx.zone_ids) == 3,
                        f"expected 3 zones, got {len(ctx.zone_ids)}")
                _assert(len(ctx.zone_obs) == 3 and len(ctx.zone_forecasts) == 3,
                        "zone_obs/zone_forecasts length mismatch")
        _assert(len(seen_dates) >= 2,
                f"sampling looks non-random across dates: only saw {seen_dates}")
        print(f"  Sampling OK (never returned a holdout date across 50 draws, "
              f"saw {len(seen_dates)}/{len(dates_eligible)} distinct eligible dates)")

        # 4. Noise injection actually perturbs values (and is reproducible
        #    given the same rng state) without touching zone_id/valid_time,
        #    AND preserves the nested precip-window ordering (the fix under
        #    test: 24h <= 7d <= 14d <= 30d must hold after perturbation,
        #    not just before it).
        rng_a = random.Random(42)
        ctx_plain = idx.sample(rng_a, n_zones=3, config=cfg, inject_noise=False)
        rng_b = random.Random(42)
        ctx_noisy = idx.sample(rng_b, n_zones=3, config=cfg,
                                inject_noise=True, noise_scale=0.1)
        _assert(ctx_plain is not None and ctx_noisy is not None,
                "noise-injection sample() returned None")
        if ctx_plain is not None and ctx_noisy is not None:
            differs = any(
                abs(a.precip_24h_mm - b.precip_24h_mm) > 1e-9
                or abs(a.rh_max_pct - b.rh_max_pct) > 1e-9
                for a, b in zip(ctx_plain.zone_obs, ctx_noisy.zone_obs)
            )
            _assert(differs, "inject_noise=True produced identical values to inject_noise=False")
            same_ids = all(
                a.zone_id == b.zone_id
                for a, b in zip(ctx_plain.zone_obs, ctx_noisy.zone_obs)
            )
            _assert(same_ids, "noise injection altered zone_id identity")
            for zo in ctx_noisy.zone_obs:
                _assert(
                    zo.precip_24h_mm <= zo.precip_7d_mm + 1e-6
                    and zo.precip_7d_mm <= zo.precip_14d_mm + 1e-6
                    and zo.precip_14d_mm <= zo.precip_30d_mm + 1e-6,
                    f"noise injection broke precip window nesting for "
                    f"{zo.zone_id}: 24h={zo.precip_24h_mm} 7d={zo.precip_7d_mm} "
                    f"14d={zo.precip_14d_mm} 30d={zo.precip_30d_mm}",
                )
        print("  Noise injection OK (perturbs continuous fields, preserves "
              "identity fields, preserves nested precip-window ordering)")

        # 5. n_zones larger than any single day's coverage falls back to None.
        ctx_toolarge = idx.sample(random.Random(1), n_zones=99, config=ForecastConfig(n_zones=99))
        _assert(ctx_toolarge is None,
                "sample() should return None when no date has enough zone coverage")
        print("  Graceful None fallback OK (n_zones exceeding available coverage)")

        # 6. Empty-after-holdout raises clearly rather than silently
        #    returning an unusable index.
        try:
            RealEpisodeIndex(pkl_path, holdout_ranges=(("2000-01-01", "2030-01-01"),))
            _assert(False, "RealEpisodeIndex should raise when holdout excludes everything")
        except ValueError:
            pass
        print("  Empty-after-holdout raises ValueError OK")

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


if __name__ == "__main__":
    import sys
    sys.exit(_self_test())