'''
models.py defines a configurable spectral world model family and its building blocks. It includes:

text encoders: DCTTextEncoder (fixed DCT basis) and MeanTextEncoder
image encoders/decoders: wavelet, convolutional, and FFT-based variants
transition modules: low-rank spectral, MLP, Koopman, and mini-transformer transitions
WorldModel, which fuses image/text latents, predicts the next latent state, and decodes next-frame images plus text logits
build_model(name, **kwargs) to instantiate named model variants and ablations by aliasing configuration names.
'''

# Andrew Kiruluta, UC Berkeley, June 2026

from __future__ import annotations

import math
import torch
import torch.nn as nn
import torch.nn.functional as F
from .haar import HaarWavelet2D
from .data import VOCAB, ACTION_NAMES


class DCTTextEncoder(nn.Module):
    """Text-to-spectral-field encoder using a fixed DCT basis over token position."""

    def __init__(self, vocab_size: int, max_len: int, emb_dim: int, out_dim: int):
        super().__init__()
        self.max_len = max_len
        self.emb = nn.Embedding(vocab_size, emb_dim, padding_idx=0)
        dct = torch.zeros(max_len, max_len)
        for k in range(max_len):
            for n in range(max_len):
                alpha = math.sqrt(1.0 / max_len) if k == 0 else math.sqrt(2.0 / max_len)
                dct[k, n] = alpha * math.cos(math.pi * (n + 0.5) * k / max_len)
        self.register_buffer("dct", dct)
        self.proj = nn.Sequential(
            nn.Linear(max_len * emb_dim, out_dim),
            nn.LayerNorm(out_dim),
            nn.GELU(),
        )

    def forward(self, tokens: torch.Tensor) -> torch.Tensor:
        e = self.emb(tokens)  # [B, L, E]
        modes = torch.einsum("kl,ble->bke", self.dct, e)
        return self.proj(modes.flatten(1))


class MeanTextEncoder(nn.Module):
    def __init__(self, vocab_size: int, emb_dim: int, out_dim: int):
        super().__init__()
        self.emb = nn.Embedding(vocab_size, emb_dim, padding_idx=0)
        self.proj = nn.Sequential(nn.Linear(emb_dim, out_dim), nn.LayerNorm(out_dim), nn.GELU())

    def forward(self, tokens: torch.Tensor) -> torch.Tensor:
        mask = tokens.ne(0).float().unsqueeze(-1)
        e = self.emb(tokens)
        pooled = (e * mask).sum(dim=1) / mask.sum(dim=1).clamp_min(1.0)
        return self.proj(pooled)


class WaveletImageEncoder(nn.Module):
    def __init__(self, image_size: int, levels: int, out_dim: int):
        super().__init__()
        self.wave = HaarWavelet2D(image_size=image_size, levels=levels)
        self.coeff_dim = image_size * image_size
        self.net = nn.Sequential(
            nn.Linear(self.coeff_dim, out_dim),
            nn.LayerNorm(out_dim),
            nn.GELU(),
        )

    def forward(self, image: torch.Tensor) -> torch.Tensor:
        coeffs = self.wave(image)
        return self.net(coeffs)


class ConvImageEncoder(nn.Module):
    def __init__(self, image_size: int, out_dim: int):
        super().__init__()
        self.net = nn.Sequential(
            nn.Conv2d(1, 16, 3, stride=2, padding=1), nn.GELU(),
            nn.Conv2d(16, 32, 3, stride=2, padding=1), nn.GELU(),
            nn.Flatten(),
            nn.Linear(32 * (image_size // 4) * (image_size // 4), out_dim),
            nn.LayerNorm(out_dim),
            nn.GELU(),
        )

    def forward(self, image: torch.Tensor) -> torch.Tensor:
        return self.net(image)


class FFTImageEncoder(nn.Module):
    """Small Fourier feature encoder used as a neural-operator-style baseline."""

    def __init__(self, image_size: int, out_dim: int, keep: int = 8):
        super().__init__()
        self.keep = keep
        # rfft width keep, height keep => real+imag features.
        in_dim = keep * keep * 2
        self.net = nn.Sequential(nn.Linear(in_dim, out_dim), nn.LayerNorm(out_dim), nn.GELU())

    def forward(self, image: torch.Tensor) -> torch.Tensor:
        fft = torch.fft.rfft2(image.squeeze(1), norm="ortho")
        x = fft[:, : self.keep, : self.keep]
        feats = torch.cat([x.real, x.imag], dim=1).flatten(1)
        return self.net(feats)


class WaveletImageDecoder(nn.Module):
    def __init__(self, image_size: int, levels: int, latent_dim: int):
        super().__init__()
        self.wave = HaarWavelet2D(image_size=image_size, levels=levels)
        self.net = nn.Sequential(
            nn.Linear(latent_dim, 256), nn.GELU(),
            nn.Linear(256, image_size * image_size),
        )

    def forward(self, z: torch.Tensor) -> torch.Tensor:
        coeffs = self.net(z)
        return torch.sigmoid(self.wave.inverse(coeffs))


class MLPImageDecoder(nn.Module):
    def __init__(self, image_size: int, latent_dim: int):
        super().__init__()
        self.image_size = image_size
        self.net = nn.Sequential(
            nn.Linear(latent_dim, 256), nn.GELU(),
            nn.Linear(256, image_size * image_size), nn.Sigmoid(),
        )

    def forward(self, z: torch.Tensor) -> torch.Tensor:
        return self.net(z).view(z.shape[0], 1, self.image_size, self.image_size)


class LowRankSpectralTransition(nn.Module):
    """Action-conditioned sparse/low-rank spectral operator dictionary.

    z_{t+1} = tanh(W0 z + sum_r beta_r(a) U_r V_r z + b(a) + nonlinear(z)).
    """

    def __init__(self, latent_dim: int, action_dim: int, n_ops: int = 4, rank: int = 16):
        super().__init__()
        self.latent_dim = latent_dim
        self.n_ops = n_ops
        self.rank = rank
        self.base = nn.Linear(latent_dim, latent_dim, bias=False)
        with torch.no_grad():
            self.base.weight.copy_(0.72 * torch.eye(latent_dim) + 0.01 * torch.randn(latent_dim, latent_dim))
        self.U = nn.Parameter(torch.randn(n_ops, latent_dim, rank) * 0.035)
        self.V = nn.Parameter(torch.randn(n_ops, rank, latent_dim) * 0.035)
        self.beta = nn.Embedding(action_dim, n_ops)
        self.bias = nn.Embedding(action_dim, latent_dim)
        self.nonlin = nn.Sequential(nn.Linear(latent_dim, latent_dim), nn.GELU(), nn.Linear(latent_dim, latent_dim))

    def forward(self, z: torch.Tensor, action: torch.Tensor) -> torch.Tensor:
        out = self.base(z)
        beta = torch.tanh(self.beta(action))
        # Low-rank operator dictionary.
        for r in range(self.n_ops):
            vrz = z @ self.V[r].t()
            urvz = vrz @ self.U[r].t()
            out = out + beta[:, r : r + 1] * urvz
        out = out + self.bias(action) + 0.10 * self.nonlin(z)
        return torch.tanh(out)

    def stability_penalty(self) -> torch.Tensor:
        # Smooth proxy for keeping operator energy bounded.
        base_energy = torch.linalg.matrix_norm(self.base.weight, ord="fro") / math.sqrt(self.latent_dim)
        op_energy = torch.stack([torch.linalg.matrix_norm(self.U[r] @ self.V[r], ord="fro") / math.sqrt(self.latent_dim) for r in range(self.n_ops)]).mean()
        return F.relu(base_energy + op_energy - 0.95) ** 2


class SpectralNormConstrainedTransition(LowRankSpectralTransition):
    """Low-rank action-conditioned transition with a hard spectral-norm cap.

    For each action, form A(a)=W0+sum_r beta_r(a) U_r V_r and rescale the
    effective matrix when its largest singular value exceeds ``rho``.  This
    variant intentionally removes the unconstrained nonlinear residual so the
    latent transition has a directly testable contraction bound before tanh.
    """
    def __init__(self, latent_dim: int, action_dim: int, n_ops: int = 4, rank: int = 16, rho: float = 0.95):
        super().__init__(latent_dim, action_dim, n_ops=n_ops, rank=rank)
        self.rho = float(rho)
        self.action_dim = action_dim

    def effective_matrix(self, action_ids: torch.Tensor) -> torch.Tensor:
        # [A,D,D] for unique action ids. Embedding coefficients remain trainable.
        beta = torch.tanh(self.beta(action_ids))
        uv = torch.einsum("odr,ore->ode", self.U, self.V)
        return self.base.weight.unsqueeze(0) + torch.einsum("ao,ode->ade", beta, uv)

    def constrained_matrix(self, action_ids: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
        A = self.effective_matrix(action_ids)
        sigma = torch.linalg.matrix_norm(A, ord=2, dim=(-2, -1))
        scale = torch.clamp(self.rho / sigma.clamp_min(1e-8), max=1.0)
        return A * scale[:, None, None], sigma

    def forward(self, z: torch.Tensor, action: torch.Tensor) -> torch.Tensor:
        unique, inverse = torch.unique(action, sorted=True, return_inverse=True)
        A, _ = self.constrained_matrix(unique)
        A_batch = A[inverse]
        out = torch.bmm(A_batch, z.unsqueeze(-1)).squeeze(-1) + self.bias(action)
        return torch.tanh(out)

    def stability_penalty(self) -> torch.Tensor:
        # Hard forward constraint already enforces the cap; this auxiliary loss
        # discourages repeatedly relying on rescaling.
        ids = torch.arange(self.action_dim, device=self.base.weight.device)
        raw = self.effective_matrix(ids)
        sigma = torch.linalg.matrix_norm(raw, ord=2, dim=(-2, -1))
        return torch.mean(F.relu(sigma - self.rho) ** 2)

    @torch.no_grad()
    def stability_diagnostics(self) -> dict[str, float]:
        ids = torch.arange(self.action_dim, device=self.base.weight.device)
        A, raw_sigma = self.constrained_matrix(ids)
        constrained_sigma = torch.linalg.matrix_norm(A, ord=2, dim=(-2, -1))
        return {
            "raw_operator_norm_max": float(raw_sigma.max().cpu()),
            "raw_operator_norm_mean": float(raw_sigma.mean().cpu()),
            "effective_operator_norm_max": float(constrained_sigma.max().cpu()),
            "effective_operator_norm_mean": float(constrained_sigma.mean().cpu()),
            "rho": self.rho,
        }


class SelectiveSpectralTransition(LowRankSpectralTransition):
    """Selective mode regularization for multimodal transport.

    Unlike the v3 hard contraction variant, this keeps singular modes at or below
    ``max_gain`` untouched and clips only expansive singular values. Near-unit
    modes are therefore preserved for persistent/action-conditioned transport.
    The nonlinear residual is retained.
    """
    def __init__(self, latent_dim: int, action_dim: int, n_ops: int = 4, rank: int = 16, max_gain: float = 1.05):
        super().__init__(latent_dim, action_dim, n_ops=n_ops, rank=rank)
        self.max_gain = float(max_gain)
        self.action_dim = action_dim

    def effective_matrix(self, action_ids: torch.Tensor) -> torch.Tensor:
        beta = torch.tanh(self.beta(action_ids))
        uv = torch.einsum("odr,ore->ode", self.U, self.V)
        return self.base.weight.unsqueeze(0) + torch.einsum("ao,ode->ade", beta, uv)

    def selective_matrix(self, action_ids: torch.Tensor):
        A = self.effective_matrix(action_ids)
        U, S, Vh = torch.linalg.svd(A, full_matrices=False)
        Sc = torch.clamp(S, max=self.max_gain)
        Ac = (U * Sc.unsqueeze(-2)) @ Vh
        return Ac, S, Sc

    def forward(self, z: torch.Tensor, action: torch.Tensor) -> torch.Tensor:
        unique, inverse = torch.unique(action, sorted=True, return_inverse=True)
        A, _, _ = self.selective_matrix(unique)
        out = torch.bmm(A[inverse], z.unsqueeze(-1)).squeeze(-1)
        out = out + self.bias(action) + 0.10 * self.nonlin(z)
        return torch.tanh(out)

    def stability_penalty(self) -> torch.Tensor:
        ids = torch.arange(self.action_dim, device=self.base.weight.device)
        A = self.effective_matrix(ids)
        sigma = torch.linalg.svdvals(A)
        # Penalize only expansive modes; do not push all modes toward contraction.
        return torch.mean(F.relu(sigma - self.max_gain) ** 2)

    @torch.no_grad()
    def stability_diagnostics(self) -> dict[str, float]:
        ids = torch.arange(self.action_dim, device=self.base.weight.device)
        _, raw, clipped = self.selective_matrix(ids)
        return {
            "raw_operator_norm_max": float(raw[:, 0].max().cpu()),
            "raw_operator_norm_mean": float(raw[:, 0].mean().cpu()),
            "effective_operator_norm_max": float(clipped[:, 0].max().cpu()),
            "effective_operator_norm_mean": float(clipped[:, 0].mean().cpu()),
            "rho": self.max_gain,
            "fraction_modes_clipped": float((raw > self.max_gain).float().mean().cpu()),
        }


class StructuredSpectralTransportTransition(nn.Module):
    """V6 structured action-conditioned transport.

    Separates persistent/energy-preserving transport, dissipative spectral modes,
    and action-driven low-rank transport instead of asking one generic operator
    to learn all three roles. The persistent block is parameterized by a
    skew-symmetric generator and a Cayley transform, giving an orthogonal
    transport matrix. Dissipation is diagonal in the learned latent basis and
    action effects are low rank. Only expansive singular modes of the composed
    operator are softly clipped at ``max_gain``.
    """
    def __init__(self, latent_dim:int, action_dim:int, rank:int=12, max_gain:float=1.05):
        super().__init__(); self.latent_dim=latent_dim; self.action_dim=action_dim; self.rank=rank; self.max_gain=float(max_gain)
        self.generator=nn.Parameter(0.01*torch.randn(latent_dim,latent_dim))
        self.logit_keep=nn.Parameter(torch.full((latent_dim,), 2.2))
        self.action_U=nn.Parameter(0.025*torch.randn(action_dim,latent_dim,rank))
        self.action_V=nn.Parameter(0.025*torch.randn(action_dim,rank,latent_dim))
        self.action_gate=nn.Embedding(action_dim,1); nn.init.zeros_(self.action_gate.weight)
        self.bias=nn.Embedding(action_dim,latent_dim)
        self.residual=nn.Sequential(nn.Linear(latent_dim,latent_dim),nn.GELU(),nn.Linear(latent_dim,latent_dim))

    def transport_matrix(self):
        G=0.25*torch.tanh(self.generator); K=0.5*(G-G.t()); I=torch.eye(self.latent_dim,device=K.device,dtype=K.dtype)
        return torch.linalg.solve(I-0.5*K, I+0.5*K)

    def effective_matrix(self, action_ids):
        Q=self.transport_matrix(); keep=0.80+0.20*torch.sigmoid(self.logit_keep)
        base=Q*keep.unsqueeze(0)
        uv=torch.einsum('adr,are->ade',self.action_U[action_ids],self.action_V[action_ids])
        gate=torch.tanh(self.action_gate(action_ids)).view(-1,1,1)
        return base.unsqueeze(0)+gate*uv

    def selective_matrix(self, action_ids):
        A=self.effective_matrix(action_ids); U,S,Vh=torch.linalg.svd(A.detach(),full_matrices=False); Sc=torch.clamp(S,max=self.max_gain)
        Ac=(U*Sc.unsqueeze(-2))@Vh
        # Straight-through projection: clipped matrix in the forward pass while
        # avoiding unstable SVD gradients during optimization.
        return A + (Ac-A.detach()),S,Sc

    def forward(self,z,action):
        unique,inverse=torch.unique(action,sorted=True,return_inverse=True); A,_,_=self.selective_matrix(unique)
        out=torch.bmm(A[inverse],z.unsqueeze(-1)).squeeze(-1)+self.bias(action)+0.05*self.residual(z)
        return torch.tanh(out)

    def stability_penalty(self):
        ids=torch.arange(self.action_dim,device=self.generator.device); A=self.effective_matrix(ids); s=torch.linalg.svdvals(A)
        expansive=F.relu(s-self.max_gain).pow(2).mean(); q=self.transport_matrix(); I=torch.eye(self.latent_dim,device=q.device)
        orth=(q.t()@q-I).pow(2).mean(); return expansive+0.01*orth

    @torch.no_grad()
    def stability_diagnostics(self):
        ids=torch.arange(self.action_dim,device=self.generator.device); _,raw,clip=self.selective_matrix(ids); q=self.transport_matrix(); I=torch.eye(self.latent_dim,device=q.device)
        return {'raw_operator_norm_max':float(raw[:,0].max().cpu()),'raw_operator_norm_mean':float(raw[:,0].mean().cpu()),
                'effective_operator_norm_max':float(clip[:,0].max().cpu()),'effective_operator_norm_mean':float(clip[:,0].mean().cpu()),
                'fraction_modes_clipped':float((raw>self.max_gain).float().mean().cpu()),'rho':self.max_gain,
                'transport_orthogonality_error':float(torch.linalg.matrix_norm(q.t()@q-I,ord='fro').cpu()/self.latent_dim**0.5),
                'mean_dissipative_keep':float((0.80+0.20*torch.sigmoid(self.logit_keep)).mean().cpu())}


class MLPTransition(nn.Module):
    def __init__(self, latent_dim: int, action_dim: int):
        super().__init__()
        self.emb = nn.Embedding(action_dim, 24)
        self.net = nn.Sequential(
            nn.Linear(latent_dim + 24, 192), nn.GELU(),
            nn.Linear(192, 192), nn.GELU(),
            nn.Linear(192, latent_dim), nn.Tanh(),
        )

    def forward(self, z: torch.Tensor, action: torch.Tensor) -> torch.Tensor:
        return self.net(torch.cat([z, self.emb(action)], dim=-1))

    def stability_penalty(self) -> torch.Tensor:
        return torch.tensor(0.0, device=next(self.parameters()).device)


class KoopmanTransition(nn.Module):
    def __init__(self, latent_dim: int, action_dim: int):
        super().__init__()
        self.A = nn.Parameter(0.75 * torch.eye(latent_dim).unsqueeze(0).repeat(action_dim, 1, 1) + 0.005 * torch.randn(action_dim, latent_dim, latent_dim))
        self.bias = nn.Embedding(action_dim, latent_dim)

    def forward(self, z: torch.Tensor, action: torch.Tensor) -> torch.Tensor:
        A = self.A[action]
        return torch.tanh(torch.bmm(A, z.unsqueeze(-1)).squeeze(-1) + self.bias(action))

    def stability_penalty(self) -> torch.Tensor:
        fro = torch.linalg.matrix_norm(self.A, ord="fro", dim=(1, 2)).mean() / math.sqrt(self.A.shape[-1])
        return F.relu(fro - 0.95) ** 2


class MiniTransformerTransition(nn.Module):
    def __init__(self, latent_dim: int, action_dim: int):
        super().__init__()
        self.action = nn.Embedding(action_dim, latent_dim)
        layer = nn.TransformerEncoderLayer(d_model=latent_dim, nhead=4, dim_feedforward=192, batch_first=True, dropout=0.0)
        self.encoder = nn.TransformerEncoder(layer, num_layers=1)
        self.out = nn.Sequential(nn.Linear(latent_dim, latent_dim), nn.Tanh())

    def forward(self, z: torch.Tensor, action: torch.Tensor) -> torch.Tensor:
        seq = torch.stack([z, self.action(action)], dim=1)
        h = self.encoder(seq)
        return self.out(h[:, 0])

    def stability_penalty(self) -> torch.Tensor:
        return torch.tensor(0.0, device=next(self.parameters()).device)


class WorldModel(nn.Module):
    def __init__(
        self,
        model_type: str = "swm",
        image_size: int = 32,
        max_text_len: int = 5,
        vocab_size: int = len(VOCAB),
        action_dim: int = len(ACTION_NAMES),
        latent_dim: int = 96,
        image_latent_dim: int = 48,
        text_latent_dim: int = 32,
        no_text: bool = False,
    ):
        super().__init__()
        self.model_type = model_type
        self.no_text = no_text
        self.latent_dim = latent_dim
        self.image_size = image_size
        if model_type in {"swm", "swm_spectral_norm", "swm_selective", "swm_structured", "swm_no_transition", "swm_no_text", "swm_no_stability"}:
            self.image_encoder = WaveletImageEncoder(image_size, levels=2, out_dim=image_latent_dim)
            self.text_encoder = DCTTextEncoder(vocab_size, max_text_len, emb_dim=24, out_dim=text_latent_dim)
            self.image_decoder = WaveletImageDecoder(image_size, levels=2, latent_dim=latent_dim)
        elif model_type == "neural_operator":
            self.image_encoder = FFTImageEncoder(image_size, out_dim=image_latent_dim)
            self.text_encoder = DCTTextEncoder(vocab_size, max_text_len, emb_dim=24, out_dim=text_latent_dim)
            self.image_decoder = MLPImageDecoder(image_size, latent_dim=latent_dim)
        else:
            self.image_encoder = ConvImageEncoder(image_size, out_dim=image_latent_dim)
            self.text_encoder = MeanTextEncoder(vocab_size, emb_dim=24, out_dim=text_latent_dim)
            self.image_decoder = MLPImageDecoder(image_size, latent_dim=latent_dim)

        if no_text or model_type == "swm_no_text":
            text_latent_dim = 0
            self.no_text = True
        self.fusion = nn.Sequential(
            nn.Linear(image_latent_dim + text_latent_dim, latent_dim),
            nn.LayerNorm(latent_dim),
            nn.GELU(),
            nn.Linear(latent_dim, latent_dim),
            nn.LayerNorm(latent_dim),
        )
        if model_type == "swm_spectral_norm":
            self.transition = SpectralNormConstrainedTransition(latent_dim, action_dim, n_ops=4, rank=12, rho=0.95)
        elif model_type == "swm_selective":
            self.transition = SelectiveSpectralTransition(latent_dim, action_dim, n_ops=4, rank=12, max_gain=1.05)
        elif model_type == "swm_structured":
            self.transition = StructuredSpectralTransportTransition(latent_dim, action_dim, rank=12, max_gain=1.05)
        elif model_type in {"swm", "swm_no_text", "swm_no_stability"}:
            self.transition = LowRankSpectralTransition(latent_dim, action_dim, n_ops=4, rank=12)
        elif model_type == "koopman":
            self.transition = KoopmanTransition(latent_dim, action_dim)
        elif model_type == "transformer":
            self.transition = MiniTransformerTransition(latent_dim, action_dim)
        else:
            self.transition = MLPTransition(latent_dim, action_dim)
        self.text_decoder = nn.Linear(latent_dim, max_text_len * vocab_size)
        self.max_text_len = max_text_len
        self.vocab_size = vocab_size

    def encode(self, image: torch.Tensor, tokens: torch.Tensor) -> torch.Tensor:
        zi = self.image_encoder(image)
        if self.no_text:
            fused = zi
        else:
            zt = self.text_encoder(tokens)
            fused = torch.cat([zi, zt], dim=-1)
        return self.fusion(fused)

    def decode_text(self, z: torch.Tensor) -> torch.Tensor:
        return self.text_decoder(z).view(z.shape[0], self.max_text_len, self.vocab_size)

    def forward(self, image: torch.Tensor, tokens: torch.Tensor, action: torch.Tensor) -> dict[str, torch.Tensor]:
        z = self.encode(image, tokens)
        z_next = self.transition(z, action)
        return {
            "z": z,
            "z_pred": z_next,
            "image_pred": self.image_decoder(z_next),
            "text_logits": self.decode_text(z_next),
        }

    def stability_penalty(self) -> torch.Tensor:
        return self.transition.stability_penalty()


def build_model(name: str, **kwargs) -> WorldModel:
    aliases = {
        "swm": "swm",
        "swm_spectral_norm": "swm_spectral_norm",
        "swm_selective": "swm_selective",
        "swm_structured": "swm_structured",
        "swm_structured_cf": "swm_structured",
        "swm_structured_cf_v9": "swm_structured",
        "swm_structured_cf_v11": "swm_structured",
        "dreamer": "dreamer",
        "transformer": "transformer",
        "koopman": "koopman",
        "neural_operator": "neural_operator",
        "no_spectral_state": "dreamer",
        "no_spectral_transition": "swm_no_transition",
        "no_text_branch": "swm_no_text",
        "no_stability_penalty": "swm_no_stability",
    }
    if name not in aliases:
        raise ValueError(f"Unknown model '{name}'. Options: {sorted(aliases)}")
    return WorldModel(model_type=aliases[name], **kwargs)