"""
Transformer Encoder 模块 — Person B 负责实现

包含:
- TransformerEncoderLayer: 单层编码器
- TransformerEncoder: 多层编码器堆叠

架构 (Pre-LayerNorm):
    x → LayerNorm → MultiHeadAttention → Residual → LayerNorm → FFN → Residual
"""

from __future__ import annotations

import torch
import torch.nn as nn
from typing import Optional


class TransformerEncoderLayer(nn.Module):
    """
    单层 Transformer Encoder。

    TODO [Person B]: 实现以下组件:
    1. Self-Attention: MultiHeadAttention (支持 Flash Attention)
    2. Feed-Forward Network: Linear → Activation → Dropout → Linear
    3. 两个 LayerNorm
    4. Residual connections
    5. Dropout

    注意: 使用 Pre-LayerNorm 架构 (先 norm 再 attention/ffn)
    """

    def __init__(
        self,
        d_model: int = 512,
        nhead: int = 8,
        dim_feedforward: int = 2048,
        dropout: float = 0.1,
        activation: str = "gelu",
        use_flash_attention: bool = True,
        use_rotary_embedding: bool = True,
        pre_norm: bool = True,
    ):
        super().__init__()
        raise NotImplementedError("TODO: Person B 实现 TransformerEncoderLayer.__init__")

    def forward(
        self,
        src: torch.Tensor,                             # [B, S, D]
        src_key_padding_mask: Optional[torch.BoolTensor] = None,  # [B, S]
    ) -> torch.Tensor:
        """
        TODO [Person B]: Pre-LayerNorm 前向传播:
        1. residual = src
        2. src = layer_norm_1(src)
        3. src = self_attention(src, src, src, key_padding_mask=src_key_padding_mask)
        4. src = residual + dropout(src)
        5. residual = src
        6. src = layer_norm_2(src)
        7. src = ffn(src)
        8. src = residual + dropout(src)
        """
        raise NotImplementedError("TODO: Person B 实现 TransformerEncoderLayer.forward")


class TransformerEncoder(nn.Module):
    """
    多层 Transformer Encoder。

    TODO [Person B]:
    1. 堆叠 N 个 TransformerEncoderLayer
    2. 最终加一个 LayerNorm (Pre-Norm 架构需要)
    """

    def __init__(self, encoder_layer: TransformerEncoderLayer, num_layers: int):
        super().__init__()
        raise NotImplementedError("TODO: Person B 实现 TransformerEncoder.__init__")

    def forward(
        self,
        src: torch.Tensor,
        src_key_padding_mask: Optional[torch.BoolTensor] = None,
    ) -> torch.Tensor:
        raise NotImplementedError("TODO: Person B 实现 TransformerEncoder.forward")