physics_dynamics.py

"""
physics_dynamics.py
===================
Physics-informed dynamics model for WeatherForecastEnv.

Architecture
------------
This is a Dyna-style learned dynamics model: trained offline on ERA5 data,
then used during PPO training to generate synthetic rollouts that augment
real environment experience. It does NOT replace the environment — it
supplements it, improving sample efficiency and generalization.

The model predicts how the zone-level forecast state evolves over time,
constrained by an advection-diffusion PDE residual that prevents physically
impossible predictions (e.g. precipitation materialising from nothing,
uncertainty decreasing without new observations).
"""

from __future__ import annotations

import logging
from dataclasses import dataclass
from pathlib import Path
from typing import List, Optional, Tuple

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader, TensorDataset

logger = logging.getLogger(__name__)


# ---------------------------------------------------------------------------
# Data container
# ---------------------------------------------------------------------------

@dataclass
class ZoneStateTensor:
    precip: torch.Tensor       # [batch, n_zones, horizon_days]
    uncertainty: torch.Tensor  # [batch, n_zones]
    belief: torch.Tensor       # [batch, n_zones]

    @property
    def batch_size(self) -> int:
        return self.precip.shape[0]

    @property
    def n_zones(self) -> int:
        return self.precip.shape[1]

    @property
    def horizon_days(self) -> int:
        return self.precip.shape[2]

    def to(self, device: torch.device) -> "ZoneStateTensor":
        return ZoneStateTensor(
            precip=self.precip.to(device),
            uncertainty=self.uncertainty.to(device),
            belief=self.belief.to(device),
        )

    def flat(self) -> torch.Tensor:
        B, Z, H = self.precip.shape
        precip_flat = self.precip.reshape(B, Z * H)
        return torch.cat([precip_flat, self.uncertainty, self.belief], dim=-1)

    @property
    def flat_dim(self) -> int:
        return self.n_zones * (self.horizon_days + 2)

    @classmethod
    def from_numpy(
        cls,
        precip: np.ndarray,
        uncertainty: np.ndarray,
        belief: np.ndarray,
    ) -> "ZoneStateTensor":
        if precip.ndim == 2:
            precip = precip[None]
        if uncertainty.ndim == 1:
            uncertainty = uncertainty[None]
        if belief.ndim == 1:
            belief = belief[None]
        return cls(
            precip=torch.from_numpy(precip.astype(np.float32)),
            uncertainty=torch.from_numpy(uncertainty.astype(np.float32)),
            belief=torch.from_numpy(belief.astype(np.float32)),
        )


# ---------------------------------------------------------------------------
# Physics residual
# ---------------------------------------------------------------------------

class PhysicsResidualLoss(nn.Module):

    def __init__(self, weight: float = 0.01):
        super().__init__()
        self.weight = weight
        self.log_v = nn.Parameter(torch.tensor(0.0))   # exp(0) = 1.0
        self.log_D = nn.Parameter(torch.tensor(-2.3))  # exp(-2.3) ≈ 0.1

    @property
    def v(self) -> torch.Tensor:
        return torch.exp(self.log_v)

    @property
    def D(self) -> torch.Tensor:
        return torch.exp(self.log_D)

    def forward(
        self,
        u_current: torch.Tensor,   # [batch, n_zones, horizon_days]
        u_next: torch.Tensor,      # [batch, n_zones, horizon_days]
        dt: float = 1.0,
    ) -> torch.Tensor:
        B, Z, H = u_current.shape

        # ∂u/∂t ≈ (u_next - u_current) / dt
        du_dt = (u_next - u_current) / dt

        # ∂u/∂τ — first derivative along horizon axis (central differences)
        # Shape: [batch, n_zones, horizon_days]
        du_dtau = torch.zeros_like(u_current)
        if H > 2:
            du_dtau[:, :, 1:-1] = (u_current[:, :, 2:] - u_current[:, :, :-2]) / 2.0
            du_dtau[:, :, 0]    = u_current[:, :, 1] - u_current[:, :, 0]
            du_dtau[:, :, -1]   = u_current[:, :, -1] - u_current[:, :, -2]

        # ∂²u/∂τ² — second derivative along horizon axis (Laplacian)
        d2u_dtau2 = torch.zeros_like(u_current)
        if H > 2:
            d2u_dtau2[:, :, 1:-1] = (
                u_current[:, :, 2:] - 2 * u_current[:, :, 1:-1] + u_current[:, :, :-2]
            )
            d2u_dtau2[:, :, 0]  = d2u_dtau2[:, :, 1]
            d2u_dtau2[:, :, -1] = d2u_dtau2[:, :, -2]

        # PDE residual: ∂u/∂t + v·∂u/∂τ - D·∂²u/∂τ² = 0
        residual = du_dt + self.v * du_dtau - self.D * d2u_dtau2

        return self.weight * torch.mean(residual ** 2)


# ---------------------------------------------------------------------------
# Core dynamics model
# ---------------------------------------------------------------------------

class TemporalDynamicsModel(nn.Module):

    def __init__(
        self,
        n_zones: int = 4,
        horizon_days: int = 14,
        latent_dim: int = 32,
        hidden_dim: int = 128,
    ):
        super().__init__()
        self.n_zones = n_zones
        self.horizon_days = horizon_days
        self.latent_dim = latent_dim
        self.hidden_dim = hidden_dim

        self.precip_encoder = nn.GRU(
            input_size=1,
            hidden_size=latent_dim,
            num_layers=1,
            batch_first=True,
        )

        self.meta_encoder = nn.Sequential(
            nn.Linear(2, latent_dim),
            nn.Tanh(),
        )

        zone_latent_dim = latent_dim * 2  # precip latent + meta latent

        full_latent_dim = n_zones * zone_latent_dim
        self.transition = nn.Sequential(
            nn.Linear(full_latent_dim, hidden_dim),
            nn.SiLU(),
            nn.Linear(hidden_dim, hidden_dim),
            nn.SiLU(),
            nn.Linear(hidden_dim, full_latent_dim),
        )

        self.precip_decoder = nn.Sequential(
            nn.Linear(zone_latent_dim, hidden_dim),
            nn.SiLU(),
            nn.Linear(hidden_dim, horizon_days),
            nn.Softplus(),  # precipitation ≥ 0
        )
        self.uncertainty_decoder = nn.Sequential(
            nn.Linear(zone_latent_dim, 32),
            nn.SiLU(),
            nn.Linear(32, 1),
            nn.Sigmoid(),  # uncertainty in [0, 1]
        )
        self.belief_decoder = nn.Sequential(
            nn.Linear(zone_latent_dim, 32),
            nn.SiLU(),
            nn.Linear(32, 1),
            nn.Sigmoid(),  # belief in [0, 1]
        )

        self.physics_loss = PhysicsResidualLoss(weight=0.01)

        logger.info(
            "TemporalDynamicsModel: n_zones=%d horizon=%d latent=%d hidden=%d",
            n_zones, horizon_days, latent_dim, hidden_dim,
        )

    def _encode(self, state: ZoneStateTensor) -> torch.Tensor:
        B, Z, H = state.precip.shape

        precip_seq = state.precip.reshape(B * Z, H, 1)
        _, h_n = self.precip_encoder(precip_seq)  # h_n: [1, B*Z, latent_dim]
        precip_latent = h_n.squeeze(0).reshape(B, Z, self.latent_dim)

        meta = torch.stack([state.uncertainty, state.belief], dim=-1)  # [B, Z, 2]
        meta_flat = meta.reshape(B * Z, 2)
        meta_latent = self.meta_encoder(meta_flat).reshape(B, Z, self.latent_dim)

        return torch.cat([precip_latent, meta_latent], dim=-1)  # [B, Z, 2*latent_dim]

    def forward(
        self,
        current: ZoneStateTensor,
        return_physics_loss: bool = True,
        dt: float = 1.0,
    ) -> Tuple[ZoneStateTensor, Optional[torch.Tensor]]:
        B, Z, H = current.precip.shape

        latent = self._encode(current)              # [B, Z, zone_latent_dim]
        latent_flat = latent.reshape(B, -1)         # [B, Z * zone_latent_dim]

        next_latent_flat = self.transition(latent_flat)
        next_latent = next_latent_flat.reshape(B, Z, -1)  # [B, Z, zone_latent_dim]

        next_latent_per_zone = next_latent.reshape(B * Z, -1)

        next_precip = self.precip_decoder(next_latent_per_zone).reshape(B, Z, H)
        next_uncertainty = self.uncertainty_decoder(next_latent_per_zone).reshape(B, Z)
        next_belief = self.belief_decoder(next_latent_per_zone).reshape(B, Z)

        next_state = ZoneStateTensor(
            precip=next_precip,
            uncertainty=next_uncertainty,
            belief=next_belief,
        )

        phys_loss = None
        if return_physics_loss:
            phys_loss = self.physics_loss(
                current.precip, next_precip, dt=float(dt),
            )

        return next_state, phys_loss

    def rollout(
        self,
        initial: ZoneStateTensor,
        steps: int = 5,
    ) -> List[ZoneStateTensor]:
        states = [initial]
        current = initial
        with torch.no_grad():
            for _ in range(steps):
                next_state, _ = self.forward(current, return_physics_loss=False)
                next_state = ZoneStateTensor(
                    precip=torch.clamp(next_state.precip, 0.0, 500.0),
                    uncertainty=torch.clamp(next_state.uncertainty, 0.0, 1.0),
                    belief=torch.clamp(next_state.belief, 0.0, 1.0),
                )
                states.append(next_state)
                current = next_state
        return states

    def save(self, path: str | Path) -> None:
        path = Path(path)
        path.parent.mkdir(parents=True, exist_ok=True)
        torch.save({
            "state_dict": self.state_dict(),
            "config": {
                "n_zones": self.n_zones,
                "horizon_days": self.horizon_days,
                "latent_dim": self.latent_dim,
                # FIX: store the actual hidden_dim integer, not a class name string
                "hidden_dim": self.hidden_dim,
            }
        }, path)
        logger.info("Saved dynamics model to %s", path)

    @classmethod
    def load(cls, path: str | Path, device: Optional[torch.device] = None) -> "TemporalDynamicsModel":
        path = Path(path)
        checkpoint = torch.load(path, map_location=device or "cpu")
        cfg = checkpoint["config"]
        model = cls(
            n_zones=cfg["n_zones"],
            horizon_days=cfg["horizon_days"],
            latent_dim=cfg.get("latent_dim", 32),
            hidden_dim=cfg.get("hidden_dim", 128),
        )
        model.load_state_dict(checkpoint["state_dict"])
        logger.info("Loaded dynamics model from %s", path)
        return model


# ---------------------------------------------------------------------------
# Offline trainer
# ---------------------------------------------------------------------------

class DynamicsTrainer:

    def __init__(
        self,
        n_zones: int = 4,
        horizon_days: int = 14,
        latent_dim: int = 32,
        hidden_dim: int = 128,
        physics_weight: float = 0.01,
        device: Optional[str] = None,
    ):
        self.device = torch.device(
            device or ("cuda" if torch.cuda.is_available() else "cpu")
        )
        self.model = TemporalDynamicsModel(
            n_zones=n_zones,
            horizon_days=horizon_days,
            latent_dim=latent_dim,
            hidden_dim=hidden_dim,
        ).to(self.device)
        self.physics_weight = physics_weight
        logger.info("DynamicsTrainer: device=%s  physics_weight=%.3f", self.device, physics_weight)

    def train(
        self,
        sequence_pairs: List[Tuple[ZoneStateTensor, ZoneStateTensor]],
        epochs: int = 50,
        batch_size: int = 64,
        lr: float = 1e-3,
        val_split: float = 0.1,
        dts: Optional[List[float]] = None,
        default_dt: float = 1.0,
    ) -> dict:
        if not sequence_pairs:
            raise ValueError("sequence_pairs is empty — provide ERA5 data")

        if dts is not None and len(dts) != len(sequence_pairs):
            raise ValueError(
                f"dts length {len(dts)} != sequence_pairs length "
                f"{len(sequence_pairs)}"
            )

        current_precips, current_uncerts, current_beliefs = [], [], []
        next_precips, next_uncerts, next_beliefs = [], [], []
        dt_list: List[float] = []

        for i, (curr, nxt) in enumerate(sequence_pairs):
            current_precips.append(curr.precip)
            current_uncerts.append(curr.uncertainty)
            current_beliefs.append(curr.belief)
            next_precips.append(nxt.precip)
            next_uncerts.append(nxt.uncertainty)
            next_beliefs.append(nxt.belief)
            dt_list.append(float(dts[i]) if dts is not None else float(default_dt))

        cp = torch.cat(current_precips, dim=0)
        cu = torch.cat(current_uncerts, dim=0)
        cb = torch.cat(current_beliefs, dim=0)
        np_ = torch.cat(next_precips, dim=0)
        nu = torch.cat(next_uncerts, dim=0)
        nb = torch.cat(next_beliefs, dim=0)
        dt_t = torch.tensor(dt_list, dtype=torch.float32)

        N = cp.shape[0]
        n_val = int(N * val_split) if val_split > 0 else 0
        if n_val >= N:
            n_val = max(0, N - 1)   # leave at least 1 sample for training
        n_train = N - n_val
        if n_train <= 0:
            raise ValueError(
                f"Dataset too small for the requested val_split: "
                f"N={N}, val_split={val_split} produces n_train={n_train}. "
                f"Reduce val_split or provide more pairs."
            )

        train_ds = TensorDataset(
            cp[:n_train], cu[:n_train], cb[:n_train],
            np_[:n_train], nu[:n_train], nb[:n_train],
            dt_t[:n_train],
        )
        val_ds = TensorDataset(
            cp[n_train:], cu[n_train:], cb[n_train:],
            np_[n_train:], nu[n_train:], nb[n_train:],
            dt_t[n_train:],
        )

        train_loader = DataLoader(train_ds, batch_size=batch_size, shuffle=True)
        val_loader   = DataLoader(val_ds,   batch_size=batch_size, shuffle=False)

        optimizer = torch.optim.AdamW(self.model.parameters(), lr=lr, weight_decay=1e-4)
        scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=epochs)

        history = {"train_loss": [], "val_loss": [], "physics_loss": [], "dt_mean": float(dt_t.mean())}

        for epoch in range(epochs):
            self.model.train()
            epoch_data_loss = 0.0
            epoch_phys_loss = 0.0

            for batch in train_loader:
                cp_b, cu_b, cb_b, np_b, nu_b, nb_b, dt_b = [
                    t.to(self.device) for t in batch
                ]

                current = ZoneStateTensor(precip=cp_b, uncertainty=cu_b, belief=cb_b)
                target  = ZoneStateTensor(precip=np_b, uncertainty=nu_b, belief=nb_b)

                batch_dt = float(dt_b.mean().item())
                pred, phys_loss = self.model(
                    current, return_physics_loss=True, dt=batch_dt,
                )

                data_loss = (
                    F.mse_loss(pred.precip / 500.0, target.precip / 500.0)
                    + F.mse_loss(pred.uncertainty, target.uncertainty)
                    + F.mse_loss(pred.belief,      target.belief)
                )

                total_loss = data_loss + self.physics_weight * phys_loss

                optimizer.zero_grad()
                total_loss.backward()
                torch.nn.utils.clip_grad_norm_(self.model.parameters(), 1.0)
                optimizer.step()

                epoch_data_loss += data_loss.item()
                epoch_phys_loss += phys_loss.item()

            scheduler.step()

            avg_data = epoch_data_loss / len(train_loader)
            avg_phys = epoch_phys_loss / len(train_loader)

            self.model.eval()
            val_loss = 0.0
            with torch.no_grad():
                for batch in val_loader:
                    cp_b, cu_b, cb_b, np_b, nu_b, nb_b, dt_b = [
                        t.to(self.device) for t in batch
                    ]
                    current = ZoneStateTensor(precip=cp_b, uncertainty=cu_b, belief=cb_b)
                    target  = ZoneStateTensor(precip=np_b, uncertainty=nu_b, belief=nb_b)
                    batch_dt = float(dt_b.mean().item()) if dt_b.numel() else 1.0
                    pred, _ = self.model(current, return_physics_loss=False)
                    val_loss += (
                        F.mse_loss(pred.precip / 500.0, target.precip / 500.0)
                        + F.mse_loss(pred.uncertainty, target.uncertainty)
                        + F.mse_loss(pred.belief,      target.belief)
                    ).item()

            avg_val = val_loss / max(len(val_loader), 1)

            history["train_loss"].append(avg_data)
            history["val_loss"].append(avg_val)
            history["physics_loss"].append(avg_phys)

            if epoch % 10 == 0 or epoch == epochs - 1:
                logger.info(
                    "Epoch %3d/%d  train=%.4f  val=%.4f  physics=%.4f  "
                    "v=%.3f  D=%.3f",
                    epoch + 1, epochs, avg_data, avg_val, avg_phys,
                    self.model.physics_loss.v.item(),
                    self.model.physics_loss.D.item(),
                )

        return history

    def save(self, path: str | Path) -> None:
        self.model.save(path)


# ---------------------------------------------------------------------------
# Ensemble for uncertainty quantification
# ---------------------------------------------------------------------------

class EnsembleDynamics:

    def __init__(self, n_models: int = 5, **model_kwargs):
        self.models = [TemporalDynamicsModel(**model_kwargs) for _ in range(n_models)]
        logger.info("EnsembleDynamics: %d models", n_models)

    def predict(
        self,
        current: ZoneStateTensor,
    ) -> Tuple[ZoneStateTensor, torch.Tensor]:
        all_precips, all_uncerts, all_beliefs = [], [], []

        for model in self.models:
            model.eval()
            with torch.no_grad():
                pred, _ = model(current, return_physics_loss=False)
            all_precips.append(pred.precip)
            all_uncerts.append(pred.uncertainty)
            all_beliefs.append(pred.belief)

        precip_stack = torch.stack(all_precips)    # [N, B, Z, H]
        uncert_stack = torch.stack(all_uncerts)    # [N, B, Z]
        belief_stack = torch.stack(all_beliefs)    # [N, B, Z]

        mean_state = ZoneStateTensor(
            precip=precip_stack.mean(0),
            uncertainty=uncert_stack.mean(0),
            belief=belief_stack.mean(0),
        )

        epistemic = (
            (precip_stack.std(0, correction=0) / 500.0).mean()
            + uncert_stack.std(0, correction=0).mean()
            + belief_stack.std(0, correction=0).mean()
        ) / 3.0

        return mean_state, epistemic

    def to(self, device: torch.device) -> "EnsembleDynamics":
        for m in self.models:
            m.to(device)
        return self


# ---------------------------------------------------------------------------
# Dyna rollout buffer
# ---------------------------------------------------------------------------

class DynaRolloutBuffer:

    def __init__(
        self,
        dynamics: TemporalDynamicsModel,
        n_synthetic_steps: int = 3,
        uncertainty_weight: float = 0.1,
    ):
        self.dynamics = dynamics
        self.n_synthetic_steps = n_synthetic_steps
        self.uncertainty_weight = uncertainty_weight

    def compute_surprise_bonus(
        self,
        obs_current: ZoneStateTensor,
        obs_actual_next: ZoneStateTensor,
    ) -> torch.Tensor:
        self.dynamics.eval()
        with torch.no_grad():
            pred_next, _ = self.dynamics(obs_current, return_physics_loss=False)

        precip_err = F.mse_loss(
            pred_next.precip / 500.0,
            obs_actual_next.precip / 500.0,
        )
        uncert_err = F.mse_loss(pred_next.uncertainty, obs_actual_next.uncertainty)
        belief_err = F.mse_loss(pred_next.belief, obs_actual_next.belief)

        surprise = (precip_err + uncert_err + belief_err) / 3.0
        return torch.clamp(surprise * self.uncertainty_weight, 0.0, 1.0)

    def generate_rollout(
        self,
        seed_state: ZoneStateTensor,
    ) -> List[ZoneStateTensor]:
        return self.dynamics.rollout(seed_state, steps=self.n_synthetic_steps)