train_curriculum.py

"""
train_curriculum.py
===================
"""

from __future__ import annotations

import argparse
import logging
import os
from dataclasses import dataclass, field
from pathlib import Path
from typing import Dict, List, Optional

import zone_observation as _zo

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

from zone_observation import ForecastConfig
from crop_risk_scorer import RiskWeights

# ---------------------------------------------------------------------------
# Optional ML imports (graceful degradation)
# ---------------------------------------------------------------------------

try:
    import torch
    _TORCH_AVAILABLE = True
except ImportError:
    _TORCH_AVAILABLE = False

try:
    from weather_forecast_env import make_weather_env
    from sb3_contrib import MaskablePPO
    from stable_baselines3.common.monitor import Monitor
    from stable_baselines3.common.callbacks import BaseCallback
    _ML_AVAILABLE = True
except ImportError as _e:
    _ML_AVAILABLE = False
    _ML_IMPORT_ERROR = str(_e)
    make_weather_env = None        # type: ignore
    MaskablePPO     = None        # type: ignore
    Monitor         = None        # type: ignore
    BaseCallback    = object      # type: ignore

# GRU policy is optional — falls back to MlpPolicy if not present
try:
    from gru_weather_policy import (
        create_gru_weather_policy_kwargs,
        ZoneEquivariantMaskablePolicy,
    )
    _GRU_AVAILABLE = True
except ImportError:
    _GRU_AVAILABLE = False
    create_gru_weather_policy_kwargs = None  # type: ignore
    ZoneEquivariantMaskablePolicy = None  # type: ignore

# Physics dynamics is optional — Dyna augmentation disabled if unavailable.
# Import failure is silent: train_phase() runs identically to the original
# when dynamics_config=None or _DYNAMICS_AVAILABLE=False.
try:
    from physics_dynamics import TemporalDynamicsModel, DynaRolloutBuffer, ZoneStateTensor
    _DYNAMICS_AVAILABLE = True
except ImportError:
    _DYNAMICS_AVAILABLE = False
    TemporalDynamicsModel = None  # type: ignore
    DynaRolloutBuffer     = None  # type: ignore
    ZoneStateTensor       = None  # type: ignore


# ---------------------------------------------------------------------------
# Logging
# ---------------------------------------------------------------------------

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s | %(levelname)s | %(name)s | %(message)s",
    handlers=[
        logging.FileHandler("training.log"),
        logging.StreamHandler(),
    ],
)
logger = logging.getLogger(__name__)


# ---------------------------------------------------------------------------
# Device selection
# ---------------------------------------------------------------------------

def _select_device(requested: str) -> str:
    """Return 'cuda' if available and requested, else 'cpu'."""
    if requested == "cuda":
        if _TORCH_AVAILABLE and torch.cuda.is_available():
            return "cuda"
        logger.warning("CUDA requested but not available — falling back to CPU.")
        return "cpu"
    return requested


# ---------------------------------------------------------------------------
# Dynamics configuration
# ---------------------------------------------------------------------------

@dataclass
class DynamicsConfig:
    """
    Configuration for Dyna-style physics dynamics augmentation.

    When dynamics_model_path is set and the model file exists, DynaCallback
    loads the pre-trained TemporalDynamicsModel and adds a surprise bonus
    to the PPO reward at each step. When dynamics_model_path is None (default),
    training is identical to the original curriculum — no overhead, no change.

    Fields
    ------
    dynamics_model_path:
        Path to a pre-trained TemporalDynamicsModel checkpoint (.pt).
        Produced by DynamicsTrainer.save() in physics_dynamics.py.
        If None or the file does not exist, DynaCallback is not attached.

    surprise_weight:
        Scalar multiplier for the surprise bonus added to the PPO reward.
        Start at 0.05. Increase to 0.1 if the agent is under-exploring;
        decrease to 0.01 if the dynamics bonus dominates task reward.
        The bonus is clipped to [0, surprise_weight] before adding, so
        this value is also the maximum bonus per step.

    update_dynamics_every_n_steps:
        Fine-tune the dynamics model on transitions collected during RL
        training every N environment steps. 0 = no fine-tuning (frozen model).
        Fine-tuning closes the Dyna loop: better policy -> richer data ->
        better dynamics -> better policy. Start with 0 until baseline training
        is stable, then enable at 50_000 steps.

    fine_tune_epochs:
        Number of gradient steps per fine-tuning update. Keep low (3-5)
        to avoid overfitting to the most recent transitions.

    transition_buffer_size:
        Maximum number of (current, next) transition pairs stored for
        fine-tuning. Ring buffer: oldest pairs dropped when full.
    """
    dynamics_model_path:           Optional[str] = None
    surprise_weight:               float         = 0.05
    update_dynamics_every_n_steps: int           = 0       # 0 = frozen
    fine_tune_epochs:              int           = 3
    transition_buffer_size:        int           = 10_000


# ---------------------------------------------------------------------------
# Curriculum definition
# ---------------------------------------------------------------------------

def resolve_phase_max_steps(n_zones: int, budget_mode: str, episode_length: int) -> int:
    """
    Map budget_mode → max_steps for a curriculum phase.

    The visit-once mask makes the structural ceiling n_zones+1. Old phases
    used episode_length of 150–300, which never forced zone selection.
    budget_mode overrides that so later phases train allocation skill.

      full    → n_zones + 1
      scarce  → n_zones
      triage  → max(1, n_zones - 1)
      legacy  → keep episode_length (old behaviour)
    """
    n = max(1, int(n_zones))
    mode = (budget_mode or "triage").strip().lower()
    if mode == "legacy":
        return max(1, int(episode_length))
    if mode == "full":
        return n + 1
    if mode == "scarce":
        return n
    if mode == "triage":
        return max(1, n - 1)
    raise ValueError(f"Unknown budget_mode {budget_mode!r}")


@dataclass
class CurriculumPhase:
    name:           str
    total_steps:    int
    episode_length: int   # used only when budget_mode="legacy"
    n_zones:        int
    risk_weights:   RiskWeights
    # Default triage: no phase can pass without learning which zones to skip.
    budget_mode:    str = "triage"  # full | scarce | triage | legacy

    learning_rate:  float = 3e-4
    n_steps:        int   = 4_096
    batch_size:     int   = 256
    n_epochs:       int   = 10
    gamma:          float = 0.995
    gae_lambda:     float = 0.95
    clip_range:     float = 0.2
    ent_coef:       float = 0.02
    vf_coef:        float = 0.5
    max_grad_norm:  float = 0.5

    def resolved_max_steps(self) -> int:
        return resolve_phase_max_steps(self.n_zones, self.budget_mode, self.episode_length)


class WeatherCurriculum:
    """Five-phase climate curriculum from baseline through stress extremes.

    Budget progression (forces zone differentiation):
      normal   → full   (learn inspection has value; 2 zones)
      monsoon  → scarce (start leaving someone out; 3 zones)
      drought  → triage (must skip ≥1; 3 zones)
      heatwave → triage (4 zones)
      humidity → triage (4 zones)
    """

    PHASES: Dict[str, CurriculumPhase] = {

        "normal": CurriculumPhase(
            name="normal",
            total_steps=200_000,
            episode_length=150,
            n_zones=2,
            budget_mode="full",
            risk_weights=RiskWeights(),
            n_steps=4_096,
            ent_coef=0.05,
        ),

        "monsoon": CurriculumPhase(
            name="monsoon",
            total_steps=150_000,
            episode_length=300,
            n_zones=3,
            budget_mode="scarce",
            risk_weights=RiskWeights(
                drought_obs_weight=0.40, drought_forecast_weight=0.60,
                flood_obs_weight=0.70,   flood_forecast_weight=0.30,
                fungi_obs_weight=0.75,   fungi_forecast_weight=0.25,
                supply_drought_weight=0.25,
                supply_flood_weight=0.50,
                supply_harvest_pressure_weight=0.25,
            ),
            n_steps=4_096,
        ),

        "drought": CurriculumPhase(
            name="drought",
            total_steps=120_000,
            episode_length=250,
            n_zones=3,
            budget_mode="triage",
            risk_weights=RiskWeights(
                drought_obs_weight=0.80, drought_forecast_weight=0.20,
                flood_obs_weight=0.30,   flood_forecast_weight=0.70,
                fungi_obs_weight=0.55,   fungi_forecast_weight=0.45,
                supply_drought_weight=0.55,
                supply_flood_weight=0.25,
                supply_harvest_pressure_weight=0.20,
            ),
            n_steps=4_096,
        ),

        "heatwave": CurriculumPhase(
            name="heatwave",
            total_steps=120_000,
            episode_length=220,
            n_zones=4,
            budget_mode="triage",
            risk_weights=RiskWeights(
                drought_obs_weight=0.75, drought_forecast_weight=0.25,
                flood_obs_weight=0.25,   flood_forecast_weight=0.75,
                fungi_obs_weight=0.50,   fungi_forecast_weight=0.50,
                supply_drought_weight=0.60,
                supply_flood_weight=0.15,
                supply_harvest_pressure_weight=0.25,
            ),
            n_steps=4_096,
        ),

        "humidity": CurriculumPhase(
            name="humidity",
            total_steps=100_000,
            episode_length=200,
            n_zones=4,
            budget_mode="triage",
            risk_weights=RiskWeights(
                drought_obs_weight=0.30, drought_forecast_weight=0.70,
                flood_obs_weight=0.50,   flood_forecast_weight=0.50,
                fungi_obs_weight=0.85,   fungi_forecast_weight=0.15,
                supply_drought_weight=0.20,
                supply_flood_weight=0.30,
                supply_harvest_pressure_weight=0.50,
                quality_fungi_weight=0.80,
                quality_delay_weight=0.20,
            ),
            n_steps=4_096,
        ),
    }

    @classmethod
    def get_phase(cls, name: str) -> CurriculumPhase:
        if name not in cls.PHASES:
            raise ValueError(
                f"Unknown phase '{name}'. Options: {sorted(cls.PHASES)}"
            )
        return cls.PHASES[name]

    @classmethod
    def phase_order(cls) -> List[str]:
        return ["normal", "monsoon", "drought", "heatwave", "humidity"]


# ---------------------------------------------------------------------------
# Checkpoint callback
# ---------------------------------------------------------------------------

class CheckpointCallback(BaseCallback):
    """Save a checkpoint every `save_freq` timesteps."""

    def __init__(self, output_dir: Path, save_freq: int = 25_000) -> None:
        super().__init__()
        self.output_dir = output_dir
        self.save_freq  = save_freq
        self._last_save = 0

    def _on_step(self) -> bool:
        if self.num_timesteps - self._last_save >= self.save_freq:
            self._last_save = self.num_timesteps
            path = self.output_dir / f"checkpoint_{self.num_timesteps}.zip"
            self.model.save(str(path))
            logger.info("Checkpoint saved: %s", path.name)
        return True



# ---------------------------------------------------------------------------
# Dyna callback
# ---------------------------------------------------------------------------

class DynaCallback(BaseCallback):
    """
    Augments PPO rewards with a physics-dynamics surprise bonus (Dyna-style).

    At each environment step, this callback:
      1. Extracts the current and next observation as ZoneStateTensor objects.
      2. Calls DynaRolloutBuffer.compute_surprise_bonus() — the normalised
         prediction error of the dynamics model for this transition.
      3. Writes the bonus directly into the PPO rollout buffer at the position
         that was just written by env.step().
      4. Optionally fine-tunes the dynamics model on accumulated transitions.

    Reward injection mechanism
    --------------------------
    SB3's RolloutBuffer stores rewards at self.model.rollout_buffer.rewards[pos-1]
    immediately after env.step() returns, where pos is the buffer write pointer.
    The callback's _on_step() runs after that write, so we can read and modify
    the reward before any PPO computation sees it.

    The pos pointer advances BEFORE _on_step() is called, so the correct
    index is (self.model.rollout_buffer.pos - 1) % n_steps.

    This is the same approach used by SB3's RND and curiosity implementations.

    Safe degradation
    ----------------
    If the dynamics model is unavailable, or if obs keys are missing (e.g.
    during the first step of an episode), the callback returns True silently
    without modifying any reward. It never raises or interrupts training.

    Args:
        dyna_buffer:    DynaRolloutBuffer wrapping the loaded dynamics model.
        dynamics_cfg:   DynamicsConfig controlling weights and fine-tuning.
        n_zones:        Must match the environment's n_zones.
        horizon_days:   Must match ForecastConfig.horizon_days.
        device:         Torch device string for tensor ops.
    """

    _OBS_KEYS = ("forecast_precip", "forecast_uncertainty", "zone_belief")

    def __init__(
        self,
        dyna_buffer:  "DynaRolloutBuffer",
        dynamics_cfg: DynamicsConfig,
        n_zones:      int,
        horizon_days: int,
        device:       str = "cpu",
    ) -> None:
        super().__init__()
        self.dyna_buffer   = dyna_buffer
        self.dynamics_cfg  = dynamics_cfg
        self.n_zones       = n_zones
        self.horizon_days  = horizon_days
        self.device        = device

        # Ring buffer for fine-tuning transitions
        # Stored as (current_ZoneStateTensor, next_ZoneStateTensor) pairs
        self._transition_buffer: list = []
        self._tb_max = dynamics_cfg.transition_buffer_size

        # Statistics logged every 10k steps
        self._bonus_sum   = 0.0
        self._bonus_count = 0
        self._log_freq    = 10_000
        self._last_log    = 0

        # Previous obs for transition construction (obs_t → obs_t+1)
        self._prev_obs: Optional[dict] = None

    def _obs_to_state_tensor(self, obs: dict) -> Optional["ZoneStateTensor"]:
        """
        Convert a raw SB3 obs dict to ZoneStateTensor.

        SB3 stores observations as numpy arrays with a leading env-count
        dimension even for a single env: shape [1, ...]. We squeeze that dim.

        Returns None if any required key is missing (safe degradation).
        """
        if not all(k in obs for k in self._OBS_KEYS):
            return None

        import torch
        import numpy as np

        try:
            # SB3 obs shapes: [n_envs, ...] — squeeze env dim (n_envs=1)
            precip = np.array(obs["forecast_precip"],      dtype=np.float32)
            uncert = np.array(obs["forecast_uncertainty"], dtype=np.float32)
            belief = np.array(obs["zone_belief"],          dtype=np.float32)

            # Handle both [1, n_zones, H] and [n_zones, H] shapes gracefully
            if precip.ndim == 2:
                precip = precip[np.newaxis]   # [n_zones, H] -> [1, n_zones, H]
            if uncert.ndim == 1:
                uncert = uncert[np.newaxis]   # [n_zones]    -> [1, n_zones]
            if belief.ndim == 1:
                belief = belief[np.newaxis]

            return ZoneStateTensor(
                precip=torch.from_numpy(precip).to(self.device),
                uncertainty=torch.from_numpy(uncert).to(self.device),
                belief=torch.from_numpy(belief).to(self.device),
            )
        except Exception as e:
            logger.debug("DynaCallback._obs_to_state_tensor failed: %s", e)
            return None

    def _on_step(self) -> bool:
        """
        Called after every env.step(). Injects surprise bonus into reward buffer.
        """
        # --- Extract current and next observations ---
        # self.locals["obs_tensor"] is the obs BEFORE the step (obs_t).
        # self.locals["new_obs"] is the obs AFTER the step (obs_t+1).
        # Both are available in SB3 >= 1.8 on_step locals.
        try:
            obs_now  = self.locals.get("obs_tensor") or self.locals.get("obs")
            obs_next = self.locals.get("new_obs")

            if obs_now is None or obs_next is None:
                return True  # safe: missing locals, skip silently

            # Convert to ZoneStateTensor
            if hasattr(obs_now, "numpy"):
                # Tensor: convert dict-of-tensors or single tensor
                obs_now_np  = {k: v.cpu().numpy() for k, v in obs_now.items()}                               if hasattr(obs_now, "items") else {"_raw": obs_now.cpu().numpy()}
            else:
                obs_now_np = obs_now

            if hasattr(obs_next, "items"):
                obs_next_np = {k: (v.cpu().numpy() if hasattr(v, "cpu") else v)
                               for k, v in obs_next.items()}
            else:
                obs_next_np = obs_next

            curr_state = self._obs_to_state_tensor(obs_now_np)
            next_state = self._obs_to_state_tensor(obs_next_np)

            if curr_state is None or next_state is None:
                return True  # safe: obs keys not present yet

            # --- Compute surprise bonus ---
            bonus = self.dyna_buffer.compute_surprise_bonus(curr_state, next_state)
            bonus_val = float(bonus.item())
            bonus_clipped = min(bonus_val, self.dynamics_cfg.surprise_weight)

            # --- Inject into PPO rollout buffer ---
            # The rollout buffer pos pointer has already advanced; the reward
            # for the current step is at (pos - 1) % n_steps.
            rb = self.model.rollout_buffer
            if rb is not None and hasattr(rb, "rewards") and rb.rewards is not None:
                idx = (rb.pos - 1) % rb.buffer_size
                rb.rewards[idx] += bonus_clipped

            # --- Accumulate for fine-tuning ---
            if self.dynamics_cfg.update_dynamics_every_n_steps > 0:
                self._transition_buffer.append((curr_state, next_state))
                if len(self._transition_buffer) > self._tb_max:
                    self._transition_buffer.pop(0)  # ring buffer: drop oldest

            # --- Statistics ---
            self._bonus_sum   += bonus_clipped
            self._bonus_count += 1

            if self.num_timesteps - self._last_log >= self._log_freq:
                avg_bonus = (
                    self._bonus_sum / self._bonus_count
                    if self._bonus_count > 0 else 0.0
                )
                logger.info(
                    "DynaCallback: step=%d  avg_surprise_bonus=%.4f  "
                    "buffer_size=%d",
                    self.num_timesteps, avg_bonus,
                    len(self._transition_buffer),
                )
                self._bonus_sum   = 0.0
                self._bonus_count = 0
                self._last_log    = self.num_timesteps

        except Exception as e:
            # Never interrupt training on callback error — log and continue
            logger.debug("DynaCallback._on_step error (non-fatal): %s", e)

        return True

    def _on_rollout_end(self) -> None:
        """
        Called at the end of each rollout collection. Optionally fine-tunes
        the dynamics model on accumulated transitions.
        """
        if (
            self.dynamics_cfg.update_dynamics_every_n_steps <= 0
            or self.num_timesteps % self.dynamics_cfg.update_dynamics_every_n_steps != 0
            or len(self._transition_buffer) < 16  # need at least one batch
        ):
            return

        try:
            from physics_dynamics import DynamicsTrainer
            # Access the dynamics model directly from the buffer
            dynamics_model = self.dyna_buffer.dynamics

            # Minimal fine-tune: a few gradient steps on recent transitions
            # We construct a temporary DynamicsTrainer around the existing model
            # rather than creating a new one, to avoid re-initialising weights.
            import torch
            import torch.nn.functional as F

            optimizer = torch.optim.AdamW(
                dynamics_model.parameters(), lr=1e-4, weight_decay=1e-4
            )
            dynamics_model.train()

            pairs = list(self._transition_buffer)  # snapshot
            batch_size = min(32, len(pairs))

            for epoch in range(self.dynamics_cfg.fine_tune_epochs):
                import random
                random.shuffle(pairs)
                total_loss = 0.0
                n_batches  = 0

                for i in range(0, len(pairs) - batch_size, batch_size):
                    batch = pairs[i : i + batch_size]
                    curr_list = [p[0] for p in batch]
                    next_list = [p[1] for p in batch]

                    import torch as _t
                    # Stack batch dimension
                    curr_b = ZoneStateTensor(
                        precip=_t.cat([s.precip for s in curr_list], dim=0),
                        uncertainty=_t.cat([s.uncertainty for s in curr_list], dim=0),
                        belief=_t.cat([s.belief for s in curr_list], dim=0),
                    )
                    next_b = ZoneStateTensor(
                        precip=_t.cat([s.precip for s in next_list], dim=0),
                        uncertainty=_t.cat([s.uncertainty for s in next_list], dim=0),
                        belief=_t.cat([s.belief for s in next_list], dim=0),
                    )

                    pred, phys_loss = dynamics_model(curr_b, return_physics_loss=True)
                    data_loss = (
                        F.mse_loss(pred.precip / 500.0, next_b.precip / 500.0)
                        + F.mse_loss(pred.uncertainty, next_b.uncertainty)
                        + F.mse_loss(pred.belief,      next_b.belief)
                    )
                    loss = data_loss + 0.01 * phys_loss

                    optimizer.zero_grad()
                    loss.backward()
                    _t.nn.utils.clip_grad_norm_(dynamics_model.parameters(), 1.0)
                    optimizer.step()

                    total_loss += loss.item()
                    n_batches  += 1

            dynamics_model.eval()
            logger.info(
                "DynaCallback: fine-tuned dynamics model at step=%d  "
                "avg_loss=%.4f  n_transitions=%d",
                self.num_timesteps,
                total_loss / max(n_batches, 1),
                len(self._transition_buffer),
            )

        except Exception as e:
            logger.warning(
                "DynaCallback._on_rollout_end fine-tune failed (non-fatal): %s", e
            )


def _build_dyna_callback(
    dynamics_cfg: Optional[DynamicsConfig],
    n_zones:      int,
    horizon_days: int,
    device:       str,
) -> Optional["DynaCallback"]:
    """
    Build a DynaCallback if dynamics are configured and available.

    Returns None (no Dyna augmentation) if:
      - dynamics_cfg is None
      - dynamics_model_path is not set
      - the model file does not exist
      - physics_dynamics module is unavailable
      - any load/init error occurs

    Callers pass the return value directly to CallbackList — None is ignored.
    """
    if dynamics_cfg is None or dynamics_cfg.dynamics_model_path is None:
        return None

    if not _DYNAMICS_AVAILABLE:
        logger.warning(
            "DynamicsConfig provided but physics_dynamics not installed — "
            "Dyna augmentation disabled."
        )
        return None

    model_path = Path(dynamics_cfg.dynamics_model_path)
    if not model_path.exists():
        logger.warning(
            "Dynamics model not found at %s — Dyna augmentation disabled.",
            model_path,
        )
        return None

    try:
        dynamics_model = TemporalDynamicsModel.load(str(model_path))
        dynamics_model.eval()
        dyna_buffer = DynaRolloutBuffer(
            dynamics=dynamics_model,
            uncertainty_weight=dynamics_cfg.surprise_weight,
        )
        callback = DynaCallback(
            dyna_buffer=dyna_buffer,
            dynamics_cfg=dynamics_cfg,
            n_zones=n_zones,
            horizon_days=horizon_days,
            device=device,
        )
        logger.info(
            "DynaCallback loaded: model=%s  surprise_weight=%.3f  "
            "fine_tune_every=%d",
            model_path.name,
            dynamics_cfg.surprise_weight,
            dynamics_cfg.update_dynamics_every_n_steps,
        )
        return callback

    except Exception as e:
        logger.warning(
            "Failed to build DynaCallback (%s) — Dyna augmentation disabled.", e
        )
        return None



# ---------------------------------------------------------------------------
# Training
# ---------------------------------------------------------------------------

def transfer_curriculum_weights(
    resume_from: str,
    model: "MaskablePPO",
    device: str = "auto",
) -> "MaskablePPO":
    """
    Warm-start `model` (freshly constructed for the CURRENT phase's env/n_zones)
    from `resume_from`'s checkpoint, transferring every parameter whose shape
    matches exactly and leaving the rest at fresh random initialization.

    Exists because MaskablePPO.load(path, env=new_env) raises "Observation
    spaces do not match" whenever n_zones changes between curriculum phases
    -- a hard SB3-level space-equality check that fires before any weight-
    shape question is even considered. Loading without `env=` sidesteps that
    (the checkpoint reconstructs against its own saved spaces); this function
    then transfers whatever's compatible directly via the two state_dicts.

    With the permutation-invariant GRUWeatherFeaturesExtractor (see
    gru_weather_policy.py), every parameter except the action_net output
    layer (shape tied to n_zones+1, the discrete action count) now matches
    across any n_zones -- verified empirically at 61/63 tensors transferred
    in a 2-zone -> 3-zone test. value_net transfers too (scalar output,
    always n_zones-independent); only action_net needs relearning.
    """
    old_model = MaskablePPO.load(resume_from, device=device)
    old_state = old_model.policy.state_dict()
    new_state = model.policy.state_dict()

    transferred, skipped = [], []
    merged = {}
    for key, new_tensor in new_state.items():
        old_tensor = old_state.get(key)
        if old_tensor is not None and old_tensor.shape == new_tensor.shape:
            merged[key] = old_tensor.clone()
            transferred.append(key)
        else:
            merged[key] = new_tensor
            skipped.append(key)

    model.policy.load_state_dict(merged)

    logger.info(
        "transfer_curriculum_weights: transferred %d/%d parameter tensors from %s "
        "(freshly initialized: %s)",
        len(transferred), len(new_state), resume_from, skipped or "none",
    )
    if not transferred:
        logger.warning(
            "transfer_curriculum_weights: transferred ZERO parameters -- the "
            "architectures are likely genuinely incompatible (e.g. resuming "
            "from a pre-permutation-invariant checkpoint), not just a normal "
            "n_zones change. Check resume_from's origin before trusting this run."
        )
    return model


def train_phase(
    phase_name:     str,
    output_dir:     Path,
    resume_from:    Optional[str]           = None,
    override_steps: Optional[int]           = None,
    hidden_size:    int                     = 64,
    device:         str                     = "auto",
    seed:           int                     = 42,
    dynamics_cfg:   Optional[DynamicsConfig] = None,
    precip_scale:      float                = 40.0,
) -> str:
    """Train one curriculum phase. Returns path to the saved final model.

    Args:
        phase_name:    One of the WeatherCurriculum phase names.
        output_dir:    Root directory for checkpoints and final model.
        resume_from:   Path to a previous phase checkpoint to resume from.
        override_steps: Override total_steps (useful for quick tests).
        hidden_size:   GRU hidden size; must match any checkpoint being resumed.
        device:        'cpu', 'cuda', or 'auto'.
        seed:          Random seed.
        dynamics_cfg:  Optional DynamicsConfig for Dyna surprise-bonus augmentation.
                       Pass None (default) for standard training with no overhead.
        precip_scale:  Fixed (non-learned) divisor applied to forecast_precip
                       before the GRU extractor. See train_kaggle.py's
                       INPUT NORMALIZATION docstring section for why this
                       exists. Default 40.0 matches the value that produced
                       the validated single-dirty selection-accuracy result
                       (see model card) -- confirmed working, not confirmed
                       optimal, and not yet validated across curriculum
                       phase transitions specifically (only within a single
                       train_kaggle.py run). Must match across resumed
                       checkpoints the same way hidden_size must.
    """

    if not _ML_AVAILABLE:
        raise RuntimeError(
            f"ML stack not available: {_ML_IMPORT_ERROR}\n"
            "Install: pip install stable-baselines3 sb3-contrib torch"
        )

    output_dir.mkdir(parents=True, exist_ok=True)
    models_dir = output_dir / "models"
    models_dir.mkdir(exist_ok=True)

    device = _select_device(
        device if device != "auto"
        else ("cuda" if _TORCH_AVAILABLE and torch.cuda.is_available() else "cpu")
    )

    phase = WeatherCurriculum.get_phase(phase_name)
    total_steps = override_steps or phase.total_steps
    max_steps = phase.resolved_max_steps()
    full_ceiling = phase.n_zones + 1

    logger.info(
        "Phase=%s  steps=%d  max_steps=%d (budget_mode=%s, full_ceiling=%d)  "
        "n_zones=%d  device=%s  must_skip=%s",
        phase.name, total_steps, max_steps, phase.budget_mode, full_ceiling,
        phase.n_zones, device,
        "yes" if max_steps < full_ceiling else "no",
    )
    if max_steps >= full_ceiling and phase.budget_mode not in ("full", "legacy"):
        logger.warning(
            "Phase %s: max_steps=%d >= full_ceiling=%d despite budget_mode=%s — "
            "check resolve_phase_max_steps.",
            phase.name, max_steps, full_ceiling, phase.budget_mode,
        )

    # --- Environment ---
    config = ForecastConfig(
        n_zones=phase.n_zones,
        seed=seed,
        soft_reset=True,
        max_steps=max_steps,
    )
    phase.risk_weights.attach_to_config(config)

    # Monitor wraps correctly: get_wrapper_attr('action_masks') walks the
    # wrapper stack and finds action_masks() on WeatherForecastEnv.
    env = Monitor(make_weather_env(config))

    # --- Policy kwargs ---
    if _GRU_AVAILABLE:
        policy_kwargs = create_gru_weather_policy_kwargs(
            hidden_size=hidden_size,
            precip_scale=precip_scale,
        )
        # ZoneEquivariantMaskablePolicy, NOT the "MultiInputPolicy" string.
        # The default policy builds action logits from action_net(latent_pi)
        # on top of the extractor's pooled (permutation-invariant) feature
        # vector -- zone identity is erased before the action head ever
        # sees it, so under triage the policy can only express a static
        # per-slot bias (empirically: always inspects action index 0,
        # regardless of which zone's content actually looks risky). This
        # was an active bug in this function: GRU features were used, but
        # every prior curriculum-trained checkpoint went through the same
        # pooled action_net as the plain-MLP fallback below and could not
        # have learned risk-conditioned zone selection. See
        # gru_weather_policy.py module docstring and train_kaggle.py's
        # POLICY section for the full diagnosis.
        policy = ZoneEquivariantMaskablePolicy
        logger.info(
            "Using zone-equivariant GRU policy (hidden_size=%d, "
            "precip_scale=%.3g)",
            hidden_size, precip_scale,
        )
    else:
        # dict with 'pi'/'vf' keys is the correct net_arch format for
        # MultiInputPolicy in SB3 >= 1.8 (validated on SB3 2.8.0)
        policy_kwargs = dict(net_arch=dict(pi=[128, 64], vf=[128, 64]))
        policy = "MultiInputPolicy"
        logger.info("GRU policy unavailable — using MLP policy (net_arch=128,64)")

    # --- Model ---
    ppo_kwargs = dict(
        learning_rate=phase.learning_rate,
        n_steps=phase.n_steps,
        batch_size=phase.batch_size,
        n_epochs=phase.n_epochs,
        gamma=phase.gamma,
        gae_lambda=phase.gae_lambda,
        clip_range=phase.clip_range,
        ent_coef=phase.ent_coef,
        vf_coef=phase.vf_coef,
        max_grad_norm=phase.max_grad_norm,
        device=device,
        verbose=1,
        seed=seed,
    )

    if resume_from:
        logger.info("Resuming from %s", resume_from)
        # Build a FRESH model for the CURRENT phase's env/n_zones (correct
        # observation/action spaces throughout), then warm-start it from the
        # checkpoint wherever shapes match. MaskablePPO.load(resume_from,
        # env=env) directly would raise "Observation spaces do not match"
        # the moment n_zones changes between phases -- see
        # transfer_curriculum_weights()'s docstring for why, and what
        # actually transfers (everything except action_net).
        model = MaskablePPO(
            policy=policy,
            env=env,
            policy_kwargs=policy_kwargs,
            **ppo_kwargs,
        )
        model = transfer_curriculum_weights(resume_from, model, device=device)
        reset_timesteps = False
    else:
        model = MaskablePPO(
            policy=policy,
            env=env,
            policy_kwargs=policy_kwargs,
            **ppo_kwargs,
        )
        reset_timesteps = True

    # --- Callbacks ---
    from stable_baselines3.common.callbacks import CallbackList
    callbacks = [CheckpointCallback(output_dir)]

    horizon_days = getattr(config, "horizon_days", 14)
    dyna_cb = _build_dyna_callback(
        dynamics_cfg=dynamics_cfg,
        n_zones=phase.n_zones,
        horizon_days=horizon_days,
        device=device,
    )
    if dyna_cb is not None:
        callbacks.append(dyna_cb)
        logger.info("Dyna augmentation active for phase=%s", phase.name)
    else:
        logger.info("Dyna augmentation inactive for phase=%s", phase.name)

    # --- Train ---
    # use_masking=True is the default; MaskablePPO calls action_masks()
    # automatically via get_wrapper_attr during rollout collection.
    # Do NOT pass action_masks= to learn() — it is not a valid parameter.
    model.learn(
        total_timesteps=total_steps,
        callback=CallbackList(callbacks),
        reset_num_timesteps=reset_timesteps,
        use_masking=True,
    )

    final_path = models_dir / f"final_{phase.name}.zip"
    model.save(str(final_path))
    logger.info("Saved final model: %s", final_path)

    return str(final_path)


def train_full_curriculum(
    output_dir:   Path,
    device:       str                      = "auto",
    seed:         int                      = 42,
    dynamics_cfg: Optional[DynamicsConfig] = None,
    precip_scale:      float               = 40.0,
) -> None:
    """Run all phases in order, chaining each phase from the previous."""
    phases = WeatherCurriculum.phase_order()
    resume = None
    for phase_name in phases:
        logger.info("=== Starting phase: %s ===", phase_name)
        resume = train_phase(
            phase_name=phase_name,
            output_dir=output_dir / phase_name,
            resume_from=resume,
            device=device,
            seed=seed,
            dynamics_cfg=dynamics_cfg,
            precip_scale=precip_scale,
        )
        logger.info("=== Completed phase: %s ===", phase_name)


# ---------------------------------------------------------------------------
# CLI
# ---------------------------------------------------------------------------

def main() -> None:
    p = argparse.ArgumentParser(
        description="MaskablePPO curriculum trainer for WeatherForecastEnv"
    )
    p.add_argument(
        "--phase",
        choices=list(WeatherCurriculum.PHASES) + ["all"],
        default="normal",
        help="Curriculum phase to run, or 'all' to run full curriculum.",
    )
    p.add_argument("--output-dir",   default="./run",  help="Root output directory")
    p.add_argument("--resume-from",  default=None,     help="Path to checkpoint .zip")
    p.add_argument("--steps",        type=int, default=None, help="Override total_steps")
    p.add_argument("--hidden-size",  type=int, default=64)
    p.add_argument("--device",       default="auto",   help="'cpu', 'cuda', or 'auto'")
    p.add_argument("--seed",                  type=int,   default=42)
    p.add_argument(
        "--dynamics-model",
        default=None,
        help="Path to pre-trained TemporalDynamicsModel .pt file. Enables Dyna augmentation.",
    )
    p.add_argument(
        "--dynamics-weight",
        type=float,
        default=0.05,
        help="Surprise bonus weight per step (only used with --dynamics-model). Default 0.05.",
    )
    p.add_argument(
        "--dynamics-finetune-every",
        type=int,
        default=0,
        help="Fine-tune dynamics model every N steps. 0=frozen (default).",
    )
    p.add_argument(
        "--precip-scale",
        type=float,
        default=40.0,
        help="Fixed (non-learned) divisor applied to forecast_precip before "
             "the GRU extractor. See train_kaggle.py's INPUT NORMALIZATION "
             "docstring section for the full rationale. 40.0 matches the "
             "value that produced the validated single-dirty "
             "selection-accuracy result (see model card). Must stay "
             "consistent across a resumed checkpoint's phases, the same "
             "way --hidden-size must.",
    )
    args = p.parse_args()

    output_dir = Path(args.output_dir)

    dynamics_cfg: Optional[DynamicsConfig] = None
    if args.dynamics_model is not None:
        dynamics_cfg = DynamicsConfig(
            dynamics_model_path=args.dynamics_model,
            surprise_weight=args.dynamics_weight,
            update_dynamics_every_n_steps=args.dynamics_finetune_every,
        )
        logger.info(
            "Dyna config: model=%s  weight=%.3f  finetune_every=%d",
            args.dynamics_model, args.dynamics_weight, args.dynamics_finetune_every,
        )

    if args.phase == "all":
        train_full_curriculum(
            output_dir=output_dir,
            device=args.device,
            seed=args.seed,
            dynamics_cfg=dynamics_cfg,
            precip_scale=args.precip_scale,
        )
    else:
        train_phase(
            phase_name=args.phase,
            output_dir=output_dir,
            resume_from=args.resume_from,
            override_steps=args.steps,
            hidden_size=args.hidden_size,
            device=args.device,
            seed=args.seed,
            dynamics_cfg=dynamics_cfg,
            precip_scale=args.precip_scale,
        )


if __name__ == "__main__":
    main()