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
# ---------------------------------------------------------------------------
# v3 migration note:
#   - ZoneObs gained two optional per-zone fields: precip_satellite_mm,
#     soil_moisture_satellite_pct (both Optional[float], default None —
#     None means "no satellite pass available", NOT zero).
#   - New BasinContext dataclass added (episode-level, NOT per-zone):
#     enso_oni, iod_dmi, itcz_latitude_deg, mslp_regional_hpa.
#     EpisodeContext gained an optional `basin_context` field.
#   - DataSource gained SATELLITE_PRECIP, SATELLITE_SOIL, PUBLISHED_INDEX.
#   - BasinContext later gained optional helio / space-weather fields
#     (solar_wind_speed_kms, kp_index, goes_xray_flux, helio_regime) with
#     quiet-Sun defaults so old records deserialise without a schema bump.
#   - ForecastConfig gained require_real_basin_context for strict real-only
#     basin/helio ingest (enforced in era5_data_pipeline.fetch_basin_context).
#   Old v2 records deserialise fine via ZoneObs.from_dict/EpisodeContext.from_dict
#   as long as the caller bumps stored "_schema_version" to 3 first (the new
#   fields all have defaults, so no other migration is required).
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):
    """Risk alert level emitted by the RL policy's terminate action."""
    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):
    """Provenance tag so downstream code can weight or distrust readings.

    Migration note (schema v2):
        DataSource.EDGE_NODE serialises as "edge_node".
        Any records previously serialised as "gnus_node" must be migrated:
            d["source"] = "edge_node"  # was "gnus_node"
    """
    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:
        """True for real-world sources (not synthetic or unknown)."""
        return self not in (DataSource.SYNTHETIC, DataSource.UNKNOWN)


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

def _clip(value: float, lo: float, hi: float) -> float:
    """stdlib-only clip. Avoids numpy dependency at this layer."""
    return max(lo, min(hi, value))


def _stable_seed(key: str) -> int:
    """Deterministic integer seed from a string key.

    Uses zlib.crc32 rather than hash() -- hash() is randomised per process
    by PYTHONHASHSEED and would produce different synthetic obs for the same
    zone_id across runs, breaking training reproducibility.
    """
    return zlib.crc32(key.encode("utf-8")) & 0x7FFFFFFF


def _copy_and_pop_schema(d: Dict[str, Any]) -> Tuple[Optional[int], Dict[str, Any]]:
    """Return (schema_version, clean_copy) without mutating the input dict.

    All from_dict() methods call this instead of d.pop() directly, which
    was the v1 caller-mutation bug.
    """
    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:
    """Raise ValueError loudly on schema version mismatch."""
    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:
    """
    Validated configuration for the weather forecasting RL environment.

    Mirrors InspectionConfig from the MEMS env: every economic parameter
    is validated in __post_init__, and the rational termination threshold
    is computed and checked explicitly so misconfiguration is loud.

    Rational termination condition (mirrors MEMS belief_floor logic):
      Issue alert when:
        P(risk_event) > false_alert_penalty / (alert_value + false_alert_penalty)
      With defaults: P > 20 / (100 + 20) = 0.1667

      belief_floor MUST be set below this threshold or early termination
      can never be EV-positive.
    """

    # --- 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

    # --- 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)

    # --- Exploration shaping ---
    # FIX (premature termination): ep_len_mean stayed at ~1.3-1.6 for the
    # entire 'normal' (n_zones=2) curriculum phase across 1M steps -- not
    # slow convergence, the actual argmax of the reward as specified.
    # Root cause: belief_map is already ~30%-informed by the true signal at
    # reset() (see _per_zone_beliefs' 0.7*prior + 0.3*signal blend), so the
    # terminate action's EV barely depends on whether the agent inspected
    # anything. Meanwhile a single inspect step costs inspection_cost=1.0
    # raw, while its info_gain term (5.0 * info_gain in
    # _compute_zone_refinement_reward) rarely exceeds ~0.2-0.4 raw for
    # realistic belief deltas, and the uncertainty-decay payoff at
    # termination is a mean over ALL zones, diluting the effect of
    # inspecting any single one. Net EV of inspecting was negative under
    # almost all conditions.
    #
    # These two fields close that gap directly, without touching the core
    # alert/false-alarm economics that crop_risk_scorer.py's threshold
    # logic depends on:
    #   zone_visit_bonus: added to the raw inspection reward on a zone's
    #     FIRST visit only (the existing revisit_penalty path is untouched,
    #     so repeat visits are still never profitable). Set > inspection_cost
    #     so a first visit is always non-negative in expectation, independent
    #     of info_gain -- exploration is never structurally punished.
    #   unvisited_zone_penalty: subtracted from the termination reward,
    #     scaled by the fraction of active zones never visited this episode.
    #     Terminating having inspected nothing costs the full penalty;
    #     terminating after visiting every zone costs nothing. Uses the same
    #     /alert_value normalization as every other termination-reward term,
    #     so it composes on the same scale as gain/cost/unc.
    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

    # --- Robustness training ---
    economic_randomization: bool  = False
    clean_episode_ratio:    float = 0.7
    # Spatial correlation of synthetic hazard events across zones in one
    # episode. 0 = independent per-zone draws (legacy); 1 = all-or-nothing
    # regional event (El Niño-style). Real Java drought/flood is highly
    # correlated; default 0.85 so multi-zone allocation sees joint events.
    event_spatial_correlation: float = 0.85
    # If True (default), training reset permutes zone slot order each episode
    # so action index i is not permanently tied to zone_i / a fixed name.
    # Prevents policies from learning "always inspect slot 1" by index alone.
    # Real-eval injection path does not shuffle (order is caller-defined).
    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 ---
    soft_reset:  bool = True

    # --- Seeding ---
    seed: Optional[int] = None

    # --- Data pipeline (era5_data_pipeline.py) ---
    # These fields are consumed by fetch_zone_obs() to control source selection
    # and noise injection. They must live here so user-set values are respected;
    # era5_data_pipeline.py previously fell back to getattr() defaults, silently
    # ignoring any values set on a ForecastConfig instance.
    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)

    # --- Satellite sources (schema v3+) ---
    # Opt-in and default False: existing callers get byte-identical behaviour
    # unless they explicitly enable these. When enabled, era5_data_pipeline.py
    # tries the satellite source first (via Google Earth Engine) and falls
    # back through the existing ERA5 -> Open-Meteo -> synthetic chain on
    # failure, exactly like every other source in fetch_zone_obs().
    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
    # Strict real-only basin context: when True, fetch_basin_context() refuses
    # synthetic / quiet-default fallbacks and raises RuntimeError if any of
    # RONI, DMI, or the required SWPC helio fields cannot be obtained from
    # live published sources. Default False preserves the soft-fallback path
    # used by training and offline runs. Implies include_basin_context in
    # spirit (callers should set both), but the fetch function itself is the
    # enforcement point.
    require_real_basin_context: bool = False

    # --- Forecast backend (timesfm_wrapper.py integration) ---
    # Which ForecastBackend fetch_episode_context() uses to build the
    # EpisodeContext's ForecastResult. Default 'synthetic' reproduces the
    # pre-integration behaviour exactly (deterministic synthetic forecast),
    # so existing callers see no change unless they opt in.
    #   'synthetic'  -- deterministic make_synthetic_forecast_result (offline)
    #   'baseline'   -- deterministic statistical forecast from ZoneObs fields
    #                   (timesfm_wrapper.BaselineBackend; offline, no deps)
    #   'openmeteo'  -- real Open-Meteo API forecast (needs network; falls
    #                   back to synthetic on any failure)
    #   'timesfm'    -- local TimesFM checkpoint (needs the checkpoint file +
    #                   timesfm package; see timesfm_wrapper.LocalTimesFMBackend)
    forecast_backend: str = "synthetic"

    # --- Climatology anomalies (climatology.py integration) ---
    # Opt-in, default False for byte-identical back-compat. When True,
    # era5_data_pipeline.fetch_zone_obs() post-processes every fetched
    # ZoneObs through climatology.apply_climatology_anomalies(), which
    # populates precip_anomaly_idx / temp_anomaly_idx / soil_moisture_anom
    # as z-scores against a per-zone day-of-year climatology. Without this,
    # ALL real fetchers leave those fields at 0.0, which zeroes out most of
    # ZoneObs.drought_signal()/flood_signal() -- i.e. real observations are
    # barely scoreable. Synthetic training episodes deliberately leave this
    # off (their anomaly fields are injected directly by the event flags).
    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 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.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)
        )

        # Threshold is a 3-way EV break-even (alert vs. no-alert at believed
        # probability p), not just alert_value vs false_alert_penalty:
        #   p > false_alert_penalty / (alert_value + false_alert_penalty + miss_penalty)
        # Single source of truth: crop_risk_scorer._alert_level and
        # weather_forecast_env._compute_termination_reward both read/derive
        # from this instead of keeping their own copy of the formula.
        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:
        """P(event) above which issuing an alert has positive expected value,
        accounting for alert_value, false_alert_penalty, AND miss_penalty
        (3-way expected-value breakeven -- see the derivation in
        __post_init__). Single source of truth: crop_risk_scorer._alert_level
        and weather_forecast_env._compute_termination_reward both read this
        property rather than recomputing the formula.
        """
        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,
            "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,
            "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:
    """Lightweight polygon -- no shapely dependency at this layer.

    Vertices are (lat, lon) pairs in decimal degrees, WGS-84.
    Vertex coordinates are coerced to float at construction so
    JSON-deserialized strings ('3.0') do not silently fail later.
    """
    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]:
        """Arithmetic centroid (lat, lon)."""
        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:
        """Shoelace formula on a flat-earth approximation.
        Accurate to ~1% for zones < 200 km across.
        """
        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:
        """Ray-casting point-in-polygon test (WGS-84 lat/lon)."""
        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:
    """
    One observation snapshot for one sourcing zone at one point in time.

    The single object that crosses all layer boundaries.
    The pipeline produces it. The env consumes it. The model reads it.
    The scorer reads it. The transport serialises it.

    Field naming convention:
      *_mm    = millimetres
      *_c     = degrees Celsius
      *_pct   = percentage 0-100
      *_idx   = dimensionless index (standardised anomaly or ratio)
      *_prob  = probability 0.0-1.0
      *_days  = duration in days
      *_ms    = metres per second
      *_anom  = z-score anomaly vs climatological mean
    """

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

    # --- Precipitation ---
    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

    # --- Temperature ---
    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

    # --- Heat stress ---
    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 ---
    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

    # --- Direct satellite retrievals (schema v3+, per-zone, optional) ---
    # These are raw provenance/audit fields, NOT yet consumed by
    # crop_risk_scorer.py's drought_signal()/flood_signal(). Wiring them in
    # requires a per-zone climatological baseline (mean/std) to convert a raw
    # satellite reading into an anomaly comparable to precip_anomaly_idx /
    # soil_moisture_anom -- that baseline does not exist anywhere in this
    # pipeline yet. Do not fabricate one; compute it from a real climatology
    # (e.g. multi-year ERA5 or IMERG mean) before using these fields for risk
    # scoring. None means "no satellite pass available", NOT zero.
    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 ---
    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)

    # --- Humidity ---
    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
    # FIX (fungi/regime-insensitivity, verified on a 3,623-point real-Indonesia
    # backtest): fungi_risk_signal() used rh_max_pct against a fixed absolute
    # threshold. In a tropical maritime climate rh_max_pct sits near
    # saturation (90-100%) on most days regardless of ENSO phase, so the
    # signal was flat across El Nino/La Nina and won the alert_level max()
    # in 73.6% of points -- masking the correctly regime-sensitive
    # drought/flood signals in the system's actual output (alert_level).
    # z-score vs climatological mean, built from rh_mean_pct (NOT rh_max_pct
    # -- Open-Meteo's archive API documents max-RH aggregation as not
    # reliably available across models, while mean-RH is; daily max and mean
    # RH anomalies are highly correlated, so mean's anomaly is used as the
    # proxy correction applied to the max-based absolute term). Default 0.0
    # is a true no-op in fungi_risk_signal() (additive, clipped, zero when
    # unset) -- synthetic data and any caller that never populates this are
    # byte-identical to the pre-fix formula. Populated for real data by
    # climatology.apply_climatology_anomalies().
    rh_anomaly_idx:      float = 0.0

    # --- Vegetation (optional -- populated when satellite pass available) ---
    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 / waterlogging ---
    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 context ---
    crop_stage:         CropStage = CropStage.UNKNOWN
    days_to_harvest:    Optional[int] = None   # None = unknown; 0 = harvest now
    planting_date:      Optional[datetime] = None

    # --- Observation quality ---
    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

    # --- Extensibility hatch ---
    extras: Dict[str, Any] = field(default_factory=dict)
    # See KNOWN_EXTRAS at module level for documented keys.

    def __post_init__(self) -> None:
        # --- Defensive copy: prevent caller mutation of extras dict ---
        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))

        # --- Anomaly clipping: prevent z-score explosions from destabilising RL ---
        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)

        # --- Monotonic precipitation aggregates sanity check ---
        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]:
        """Return a list of validation warnings (empty = clean).

        Called by era5_data_pipeline.py before inserting obs into the env.
        strict=True raises ValueError instead of returning issues.
        """
        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]:
        """Lossless serialisation to JSON-safe dict. Never mutates self."""
        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":
        """Reconstruct from to_dict(). Never mutates the caller's dict."""
        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)

    # --- Diagnostic helpers ---

    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:
        """Composite drought signal [0, 1]. Heuristic for diagnostics only."""
        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:
        """Composite flood signal [0, 1]."""
        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:
        """Composite fungi / post-harvest moisture risk [0, 1].

        FIX (fungi/regime-insensitivity -- see rh_anomaly_idx's field
        comment for the verified root cause): the pure absolute-threshold
        term below (rh_s) is UNCHANGED -- it stays the sole driver whenever
        rh_anomaly_idx is at its 0.0 default, which is every synthetic
        ZoneObs and any real-data caller that hasn't run climatology
        anomaly application. This is a true no-op for existing behavior,
        not a reweighting: anomaly_adj is additive and clipped to +/-0.3,
        so it is exactly 0 when rh_anomaly_idx is exactly 0.

        When populated (real data via climatology.apply_climatology_anomalies),
        anomaly_adj pulls the absolute term down during anomalously DRY
        periods (e.g. El Nino, even though rh_max_pct often stays nominally
        high in absolute terms in a tropical climate) and up during
        anomalously WET periods (La Nina) -- restoring the regime
        sensitivity the pure absolute threshold structurally couldn't have.
        The absolute floor is kept deliberately: fungal risk still needs a
        real minimum humidity regardless of how anomalous conditions are.
        """
        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:
        """Single-number composite risk [0, 1] for belief map initialisation."""
        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:
    """
    Probabilistic forecast for one zone over horizon_days.

    Frozen: produced by timesfm_wrapper.py, never mutated downstream.
    All sequences are horizon_days elements (default 30).
    """
    zone_id:       str
    forecast_time: datetime
    horizon_days:  int = 30

    # Point forecasts (median)
    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)

    # Uncertainty (p10 to p90)
    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)

    # Exceedance probabilities, validated to [0, 1]
    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:
        # --- Defensive copy: prevent caller mutation of extras dict ---
        object.__setattr__(self, "extras", dict(self.extras))

        # All non-empty sequences must be the same length
        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} }"
                )
            # Enforce that all sequences match horizon_days exactly
            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."
                    )

        # Probability values must be in [0, 1]
        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."
                    )

        # p10 <= p90
        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:
        """Mean probability of exceedance over window_days (default: full horizon)."""
        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:
    """
    Composite risk assessment for one sourcing zone.

    Frozen: terminal output of crop_risk_scorer.py. Never mutated.
    All scalar scores in [0.0, 1.0] unless noted.
    """
    zone_id:   str
    scored_at: datetime

    # Supply risk
    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]

    # Quality risk
    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]

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

    # Decision output
    alert_level:  AlertLevel = AlertLevel.NONE
    action_notes: str = ""

    # Confidence
    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:
        # --- Defensive copy: prevent caller mutation of extras dict ---
        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)

        # RiskScore is fully internal (produced by crop_risk_scorer.py only),
        # so we enforce timezone-awareness strictly here rather than silently fixing.
        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:
        """
        Elevated attention: ADVISORY and above (alert_level > WATCH).

        Historical name — prefer is_elevated() / is_product_actionable() for
        new code. Product bus emission uses product_alert_service gate
        (WARNING+ or hazard thresholds), not this method alone.
        """
        return self.alert_level > AlertLevel.WATCH

    def is_elevated(self) -> bool:
        """Internal / elevated attention: ADVISORY and above."""
        return self.alert_level.severity() >= AlertLevel.ADVISORY.severity()

    def is_product_actionable(self) -> bool:
        """
        Client-facing product floor on AlertLevel alone: WARNING and above.

        Hazard-threshold override (drought/flood freeze bars) is applied in
        product_alert_service.is_product_actionable(score, gate), not here.
        """
        return self.alert_level.severity() >= AlertLevel.WARNING.severity()

    def harvest_window_days(self) -> Optional[int]:
        """Duration of optimal harvest window in days. None if not set."""
        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:
    """Map Kp + GOES X-ray flux to a coarse helio regime label.

    Thresholds chosen so the *default / typical* state is ``quiet`` and
    storms are rare — the same anti-saturation principle used for the
    humidity anomaly fix. These labels are context features only; they do
    not directly drive crop-risk scores until a later policy/scorer change
    explicitly consumes them.

    Rules (NOAA-style G-scale / GOES class approximation):
      storm  — Kp >= 5 (G1+) OR X-ray >= 1e-5 (M-class+)
      active — Kp >= 3 (unsettled) OR X-ray >= 1e-6 (C-class+)
      quiet  — otherwise
    """
    if kp_index >= 5.0 or goes_xray_flux >= 1e-5:
        return "storm"
    if kp_index >= 3.0 or goes_xray_flux >= 1e-6:
        return "active"
    return "quiet"


@dataclass(frozen=True)
class BasinContext:
    """
    Basin-scale / teleconnection context for one episode.

    Deliberately NOT a ZoneObs field. ENSO, IOD, ITCZ position, monsoon
    pressure systems, and heliophysics indices are not zone-specific --
    they are the same value for every zone in an episode. Storing them
    per-zone would duplicate an identical scalar across every zone and
    risk the copies drifting apart across a training run. One BasinContext
    per EpisodeContext, consumed as a distinct input channel (not folded
    into per-zone belief/forecast arrays) by weather_forecast_env.py and
    gru_weather_policy.py.

    Field bounds are generous (not tight physical limits) since these are
    RL policy inputs -- clipping only guards against clearly corrupt values,
    not against genuine extremes (e.g. an ONI of 2.8 during a strong
    El Nino is real, not an error).

    Helio fields (solar_wind_speed_kms, kp_index, goes_xray_flux,
    helio_regime) default to quiet-Sun values so missing / synthetic /
    offline paths never inject artificial "storm" context the way the old
    absolute humidity term saturated fungi risk. They ride along as
    optional context; the crop-risk scorer does not consume them until an
    explicit later integration step.
    """
    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"))
        # Helio fields optional for backward compatibility with records
        # serialised before they existed — dataclass defaults apply when
        # keys are absent.
        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:
    """Deterministic synthetic BasinContext for tests and default episodes.

    Helio draws are deliberately *quiet-biased* (Kp mostly 0.5–3.5, X-ray
    background-to-low-C) so synthetic training trajectories do not inject
    a constant "active/storm" prior the way the old absolute humidity term
    saturated fungi risk. Occasional higher draws still occur so the
    policy can see regime transitions.
    """
    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:
    """
    Typed episode initialiser for weather_forecast_env.py.

    Analogous to the wafer_data dict in the MEMS env, but fully typed
    and validated. The pipeline produces this; the env consumes it at
    _soft_reset(). Economic parameters live in config (ForecastConfig)
    rather than as loose floats, so the rational termination threshold
    is always validated.

    Multi-zone real eval (Blocker B)
    --------------------------------
    Prefer ``zone_obs`` / ``zone_forecasts`` parallel lists (same length as
    ``zone_ids``). Legacy single-zone callers may leave those lists empty and
    only set ``obs`` / ``forecast``; ``resolved_zone_obs()`` then returns
    ``[obs]``. Injecting multi-zone real data without the lists is refused
    by WeatherForecastEnv (padding one zone into N slots is not multi-zone).
    """
    obs:      ZoneObs
    forecast: ForecastResult
    config:   ForecastConfig = field(default_factory=ForecastConfig)

    ground_truth: Optional[RiskScore] = None
    # None during live deployment; set during backtesting / eval

    # Spatial context -- mirrors cluster_model in the MEMS env
    zone_ids:  List[str]            = field(default_factory=list)
    adjacency: Dict[str, List[str]] = field(default_factory=dict)
    # adjacency[zone_id] = [neighbouring zone_ids]

    data_source: DataSource = DataSource.SYNTHETIC

    # Basin-scale / teleconnection context (schema v3+). One per episode,
    # shared across all zones -- see BasinContext docstring. None is a valid
    # and expected value (e.g. legacy v2-shaped callers, or callers that
    # haven't wired up basin data yet); consumers must supply their own
    # neutral/climatological default when it is None, not treat it as an error.
    basin_context: Optional[BasinContext] = None

    # Per-zone real/synthetic payloads for multi-zone injection (optional).
    # When empty, resolved_* fall back to [obs] / [forecast] (single-zone).
    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"
                    )
            # Keep primary obs/forecast aligned with slot 0 for legacy readers.
            if self.zone_obs[0].zone_id != self.obs.zone_id:
                # Prefer list as authority when multi-zone lists are present.
                self.obs = self.zone_obs[0]
                self.forecast = self.zone_forecasts[0]

        # --- Adjacency integrity: all keys and neighbours must be in zone_ids ---
        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]:
        """Per-zone obs for env injection; falls back to [obs] if lists empty."""
        if self.zone_obs:
            return list(self.zone_obs)
        return [self.obs]

    def resolved_zone_forecasts(self) -> List[ForecastResult]:
        """Per-zone forecasts for env injection; falls back to [forecast]."""
        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:
    """Deterministic synthetic ZoneObs for unit tests and env smoke tests.

    Seeded via zlib.crc32 (not hash()) for process-stable reproducibility.
    In production, era5_data_pipeline.py replaces this entirely.
    """
    rng = random.Random(seed if seed is not None else _stable_seed(zone_id))

    # Deterministic timestamp: seeded offset from a fixed base, not datetime.now().
    # datetime.now() breaks training reproducibility -- two runs with the same seed
    # produce ZoneObs objects that compare unequal on valid_time.
    _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:
    """Deterministic synthetic ForecastResult for testing."""
    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:
    """Construct a complete synthetic EpisodeContext for env smoke tests."""
    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")
    # Verify force_data_source serialises/deserialises correctly when set
    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"
    )
    # New pipeline-integration fields round-trip + validation
    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
    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.")