test_weather_forecast_env_outcome_codes.py

"""
test_weather_forecast_env_outcome_codes.py
===========================================

Spec suite for the LOCKED episode-lifecycle outcome-code contract on
WeatherForecastEnv. Written against the contract, not against an
implementation that exists yet: `info["outcome_code"]` is not populated
by the current step()/reset() code, so every RESET_OK / INSPECT_OK /
TERMINATED_* / PENALTY_* test in this file is EXPECTED TO FAIL until the
outcome_code patch lands. That failure is the point (build-order rule:
contract -> tests -> code). The four ERROR_* tests assert on raised
exceptions, which already exist in the current code and should pass today.

Locked contract this file tests against
----------------------------------------
Returned paths (info["outcome_code"]):
    RESET_OK
    INSPECT_OK                      -- valid inspect, episode continues
    INSPECT_AND_BUDGET_EXHAUSTED    -- valid inspect that also hits max_steps
    TERMINATED_EARLY                -- explicit terminate action
    TERMINATED_BUDGET               -- budget exhausted via terminate action
                                        OR via a penalty branch (accepted
                                        documented asymmetry vs the inspect
                                        case above)
    PENALTY_REVISIT                 -- episode not yet terminated
    PENALTY_PADDING                 -- episode not yet terminated
    PENALTY_INVALID                 -- episode not yet terminated

Raised, never coded (real exceptions, not outcome codes):
    step() before reset()                              -> RuntimeError
    injected context length / missing zone_obs lists   -> ValueError
    n_active > max_zones                                -> ValueError
    NaNSafetyWrapper nan_limit exceeded                 -> RuntimeError

UNDEFINED, intentionally not tested here (per the contract):
    behaviour of a second step() after any terminal outcome_code
    whether a future wrapper ever converts raises into returned codes

FIXTURES
--------
Grounded directly in zone_observation.py (read in full, not inferred):
ForecastConfig, EpisodeContext, ZoneObs, ForecastResult, and the
make_synthetic_* helpers. Fixture-construction bugs found and fixed while
cross-checking against the real source (see conversation notes):
  - EpisodeContext requires real obs/forecast objects; obs=None fails in
    __post_init__ with AttributeError, not the ValueError being tested for.
  - Injecting a single-zone EpisodeContext (zone_ids length 1, zone_obs/
    zone_forecasts provided) into a 2-slot env is what actually reaches the
    padding branch -- previously skipped as "not constructible".
  - A legacy EpisodeContext with zone_ids length 2 but empty zone_obs/
    zone_forecasts lists reaches WeatherForecastEnv's
    "zone payload length mismatch" ValueError via resolved_zone_obs()
    falling back to [obs] (length 1 != n_active).
  - Corrupting belief_map[action] before stepping into that same action
    gets silently healed by _compute_zone_refinement_reward (NaN
    comparisons are always False, so max(nan - rate, signal) == signal).
    Must corrupt a zone other than the one being inspected.
"""

from __future__ import annotations

import numpy as np
import pytest

from zone_observation import (
    EpisodeContext,
    ForecastConfig,
    make_synthetic_forecast_result,
    make_synthetic_zone_obs,
)
from weather_forecast_env import (
    WeatherForecastEnv,
    NaNSafetyWrapper,
    make_weather_env,
)


# ---------------------------------------------------------------------------
# Fixtures
# ---------------------------------------------------------------------------

def _cfg(**overrides) -> ForecastConfig:
    base = dict(
        n_zones=2,
        horizon_days=3,
        max_steps=1,
        prior_belief=0.12,
        clean_episode_ratio=0.0,   # force a hazard every episode -> deterministic tests
        event_spatial_correlation=0.85,
        shuffle_zone_order=False,  # deterministic action->zone mapping for these tests
        seed=7,
    )
    base.update(overrides)
    return ForecastConfig(**base)


def _zone_obs_and_forecast(zone_id: str, cfg: ForecastConfig, seed: int, **event_flags):
    obs = make_synthetic_zone_obs(zone_id, seed=seed, **event_flags)
    fc = make_synthetic_forecast_result(
        zone_id,
        valid_time=obs.valid_time,
        horizon_days=cfg.horizon_days,
        seed=seed,
        **{k: v for k, v in event_flags.items() if k in ("drought", "flood")},
    )
    return obs, fc


@pytest.fixture
def env_triage():
    return WeatherForecastEnv(_cfg(n_zones=2, max_steps=1))


@pytest.fixture
def env_full_tour():
    return WeatherForecastEnv(_cfg(n_zones=2, max_steps=5))


def _outcome(info: dict) -> str:
    assert "outcome_code" in info, (
        "info['outcome_code'] is not populated by the current code. "
        "This is expected until the outcome-code patch is implemented; "
        "this suite is the spec for that patch, not a report of current "
        "passing behaviour."
    )
    return info["outcome_code"]


# ---------------------------------------------------------------------------
# RESET_OK
# ---------------------------------------------------------------------------

def test_reset_ok(env_full_tour):
    obs, info = env_full_tour.reset(seed=1)
    assert _outcome(info) == "RESET_OK"


def test_reset_ok_is_deterministic_for_same_seed_and_config():
    e1 = WeatherForecastEnv(_cfg(seed=42))
    e2 = WeatherForecastEnv(_cfg(seed=42))
    obs1, info1 = e1.reset(seed=42)
    obs2, info2 = e2.reset(seed=42)
    np.testing.assert_array_equal(obs1["zone_belief"], obs2["zone_belief"])
    np.testing.assert_array_equal(obs1["forecast_precip"], obs2["forecast_precip"])
    assert info1["zone_ids"] == info2["zone_ids"]


# ---------------------------------------------------------------------------
# INSPECT_OK vs INSPECT_AND_BUDGET_EXHAUSTED  (the case Option 2 exists for)
# ---------------------------------------------------------------------------

def test_inspect_ok_when_episode_continues(env_full_tour):
    env_full_tour.reset(seed=1)
    obs, reward, terminated, truncated, info = env_full_tour.step(0)
    assert terminated is False
    assert truncated is False
    assert _outcome(info) == "INSPECT_OK"


def test_inspect_and_budget_exhausted_under_triage(env_triage):
    env_triage.reset(seed=1)
    obs, reward, terminated, truncated, info = env_triage.step(0)
    assert terminated is True
    assert truncated is False
    assert _outcome(info) == "INSPECT_AND_BUDGET_EXHAUSTED"
    assert info["zone_id"] is not None
    assert info["visited"][0] == True  # noqa: E712 (explicit bool array check)


def test_inspect_ok_zero_inspect_null_control(env_full_tour):
    obs0, _ = env_full_tour.reset(seed=1)
    belief_before = obs0["zone_belief"].copy()
    obs1, reward, terminated, truncated, info = env_full_tour.step(
        env_full_tour.terminate_action
    )
    np.testing.assert_array_equal(obs1["zone_belief"], belief_before)
    assert _outcome(info) == "TERMINATED_EARLY"


# ---------------------------------------------------------------------------
# TERMINATED_EARLY / TERMINATED_BUDGET
# ---------------------------------------------------------------------------

def test_terminated_early_on_explicit_terminate(env_full_tour):
    env_full_tour.reset(seed=3)
    _, _, terminated, _, info = env_full_tour.step(env_full_tour.terminate_action)
    assert terminated is True
    assert _outcome(info) == "TERMINATED_EARLY"


def test_terminated_budget_via_penalty_branch(env_triage):
    env_triage.reset(seed=5)
    invalid_action = env_triage.max_zones + 5  # out of legal range
    _, _, terminated, _, info = env_triage.step(invalid_action)
    assert terminated is True
    assert _outcome(info) == "TERMINATED_BUDGET"
    assert info.get("invalid_action") is True


def test_terminate_and_budget_exhaust_agree_on_reward_shape():
    e1 = WeatherForecastEnv(_cfg(n_zones=2, max_steps=5, seed=9))
    e1.reset(seed=9)
    e1.step(0)
    _, r1, term1, _, info1 = e1.step(e1.terminate_action)

    e2 = WeatherForecastEnv(_cfg(n_zones=2, max_steps=2, seed=9))
    e2.reset(seed=9)
    e2.step(0)
    _, r2, term2, _, info2 = e2.step(1) 

    assert term1 is True and term2 is True
    assert _outcome(info1) == "TERMINATED_EARLY"
    assert _outcome(info2) == "INSPECT_AND_BUDGET_EXHAUSTED"
    assert "believed_p" in info1 and "believed_p" in info2


# ---------------------------------------------------------------------------
# PENALTY_REVISIT / PENALTY_PADDING / PENALTY_INVALID (episode continues)
# ---------------------------------------------------------------------------

def test_penalty_revisit_when_not_terminated(env_full_tour):
    env_full_tour.reset(seed=11)
    env_full_tour.step(0)  # first visit, legal
    _, reward, terminated, _, info = env_full_tour.step(0)  # revisit
    assert terminated is False
    assert _outcome(info) == "PENALTY_REVISIT"
    assert reward < 0


def test_penalty_padding_when_not_terminated():
    cfg = _cfg(n_zones=2, max_steps=5)
    env = WeatherForecastEnv(cfg)
    obs0, fc0 = _zone_obs_and_forecast("z0", cfg, seed=31, drought=True)
    single_zone_context = EpisodeContext(
        obs=obs0,
        forecast=fc0,
        config=cfg,
        zone_ids=["z0"],
        zone_obs=[obs0],
        zone_forecasts=[fc0],
    )
    env.reset(seed=1, options={"context": single_zone_context})
    _, reward, terminated, _, info = env.step(1)  # slot 1: padding, n_active=1
    assert terminated is False
    assert _outcome(info) == "PENALTY_PADDING"
    assert reward < 0


def test_penalty_invalid_when_not_terminated(env_full_tour):
    env_full_tour.reset(seed=17)
    invalid_action = env_full_tour.max_zones + 5
    _, reward, terminated, _, info = env_full_tour.step(invalid_action)
    assert terminated is False
    assert _outcome(info) == "PENALTY_INVALID"
    assert reward < 0


# ---------------------------------------------------------------------------
# Shuffle mapping (test scenario #7)
# ---------------------------------------------------------------------------

def test_shuffle_changes_action_to_zone_mapping_not_belief_semantics():
    cfg_shuffled = _cfg(n_zones=2, max_steps=5, shuffle_zone_order=True, seed=21)
    env = WeatherForecastEnv(cfg_shuffled)
    obs, info = env.reset(seed=21)
    zone_ids_at_reset = list(info["zone_ids"])
    _, _, _, _, info_after_step = env.step(0)
    assert info_after_step["zone_ids"] == zone_ids_at_reset


# ---------------------------------------------------------------------------
# Raised, never coded (real exceptions -- should pass against current code)
# ---------------------------------------------------------------------------

def test_step_before_reset_raises_runtime_error():
    env = WeatherForecastEnv(_cfg())
    with pytest.raises(RuntimeError):
        env.step(0)


def test_zone_overflow_raises_value_error():
    cfg = _cfg(n_zones=2, max_steps=5)
    env = WeatherForecastEnv(cfg)
    obs0, fc0 = _zone_obs_and_forecast("z0", cfg, seed=33)
    bad_context = EpisodeContext(
        obs=obs0,
        forecast=fc0,
        config=cfg,
        zone_ids=["z0", "z1", "z2"],  
    )
    with pytest.raises(ValueError):
        env.reset(seed=1, options={"context": bad_context})


def test_injection_length_mismatch_raises_value_error():
    cfg = _cfg(n_zones=2, max_steps=5)
    env = WeatherForecastEnv(cfg)
    obs0, fc0 = _zone_obs_and_forecast("z0", cfg, seed=35)
    legacy_multi_zone_context = EpisodeContext(
        obs=obs0,
        forecast=fc0,
        config=cfg,
        zone_ids=["z0", "z1"],
    )
    with pytest.raises(ValueError):
        env.reset(seed=1, options={"context": legacy_multi_zone_context})


def test_nan_limit_exceeded_raises_runtime_error():
    cfg = _cfg(n_zones=2, max_steps=5)
    inner = WeatherForecastEnv(cfg)
    wrapped = NaNSafetyWrapper(inner, nan_limit=0)
    wrapped.reset(seed=1)

    inner._belief_map[1] = float("nan")
    with pytest.raises(RuntimeError):
        wrapped.step(0)


# ---------------------------------------------------------------------------
# max_steps=1 (triage) vs full tour -- test scenario #8
# ---------------------------------------------------------------------------

def test_triage_vs_full_tour_are_not_interchangeable(env_triage, env_full_tour):
    assert env_triage.config.max_steps < env_triage.max_zones
    assert env_full_tour.config.max_steps >= env_full_tour.max_zones