MoFlow/models/graph_interaction_nba_v5.py

"""
FutureInteractionGraphV5 — uncertainty-aware semantic scorer + sparse RelTrajEncoder.

Builds on V4 (learnable top-N sparse graph) with an enhanced edge scorer:

  V4 scorer:  [mean_rel, std_rel, min_dist]          → [E, 5]   (geometric only)
  V5 scorer:  [y0_emb_i, y0_emb_j_scaled,            → [E, 2*D_score + 7]
               σ_i, σ_j,
               mean_rel, std_rel, min_dist]

Where:
  y0_emb_{i,j}  — trajectory embedding from clean y_0_hat (not noisy y_t).
                   Projected from y_abs [B,K,A,T,2] which is already y_0_hat
                   unnormalised — no noise, same source as edge geometry.
  σ_i, σ_j     — explicit per-agent mean uncertainty, allowing the scorer to
                   learn "uncertain target prefers certain source."

When sigma_agent is None (pass 1 of the two-pass forward), certainty scaling
is skipped and σ features are zeroed — scorer still runs on geometric +
semantic features, gracefully degrading to a noisier but functional signal.
"""

import torch
import torch.nn as nn
from models.graph_interaction_nba_v4 import FutureInteractionGraphV4
from models.graph_interaction_nba_v3 import RelTrajEncoder


def _heading_diff(pos_i: torch.Tensor,
                  pos_j: torch.Tensor) -> torch.Tensor:
    """Per-timestep relative heading angle difference.

    Args:
        pos_i, pos_j: [E, T, 2]  absolute positions of target and source
    Returns:
        [E, T, 1]  angle difference wrapped to [-π, π]
                   — positive: j is turning more CCW than i
                   — ≈ 0: parallel motion; ≈ ±π: opposing motion
    """
    vel_i = pos_i[:, 1:] - pos_i[:, :-1]          # [E, T-1, 2]
    vel_j = pos_j[:, 1:] - pos_j[:, :-1]          # [E, T-1, 2]

    h_i = torch.atan2(vel_i[..., 1], vel_i[..., 0])   # [E, T-1]
    h_j = torch.atan2(vel_j[..., 1], vel_j[..., 0])   # [E, T-1]

    diff = torch.atan2(torch.sin(h_j - h_i),
                       torch.cos(h_j - h_i))           # [E, T-1] wrapped

    # Repeat last step so length matches T
    diff = torch.cat([diff, diff[:, -1:]], dim=1)      # [E, T]
    return diff.unsqueeze(-1)                          # [E, T, 1]


class FutureInteractionGraphV5(FutureInteractionGraphV4):
    """Sparse interaction graph with uncertainty-aware semantic edge scorer.

    Extra constructor kwarg (beyond V4):
        y0_score_dim (int, default 32): projection dim for y_0_hat trajectory
                                        embedding used in the scorer.
    """

    def __init__(self, embed_dim: int, future_steps: int, num_agents: int,
                 num_heads: int = 4, dropout: float = 0.1,
                 num_gnn_layers: int = 2, time_dim: int = 128,
                 top_n_neighbors: int = 5, rel_traj_hidden: int = 32,
                 y0_score_dim: int = 32):
        super().__init__(
            embed_dim       = embed_dim,
            future_steps    = future_steps,
            num_agents      = num_agents,
            num_heads       = num_heads,
            dropout         = dropout,
            num_gnn_layers  = num_gnn_layers,
            time_dim        = time_dim,
            top_n_neighbors = top_n_neighbors,
            rel_traj_hidden = rel_traj_hidden,
        )
        self.y0_score_dim = y0_score_dim

        # Replace V4's encoder (in_channels=2) with one that also accepts
        # the relative heading angle channel → in_channels=3.
        self.rel_traj_encoder = RelTrajEncoder(
            out_dim     = embed_dim,
            T           = future_steps,
            D_hidden    = rel_traj_hidden,
            num_heads   = 4,
            in_channels = 3,           # rel_pos(2) + heading_diff(1)
        )

        # Project each agent's clean y_0_hat trajectory to scoring space.
        # Input: flattened future trajectory [T*2].
        self.y0_score_proj = nn.Sequential(
            nn.Linear(future_steps * 2, y0_score_dim),
            nn.ReLU(inplace=True),
        )

        # Pairwise interaction scorer.
        #
        # Step 1 — separate projections for i and j:
        self.score_proj_i = nn.Linear(y0_score_dim, y0_score_dim)
        self.score_proj_j = nn.Linear(y0_score_dim, y0_score_dim)
        #
        # Step 2 — explicit pairwise interaction terms:
        #   h_i * h_j  [D_s] — element-wise similarity
        #   h_i - h_j  [D_s] — directional asymmetry (i doing X, j doing Y)
        #   → cat → [2*D_s]
        #
        # Step 3 — head combines interaction + uncertainty + geometry:
        #   [2*D_s + σ_i(1) + σ_j(1) + mean_rel(2) + std_rel(2) + min_dist(1)]
        head_in = y0_score_dim * 2 + 7
        self.edge_scorer = nn.Sequential(
            nn.Linear(head_in, 32),
            nn.ReLU(inplace=True),
            nn.Linear(32, 1),
        )

    # ------------------------------------------------------------------
    # Forward
    # ------------------------------------------------------------------

    def forward(
        self,
        y_emb:       torch.Tensor,             # [B, K, A, D]
        y_abs:       torch.Tensor,             # [B, K, A, T, 2]  ← y_0_hat unnorm
        t_emb:       torch.Tensor,             # [B, D]
        tau:         torch.Tensor,             # [B] ∈ [0, 1]
        sigma_agent: torch.Tensor = None,      # [B, K, A, T] or None
    ) -> torch.Tensor:                         # [B, K, A, D]
        B, K, A, D = y_emb.shape
        T  = y_abs.shape[3]
        E0 = self._E0
        N  = self.top_n

        # ---- Full per-mode relative trajectories  [B*K*E0, T, 2] ----------
        pos_bk        = y_abs.reshape(B * K * A, T, 2)
        edge_index_bk = self._make_batched_edge_index(B * K)    # [2, B*K*E0]

        pos_i_t   = pos_bk[edge_index_bk[1]]                    # [B*K*E0, T, 2]
        pos_j_t   = pos_bk[edge_index_bk[0]]                    # [B*K*E0, T, 2]
        rel_pos_t = pos_j_t - pos_i_t                           # [B*K*E0, T, 2]

        # ---- Geometric features (all edges) ------------------------------
        mean_rel  = rel_pos_t.mean(dim=1)                        # [E, 2]
        std_rel   = rel_pos_t.std(dim=1)                         # [E, 2]
        min_dist  = (rel_pos_t.norm(dim=-1)
                              .min(dim=1).values
                              .unsqueeze(-1))                    # [E, 1]

        # ---- Semantic features from clean y_0_hat  -----------------------
        y0_flat   = y_abs.reshape(B * K * A, T * 2)             # [B*K*A, T*2]
        y0_emb    = self.y0_score_proj(y0_flat)                  # [B*K*A, D_s]

        y0_emb_i  = y0_emb[edge_index_bk[1]]                    # [E, D_s]  target
        y0_emb_j  = y0_emb[edge_index_bk[0]]                    # [E, D_s]  source

        # ---- Uncertainty features ----------------------------------------
        if sigma_agent is not None:
            sigma_mean  = sigma_agent.mean(dim=-1)               # [B, K, A]
            sigma_bka   = sigma_mean.reshape(B * K * A)          # [B*K*A]
            sigma_i     = sigma_bka[edge_index_bk[1]].unsqueeze(-1)  # [E, 1]
            sigma_j     = sigma_bka[edge_index_bk[0]].unsqueeze(-1)  # [E, 1]
            tau_bka     = sigma_bka
        else:
            # Pass 1: no uncertainty available — zero sigma features
            sigma_i = torch.zeros(rel_pos_t.size(0), 1, device=y_abs.device)
            sigma_j = torch.zeros_like(sigma_i)
            tau_bka = (tau
                       .unsqueeze(1).unsqueeze(2)
                       .expand(-1, K, A)
                       .reshape(B * K * A))

        # ---- Enhanced scorer  --------------------------------------------
        # Pairwise interaction: separate projections then explicit relation terms
        h_i = self.score_proj_i(y0_emb_i)                       # [E, D_s]
        h_j = self.score_proj_j(y0_emb_j)                       # [E, D_s]
        interact = torch.cat([h_i * h_j,                         # similarity
                               h_i - h_j], dim=-1)               # [E, 2*D_s]

        score_feat = torch.cat([
            interact,   # pairwise relation                [E, 2*D_s]
            sigma_i,    # target uncertainty                [E, 1]
            sigma_j,    # source uncertainty                [E, 1]
            mean_rel,   # mean relative position            [E, 2]
            std_rel,    # spatial spread over time          [E, 2]
            min_dist,   # closest approach distance         [E, 1]
        ], dim=-1)                                               # [E, 2*D_s+7]
        scores = self.edge_scorer(score_feat).squeeze(-1)        # [E]

        # ---- Top-N selection per target agent  ---------------------------
        scores_grouped = scores.view(B * K * A, A - 1)
        _, top_idx     = scores_grouped.topk(N, dim=-1, sorted=False)
        mask           = torch.zeros(B * K * A, A - 1,
                                     device=scores.device, dtype=torch.bool)
        mask.scatter_(1, top_idx, True)
        mask_flat      = mask.view(-1)                           # [B*K*E0]

        # ---- RelTrajEncoder on selected edges only  ----------------------
        # ---- Heading diff on selected edges only  ------------------------
        heading   = _heading_diff(pos_i_t[mask_flat],
                                  pos_j_t[mask_flat])            # [B*K*A*N, T, 1]
        rel_pos_sparse = torch.cat(
            [rel_pos_t[mask_flat], heading], dim=-1
        )                                                        # [B*K*A*N, T, 3]

        if sigma_agent is not None:
            sigma_full  = sigma_agent.reshape(B * K * A, T)
            sigma_i_t   = sigma_full[edge_index_bk[1][mask_flat]]
            sigma_j_t   = sigma_full[edge_index_bk[0][mask_flat]]
            sigma_bias  = sigma_i_t - sigma_j_t                 # [B*K*A*N, T]
        else:
            sigma_bias  = None

        edge_attr_sparse = self.rel_traj_encoder(
            rel_pos_sparse, sigma_bias
        )                                                        # [B*K*A*N, D]

        # ---- Sparse GNN pass  --------------------------------------------
        edge_index_sparse = edge_index_bk[:, mask_flat]

        temb_bka = (t_emb
                    .unsqueeze(1).unsqueeze(2)
                    .expand(-1, K, A, -1)
                    .reshape(B * K * A, D))

        nodes = y_emb.reshape(B * K * A, D)
        for layer in self.gnn_layers:
            nodes = layer(nodes, edge_index_sparse, edge_attr_sparse,
                          temb_agent=temb_bka, tau=tau_bka)

        # ---- Gated residual  ---------------------------------------------
        orig = y_emb.reshape(B * K * A, D)
        gate = self.gate_proj(torch.cat([orig, nodes], dim=-1))
        out  = orig + gate * self.out_proj(nodes)

        return out.view(B, K, A, D)