LED/models/future_interaction_graph_v6.py

"""
V6-style future-trajectory interaction graph for LED.

Re-uses MoFlow's `FutureInteractionGraphV6` verbatim (imported via sys.path)
so that LED, MID, and MoFlow share the identical neighbor-scoring and
message-passing component:

  Per-node: q_i = W_q([y0_emb, sigma_mean]), k_j = W_k([y0_emb, sigma_mean])
  Per-edge: semantic = (q_i · k_j)/sqrt(D_s), geo_bias = geo_mlp(mean_rel, std_rel, min_dist, heading_diff_mean)
  Top-N by (semantic + geo_bias) → RelTrajEncoder([rel_pos, heading_diff], sigma_bias)
  GNN layers → gated residual → refined node embedding

This file is a thin adapter: it produces the V6 inputs from LED's native call
signature and decodes the refined node embeddings into a per-agent, per-mode
residual trajectory that LED adds to its frozen core denoiser's epsilon
output (or to the implied y_0 estimate, depending on --residual_on).

The decoder's final layer is zero-initialized so the overall residual starts
as a no-op; V6's own gated residual inside is additionally safe at init
because its gate_proj is randomly initialized but small.
"""

import os, sys
import torch
import torch.nn as nn

# Import MoFlow's V6 graph module (and dependencies) directly.
_MOFLOW_ROOT = os.path.abspath(
    os.path.join(os.path.dirname(__file__), '..', '..', 'MoFlow'))
if _MOFLOW_ROOT not in sys.path:
    sys.path.insert(0, _MOFLOW_ROOT)

from models.graph_interaction_nba_v6 import FutureInteractionGraphV6
from models.interaction_baselines import build_interaction_module  # E4: SRA_MODULE factory


class FutureInteractionGraphV6Wrapper(nn.Module):
    """Adapter that exposes MoFlow's V6 graph to LED's call signature.

    LED call:
        delta = graph(y0_hat, past, step_idx, sigma=sigma)
        y0_hat [B*A, K, T, 2], past [B*A, T_h, 6], step_idx int, sigma [B*A, 1] or None
        delta  [B*A, K, T, 2]

    V6 call (underneath):
        refined = future_graph(y_emb, y_abs, t_emb, tau, sigma_agent)
        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] or None
    """

    def __init__(
        self,
        num_agents: int = 11,
        future_steps: int = 20,
        past_steps: int = 10,
        past_channels: int = 6,
        node_dim: int = 128,
        top_n: int = 5,                  # matches MoFlow default
        num_denoise_steps: int = 5,
        num_gnn_layers: int = 2,
        rel_traj_hidden: int = 32,
        y0_score_dim: int = 32,
        num_heads: int = 4,
        dropout: float = 0.1,
        edge_mode: str = 'full',
        neighbor_mode: str = 'rag',
    ):
        super().__init__()
        self.A = num_agents
        self.T = future_steps
        self.D = node_dim

        # ---- Input encoders for the node embedding fed to V6 ----
        self.past_proj = nn.Sequential(
            nn.Linear(past_steps * past_channels, node_dim),
            nn.ReLU(inplace=True),
            nn.Linear(node_dim, node_dim),
        )
        self.y0_proj = nn.Sequential(
            nn.Linear(future_steps * 2, node_dim),
            nn.ReLU(inplace=True),
            nn.Linear(node_dim, node_dim),
        )
        self.step_emb = nn.Embedding(num_denoise_steps, node_dim)

        # ---- MoFlow V6 graph (scoring, edge encoding, GNN, gated residual) ----
        self.future_graph = build_interaction_module(   # E4: sra|gameformer|c2f via env SRA_MODULE
            embed_dim        = node_dim,
            future_steps     = future_steps,
            num_agents       = num_agents,
            num_heads        = num_heads,
            dropout          = dropout,
            num_gnn_layers   = num_gnn_layers,
            time_dim         = node_dim,
            top_n_neighbors  = min(top_n, num_agents - 1),
            rel_traj_hidden  = rel_traj_hidden,
            y0_score_dim     = y0_score_dim,
            edge_mode        = edge_mode,
            neighbor_mode    = neighbor_mode,
        )

        # ---- Decoder: refined node embedding -> per-timestep trajectory residual ----
        self.decode = nn.Sequential(
            nn.Linear(node_dim, node_dim),
            nn.ReLU(inplace=True),
            nn.Linear(node_dim, future_steps * 2),
        )
        nn.init.zeros_(self.decode[-1].weight)
        nn.init.zeros_(self.decode[-1].bias)

    def forward(
        self,
        y0_hat: torch.Tensor,   # [B*A, K, T, 2]
        past: torch.Tensor,     # [B*A, T_h, 6]
        step_idx: int,
        sigma: torch.Tensor = None,  # [B*A, 1] (per-agent scalar uncertainty, optional)
        A_override: int = None,
    ) -> torch.Tensor:          # [B*A, K, T, 2]
        BA, K, T, _ = y0_hat.shape
        # If A_override is given (variable-A scenes), use it; otherwise use init A.
        if A_override is not None:
            A = A_override
        else:
            A = self.A
        if BA % A != 0:
            # Fall back to treating whole batch as a single scene (B=1, A=BA)
            A = BA
        B = BA // A
        device = y0_hat.device
        # If A < 2, graph is meaningless — return zero residual
        if A < 2:
            return torch.zeros_like(y0_hat)
        # Rebuild graph's edge index for this A
        fg = self.future_graph
        if fg.num_agents != A:
            fg.num_agents = A
            fg._E0 = A * (A - 1)
            fg.top_n = max(1, min(fg.top_n, A - 1))
            src, dst = [], []
            for i in range(A):
                for j in range(A):
                    if i != j: src.append(j); dst.append(i)
            fg._single_edge_index = torch.tensor([src, dst], dtype=torch.long, device=device)

        # --- Build y_emb [B, K, A, D] from past + y0 + step ---
        past_feat = self.past_proj(past.reshape(BA, -1))             # [BA, D]
        past_feat = past_feat.unsqueeze(1).expand(BA, K, self.D)     # [BA, K, D]
        y0_feat   = self.y0_proj(y0_hat.reshape(BA, K, T * 2))       # [BA, K, D]
        step      = self.step_emb(torch.tensor(step_idx, device=device))  # [D]
        nodes_bka = past_feat + y0_feat + step                       # [BA, K, D]
        y_emb     = nodes_bka.view(B, A, K, self.D).permute(0, 2, 1, 3).contiguous()  # [B, K, A, D]

        # --- y_abs [B, K, A, T, 2] (unnormalized future positions) ---
        y_abs = y0_hat.view(B, A, K, T, 2).permute(0, 2, 1, 3, 4).contiguous()

        # --- t_emb [B, D] and tau [B] ---
        t_emb = self.step_emb(torch.tensor(step_idx, device=device)).unsqueeze(0).expand(B, -1)  # [B, D]
        tau   = torch.full((B,), float(step_idx), dtype=torch.float32, device=device)

        # --- sigma_agent [B, K, A, T] or None ---
        # LED's variance_estimation is per-agent scalar [BA, 1]; broadcast over K and T.
        sigma_agent = None
        if sigma is not None:
            sigma_ba = sigma.view(B, A)                              # [B, A]
            sigma_agent = sigma_ba.view(B, 1, A, 1).expand(B, K, A, T)

        # --- Run V6 graph: scoring, sparse edges, RelTrajEncoder, GNN, gated residual ---
        y_emb_refined = self.future_graph(y_emb, y_abs, t_emb, tau, sigma_agent=sigma_agent)
        # y_emb_refined: [B, K, A, D]

        # --- Decode to per-mode per-agent trajectory residual ---
        refined_bka = y_emb_refined.permute(0, 2, 1, 3).contiguous().view(BA, K, self.D)
        residual = self.decode(refined_bka).view(BA, K, T, 2)
        return residual