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):
        """
        旋转位置编码 (RoPE)
        dim: 必须是 head_dim (token_dim // num_heads)，通常为偶数
        max_seq_len: 预设的最大序列长度
        """
        super().__init__()
        self.dim = dim
        self.max_seq_len = max_seq_len
        
        # 计算旋转频率项
        inv_freq = 1.0 / (10000 ** (torch.arange(0, dim, 2).float() / dim))
        self.register_buffer("inv_freq", inv_freq, persistent=True)
        
        # 初始化缓存
        self._set_cos_sin_cache(max_seq_len, device=torch.device("cpu"))
        

    def _set_cos_sin_cache(self, seq_len, device):
        self.max_seq_len = seq_len
        t = torch.arange(seq_len, device=device, dtype=torch.float32)
        freqs = torch.einsum("i,j->ij", t, self.inv_freq.to(device))
        emb = torch.cat((freqs, freqs), dim=-1)
        
        self.register_buffer("cos_cached", emb.cos()[None, None, :, :], persistent=False)
        self.register_buffer("sin_cached", 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):
        if seq_len > self.max_seq_len or self.cos_cached.device != x.device:
            self._set_cos_sin_cache(seq_len=max(seq_len, self.max_seq_len), device=x.device)
            
        cos = self.cos_cached[..., :seq_len, :].to(dtype=x.dtype)
        sin = self.sin_cached[..., :seq_len, :].to(dtype=x.dtype)
        return (x * cos) + (self._rotate_half(x) * sin)


class TransformerLayer(nn.Module):
    """标准的自注意力时序层 (Self-Attention)"""
    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, rope=None, mask=None):
        residual = x
        x = self.norm1(x)
        B_N, T, _ = x.shape
        
        qkv = self.qkv(x).chunk(3, dim=-1)
        q, k, v = map(
            lambda t: t.view(B_N, T, self.num_heads, self.head_dim).transpose(1, 2).contiguous(), 
            qkv
        )
        
        if rope is not None:
            q = rope(q, seq_len=T)
            k = rope(k, seq_len=T)
        
        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(B_N, T, -1).contiguous()
        x = self.out_proj(x) + residual
        
        x = x + self.ffn(self.norm2(x))
        return x


class CrossAttentionLayer(nn.Module):
    """专门用于时序特征与文本条件对齐的交叉注意力层 (Cross-Attention)"""
    def __init__(self, token_dim, text_dim, num_heads, dim_feedforward, dropout):
        super().__init__()
        self.num_heads = num_heads
        self.head_dim = token_dim // num_heads
        
        self.q_proj = nn.Linear(token_dim, token_dim)
        self.kv_proj = nn.Linear(text_dim, token_dim * 2)
        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, text_embed):
        # x 形状: (B_N, T, token_dim)
        # text_embed 形状: (B_N, L_text, text_dim) -> 外部需保证已对齐至 B_N 轴
        residual = x
        x = self.norm1(x)
        B_N, T, _ = x.shape
        
        q = self.q_proj(x).view(B_N, T, self.num_heads, self.head_dim).transpose(1, 2).contiguous()
        
        k_v = self.kv_proj(text_embed)
        k, v = k_v.chunk(2, dim=-1)
        k = k.view(B_N, -1, self.num_heads, self.head_dim).transpose(1, 2).contiguous()
        v = v.view(B_N, -1, self.num_heads, self.head_dim).transpose(1, 2).contiguous()
        
        # 文本条件通常不需要时序的因果 Mask
        x = F.scaled_dot_product_attention(
            q, k, v, 
            attn_mask=None, 
            dropout_p=self.dropout if self.training else 0.0
        )
        
        x = x.transpose(1, 2).reshape(B_N, T, -1).contiguous()
        x = self.out_proj(x) + residual
        
        x = x + self.ffn(self.norm2(x))
        return x


class TemporalTrajectoryPredictor(nn.Module):
    def __init__(
        self,
        num_tokens=512,
        token_dim=768,
        text_dim=64,      # 新增 text_dim 输入参数
        max_seq_len=64,
        num_blocks=5,       # 重构为大块数量
        layers_per_block=4, # 每个大块内部自注意力层数
        num_heads=8,
        dim_feedforward=2048,
        dropout=0.05
    ):
        super().__init__()
        self.num_tokens = num_tokens
        self.token_dim = token_dim
        self.max_seq_len = max_seq_len
        self.num_heads = num_heads
        self.head_dim = token_dim // num_heads

        self.rope = RotaryEmbedding(self.head_dim, max_seq_len)
        
        # 动态构建层序列
        self.layers = nn.ModuleList()
        for _ in range(num_blocks):
            # 添加 4 层 Self-Attention 层
            for _ in range(layers_per_block):
                self.layers.append(
                    TransformerLayer(token_dim, num_heads, dim_feedforward, dropout)
                )
            # 紧接着添加 1 层 Cross-Attention 层
            self.layers.append(
                CrossAttentionLayer(token_dim, text_dim, num_heads, dim_feedforward, dropout)
            )
        
        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, text_embed, mask=None, predict_next_only=False):
        """
        Args:
            history_frames: 4D 历史帧输入，形状为 (B, t, N, dim)
            text_embed: 文本条件，形状为 (B, 1, text_dim) 或 (B, L_text, text_dim)
            mask: 可选。自定义注意力掩码
            predict_next_only: 是否只预测未来下一帧
        """
        target_dtype = self.output_projection.weight.dtype 
        
        history_frames = history_frames.to(dtype=target_dtype)
        text_embed = text_embed.to(dtype=target_dtype)

        B, t, N, dim = history_frames.shape
        
        # 空间维度 N 转入 Batch 轴，暴露出纯粹的时间序列维度 t: (B * N, t, dim)
        x = history_frames.transpose(1, 2).reshape(B * N, t, dim).contiguous()
        
        # 同步将 text_embed 扩展到与 x 相同的 Batch 维度 (B * N)
        # text_embed 原始: (B, L, text_dim) -> 扩展后: (B, N, L, text_dim) -> 重塑为: (B * N, L, text_dim)
        L_text = text_embed.shape[1]
        text_dim_size = text_embed.shape[-1]
        text_embed_expanded = text_embed.unsqueeze(1).expand(-1, N, -1, -1).reshape(B * N, L_text, text_dim_size).contiguous()
        
        if mask is None:
            mask = torch.ones((t, t), device=x.device, dtype=torch.bool).tril()
            mask = mask.unsqueeze(0).unsqueeze(1)
        elif mask.dtype != torch.bool:
            mask = (mask < 0) if mask.dtype.is_floating_point else mask.bool()

        # 遍历执行所有层
        for layer in self.layers:
            if isinstance(layer, TransformerLayer):
                # 自注意力机制：传入 rope 和时序因果掩码 mask
                x = layer(x, rope=self.rope, mask=mask)
            elif isinstance(layer, CrossAttentionLayer):
                # 交叉注意力机制：传入对齐后的 text_embed_expanded
                x = layer(x, text_embed=text_embed_expanded)
            
        if predict_next_only:
            x_last = x[:, -1:, :]  
            next_frame = self.output_projection(x_last)  
            return next_frame.reshape(B, N, 1, dim).transpose(1, 2).contiguous()
        else:
            all_preds = self.output_projection(x)  
            return all_preds.reshape(B, N, t, dim).transpose(1, 2).contiguous()


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.05
    ):
        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)
        ])

        # self.delta_proj = nn.Linear(token_dim, token_dim)
        # nn.init.zeros_(self.delta_proj.weight)
        # nn.init.zeros_(self.delta_proj.bias)

    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) 
        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.05 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  
        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)
            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))
            
        # delta = self.delta_proj(x).reshape(B, T, N, dim)
        return x.reshape(B, T, N, dim)