era5_data_pipeline.py

"""
era5_data_pipeline.py
=====================
"""

from __future__ import annotations

import hashlib
import json
import logging
import math
import os
import random
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"Schema mismatch: expected 3, got {_zo.SCHEMA_VERSION}"
)

from zone_observation import (
    BasinContext,
    DataSource,
    EpisodeContext,
    ForecastConfig,
    GeoPolygon,
    ZoneObs,
    derive_helio_regime,
    make_synthetic_basin_context,
    make_synthetic_episode_context,
    make_synthetic_zone_obs,
    _stable_seed,
)

logger = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# Optional dependencies
# ---------------------------------------------------------------------------
try:
    import requests
    REQUESTS_AVAILABLE = True
except ImportError:
    REQUESTS_AVAILABLE = False
    logger.warning("requests not installed — Open-Meteo unavailable")

try:
    import cdsapi
    CDSAPI_AVAILABLE = True
except ImportError:
    CDSAPI_AVAILABLE = False
    logger.info("cdsapi not installed — ERA5 will fall back to synthetic")

try:
    import numpy as np
    NUMPY_AVAILABLE = True
except ImportError:
    NUMPY_AVAILABLE = False

try:
    import ee  # Google Earth Engine — IMERG / CHIRPS / SMAP access
    EE_AVAILABLE = True
except ImportError:
    EE_AVAILABLE = False
    logger.info("earthengine-api not installed — satellite sources unavailable")

_EE_INITIALIZED = False


def _ensure_ee_initialized() -> None:
    """Lazily call ee.Initialize() once per process."""
    global _EE_INITIALIZED
    if _EE_INITIALIZED:
        return
    if not EE_AVAILABLE:
        raise RuntimeError("earthengine-api not installed")
    project = os.environ.get("EARTHENGINE_PROJECT")
    if project:
        ee.Initialize(project=project)
    else:
        ee.Initialize()
    _EE_INITIALIZED = True

# ---------------------------------------------------------------------------
# Zone registry
# ---------------------------------------------------------------------------
_ZONE_REGISTRY: Dict[str, GeoPolygon] = {}


def register_zone(polygon: GeoPolygon) -> None:
    """Register a sourcing zone polygon for lat/lon resolution."""
    _ZONE_REGISTRY[polygon.zone_id] = polygon
    logger.info(
        f"Registered zone {polygon.zone_id} centroid={polygon.centroid}"
    )


def _resolve_latlon(zone_id: str) -> Tuple[float, float]:
    if zone_id not in _ZONE_REGISTRY:
        raise KeyError(
            f"Zone '{zone_id}' not registered. "
            f"Call register_zone() before fetching data."
        )
    return _ZONE_REGISTRY[zone_id].centroid


# ---------------------------------------------------------------------------
# Config / cache
# ---------------------------------------------------------------------------
_CACHE_DIR = Path(os.environ.get("WEATHER_CACHE_DIR", ".cache/era5"))
_CACHE_DIR.mkdir(parents=True, exist_ok=True)

_ERA5_CACHE_DIR  = _CACHE_DIR / "era5_nc"
_ERA5_CACHE_DIR.mkdir(parents=True, exist_ok=True)

_OPENMETEO_URL   = "https://api.open-meteo.com/v1/forecast"
_OPENMETEO_ARCHIVE_URL = "https://archive-api.open-meteo.com/v1/archive"

_TIMEOUT_S       = int(os.environ.get("WEATHER_HTTP_TIMEOUT",    "30"))
_CACHE_TTL_DAYS  = int(os.environ.get("WEATHER_CACHE_TTL_DAYS",  "7"))
_ERA5_TTL_DAYS   = int(os.environ.get("WEATHER_ERA5_TTL_DAYS",   "30"))

# ERA5 bounding box padding in degrees around zone centroid
_ERA5_BOX_PAD    = float(os.environ.get("WEATHER_ERA5_BOX_PAD",  "0.5"))


# ---------------------------------------------------------------------------
# Cache utilities
# ---------------------------------------------------------------------------
def _cache_key(url: str, params: Dict[str, Any]) -> str:
    payload = f"{url}{json.dumps(params, sort_keys=True)}"
    return hashlib.sha256(payload.encode()).hexdigest()


def _cached_get(url: str, params: Dict[str, Any]) -> Dict[str, Any]:
    """HTTP GET with file-based JSON cache."""
    if not REQUESTS_AVAILABLE:
        raise RuntimeError("requests not installed — cannot fetch HTTP data")

    path = _CACHE_DIR / f"{_cache_key(url, params)}.json"

    if path.exists():
        age = (
            datetime.now(timezone.utc)
            - datetime.fromtimestamp(path.stat().st_mtime, tz=timezone.utc)
        ).days
        if age < _CACHE_TTL_DAYS:
            with open(path) as f:
                return json.load(f)

    resp = requests.get(url, params=params, timeout=_TIMEOUT_S)
    resp.raise_for_status()
    data = resp.json()

    with open(path, "w") as f:
        json.dump(data, f)

    return data


def _era5_cache_path(zone_id: str, date_range: Tuple[datetime, datetime]) -> Path:
    """Deterministic cache file path for an ERA5 download."""
    key = _stable_seed(
        zone_id
        + date_range[0].date().isoformat()
        + date_range[1].date().isoformat()
    )
    return _ERA5_CACHE_DIR / f"{zone_id}_{key}.nc"


# ---------------------------------------------------------------------------
# Timezone helper
# ---------------------------------------------------------------------------
def _ensure_utc(dt: datetime) -> datetime:
    """Normalise a datetime to UTC at the pipeline boundary."""
    if dt.tzinfo is None:
        return dt.replace(tzinfo=timezone.utc)
    return dt.astimezone(timezone.utc)


_ERA5_VARIABLES: List[str] = [
    "total_precipitation",
    "2m_temperature",
    "2m_dewpoint_temperature",
    "10m_u_component_of_wind",
    "10m_v_component_of_wind",
    "surface_pressure",
    "volumetric_soil_water_layer_1",
    "potential_evaporation",
]


def _dewpoint_to_rh(temp_c: float, dewpoint_c: float) -> float:
    """Magnus formula: relative humidity from temperature and dewpoint (%)."""
    a, b = 17.625, 243.04  # Magnus coefficients
    rh = 100.0 * math.exp(
        (a * dewpoint_c / (b + dewpoint_c))
        - (a * temp_c / (b + temp_c))
    )
    return max(0.0, min(100.0, rh))


def _build_era5_obs(
    zone_id:    str,
    start:      datetime,
    nc_path:    Path,
) -> ZoneObs:
    """Map a downloaded ERA5 NetCDF file to a ZoneObs."""
    try:
        import netCDF4 as nc  # type: ignore
    except ImportError:
        try:
            import xarray as xr  # type: ignore
            return _build_era5_obs_xarray(zone_id, start, nc_path)
        except ImportError:
            raise ImportError(
                "netCDF4 or xarray required for ERA5 ingestion. "
                "Install with: pip install netCDF4  or  pip install xarray"
            )

    import numpy as np

    ds = nc.Dataset(str(nc_path), "r")
    try:
        def _daily_mean(var: str) -> float:
            if var not in ds.variables:
                return 0.0
            data = ds.variables[var][:24].flatten()
            data = np.ma.filled(data, np.nan)
            valid = data[np.isfinite(data)]
            return float(np.mean(valid)) if len(valid) > 0 else 0.0

        def _daily_max(var: str) -> float:
            if var not in ds.variables:
                return 0.0
            data = ds.variables[var][:24].flatten()
            data = np.ma.filled(data, np.nan)
            valid = data[np.isfinite(data)]
            return float(np.max(valid)) if len(valid) > 0 else 0.0

        def _daily_sum(var: str, scale: float = 1.0) -> float:
            if var not in ds.variables:
                return 0.0
            data = ds.variables[var][:24].flatten()
            data = np.ma.filled(data, np.nan)
            valid = data[np.isfinite(data)]
            return float(np.sum(valid) * scale) if len(valid) > 0 else 0.0

        temp_mean_c = _daily_mean("t2m") - 273.15
        temp_max_c  = _daily_max("t2m")  - 273.15
        if "t2m" in ds.variables:
            t_data = np.ma.filled(ds.variables["t2m"][:24].flatten(), np.nan)
            valid  = t_data[np.isfinite(t_data)]
            temp_min_c = float(np.min(valid)) - 273.15 if len(valid) > 0 else temp_mean_c
        else:
            temp_min_c = temp_mean_c

        dewpoint_mean_c = _daily_mean("d2m") - 273.15
        rh_mean = _dewpoint_to_rh(temp_mean_c, dewpoint_mean_c)
        dewpoint_max_c  = _daily_max("d2m") - 273.15
        rh_max  = _dewpoint_to_rh(temp_min_c, dewpoint_max_c)

        precip_24h = _daily_sum("tp", scale=1000.0)

        total_hours = len(ds.variables.get("tp", [])) if "tp" in ds.variables else 24
        if total_hours >= 24 * 7:
            def _window_sum(var: str, hours: int, scale: float = 1.0) -> float:
                if var not in ds.variables:
                    return 0.0
                data = np.ma.filled(
                    ds.variables[var][:hours].flatten(), np.nan
                )
                valid = data[np.isfinite(data)]
                return float(np.sum(valid) * scale) if len(valid) > 0 else 0.0
            precip_7d  = _window_sum("tp", 24*7,  scale=1000.0)
            precip_14d = _window_sum("tp", 24*14, scale=1000.0)
            precip_30d = _window_sum("tp", 24*30, scale=1000.0)
        else:
            precip_7d  = precip_24h * 7.0
            precip_14d = precip_24h * 14.0
            precip_30d = precip_24h * 30.0

        u = _daily_mean("u10")
        v = _daily_mean("v10")
        wind_mean = math.sqrt(u**2 + v**2)
        u_max = _daily_max("u10")
        v_max = _daily_max("v10")
        wind_max = math.sqrt(u_max**2 + v_max**2)

        soil_pct = _daily_mean("swvl1") * 100.0

        pe_m = _daily_sum("pev", scale=1.0)
        et0_mm = abs(pe_m) * 1000.0

        return ZoneObs(
            zone_id=zone_id,
            valid_time=start,
            source=DataSource.ERA5_REANALYSIS,
            precip_24h_mm=max(0.0, precip_24h),
            precip_7d_mm=max(0.0, precip_7d),
            precip_14d_mm=max(0.0, precip_14d),
            precip_30d_mm=max(0.0, precip_30d),
            temp_mean_c=temp_mean_c,
            temp_max_c=max(temp_mean_c, temp_max_c),
            temp_min_c=min(temp_mean_c, temp_min_c),
            temp_anomaly_idx=0.0,
            precip_anomaly_idx=0.0,
            evapotranspiration_mm=max(0.0, et0_mm),
            wind_speed_mean_ms=max(0.0, wind_mean),
            wind_speed_max_ms=max(0.0, wind_max),
            rh_mean_pct=rh_mean,
            rh_max_pct=max(rh_mean, rh_max),
            soil_moisture_pct=max(0.0, min(100.0, soil_pct)),
            soil_moisture_anom=0.0,
            quality_flag=0,
        )
    finally:
        ds.close()


def _build_era5_obs_xarray(
    zone_id: str,
    start:   datetime,
    nc_path: Path,
) -> ZoneObs:
    import xarray as xr
    import numpy as np

    ds = xr.open_dataset(str(nc_path))
    try:
        def _mean(var: str) -> float:
            if var not in ds:
                return 0.0
            return float(ds[var].values.flatten()[np.isfinite(
                ds[var].values.flatten()
            )].mean()) if len(ds[var].values.flatten()) > 0 else 0.0

        def _max(var: str) -> float:
            if var not in ds:
                return 0.0
            vals = ds[var].values.flatten()
            valid = vals[np.isfinite(vals)]
            return float(valid.max()) if len(valid) > 0 else 0.0

        def _sum(var: str, scale: float = 1.0) -> float:
            if var not in ds:
                return 0.0
            vals = ds[var].values.flatten()
            valid = vals[np.isfinite(vals)]
            return float(valid.sum() * scale) if len(valid) > 0 else 0.0

        temp_mean_c    = _mean("t2m") - 273.15
        temp_max_c     = _max("t2m")  - 273.15
        if "t2m" in ds:
            _t2m_flat = ds["t2m"].values.flatten()
            _t2m_valid = _t2m_flat[np.isfinite(_t2m_flat)]
            temp_min_c = float(np.min(_t2m_valid)) - 273.15 if len(_t2m_valid) > 0 else temp_mean_c
        else:
            temp_min_c = temp_mean_c

        dewpoint_mean_c = _mean("d2m") - 273.15
        dewpoint_max_c  = _max("d2m")  - 273.15
        rh_mean = _dewpoint_to_rh(temp_mean_c, dewpoint_mean_c)
        rh_max  = _dewpoint_to_rh(temp_min_c, dewpoint_max_c)

        precip_24h = max(0.0, _sum("tp", scale=1000.0))

        if "tp" in ds:
            total_hours = len(ds["tp"].values.flatten())
        else:
            total_hours = 0

        if total_hours >= 24 * 7:
            def _window_sum_xr(var: str, hours: int, scale: float = 1.0) -> float:
                if var not in ds:
                    return 0.0
                vals = ds[var].values.flatten()[:hours]
                valid = vals[np.isfinite(vals)]
                return float(valid.sum() * scale) if len(valid) > 0 else 0.0
            precip_7d  = max(0.0, _window_sum_xr("tp", 24 * 7,  scale=1000.0))
            precip_14d = max(0.0, _window_sum_xr("tp", 24 * 14, scale=1000.0))
            precip_30d = max(0.0, _window_sum_xr("tp", 24 * 30, scale=1000.0))
        else:
            precip_7d  = precip_24h * 7.0
            precip_14d = precip_24h * 14.0
            precip_30d = precip_24h * 30.0
        u = _mean("u10")
        v = _mean("v10")
        wind_mean = math.sqrt(u**2 + v**2)
        wind_max  = math.sqrt(_max("u10")**2 + _max("v10")**2)
        soil_pct  = _mean("swvl1") * 100.0
        et0_mm    = abs(_sum("pev")) * 1000.0

        return ZoneObs(
            zone_id=zone_id,
            valid_time=start,
            source=DataSource.ERA5_REANALYSIS,
            precip_24h_mm=precip_24h,
            precip_7d_mm=precip_7d,
            precip_14d_mm=precip_14d,
            precip_30d_mm=precip_30d,
            temp_mean_c=temp_mean_c,
            temp_max_c=max(temp_mean_c, temp_max_c),
            temp_min_c=min(temp_mean_c, temp_min_c),
            evapotranspiration_mm=max(0.0, et0_mm),
            wind_speed_mean_ms=max(0.0, wind_mean),
            wind_speed_max_ms=max(0.0, wind_max),
            rh_mean_pct=rh_mean,
            rh_max_pct=max(rh_mean, rh_max),
            soil_moisture_pct=max(0.0, min(100.0, soil_pct)),
            quality_flag=0,
        )
    finally:
        ds.close()


def _fetch_openmeteo(zone_id: str, date_range: Tuple[datetime, datetime]) -> ZoneObs:
    """Fetch from Open-Meteo API with full variable coverage."""
    lat, lon = _resolve_latlon(zone_id)
    start    = _ensure_utc(date_range[0])
    end      = _ensure_utc(date_range[1])

    today = datetime.now(timezone.utc).date()
    use_archive = start.date() < today - timedelta(days=5)
    url = _OPENMETEO_ARCHIVE_URL if use_archive else _OPENMETEO_URL

    daily_vars = ",".join([
        "temperature_2m_mean",
        "temperature_2m_max",
        "temperature_2m_min",
        "precipitation_sum",
        "precipitation_hours",
        "wind_speed_10m_mean",
        "wind_speed_10m_max",
        "relative_humidity_2m_mean",
        "relative_humidity_2m_max",
        "et0_fao_evapotranspiration",
        "shortwave_radiation_sum",
        "soil_moisture_0_to_7cm_mean",
    ])

    params = {
        "latitude":   lat,
        "longitude":  lon,
        "daily":      daily_vars,
        "start_date": start.date().isoformat(),
        "end_date":   end.date().isoformat(),
        "timezone":   "UTC",
    }

    data  = _cached_get(url, params)
    daily = data.get("daily", {})

    def _safe(key: str, i: int = 0) -> float:
        vals = daily.get(key, [])
        return float(vals[i]) if i < len(vals) and vals[i] is not None else 0.0

    def _safe_sum(key: str, n: int) -> float:
        vals = daily.get(key, [])
        return float(sum(
            v for v in vals[:n] if v is not None
        ))

    temp_mean = _safe("temperature_2m_mean")
    temp_max  = _safe("temperature_2m_max")
    temp_min  = _safe("temperature_2m_min")

    temp_max  = max(temp_mean, temp_max)
    temp_min  = min(temp_mean, temp_min)

    rh_mean = _safe("relative_humidity_2m_mean")
    rh_max  = _safe("relative_humidity_2m_max")
    rh_max  = max(rh_mean, rh_max)

    wind_mean = _safe("wind_speed_10m_mean")
    wind_max  = _safe("wind_speed_10m_max")
    wind_max  = max(wind_mean, wind_max)

    return ZoneObs(
        zone_id=zone_id,
        valid_time=start,
        source=DataSource.OPENMETEO_LIVE,
        precip_24h_mm=_safe("precipitation_sum"),
        precip_7d_mm=_safe_sum("precipitation_sum", 7),
        precip_14d_mm=_safe_sum("precipitation_sum", 14),
        precip_30d_mm=_safe_sum("precipitation_sum", 30),
        temp_mean_c=temp_mean,
        temp_max_c=temp_max,
        temp_min_c=temp_min,
        evapotranspiration_mm=_safe("et0_fao_evapotranspiration"),
        wind_speed_mean_ms=wind_mean / 3.6,
        wind_speed_max_ms=wind_max  / 3.6,
        rh_mean_pct=rh_mean,
        rh_max_pct=rh_max,
        soil_moisture_pct=max(0.0, min(100.0,
                               _safe("soil_moisture_0_to_7cm_mean") * 100.0)),
        quality_flag=0,
    )


def _fetch_era5(zone_id: str, date_range: Tuple[datetime, datetime]) -> ZoneObs:
    if not CDSAPI_AVAILABLE:
        logger.debug("cdsapi unavailable — falling back to synthetic for %s", zone_id)
        return _fetch_synthetic(zone_id, date_range)

    lat, lon = _resolve_latlon(zone_id)
    start    = _ensure_utc(date_range[0])
    end      = _ensure_utc(date_range[1])

    bbox = [
        round(lat + _ERA5_BOX_PAD, 2),
        round(lon - _ERA5_BOX_PAD, 2),
        round(lat - _ERA5_BOX_PAD, 2),
        round(lon + _ERA5_BOX_PAD, 2),
    ]

    nc_path = _era5_cache_path(zone_id, date_range)

    if nc_path.exists():
        age_days = (
            datetime.now(timezone.utc)
            - datetime.fromtimestamp(nc_path.stat().st_mtime, tz=timezone.utc)
        ).days
        if age_days < _ERA5_TTL_DAYS:
            logger.debug("ERA5 cache hit for %s", zone_id)
        else:
            nc_path.unlink()

    if not nc_path.exists():
        years:  set = set()
        months: set = set()
        days:   set = set()
        cursor = start.date()
        end_date = end.date()
        while cursor <= end_date:
            years.add(cursor.year)
            months.add(cursor.month)
            days.add(cursor.day)
            cursor += timedelta(days=1)

        logger.info(
            "Requesting ERA5 data for %s  bbox=%s  dates=%s to %s",
            zone_id, bbox,
            start.date().isoformat(),
            end.date().isoformat(),
        )
        try:
            client = cdsapi.Client(quiet=True)
            client.retrieve(
                "reanalysis-era5-single-levels",
                {
                    "product_type": "reanalysis",
                    "variable":     _ERA5_VARIABLES,
                    "year":         sorted(str(y) for y in years),
                    "month":        sorted(f"{m:02d}" for m in months),
                    "day":          sorted(f"{d:02d}" for d in days),
                    "time":         [f"{h:02d}:00" for h in range(24)],
                    "area":         bbox,
                    "format":       "netcdf",
                },
                str(nc_path),
            )
            logger.info("ERA5 download complete: %s", nc_path)
        except Exception as e:
            logger.warning(
                "ERA5 CDS request failed for %s: %s — falling back to synthetic",
                zone_id, e
            )
            if nc_path.exists():
                nc_path.unlink()
            return _fetch_synthetic(zone_id, date_range)

    try:
        return _build_era5_obs(zone_id, start, nc_path)
    except Exception as e:
        logger.warning(
            "ERA5 NetCDF parse failed for %s: %s — falling back to synthetic",
            zone_id, e
        )
        return _fetch_synthetic(zone_id, date_range)


def _fetch_imerg(zone_id: str, date_range: Tuple[datetime, datetime]) -> ZoneObs:
    _ensure_ee_initialized()

    lat, lon = _resolve_latlon(zone_id)
    start = _ensure_utc(date_range[0])

    region = ee.Geometry.Point([lon, lat]).buffer(_ERA5_BOX_PAD * 111_000)

    def _window_sum_mm(n_days: int) -> float:
        window_end = start + timedelta(days=n_days)
        coll = (
            ee.ImageCollection("NASA/GPM_L3/IMERG_V07")
            .filterDate(start.isoformat(), window_end.isoformat())
            .filterBounds(region)
            .select("precipitation")
        )
        total_mm_image = coll.sum().multiply(0.5)
        stats = total_mm_image.reduceRegion(
            reducer=ee.Reducer.mean(), geometry=region, scale=11_000, bestEffort=True,
        ).getInfo()
        return float(stats.get("precipitation", 0.0) or 0.0)

    precip_24h = max(0.0, _window_sum_mm(1))
    precip_7d  = max(precip_24h, _window_sum_mm(7))
    precip_14d = max(precip_7d,  _window_sum_mm(14))
    precip_30d = max(precip_14d, _window_sum_mm(30))

    return ZoneObs(
        zone_id=zone_id,
        valid_time=start,
        source=DataSource.SATELLITE_PRECIP,
        precip_24h_mm=precip_24h,
        precip_7d_mm=precip_7d,
        precip_14d_mm=precip_14d,
        precip_30d_mm=precip_30d,
        precip_satellite_mm=precip_24h,
        quality_flag=0,
    )


def _fetch_smap(zone_id: str, date_range: Tuple[datetime, datetime]) -> ZoneObs:
    _ensure_ee_initialized()

    lat, lon = _resolve_latlon(zone_id)
    start = _ensure_utc(date_range[0])
    end = start + timedelta(days=1)

    region = ee.Geometry.Point([lon, lat]).buffer(_ERA5_BOX_PAD * 111_000)

    coll = (
        ee.ImageCollection("NASA/SMAP/SPL4SMGP/007")
        .filterDate(start.isoformat(), end.isoformat())
        .filterBounds(region)
        .select("sm_surface")
    )
    stats = coll.mean().reduceRegion(
        reducer=ee.Reducer.mean(), geometry=region, scale=9_000, bestEffort=True,
    ).getInfo()

    vwc = stats.get("sm_surface")
    if vwc is None:
        raise RuntimeError(f"SMAP: no data returned for zone={zone_id} date={start.date()}")
    soil_pct = max(0.0, min(100.0, float(vwc) * 100.0))

    return ZoneObs(
        zone_id=zone_id,
        valid_time=start,
        source=DataSource.SATELLITE_SOIL,
        soil_moisture_pct=soil_pct,
        soil_moisture_satellite_pct=soil_pct,
        quality_flag=0,
    )


def _fetch_synthetic(zone_id: str, date_range: Tuple[datetime, datetime]) -> ZoneObs:
    seed = _stable_seed(zone_id + date_range[0].isoformat())
    return make_synthetic_zone_obs(zone_id=zone_id, seed=seed)


def _inject_noise(obs: ZoneObs, rng: random.Random, scale: float) -> ZoneObs:
    def perturb(x: float) -> float:
        return x * (1.0 + rng.uniform(-scale, scale))

    base = obs.to_dict()
    base.pop("_schema_version", None)

    base["precip_24h_mm"]      = max(0.0, perturb(obs.precip_24h_mm))
    base["precip_7d_mm"]       = max(0.0, perturb(obs.precip_7d_mm))
    base["precip_14d_mm"]      = max(0.0, perturb(obs.precip_14d_mm))
    base["precip_30d_mm"]      = max(0.0, perturb(obs.precip_30d_mm))
    base["temp_mean_c"]        = perturb(obs.temp_mean_c)
    base["temp_max_c"]         = perturb(obs.temp_max_c)
    base["temp_min_c"]         = perturb(obs.temp_min_c)
    base["rh_mean_pct"]        = min(100.0, max(0.0, perturb(obs.rh_mean_pct)))
    base["rh_max_pct"]         = min(100.0, max(0.0, perturb(obs.rh_max_pct)))
    base["wind_speed_mean_ms"] = max(0.0, perturb(obs.wind_speed_mean_ms))
    base["wind_speed_max_ms"]  = max(0.0, perturb(obs.wind_speed_max_ms))

    p24  = base["precip_24h_mm"]
    p7   = max(base["precip_7d_mm"],  p24)
    p14  = max(base["precip_14d_mm"], p7)
    p30  = max(base["precip_30d_mm"], p14)
    base["precip_7d_mm"]  = p7
    base["precip_14d_mm"] = p14
    base["precip_30d_mm"] = p30

    t_mean = base["temp_mean_c"]
    base["temp_max_c"] = max(t_mean, base["temp_max_c"])
    base["temp_min_c"] = min(t_mean, base["temp_min_c"])

    base["rh_max_pct"] = max(base["rh_mean_pct"], base["rh_max_pct"])

    base["wind_speed_max_ms"] = max(
        base["wind_speed_mean_ms"], base["wind_speed_max_ms"]
    )

    return ZoneObs.from_dict(base)


_NOAA_RONI_URL = "https://www.cpc.ncep.noaa.gov/data/indices/RONI.ascii.txt"

_PSL_DMI_URLS: Tuple[str, ...] = (
    "https://psl.noaa.gov/gcos_wgsp/Timeseries/Data/dmi.had.long.data",
    "https://psl.noaa.gov/data/timeseries/month/data/dmi.had.long.data",
    "https://psl.noaa.gov/data/correlation/dmi.data",
)

_SWPC_PLASMA_URLS: Tuple[str, ...] = (
    "https://services.swpc.noaa.gov/json/rtsw/rtsw_wind_1m.json",
    "https://services.swpc.noaa.gov/products/summary/solar-wind-speed.json",
    "https://services.swpc.noaa.gov/products/solar-wind/plasma-1-day.json",
)
_SWPC_KP_URL   = "https://services.swpc.noaa.gov/json/planetary_k_index_1m.json"
_SWPC_XRAY_URL = "https://services.swpc.noaa.gov/json/goes/primary/xrays-1-day.json"

_SEASON_FOR_MONTH: Dict[int, str] = {
    1: "DJF", 2: "JFM", 3: "FMA", 4: "MAM", 5: "AMJ", 6: "MJJ",
    7: "JJA", 8: "JAS", 9: "ASO", 10: "SON", 11: "OND", 12: "NDJ",
}
_SEASON_ORDER: Tuple[str, ...] = (
    "DJF", "JFM", "FMA", "MAM", "AMJ", "MJJ",
    "JJA", "JAS", "ASO", "SON", "OND", "NDJ",
)


def _parse_cpc_seasonal_ascii(text: str) -> Dict[Tuple[int, str], float]:
    """Parse NOAA CPC's seasonal index ASCII format."""
    out: Dict[Tuple[int, str], float] = {}
    lines = text.strip().splitlines()
    for line in lines[1:]:
        parts = line.split()
        if len(parts) < 3:
            continue
        season = parts[0]
        try:
            year = int(parts[1])
            value = float(parts[-1])
        except (ValueError, IndexError):
            continue
        out[(year, season)] = value
    return out


def _parse_psl_monthly_ascii(
    text: str, year: int, month: int, missing_below: float = -90.0
) -> Optional[float]:
    for line in text.strip().splitlines():
        parts = line.split()
        if len(parts) != 13:
            continue
        try:
            row_year = int(parts[0])
            values = [float(v) for v in parts[1:]]
        except ValueError:
            continue
        if row_year == year:
            v = values[month - 1]
            return None if v <= missing_below else v
    return None


def _lookup_roni_with_lag(
    table: Dict[Tuple[int, str], float], year: int, season: str
) -> Tuple[float, str]:
    if (year, season) in table:
        return table[(year, season)], f"{season} {year}"

    try:
        idx = _SEASON_ORDER.index(season)
    except ValueError:
        idx = 0

    y, i = year, idx
    for _ in range(24):
        i -= 1
        if i < 0:
            i = len(_SEASON_ORDER) - 1
            y -= 1
        key = (y, _SEASON_ORDER[i])
        if key in table:
            return table[key], f"{_SEASON_ORDER[i]} {y} (lagged from {season} {year})"
    raise ValueError(f"No RONI value for {season} {year} or any prior season in table")


def _lookup_dmi_with_lag(
    text: str, year: int, month: int
) -> Tuple[float, str]:
    y, m = year, month
    for _ in range(24):
        v = _parse_psl_monthly_ascii(text, y, m)
        if v is not None:
            label = f"{y}-{m:02d}"
            if (y, m) != (year, month):
                label += f" (lagged from {year}-{month:02d})"
            return v, label
        m -= 1
        if m < 1:
            m = 12
            y -= 1
    raise ValueError(f"No DMI value for {year}-{month:02d} or any prior month")


def _extract_solar_wind_speeds(payload: Any) -> List[float]:
    speeds: List[float] = []

    def _from_row(row: Any) -> None:
        if not isinstance(row, dict):
            return
        for key in ("proton_speed", "speed", "wind_speed", "value"):
            sp = row.get(key)
            if sp is None:
                continue
            try:
                f = float(sp)
            except (TypeError, ValueError):
                continue
            if f > 0:
                speeds.append(f)
                return

    if isinstance(payload, list):
        for row in payload:
            if isinstance(row, dict):
                _from_row(row)
            elif isinstance(row, (list, tuple)) and len(row) >= 2:
                try:
                    f = float(row[1])
                    if f > 0:
                        speeds.append(f)
                except (TypeError, ValueError):
                    continue
    elif isinstance(payload, dict):
        _from_row(payload)
        for v in payload.values():
            if isinstance(v, list):
                speeds.extend(_extract_solar_wind_speeds(v))

    return speeds


def _fetch_swpc_helio(valid_date: datetime) -> Dict[str, Any]:
    out: Dict[str, Any] = {}
    if not REQUESTS_AVAILABLE:
        return out

    plasma_err: Optional[Exception] = None
    for url in _SWPC_PLASMA_URLS:
        try:
            resp = requests.get(url, timeout=_TIMEOUT_S)
            resp.raise_for_status()
            speeds = _extract_solar_wind_speeds(resp.json())
            if speeds:
                out["solar_wind_speed_kms"] = speeds[-1]
                plasma_err = None
                break
            plasma_err = ValueError(f"no speed samples in {url}")
        except Exception as e:
            plasma_err = e
            continue
    if "solar_wind_speed_kms" not in out and plasma_err is not None:
        logger.warning(
            "_fetch_swpc_helio: plasma fetch failed (%s) for %s",
            plasma_err, valid_date.date(),
        )

    try:
        resp = requests.get(_SWPC_KP_URL, timeout=_TIMEOUT_S)
        resp.raise_for_status()
        rows = resp.json()
        kps = []
        for row in rows:
            try:
                kp = row.get("kp_index", row.get("kp"))
                if kp is not None:
                    kps.append(float(kp))
            except (TypeError, ValueError, AttributeError):
                continue
        if kps:
            out["kp_index"] = kps[-1]
    except Exception as e:
        logger.warning(
            "_fetch_swpc_helio: Kp fetch failed (%s) for %s",
            e, valid_date.date(),
        )

    try:
        resp = requests.get(_SWPC_XRAY_URL, timeout=_TIMEOUT_S)
        resp.raise_for_status()
        rows = resp.json()
        fluxes = []
        for row in rows:
            try:
                flux = row.get("flux")
                energy = str(row.get("energy", "")).lower()
                if flux is not None and float(flux) > 0:
                    if "0.1-0.8" in energy or "long" in energy or not energy:
                        fluxes.append(float(flux))
            except (TypeError, ValueError, AttributeError):
                continue
        if fluxes:
            out["goes_xray_flux"] = fluxes[-1]
    except Exception as e:
        logger.warning(
            "_fetch_swpc_helio: GOES X-ray fetch failed (%s) for %s",
            e, valid_date.date(),
        )

    if out:
        kp = float(out.get("kp_index", 2.0))
        xray = float(out.get("goes_xray_flux", 1e-7))
        out["helio_regime"] = derive_helio_regime(kp, xray)

    return out


def fetch_basin_context(
    valid_date: datetime,
    config: Optional[ForecastConfig] = None,
) -> BasinContext:
    cfg = config or ForecastConfig()
    valid_date = _ensure_utc(valid_date)
    season = _SEASON_FOR_MONTH[valid_date.month]
    strict = bool(getattr(cfg, "require_real_basin_context", False))

    if not REQUESTS_AVAILABLE:
        if strict:
            raise RuntimeError(
                "require_real_basin_context=True but the 'requests' package "
                "is not installed — cannot fetch live basin/helio indices."
            )
        logger.info("requests not installed — basin context falls back to synthetic")
        return make_synthetic_basin_context(valid_date=valid_date)

    try:
        resp = requests.get(_NOAA_RONI_URL, timeout=_TIMEOUT_S)
        resp.raise_for_status()
        table = _parse_cpc_seasonal_ascii(resp.text)
        enso_oni, roni_label = _lookup_roni_with_lag(table, valid_date.year, season)
        if "lagged" in roni_label:
            logger.info(
                "fetch_basin_context: RONI using %s for %s",
                roni_label, valid_date.date(),
            )
    except Exception as e:
        if strict:
            raise RuntimeError(
                f"require_real_basin_context=True: RONI fetch failed for "
                f"{valid_date.date()}: {e}"
            ) from e
        logger.warning(
            "fetch_basin_context: RONI fetch/parse failed (%s) — synthetic ENSO for %s",
            e, valid_date.date(),
        )
        enso_oni = make_synthetic_basin_context(
            valid_date=valid_date, seed=_stable_seed(f"enso_{valid_date.date().isoformat()}")
        ).enso_oni

    iod_dmi: Optional[float] = None
    dmi_err: Optional[Exception] = None
    for dmi_url in _PSL_DMI_URLS:
        try:
            resp = requests.get(dmi_url, timeout=_TIMEOUT_S)
            resp.raise_for_status()
            iod_dmi, dmi_label = _lookup_dmi_with_lag(
                resp.text, valid_date.year, valid_date.month
            )
            if "lagged" in dmi_label:
                logger.info(
                    "fetch_basin_context: DMI using %s for %s (from %s)",
                    dmi_label, valid_date.date(), dmi_url,
                )
            dmi_err = None
            break
        except Exception as e:
            dmi_err = e
            continue
    if iod_dmi is None:
        if strict:
            raise RuntimeError(
                f"require_real_basin_context=True: DMI fetch failed for "
                f"{valid_date.date()}: {dmi_err}"
            )
        logger.warning(
            "fetch_basin_context: DMI fetch/parse failed (%s) — synthetic IOD for %s",
            dmi_err, valid_date.date(),
        )
        iod_dmi = make_synthetic_basin_context(
            valid_date=valid_date, seed=_stable_seed(f"iod_{valid_date.date().isoformat()}")
        ).iod_dmi

    helio = _fetch_swpc_helio(valid_date)
    required_helio = ("solar_wind_speed_kms", "kp_index", "goes_xray_flux")
    missing_helio = [k for k in required_helio if k not in helio]

    if strict and missing_helio:
        raise RuntimeError(
            f"require_real_basin_context=True: SWPC helio incomplete for "
            f"{valid_date.date()} — missing {missing_helio}. "
            f"Got keys: {sorted(helio.keys())}"
        )

    return BasinContext(
        valid_date=valid_date,
        enso_oni=enso_oni,
        iod_dmi=iod_dmi,
        solar_wind_speed_kms=float(helio.get("solar_wind_speed_kms", 400.0)),
        kp_index=float(helio.get("kp_index", 2.0)),
        goes_xray_flux=float(helio.get("goes_xray_flux", 1e-7)),
        helio_regime=str(helio.get("helio_regime", "quiet")),
        source=DataSource.PUBLISHED_INDEX,
    )


def _select_source(cfg: ForecastConfig, rng: random.Random) -> DataSource:
    real_ratio = getattr(cfg, "real_data_ratio", 0.7)
    era5_ratio = getattr(cfg, "era5_ratio",      0.5)

    if rng.random() > real_ratio:
        return DataSource.SYNTHETIC

    satellite_options = []
    if getattr(cfg, "use_satellite_precip", False):
        satellite_options.append(DataSource.SATELLITE_PRECIP)
    if getattr(cfg, "use_satellite_soil", False):
        satellite_options.append(DataSource.SATELLITE_SOIL)
    if satellite_options:
        return satellite_options[rng.randrange(len(satellite_options))]

    return (
        DataSource.ERA5_REANALYSIS
        if rng.random() < era5_ratio
        else DataSource.OPENMETEO_LIVE
    )


def _fallback_chain(primary: DataSource) -> List[DataSource]:
    if primary in (DataSource.SATELLITE_PRECIP, DataSource.SATELLITE_SOIL):
        return [primary, DataSource.ERA5_REANALYSIS, DataSource.SYNTHETIC]
    return [primary, DataSource.SYNTHETIC]


def fetch_zone_obs(
    zone_id:    str,
    date_range: Tuple[datetime, datetime],
    config:     Optional[ForecastConfig] = None,
) -> ZoneObs:
    cfg  = config or ForecastConfig()
    seed = _stable_seed(zone_id + date_range[0].isoformat())
    rng  = random.Random(seed)

    if getattr(cfg, "force_data_source", None) is not None:
        primary = cfg.force_data_source
    else:
        primary = _select_source(cfg, rng)

    fetchers = {
        DataSource.OPENMETEO_LIVE:   _fetch_openmeteo,
        DataSource.ERA5_REANALYSIS:  _fetch_era5,
        DataSource.SATELLITE_PRECIP: _fetch_imerg,
        DataSource.SATELLITE_SOIL:   _fetch_smap,
        DataSource.SYNTHETIC:        _fetch_synthetic,
    }

    obs: Optional[ZoneObs] = None
    for source in _fallback_chain(primary):
        try:
            obs = fetchers[source](zone_id, date_range)
            ZoneObs.validate(obs, strict=True)
            break
        except Exception as e:
            logger.warning("%s failed for %s: %s", source.value, zone_id, e)
            obs = None

    if obs is None:
        raise RuntimeError(
            f"All data sources failed for zone '{zone_id}'. "
            f"Check zone registration and network connectivity."
        )

    if getattr(cfg, "inject_noise", False):
        obs = _inject_noise(obs, rng, cfg.noise_scale)
        ZoneObs.validate(obs, strict=True)

    if getattr(cfg, "use_climatology_anomalies", False):
        try:
            _lat, _lon = _resolve_latlon(zone_id)
        except KeyError:
            logger.warning(
                "use_climatology_anomalies=True but zone '%s' is not "
                "registered -- cannot locate the zone for its climatology. "
                "Anomaly fields left at fetcher defaults. Call "
                "register_zone() first.", zone_id,
            )
        else:
            from climatology import (
                apply_climatology_anomalies,
                get_zone_climatology,
            )
            _clim = get_zone_climatology(
                zone_id, _lat, _lon,
                years=getattr(cfg, "climatology_years", 10),
                end_year=date_range[0].year - 1,
            )
            obs = apply_climatology_anomalies(obs, _clim)
            ZoneObs.validate(obs, strict=True)

    return obs


def _build_context_forecast(
    zone_id: str,
    obs: ZoneObs,
    cfg: ForecastConfig,
    seed: int,
) -> "ForecastResult":
    from zone_observation import ForecastResult  # noqa: F401  (type hint only)

    mode = getattr(cfg, "forecast_backend", "synthetic")

    def _synthetic() -> "ForecastResult":
        from zone_observation import make_synthetic_forecast_result
        return make_synthetic_forecast_result(
            zone_id=zone_id,
            valid_time=obs.valid_time,
            horizon_days=cfg.horizon_days,
            seed=seed,
        )

    if mode == "synthetic":
        return _synthetic()

    if mode == "timesfm":
        logger.warning(
            "forecast_backend='timesfm' cannot be constructed inside "
            "fetch_episode_context (no checkpoint/sha available at this "
            "layer) -- falling back to synthetic for zone=%s. Build the "
            "backend via timesfm_wrapper.create_forecast_backend() and "
            "call backend.forecast(obs) directly if you need TimesFM.",
            zone_id,
        )
        return _synthetic()

    try:
        from timesfm_wrapper import create_forecast_backend
        lat = lon = None
        if mode == "openmeteo":
            lat, lon = _resolve_latlon(zone_id)
        backend = create_forecast_backend(
            mode=mode, lat=lat, lon=lon, horizon_days=cfg.horizon_days,
        )
        return backend.forecast(obs)
    except Exception as e:
        logger.warning(
            "forecast backend %r failed for zone=%s (%s) -- synthetic "
            "forecast fallback", mode, zone_id, e,
        )
        return _synthetic()


def fetch_episode_context(
    zone_id:    str,
    date_range: Tuple[datetime, datetime],
    config:     Optional[ForecastConfig] = None,
) -> EpisodeContext:
    cfg = config or ForecastConfig()
    obs = fetch_zone_obs(zone_id, date_range, cfg)

    seed  = _stable_seed(zone_id + date_range[0].isoformat() + "_forecast")
    fcast = _build_context_forecast(zone_id, obs, cfg, seed)

    basin_context = None
    if getattr(cfg, "include_basin_context", False):
        strict_basin = bool(getattr(cfg, "require_real_basin_context", False))
        try:
            basin_context = fetch_basin_context(date_range[0], cfg)
        except Exception as e:
            if strict_basin:
                raise
            logger.warning(
                "fetch_episode_context: basin context fetch failed (%s) — "
                "leaving basin_context=None for zone=%s", e, zone_id,
            )
            basin_context = None

    return EpisodeContext(
        obs=obs,
        forecast=fcast,
        config=cfg,
        zone_ids=[zone_id],
        data_source=obs.source,
        basin_context=basin_context,
    )


try:
    import numpy as _np
    _NUMPY_FOR_DYNAMICS = True
except ImportError:
    _NUMPY_FOR_DYNAMICS = False

try:
    import torch as _torch
    _TORCH_FOR_DYNAMICS = True
except ImportError:
    _TORCH_FOR_DYNAMICS = False


def _forecast_result_to_arrays(
    forecast_result: Any,
    horizon_days: int,
) -> Tuple["np.ndarray", float]:
    import numpy as np

    precip_seq = list(forecast_result.precip_mm) if forecast_result.precip_mm else []
    if len(precip_seq) < horizon_days:
        precip_seq = precip_seq + [0.0] * (horizon_days - len(precip_seq))
    precip_arr = np.clip(
        np.array(precip_seq[:horizon_days], dtype=np.float32),
        0.0, 500.0,
    )

    uncertainty = 0.5
    if forecast_result.precip_p90 and forecast_result.precip_p10:
        p90 = np.array(list(forecast_result.precip_p90)[:horizon_days], dtype=np.float32)
        p10 = np.array(list(forecast_result.precip_p10)[:horizon_days], dtype=np.float32)
        spread = np.clip(
            (p90 - p10) / np.maximum(np.abs(p90), 1e-6),
            0.0, 1.0,
        )
        uncertainty = float(np.mean(spread))

    return precip_arr, uncertainty


def _episode_context_to_state_tensor(
    contexts: List[Any],
    horizon_days: int,
    prior_belief: float = 0.5,
) -> "ZoneStateTensor":

    import numpy as np
    from physics_dynamics import ZoneStateTensor

    n_zones = len(contexts)
    precip_arr   = np.zeros((1, n_zones, horizon_days), dtype=np.float32)
    uncert_arr   = np.zeros((1, n_zones), dtype=np.float32)
    belief_arr   = np.zeros((1, n_zones), dtype=np.float32)

    for zi, ctx in enumerate(contexts):
        p_arr, unc = _forecast_result_to_arrays(ctx.forecast, horizon_days)
        precip_arr[0, zi, :] = p_arr
        uncert_arr[0, zi]    = unc

        try:
            signal = float(ctx.obs.composite_risk())
            belief = 0.7 * prior_belief + 0.3 * signal
        except Exception:
            belief = prior_belief
        belief_arr[0, zi] = float(np.clip(belief, 0.0, 1.0))

    import torch
    return ZoneStateTensor(
        precip=torch.from_numpy(precip_arr),
        uncertainty=torch.from_numpy(uncert_arr),
        belief=torch.from_numpy(belief_arr),
    )


def _build_synthetic_sequence(
    zone_ids: List[str],
    start_date: "datetime",
    n_days: int,
    horizon_days: int,
    config: Optional["ForecastConfig"] = None,
) -> List["ZoneStateTensor"]:
    from zone_observation import make_synthetic_episode_context, _stable_seed
    from physics_dynamics import ZoneStateTensor
    import numpy as np
    import torch

    cfg = config or ForecastConfig()
    sequence: List[ZoneStateTensor] = []

    for day_offset in range(n_days):
        current_date = start_date + timedelta(days=day_offset)
        contexts = []
        for zi, zid in enumerate(zone_ids):
            seed = _stable_seed(
                zid + current_date.isoformat() + str(day_offset)
            )
            ctx = make_synthetic_episode_context(zone_id=zid, seed=seed)
            contexts.append(ctx)

        state = _episode_context_to_state_tensor(
            contexts, horizon_days, prior_belief=cfg.prior_belief
        )
        sequence.append(state)

    return sequence


def _build_era5_sequence(
    zone_ids: List[str],
    start_date: "datetime",
    n_days: int,
    horizon_days: int,
    config: Optional["ForecastConfig"] = None,
) -> List["ZoneStateTensor"]:
    cfg = config or ForecastConfig()
    sequence = []

    for day_offset in range(n_days):
        current_date = _ensure_utc(start_date + timedelta(days=day_offset))
        date_range = (current_date, current_date + timedelta(days=horizon_days))

        contexts = []
        for zid in zone_ids:
            try:
                ctx = fetch_episode_context(zid, date_range, cfg)
            except Exception as e:
                logger.warning(
                    "_build_era5_sequence: failed for zone=%s date=%s: %s — using synthetic",
                    zid, current_date.date().isoformat(), e,
                )
                from zone_observation import make_synthetic_episode_context, _stable_seed
                seed = _stable_seed(zid + current_date.isoformat())
                ctx = make_synthetic_episode_context(zone_id=zid, seed=seed)
            contexts.append(ctx)

        state = _episode_context_to_state_tensor(
            contexts, horizon_days, prior_belief=cfg.prior_belief
        )
        sequence.append(state)

    return sequence


def get_consecutive_pairs(
    zone_ids: List[str],
    start_date: "datetime",
    n_days: int = 365,
    horizon_days: int = 14,
    config: Optional["ForecastConfig"] = None,
    use_real_data: bool = True,
    synthetic_fallback: bool = True,
) -> List[Tuple["ZoneStateTensor", "ZoneStateTensor"]]:
    if not _NUMPY_FOR_DYNAMICS:
        raise ImportError("numpy required for get_consecutive_pairs(). pip install numpy")
    if not _TORCH_FOR_DYNAMICS:
        raise ImportError("torch required for get_consecutive_pairs(). pip install torch")

    if n_days < 2:
        raise ValueError(f"n_days must be >= 2 to produce at least one pair, got {n_days}")

    start_date = _ensure_utc(start_date)

    logger.info(
        "get_consecutive_pairs: zones=%s  start=%s  n_days=%d  horizon=%d  real=%s",
        zone_ids,
        start_date.date().isoformat(),
        n_days,
        horizon_days,
        use_real_data,
    )

    if use_real_data:
        try:
            sequence = _build_era5_sequence(
                zone_ids, start_date, n_days, horizon_days, config
            )
        except Exception as e:
            if not synthetic_fallback:
                raise RuntimeError(
                    f"ERA5 sequence build failed and synthetic_fallback=False: {e}"
                ) from e
            logger.warning(
                "ERA5 sequence build failed (%s) — falling back to full synthetic sequence",
                e,
            )
            sequence = _build_synthetic_sequence(
                zone_ids, start_date, n_days, horizon_days, config
            )
    else:
        sequence = _build_synthetic_sequence(
            zone_ids, start_date, n_days, horizon_days, config
        )

    pairs = [
        (sequence[i], sequence[i + 1])
        for i in range(len(sequence) - 1)
    ]

    logger.info(
        "get_consecutive_pairs: built %d pairs from %d-day sequence",
        len(pairs), n_days,
    )
    return pairs


def get_consecutive_pairs_multi_year(
    zone_ids: List[str],
    years: List[int],
    horizon_days: int = 14,
    config: Optional["ForecastConfig"] = None,
    use_real_data: bool = True,
    skip_on_failure: bool = True,
) -> List[Tuple["ZoneStateTensor", "ZoneStateTensor"]]:
    all_pairs: List[Tuple["ZoneStateTensor", "ZoneStateTensor"]] = []

    for year in years:
        start = _ensure_utc(datetime(year, 1, 1, tzinfo=timezone.utc))
        n_days = 366 if (year % 4 == 0 and (year % 100 != 0 or year % 400 == 0)) else 365

        try:
            year_pairs = get_consecutive_pairs(
                zone_ids=zone_ids,
                start_date=start,
                n_days=n_days,
                horizon_days=horizon_days,
                config=config,
                use_real_data=use_real_data,
                synthetic_fallback=True,
            )
            all_pairs.extend(year_pairs)
            logger.info("Year %d: added %d pairs (total=%d)", year, len(year_pairs), len(all_pairs))

        except Exception as e:
            if not skip_on_failure:
                raise
            logger.warning("Year %d failed (%s) — skipped", year, e)

    if not all_pairs:
        raise RuntimeError(
            f"No pairs collected across years {years}. "
            "Check zone registration and data availability."
        )

    logger.info(
        "get_consecutive_pairs_multi_year: %d total pairs from %d years",
        len(all_pairs), len(years),
    )
    return all_pairs


def compute_dataset_statistics(
    pairs: List[Tuple["ZoneStateTensor", "ZoneStateTensor"]],
) -> Dict[str, Any]:
    if not pairs:
        raise ValueError("pairs is empty")

    import numpy as np
    import torch

    all_precip = []
    all_uncert = []
    all_belief = []

    for curr, nxt in pairs:
        for state in (curr, nxt):
            all_precip.append(state.precip.numpy().flatten())
            all_uncert.append(state.uncertainty.numpy().flatten())
            all_belief.append(state.belief.numpy().flatten())

    precip_all = np.concatenate(all_precip)
    uncert_all = np.concatenate(all_uncert)
    belief_all = np.concatenate(all_belief)

    stats = {
        "precip_mean":  float(np.mean(precip_all)),
        "precip_std":   float(np.std(precip_all)) + 1e-8,
        "uncert_mean":  float(np.mean(uncert_all)),
        "uncert_std":   float(np.std(uncert_all)) + 1e-8,
        "belief_mean":  float(np.mean(belief_all)),
        "belief_std":   float(np.std(belief_all)) + 1e-8,
        "n_pairs":      len(pairs),
        "n_zones":      pairs[0][0].n_zones,
        "horizon_days": pairs[0][0].horizon_days,
        "precip_p95":   float(np.percentile(precip_all, 95)),
        "precip_max":   float(np.max(precip_all)),
    }

    logger.info(
        "Dataset stats: n_pairs=%d  precip_mean=%.1f±%.1f mm  "
        "uncert_mean=%.3f  belief_mean=%.3f",
        stats["n_pairs"],
        stats["precip_mean"], stats["precip_std"],
        stats["uncert_mean"],
        stats["belief_mean"],
    )
    return stats