import torch
import torch.nn as nn
import torch.nn.functional as F
import math
import matplotlib.pyplot as plt

class ALiBiPositionalBias(nn.Module):
    """
    ALiBi (Attention with Linear Biases) 位置编码实现
    """
    def __init__(self, num_heads, max_positions=512):
        super().__init__()
        self.num_heads = num_heads
        self.max_positions = max_positions

        # 为每个头生成斜率 (slope)
        slopes = torch.Tensor(self._get_slopes(num_heads))
        self.register_buffer('slopes', slopes)

        # 预计算偏置矩阵
        self.register_buffer('bias_matrix', self._precompute_bias_matrix())

    def _get_slopes(self, n):
        """
        为 n 个头生成斜率，按照论文中的公式
        """
        def get_slopes_power_of_2(n):
            start = (2**(-2**-(math.log2(n)-3)))
            ratio = start
            return [start*ratio**i for i in range(n)]

        if math.log2(n).is_integer():
            return get_slopes_power_of_2(n)
        else:
            # 最接近的2的幂
            closest_power_of_2 = 2 ** math.floor(math.log2(n))
            return (get_slopes_power_of_2(closest_power_of_2) + 
                    self._get_slopes(2 * closest_power_of_2)[0::2][:n - closest_power_of_2])

    def _precompute_bias_matrix(self):
        """
        预计算偏置矩阵
        """
        # 创建相对位置矩阵
        context_position = torch.arange(self.max_positions).unsqueeze(0)  # [1, seq_len]
        memory_position = torch.arange(self.max_positions).unsqueeze(1)  # [seq_len, 1]

        # 相对位置 = 目标位置 - 源位置
        relative_position = memory_position - context_position  # [seq_len, seq_len]
        relative_position = torch.abs(relative_position)

        # 为每个头创建偏置矩阵 [num_heads, seq_len, seq_len]
        bias_matrix = torch.zeros(self.num_heads, self.max_positions, self.max_positions)

        for h in range(self.num_heads):
            # 每个头使用自己的斜率
            bias_matrix[h] = -relative_position * self.slopes[h]

        return bias_matrix

    def forward(self, seq_len):
        """
        返回对应序列长度的偏置矩阵
        Args:
            seq_len: 序列长度
        Returns:
            bias: [num_heads, seq_len, seq_len] 的偏置矩阵
        """
        if seq_len <= self.max_positions:
            return self.bias_matrix[:, :seq_len, :seq_len]
        else:
            # 如果序列长度超过预计算的最大长度，动态计算
            context_position = torch.arange(seq_len).unsqueeze(0)
            memory_position = torch.arange(seq_len).unsqueeze(1)
            relative_position = memory_position - context_position
            relative_position = torch.abs(relative_position)

            bias_matrix = torch.zeros(self.num_heads, seq_len, seq_len)
            for h in range(self.num_heads):
                bias_matrix[h] = -relative_position * self.slopes[h]

            return bias_matrix.to(self.bias_matrix.device)

class ALiBiAttention(nn.Module):
    """
    集成 ALiBi 的自注意力层
    """
    def __init__(self, d_model, num_heads, max_positions=512):
        super().__init__()
        self.d_model = d_model
        self.num_heads = num_heads
        self.head_dim = d_model // num_heads

        assert d_model % num_heads == 0, "d_model 必须能被 num_heads 整除"

        # 线性变换层
        self.wq = nn.Linear(d_model, d_model)
        self.wk = nn.Linear(d_model, d_model)
        self.wv = nn.Linear(d_model, d_model)
        self.wo = nn.Linear(d_model, d_model)

        # ALiBi 偏置
        self.alibi = ALiBiPositionalBias(num_heads, max_positions)

    def forward(self, x, mask=None):
        """
        Args:
            x: [batch_size, seq_len, d_model]
            mask: 可选的注意力掩码 [batch_size, seq_len, seq_len]
        Returns:
            output: [batch_size, seq_len, d_model]
        """
        batch_size, seq_len, d_model = x.shape

        # 线性变换
        Q = self.wq(x)  # [batch_size, seq_len, d_model]
        K = self.wk(x)  # [batch_size, seq_len, d_model]
        V = self.wv(x)  # [batch_size, seq_len, d_model]

        # 重塑为多头
        Q = Q.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
        K = K.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
        V = V.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)

        # 计算注意力分数
        scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.head_dim)
        # scores: [batch_size, num_heads, seq_len, seq_len]

        # 添加 ALiBi 偏置
        alibi_bias = self.alibi(seq_len)  # [num_heads, seq_len, seq_len]
        scores = scores + alibi_bias.unsqueeze(0)  # 添加batch维度

        # 应用可选的注意力掩码
        if mask is not None:
            scores = scores.masked_fill(mask == 0, -1e9)

        # 计算注意力权重
        attention_weights = F.softmax(scores, dim=-1)

        # 应用注意力权重到值向量
        output = torch.matmul(attention_weights, V)  # [batch_size, num_heads, seq_len, head_dim]

        # 重塑回原始形状
        output = output.transpose(1, 2).contiguous().view(batch_size, seq_len, d_model)

        # 最终线性变换
        output = self.wo(output)

        return output, attention_weights

def demo_alibi_attention():
    """演示 ALiBi 注意力的使用"""
    print("=== ALiBi 位置编码 Demo ===\n")

    # 设置参数
    d_model = 64
    num_heads = 4
    batch_size = 2
    seq_len = 8

    # 创建模型和随机输入
    attention_layer = ALiBiAttention(d_model, num_heads)
    x = torch.randn(batch_size, seq_len, d_model)

    print(f"输入形状: {x.shape}")
    print(f"模型参数: d_model={d_model}, num_heads={num_heads}")

    # 前向传播
    output, attn_weights = attention_layer(x)
    print(f"输出形状: {output.shape}")
    print(f"注意力权重形状: {attn_weights.shape}")

    return attention_layer, x, output, attn_weights

def visualize_alibi_bias():
    """可视化 ALiBi 偏置矩阵"""
    print("\n=== 可视化 ALiBi 偏置矩阵 ===")

    num_heads = 4
    seq_len = 10
    alibi = ALiBiPositionalBias(num_heads)

    # 获取偏置矩阵
    bias_matrix = alibi(seq_len)

    # 打印斜率
    print(f"各头的斜率: {alibi.slopes.tolist()}")

    # 可视化每个头的偏置矩阵
    fig, axes = plt.subplots(1, num_heads, figsize=(16, 4))

    for i in range(num_heads):
        im = axes[i].imshow(bias_matrix[i].detach().numpy(), cmap='coolwarm', aspect='auto')
        axes[i].set_title(f'头 {i}\n斜率: {alibi.slopes[i]:.4f}')
        axes[i].set_xlabel('目标位置')
        axes[i].set_ylabel('源位置')
        plt.colorbar(im, ax=axes[i])

    plt.tight_layout()
    plt.savefig('alibi_bias_matrix.png', dpi=300, bbox_inches='tight')
    print("偏置矩阵可视化已保存为 'alibi_bias_matrix.png'")

    return bias_matrix

def test_extrapolation():
    """测试 ALiBi 的外推能力"""
    print("\n=== 测试外推能力 ===")

    # 训练时的序列长度
    train_seq_len = 8
    # 推理时的序列长度（比训练时长）
    eval_seq_len = 16

    num_heads = 4
    alibi = ALiBiPositionalBias(num_heads, max_positions=eval_seq_len)

    # 获取训练和推理时的偏置矩阵
    train_bias = alibi(train_seq_len)
    eval_bias = alibi(eval_seq_len)

    print(f"训练序列长度: {train_seq_len}")
    print(f"推理序列长度: {eval_seq_len}")
    print(f"偏置矩阵可以无缝扩展到更长的序列")

    # 检查训练部分的偏置是否一致
    train_part_from_eval = eval_bias[:, :train_seq_len, :train_seq_len]
    is_consistent = torch.allclose(train_bias, train_part_from_eval, atol=1e-6)

    print(f"训练部分在扩展后保持一致: {is_consistent}")

    return train_bias, eval_bias

if __name__ == "__main__":
    # 设置随机种子以便复现
    torch.manual_seed(42)

    # 运行演示
    attention_layer, x, output, attn_weights = demo_alibi_attention()

    # 可视化偏置矩阵
    bias_matrix = visualize_alibi_bias()

    # 测试外推能力
    train_bias, eval_bias = test_extrapolation()

    print("\n=== Demo 完成 ===")
    print("这个 Demo 展示了:")
    print("1. ALiBi 位置编码的实现")
    print("2. 如何集成到自注意力机制中")
    print("3. 偏置矩阵的可视化")
    print("4. ALiBi 强大的外推能力")