import numpy as np
import matplotlib.pyplot as plt

class PositionalEmbedding:
    def __init__(self, d_model, max_seq_len):
        """
        初始化位置嵌入
        
        参数:
            d_model: 嵌入维度
            max_seq_len: 最大序列长度
        """
        self.d_model = d_model
        self.max_seq_len = max_seq_len
        self.pos_embedding = self._create_positional_embedding()
        
    def _create_positional_embedding(self):
        """创建正弦余弦位置嵌入"""
        # 初始化位置嵌入矩阵
        pos_embedding = np.zeros((self.max_seq_len, self.d_model))
        
        # 遍历每个位置
        for pos in range(self.max_seq_len):
            # 遍历每个维度
            for i in range(self.d_model // 2):
                # 计算正弦和余弦值
                angle = pos / np.power(10000, 2 * i / self.d_model)
                pos_embedding[pos, 2*i] = np.sin(angle)
                pos_embedding[pos, 2*i + 1] = np.cos(angle)
                
        return pos_embedding
    
    def get_embedding(self, pos):
        """获取指定位置的嵌入向量"""
        if pos < 0 or pos >= self.max_seq_len:
            raise ValueError(f"位置必须在[0, {self.max_seq_len-1}]范围内")
        return self.pos_embedding[pos]
    
    def visualize_embedding(self, num_positions=10, figsize=(12, 8)):
        """可视化位置嵌入"""
        plt.figure(figsize=figsize)
        # 只显示前num_positions个位置和所有维度
        plt.imshow(self.pos_embedding[:num_positions, :], cmap='viridis')
        plt.xlabel('嵌入维度')
        plt.ylabel('位置')
        plt.title('Transformer位置嵌入可视化')
        plt.colorbar()
        plt.show()
        
    def visualize_dimensions(self, dim1=0, dim2=1, num_positions=50, figsize=(10, 10)):
        """可视化不同位置在两个维度上的分布"""
        plt.figure(figsize=figsize)
        x = self.pos_embedding[:num_positions, dim1]
        y = self.pos_embedding[:num_positions, dim2]
        plt.scatter(x, y)
        
        # 为每个点添加位置标签
        for i in range(num_positions):
            plt.annotate(str(i), (x[i], y[i]))
            
        plt.xlabel(f'维度 {dim1}')
        plt.ylabel(f'维度 {dim2}')
        plt.title(f'位置在维度 {dim1} 和 {dim2} 上的分布')
        plt.grid(True)
        plt.show()

# 示例用法
if __name__ == "__main__":
    # 创建位置嵌入实例，模型维度为512，最大序列长度为100
    pos_embed = PositionalEmbedding(d_model=512, max_seq_len=100)
    
    # 获取位置10的嵌入向量
    position_10_embedding = pos_embed.get_embedding(10)
    print(f"位置10的嵌入向量形状: {position_10_embedding.shape}")
    
    # 可视化位置嵌入
    pos_embed.visualize_embedding(num_positions=20)
    
    # 可视化不同位置在两个维度上的分布
    pos_embed.visualize_dimensions(dim1=0, dim2=1)
    pos_embed.visualize_dimensions(dim1=2, dim2=3)