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
    MaskablePPO     = None
    Monitor         = None
    BaseCallback    = object

try:
    from gru_weather_policy import create_gru_weather_policy_kwargs, get_equivariant_policy_class
    _GRU_AVAILABLE = True
except ImportError:
    _GRU_AVAILABLE = False
    create_gru_weather_policy_kwargs = None
    get_equivariant_policy_class = None

try:
    from physics_dynamics import TemporalDynamicsModel, DynaRolloutBuffer, ZoneStateTensor
    _DYNAMICS_AVAILABLE = True
except ImportError:
    _DYNAMICS_AVAILABLE = False
    TemporalDynamicsModel = None
    DynaRolloutBuffer     = None
    ZoneStateTensor       = None


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

def set_global_seeds(seed: int) -> None:
    import random
    import numpy as np
    random.seed(seed)
    np.random.seed(seed)
    try:
        import torch
        torch.manual_seed(seed)
        if torch.cuda.is_available():
            torch.cuda.manual_seed_all(seed)
    except ImportError:
        pass




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

def _select_device(requested: str) -> str:
    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:
    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:
    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
    n_zones:        int
    risk_weights:   RiskWeights
    budget_mode:    str = "triage"

    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:

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

    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


class RealDataUsageCallback(BaseCallback):
    """Confirms at runtime -- not just at startup -- what fraction of
    episodes actually drew real vs synthetic data. info["context_source"]
    (set by WeatherForecastEnv.reset(), one of "injected"/"real_sampled"/
    "synthetic") is only present on the step where an episode boundary
    triggered an autoreset, so this counts opportunistically rather than
    on every step. A silent zero real_sampled count with real_data_pkl_path
    configured is exactly the kind of thing that goes unnoticed for a long
    time otherwise -- this exists so it can't.
    """

    def __init__(self, log_every: int = 10_000) -> None:
        super().__init__()
        self.log_every = log_every
        self._counts: Dict[str, int] = {}
        self._last_log = 0

    def _on_step(self) -> bool:
        infos = self.locals.get("infos")
        if infos:
            for info in infos:
                src = info.get("context_source") if isinstance(info, dict) else None
                if src is not None:
                    self._counts[src] = self._counts.get(src, 0) + 1
        if self.num_timesteps - self._last_log >= self.log_every:
            total = sum(self._counts.values())
            if total > 0:
                logger.info(
                    "RealDataUsageCallback: step=%d episode context_source "
                    "counts so far: %s", self.num_timesteps, dict(self._counts),
                )
            self._last_log = self.num_timesteps
        return True

    def _on_training_end(self) -> None:
        total = sum(self._counts.values())
        if total == 0:
            logger.warning(
                "RealDataUsageCallback: never observed a context_source in "
                "any info dict this run -- could not confirm real-vs-synthetic "
                "mix (this can happen with some VecEnv wrapping; it does not "
                "necessarily mean sampling failed, but it also cannot confirm "
                "it succeeded -- verify with a direct env.reset() smoke test "
                "if this matters for the run)."
            )
            return
        logger.info(
            "RealDataUsageCallback: final episode context_source counts: "
            "%s (%d total, real fraction=%.3f)",
            dict(self._counts), total,
            self._counts.get("real_sampled", 0) / total,
        )



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

class DynaCallback(BaseCallback):

    _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

        self._transition_buffer: list = []
        self._tb_max = dynamics_cfg.transition_buffer_size

        self._bonus_sum   = 0.0
        self._bonus_count = 0
        self._log_freq    = 10_000
        self._last_log    = 0

        self._prev_obs: Optional[dict] = None

    def _obs_to_state_tensor(self, obs: dict) -> Optional["ZoneStateTensor"]:
        if not all(k in obs for k in self._OBS_KEYS):
            return None

        import torch
        import numpy as np

        try:
            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)

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

            if hasattr(obs_now, "numpy"):
                if hasattr(obs_now, "items"):
                    obs_now_np = {k: v.cpu().numpy() for k, v in obs_now.items()}
                else:
                    obs_now_np = {"_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

            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)

            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

            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)

            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:
            logger.debug("DynaCallback._on_step error (non-fatal): %s", e)

        return True

    def _on_rollout_end(self) -> None:
        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
        ):
            return

        try:
            from physics_dynamics import DynamicsTrainer
            dynamics_model = self.dyna_buffer.dynamics

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

                # FIX: iterate over all pairs, including the final partial batch
                for i in range(0, len(pairs), 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
                    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"]:
    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:
        # FIX: load onto the same device as training to avoid CPU/CUDA mismatch
        import torch as _torch
        dynamics_model = TemporalDynamicsModel.load(
            str(model_path), device=_torch.device(device)
        )
        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":
    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,
    real_data_pkl:         Optional[str]   = None,
    real_data_ratio:       Optional[float] = None,
    real_data_noise_scale: Optional[float] = None,
) -> str:

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

    config_kwargs: Dict[str, object] = dict(
        n_zones=phase.n_zones,
        seed=seed,
        soft_reset=True,
        max_steps=max_steps,
        real_data_pkl_path=real_data_pkl,
        inject_noise=bool(real_data_pkl),
    )
    if real_data_ratio is not None:
        config_kwargs["real_data_ratio"] = real_data_ratio
    if real_data_noise_scale is not None:
        config_kwargs["noise_scale"] = real_data_noise_scale
    config = ForecastConfig(**config_kwargs)
    # Train cube v4 TIGGE control: precip_mm length 15.
    # Holdout persistence rows are still horizon 30. Do NOT copy this
    # into evaluate_checkpoint_real / backtest_indonesia.
    if config.real_data_pkl_path:
        config.horizon_days = 15
    phase.risk_weights.attach_to_config(config)

    if config.real_data_pkl_path:
        logger.info(
            "=" * 72 + "\n"
            "REAL-DATA TRAINING ENABLED for phase=%s\n"
            "  pkl              = %s\n"
            "  real_data_ratio  = %.2f\n"
            "  inject_noise     = %s   noise_scale = %.3f\n"
            "  holdout          = enforced by RealEpisodeIndex "
            "(real_episode_sampler.DEFAULT_HOLDOUT_RANGES)\n"
            + "=" * 72,
            phase.name, config.real_data_pkl_path, config.real_data_ratio,
            config.inject_noise, config.noise_scale,
        )
    else:
        logger.warning(
            "=" * 72 + "\n"
            "Phase=%s is 100%% SYNTHETIC -- no --real-data-pkl was provided.\n"
            "Pass --real-data-pkl <path-to-cache.pkl> to change that.\n"
            + "=" * 72,
            phase.name,
        )

    env = Monitor(make_weather_env(config))

    # FIX: use the zone-equivariant policy class instead of the generic string
    if _GRU_AVAILABLE:
        # See train_kaggle.py's build_model() for the full explanation --
        # confirmed empirically this session: features_dim must equal
        # phase.n_zones * hidden_size * 2 to match GRUWeatherFeaturesExtractor's
        # actual output shape, or policy construction crashes on the first
        # forward pass. The factory's own default (128, no n_zones scaling)
        # only happens to be correct for n_zones == 1.
        features_dim = phase.n_zones * hidden_size * 2
        policy_kwargs = create_gru_weather_policy_kwargs(
            hidden_size=hidden_size, features_dim=features_dim,
        )
        policy = get_equivariant_policy_class()
        logger.info(
            "Using GRU policy (hidden_size=%d, n_zones=%d, features_dim=%d)",
            hidden_size, phase.n_zones, features_dim,
        )
    else:
        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)")

    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)
        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), RealDataUsageCallback()]

    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)

    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,
    real_data_pkl:         Optional[str]   = None,
    real_data_ratio:       Optional[float] = None,
    real_data_noise_scale: Optional[float] = None,
) -> 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,
            real_data_pkl=real_data_pkl,
            real_data_ratio=real_data_ratio,
            real_data_noise_scale=real_data_noise_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(
        "--real-data-pkl", default=None,
        help=(
            "Path to a historical trajectory cache (see real_episode_sampler.py "
            "/ RealEpisodeIndex). When set, applies to every phase run this "
            "invocation (including --phase all): training samples real "
            "historical episodes with probability --real-data-ratio each "
            "reset, respecting RealEpisodeIndex's built-in eval-window "
            "holdout. When NOT set (the default), training is 100%% "
            "synthetic -- this was true of every checkpoint in this project "
            "before 2026-08-24; see technical_details.md. Any cache file "
            "works here regardless of name -- naming convention is "
            "documentation only, see real_episode_sampler.py."
        ),
    )
    p.add_argument(
        "--real-data-ratio", type=float, default=None,
        help=(
            "Per-episode probability of sampling real data when "
            "--real-data-pkl is set. Omit to use "
            "ForecastConfig.real_data_ratio's own default (0.7)."
        ),
    )
    p.add_argument(
        "--real-data-noise-scale", type=float, default=None,
        help=(
            "Perturbation magnitude applied to sampled real data. Omit to "
            "use ForecastConfig.noise_scale's own default (0.05). Noise "
            "injection is enabled automatically whenever --real-data-pkl "
            "is set."
        ),
    )
    args = p.parse_args()
    set_global_seeds(args.seed)
    logger.info("Global seeds set to %s (Python / NumPy / PyTorch)", args.seed)

    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,
            real_data_pkl=args.real_data_pkl,
            real_data_ratio=args.real_data_ratio,
            real_data_noise_scale=args.real_data_noise_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,
            real_data_pkl=args.real_data_pkl,
            real_data_ratio=args.real_data_ratio,
            real_data_noise_scale=args.real_data_noise_scale,
        )


if __name__ == "__main__":
    main()