"""注意力与 Transformer 原语(从 pure_iflame_with_condition 提取, 只保留 MLA 及其依赖)。

MLA 是细分器 subdivider.DiscreteDiffusionTransformer 使用的多头潜在注意力块。
原目录里的 MeshRes / MLCA / LinearDownsample / LinearUpsample / MultiheadCrossFlashrope
都不在推理路径上, 已删除。
"""

import torch
import torch.nn as nn
import math
import torch.nn.functional as F
from torch import nn
import time #cal_time new at 10.21
import numpy as np
from typing import List, Optional, Tuple
from tqdm import tqdm
from einops import rearrange, repeat, pack


# ---- 原 rms_norm.py ----

class RMSNorm(nn.Module):
    def __init__(self, dim: int, eps: float = 1e-6, elementwise_affine=True, memory_efficient=False):
        super().__init__()
        self.dim = dim
        self.eps = eps
        self.elementwise_affine = elementwise_affine
        if self.elementwise_affine:
            self.weight = nn.Parameter(torch.ones(dim))
        else:
            self.register_parameter('weight', None)

    def _norm(self, x):
        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)

    def forward(self, x):
        output = self._norm(x.float()).type_as(x)
        if self.weight is not None:
            output = output * self.weight
        return output

    def extra_repr(self) -> str:
        return f'dim={self.dim}, eps={self.eps}, elementwise_affine={self.elementwise_affine}'

# ---- 原 multihead_flashdiff_2.py ----

class Rotary(nn.Module):
    def __init__(self, dim: int, max_seq_len=9*50000, dtype=torch.float16):
        super().__init__()
        # half-truncate RoPE by @YouJiacheng (w/ base freq tuning)
        angular_freq = (1 / 1024) ** torch.linspace(0, 1, steps=dim//4, dtype=dtype)
        angular_freq = torch.cat([angular_freq, angular_freq.new_zeros(dim//4)])
        t = torch.arange(max_seq_len, dtype=dtype)
        theta = torch.einsum("i,j -> ij", t, angular_freq)
        self.register_buffer('cos', theta.cos(), persistent=False)
        self.register_buffer('sin', theta.sin(), persistent=False)
        self.max_seq_len = max_seq_len
        self.dtype = dtype

    def forward(self, x_BTHD, offset=0):
        """
        Args:
            x_BTHD: Input tensor of shape [batch, seq_len, num_heads, dim]
            offset: Position offset for the sequence
        """
        batch_size, seq_len = x_BTHD.size(0), x_BTHD.size(-3)
        assert torch.max(offset + seq_len) <= self.max_seq_len, f"Sequence length {offset + seq_len} exceeds maximum length {self.max_seq_len}"

        # Get the appropriate slice of cos/sin based on offset
        positions = torch.arange(seq_len, device=x_BTHD.device).expand(batch_size, -1) + offset.unsqueeze(-1)
        cos = self.cos[positions].unsqueeze(-2)
        sin = self.sin[positions].unsqueeze(-2)

        x1, x2 = x_BTHD.to(dtype=self.dtype).chunk(2, dim=-1)
        y1 = x1 * cos - x2 * sin
        y2 = x1 * sin + x2 * cos
        return torch.cat((y1, y2), dim=-1)

class MultiheadFlashrope(nn.Module):
    def __init__(self, args, embed_dim, num_heads):
        super().__init__()
        self.args = args
        kvcache_method, kv_cache_window_size, window_stride = args
        self.embed_dim = embed_dim
        self.num_heads = num_heads  
        self.num_kv_heads = num_heads
        self.head_dim = embed_dim // num_heads 
        self.scaling = self.head_dim ** -0.5
        self.kv_cache_window_size = kv_cache_window_size
        self.window_stride = window_stride
        self.kvcache_method = kvcache_method
        # KV cache related attributes
        self.kv_cache_enabled = False
        self.k_cache = None
        self.v_cache = None
        self.cache_pos = 0  # Track position in cache
        
        self.q_proj = nn.Linear(embed_dim, embed_dim, bias=False)
        self.k_proj = nn.Linear(embed_dim, embed_dim, bias=False)
        self.v_proj = nn.Linear(embed_dim, embed_dim, bias=False)
        self.out_proj = nn.Linear(embed_dim, embed_dim, bias=False)
        self.rotary = Rotary(self.head_dim)
        self.sliding_window_cur_len = 0
        # self.attention_scores = None

    def forward(self, x, start_pos, context=None, attention_mask=None, causal=False):
        """
        Args:
            x: [batch_size, tgt_len, embed_dim]
            start_pos: [batch_size]
            context: Optional context for cross-attention [batch_size, src_len, embed_dim]
        """
        bsz, tgt_len, embed_dim = x.size()
        # Project inputs to q, k, v
        q = self.q_proj(x)
        k = self.k_proj(context if context is not None else x)
        v = self.v_proj(context if context is not None else x)

        # Reshape to [batch_size, seq_len, num_heads, head_dim]
        q = q.view(bsz, -1, self.num_heads, self.head_dim)
        k = k.view(bsz, -1, self.num_heads, self.head_dim)
        v = v.view(bsz, -1, self.num_heads, self.head_dim)
        q = self.rotary(q, offset=start_pos)
        # Normal operation without cache
        k = self.rotary(k, offset=start_pos)
        if attention_mask is not None:
            attn_output = self.attention_pytorch_optimized(q, k, v, mask=attention_mask)
        else:
            attn_output = self.attention_pytorch_optimized(q, k, v, causal=causal)

        # Reshape output
        attn_output = attn_output.contiguous().view(bsz, tgt_len, embed_dim)
        attn_output = self.out_proj(attn_output)

        return attn_output
    
    def attention_pytorch_optimized(self, q, k, v, mask=None, causal=False):
        b, seq_q, head_num, c = q.shape
        seq_k = k.shape[1]
        q = q.permute(0, 2, 1, 3)
        k = k.permute(0, 2, 1, 3)
        v = v.permute(0, 2, 1, 3)
        if mask is not None:
            atten_mask_npu = mask.unsqueeze(1).expand(b, head_num, seq_q, seq_k)
            atten_mask_npu = torch.logical_not(atten_mask_npu)
            attn_output = F.scaled_dot_product_attention(q, k, v, attn_mask=atten_mask_npu)
        else:
            if causal:
                attn_output = F.scaled_dot_product_attention(q, k, v, is_causal=True)
            else:
                attn_output = F.scaled_dot_product_attention(q, k, v)

        return attn_output.permute(0, 2, 1, 3)

    def attention_pytorch(self, q, k, v, causal=False):
        # [batch_size, num_heads, seq_len, head_dim]
        q = q.permute(0, 2, 1, 3)
        k = k.permute(0, 2, 1, 3)
        v = v.permute(0, 2, 1, 3)

        attn_weights = torch.matmul(q, k.transpose(-2, -1)) * self.scaling

        if causal:
            seq_len = q.size(-2)
            causal_mask = torch.triu(
                torch.ones(seq_len, seq_len, dtype=torch.bool, device=q.device),
                diagonal=1
            )
            attn_weights.masked_fill_(causal_mask, float('-inf'))
  
        attn_weights = F.softmax(attn_weights, dim=-1)
        
  
        attn_output = torch.matmul(attn_weights, v)
        

        return attn_output.permute(0, 2, 1, 3)

    def empty_kv_cache(self, batch_size: int, dtype: torch.dtype):
        """Initialize empty KV cache"""
        self.kv_cache_enabled=True

        device = next(self.parameters()).device
        kv_cache_maxlen = self.kv_cache_window_size
        # Initialize empty caches
        k_cache = torch.zeros(
            batch_size, 
            kv_cache_maxlen, 
            self.num_heads,
            self.head_dim,
            dtype=dtype,
            device=device
        )
        self.k_cache=k_cache

        v_cache = torch.zeros(
            batch_size,
            kv_cache_maxlen,
            self.num_heads,
            self.head_dim,
            dtype=dtype,
            device=device
        )       
        self.v_cache=v_cache
        self.cache_pos = 0  # Reset cache position

        element_size = {
            torch.float32: 4,
            torch.float16: 2,
            torch.bfloat16: 2,
        }[dtype]
        k_cache_size = (batch_size * 
                    kv_cache_maxlen * 
                    self.num_heads * 
                    self.head_dim * 
                    element_size)
        v_cache_size = k_cache_size 
        total_cache_size = k_cache_size + v_cache_size
        cache_size_gb = total_cache_size / (1024**3)
        return cache_size_gb
        

    def reset_cache(self):
        """Reset KV cache state"""
        self.k_cache = None
        self.v_cache = None
        self.cache_pos = 0

class CrossAttention(nn.Module):
    def __init__(self, dim, context_dim, n_heads, dropout=0.0):
        super().__init__()
        self.n_heads = n_heads
        self.scale = (dim // n_heads) ** -0.5

        self.q_proj = nn.Linear(dim, dim, bias=False)
        self.kv_proj = nn.Linear(context_dim, 2 * dim, bias=False)
        self.out_proj = nn.Linear(dim, dim, bias=False)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x, context):
        # x: [batch, seq_len, dim], context: [batch, context_len, context_dim]
        B, N, C = x.shape
        _, M, _ = context.shape
        H = self.n_heads

        # Linear projections
        q = self.q_proj(x).view(B, N, H, C // H)  # [B, H, seq_len, dim//H]
        k, v = self.kv_proj(context).chunk(2, dim=-1)
        k = k.view(B, M, H, C // H)  # [B, context_len, H, dim//H]
        v = v.view(B, M, H, C // H)  # [B, context_len, H, dim//H]
        
        attn_output = self.attention_pytorch_optimized(q,k,v)
        attn_output = attn_output.contiguous().view(B, N, C)
        
        # q = self.q_proj(x).view(B, N, H, C // H).transpose(1, 2)  # [B, H, seq_len, dim//H]
        # k, v = self.kv_proj(context).chunk(2, dim=-1)
        # k = k.view(B, M, H, C // H).transpose(1, 2)  # [B, H, context_len, dim//H]
        # v = v.view(B, M, H, C // H).transpose(1, 2)  # [B, H, context_len, dim//H]
        # # Scaled dot-product attention
        # attn_weights = (q @ k.transpose(-2, -1)) * self.scale  # [B, H, seq_len, context_len]
        # attn_weights = attn_weights.softmax(dim=-1)
        # attn_output = (attn_weights @ v).transpose(1, 2).contiguous().view(B, N, C)  # [B, seq_len, dim]
        out = self.out_proj(attn_output)
        out = self.dropout(out)

        return out
    
    def attention_pytorch_optimized(self, q, k, v, mask=None, causal=False):
        b, seq_q, head_num, c = q.shape
        seq_k = k.shape[1]
        q = q.permute(0, 2, 1, 3)
        k = k.permute(0, 2, 1, 3)
        v = v.permute(0, 2, 1, 3)
        if mask is not None:
            atten_mask_npu = mask.unsqueeze(1).expand(b, head_num, seq_q, seq_k)
            atten_mask_npu = torch.logical_not(atten_mask_npu)
            attn_output = F.scaled_dot_product_attention(q, k, v, attn_mask=atten_mask_npu)
        else:
            if causal:
                attn_output = F.scaled_dot_product_attention(q, k, v, is_causal=True)
            else:
                attn_output = F.scaled_dot_product_attention(q, k, v)

        return attn_output.permute(0, 2, 1, 3)

# ---- 原 transformer.py ----

class SwiGLU(nn.Module):
    def __init__(self, embed_dim):
        super(SwiGLU, self).__init__()
        self.proj1 = nn.Linear(embed_dim, int(8 / 3 * embed_dim))  # XWG
        self.proj2 = nn.Linear(embed_dim, int(8 / 3 * embed_dim))  # XW1
        self.proj3 = nn.Linear(int(8 / 3 * embed_dim), embed_dim)  # W2

    def forward(self, x):
        # Apply Swish (SiLU in PyTorch) to the first projection
        x_proj1 = F.silu(self.proj1(x))  # Swish(XWG)
        # Apply the second projection
        x_proj2 = self.proj2(x)  # XW1
        # Element-wise multiplication
        x_glu = x_proj1 * x_proj2  # (swish(XWG) ⊙ XW1)
        # Final projection back to the original dimension
        output = self.proj3(x_glu)  # (swish(XWG) ⊙ XW1)W2
        
        return output

class MLA(nn.Module):
    def __init__(self, embed_dim, num_heads, args, cross_embed_dim, causal=True, dropout=0.):
        """
        Differential Transformer Block with optional causal self-attention or cross-attention
        and automatic application of ROPE.

        Args:
            embed_dim (int): Embedding dimension.
            num_heads (int): Number of attention heads.
            args (Namespace): Arguments for attention settings.
            causal (bool): Whether to use causal masking by default for self-attention.
        """
        super(MLA, self).__init__()
        # Multihead differential attention module
        self.attn = MultiheadFlashrope(args, embed_dim, num_heads)
        self.causal = causal

        self.feed_forward = SwiGLU(embed_dim)
        self.norm1 = RMSNorm(embed_dim, eps=1e-5)
        self.norm2 = RMSNorm(embed_dim, eps=1e-5)
        self.attn_dropout = nn.Dropout(dropout)
        self.ff_dropout = nn.Dropout(dropout)
        self.residual_dropout = nn.Dropout(dropout) 
        if cross_embed_dim != -1:
            self.cross_norm = RMSNorm(embed_dim, eps=1e-5)
            self.cross_attn = CrossAttention(embed_dim, cross_embed_dim, num_heads)

    def forward(
        self, x, start, conds=None, attention_mask=None, use_cache=False
    ):
        """
        Args:
            x: Input tensor
            use_cache: Whether to use KV cache
        """
        # Enable/disable KV cache in attention module
        self.attn.kv_cache_enabled = use_cache

        # 第一个残差连接：注意力 + dropout
        residual = x
        x = self.norm1(x)
        attn_out = self.attn(
            x,
            start_pos=start,
            attention_mask=attention_mask,
            causal=self.causal,
        )
        attn_out = self.attn_dropout(attn_out)
        x = residual + attn_out
        if conds is not None:
            x_skip = x
            x = self.cross_norm(x)
            x = self.cross_attn(x, conds) + x_skip 

        # 第二个残差连接：前馈网络 + dropout
        residual = x
        ff_out = self.ff_dropout((self.feed_forward(self.norm2(x))))
        x = residual + ff_out

        return x

    def init_kv_cache(self, batch_size, dtype=torch.bfloat16):
        """Initialize KV cache for this transformer block"""
        return self.attn.empty_kv_cache(
            batch_size=batch_size, dtype=dtype
        )

    def reset_kv_cache(self):
        """Reset KV cache for this transformer block"""
        self.attn.reset_cache()