import torch
import torch.nn as nn
import copy


def clones(module, N):
    """克隆模块N次"""
    return nn.ModuleList([copy.deepcopy(module) for _ in range(N)])


class LayerNorm(nn.Module):
    """层归一化"""
    def __init__(self, features, eps=1e-6):
        super().__init__()
        self.a_2 = nn.Parameter(torch.ones(features))
        self.b_2 = nn.Parameter(torch.zeros(features))
        self.eps = eps
    
    def forward(self, x):
        mean = x.mean(-1, keepdim=True)
        std = x.std(-1, keepdim=True)
        return self.a_2 * (x - mean) / (std + self.eps) + self.b_2


class SublayerConnection(nn.Module):
    """子层连接: LayerNorm + sublayer + Dropout + 残差"""
    def __init__(self, size, dropout):
        super().__init__()
        self.norm = LayerNorm(size)
        self.dropout = nn.Dropout(dropout)
    
    def forward(self, x, sublayer):
        return x + self.dropout(sublayer(self.norm(x)))


class EncoderLayer(nn.Module):
    """单个编码器层 = 2个子层"""
    def __init__(self, size, self_attn, feed_forward, dropout):
        """
        size: d_model
        self_attn: 多头自注意力模块
        feed_forward: 前馈神经网络模块
        dropout: dropout比例
        """
        super().__init__()
        self.self_attn = self_attn
        self.feed_forward = feed_forward
        # 克隆2个SublayerConnection (编码器有2个子层)
        self.sublayer = clones(SublayerConnection(size, dropout), 2)
        self.size = size
    
    def forward(self, x, mask):
        """
        x: [batch, seq_len, d_model]
        mask: 填充掩码
        return: [batch, seq_len, d_model]
        """
        # 子层1: 多头自注意力 (Q=K=V=x)
        x = self.sublayer[0](
            x, lambda x: self.self_attn(x, x, x, mask)
        )
        # 子层2: 前馈神经网络
        x = self.sublayer[1](x, self.feed_forward)
        return x


class Encoder(nn.Module):
    """完整编码器 = N层EncoderLayer堆叠"""
    def __init__(self, layer, N):
        """
        layer: 单个EncoderLayer
        N: 层数 (通常为6)
        """
        super().__init__()
        # 克隆N个EncoderLayer
        self.layers = clones(layer, N)
        self.norm = LayerNorm(layer.size)  # 最后的层归一化
    
    def forward(self, x, mask):
        """
        x: [batch, seq_len, d_model] (词嵌入+位置编码)
        mask: 填充掩码
        return: [batch, seq_len, d_model] 编码器输出
        """
        # 逐层传递
        for layer in self.layers:
            x = layer(x, mask)
        # 最后的层归一化
        return self.norm(x)


# 使用示例
if __name__ == "__main__":
    from transformer_components import MultiHeadedAttention, PositionwiseFeedForward
    
    d_model = 512
    h = 8
    d_ff = 2048
    N = 6
    dropout = 0.1
    batch_size = 2
    seq_len = 10
    
    # 创建组件
    attn = MultiHeadedAttention(h, d_model, dropout)
    ff = PositionwiseFeedForward(d_model, d_ff, dropout)
    
    # 创建单个编码器层
    encoder_layer = EncoderLayer(d_model, attn, ff, dropout)
    
    # 创建完整编码器 (6层)
    encoder = Encoder(encoder_layer, N)
    
    # 前向传播
    x = torch.randn(batch_size, seq_len, d_model)  # 词嵌入+位置编码
    mask = None  # 或填充掩码
    output = encoder(x, mask)
    print(f"输入: {x.shape}")
    print(f"输出: {output.shape}")  # [2, 10, 512]