weather_forecast_env.py

"""
weather_forecast_env.py
=======================
"""

from __future__ import annotations

import logging
import random
from typing import Any, Dict, List, Optional, Tuple

import numpy as np
import gymnasium as gym
from gymnasium import spaces

import zone_observation as _zo

assert _zo.SCHEMA_VERSION == 3, (
    f"weather_forecast_env: zone_observation schema mismatch "
    f"(expected 3, got {_zo.SCHEMA_VERSION})"
)

from zone_observation import (
    BasinContext,
    DataSource,
    EpisodeContext,
    ForecastConfig,
    RiskScore,
    make_synthetic_forecast_result,
    make_synthetic_zone_obs,
    _stable_seed,
)

_BASIN_CONTEXT_NEUTRAL = (0.0, 0.0, 0.0, 1013.25, 400.0, 2.0, -7.0, 0.0)
_HELIO_REGIME_ORD = {"quiet": 0.0, "active": 1.0, "storm": 2.0}


def basin_context_vector(bc: Optional[BasinContext]) -> np.ndarray:
    if bc is None:
        return np.array(_BASIN_CONTEXT_NEUTRAL, dtype=np.float32)
    goes = float(getattr(bc, "goes_xray_flux", 1e-7) or 1e-7)
    goes = max(goes, 1e-12)
    regime = str(getattr(bc, "helio_regime", "quiet") or "quiet").lower()
    return np.array(
        [
            float(getattr(bc, "enso_oni", 0.0) or 0.0),
            float(getattr(bc, "iod_dmi", 0.0) or 0.0),
            float(getattr(bc, "itcz_latitude_deg", 0.0) or 0.0),
            float(getattr(bc, "mslp_regional_hpa", 1013.25) or 1013.25),
            float(getattr(bc, "solar_wind_speed_kms", 400.0) or 400.0),
            float(getattr(bc, "kp_index", 2.0) or 2.0),
            float(np.log10(goes)),
            float(_HELIO_REGIME_ORD.get(regime, 0.0)),
        ],
        dtype=np.float32,
    )


from crop_risk_scorer import compute_risk_score

try:
    from product_alert_service import (
        DEFAULT_PRODUCT_GATE,
        is_product_actionable as _product_gate_fn,
    )
except ImportError:
    DEFAULT_PRODUCT_GATE = None
    _product_gate_fn = None

try:
    from real_episode_sampler import RealEpisodeIndex
except ImportError:
    RealEpisodeIndex = None  # training real-data sampling unavailable;
                              # ForecastConfig.real_data_pkl_path is then
                              # ignored and training stays fully synthetic,
                              # same as before this module existed.

logger = logging.getLogger(__name__)


def _info_product_flags(risk_score: RiskScore) -> Dict[str, Any]:
    elevated = bool(risk_score.is_elevated())
    if _product_gate_fn is not None and DEFAULT_PRODUCT_GATE is not None:
        product = bool(_product_gate_fn(risk_score, DEFAULT_PRODUCT_GATE))
    else:
        product = bool(risk_score.is_product_actionable())
    return {
        "elevated": elevated,
        "product_actionable": product,
    }


class NaNSafetyWrapper(gym.Wrapper):
    def __init__(
        self,
        env: gym.Env,
        default_value: float = 0.0,
        reward_default: float = -100.0,
        nan_limit: int = 5000,
        verbose: bool = False,
    ) -> None:
        super().__init__(env)
        self.default_value = float(default_value)
        self.reward_default = float(reward_default)
        self.nan_count = 0
        self.nan_limit = nan_limit
        self.verbose = verbose

    def _sanitize(self, value: Any, name: str = "") -> Any:
        if isinstance(value, dict):
            return {k: self._sanitize(v, f"{name}.{k}") for k, v in value.items()}
        if isinstance(value, np.ndarray):
            if not np.all(np.isfinite(value)):
                n = int((~np.isfinite(value)).sum())
                self.nan_count += n
                if self.verbose:
                    logger.warning("NaNSafetyWrapper: %d NaN/Inf in %s", n, name)
                return np.nan_to_num(
                    value,
                    nan=self.default_value,
                    posinf=self.default_value,
                    neginf=self.default_value,
                )
            return value
        if isinstance(value, (float, np.floating)):
            if not np.isfinite(value):
                self.nan_count += 1
                return self.default_value
            return float(value)
        if isinstance(value, bool):
            return value
        if isinstance(value, (int, np.integer)):
            return int(value)
        return value

    def step(self, action):
        obs, reward, terminated, truncated, info = self.env.step(action)
        obs = self._sanitize(obs, "obs")
        info = self._sanitize(info, "info")
        reward = reward if np.isfinite(reward) else self.reward_default
        reward = float(np.clip(reward, -31_000.0, 15_500.0))
        if self.nan_count > self.nan_limit:
            raise RuntimeError(
                f"NaNSafetyWrapper: {self.nan_count} invalid values -- "
                f"check data pipeline."
            )
        return obs, reward, terminated, truncated, info

    def reset(self, **kwargs):
        obs, info = self.env.reset(**kwargs)
        return self._sanitize(obs, "obs"), self._sanitize(info, "info")

    @property
    def terminate_action(self) -> int:
        return self.env.terminate_action

    def action_masks(self) -> np.ndarray:
        return self.env.action_masks()


def _zone_event_flags(seed: int, clean_ratio: float) -> Dict[str, bool]:
    rng = random.Random(seed)
    if rng.random() < clean_ratio:
        return {"drought": False, "flood": False, "fungi": False}
    kind = rng.choice(["drought", "flood", "fungi"])
    return {
        "drought": kind == "drought",
        "flood":   kind == "flood",
        "fungi":   kind == "fungi",
    }


def _episode_event_plan(
    n_zones: int,
    base_seed: int,
    clean_ratio: float,
    spatial_corr: float = 0.85,
) -> List[Dict[str, bool]]:
    n = max(1, int(n_zones))
    clean = {"drought": False, "flood": False, "fungi": False}
    rng = random.Random(base_seed + 7919)
    rho = float(max(0.0, min(1.0, spatial_corr)))
    p_event = 1.0 - float(max(0.0, min(1.0, clean_ratio)))

    if rng.random() >= p_event:
        return [dict(clean) for _ in range(n)]

    kind = rng.choice(["drought", "flood", "fungi"])
    shared = {
        "drought": kind == "drought",
        "flood":   kind == "flood",
        "fungi":   kind == "fungi",
    }
    plan: List[Dict[str, bool]] = []
    for i in range(n):
        if rng.random() < rho:
            plan.append(dict(shared))
        else:
            plan.append(dict(clean))
    if not any(any(f.values()) for f in plan):
        plan[rng.randrange(n)] = dict(shared)
    return plan


def _make_multi_zone_context(
    cfg: ForecastConfig,
    seed: Optional[int],
) -> EpisodeContext:
    n = cfg.n_zones
    base = seed if seed is not None else 0

    zone_ids = [f"zone_{i}" for i in range(n)]
    rho = float(getattr(cfg, "event_spatial_correlation", 0.85))
    plan = _episode_event_plan(n, base, cfg.clean_episode_ratio, rho)
    ev0 = plan[0]
    obs_0 = make_synthetic_zone_obs(zone_ids[0], seed=base, **ev0)
    fc_0 = make_synthetic_forecast_result(
        zone_ids[0], valid_time=obs_0.valid_time, seed=base,
        drought=ev0["drought"], flood=ev0["flood"],
    )

    return EpisodeContext(
        obs=obs_0,
        forecast=fc_0,
        config=cfg,
        zone_ids=zone_ids,
        data_source=DataSource.SYNTHETIC,
    )


def _make_per_zone_forecasts(
    zone_ids: List[str],
    valid_time: Any,
    horizon: int,
    base_seed: int,
    clean_ratio: float,
    spatial_corr: float = 0.85,
    event_plan: Optional[List[Dict[str, bool]]] = None,
) -> List[Any]:
    if event_plan is None:
        event_plan = _episode_event_plan(
            len(zone_ids), base_seed, clean_ratio, spatial_corr
        )
    forecasts = []
    for i, zid in enumerate(zone_ids):
        ev = event_plan[i] if i < len(event_plan) else _zone_event_flags(
            base_seed + i, clean_ratio
        )
        fc = make_synthetic_forecast_result(
            zid,
            horizon_days=horizon,
            valid_time=valid_time,
            seed=base_seed + i * 31,
            drought=ev["drought"],
            flood=ev["flood"],
        )
        forecasts.append(fc)
    return forecasts


def _per_zone_beliefs(
    zone_ids: List[str],
    base_seed: int,
    prior: float,
    clean_ratio: float,
    spatial_corr: float = 0.85,
    event_plan: Optional[List[Dict[str, bool]]] = None,
    prior_weight: float = 0.70,
) -> np.ndarray:
    if event_plan is None:
        event_plan = _episode_event_plan(
            len(zone_ids), base_seed, clean_ratio, spatial_corr
        )
    beliefs = []
    for i, zid in enumerate(zone_ids):
        ev = event_plan[i] if i < len(event_plan) else _zone_event_flags(
            base_seed + i, clean_ratio
        )
        obs = make_synthetic_zone_obs(zid, seed=base_seed + i * 17, **ev)
        signal = float(obs.composite_risk())
        belief = prior_weight * prior + (1.0 - prior_weight) * signal
        beliefs.append(belief)
    return np.clip(np.array(beliefs, dtype=np.float32), 0.0, 1.0)


def _resolved_zone_beliefs(
    zone_obs: List[Any],
    prior: float,
    prior_weight: float,
) -> np.ndarray:
    """Belief-map values for zone_obs that are already resolved (real data,
    whether externally injected via options={"context": ...} or sampled by
    RealEpisodeIndex during training) -- as opposed to _per_zone_beliefs()
    above, which generates synthetic obs itself from a seed.

    Single source of truth for the prior/signal blend across both
    real-data code paths in reset(), so they cannot drift apart the way
    the injected path and the synthetic path's hard-coded 0.70 literal
    once did.
    """
    beliefs = np.zeros(len(zone_obs), dtype=np.float32)
    for i, zo in enumerate(zone_obs):
        signal = float(zo.composite_risk())
        beliefs[i] = prior_weight * prior + (1.0 - prior_weight) * signal
    return np.clip(beliefs, 0.0, 1.0)


def _make_per_zone_obs(
    zone_ids: List[str],
    base_seed: int,
    clean_ratio: float,
    spatial_corr: float = 0.85,
    event_plan: Optional[List[Dict[str, bool]]] = None,
) -> List[Any]:
    if event_plan is None:
        event_plan = _episode_event_plan(
            len(zone_ids), base_seed, clean_ratio, spatial_corr
        )
    out = []
    for i, zid in enumerate(zone_ids):
        ev = event_plan[i] if i < len(event_plan) else _zone_event_flags(
            base_seed + i, clean_ratio
        )
        out.append(make_synthetic_zone_obs(zid, seed=base_seed + i * 17, **ev))
    return out


class WeatherForecastEnv(gym.Env):
    metadata = {"render_modes": ["human"]}

    def __init__(self, config: Optional[ForecastConfig] = None) -> None:
        super().__init__()
        self.config = config or ForecastConfig()
        self.max_zones = self.config.n_zones
        self.terminate_action = self.max_zones

        self.action_space = spaces.Discrete(self.max_zones + 1)
        self.observation_space = spaces.Dict({
            "zone_belief": spaces.Box(
                0.0, 1.0, (self.max_zones,), np.float32,
            ),
            "forecast_precip": spaces.Box(
                0.0, 500.0, (self.max_zones, self.config.horizon_days), np.float32,
            ),
            "forecast_uncertainty": spaces.Box(
                0.0, 1.0, (self.max_zones,), np.float32,
            ),
            "action_mask": spaces.Box(
                0, 1, (self.max_zones + 1,), bool,
            ),
            "prior_belief": spaces.Box(
                0.0, 1.0, (1,), np.float32,
            ),
            "basin_context": spaces.Box(
                low=np.array(
                    [-5.0, -5.0, -30.0, 900.0, 200.0, 0.0, -9.0, 0.0],
                    dtype=np.float32,
                ),
                high=np.array(
                    [5.0, 5.0, 30.0, 1100.0, 900.0, 9.0, -3.0, 2.0],
                    dtype=np.float32,
                ),
                dtype=np.float32,
            ),
        })

        self._belief_map = np.full(self.max_zones, self.config.prior_belief, np.float32)
        self._action_mask = np.ones(self.max_zones + 1, dtype=bool)
        self._forecast_arr = np.zeros((self.max_zones, self.config.horizon_days), np.float32)
        self._uncertainty = np.zeros(self.max_zones, np.float32)
        self._uncertainty_init = np.zeros(self.max_zones, np.float32)
        self._visited = np.zeros(self.max_zones, dtype=bool)
        self._basin_context = np.array(_BASIN_CONTEXT_NEUTRAL, dtype=np.float32)

        self._context: Optional[EpisodeContext] = None
        self._zone_forecasts: List[Any] = []
        self._zone_obs: List[Any] = []
        self._zone_ids: List[str] = []
        self._steps_taken: int = 0
        self._cum_reward: float = 0.0
        self._episode_count: int = 0
        self._episode_seed: int = 0
        self._context_source: str = "synthetic"

        self._real_index = None
        pkl_path = getattr(self.config, "real_data_pkl_path", None)
        if pkl_path:
            if RealEpisodeIndex is None:
                logger.warning(
                    "WeatherForecastEnv: config.real_data_pkl_path=%r set, "
                    "but real_episode_sampler.py is not importable -- "
                    "training will remain fully synthetic.",
                    pkl_path,
                )
            else:
                try:
                    self._real_index = RealEpisodeIndex(pkl_path)
                    logger.info(
                        "WeatherForecastEnv: real-data training sampling "
                        "ENABLED from %s (%d eligible dates, "
                        "real_data_ratio=%.2f, inject_noise=%s)",
                        pkl_path, self._real_index.n_eligible_dates,
                        self.config.real_data_ratio, self.config.inject_noise,
                    )
                except Exception as e:
                    logger.error(
                        "WeatherForecastEnv: failed to load "
                        "real_data_pkl_path=%r (%s) -- training will "
                        "remain fully synthetic.", pkl_path, e,
                    )
                    self._real_index = None

        logger.info(
            "WeatherForecastEnv: n_zones=%d  horizon=%dd  max_steps=%d  "
            "terminate_action=%d",
            self.max_zones, self.config.horizon_days,
            self.config.max_steps, self.terminate_action,
        )

    def action_masks(self) -> np.ndarray:
        return self._action_mask.copy()

    def reset(
        self,
        *,
        seed: Optional[int] = None,
        options: Optional[Dict] = None,
    ) -> Tuple[Dict[str, np.ndarray], Dict[str, Any]]:
        super().reset(seed=seed)

        effective_seed = seed if seed is not None else (
            (self.config.seed or 0) + self._episode_count
        )
        self._episode_seed = effective_seed

        if options and "context" in options:
            ctx = options["context"]
            if not isinstance(ctx, EpisodeContext):
                raise TypeError(
                    f"options['context'] must be EpisodeContext, "
                    f"got {type(ctx).__name__}"
                )
            self._context = ctx
            context_source = "injected"
        else:
            self._context = None
            context_source = "synthetic"
            if (
                self._real_index is not None
                and float(self.config.real_data_ratio) > 0.0
            ):
                # Dedicated RNG stream, offset from the streams
                # _maybe_shuffle_zone_order and other reset-time randomness
                # use, so enabling real-data sampling doesn't change the
                # synthetic path's own random draws when it IS taken.
                real_rng = random.Random(int(effective_seed) + 9001)
                if real_rng.random() < float(self.config.real_data_ratio):
                    sampled = self._real_index.sample(
                        real_rng, self.config.n_zones, self.config,
                        inject_noise=bool(self.config.inject_noise),
                        noise_scale=float(self.config.noise_scale),
                    )
                    if sampled is not None:
                        self._context = sampled
                        context_source = "real_sampled"
            if self._context is None:
                self._context = _make_multi_zone_context(self.config, effective_seed)
                context_source = "synthetic"

        self._zone_ids = list(self._context.zone_ids)
        n_active = len(self._zone_ids)
        self._context_source = context_source  # stamped onto every step's
                                                  # info via _get_info(), not
                                                  # just reset's -- SB3's
                                                  # DummyVecEnv stores
                                                  # auto-reset info separately
                                                  # from what step() returns,
                                                  # so anything only set in
                                                  # reset()'s info dict is
                                                  # invisible to a callback's
                                                  # self.locals["infos"].
                                                  # Confirmed empirically.

        if n_active > self.max_zones:
            raise ValueError(
                f"EpisodeContext has {n_active} zones but env supports "
                f"max {self.max_zones}."
            )

        # "injected" (explicit real-eval context) and "real_sampled"
        # (RealEpisodeIndex draw during training) both arrive as a fully
        # resolved EpisodeContext with real zone_obs/zone_forecasts already
        # attached -- they share the same validation and belief-computation
        # path below. Only "synthetic" generates obs/forecasts itself.
        use_resolved_path = context_source in ("injected", "real_sampled")

        self._forecast_arr.fill(0.0)
        self._uncertainty.fill(0.0)
        self._uncertainty_init.fill(0.0)
        self._visited.fill(False)
        self._steps_taken = 0
        self._cum_reward = 0.0
        self._episode_count += 1

        self._basin_context[:] = basin_context_vector(
            self._context.basin_context if self._context is not None else None
        )

        # Single source of truth for the prior/signal belief-blend weight --
        # the injected, real-sampled, and synthetic reset paths below must
        # all read this same value, not their own copies of it.
        prior_w = float(getattr(self.config, "belief_prior_weight", 0.70))

        if use_resolved_path:
            z_obs = self._context.resolved_zone_obs()
            z_fc = self._context.resolved_zone_forecasts()
            if len(z_obs) != n_active or len(z_fc) != n_active:
                raise ValueError(
                    f"{context_source} EpisodeContext zone payload length "
                    f"mismatch: n_active={n_active} resolved_obs={len(z_obs)} "
                    f"resolved_fc={len(z_fc)} zone_ids={self._zone_ids}. "
                    f"For multi-zone episodes supply zone_obs/zone_forecasts "
                    f"parallel to zone_ids (do not pad one zone into N slots)."
                )
            if n_active > 1 and not self._context.zone_obs:
                raise ValueError(
                    f"{context_source} EpisodeContext refuses n_zones>1 "
                    f"without explicit zone_obs lists (got n_active={n_active}, "
                    f"zone_ids={self._zone_ids}). Padding the primary obs into "
                    f"multiple slots is not multi-zone data."
                )
            for i, zid in enumerate(self._zone_ids):
                if z_obs[i].zone_id != zid or z_fc[i].zone_id != zid:
                    raise ValueError(
                        f"{context_source} context zone mismatch at slot {i}: "
                        f"zone_ids[{i}]={zid!r} obs={z_obs[i].zone_id!r} "
                        f"fc={z_fc[i].zone_id!r}"
                    )
            self._zone_obs = list(z_obs)
            self._zone_forecasts = list(z_fc)
            prior = float(self.config.prior_belief)
            risks = _resolved_zone_beliefs(self._zone_obs, prior, prior_w)
            self._belief_map[:n_active] = risks[:n_active]
            self._belief_map[n_active:] = 0.0
        else:
            rho = float(getattr(self.config, "event_spatial_correlation", 0.85))
            plan = _episode_event_plan(
                n_active,
                effective_seed,
                self.config.clean_episode_ratio,
                rho,
            )
            risks = _per_zone_beliefs(
                self._zone_ids, effective_seed, self.config.prior_belief,
                self.config.clean_episode_ratio,
                spatial_corr=rho,
                event_plan=plan,
                prior_weight=prior_w,
            )
            self._belief_map[:n_active] = risks[:n_active]
            self._belief_map[n_active:] = 0.0
            self._zone_forecasts = _make_per_zone_forecasts(
                self._zone_ids,
                self._context.obs.valid_time,
                self.config.horizon_days,
                effective_seed,
                self.config.clean_episode_ratio,
                spatial_corr=rho,
                event_plan=plan,
            )
            self._zone_obs = _make_per_zone_obs(
                self._zone_ids, effective_seed, self.config.clean_episode_ratio,
                spatial_corr=rho,
                event_plan=plan,
            )

        self._maybe_shuffle_zone_order(effective_seed, n_active)

        self._update_forecast_arrays()
        self._uncertainty_init[:] = self._uncertainty

        info = self._get_info()
        info["outcome_code"] = "RESET_OK"
        info["context_source"] = context_source
        return self._get_obs(), info

    def _maybe_shuffle_zone_order(self, effective_seed: int, n_active: int) -> None:
        if n_active <= 1:
            return
        if not bool(getattr(self.config, "shuffle_zone_order", True)):
            return
        if not self._zone_obs or not self._zone_forecasts:
            return
        if len(self._zone_obs) < n_active or len(self._zone_forecasts) < n_active:
            return

        order = list(range(n_active))
        random.Random(int(effective_seed) + 4242).shuffle(order)

        self._zone_ids = [self._zone_ids[i] for i in order]
        self._zone_obs = [self._zone_obs[i] for i in order]
        self._zone_forecasts = [self._zone_forecasts[i] for i in order]

        reordered = np.zeros_like(self._belief_map)
        for new_i, old_i in enumerate(order):
            reordered[new_i] = self._belief_map[old_i]
        self._belief_map[:n_active] = reordered[:n_active]
        self._belief_map[n_active:] = 0.0

        if self._context is not None:
            self._context = EpisodeContext(
                obs=self._zone_obs[0],
                forecast=self._zone_forecasts[0],
                config=self.config,
                zone_ids=list(self._zone_ids),
                basin_context=self._context.basin_context,
                data_source=getattr(
                    self._context, "data_source", DataSource.SYNTHETIC
                ),
                ground_truth=getattr(self._context, "ground_truth", None),
                zone_obs=list(self._zone_obs),
                zone_forecasts=list(self._zone_forecasts),
            )

    def _update_forecast_arrays(self) -> None:
        if self._context is None:
            return

        h = self.config.horizon_days
        n = min(len(self._zone_ids), self.max_zones)

        for zi in range(n):
            fr = (
                self._zone_forecasts[zi]
                if zi < len(self._zone_forecasts)
                else self._context.forecast
            )

            if fr.precip_mm and len(fr.precip_mm) != h:
                logger.warning(
                    "_update_forecast_arrays: zone %d forecast length=%d != %d",
                    zi, len(fr.precip_mm), h,
                )

            if fr.precip_mm:
                vals = np.asarray(fr.precip_mm[:h], dtype=np.float32)
                self._forecast_arr[zi, :len(vals)] = vals

            if fr.precip_p90 and fr.precip_p10:
                p90 = np.asarray(fr.precip_p90[:h], dtype=np.float32)
                p10 = np.asarray(fr.precip_p10[:h], dtype=np.float32)
                spread = np.clip(
                    (p90 - p10) / np.maximum(np.abs(p90), 1e-6), 0.0, 1.0
                )
                self._uncertainty[zi] = float(np.mean(spread))

    def _compute_multi_zone_risk(self) -> RiskScore:
        if not self._zone_obs or not self._zone_ids:
            return compute_risk_score(
                self._context.obs,
                self._context.forecast,
                self._context.config,
            )

        n_active = min(len(self._zone_ids), len(self._zone_obs), self.max_zones)
        best_score: Optional[RiskScore] = None

        for zi in range(n_active):
            zone_obs = self._zone_obs[zi]
            zone_fc = (
                self._zone_forecasts[zi]
                if zi < len(self._zone_forecasts)
                else self._context.forecast
            )
            score = compute_risk_score(
                zone_obs,
                zone_fc,
                self._context.config,
            )
            if best_score is None or score.supply_shortfall_prob > best_score.supply_shortfall_prob:
                best_score = score

        return best_score if best_score is not None else compute_risk_score(
            self._context.obs,
            self._context.forecast,
            self._context.config,
        )

    def step(
        self, action: int
    ) -> Tuple[Dict[str, np.ndarray], float, bool, bool, Dict[str, Any]]:
        if self._context is None:
            raise RuntimeError("WeatherForecastEnv: call reset() before step()")

        info: Dict[str, Any] = {}

        if action == self.terminate_action:
            risk_score = self._compute_multi_zone_risk()
            believed_p = float(np.max(self._belief_map[:len(self._zone_ids)]))
            reward = self._compute_termination_reward(risk_score, believed_p)
            self._steps_taken += 1
            self._cum_reward += reward

            info["early_termination"] = True
            info["budget_saved"] = self.config.max_steps - self._steps_taken
            info["alert_level"] = risk_score.alert_level.value
            info.update(_info_product_flags(risk_score))
            info["believed_p"] = believed_p
            info["supply_risk"] = risk_score.supply_shortfall_prob
            info["flood_risk"] = risk_score.flood_risk
            info["drought_risk"] = risk_score.drought_risk
            info.update(self._get_info())
            info["outcome_code"] = "TERMINATED_EARLY"

            return self._get_obs(), float(reward), True, False, info

        if action < 0 or action >= self.max_zones:
            reward = -self.config.inspection_cost * 2.0 / self.config.alert_value
            info["invalid_action"] = True
            _branch = "invalid"
        elif action >= len(self._zone_ids):
            reward = -self.config.inspection_cost * 0.5 / self.config.alert_value
            info["padding_action"] = True
            _branch = "padding"
        elif self._visited[action]:
            reward = -self.config.inspection_cost * 3.0 / self.config.alert_value
            info["revisit_penalty"] = True
            _branch = "revisit"
        else:
            zone_id = self._zone_ids[action]
            reward = self._compute_zone_refinement_reward(action, zone_id)
            self._visited[action] = True
            info["zone_id"] = zone_id
            _branch = "inspect"

        self._steps_taken += 1
        self._cum_reward += reward

        terminated = self._steps_taken >= self.config.max_steps

        if terminated:
            risk_score = self._compute_multi_zone_risk()
            believed_p = float(np.max(self._belief_map[:len(self._zone_ids)]))
            term_r = self._compute_termination_reward(risk_score, believed_p)
            reward += term_r
            self._cum_reward += term_r
            info["budget_exhausted"] = True
            info["believed_p"] = believed_p
            info["alert_level"] = risk_score.alert_level.value
            info.update(_info_product_flags(risk_score))
            info["drought_risk"] = risk_score.drought_risk
            info["flood_risk"] = risk_score.flood_risk
            info["supply_risk"] = risk_score.supply_shortfall_prob

        if terminated:
            outcome_code = (
                "INSPECT_AND_BUDGET_EXHAUSTED" if _branch == "inspect"
                else "TERMINATED_BUDGET"
            )
        else:
            outcome_code = {
                "invalid": "PENALTY_INVALID",
                "padding": "PENALTY_PADDING",
                "revisit": "PENALTY_REVISIT",
                "inspect": "INSPECT_OK",
            }[_branch]

        info.update(self._get_info())
        info["outcome_code"] = outcome_code
        return (
            self._get_obs(),
            float(np.clip(reward, -31_000.0, 15_500.0)),
            terminated, False, info,
        )

    def _belief_entropy(self, p: float) -> float:
        p = max(1e-9, min(1.0 - 1e-9, float(p)))
        return -(p * np.log(p) + (1.0 - p) * np.log(1.0 - p))

    def _compute_zone_refinement_reward(self, idx: int, zone_id: str) -> float:
        current = float(self._belief_map[idx])
        if idx < len(self._zone_obs):
            signal = float(self._zone_obs[idx].composite_risk())
        else:
            signal = self.config.prior_belief

        if signal > current:
            updated = min(current + self.config.belief_increase_rate, signal)
        else:
            updated = max(current - self.config.belief_decrease_rate, signal)

        self._belief_map[idx] = float(np.clip(
            updated, self.config.belief_floor, 1.0,
        ))

        before_err = (current - signal) ** 2
        after_err = (float(self._belief_map[idx]) - signal) ** 2
        info_gain = max(0.0, before_err - after_err)

        decay = float(getattr(self.config, "uncertainty_decay", 0.70))
        self._uncertainty[idx] = float(
            np.clip(self._uncertainty[idx] * decay, 0.0, 1.0)
        )

        info_scale = float(getattr(self.config, "info_gain_scale", 5.0))
        raw_reward = (
            info_scale * info_gain
            - self.config.inspection_cost
            + self.config.zone_visit_bonus
        )
        return float(raw_reward / self.config.alert_value)

    def _compute_termination_reward(self, risk_score: RiskScore, believed_p: float) -> float:
        cfg = self._context.config
        p_event = believed_p
        gain = p_event * (cfg.alert_value + cfg.miss_penalty)
        cost = (1.0 - p_event) * cfg.false_alert_penalty
        unc_scale = float(getattr(cfg, "uncertainty_penalty_scale", 5.0))
        unc = unc_scale * float(np.mean(self._uncertainty))

        n_active = len(self._zone_ids)
        n_visited = int(np.sum(self._visited[:n_active])) if n_active > 0 else 0
        budget = max(1, min(n_active, int(self.config.max_steps)))
        unvisited_frac = (1.0 - n_visited / budget) if budget > 0 else 0.0
        exploration_penalty = cfg.unvisited_zone_penalty * unvisited_frac

        base = (gain - cost - unc - exploration_penalty) / cfg.alert_value

        if self._context.ground_truth is not None:
            gt = self._context.ground_truth.supply_shortfall_prob
            if (gt > cfg.rational_termination_threshold
                    and p_event < cfg.rational_termination_threshold):
                base -= cfg.miss_penalty / cfg.alert_value

        return float(base)

    def _get_obs(self) -> Dict[str, np.ndarray]:
        n_active = len(self._zone_ids)

        self._action_mask[:n_active] = ~self._visited[:n_active]
        self._action_mask[n_active:self.max_zones] = False
        self._action_mask[self.terminate_action] = True

        return {
            "zone_belief": self._belief_map.copy(),
            "forecast_precip": self._forecast_arr.copy(),
            "forecast_uncertainty": self._uncertainty.copy(),
            "action_mask": self._action_mask.copy(),
            "prior_belief": np.array([self.config.prior_belief], dtype=np.float32),
            "basin_context": self._basin_context.copy(),
        }

    def _get_info(self) -> Dict[str, Any]:
        n_active = len(self._zone_ids)
        active_belief = self._belief_map[:n_active]
        active_unc = self._uncertainty[:n_active]
        return {
            "episode_num": self._episode_count,
            "steps_taken": self._steps_taken,
            "cumulative_reward": float(self._cum_reward),
            "n_zones": n_active,
            "zone_ids": list(self._zone_ids),
            "visited": self._visited.copy(),
            "believed_p": float(np.max(active_belief)) if n_active else 0.0,
            "mean_belief": float(np.mean(active_belief)) if n_active else 0.0,
            "mean_uncertainty": float(np.mean(active_unc)) if n_active else 0.0,
            "n_visited": int(self._visited[:n_active].sum()) if n_active else 0,
            # Stamped on EVERY step's info via this method, not just reset's --
            # see the comment in reset() where self._context_source is set for
            # why that distinction matters under SB3's VecEnv autoreset.
            "context_source": self._context_source,
        }

    def render(self, mode: Optional[str] = None) -> None:
        pass

    def close(self) -> None:
        pass


def make_weather_env(
    config: Optional[ForecastConfig] = None,
    use_nan_wrapper: bool = True,
) -> gym.Env:
    env = WeatherForecastEnv(config)
    if use_nan_wrapper:
        env = NaNSafetyWrapper(env, verbose=False)
    return env


def register_weather_environments() -> None:
    gym.register(
        id="WeatherForecast-v1",
        entry_point=lambda: make_weather_env(ForecastConfig()),
        max_episode_steps=None,
    )
    logger.info("Registered WeatherForecast-v1 with Gymnasium")