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

class TemporalTrajectoryPredictor(nn.Module):
    def __init__(
        self,
        num_tokens=512,
        token_dim=768,
        max_seq_len=64,
        num_layers=12,
        num_heads=8,
        dim_feedforward=2048,
        dropout=0.1
    ):
        super().__init__()
        self.num_tokens = num_tokens
        self.token_dim = token_dim
        self.max_seq_len = max_seq_len

        self.time_pos_embed = nn.Parameter(torch.zeros(1, max_seq_len, token_dim))
        
        encoder_layer = nn.TransformerEncoderLayer(
            d_model=token_dim,
            nhead=num_heads,
            dim_feedforward=dim_feedforward,
            dropout=dropout,
            activation='gelu',
            batch_first=True
        )
        self.temporal_transformer = nn.TransformerEncoder(encoder_layer, num_layers=num_layers)
        self.output_projection = nn.Linear(token_dim, token_dim)
        
        self._init_weights()

    def _init_weights(self):
        nn.init.normal_(self.time_pos_embed, std=0.02)
        nn.init.xavier_uniform_(self.output_projection.weight)
        if self.output_projection.bias is not None:
            nn.init.constant_(self.output_projection.bias, 0)

    def forward(self, history_frames, mask=None):
        B, t, N, dim = history_frames.shape
        
        x = history_frames.transpose(1, 2).reshape(B * N, t, dim)
        x = x + self.time_pos_embed[:, :t, :]
        
        if mask is None:
            mask = torch.triu(torch.full((t, t), float('-inf'), device=x.device), diagonal=1)
        
        feat = self.temporal_transformer(x, mask=mask)
        next_frames = self.output_projection(feat)
        
        return next_frames.reshape(B, N, t, dim).transpose(1, 2)


class SpatialGraphRefiner(nn.Module):
    def __init__(
        self,
        num_tokens=512,
        token_dim=768,
        text_dim=1024,
        max_graph_dist=10,
        num_blocks=2,
        num_heads=8,
        dim_feedforward=2048,
        dropout=0.1
    ):
        super().__init__()
        self.num_tokens = num_tokens
        self.token_dim = token_dim
        self.max_graph_dist = max_graph_dist
        self.num_heads = num_heads
        
        distances = torch.arange(max_graph_dist + 1, dtype=torch.float32)
        scales = torch.linspace(0.1, 1.0, num_heads)
        self.register_buffer("graph_bias_table", -1.0 * (distances.unsqueeze(1) * scales.unsqueeze(0)))
        
        self.blocks = nn.ModuleList([
            nn.ModuleDict({
                "sa1": nn.MultiheadAttention(token_dim, num_heads, dropout=dropout, batch_first=True),
                "ln1": nn.LayerNorm(token_dim),
                "sa2": nn.MultiheadAttention(token_dim, num_heads, dropout=dropout, batch_first=True),
                "ln2": nn.LayerNorm(token_dim),
                "ca": nn.MultiheadAttention(token_dim, num_heads, kdim=text_dim, vdim=text_dim, dropout=dropout, batch_first=True),
                "ln3": nn.LayerNorm(token_dim),
                "ffn": nn.Sequential(
                    nn.Linear(token_dim, dim_feedforward),
                    nn.GELU(),
                    nn.Dropout(dropout),
                    nn.Linear(dim_feedforward, token_dim),
                    nn.Dropout(dropout)
                ),
                "ln4": nn.LayerNorm(token_dim)
            }) for _ in range(num_blocks)
        ])

    def _compute_graph_bias(self, shortest_path_matrix, B, t):
        N = shortest_path_matrix.shape[-1]
        dist_clipped = torch.clamp(shortest_path_matrix, 0, self.max_graph_dist).long()
        bias = F.embedding(dist_clipped, self.graph_bias_table) # (B, N, N, H)
        return bias.unsqueeze(1).expand(-1, t, -1, -1, -1).permute(0, 1, 4, 2, 3).reshape(B * t * self.num_heads, N, N)

    def forward(self, next_frame_init, text_embed, shortest_path_matrix):
        B, t, N, dim = next_frame_init.shape
        graph_attn_bias = self._compute_graph_bias(shortest_path_matrix, B, t)
        
        x = next_frame_init.reshape(B * t, N, dim)
        text_embed_expanded = text_embed.unsqueeze(1).expand(-1, t, -1, -1).reshape(B * t, -1, text_embed.shape[-1])
        
        for b in self.blocks:
            attn, _ = b["sa1"](x, x, x, attn_mask=graph_attn_bias)
            x = b["ln1"](x + attn)
            
            attn, _ = b["sa2"](x, x, x, attn_mask=graph_attn_bias)
            x = b["ln2"](x + attn)
            
            attn, _ = b["ca"](x, text_embed_expanded, text_embed_expanded)
            x = b["ln3"](x + attn)
            
            x = b["ln4"](x + b["ffn"](x))
            
        return x.reshape(B, t, N, dim)