"""
Future-trajectory interaction graph for LED.

Idea (ported from MoFlow's FutureInteractionGraphV6):
  At each leapfrog denoising step, convert the frozen core denoiser's
  epsilon prediction into an implied y_0 estimate, then run a small
  agent-to-agent graph on that future geometry and emit a residual
  correction added back to epsilon. The module is zero-initialized at
  the output so it begins as a no-op and only nudges the reverse chain
  as it learns.
"""

import torch
import torch.nn as nn
import torch.nn.functional as F


class FutureInteractionGraph(nn.Module):
    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,
        num_denoise_steps: int = 5,
        use_sigma: bool = False,
    ):
        super().__init__()
        self.A = num_agents
        self.T = future_steps
        self.D = node_dim
        self.top_n = min(top_n, num_agents - 1)
        self.use_sigma = use_sigma

        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)

        if use_sigma:
            self.sigma_proj = nn.Sequential(
                nn.Linear(1, node_dim // 4),
                nn.ReLU(inplace=True),
                nn.Linear(node_dim // 4, node_dim),
            )

        # edge feature: base 5 dims + 1 sigma_diff if use_sigma
        edge_in_dim = 6 if use_sigma else 5
        self.edge_mlp = nn.Sequential(
            nn.Linear(edge_in_dim, 32),
            nn.ReLU(inplace=True),
            nn.Linear(32, node_dim),
        )

        self.attn_q = nn.Linear(node_dim, node_dim)
        self.attn_k = nn.Linear(node_dim * 2, node_dim)
        self.attn_v = nn.Linear(node_dim * 2, node_dim)
        self.scale = node_dim ** -0.5

        self.fuse = nn.Sequential(
            nn.Linear(node_dim * 2, node_dim),
            nn.ReLU(inplace=True),
        )

        self.decode = nn.Sequential(
            nn.Linear(node_dim, node_dim),
            nn.ReLU(inplace=True),
            nn.Linear(node_dim, future_steps * 2),
        )
        # zero-init output so residual starts as a no-op
        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,          # current denoising step in [0, num_denoise_steps)
        sigma: torch.Tensor = None,  # [B*A, 1] per-agent uncertainty (optional)
    ) -> torch.Tensor:          # [B*A, K, T, 2]
        BA, K, T, _ = y0_hat.shape
        A = self.A
        assert BA % A == 0, f"expected B*A agents with A={A}, got {BA}"
        B = BA // A
        device = y0_hat.device

        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 = past_feat + y0_feat + step                                # [BA, K, D]

        if self.use_sigma and sigma is not None:
            sigma_feat = self.sigma_proj(sigma)                            # [BA, D]
            nodes = nodes + sigma_feat.unsqueeze(1)                        # broadcast over K

        nodes = nodes.view(B, A, K, self.D).permute(0, 2, 1, 3).contiguous()
        # nodes: [B, K, A, D]

        pos = y0_hat.view(B, A, K, T, 2).permute(0, 2, 1, 3, 4).contiguous()
        # pos: [B, K, A, T, 2]

        # pairwise relative future trajectories: row i = receiver, col j = source
        # rel[:, :, i, j] = pos_j - pos_i (where j is relative to i)
        rel = pos.unsqueeze(2) - pos.unsqueeze(3)                          # [B,K,A_i,A_j,T,2]
        mean_rel = rel.mean(dim=-2)                                        # [B,K,A,A,2]
        std_rel = rel.std(dim=-2)                                          # [B,K,A,A,2]
        dist = rel.norm(dim=-1)                                            # [B,K,A,A,T]
        min_dist = dist.min(dim=-1).values.unsqueeze(-1)                   # [B,K,A,A,1]

        if self.use_sigma and sigma is not None:
            # Per-agent sigma → pairwise sigma difference as edge feature
            sigma_bka = sigma.view(B, A, 1).permute(0, 2, 1).contiguous()  # [B, 1, A]
            sigma_i = sigma_bka.unsqueeze(3).expand(-1, K, A, A)            # [B,K,A,A]
            sigma_j = sigma_bka.unsqueeze(2).expand(-1, K, A, A)            # [B,K,A,A]
            sigma_diff = (sigma_i - sigma_j).unsqueeze(-1)                  # [B,K,A,A,1]
            edge_raw = torch.cat([mean_rel, std_rel, min_dist, sigma_diff], dim=-1)  # [B,K,A,A,6]
        else:
            edge_raw = torch.cat([mean_rel, std_rel, min_dist], dim=-1)    # [B,K,A,A,5]
        edge_feat = self.edge_mlp(edge_raw)                                # [B,K,A,A,D]

        # top-N neighbor selection by closest min-distance over future horizon
        score = -min_dist.squeeze(-1)                                      # [B,K,A,A]
        diag = torch.eye(A, dtype=torch.bool, device=device)
        score = score.masked_fill(diag, float('-inf'))
        _, top_idx = score.topk(self.top_n, dim=-1)                        # [B,K,A,N]

        idx_node = top_idx.unsqueeze(-1).expand(-1, -1, -1, -1, self.D)
        nodes_j = nodes.unsqueeze(2).expand(-1, -1, A, -1, -1)             # [B,K,A_i,A_j,D]
        neigh_nodes = torch.gather(nodes_j, 3, idx_node)                   # [B,K,A,N,D]

        idx_edge = top_idx.unsqueeze(-1).expand(-1, -1, -1, -1, self.D)
        neigh_edges = torch.gather(edge_feat, 3, idx_edge)                 # [B,K,A,N,D]

        q = self.attn_q(nodes).unsqueeze(-2)                               # [B,K,A,1,D]
        kv_in = torch.cat([neigh_nodes, neigh_edges], dim=-1)              # [B,K,A,N,2D]
        k = self.attn_k(kv_in)
        v = self.attn_v(kv_in)
        attn = (q * k).sum(dim=-1) * self.scale                            # [B,K,A,N]
        attn = F.softmax(attn, dim=-1)
        msg = (attn.unsqueeze(-1) * v).sum(dim=-2)                         # [B,K,A,D]

        fused = self.fuse(torch.cat([nodes, msg], dim=-1))                 # [B,K,A,D]
        residual = self.decode(fused).view(B, K, A, T, 2)
        residual = residual.permute(0, 2, 1, 3, 4).contiguous().view(BA, K, T, 2)
        return residual