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


class RotaryEmbedding(nn.Module):
    def __init__(self, dim, max_seq_len=2048):
        super().__init__()
        inv_freq = 1.0 / (10000 ** (torch.arange(0, dim, 2).float() / dim))
        self.register_buffer("inv_freq", inv_freq)
        self._set_cos_sin_cache(max_seq_len)

    def _set_cos_sin_cache(self, seq_len):
        t = torch.arange(seq_len, device=self.inv_freq.device)
        freqs = torch.einsum("i,j->ij", t, self.inv_freq)
        emb = torch.cat((freqs, freqs), dim=-1)
        self.register_buffer("cos", emb.cos()[None, None, :, :], persistent=False)
        self.register_buffer("sin", emb.sin()[None, None, :, :], persistent=False)

    def _rotate_half(self, x):
        x1 = x[..., :x.shape[-1] // 2]
        x2 = x[..., x.shape[-1] // 2:]
        return torch.cat((-x2, x1), dim=-1)

    def forward(self, x):
        seq_len = x.shape[-2]
        return (x * self.cos[:, :, :seq_len, :] + 
                self._rotate_half(x) * self.sin[:, :, :seq_len, :])

class TransformerLayer(nn.Module):
    def __init__(self, token_dim, num_heads, dim_feedforward, dropout):
        super().__init__()
        self.num_heads = num_heads
        self.head_dim = token_dim // num_heads
        self.qkv = nn.Linear(token_dim, token_dim * 3)
        self.out_proj = nn.Linear(token_dim, token_dim)
        self.norm1 = nn.LayerNorm(token_dim)
        self.norm2 = nn.LayerNorm(token_dim)
        self.ffn = nn.Sequential(
            nn.Linear(token_dim, dim_feedforward),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(dim_feedforward, token_dim),
            nn.Dropout(dropout)
        )
        self.dropout = dropout

    def forward(self, x, mask=None):
        # Attention
        residual = x
        x = self.norm1(x)
        qkv = self.qkv(x).chunk(3, dim=-1)
        q, k, v = map(lambda t: t.view(t.shape[0], t.shape[1], self.num_heads, self.head_dim).transpose(1, 2), qkv)
        
        # 使用 SDPA
        x = F.scaled_dot_product_attention(q, k, v, attn_mask=mask, dropout_p=self.dropout if self.training else 0.0)
        x = x.transpose(1, 2).reshape(x.shape[0], x.shape[2], -1)
        x = self.out_proj(x) + residual
        
        # FFN
        x = x + self.ffn(self.norm2(x))
        return x

class TemporalTrajectoryPredictor(nn.Module):
    def __init__(
        self,
        num_tokens=512,
        token_dim=768,
        max_seq_len=64,
        num_layers=16,
        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.rope = RotaryEmbedding(token_dim, max_seq_len)
        self.layers = nn.ModuleList([
            TransformerLayer(token_dim, num_heads, dim_feedforward, dropout) 
            for _ in range(num_layers)
        ])
        self.output_projection = nn.Linear(token_dim, token_dim)
        self._init_weights()

    def _init_weights(self):
        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)
        
        # 应用 RoPE
        x = self.rope(x.unsqueeze(1)).squeeze(1)
        
        # if mask is None:
        #     mask = torch.triu(torch.full((t, t), float('-inf'), device=x.device), diagonal=1)
        #     # SDPA 的 mask 如果是 boolean，True 表示要 mask 掉；如果是 float，-inf 表示要 mask 掉
        #     mask = (mask == float('-inf')) 
        if mask is None:
            mask = torch.ones((t, t), device=x.device, dtype=torch.bool).tril()
            mask = mask.unsqueeze(0).unsqueeze(1)

        for layer in self.layers:
            x = layer(x, mask=mask)
            
        next_frames = self.output_projection(x)
        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=3,
        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
        self.head_dim = token_dim // 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({
                "qkv1": nn.Linear(token_dim, token_dim * 3),
                "out1": nn.Linear(token_dim, token_dim),
                "ln1": nn.LayerNorm(token_dim),
                "qkv2": nn.Linear(token_dim, token_dim * 3),
                "out2": nn.Linear(token_dim, token_dim),
                "ln2": nn.LayerNorm(token_dim),
                "q_ca": nn.Linear(token_dim, token_dim),
                "kv_ca": nn.Linear(text_dim, token_dim * 2),
                "out_ca": nn.Linear(token_dim, token_dim),
                "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)
        # 调整形状以匹配 SDPA 的 attention_mask: (B*t, H, N, N)
        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 _sdpa_attention(self, q, k, v, bias=None):
        return F.scaled_dot_product_attention(q, k, v, attn_mask=bias, dropout_p=0.1 if self.training else 0.0)

    def forward(self, next_frame_init, text_embed, shortest_path_matrix):
        B, T, N, dim = next_frame_init.shape  # 使用大写 T 避免冲突
        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:
            # 1. Pre-LN SA1
            residual = x
            x_norm = b["ln1"](x)
            q, k, v = b["qkv1"](x_norm).chunk(3, dim=-1)
            # 显式重塑，避免使用 lambda 带来的变量遮蔽问题
            q = q.view(B * T, N, self.num_heads, self.head_dim).transpose(1, 2)
            k = k.view(B * T, N, self.num_heads, self.head_dim).transpose(1, 2)
            v = v.view(B * T, N, self.num_heads, self.head_dim).transpose(1, 2)
            attn = self._sdpa_attention(q, k, v, bias=graph_attn_bias)
            x = residual + b["out1"](attn.transpose(1, 2).reshape(B * T, N, dim))
            
            # 2. Pre-LN SA2
            residual = x
            x_norm = b["ln2"](x)
            q, k, v = b["qkv2"](x_norm).chunk(3, dim=-1)
            q = q.view(B * T, N, self.num_heads, self.head_dim).transpose(1, 2)
            k = k.view(B * T, N, self.num_heads, self.head_dim).transpose(1, 2)
            v = v.view(B * T, N, self.num_heads, self.head_dim).transpose(1, 2)
            attn = self._sdpa_attention(q, k, v, bias=graph_attn_bias)
            x = residual + b["out2"](attn.transpose(1, 2).reshape(B * T, N, dim))
            
            # 3. Pre-LN CA
            residual = x
            x_norm = b["ln3"](x)
            q = b["q_ca"](x_norm).view(B * T, N, self.num_heads, self.head_dim).transpose(1, 2)
            k_v = b["kv_ca"](text_embed_expanded)
            k, v = k_v.chunk(2, dim=-1)
            k = k.view(B * T, -1, self.num_heads, self.head_dim).transpose(1, 2)
            v = v.view(B * T, -1, self.num_heads, self.head_dim).transpose(1, 2)
            attn = self._sdpa_attention(q, k, v)
            x = residual + b["out_ca"](attn.transpose(1, 2).reshape(B * T, N, dim))
            
            # 4. Pre-LN FFN
            x = x + b["ffn"](b["ln4"](x))
            
        return x.reshape(B, T, N, dim)