zone_observation.py

"""
zone_observation.py
===================
"""

from __future__ import annotations

import json
import logging
import math
import random
import zlib
from dataclasses import dataclass, field, asdict
from datetime import datetime, timedelta, timezone
from enum import Enum, unique
from typing import Any, ClassVar, Dict, List, Optional, Tuple

logger = logging.getLogger(__name__)

SCHEMA_VERSION: int = 3

# ---------------------------------------------------------------------------
# Known extras keys
# ---------------------------------------------------------------------------
KNOWN_EXTRAS: Dict[str, str] = {
    "gdd_base_c":           "float -- crop-specific GDD base temperature (degrees C)",
    "crop_substage":        "str   -- variety-specific growth substage",
    "export_grade_risk":    "float -- pre-computed quality risk from field notes [0,1]",
    "edge_node_id":         "str   -- edge sensor network node that sourced this observation",
    "contract_volume_mt":   "float -- contracted volume for this zone (metric tonnes)",
    "soil_type":            "str   -- FAO soil classification string",
    "irrigation_source":    "str   -- 'rainfed' | 'irrigated' | 'supplemental'",
    "sar_flood_date":       "str   -- ISO-8601 date of most recent SAR flood detection",
}


# ---------------------------------------------------------------------------
# Enumerations
# ---------------------------------------------------------------------------

@unique
class CropStage(Enum):
    """Generalised crop growth stage.

    Kept coarse deliberately — variety-specific substages can be added
    via ZoneObs.extras['crop_substage'] without a schema bump.
    """
    UNKNOWN       = "unknown"
    LAND_PREP     = "land_prep"       # tillage, flooding (rice), bed preparation
    PLANTING      = "planting"        # transplanting / direct seeding
    VEGETATIVE    = "vegetative"      # tillering (rice), canopy closure
    REPRODUCTIVE  = "reproductive"    # booting -> heading -> flowering
    GRAIN_FILLING = "grain_filling"   # dough stage -- highest moisture risk
    MATURATION    = "maturation"      # drying down, harvest window opens
    HARVEST       = "harvest"         # active harvest, logistics pressure
    FALLOW        = "fallow"          # between seasons


@unique
class AlertLevel(Enum):
    NONE     = "none"      # no alert warranted
    WATCH    = "watch"     # monitor closely -- conditions developing
    ADVISORY = "advisory"  # elevated risk -- recommend pre-emptive action
    WARNING  = "warning"   # high probability of supply/quality impact
    CRITICAL = "critical"  # immediate action required

    def severity(self) -> int:
        """Integer severity: NONE=0, WATCH=1, ADVISORY=2, WARNING=3, CRITICAL=4."""
        return {"none": 0, "watch": 1, "advisory": 2, "warning": 3, "critical": 4}[self.value]

    def __lt__(self, other: "AlertLevel") -> bool:  # type: ignore[override]
        return self.severity() < other.severity()

    def __le__(self, other: "AlertLevel") -> bool:  # type: ignore[override]
        return self.severity() <= other.severity()

    def __gt__(self, other: "AlertLevel") -> bool:  # type: ignore[override]
        return self.severity() > other.severity()

    def __ge__(self, other: "AlertLevel") -> bool:  # type: ignore[override]
        return self.severity() >= other.severity()


@unique
class DataSource(Enum):
    ERA5_REANALYSIS  = "era5_reanalysis"   # ECMWF ERA5 via CDS or Open-Meteo
    OPENMETEO_LIVE   = "openmeteo_live"    # Open-Meteo forecast API (free tier)
    BMKG_STATION     = "bmkg_station"      # Indonesian met agency station data
    SATELLITE_NDVI   = "satellite_ndvi"    # Sentinel-2 / Landsat NDVI tile
    SATELLITE_PRECIP = "satellite_precip"  # IMERG / CHIRPS retrieval (schema v3+)
    SATELLITE_SOIL   = "satellite_soil"    # SMAP L3/L4 retrieval (schema v3+)
    PUBLISHED_INDEX  = "published_index"   # NOAA/BOM basin-scale index, e.g. ONI/DMI (v3+)
    EDGE_NODE        = "edge_node"         # Distributed edge sensor network node
    SYNTHETIC        = "synthetic"         # Generated by make_synthetic_* for training
    UNKNOWN          = "unknown"

    def is_observational(self) -> bool:
        return self not in (DataSource.SYNTHETIC, DataSource.UNKNOWN)


# ---------------------------------------------------------------------------
# Utilities
# ---------------------------------------------------------------------------

def _clip(value: float, lo: float, hi: float) -> float:
    return max(lo, min(hi, value))


def _stable_seed(key: str) -> int:
    return zlib.crc32(key.encode("utf-8")) & 0x7FFFFFFF


def _copy_and_pop_schema(d: Dict[str, Any]) -> Tuple[Optional[int], Dict[str, Any]]:
    d_copy = dict(d)
    sv = d_copy.pop("_schema_version", None)
    return sv, d_copy


def _check_schema(sv: Optional[int], class_name: str) -> None:
    if sv is not None and sv != SCHEMA_VERSION:
        raise ValueError(
            f"{class_name}.from_dict: schema version mismatch -- "
            f"stored={sv}, current={SCHEMA_VERSION}. "
            f"Run migration script or increment SCHEMA_VERSION."
        )


# ---------------------------------------------------------------------------
# ForecastConfig
# ---------------------------------------------------------------------------

@dataclass
class ForecastConfig:

    # --- Spatial ---
    n_zones:              int   = 1      # number of sourcing zones per episode
    horizon_days:         int   = 30     # forecast horizon (days)
    max_steps:            int   = 200    # max env steps per episode

    # --- Belief map ---
    prior_belief:         float = 0.12   # initial P(risk_event) per zone
    belief_floor:         float = 0.005  # minimum belief after decay
    belief_update_radius: int   = 2      # spatial propagation radius (zone cells)
    belief_increase_rate: float = 0.30   # update magnitude when event confirmed
    belief_decrease_rate: float = 0.05   # update magnitude when event absent
    belief_prior_weight:  float = 0.70   # weight on episode prior vs. zone signal when
                                          # seeding zone_belief at reset:
                                          # prior_w*prior + (1-prior_w)*signal

    # --- Economics ---
    alert_value:          float = 100.0  # reward for correct advisory issuance
    false_alert_penalty:  float = 20.0   # penalty for unnecessary advisory
    miss_penalty:         float = 200.0  # penalty per missed true risk event
    inspection_cost:      float = 1.0    # cost per zone-step (resource use)


    zone_visit_bonus:       float = 1.5   # raw bonus, first visit to a zone only
    unvisited_zone_penalty: float = 40.0  # raw penalty * (unvisited/active) at terminate

    uncertainty_decay:         float = 0.70  # forecast_uncertainty[zone] *= this on inspection
    info_gain_scale:           float = 5.0   # multiplier on belief info-gain in step reward
    uncertainty_penalty_scale: float = 5.0   # multiplier on mean uncertainty at termination

    economic_randomization: bool  = False
    clean_episode_ratio:    float = 0.7
    event_spatial_correlation: float = 0.85
    shuffle_zone_order: bool = True
    alert_value_range:      Tuple[float, float] = (0.8, 1.2)
    miss_penalty_range:     Tuple[float, float] = (0.9, 1.1)

    soft_reset:  bool = True

    seed: Optional[int] = None

    real_data_ratio:   float               = 0.7   # fraction of episodes that attempt real data
    era5_ratio:        float               = 0.5   # of real-data attempts, fraction using ERA5
    force_data_source: Optional[DataSource] = None  # pin source for debug/test (overrides above)
    inject_noise:      bool                = False  # apply stochastic noise after fetch
    noise_scale:       float               = 0.05   # noise magnitude (fraction of field range)

    use_satellite_precip: bool  = False  # prefer IMERG/CHIRPS over ERA5 precip
    use_satellite_soil:   bool  = False  # prefer SMAP over ERA5 soil moisture
    include_basin_context: bool = False  # attach ENSO/IOD/monsoon/helio context to EpisodeContext
    require_real_basin_context: bool = False

    forecast_backend: str = "synthetic"

    use_climatology_anomalies: bool = False
    climatology_years:        int   = 10  # years of history for the climatology

    def __post_init__(self) -> None:
        if self.alert_value <= 0:
            raise ValueError(f"ForecastConfig: alert_value={self.alert_value} must be > 0")
        if self.false_alert_penalty < 0:
            raise ValueError(f"ForecastConfig: false_alert_penalty must be >= 0")
        if self.miss_penalty <= 0:
            raise ValueError(f"ForecastConfig: miss_penalty must be > 0")
        if self.inspection_cost <= 0:
            raise ValueError(f"ForecastConfig: inspection_cost must be > 0")
        if self.zone_visit_bonus < 0:
            raise ValueError(f"ForecastConfig: zone_visit_bonus must be >= 0")
        if self.unvisited_zone_penalty < 0:
            raise ValueError(f"ForecastConfig: unvisited_zone_penalty must be >= 0")
        if not (0.0 <= self.belief_prior_weight <= 1.0):
            raise ValueError(
                f"ForecastConfig: belief_prior_weight={self.belief_prior_weight} "
                f"must be in [0, 1]"
            )
        if not (0.0 <= self.uncertainty_decay <= 1.0):
            raise ValueError(
                f"ForecastConfig: uncertainty_decay={self.uncertainty_decay} "
                f"must be in [0, 1]"
            )
        if self.info_gain_scale < 0:
            raise ValueError(f"ForecastConfig: info_gain_scale must be >= 0")
        if self.uncertainty_penalty_scale < 0:
            raise ValueError(f"ForecastConfig: uncertainty_penalty_scale must be >= 0")
        if self.horizon_days < 1:
            raise ValueError(f"ForecastConfig: horizon_days must be >= 1")
        if self.n_zones < 1:
            raise ValueError(f"ForecastConfig: n_zones must be >= 1")
        if self.max_steps < 1:
            raise ValueError(f"ForecastConfig: max_steps must be >= 1")
        if not (0.0 <= self.real_data_ratio <= 1.0):
            raise ValueError(
                f"ForecastConfig: real_data_ratio={self.real_data_ratio} must be in [0, 1]"
            )
        if not (0.0 <= self.era5_ratio <= 1.0):
            raise ValueError(
                f"ForecastConfig: era5_ratio={self.era5_ratio} must be in [0, 1]"
            )
        if self.noise_scale < 0.0:
            raise ValueError(
                f"ForecastConfig: noise_scale={self.noise_scale} must be >= 0"
            )
        _VALID_BACKENDS = ("synthetic", "baseline", "openmeteo", "timesfm")
        if self.forecast_backend not in _VALID_BACKENDS:
            raise ValueError(
                f"ForecastConfig: forecast_backend={self.forecast_backend!r} "
                f"must be one of {_VALID_BACKENDS}"
            )
        if not (1 <= self.climatology_years <= 30):
            raise ValueError(
                f"ForecastConfig: climatology_years={self.climatology_years} "
                f"must be in [1, 30]"
            )

        self.alert_value         = float(_clip(self.alert_value,         0.1,   10_000.0))
        self.false_alert_penalty = float(_clip(self.false_alert_penalty, 0.0,   10_000.0))
        self.miss_penalty        = float(_clip(self.miss_penalty,        0.1,  100_000.0))
        self.inspection_cost     = float(_clip(self.inspection_cost,     0.01,  1_000.0))
        self.zone_visit_bonus       = float(_clip(self.zone_visit_bonus,       0.0, 1_000.0))
        self.unvisited_zone_penalty = float(_clip(self.unvisited_zone_penalty, 0.0, 10_000.0))
        self.belief_prior_weight       = float(_clip(self.belief_prior_weight,       0.0, 1.0))
        self.uncertainty_decay         = float(_clip(self.uncertainty_decay,         0.0, 1.0))
        self.info_gain_scale           = float(_clip(self.info_gain_scale,           0.0, 1_000.0))
        self.uncertainty_penalty_scale = float(_clip(self.uncertainty_penalty_scale, 0.0, 1_000.0))
        self.prior_belief        = float(_clip(self.prior_belief,        0.001, 0.999))
        self.belief_floor        = float(_clip(self.belief_floor,        0.001, 0.5))
        self.clean_episode_ratio = float(_clip(self.clean_episode_ratio, 0.0,   1.0))
        self.event_spatial_correlation = float(
            _clip(self.event_spatial_correlation, 0.0, 1.0)
        )

        rational = self.false_alert_penalty / max(
            self.alert_value + self.false_alert_penalty + self.miss_penalty, 1e-9
        )
        if self.belief_floor >= rational:
            logger.warning(
                f"ForecastConfig: belief_floor={self.belief_floor:.4f} >= "
                f"rational_termination_threshold={rational:.4f}. "
                f"Early termination will never be EV-positive. "
                f"Set belief_floor < {rational:.4f}."
            )
        self._rational_threshold: float = rational

    @property
    def rational_termination_threshold(self) -> float:
        return self._rational_threshold

    def to_dict(self) -> Dict[str, Any]:
        return {
            "n_zones": self.n_zones,
            "horizon_days": self.horizon_days,
            "max_steps": self.max_steps,
            "prior_belief": self.prior_belief,
            "belief_floor": self.belief_floor,
            "belief_update_radius": self.belief_update_radius,
            "belief_increase_rate": self.belief_increase_rate,
            "belief_decrease_rate": self.belief_decrease_rate,
            "belief_prior_weight": self.belief_prior_weight,
            "alert_value": self.alert_value,
            "false_alert_penalty": self.false_alert_penalty,
            "miss_penalty": self.miss_penalty,
            "inspection_cost": self.inspection_cost,
            "zone_visit_bonus": self.zone_visit_bonus,
            "unvisited_zone_penalty": self.unvisited_zone_penalty,
            "uncertainty_decay": self.uncertainty_decay,
            "info_gain_scale": self.info_gain_scale,
            "uncertainty_penalty_scale": self.uncertainty_penalty_scale,
            "economic_randomization": self.economic_randomization,
            "clean_episode_ratio": self.clean_episode_ratio,
            "event_spatial_correlation": self.event_spatial_correlation,
            "shuffle_zone_order": self.shuffle_zone_order,
            "alert_value_range": list(self.alert_value_range),
            "miss_penalty_range": list(self.miss_penalty_range),
            "soft_reset": self.soft_reset,
            "seed": self.seed,
            "real_data_ratio": self.real_data_ratio,
            "era5_ratio": self.era5_ratio,
            "force_data_source": (
                self.force_data_source.value if self.force_data_source is not None else None
            ),
            "inject_noise": self.inject_noise,
            "noise_scale": self.noise_scale,
            "use_satellite_precip": self.use_satellite_precip,
            "use_satellite_soil": self.use_satellite_soil,
            "include_basin_context": self.include_basin_context,
            "require_real_basin_context": self.require_real_basin_context,
            "forecast_backend": self.forecast_backend,
            "use_climatology_anomalies": self.use_climatology_anomalies,
            "climatology_years": self.climatology_years,
            "_schema_version": SCHEMA_VERSION,
        }

    @classmethod
    def from_dict(cls, d: Dict[str, Any]) -> "ForecastConfig":
        sv, d = _copy_and_pop_schema(d)
        _check_schema(sv, "ForecastConfig")
        if "alert_value_range" in d and isinstance(d["alert_value_range"], list):
            d["alert_value_range"] = tuple(d["alert_value_range"])
        if "miss_penalty_range" in d and isinstance(d["miss_penalty_range"], list):
            d["miss_penalty_range"] = tuple(d["miss_penalty_range"])
        if "force_data_source" in d and d["force_data_source"] is not None:
            d["force_data_source"] = DataSource(d["force_data_source"])
        return cls(**{k: v for k, v in d.items() if not k.startswith("_")})


# ---------------------------------------------------------------------------
# GeoPolygon
# ---------------------------------------------------------------------------

@dataclass
class GeoPolygon:
    vertices: List[Tuple[float, float]]   # [(lat degrees, lon degrees), ...]
    zone_id:  str
    label:    str = ""

    def __post_init__(self) -> None:
        self.vertices = [(float(v[0]), float(v[1])) for v in self.vertices]
        if len(self.vertices) < 3:
            raise ValueError(
                f"GeoPolygon '{self.zone_id}' needs >= 3 vertices, "
                f"got {len(self.vertices)}"
            )
        for lat, lon in self.vertices:
            if not (-90.0 <= lat <= 90.0):
                raise ValueError(
                    f"GeoPolygon '{self.zone_id}': latitude {lat} out of [-90, 90]"
                )
            if not (-180.0 <= lon <= 180.0):
                raise ValueError(
                    f"GeoPolygon '{self.zone_id}': longitude {lon} out of [-180, 180]"
                )

    @property
    def centroid(self) -> Tuple[float, float]:
        lats = [v[0] for v in self.vertices]
        lons = [v[1] for v in self.vertices]
        return (sum(lats) / len(lats), sum(lons) / len(lons))

    @property
    def approx_area_km2(self) -> float:
        lat_c, _ = self.centroid
        km_per_deg_lat = 111.0
        km_per_deg_lon = 111.0 * math.cos(math.radians(lat_c))
        n = len(self.vertices)
        area = 0.0
        for i in range(n):
            x0 = self.vertices[i][1] * km_per_deg_lon
            y0 = self.vertices[i][0] * km_per_deg_lat
            x1 = self.vertices[(i + 1) % n][1] * km_per_deg_lon
            y1 = self.vertices[(i + 1) % n][0] * km_per_deg_lat
            area += x0 * y1 - x1 * y0
        return abs(area) / 2.0

    def contains_point(self, lat: float, lon: float) -> bool:
        n = len(self.vertices)
        inside = False
        j = n - 1
        for i in range(n):
            xi, yi = self.vertices[i][1], self.vertices[i][0]
            xj, yj = self.vertices[j][1], self.vertices[j][0]
            if ((yi > lat) != (yj > lat)) and (
                lon < (xj - xi) * (lat - yi) / (yj - yi + 1e-12) + xi
            ):
                inside = not inside
            j = i
        return inside

    def to_dict(self) -> Dict[str, Any]:
        return {
            "vertices": [list(v) for v in self.vertices],
            "zone_id":  self.zone_id,
            "label":    self.label,
        }

    @classmethod
    def from_dict(cls, d: Dict[str, Any]) -> "GeoPolygon":
        return cls(
            vertices=[(float(v[0]), float(v[1])) for v in d["vertices"]],
            zone_id=d["zone_id"],
            label=d.get("label", ""),
        )


# ---------------------------------------------------------------------------
# ZoneObs
# ---------------------------------------------------------------------------

@dataclass
class ZoneObs:

    zone_id:    str
    valid_time: datetime               # UTC timestamp of observation window start
    source:     DataSource = DataSource.UNKNOWN

    precip_24h_mm:      float = 0.0   # total precip last 24 h (mm)
    precip_7d_mm:       float = 0.0   # total precip last 7 days (mm)
    precip_14d_mm:      float = 0.0   # total precip last 14 days (mm)
    precip_30d_mm:      float = 0.0   # total precip last 30 days (mm)
    precip_anomaly_idx: float = 0.0   # z-score vs ERA5 climatological mean
                                      # negative = drought, positive = excess

    temp_mean_c:        float = 0.0   # daily mean (degrees C)
    temp_max_c:         float = 0.0   # daily maximum (degrees C)
    temp_min_c:         float = 0.0   # daily minimum (degrees C)
    temp_anomaly_idx:   float = 0.0   # z-score vs climatological mean

    gdd_accumulated:    float = 0.0   # growing degree-days since sowing
                                      # base temp is crop-specific -> extras['gdd_base_c']
    heat_stress_days:   int   = 0     # days where temp_max_c > threshold (default 35 C)
    cold_stress_days:   int   = 0     # days where temp_min_c < threshold (default 15 C)

    soil_moisture_pct:     float = 0.0   # volumetric water content top 10cm, 0-100
    soil_moisture_anom:    float = 0.0   # z-score vs climatological mean
    evapotranspiration_mm: float = 0.0   # reference ET0 (FAO-56 Penman-Monteith), mm/day

    precip_satellite_mm:        Optional[float] = None  # IMERG/CHIRPS daily total (mm)
    soil_moisture_satellite_pct: Optional[float] = None  # SMAP L3/L4 retrieval (%, 0-100)

    wind_speed_max_ms:  float = 0.0   # maximum gust in window (m/s)
    wind_speed_mean_ms: float = 0.0   # mean 10m wind speed (m/s)

    rh_mean_pct:        float = 0.0   # relative humidity daily mean, 0-100
    rh_max_pct:         float = 0.0   # daily maximum, 0-100; key fungi risk driver
    rh_anomaly_idx:      float = 0.0

    ndvi:               Optional[float] = None  # NDVI -1.0 to 1.0; None if no recent pass
    ndvi_anomaly_idx:   Optional[float] = None  # z-score vs same-DOY climatology
    ndvi_trend_14d:     Optional[float] = None  # linear slope over 14 days (NDVI/day)

    flood_extent_pct:   float = 0.0   # % of zone with standing water (SAR-derived), 0-100
    drainage_risk_idx:  float = 0.0   # composite: slope + soil type + recent precip, 0-1

    crop_stage:         CropStage = CropStage.UNKNOWN
    days_to_harvest:    Optional[int] = None   # None = unknown; 0 = harvest now
    planting_date:      Optional[datetime] = None

    quality_flag:       int   = 0     # 0=good, 1=interpolated, 2=gap-filled, 3=synthetic
    cloud_cover_pct:    float = 0.0   # cloud fraction 0-100; high values degrade NDVI

    extras: Dict[str, Any] = field(default_factory=dict)

    def __post_init__(self) -> None:
        self.extras = dict(self.extras)

        if self.valid_time.tzinfo is None:
            logger.warning(
                f"ZoneObs('{self.zone_id}'): valid_time has no timezone, assuming UTC."
            )
            self.valid_time = self.valid_time.replace(tzinfo=timezone.utc)

        self.soil_moisture_pct  = float(_clip(self.soil_moisture_pct,  0.0, 100.0))
        self.rh_mean_pct        = float(_clip(self.rh_mean_pct,        0.0, 100.0))
        self.rh_max_pct         = float(_clip(self.rh_max_pct,         0.0, 100.0))
        self.flood_extent_pct   = float(_clip(self.flood_extent_pct,   0.0, 100.0))
        self.cloud_cover_pct    = float(_clip(self.cloud_cover_pct,    0.0, 100.0))
        self.drainage_risk_idx  = float(_clip(self.drainage_risk_idx,  0.0, 1.0))

        self.precip_anomaly_idx = float(_clip(self.precip_anomaly_idx, -5.0, 5.0))
        self.temp_anomaly_idx   = float(_clip(self.temp_anomaly_idx,   -5.0, 5.0))
        self.soil_moisture_anom = float(_clip(self.soil_moisture_anom, -5.0, 5.0))
        self.rh_anomaly_idx     = float(_clip(self.rh_anomaly_idx,     -5.0, 5.0))

        if self.ndvi is not None:
            self.ndvi = float(_clip(self.ndvi, -1.0, 1.0))

        if self.soil_moisture_satellite_pct is not None:
            self.soil_moisture_satellite_pct = float(
                _clip(self.soil_moisture_satellite_pct, 0.0, 100.0)
            )
        if self.precip_satellite_mm is not None and self.precip_satellite_mm < 0.0:
            logger.warning(
                f"ZoneObs('{self.zone_id}'): precip_satellite_mm="
                f"{self.precip_satellite_mm:.4f} < 0, clipping to 0."
            )
            self.precip_satellite_mm = 0.0

        for attr in (
            "precip_24h_mm", "precip_7d_mm", "precip_14d_mm", "precip_30d_mm",
            "evapotranspiration_mm", "wind_speed_max_ms", "wind_speed_mean_ms",
            "gdd_accumulated",
        ):
            val = getattr(self, attr)
            if val < 0.0:
                logger.warning(
                    f"ZoneObs('{self.zone_id}'): {attr}={val:.4f} < 0, clipping to 0."
                )
                setattr(self, attr, 0.0)

        if self.heat_stress_days < 0:
            self.heat_stress_days = 0
        if self.cold_stress_days < 0:
            self.cold_stress_days = 0
        if self.quality_flag not in (0, 1, 2, 3):
            logger.warning(
                f"ZoneObs('{self.zone_id}'): quality_flag={self.quality_flag} "
                f"not in {{0,1,2,3}}, setting to 3."
            )
            self.quality_flag = 3

        if self.planting_date is not None and self.planting_date.tzinfo is None:
            self.planting_date = self.planting_date.replace(tzinfo=timezone.utc)

        if not (
            self.precip_30d_mm >= self.precip_14d_mm
            >= self.precip_7d_mm >= self.precip_24h_mm
        ):
            logger.warning(
                f"ZoneObs('{self.zone_id}'): non-monotonic precipitation aggregates "
                f"(24h={self.precip_24h_mm:.2f}, 7d={self.precip_7d_mm:.2f}, "
                f"14d={self.precip_14d_mm:.2f}, 30d={self.precip_30d_mm:.2f})"
            )

    @classmethod
    def validate(cls, obs: "ZoneObs", strict: bool = False) -> List[str]:
        issues: List[str] = []
        if not obs.zone_id:
            issues.append("zone_id is empty")
        if obs.temp_max_c < obs.temp_min_c:
            issues.append(f"temp_max_c={obs.temp_max_c} < temp_min_c={obs.temp_min_c}")
        if obs.precip_14d_mm < obs.precip_7d_mm:
            issues.append(f"precip_14d_mm < precip_7d_mm")
        if obs.precip_30d_mm < obs.precip_14d_mm:
            issues.append(f"precip_30d_mm < precip_14d_mm")
        if obs.wind_speed_max_ms < obs.wind_speed_mean_ms:
            issues.append(f"wind_speed_max_ms < wind_speed_mean_ms")
        if obs.rh_max_pct < obs.rh_mean_pct:
            issues.append(f"rh_max_pct < rh_mean_pct")
        if obs.days_to_harvest is not None and obs.days_to_harvest < 0:
            issues.append(f"days_to_harvest={obs.days_to_harvest} < 0")
        if obs.quality_flag >= 2 and obs.source.is_observational():
            issues.append(
                f"quality_flag={obs.quality_flag} (gap-filled/synthetic) "
                f"but source={obs.source.value} is observational"
            )
        if strict and issues:
            raise ValueError(f"ZoneObs('{obs.zone_id}') strict validation failed: {issues}")
        return issues

    def to_dict(self) -> Dict[str, Any]:
        d = asdict(self)
        d["source"]        = self.source.value
        d["crop_stage"]    = self.crop_stage.value
        d["valid_time"]    = self.valid_time.isoformat()
        d["planting_date"] = self.planting_date.isoformat() if self.planting_date else None
        d["_schema_version"] = SCHEMA_VERSION
        return d

    @classmethod
    def from_dict(cls, d: Dict[str, Any]) -> "ZoneObs":
        sv, d = _copy_and_pop_schema(d)
        _check_schema(sv, "ZoneObs")
        d["valid_time"]    = datetime.fromisoformat(d["valid_time"])
        d["source"]        = DataSource(d["source"])
        d["crop_stage"]    = CropStage(d["crop_stage"])
        d["planting_date"] = (
            datetime.fromisoformat(d["planting_date"]) if d.get("planting_date") else None
        )
        return cls(**d)


    def is_harvest_window(self, lookahead_days: int = 21) -> bool:
        if self.days_to_harvest is None:
            return self.crop_stage in (CropStage.MATURATION, CropStage.HARVEST)
        return 0 <= self.days_to_harvest <= lookahead_days

    def has_reliable_ndvi(self) -> bool:
        return self.ndvi is not None and self.cloud_cover_pct < 30.0

    def drought_signal(self) -> float:
        return float(
            0.6 * _clip(-self.precip_anomaly_idx / 3.0, 0.0, 1.0)
            + 0.4 * _clip(-self.soil_moisture_anom / 3.0, 0.0, 1.0)
        )

    def flood_signal(self) -> float:
        return float(
            0.4 * _clip(self.precip_anomaly_idx / 3.0, 0.0, 1.0)
            + 0.4 * (self.flood_extent_pct / 100.0)
            + 0.2 * _clip(self.drainage_risk_idx, 0.0, 1.0)
        )

    def fungi_risk_signal(self) -> float:
        rh_s = _clip((self.rh_max_pct - 70.0) / 30.0, 0.0, 1.0)
        anomaly_adj = _clip(self.rh_anomaly_idx / 3.0, -0.3, 0.3)
        rh_s_adjusted = _clip(rh_s + anomaly_adj, 0.0, 1.0)
        mult = (
            1.0 if self.crop_stage in (CropStage.GRAIN_FILLING, CropStage.MATURATION)
            else 0.5
        )
        return float(rh_s_adjusted * mult)

    def composite_risk(self) -> float:
        return float(_clip(
            0.35 * self.drought_signal()
            + 0.40 * self.flood_signal()
            + 0.25 * self.fungi_risk_signal(),
            0.0, 1.0,
        ))


# ---------------------------------------------------------------------------
# ForecastResult
# ---------------------------------------------------------------------------

@dataclass(frozen=True)
class ForecastResult:
    zone_id:       str
    forecast_time: datetime
    horizon_days:  int = 30

    precip_mm:     Tuple[float, ...] = field(default_factory=tuple)
    temp_mean_c:   Tuple[float, ...] = field(default_factory=tuple)
    rh_mean_pct:   Tuple[float, ...] = field(default_factory=tuple)

    precip_p10:    Tuple[float, ...] = field(default_factory=tuple)
    precip_p90:    Tuple[float, ...] = field(default_factory=tuple)
    temp_p10:      Tuple[float, ...] = field(default_factory=tuple)
    temp_p90:      Tuple[float, ...] = field(default_factory=tuple)

    prob_heavy_rain:    Tuple[float, ...] = field(default_factory=tuple)
    prob_drought_day:   Tuple[float, ...] = field(default_factory=tuple)
    prob_high_humidity: Tuple[float, ...] = field(default_factory=tuple)

    model_id:   str            = "timesfm-2.5-200m"
    crps_score: Optional[float] = None
    source:     DataSource     = DataSource.SYNTHETIC
    extras:     Dict[str, Any] = field(default_factory=dict)

    _SEQUENCE_FIELDS: ClassVar[Tuple[str, ...]] = (
        "precip_mm", "temp_mean_c", "rh_mean_pct",
        "precip_p10", "precip_p90", "temp_p10", "temp_p90",
        "prob_heavy_rain", "prob_drought_day", "prob_high_humidity",
    )
    _PROB_FIELDS: ClassVar[Tuple[str, ...]] = (
        "prob_heavy_rain", "prob_drought_day", "prob_high_humidity",
    )

    def __post_init__(self) -> None:
        object.__setattr__(self, "extras", dict(self.extras))

        seqs = [
            (name, getattr(self, name))
            for name in self._SEQUENCE_FIELDS
            if getattr(self, name)
        ]
        if seqs:
            lengths = {len(s) for _, s in seqs}
            if len(lengths) > 1:
                raise ValueError(
                    f"ForecastResult('{self.zone_id}'): sequence length mismatch: "
                    f"{ {n: len(s) for n, s in seqs} }"
                )
            expected = self.horizon_days
            for name, seq in seqs:
                if len(seq) != expected:
                    raise ValueError(
                        f"ForecastResult('{self.zone_id}'): {name} length={len(seq)} "
                        f"!= horizon_days={expected}. Truncate or pad before constructing."
                    )

        for fname in self._PROB_FIELDS:
            for i, v in enumerate(getattr(self, fname)):
                if not (0.0 <= v <= 1.0):
                    raise ValueError(
                        f"ForecastResult('{self.zone_id}'): "
                        f"{fname}[{i}]={v:.4f} outside [0, 1]. "
                        f"Clip before constructing ForecastResult."
                    )

        for i, (lo, hi) in enumerate(zip(self.precip_p10, self.precip_p90)):
            if lo > hi:
                raise ValueError(
                    f"ForecastResult('{self.zone_id}'): "
                    f"precip_p10[{i}]={lo} > precip_p90[{i}]={hi}"
                )

    def peak_precip_day(self) -> Optional[int]:
        if not self.precip_mm:
            return None
        return int(max(range(len(self.precip_mm)), key=lambda i: self.precip_mm[i]))

    def cumulative_precip_mm(self, window_days: int = 14) -> float:
        return float(sum(self.precip_mm[:window_days]))

    def max_consecutive_rain_days(self, threshold_mm: float = 10.0) -> int:
        max_run = run = 0
        for p in self.precip_mm:
            run = run + 1 if p > threshold_mm else 0
            max_run = max(max_run, run)
        return max_run

    def mean_exceedance_prob(
        self, field_name: str, window_days: Optional[int] = None
    ) -> float:
        seq = getattr(self, field_name, ())
        if not seq:
            return 0.0
        window = seq[:window_days] if window_days else seq
        return float(sum(window) / len(window))

    def to_dict(self) -> Dict[str, Any]:
        d = asdict(self)
        d["forecast_time"] = self.forecast_time.isoformat()
        d["source"]        = self.source.value
        d["_schema_version"] = SCHEMA_VERSION
        return d

    @classmethod
    def from_dict(cls, d: Dict[str, Any]) -> "ForecastResult":
        sv, d = _copy_and_pop_schema(d)
        _check_schema(sv, "ForecastResult")
        d["forecast_time"] = datetime.fromisoformat(d["forecast_time"])
        d["source"]        = DataSource(d["source"])
        for k in cls._SEQUENCE_FIELDS:
            if k in d and isinstance(d[k], list):
                d[k] = tuple(float(v) for v in d[k])
        return cls(**{k: v for k, v in d.items() if not k.startswith("_")})


# ---------------------------------------------------------------------------
# RiskScore
# ---------------------------------------------------------------------------

@dataclass(frozen=True)
class RiskScore:
    zone_id:   str
    scored_at: datetime

    supply_shortfall_prob:  float = 0.0  # P(zone delivers < 80% of contracted volume)
    drought_risk:           float = 0.0  # [0, 1]
    flood_risk:             float = 0.0  # [0, 1]
    supply_risk_composite:  float = 0.0  # weighted dashboard score [0, 1]

    fungi_contamination_prob: float = 0.0  # P(moisture-related quality downgrade)
    harvest_delay_days:       float = 0.0  # expected delay in days; >= 0
    quality_risk_composite:   float = 0.0  # [0, 1]

    optimal_harvest_window_start: Optional[datetime] = None
    optimal_harvest_window_end:   Optional[datetime] = None

    alert_level:  AlertLevel = AlertLevel.NONE
    action_notes: str = ""

    confidence: float = 0.5  # [0, 1]

    extras: Dict[str, Any] = field(default_factory=dict)

    _PROB_FIELDS: ClassVar[Tuple[str, ...]] = (
        "supply_shortfall_prob", "drought_risk", "flood_risk",
        "supply_risk_composite", "fungi_contamination_prob",
        "quality_risk_composite", "confidence",
    )

    def __post_init__(self) -> None:
        object.__setattr__(self, "extras", dict(self.extras))

        for attr in self._PROB_FIELDS:
            val = getattr(self, attr)
            clipped = _clip(val, 0.0, 1.0)
            if abs(clipped - val) > 1e-9:
                logger.warning(
                    f"RiskScore('{self.zone_id}'): {attr}={val:.4f} "
                    f"outside [0,1], clipped to {clipped:.4f}."
                )
            object.__setattr__(self, attr, float(clipped))

        if self.harvest_delay_days < 0.0:
            object.__setattr__(self, "harvest_delay_days", 0.0)

        if self.scored_at.tzinfo is None:
            raise ValueError(
                f"RiskScore('{self.zone_id}'): scored_at must be timezone-aware (UTC). "
                f"Use datetime.now(tz=timezone.utc) or .replace(tzinfo=timezone.utc)."
            )
        for dt_attr in ("optimal_harvest_window_start", "optimal_harvest_window_end"):
            dt = getattr(self, dt_attr)
            if dt is not None and dt.tzinfo is None:
                raise ValueError(
                    f"RiskScore('{self.zone_id}'): {dt_attr} must be timezone-aware (UTC)."
                )

    def is_actionable(self) -> bool:
        return self.alert_level > AlertLevel.WATCH

    def is_elevated(self) -> bool:
        return self.alert_level.severity() >= AlertLevel.ADVISORY.severity()

    def is_product_actionable(self) -> bool:
        return self.alert_level.severity() >= AlertLevel.WARNING.severity()

    def harvest_window_days(self) -> Optional[int]:
        if self.optimal_harvest_window_start and self.optimal_harvest_window_end:
            return max(
                0,
                (self.optimal_harvest_window_end
                 - self.optimal_harvest_window_start).days,
            )
        return None

    def to_dict(self) -> Dict[str, Any]:
        d = asdict(self)
        d["scored_at"]   = self.scored_at.isoformat()
        d["alert_level"] = self.alert_level.value
        d["optimal_harvest_window_start"] = (
            self.optimal_harvest_window_start.isoformat()
            if self.optimal_harvest_window_start else None
        )
        d["optimal_harvest_window_end"] = (
            self.optimal_harvest_window_end.isoformat()
            if self.optimal_harvest_window_end else None
        )
        d["_schema_version"] = SCHEMA_VERSION
        return d

    @classmethod
    def from_dict(cls, d: Dict[str, Any]) -> "RiskScore":
        sv, d = _copy_and_pop_schema(d)
        _check_schema(sv, "RiskScore")
        d["scored_at"]   = datetime.fromisoformat(d["scored_at"])
        d["alert_level"] = AlertLevel(d["alert_level"])
        d["optimal_harvest_window_start"] = (
            datetime.fromisoformat(d["optimal_harvest_window_start"])
            if d.get("optimal_harvest_window_start") else None
        )
        d["optimal_harvest_window_end"] = (
            datetime.fromisoformat(d["optimal_harvest_window_end"])
            if d.get("optimal_harvest_window_end") else None
        )
        return cls(**{k: v for k, v in d.items() if not k.startswith("_")})


# ---------------------------------------------------------------------------
# BasinContext  (schema v3+)
# ---------------------------------------------------------------------------

_HELIO_REGIMES = frozenset({"quiet", "active", "storm"})


def derive_helio_regime(kp_index: float, goes_xray_flux: float) -> str:
    """Classify heliophysical regime from Kp and GOES X-ray flux."""
    kp = float(kp_index)
    xray = float(goes_xray_flux) if goes_xray_flux is not None else 1e-7
    if kp >= 5.0 or xray >= 1e-5:
        return "storm"
    if kp >= 3.0 or xray >= 5e-7:
        return "active"
    return "quiet"


@dataclass
class BasinContext:
    """Basin-scale teleconnections + heliophysical context (schema v3+)."""
    valid_date: datetime

    enso_oni:          float = 0.0      # Oceanic (or Relative Oceanic) Nino Index, degrees C anomaly
    iod_dmi:            float = 0.0      # Indian Ocean Dipole Mode Index, degrees C
    itcz_latitude_deg:  float = 0.0      # approximate ITCZ position, degrees N (negative = south)
    mslp_regional_hpa:  float = 1013.25  # area-averaged regional MSLP, hPa (monsoon high/low proxy)

    # Helio / space-weather (quiet-Sun defaults — anti-saturation design)
    solar_wind_speed_kms: float = 400.0  # typical quiet-Sun ~300–450 km/s
    kp_index:             float = 2.0    # planetary K-index [0, 9]; ~2 is quiet
    goes_xray_flux:       float = 1e-7   # W/m²; background / low-C floor
    helio_regime:         str   = "quiet"  # "quiet" | "active" | "storm"

    source:  DataSource = DataSource.SYNTHETIC
    extras:  Dict[str, Any] = field(default_factory=dict)

    def __post_init__(self) -> None:
        object.__setattr__(self, "extras", dict(self.extras))

        if self.valid_date.tzinfo is None:
            object.__setattr__(
                self, "valid_date", self.valid_date.replace(tzinfo=timezone.utc)
            )

        object.__setattr__(self, "enso_oni",         float(_clip(self.enso_oni, -5.0, 5.0)))
        object.__setattr__(self, "iod_dmi",           float(_clip(self.iod_dmi, -5.0, 5.0)))
        object.__setattr__(self, "itcz_latitude_deg", float(_clip(self.itcz_latitude_deg, -30.0, 30.0)))
        object.__setattr__(self, "mslp_regional_hpa", float(_clip(self.mslp_regional_hpa, 900.0, 1100.0)))

        # Helio clipping — physical ranges, not risk-amplifying floors.
        object.__setattr__(
            self, "solar_wind_speed_kms",
            float(_clip(self.solar_wind_speed_kms, 200.0, 1200.0)),
        )
        object.__setattr__(self, "kp_index", float(_clip(self.kp_index, 0.0, 9.0)))
        object.__setattr__(
            self, "goes_xray_flux",
            float(_clip(self.goes_xray_flux, 1e-9, 1e-3)),
        )
        regime = self.helio_regime if self.helio_regime in _HELIO_REGIMES else "quiet"
        object.__setattr__(self, "helio_regime", regime)

    def to_dict(self) -> Dict[str, Any]:
        return {
            "valid_date":            self.valid_date.isoformat(),
            "enso_oni":              self.enso_oni,
            "iod_dmi":               self.iod_dmi,
            "itcz_latitude_deg":     self.itcz_latitude_deg,
            "mslp_regional_hpa":     self.mslp_regional_hpa,
            "solar_wind_speed_kms":  self.solar_wind_speed_kms,
            "kp_index":              self.kp_index,
            "goes_xray_flux":        self.goes_xray_flux,
            "helio_regime":          self.helio_regime,
            "source":                self.source.value,
            "extras":                self.extras,
            "_schema_version":       SCHEMA_VERSION,
        }

    @classmethod
    def from_dict(cls, d: Dict[str, Any]) -> "BasinContext":
        sv, d = _copy_and_pop_schema(d)
        _check_schema(sv, "BasinContext")
        d["valid_date"] = datetime.fromisoformat(d["valid_date"])
        d["source"] = DataSource(d.get("source", "synthetic"))
        return cls(**{k: v for k, v in d.items() if not k.startswith("_")})
        
def make_synthetic_basin_context(
    valid_date: Optional[datetime] = None,
    seed:       Optional[int]      = None,
) -> BasinContext:
    rng = random.Random(seed if seed is not None else 0)
    if valid_date is None:
        valid_date = datetime(2020, 1, 1, tzinfo=timezone.utc)

    kp = float(rng.uniform(0.5, 3.5))
    if rng.random() < 0.10:
        kp = float(rng.uniform(4.0, 7.0))
    sw = float(rng.uniform(320.0, 480.0))
    if kp >= 5.0:
        sw = float(rng.uniform(500.0, 800.0))
    log_xray = rng.uniform(-8.0, -6.5)
    if kp >= 5.0:
        log_xray = rng.uniform(-5.5, -4.5)
    xray = float(10.0 ** log_xray)
    regime = derive_helio_regime(kp, xray)

    return BasinContext(
        valid_date=valid_date,
        enso_oni=rng.uniform(-1.5, 1.5),
        iod_dmi=rng.uniform(-1.0, 1.0),
        itcz_latitude_deg=rng.uniform(-10.0, 10.0),
        mslp_regional_hpa=rng.uniform(1005.0, 1020.0),
        solar_wind_speed_kms=sw,
        kp_index=kp,
        goes_xray_flux=xray,
        helio_regime=regime,
        source=DataSource.SYNTHETIC,
    )


# ---------------------------------------------------------------------------
# EpisodeContext
# ---------------------------------------------------------------------------

@dataclass
class EpisodeContext:
    obs:      ZoneObs
    forecast: ForecastResult
    config:   ForecastConfig = field(default_factory=ForecastConfig)

    ground_truth: Optional[RiskScore] = None

    zone_ids:  List[str]            = field(default_factory=list)
    adjacency: Dict[str, List[str]] = field(default_factory=dict)

    data_source: DataSource = DataSource.SYNTHETIC

    basin_context: Optional[BasinContext] = None

    zone_obs: List[ZoneObs] = field(default_factory=list)
    zone_forecasts: List[ForecastResult] = field(default_factory=list)

    def __post_init__(self) -> None:
        if not self.obs.zone_id:
            raise ValueError("EpisodeContext: obs.zone_id is empty")
        if self.obs.zone_id != self.forecast.zone_id:
            raise ValueError(
                f"EpisodeContext: obs.zone_id='{self.obs.zone_id}' != "
                f"forecast.zone_id='{self.forecast.zone_id}'"
            )
        if (
            self.ground_truth is not None
            and self.ground_truth.zone_id != self.obs.zone_id
        ):
            raise ValueError(
                f"EpisodeContext: ground_truth.zone_id='{self.ground_truth.zone_id}'"
                f" != obs.zone_id='{self.obs.zone_id}'"
            )
        if self.obs.zone_id not in self.zone_ids:
            self.zone_ids = [self.obs.zone_id] + list(self.zone_ids)

        # --- Multi-zone list integrity (optional fields) ---
        if self.zone_obs or self.zone_forecasts:
            if len(self.zone_obs) != len(self.zone_forecasts):
                raise ValueError(
                    f"EpisodeContext: len(zone_obs)={len(self.zone_obs)} != "
                    f"len(zone_forecasts)={len(self.zone_forecasts)}"
                )
            if len(self.zone_obs) != len(self.zone_ids):
                raise ValueError(
                    f"EpisodeContext: len(zone_obs)={len(self.zone_obs)} != "
                    f"len(zone_ids)={len(self.zone_ids)}"
                )
            for i, (zo, zf, zid) in enumerate(
                zip(self.zone_obs, self.zone_forecasts, self.zone_ids)
            ):
                if zo.zone_id != zid:
                    raise ValueError(
                        f"EpisodeContext: zone_obs[{i}].zone_id={zo.zone_id!r} "
                        f"!= zone_ids[{i}]={zid!r}"
                    )
                if zf.zone_id != zid:
                    raise ValueError(
                        f"EpisodeContext: zone_forecasts[{i}].zone_id={zf.zone_id!r} "
                        f"!= zone_ids[{i}]={zid!r}"
                    )
                if zo.zone_id != zf.zone_id:
                    raise ValueError(
                        f"EpisodeContext: zone_obs[{i}] / zone_forecasts[{i}] "
                        f"zone_id mismatch"
                    )
            if self.zone_obs[0].zone_id != self.obs.zone_id:
                self.obs = self.zone_obs[0]
                self.forecast = self.zone_forecasts[0]

        for z, neighbours in self.adjacency.items():
            if z not in self.zone_ids:
                raise ValueError(
                    f"EpisodeContext: adjacency key '{z}' not in zone_ids={self.zone_ids}"
                )
            for n in neighbours:
                if n not in self.zone_ids:
                    raise ValueError(
                        f"EpisodeContext: adjacency neighbour '{n}' (of '{z}') "
                        f"not in zone_ids={self.zone_ids}"
                    )

    @property
    def n_zones(self) -> int:
        return len(self.zone_ids)

    def resolved_zone_obs(self) -> List[ZoneObs]:
        if self.zone_obs:
            return list(self.zone_obs)
        return [self.obs]

    def resolved_zone_forecasts(self) -> List[ForecastResult]:
        if self.zone_forecasts:
            return list(self.zone_forecasts)
        return [self.forecast]

    def to_dict(self) -> Dict[str, Any]:
        return {
            "obs":          self.obs.to_dict(),
            "forecast":     self.forecast.to_dict(),
            "config":       self.config.to_dict(),
            "ground_truth": self.ground_truth.to_dict() if self.ground_truth else None,
            "zone_ids":     self.zone_ids,
            "adjacency":    self.adjacency,
            "data_source":  self.data_source.value,
            "basin_context": self.basin_context.to_dict() if self.basin_context else None,
            "zone_obs": [z.to_dict() for z in self.zone_obs],
            "zone_forecasts": [f.to_dict() for f in self.zone_forecasts],
            "_schema_version": SCHEMA_VERSION,
        }

    @classmethod
    def from_dict(cls, d: Dict[str, Any]) -> "EpisodeContext":
        sv, d = _copy_and_pop_schema(d)
        _check_schema(sv, "EpisodeContext")
        zone_obs_raw = d.get("zone_obs") or []
        zone_fc_raw = d.get("zone_forecasts") or []
        return cls(
            obs=ZoneObs.from_dict(d["obs"]),
            forecast=ForecastResult.from_dict(d["forecast"]),
            config=ForecastConfig.from_dict(d["config"]),
            ground_truth=(
                RiskScore.from_dict(d["ground_truth"])
                if d.get("ground_truth") else None
            ),
            zone_ids=d.get("zone_ids", []),
            adjacency=d.get("adjacency", {}),
            data_source=DataSource(d.get("data_source", "synthetic")),
            basin_context=(
                BasinContext.from_dict(d["basin_context"])
                if d.get("basin_context") else None
            ),
            zone_obs=[ZoneObs.from_dict(x) for x in zone_obs_raw],
            zone_forecasts=[ForecastResult.from_dict(x) for x in zone_fc_raw],
        )


# ---------------------------------------------------------------------------
# Synthetic generators
# ---------------------------------------------------------------------------

def make_synthetic_zone_obs(
    zone_id:    str            = "synthetic_zone_0",
    crop_stage: CropStage      = CropStage.GRAIN_FILLING,
    drought:    bool           = False,
    flood:      bool           = False,
    fungi:      bool           = False,
    seed:       Optional[int]  = None,
) -> ZoneObs:
    rng = random.Random(seed if seed is not None else _stable_seed(zone_id))

    _BASE_TIME = datetime(2020, 1, 1, tzinfo=timezone.utc)
    _synthetic_valid_time = _BASE_TIME + timedelta(days=rng.randint(0, 3650))

    base_precip = (
        rng.uniform(60.0, 120.0) if flood
        else rng.uniform(0.0, 2.0) if drought
        else 5.0
    )
    rh = rng.uniform(85.0, 98.0) if fungi else rng.uniform(55.0, 75.0)

    return ZoneObs(
        zone_id=zone_id,
        valid_time=_synthetic_valid_time,
        source=DataSource.SYNTHETIC,
        precip_24h_mm=base_precip,
        precip_7d_mm=base_precip * 6.5,
        precip_14d_mm=base_precip * 12.0,
        precip_30d_mm=base_precip * 24.0,
        precip_anomaly_idx=3.0 if flood else (-2.5 if drought else rng.uniform(-0.5, 0.5)),
        temp_mean_c=rng.uniform(26.0, 32.0),
        temp_max_c=rng.uniform(31.0, 36.0),
        temp_min_c=rng.uniform(22.0, 26.0),
        temp_anomaly_idx=rng.uniform(-0.5, 0.5),
        gdd_accumulated=rng.uniform(400.0, 900.0),
        heat_stress_days=rng.randint(0, 5),
        cold_stress_days=0,
        soil_moisture_pct=rng.uniform(10.0, 25.0) if drought else rng.uniform(40.0, 70.0),
        soil_moisture_anom=-2.0 if drought else rng.uniform(-0.5, 0.5),
        evapotranspiration_mm=rng.uniform(4.0, 7.0),
        wind_speed_max_ms=rng.uniform(3.0, 8.0),
        wind_speed_mean_ms=rng.uniform(1.0, 3.5),
        rh_mean_pct=rh * 0.9,
        rh_max_pct=rh,
        ndvi=rng.uniform(0.35, 0.80),
        ndvi_anomaly_idx=rng.uniform(-0.3, 0.3),
        flood_extent_pct=rng.uniform(20.0, 60.0) if flood else 0.0,
        drainage_risk_idx=rng.uniform(0.5, 0.9) if flood else rng.uniform(0.0, 0.3),
        crop_stage=crop_stage,
        days_to_harvest=rng.randint(7, 45),
        quality_flag=3,
        cloud_cover_pct=rng.uniform(0.0, 20.0),
    )


def make_synthetic_forecast_result(
    zone_id:      str                 = "synthetic_zone_0",
    valid_time:   Optional[datetime]  = None,
    horizon_days: int                 = 30,
    drought:      bool                = False,
    flood:        bool                = False,
    seed:         Optional[int]       = None,
) -> ForecastResult:
    rng = random.Random(
        seed if seed is not None else _stable_seed(zone_id + "_forecast")
    )
    if valid_time is None:
        _BASE_TIME = datetime(2020, 1, 1, tzinfo=timezone.utc)
        t = _BASE_TIME + timedelta(days=rng.randint(0, 3650))
    else:
        t = valid_time

    precip = tuple(
        max(0.0, rng.uniform(30.0, 80.0) if flood
            else rng.uniform(0.0, 3.0) if drought
            else max(0.0, rng.gauss(8.0, 5.0)))
        for _ in range(horizon_days)
    )
    temp = tuple(rng.uniform(26.0, 32.0) for _ in range(horizon_days))
    rh   = tuple(rng.uniform(60.0, 90.0) for _ in range(horizon_days))
    p10  = tuple(max(0.0, p * rng.uniform(0.3, 0.7)) for p in precip)
    p90  = tuple(p * rng.uniform(1.3, 2.0) for p in precip)

    prob_rain = tuple(
        float(_clip(p / 60.0 + rng.uniform(-0.05, 0.05), 0.0, 1.0))
        for p in precip
    )
    prob_drought = tuple(
        float(_clip(0.8 if drought else rng.uniform(0.0, 0.15), 0.0, 1.0))
        for _ in range(horizon_days)
    )
    prob_humid = tuple(
        float(_clip((r - 70.0) / 30.0 + rng.uniform(-0.05, 0.05), 0.0, 1.0))
        for r in rh
    )

    return ForecastResult(
        zone_id=zone_id,
        forecast_time=t,
        horizon_days=horizon_days,
        precip_mm=precip,
        temp_mean_c=temp,
        rh_mean_pct=rh,
        precip_p10=p10,
        precip_p90=p90,
        temp_p10=tuple(v - rng.uniform(1.0, 3.0) for v in temp),
        temp_p90=tuple(v + rng.uniform(1.0, 3.0) for v in temp),
        prob_heavy_rain=prob_rain,
        prob_drought_day=prob_drought,
        prob_high_humidity=prob_humid,
        source=DataSource.SYNTHETIC,
    )


def make_synthetic_episode_context(
    zone_id: str                 = "synthetic_zone_0",
    config:  Optional[ForecastConfig] = None,
    drought: bool                = False,
    flood:   bool                = False,
    fungi:   bool                = False,
    seed:    Optional[int]       = None,
) -> EpisodeContext:
    cfg  = config or ForecastConfig()
    obs  = make_synthetic_zone_obs(zone_id, drought=drought, flood=flood,
                                   fungi=fungi, seed=seed)
    fcast = make_synthetic_forecast_result(zone_id, valid_time=obs.valid_time,
                                           drought=drought, flood=flood, seed=seed)
    basin = (
        make_synthetic_basin_context(valid_date=obs.valid_time, seed=seed)
        if cfg.include_basin_context else None
    )
    return EpisodeContext(
        obs=obs,
        forecast=fcast,
        config=cfg,
        zone_ids=[zone_id],
        data_source=DataSource.SYNTHETIC,
        basin_context=basin,
    )


# ---------------------------------------------------------------------------
# Self-test  (python zone_observation.py)
# ---------------------------------------------------------------------------

if __name__ == "__main__":
    import sys
    logging.basicConfig(level=logging.WARNING)
    print(f"zone_observation.py  schema_version={SCHEMA_VERSION}\n")

    failures: List[str] = []

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

    # 1. ZoneObs round-trip + non-mutation
    obs = make_synthetic_zone_obs("test_flood", flood=True, seed=42)
    d = obs.to_dict()
    had_sv = "_schema_version" in d
    obs2 = ZoneObs.from_dict(d)
    still_has_sv = "_schema_version" in d
    _assert(had_sv and still_has_sv,   "ZoneObs.from_dict mutated caller dict")
    _assert(obs.zone_id == obs2.zone_id,          "ZoneObs zone_id round-trip")
    _assert(abs(obs.precip_24h_mm - obs2.precip_24h_mm) < 1e-9, "ZoneObs precip precision")
    _assert(obs.crop_stage == obs2.crop_stage,    "ZoneObs crop_stage round-trip")
    _assert(obs.source == obs2.source,            "ZoneObs source round-trip")
    print(f"  ZoneObs  flood={obs.flood_signal():.3f}  drought={obs.drought_signal():.3f}"
          f"  fungi={obs.fungi_risk_signal():.3f}  composite={obs.composite_risk():.3f}")

    # 2. Deterministic seeding
    a = make_synthetic_zone_obs("stable", seed=99)
    b = make_synthetic_zone_obs("stable", seed=99)
    _assert(a.precip_24h_mm == b.precip_24h_mm, "Explicit seed not deterministic")
    c = make_synthetic_zone_obs("crc_zone")
    d2 = make_synthetic_zone_obs("crc_zone")
    _assert(c.precip_24h_mm == d2.precip_24h_mm, "zlib.crc32 seed not stable")
    print("  Deterministic seeding OK")

    # 3. ZoneObs.validate()
    obs_v = make_synthetic_zone_obs("val_zone", seed=1)
    obs_v.precip_7d_mm = obs_v.precip_14d_mm + 50.0
    issues = ZoneObs.validate(obs_v)
    _assert(len(issues) > 0, "validate() missed precip_14d < precip_7d")
    print(f"  ZoneObs.validate() caught {len(issues)} issue(s)")

    # 4. GeoPolygon string-vertex coercion + contains_point
    poly = GeoPolygon(
        vertices=[("3.0", "101.0"), (3.1, 101.0), (3.1, 101.1), (3.0, 101.1)],
        zone_id="sel_A1",
    )
    _assert(isinstance(poly.centroid[0], float), "GeoPolygon centroid not float")
    _assert(poly.contains_point(3.05, 101.05),   "GeoPolygon inside point")
    _assert(not poly.contains_point(4.0, 102.0), "GeoPolygon outside point")
    poly2 = GeoPolygon.from_dict(poly.to_dict())
    _assert(poly.zone_id == poly2.zone_id, "GeoPolygon round-trip")
    print(f"  GeoPolygon  area={poly.approx_area_km2:.1f} km2  contains_point OK")

    # 5. ForecastResult round-trip + prob validation + non-mutation
    fr = make_synthetic_forecast_result("test_flood", flood=True, seed=42)
    d_fr = fr.to_dict()
    had_sv_fr = "_schema_version" in d_fr
    fr2 = ForecastResult.from_dict(d_fr)
    _assert(had_sv_fr and "_schema_version" in d_fr, "ForecastResult.from_dict mutated dict")
    _assert(fr.precip_mm == fr2.precip_mm,  "ForecastResult precip round-trip")
    _assert(fr.source    == fr2.source,     "ForecastResult source round-trip")
    try:
        ForecastResult(
            zone_id="x", forecast_time=datetime.now(tz=timezone.utc),
            precip_mm=tuple([0.0]*30), temp_mean_c=tuple([29.0]*30),
            rh_mean_pct=tuple([70.0]*30), precip_p10=tuple([0.0]*30),
            precip_p90=tuple([1.0]*30), prob_heavy_rain=tuple([5.0]*30),
            prob_drought_day=tuple([0.0]*30), prob_high_humidity=tuple([0.0]*30),
        )
        _assert(False, "ForecastResult accepted prob > 1.0")
    except ValueError:
        pass
    print(f"  ForecastResult  peak_day={fr.peak_precip_day()}"
          f"  cumul14d={fr.cumulative_precip_mm(14):.1f}mm  prob_validation OK")

    # 6. RiskScore round-trip + ordering + harvest_window_days
    now = datetime.now(tz=timezone.utc)
    rs = RiskScore(
        zone_id="test_flood", scored_at=now,
        supply_shortfall_prob=0.35, drought_risk=0.05, flood_risk=0.78,
        supply_risk_composite=0.55, fungi_contamination_prob=0.42,
        harvest_delay_days=6.0, quality_risk_composite=0.42,
        optimal_harvest_window_start=now + timedelta(days=14),
        optimal_harvest_window_end=now + timedelta(days=21),
        alert_level=AlertLevel.WARNING, confidence=0.80,
    )
    d_rs = rs.to_dict()
    rs2 = RiskScore.from_dict(d_rs)
    _assert("_schema_version" in d_rs,         "RiskScore.from_dict mutated dict")
    _assert(rs.alert_level == rs2.alert_level, "RiskScore alert_level round-trip")
    _assert(rs.harvest_window_days() == 7,     "RiskScore harvest_window_days")
    _assert(rs.is_actionable(),                "RiskScore.is_actionable() for WARNING")
    _assert(AlertLevel.WARNING > AlertLevel.WATCH,   "AlertLevel ordering >")
    _assert(AlertLevel.NONE    < AlertLevel.CRITICAL, "AlertLevel ordering <")
    print(f"  RiskScore  alert={rs.alert_level.value}  window={rs.harvest_window_days()}d"
          f"  actionable={rs.is_actionable()}")

    # 7. ForecastConfig rational threshold + new pipeline fields round-trip
    cfg = ForecastConfig()
    _assert(cfg.belief_floor < cfg.rational_termination_threshold,
            "Default ForecastConfig: belief_floor >= rational_threshold")
    cfg2 = ForecastConfig.from_dict(cfg.to_dict())
    _assert(cfg.alert_value == cfg2.alert_value, "ForecastConfig round-trip")
    _assert(cfg2.real_data_ratio == 0.7,         "ForecastConfig real_data_ratio round-trip")
    _assert(cfg2.era5_ratio == 0.5,              "ForecastConfig era5_ratio round-trip")
    _assert(cfg2.force_data_source is None,      "ForecastConfig force_data_source round-trip")
    _assert(cfg2.inject_noise is False,          "ForecastConfig inject_noise round-trip")
    _assert(cfg2.noise_scale == 0.05,            "ForecastConfig noise_scale round-trip")
    cfg_era5 = ForecastConfig(force_data_source=DataSource.ERA5_REANALYSIS)
    cfg_era5_back = ForecastConfig.from_dict(cfg_era5.to_dict())
    _assert(
        cfg_era5_back.force_data_source == DataSource.ERA5_REANALYSIS,
        "ForecastConfig force_data_source=ERA5 round-trip"
    )
    cfg_new = ForecastConfig(
        forecast_backend="openmeteo",
        use_climatology_anomalies=True,
        climatology_years=15,
    )
    cfg_new_back = ForecastConfig.from_dict(cfg_new.to_dict())
    _assert(cfg_new_back.forecast_backend == "openmeteo",
            "forecast_backend round-trip")
    _assert(cfg_new_back.use_climatology_anomalies is True,
            "use_climatology_anomalies round-trip")
    _assert(cfg_new_back.climatology_years == 15,
            "climatology_years round-trip")
    _assert(cfg2.forecast_backend == "synthetic",
            "forecast_backend default should be 'synthetic' (back-compat)")
    try:
        ForecastConfig(forecast_backend="not_a_backend")
        _assert(False, "ForecastConfig accepted invalid forecast_backend")
    except ValueError:
        pass

    _assert(cfg2.belief_prior_weight == 0.70,       "belief_prior_weight default round-trip")
    _assert(cfg2.uncertainty_decay == 0.70,         "uncertainty_decay default round-trip")
    _assert(cfg2.info_gain_scale == 5.0,            "info_gain_scale default round-trip")
    _assert(cfg2.uncertainty_penalty_scale == 5.0,  "uncertainty_penalty_scale default round-trip")
    cfg_belief = ForecastConfig(
        belief_prior_weight=0.35,
        uncertainty_decay=0.5,
        info_gain_scale=2.0,
        uncertainty_penalty_scale=8.0,
    )
    cfg_belief_back = ForecastConfig.from_dict(cfg_belief.to_dict())
    _assert(cfg_belief_back.belief_prior_weight == 0.35,
            "non-default belief_prior_weight round-trip")
    _assert(cfg_belief_back.uncertainty_decay == 0.5,
            "non-default uncertainty_decay round-trip")
    _assert(cfg_belief_back.info_gain_scale == 2.0,
            "non-default info_gain_scale round-trip")
    _assert(cfg_belief_back.uncertainty_penalty_scale == 8.0,
            "non-default uncertainty_penalty_scale round-trip")
    try:
        ForecastConfig(belief_prior_weight=1.5)
        _assert(False, "ForecastConfig accepted belief_prior_weight out of [0,1]")
    except ValueError:
        pass
    try:
        ForecastConfig(info_gain_scale=-1.0)
        _assert(False, "ForecastConfig accepted negative info_gain_scale")
    except ValueError:
        pass
    print(f"  ForecastConfig  rational_threshold={cfg.rational_termination_threshold:.4f}"
          f"  belief_floor={cfg.belief_floor:.4f}  pipeline fields OK")

    # 8. EpisodeContext round-trip + validation
    ec = make_synthetic_episode_context("zone_A", seed=7)
    d_ec = ec.to_dict()
    ec2 = EpisodeContext.from_dict(d_ec)
    _assert(ec.obs.zone_id == ec2.obs.zone_id,        "EpisodeContext zone_id round-trip")
    _assert(ec.config.alert_value == ec2.config.alert_value, "EpisodeContext config round-trip")
    _assert(ec.n_zones == 1,                          "EpisodeContext n_zones")
    try:
        EpisodeContext(
            obs=make_synthetic_zone_obs("zone_A"),
            forecast=make_synthetic_forecast_result("zone_B"),
            config=ForecastConfig(),
        )
        _assert(False, "EpisodeContext accepted zone_id mismatch")
    except ValueError:
        pass
    print(f"  EpisodeContext  n_zones={ec.n_zones}  zone_mismatch_check OK")

    # 9. Full JSON round-trip
    ec_json = json.dumps(ec.to_dict())
    ec_back = EpisodeContext.from_dict(json.loads(ec_json))
    _assert(ec.obs.zone_id == ec_back.obs.zone_id,
            "EpisodeContext JSON zone_id round-trip")
    _assert(ec.forecast.precip_mm == ec_back.forecast.precip_mm,
            "ForecastResult precip JSON round-trip")
    print("  Full JSON serialisation round-trip OK")

    # 10. BasinContext round-trip + clipping + helio + EpisodeContext integration
    bc = make_synthetic_basin_context(seed=3)
    d_bc = bc.to_dict()
    bc2 = BasinContext.from_dict(d_bc)
    _assert("_schema_version" in d_bc,              "BasinContext.from_dict mutated dict")
    _assert(abs(bc.enso_oni - bc2.enso_oni) < 1e-9,  "BasinContext enso_oni round-trip")
    _assert(abs(bc.iod_dmi - bc2.iod_dmi) < 1e-9,    "BasinContext iod_dmi round-trip")
    _assert(bc.source == bc2.source,                 "BasinContext source round-trip")
    _assert(abs(bc.kp_index - bc2.kp_index) < 1e-9,  "BasinContext kp_index round-trip")
    _assert(bc.helio_regime == bc2.helio_regime,     "BasinContext helio_regime round-trip")
    _assert(bc.helio_regime in ("quiet", "active", "storm"),
            f"invalid helio_regime {bc.helio_regime!r}")
    bc_extreme = BasinContext(valid_date=now, enso_oni=99.0, mslp_regional_hpa=1.0)
    _assert(bc_extreme.enso_oni <= 5.0,              "BasinContext enso_oni not clipped")
    _assert(bc_extreme.mslp_regional_hpa >= 900.0,   "BasinContext mslp_regional_hpa not clipped")
    _assert(bc_extreme.kp_index == 2.0,              "BasinContext kp default should be quiet-Sun 2.0")
    _assert(bc_extreme.helio_regime == "quiet",      "BasinContext helio default should be quiet")
    _assert(derive_helio_regime(6.0, 1e-7) == "storm", "derive_helio_regime storm by Kp")
    _assert(derive_helio_regime(1.0, 2e-5) == "storm", "derive_helio_regime storm by X-ray")
    _assert(derive_helio_regime(3.5, 1e-7) == "active", "derive_helio_regime active")
    _assert(derive_helio_regime(1.0, 1e-8) == "quiet", "derive_helio_regime quiet")

    cfg_basin = ForecastConfig(include_basin_context=True)
    _assert(cfg_basin.require_real_basin_context is False,
            "require_real_basin_context should default False")
    ec_basin = make_synthetic_episode_context("zone_basin", config=cfg_basin, seed=11)
    _assert(ec_basin.basin_context is not None,
            "make_synthetic_episode_context did not attach basin_context when opted in")
    d_ec_basin = ec_basin.to_dict()
    ec_basin2 = EpisodeContext.from_dict(d_ec_basin)
    _assert(ec_basin2.basin_context is not None,
            "EpisodeContext.basin_context lost in round-trip")
    _assert(
        abs(ec_basin.basin_context.enso_oni - ec_basin2.basin_context.enso_oni) < 1e-9,
        "EpisodeContext.basin_context.enso_oni round-trip"
    )
    _assert(
        ec_basin.basin_context.helio_regime == ec_basin2.basin_context.helio_regime,
        "EpisodeContext.basin_context.helio_regime round-trip"
    )
    ec_no_basin = make_synthetic_episode_context("zone_no_basin", seed=11)
    _assert(ec_no_basin.basin_context is None,
            "basin_context should default to None when include_basin_context=False")
    print(f"  BasinContext  oni={bc.enso_oni:.2f}  dmi={bc.iod_dmi:.2f}  "
          f"kp={bc.kp_index:.1f}  regime={bc.helio_regime}  "
          f"round-trip OK, EpisodeContext integration OK")

    # 11. New optional ZoneObs satellite fields: None-by-default, clipping, round-trip
    obs_sat = ZoneObs(
        zone_id="sat_zone", valid_time=now,
        soil_moisture_satellite_pct=150.0,   # out of range -> should clip to 100
        precip_satellite_mm=12.5,
    )
    _assert(obs_sat.soil_moisture_satellite_pct == 100.0,
            "soil_moisture_satellite_pct not clipped to 100")
    _assert(obs_sat.precip_satellite_mm == 12.5,
            "precip_satellite_mm unexpectedly altered")
    obs_plain = make_synthetic_zone_obs("plain_zone", seed=5)
    _assert(obs_plain.precip_satellite_mm is None,
            "precip_satellite_mm should default to None, not 0.0")
    _assert(obs_plain.soil_moisture_satellite_pct is None,
            "soil_moisture_satellite_pct should default to None, not 0.0")
    d_sat = obs_sat.to_dict()
    obs_sat2 = ZoneObs.from_dict(d_sat)
    _assert(obs_sat2.precip_satellite_mm == obs_sat.precip_satellite_mm,
            "precip_satellite_mm round-trip")
    print("  ZoneObs satellite fields: None-default, clipping, round-trip OK")

    print()
    if failures:
        print(f"FAILED  {len(failures)} test(s):")
        for f in failures:
            print(f"  - {f}")
        sys.exit(1)
    else:
        print(f"All {11} test groups passed.")