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 DecoderLayer(nn.Module):
    """单个解码器层 = 3个子层"""
    def __init__(self, size, self_attn, src_attn, feed_forward, dropout):
        """
        size: d_model
        self_attn: 因果掩码多头自注意力
        src_attn: 交叉多头注意力
        feed_forward: 前馈神经网络
        dropout: dropout比例
        """
        super().__init__()
        self.size = size
        self.self_attn = self_attn    # 因果掩码自注意力
        self.src_attn = src_attn      # 交叉注意力
        self.feed_forward = feed_forward
        # 克隆3个SublayerConnection (解码器有3个子层)
        self.sublayer = clones(SublayerConnection(size, dropout), 3)
    
    def forward(self, x, memory, src_mask, tgt_mask):
        """
        x: [batch, tgt_len, d_model] 解码器输入
        memory: [batch, src_len, d_model] 编码器输出
        src_mask: 编码器的填充掩码 (用于交叉注意力)
        tgt_mask: 解码器的因果掩码 (用于自注意力)
        return: [batch, tgt_len, d_model]
        """
        # 子层1: 因果掩码多头自注意力 (Q=K=V=x)
        x = self.sublayer[0](
            x, lambda x: self.self_attn(x, x, x, tgt_mask)
        )
        
        # 子层2: 交叉注意力 (Q=x, K=V=memory)
        x = self.sublayer[1](
            x, lambda x: self.src_attn(x, memory, memory, src_mask)
        )
        
        # 子层3: 前馈神经网络
        x = self.sublayer[2](x, self.feed_forward)
        return x


class Decoder(nn.Module):
    """完整解码器 = N层DecoderLayer堆叠"""
    def __init__(self, layer, N):
        """
        layer: 单个DecoderLayer
        N: 层数 (通常为6)
        """
        super().__init__()
        self.layers = clones(layer, N)
        self.norm = LayerNorm(layer.size)
    
    def forward(self, x, memory, src_mask, tgt_mask):
        """
        x: [batch, tgt_len, d_model] 解码器输入(词嵌入+位置编码)
        memory: [batch, src_len, d_model] 编码器输出
        src_mask: 编码器填充掩码
        tgt_mask: 解码器因果掩码
        return: [batch, tgt_len, d_model] 解码器输出
        """
        for layer in self.layers:
            x = layer(x, memory, src_mask, tgt_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
    src_len = 10  # 源语言长度
    tgt_len = 8   # 目标语言长度
    
    # 创建组件
    self_attn = MultiHeadedAttention(h, d_model, dropout)  # 因果自注意力
    src_attn = MultiHeadedAttention(h, d_model, dropout)   # 交叉注意力
    ff = PositionwiseFeedForward(d_model, d_ff, dropout)
    
    # 创建单个解码器层
    decoder_layer = DecoderLayer(d_model, self_attn, src_attn, ff, dropout)
    
    # 创建完整解码器 (6层)
    decoder = Decoder(decoder_layer, N)
    
    # 前向传播
    x = torch.randn(batch_size, tgt_len, d_model)        # 解码器输入
    memory = torch.randn(batch_size, src_len, d_model)   # 编码器输出
    src_mask = None  # 编码器填充掩码
    tgt_mask = None  # 解码器因果掩码
    
    output = decoder(x, memory, src_mask, tgt_mask)
    print(f"解码器输入: {x.shape}")
    print(f"编码器输出: {memory.shape}")
    print(f"解码器输出: {output.shape}")  # [2, 8, 512]