"""
FutureInteractionGraphV2 — mode-specific edge features.

Difference from V1 (graph_interaction_nba.py):
  V1 averages the K predicted trajectories into one representative trajectory
  per scene, computes edge features from that average, and then replicates them
  across all K modes.  Every mode therefore sees identical inter-agent geometry.

  V2 computes edge features *separately for each mode*, so agent i's message
  from agent j in mode k reflects that mode's specific predicted positions,
  not the cross-mode mean.

Architecture and GNN code (FutureGeomAttnLayer) are unchanged — only the
edge-feature-construction step in forward() is different.

Memory note:
  V1 computed edges at B scale (27.5 K edges) then replicated.
  V2 computes edges directly at B*K scale (550 K edges) in one step.
  The GNN forward pass size is identical in both versions (55 K nodes,
  550 K edges); only the edge-feature tensor construction is larger in V2
  (B*K*E0*T*2 ≈ 88 MB extra intermediate before projection — negligible
  on an A6000).
"""

import torch
from models.graph_interaction_nba import FutureInteractionGraph


class FutureInteractionGraphV2(FutureInteractionGraph):
    """FutureInteractionGraph with per-mode (mode-specific) edge features.

    Constructor arguments are identical to FutureInteractionGraph.
    Only forward() is overridden.
    """

    def forward(
        self,
        y_emb: torch.Tensor,             # [B, K, A, D]
        y_abs: torch.Tensor,             # [B, K, A, T, 2]  absolute-ish future pos
        t_emb: torch.Tensor,             # [B, D]  time embedding from backbone
        tau:   torch.Tensor,             # [B] ∈ [0, 1]
        sigma_agent: torch.Tensor = None,# [B, K, A] per-agent uncertainty, or None
    ) -> torch.Tensor:                   # [B, K, A, D]
        B, K, A, D = y_emb.shape
        T  = y_abs.shape[3]
        E0 = self._E0              # A*(A-1) edges per scene

        # ---- Per-mode edge features (B*K scenes)  [B*K*E0, D] -------------
        # Flatten all B*K mode-trajectories together, then compute relative
        # positions per directed edge for each mode independently.
        pos_bk        = y_abs.reshape(B * K * A, T, 2)          # [B*K*A, T, 2]
        edge_index_bk = self._make_batched_edge_index(B * K)     # [2, B*K*E0]

        # Mean relative position over future timesteps for each (mode, edge)
        rel_pos_bk = (pos_bk[edge_index_bk[0]] -
                      pos_bk[edge_index_bk[1]]).mean(dim=1)      # [B*K*E0, 2]
        edge_attr_bk = self.rel_pos_proj(rel_pos_bk)             # [B*K*E0, D]

        # ---- Per-agent time / uncertainty embeddings  [B*K*A, D] ----------
        temb_bka = (t_emb
                    .unsqueeze(1).unsqueeze(2)
                    .expand(-1, K, A, -1)
                    .reshape(B * K * A, D))
        if sigma_agent is not None:
            s = sigma_agent.mean(dim=-1) if sigma_agent.dim() == 4 else sigma_agent
            tau_bka = s.reshape(B * K * A)
        else:
            tau_bka  = (tau
                        .unsqueeze(1).unsqueeze(2)
                        .expand(-1, K, A)
                        .reshape(B * K * A))

        # ---- GNN pass on B*K*A nodes  -------------------------------------
        nodes = y_emb.reshape(B * K * A, D)
        for layer in self.gnn_layers:
            nodes = layer(nodes, edge_index_bk, edge_attr_bk,
                          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)