LED/models/interaction_baselines.py

"""
E4 plug-in baselines — alternative interaction/refinement modules that honor the EXACT
FutureInteractionGraphV6 forward contract, for a protocol-matched comparison vs SRA:

    forward(y_emb[B,K,A,D], y_abs[B,K,A,T,2], t_emb[B,D], tau[B],
            sigma_agent[B,K,A,T]|None, agent_mask[B,A]|None=None) -> [B,K,A,D]
    (gated residual, embedding->embedding, NO sigma out; robust to K in {1,10,20}, any A)

Modules:
  * GameFormerInteraction  [ref 20]: FAITHFUL GameFormer level-k InteractionDecoder —
    host denoiser = InitialDecoder (level 0); denoiser K samples = GameFormer K modes;
    FutureEncoder + score-weighted mode aggregation + interaction self-attention +
    own-future-masked cross-attention + L-level GMM refinement. No neighbor selection,
    no uncertainty (that is SRA). See the detailed faithfulness note above the class.
  * CoarseToFineRefine     [ref 22]: Coarse-to-Fine (Jia et al. 2022) TEMPORAL refinement —
    a per-agent temporal module (GRU / 1D-CNN over the future horizon). No interaction.
    Orthogonal axis (temporal).

Both accept the SAME __init__ kwargs as V6 (extras absorbed by **kwargs) and return the SAME
gated-residual embedding, so they are true drop-ins at each host's V6 instantiation line.
Selection is via env var SRA_MODULE in {sra, gameformer, c2f} (see build_interaction_module).
"""
import os
import torch
import torch.nn as nn


class _GatedResidualHead(nn.Module):
    """V6's output convention: out = input + gate(cat[input,refined]) * out_proj(refined).
    Random-initialised (active at init), exactly like V6 — hosts that need no-op-at-init
    (MID/LED) supply their own external zero-init projection; MoFlow uses the output directly."""
    def __init__(self, embed_dim):
        super().__init__()
        self.gate_proj = nn.Sequential(nn.Linear(2 * embed_dim, embed_dim), nn.Sigmoid())
        self.out_proj = nn.Linear(embed_dim, embed_dim)
        nn.init.zeros_(self.out_proj.weight)               # NO-OP AT INIT: module starts as identity
        nn.init.zeros_(self.out_proj.bias)                 # (stabilizes MoFlow, which uses output directly)

    def forward(self, orig, refined):                      # both [N, D]
        gate = self.gate_proj(torch.cat([orig, refined], dim=-1))
        res = self.out_proj(refined)
        _cap = float(os.environ.get('MOFLOW_DAMP', 0) or 0)  # cap residual norm -> MoFlow flow-field stability
        if _cap > 0:
            rn = res.norm(dim=-1, keepdim=True)
            res = res * (rn.clamp(max=_cap) / (rn + 1e-6))
        return orig + gate * res


# ============================================================================
#  GameFormer [ref 20] -- FAITHFUL level-k InteractionDecoder as a denoiser plug-in
# ============================================================================
# Faithful to Liu et al. ICCV'23 (github.com/MCZhi/GameFormer, model/modules.py):
#   * The host DENOISER plays the InitialDecoder (level 0): its current estimate
#     (y_emb, y_abs) IS the level-0 prediction, and the denoiser's K SAMPLES are
#     treated as GameFormer's K MODES (the modal set the level-k step reasons over).
#   * Each InteractionDecoder level implements the real mechanism:
#       - FutureEncoder: MLP on per-mode [x,y,heading,vx,vy], max-pooled over time
#         (modules.FutureEncoder; box-size channels dropped -- no size in NBA/sport).
#       - score-softmax-weighted aggregation of the K modes -> one future token/agent.
#       - interaction SelfTransformer over agents (game-theoretic "condition on all
#         other agents' level-(k-1) futures").
#       - cross-attention of each mode's content (prev content + own future) to
#         [interaction ; scene-context], with the agent's OWN future token MASKED
#         (GameFormer's own-future masking).
#       - a GMM-mu trajectory decoder emits the level-k refinement; iterate L levels.
#   * The final level's content -> gated residual to the host embedding.
# ONE gap vs the paper: per-level imitation SUPERVISION (GameFormer sums a GMM loss
# over every level). A plug-in is trained only through the host denoising loss, so
# the levels are learned end-to-end rather than level-wise-supervised.


class _AgentSelfAttn(nn.Module):
    """SelfTransformer over agents (modules.SelfTransformer). x: [B, A, D]."""
    def __init__(self, dim, heads, dropout):
        super().__init__()
        self.heads = heads
        self.attn = nn.MultiheadAttention(dim, heads, dropout=dropout, batch_first=True)
        self.n1 = nn.LayerNorm(dim); self.n2 = nn.LayerNorm(dim)
        self.ffn = nn.Sequential(nn.Linear(dim, dim * 4), nn.GELU(),
                                 nn.Dropout(dropout), nn.Linear(dim * 4, dim))

    def forward(self, x, key_padding_mask=None):
        a, _ = self.attn(x, x, x, key_padding_mask=key_padding_mask, need_weights=False)
        x = self.n1(a + x)
        return self.n2(self.ffn(x) + x)


class _CrossAttn(nn.Module):
    """CrossTransformer (modules.CrossTransformer). q:[N,Lq,D] k/v:[N,Lk,D]."""
    def __init__(self, dim, heads, dropout):
        super().__init__()
        self.heads = heads
        self.attn = nn.MultiheadAttention(dim, heads, dropout=dropout, batch_first=True)
        self.n1 = nn.LayerNorm(dim); self.n2 = nn.LayerNorm(dim)
        self.ffn = nn.Sequential(nn.Linear(dim, dim * 4), nn.GELU(),
                                 nn.Dropout(dropout), nn.Linear(dim * 4, dim))

    def forward(self, q, k, v, attn_mask=None):
        a, _ = self.attn(q, k, v, attn_mask=attn_mask, need_weights=False)
        a = self.n1(a)
        return self.n2(self.ffn(a) + a)


class _FutureEncoder(nn.Module):
    """modules.FutureEncoder: per-mode future -> vector via max-pool over an MLP on
    [x, y, heading, vx, vy].  trajs [B,K,A,T,2] (K modes), cur_xy [B,K,A,2]."""
    def __init__(self, dim, dt=0.2):
        super().__init__()
        self.dt = dt
        self.mlp = nn.Sequential(nn.Linear(5, 64), nn.ReLU(inplace=True), nn.Linear(64, dim))

    def forward(self, trajs, cur_xy):
        pos = trajs - cur_xy.unsqueeze(-2)                             # centered position
        xy = torch.cat([cur_xy.unsqueeze(-2), trajs], dim=-2)          # [B,K,A,T+1,2]
        dxy = torch.diff(xy, dim=-2)
        v = dxy / self.dt
        theta = torch.atan2(dxy[..., 1], dxy[..., 0].clamp(min=1e-3)).unsqueeze(-1)
        state = torch.cat([pos, theta, v], dim=-1)                     # [B,K,A,T,5]
        h = self.mlp(state.detach())                                   # GameFormer DETACHES future feats
        return h.max(dim=-2).values                                    # [B,K,A,D]


class GameFormerInteraction(nn.Module):
    """FAITHFUL GameFormer level-k InteractionDecoder as a denoiser plug-in (see header)."""

    def __init__(self, embed_dim, future_steps, num_agents,
                 num_heads=4, dropout=0.1, num_gnn_layers=2, time_dim=128,
                 top_n_neighbors=5, rel_traj_hidden=32, y0_score_dim=32,
                 edge_mode='full', neighbor_mode='rag', num_levels=None, **kwargs):
        super().__init__()
        self.embed_dim = embed_dim
        self.num_agents = num_agents            # LED wrapper reads this
        L = int(os.environ.get('GF_LEVELS', num_levels if num_levels is not None else 3))
        self.num_levels = L
        self.t_proj = nn.Linear(embed_dim, embed_dim)
        self.future_encoder = _FutureEncoder(embed_dim)                # SHARED across levels
        self.score_head = nn.ModuleList([                              # per-mode score -> softmax weight
            nn.Sequential(nn.Linear(embed_dim, 64), nn.ELU(), nn.Linear(64, 1)) for _ in range(L)])
        self.interaction_enc = nn.ModuleList([_AgentSelfAttn(embed_dim, num_heads, dropout) for _ in range(L)])
        self.query_enc = nn.ModuleList([_CrossAttn(embed_dim, num_heads, dropout) for _ in range(L)])
        self.traj_decode = nn.ModuleList([nn.Linear(embed_dim, future_steps * 2) for _ in range(L)])
        for dec in self.traj_decode:                                   # each level starts as identity (stable)
            nn.init.zeros_(dec.weight); nn.init.zeros_(dec.bias)
        self.head = _GatedResidualHead(embed_dim)

    def forward(self, y_emb, y_abs, t_emb, tau, sigma_agent=None, agent_mask=None):
        B, K, A, D = y_emb.shape
        T = y_abs.shape[3]
        if A <= 1:
            return y_emb
        dev = y_emb.device
        cur_xy = y_abs[..., 0, :]                            # [B,K,A,2] level-0 reference (1st future step)
        y_traj = y_abs                                       # [B,K,A,T,2] current (level-0) future
        t = self.t_proj(t_emb).view(B, 1, 1, D)              # denoiser timestep context
        agent_ctx = y_emb.mean(dim=1) + t.squeeze(1)         # [B,A,D] fixed scene context (over modes)
        kp = (~agent_mask.bool()) if agent_mask is not None else None   # [B,A] True=pad
        ar = torch.arange(A, device=dev)
        heads = self.query_enc[0].heads

        content = y_emb                                      # level-0 content = host embedding
        for l in range(self.num_levels):
            multi_fut = self.future_encoder(y_traj, cur_xy)             # [B,K,A,D] per-mode future
            w = self.score_head[l](multi_fut).softmax(dim=1)           # [B,K,A,1] over K modes
            agg_fut = (multi_fut * w).sum(dim=1)                        # [B,A,D] aggregated future
            interaction = self.interaction_enc[l](agg_fut, kp)         # [B,A,D] game-theoretic
            ctx = torch.cat([interaction, agent_ctx], dim=1)           # [B,2A,D]

            q = (content + multi_fut + t).reshape(B, K * A, D)         # prev content + own future + t
            am = torch.zeros(B, K, A, 2 * A, dtype=torch.bool, device=dev)
            am[:, :, ar, ar] = True                                     # mask OWN future (interaction block)
            if agent_mask is not None:
                pad = (~agent_mask.bool())[:, None, None, :]           # [B,1,1,A]
                am[..., :A] = am[..., :A] | pad
                am[..., A:] = am[..., A:] | pad
            am = am.reshape(B, K * A, 2 * A)[:, None].expand(
                B, heads, K * A, 2 * A).reshape(B * heads, K * A, 2 * A)
            content = self.query_enc[l](q, ctx, ctx, am).reshape(B, K, A, D)
            y_traj = y_traj + self.traj_decode[l](content).view(B, K, A, T, 2)   # level-k refinement
        refined = content.reshape(B * K * A, D)
        orig = y_emb.reshape(B * K * A, D)
        return self.head(orig, refined).view(B, K, A, D)


class CoarseToFineRefine(nn.Module):
    """Coarse-to-Fine temporal refiner: per-agent GRU/1D-CNN over the future horizon (no interaction)."""

    def __init__(self, embed_dim, future_steps, num_agents,
                 num_heads=4, dropout=0.1, num_gnn_layers=2, time_dim=128,
                 top_n_neighbors=5, rel_traj_hidden=32, y0_score_dim=32,
                 edge_mode='full', neighbor_mode='rag', refine_type='gru', hidden=200, **kwargs):
        super().__init__()
        self.embed_dim = embed_dim
        self.num_agents = num_agents            # LED wrapper reads this (skips variable-A rebuild when ==A)
        self.refine_type = os.environ.get('C2F_REFINE', refine_type)
        self.pos_emb = nn.Linear(2, hidden)
        if self.refine_type == 'cnn':                                   # 1D-CNN temporal refiner (per-timestep output)
            self.temporal = nn.Sequential(
                nn.Conv1d(hidden, hidden, 3, padding=1), nn.ReLU(inplace=True),
                nn.Conv1d(hidden, hidden, 3, padding=1), nn.ReLU(inplace=True),
                nn.Conv1d(hidden, hidden, 3, padding=1), nn.ReLU(inplace=True))
        else:                                                           # AUTOREGRESSIVE (unidirectional) GRU
            self.temporal = nn.GRU(hidden, hidden, num_layers=2, batch_first=True, dropout=dropout)
        self.delta_head = nn.Linear(hidden, 2)                          # per-timestep coarse->fine correction Δ_t
        self.traj_encode = nn.Sequential(                              # re-encode the FINE trajectory -> host feature
            nn.Linear(future_steps * 2, embed_dim), nn.ReLU(inplace=True),
            nn.Linear(embed_dim, embed_dim))
        self.t_proj = nn.Linear(embed_dim, embed_dim)
        self.head = _GatedResidualHead(embed_dim)

    def forward(self, y_emb, y_abs, t_emb, tau, sigma_agent=None, agent_mask=None):
        # Faithful coarse-to-fine: host prediction = COARSE trajectory; walk it temporally
        # (autoregressive GRU / 1D-CNN) emitting a per-timestep correction Δ_t -> FINE trajectory,
        # then re-encode the fine trajectory into the host residual. Per-agent (no interaction).
        B, K, A, D = y_emb.shape
        T = y_abs.shape[3]
        N = B * K * A
        y_coarse = (y_abs - y_abs[..., :1, :]).reshape(N, T, 2)        # coarse trajectory (centered)
        seq = self.pos_emb(y_coarse)                                    # [N, T, H]
        if self.refine_type == 'cnn':
            h = self.temporal(seq.transpose(1, 2)).transpose(1, 2)      # [N, T, H]
        else:
            h, _ = self.temporal(seq)                                   # [N, T, H] per-timestep (autoregressive)
        delta = self.delta_head(h)                                      # [N, T, 2] coarse->fine correction
        y_fine = y_coarse + delta                                       # refined (fine) trajectory
        refined = self.traj_encode(y_fine.reshape(N, T * 2))            # re-encode fine traj -> feature
        t = self.t_proj(t_emb).view(B, 1, 1, D).expand(B, K, A, D).reshape(N, D)
        refined = refined + t
        orig = y_emb.reshape(N, D)
        return self.head(orig, refined).view(B, K, A, D)


def build_interaction_module(name=None, **kwargs):
    """Factory used at each host's V6 instantiation line. `name` defaults to env SRA_MODULE
    (then 'sra'). SRA branch imports the real V6 and forwards only kwargs it accepts."""
    name = (name or os.environ.get('SRA_MODULE') or 'sra').lower()
    if name in ('sra', 'v6', 'graph'):
        import inspect
        from models.graph_interaction_nba_v6 import FutureInteractionGraphV6
        allowed = set(inspect.signature(FutureInteractionGraphV6.__init__).parameters)
        v6kw = {k: v for k, v in kwargs.items() if k in allowed}
        return FutureInteractionGraphV6(**v6kw)
    if name in ('gameformer', 'gf', 'gameformer_int'):
        return GameFormerInteraction(**kwargs)
    if name in ('c2f', 'coarse2fine', 'coarsetofine'):
        return CoarseToFineRefine(**kwargs)
    raise ValueError(f"unknown SRA_MODULE / interaction_module: {name!r}")