>>> chr(0)
'\x00'

>>> chr(0)
'\x00'
>>> print(chr(0))

>>> "this is a test" + chr(0) + "string"
'this is a test\x00string'
>>> print("this is a test" + chr(0) + "string")
this is a teststring

UnicodeDecodeError: 'utf-8' codec can't decode byte 0xe4 in position 0: unexpected end of data

vocab, merges = train_bpe(input_path, vocab_size, special_tokens)

from typing import Dict, Tuple, List, Iterable, Iterator
from collections import Counter
import regex as re

from queue import Empty
from multiprocessing import Process, Queue, Manager


from tqdm import trange, tqdm
import pickle

import os 


from cs336_basics.tokenizer.pretokenization_regular_pattern import PAT
from cs336_basics.tokenizer.utils import find_chunk_boundaries 


def initialize_vocab(special_tokens: List[bytes]) -> Dict[int, bytes]:
    vocab = {i: bytes([i]) for i in range(256)}  # ASCII characters
    for i, token in enumerate(special_tokens, start=256):
        vocab[i] = token

    return vocab

def word_to_bytes(word: str) -> List[bytes]:
    """
    Convert a word to bytes.
    """
    byte_ids = [bytes([b]) for b in word.encode("utf-8")]

    return byte_ids


def split_by_special_tokens(
    text: str, special_tokens: list[str]
) -> List[str]:
    special_tokens_sorted = sorted(special_tokens, key=len, reverse=True)
    if not special_tokens_sorted:
        return [text]
    pattern = "|".join(re.escape(t) for t in special_tokens_sorted)
    special_chunks = re.split(f"({pattern})", text)

    return special_chunks



def pre_tokenize_string(text: str, special_tokens: List[str], include_special: bool = False) -> Counter:
    word_counter = Counter()
    special_chunks = split_by_special_tokens(text, special_tokens)

    for chunk in special_chunks:
        if chunk in special_tokens:
            if include_special:
                token = tuple(word_to_bytes(chunk))
                word_counter[token] += 1
        else:
            for match in re.finditer(PAT, chunk):
                word = match.group(0)
                token = tuple(word_to_bytes(word))
                word_counter[token] += 1

    return word_counter


# TODO: Implement the worker for this.
def pre_tokenize_string_worker(
    input_path: str | os.PathLike, special_tokens: list[str], queue: Queue, start: int, end: int, include_special: bool = False,
):
    """
    Pre-tokenize a string into bytes.
    """
    with open(input_path, "rb") as f:
        f.seek(start)
        chunk = f.read(end - start).decode("utf-8", errors="ignore")
        
    word_counter = pre_tokenize_string(chunk, special_tokens, include_special)
                
    # Put the result in the queue
    queue.put(word_counter)



def pair_counts(
    word_counter: Dict[Tuple[bytes], int],
) -> Dict[Tuple[bytes, bytes], int]:
    """
    Count pairs of bytes in the word counter.
    """
    pairs: Dict[Tuple[bytes, bytes], int] = {}
    for token, freq in word_counter.items():
        for i in range(len(token) - 1):
            pair = (token[i], token[i + 1])
            pairs[pair] = pairs.get(pair, 0) + freq

    return pairs


def get_most_frequent_pair(
    pairs: Dict[Tuple[bytes, bytes], int],
) -> Tuple[bytes, bytes]:
    max_freq = max(pairs.values())
    candidates = [pair for pair, freq in pairs.items() if freq == max_freq]
    res = max(candidates)

    return res


def add_pair_to_vocab(
    vocab: Dict[int, bytes], pair: Tuple[bytes, bytes], vocab_inv: Dict[bytes, int]
) -> int:
    """
    Add a new pair to the vocabulary.
    """
    index = len(vocab)
    s = vocab[vocab_inv[pair[0]]] + vocab[vocab_inv[pair[1]]]
    vocab[index] = s
    vocab_inv[vocab[index]] = index

    return index

from collections import Counter, defaultdict


def merge_pair(
    word_counter: Dict[Tuple[bytes], int], pair: Tuple[bytes, bytes]
) -> Tuple[Dict[Tuple[bytes], int], Dict]:
    """
    Merge a pair of bytes in the word counter.
    """
    new_word_counter = Counter()
    updated_pair_counts = defaultdict(int)

    for token, freq in word_counter.items():
        new_token = []
        i = 0
        while i < len(token):
            if i < len(token) - 1 and (token[i], token[i + 1]) == pair:
                new_token.append(token[i] + token[i + 1])
                i += 2
            else:
                new_token.append(token[i])
                i += 1

        new_word_counter[tuple(new_token)] += freq

        for j in range(len(new_token) - 1):
            new_pair = (new_token[j], new_token[j + 1])
            updated_pair_counts[new_pair] += freq

    return new_word_counter, updated_pair_counts


def check_and_convert_special_tokens(
    special_tokens: List[str] | List[bytes],
) -> List[bytes]:
    """
    Check if special tokens are in the vocabulary and convert them to bytes.
    """
    if not all(isinstance(token, bytes) for token in special_tokens):
        special_tokens_bytes = [
            token.encode("utf-8") for token in special_tokens if isinstance(token, str)
        ]

    return special_tokens_bytes


def train_bpe(
    input_path: str | os.PathLike ,
    vocab_size=10_000,
    special_tokens: List[str] = [],
    **kwargs,
):
    special_tokens_bytes = check_and_convert_special_tokens(special_tokens)

    vocab = initialize_vocab(special_tokens_bytes)
    vocab_inv = {v: k for k, v in vocab.items()}
    merges: List[Tuple[bytes, bytes]] = []
    
    
    # Pre-tokenization
    with open(input_path, "rb") as f:
        chunk_boundaries = find_chunk_boundaries(
            f, kwargs.get("num_processes", 8), special_tokens_bytes[0]
        )

    
    manager = Manager()
    queue = manager.Queue()
    processes = []
    
    for start, end in zip(chunk_boundaries[:-1], chunk_boundaries[1:]):
        p = Process(
            target=pre_tokenize_string_worker,
            args=(input_path, special_tokens, queue, start, end, False),
        )
        processes.append(p)
        p.start()

    for p in processes:
        p.join()
        
    word_counter = Counter()
    for _ in range(len(processes)):
        try:
            word_counter += queue.get(timeout=10)  # Wait up to 10 seconds for results
        except Empty:
            print("⚠️ Warning: A subprocess did not return a result!")

    # End Pre-tokenization

    pairs_freqs = pair_counts(word_counter)
    
    num_merges = vocab_size - len(vocab)
    for _ in trange(num_merges):

        most_common_pair = get_most_frequent_pair(pairs_freqs)

        new_index = add_pair_to_vocab(vocab, most_common_pair, vocab_inv)
        merges.append(most_common_pair)

        word_counter, pairs_freqs = merge_pair(word_counter, most_common_pair)

    return vocab, merges



class Tokenizer:
    def __init__(
        self, 
        vocab: Dict[int, bytes], 
        merges: List[Tuple[bytes, bytes]],
        special_tokens: List[str] | None = []
    ):
        self.vocab = vocab
        self.merges = merges
        
        # self.register_special_tokens(special_tokens)
        self.vocab_inv = {v: k for k, v in self.vocab.items()}
        
        if special_tokens is None:
            self.special_tokens = {}
            self.bytes_special_tokens = []
        else:
            self.special_tokens = {token: i for i, token in enumerate(special_tokens, start=len(self.vocab))}
            self.bytes_special_tokens = [token.encode("utf-8") for token in special_tokens if isinstance(token, str)]

        
        
    
    def register_special_tokens(self, special_tokens):
        if special_tokens is None:
            self.special_tokens = {}
            self.bytes_special_tokens = []
            return
        
        if not all(isinstance(token, bytes) for token in special_tokens):
            bytes_special_tokens = [token.encode("utf-8") for token in special_tokens if isinstance(token, str)]
            
        for i, token in enumerate(bytes_special_tokens, start=len(self.vocab)):
            # Add special tokens to the vocabulary
            self.vocab[i] = token
            
        # self.bytes_special_tokens = bytes_special_tokens
        # self.special_tokens = {token: i for i, token in enumerate(special_tokens, start=len(self.vocab))}
        
    def _pre_tokenize(self, text) -> List[bytes]:
        """
        Pre-tokenize the input text into bytes.
        """
        parts = split_by_special_tokens(text, list(self.special_tokens.keys()))
        token_list = []
        
        for part in parts:
            if part in self.special_tokens.keys():
                token_list.append(part.encode("utf-8"))
            else:
                tokens = re.findall(PAT, part)
                token_list.extend(word_to_bytes(token) for token in tokens)

        return token_list
    

    
    def encode(self, text: str) -> List[int]:
        byte_tokens = self._pre_tokenize(text)
        

        # Convert byte tokens to indices
        token_ids = []
        for byte_token in byte_tokens:
            # print(f"Processing byte token: {byte_token}")
            if byte_token in self.bytes_special_tokens:
                token_ids.append([self.vocab_inv[byte_token]])
            else:
                token_ids.append([self.vocab_inv[b] for b in byte_token]) #type: ignore

        for i, pretoken in enumerate(token_ids):
            for merge in self.merges:
                new_index = self.vocab_inv.get(merge[0] + merge[1], None)
                if new_index is None:
                    continue

                merged = []
                j = 0
                while j < len(pretoken):
                    if (
                        j < len(pretoken) - 1
                        and (self.vocab[pretoken[j]], self.vocab[pretoken[j + 1]]) == merge
                    ):
                        merged.append(new_index)
                        j += 2
                    else:
                        merged.append(pretoken[j])
                        j += 1
                        
                pretoken = merged
            token_ids[i] = pretoken[:]

        return [i for pre in token_ids for i in pre]
        

    

    def encode_iterable(self, iterable: Iterable[str], batch_size: int = 1024) -> Iterator[int]:
        """
        Encode lines of text from an iterable using buffered batching.
        This version preserves newlines by assuming the input was split with `splitlines(keepends=True)`.
        """
        batch = []
        for line in tqdm(iterable):
            if not line:
                continue
            batch.append(line)
            if len(batch) >= batch_size:
                for encoded in map(self.encode, batch):
                    yield from encoded
                batch.clear()
                
        if batch:
            for encoded in map(self.encode, batch):
                yield from encoded
    
    def decode(self, ids: list[int]) -> str:
        # https://en.wikipedia.org/wiki/Specials_(Unicode_block)#Replacement_character
        
        tokens = b"".join(self.vocab.get(i, b"\xef\xbf\xbd") for i in ids)
        return tokens.decode("utf-8", errors="replace")
    
    @classmethod
    def from_files(
        cls, vocab_path: str, merges_path: str, special_tokens: list[str] | None = None
    ):
        with open(vocab_path, 'rb') as vf:
            raw_vocab = pickle.load(vf)

        vocab = {int(k): (v.encode("utf-8") if isinstance(v, str) else v)
                for k, v in raw_vocab.items()}

        with open(merges_path, 'rb') as mf:
            raw_merges = pickle.load(mf)

        merges = []
        for a, b in raw_merges:
            merges.append((
                a.encode("utf-8") if isinstance(a, str) else a,
                b.encode("utf-8") if isinstance(b, str) else b
            ))
        return cls(vocab, merges, special_tokens)

string = """
low low low low low <|endoftext|>
lower lower widest widest widest <|endoftext|>
newest newest newest newest newest newest 
"""
special_tokens = ["<|endoftext|>"]
PAT = r"\S+"

vocab, merge = train_bpe(vocab_size=269)

[(b's', b't'),
 (b'e', b'st'),
 (b'o', b'w'),
 (b'l', b'ow'),
 (b'w', b'est'),
 (b'n', b'e'),
 (b'ne', b'west'),
 (b'w', b'i'),
 (b'wi', b'd'),
 (b'wid', b'est'),
 (b'low', b'e'),
 (b'lowe', b'r')]

class Tokenizer:
    def __init__(
        self, 
        vocab: Dict[int, bytes],  # 词汇表字典，将整数ID映射到字节token
        merges: List[Tuple[bytes, bytes]],  # BPE合并规则列表，每个元组表示要合并的两个字节序列
        special_tokens: List[str] | None = []  # 特殊token列表（如[CLS]、[SEP]等）
    ):
        self.vocab = vocab  # 存储词汇表
        self.merges = merges  # 存储合并规则
        
        # 创建反向词汇表，用于从字节token查找对应的整数ID
        self.vocab_inv = {v: k for k, v in self.vocab.items()}
        
        # 处理特殊token
        if special_tokens is None:
            self.special_tokens = {}  # 特殊token字典（token到ID的映射）
            self.bytes_special_tokens = []  # 特殊token的字节表示列表
        else:
            # 为特殊token分配ID（从当前词汇表长度开始）
            self.special_tokens = {token: i for i, token in enumerate(special_tokens, start=len(self.vocab))}
            # 将特殊token转换为字节表示
            self.bytes_special_tokens = [token.encode("utf-8") for token in special_tokens if isinstance(token, str)]

    def register_special_tokens(self, special_tokens):
        """注册特殊token到词汇表中"""
        if special_tokens is None:
            self.special_tokens = {}
            self.bytes_special_tokens = []
            return
        
        # 将字符串类型的特殊token转换为字节表示
        if not all(isinstance(token, bytes) for token in special_tokens):
            bytes_special_tokens = [token.encode("utf-8") for token in special_tokens if isinstance(token, str)]
            
        # 将特殊token添加到词汇表中
        for i, token in enumerate(bytes_special_tokens, start=len(self.vocab)):
            self.vocab[i] = token  # 将特殊token添加到词汇表
            
    def _pre_tokenize(self, text) -> List[bytes]:
        """
        预分词函数：将输入文本分割成字节token
        
        参数:
            text: 输入文本字符串
            
        返回:
            字节token列表
        """
        # 首先按特殊token分割文本
        parts = split_by_special_tokens(text, list(self.special_tokens.keys()))
        token_list = []
        
        # 处理每个分割部分
        for part in parts:
            if part in self.special_tokens.keys():
                # 如果是特殊token，直接编码为字节
                token_list.append(part.encode("utf-8"))
            else:
                # 使用正则表达式分割普通文本
                tokens = re.findall(PAT, part)
                # 将每个单词转换为字节表示
                token_list.extend(word_to_bytes(token) for token in tokens)

        return token_list
    
    def encode(self, text: str) -> List[int]:
        """将文本编码为整数ID序列"""
        # 预分词获取字节token
        byte_tokens = self._pre_tokenize(text)
        
        # 将字节token转换为初始ID序列
        token_ids = []
        for byte_token in byte_tokens:
            if byte_token in self.bytes_special_tokens:
                # 处理特殊token
                token_ids.append([self.vocab_inv[byte_token]])
            else:
                # 将普通token拆分为单个字节并查找对应ID
                token_ids.append([self.vocab_inv[b] for b in byte_token]) #type: ignore

        # 应用BPE合并规则
        for i, pretoken in enumerate(token_ids):
            for merge in self.merges:
                # 检查合并后的token是否在词汇表中
                new_index = self.vocab_inv.get(merge[0] + merge[1], None)
                if new_index is None:
                    continue

                # 尝试应用合并规则
                merged = []
                j = 0
                while j < len(pretoken):
                    if (
                        j < len(pretoken) - 1
                        and (self.vocab[pretoken[j]], self.vocab[pretoken[j + 1]]) == merge
                    ):
                        # 找到可合并的相邻token，应用合并
                        merged.append(new_index)
                        j += 2  # 跳过已合并的token
                    else:
                        # 无法合并，保留原token
                        merged.append(pretoken[j])
                        j += 1
                        
                pretoken = merged  # 更新当前预分词序列
            token_ids[i] = pretoken[:]  # 保存处理后的序列

        # 展平所有预分词序列并返回
        return [i for pre in token_ids for i in pre]
    
    def encode_iterable(self, iterable: Iterable[str], batch_size: int = 1024) -> Iterator[int]:
        """
        批量编码可迭代对象中的文本行
        
        参数:
            iterable: 可迭代的文本行
            batch_size: 批处理大小
            
        返回:
            整数ID的迭代器
        """
        batch = []
        for line in tqdm(iterable):
            if not line:
                continue
            batch.append(line)
            if len(batch) >= batch_size:
                # 处理完整批次
                for encoded in map(self.encode, batch):
                    yield from encoded  # 逐个生成ID
                batch.clear()  # 清空批次
                
        # 处理剩余的不完整批次
        if batch:
            for encoded in map(self.encode, batch):
                yield from encoded
    
    def decode(self, ids: list[int]) -> str:
        """
        将整数ID序列解码回文本
        
        参数:
            ids: 整数ID列表
            
        返回:
            解码后的文本字符串
        """
        # 将ID转换回字节token，使用替换字符处理未知ID
        tokens = b"".join(self.vocab.get(i, b"\xef\xbf\xbd") for i in ids)
        # 将字节解码为UTF-8字符串
        return tokens.decode("utf-8", errors="replace")
    
    @classmethod
    def from_files(
        cls, vocab_path: str, merges_path: str, special_tokens: list[str] | None = None
    ):
        """从文件加载词汇表和合并规则创建Tokenizer实例"""
        # 加载词汇表文件
        with open(vocab_path, 'rb') as vf:
            raw_vocab = pickle.load(vf)

        # 处理词汇表数据，确保值为字节类型
        vocab = {int(k): (v.encode("utf-8") if isinstance(v, str) else v)
                for k, v in raw_vocab.items()}

        # 加载合并规则文件
        with open(merges_path, 'rb') as mf:
            raw_merges = pickle.load(mf)

        # 处理合并规则数据，确保为字节类型
        merges = []
        for a, b in raw_merges:
            merges.append((
                a.encode("utf-8") if isinstance(a, str) else a,
                b.encode("utf-8") if isinstance(b, str) else b
            ))
            
        # 创建并返回Tokenizer实例
        return cls(vocab, merges, special_tokens)

class Linear(nn.Module):
    def __init__(
        self, in_features: int, out_features: int, device= None, dtype = None 
    ):
        super().__init__()
        
        self.in_features = in_features
        self.out_features = out_features
        
        self.weight = nn.Parameter(torch.empty(in_features, out_features, device=device, dtype=dtype))
        self._init_weight()

    def _init_weight(self):
        std = (2 / (self.in_features + self.out_features)) ** 0.5
        nn.init.trunc_normal_(self.weight, mean=0.0, std=std, a=-3 * std, b=3 * std)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return x @ self.weight

import torch
import torch.nn as nn

class Embedding(nn.Module):
    def __init__(self, num_embedding: int, embedding_dim: int, device = None, dtype = None):
        """
        嵌入层初始化
        
        参数:
            num_embedding: int - 词表大小，即需要嵌入的不同token数量
            embedding_dim: int - 嵌入向量的维度
            device: torch.device - 计算设备 (CPU/GPU)
            dtype: torch.dtype - 数据类型
        """
        super().__init__()  # 调用父类nn.Module的初始化方法
        
        # 保存参数
        self.num_embedding = num_embedding  # 词表大小
        self.embedding_dim = embedding_dim  # 嵌入维度
        
        # 定义可学习的权重参数，形状为 (词表大小, 嵌入维度)
        self.weight = nn.Parameter(
            torch.empty(num_embedding, embedding_dim, device=device, dtype=dtype)
        )
        # 初始化权重
        self._init_weight()
    
    def _init_weight(self):
        """使用截断正态分布初始化权重"""
        # trunc_normal_: 截断正态分布初始化，将值限制在[a, b]范围内
        # mean=0.0, std=1.0: 均值为0，标准差为1
        # a=-3.0, b=3.0: 将值限制在[-3, 3]标准差范围内，避免极端值
        nn.init.trunc_normal_(self.weight, mean=0.0, std=1.0, a=-3.0, b=3.0)

    def forward(self, token_ids: torch.Tensor) -> torch.Tensor:
        """
        前向传播过程
        
        参数:
            token_ids: torch.Tensor - 输入token索引，形状为 (B, S)
                B: batch size (批大小)
                S: sequence length (序列长度)
                
        返回:
            torch.Tensor - 嵌入向量，形状为 (B, S, D)
                D: embedding dimension (嵌入维度)
        """
        # 使用索引查找：根据token_ids中的索引从weight中提取对应的嵌入向量
        # token_ids中的每个整数索引对应weight矩阵中的一行
        # 结果形状从 (B, S) 变为 (B, S, D)
        return self.weight[token_ids]

class RMSNorm(nn.Module):
    def __init__(self, d_model: int, eps: float = 1e-5, device=None, dtype=None):
        super().__init__()

        self.d_model = d_model
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(d_model, device=device, dtype=dtype))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = x.float()  # Ensure x is float for numerical stability
        norm = torch.sqrt(torch.mean(x**2, dim=-1, keepdim=True) + self.eps)
        normalized_x = x / norm

        return normalized_x * self.weight

def SiLU(x: torch.Tensor) -> torch.Tensor:
    """
    Swish activation function.
    """
    return x * torch.sigmoid(x)


class FFN(nn.Module):
    def __init__(self, d_model: int, d_ff: int, device=None, dtype=None):
        super().__init__()
        self.d_model = d_model
        self.d_ff = d_ff

        self.w1 = Linear(d_model, d_ff)
        self.w2 = Linear(d_ff, d_model)
        self.w3 = Linear(d_model, d_ff)

    def forward(self, x: torch.Tensor) -> torch.Tensor:

        output = self.w2(SiLU(self.w1(x)) * self.w3(x))
        return output

import torch 
import torch.nn as nn 

# einops 是一个用于张量操作的库，提供更清晰、更易读的张量重排语法
import einops

class RotaryPositionalEmbedding(nn.Module):
    """
    旋转位置编码（Rotary Position Embedding, RoPE）模块。
    与传统的绝对位置编码不同，RoPE 通过将位置信息编码为旋转矩阵，
    使得注意力机制能够感知 token 之间的相对位置关系。
    """
    
    def __init__(self, theta: float, d_k: float, max_seq_len: int, device=None):
        """
        初始化 RoPE 模块。
        
        参数:
        - theta: 基础频率缩放因子（通常为 10000.0），控制频率衰减速度。
        - d_k: 注意力头的维度（即每个 token 的嵌入维度）。
        - max_seq_len: 支持的最大序列长度（用于预计算，但本实现中未显式使用）。
        - device: 计算设备（如 'cuda' 或 'cpu'）。
        """
        super().__init__()
        
        self.theta = theta
        self.d_k = d_k
        self.max_seq_len = max_seq_len
        
        # 设置设备，默认为 CPU
        self.device = device if device else torch.device("cpu")
        
        # 计算逆频率（inverse frequencies）：
        # 对于维度 0, 2, 4, ..., d_k-2，计算 1 / (theta^(i / d_k))
        # 这些频率用于构建旋转角度
        # torch.arange(0, d_k, 2) 生成 [0, 2, 4, ..., d_k-2]（假设 d_k 为偶数）
        inv_freq = 1.0 / (self.theta ** (torch.arange(0, d_k, 2).float() / d_k))
        
        # 使用 register_buffer 注册为非参数缓冲区（不会被优化器更新，但会随模型移动设备）
        self.register_buffer("inv_freq", inv_freq)

    def _rotate_half(self, x: torch.Tensor) -> torch.Tensor:
        """
        对输入张量进行“旋转一半”操作：
        将最后一维分成两半，交换并取负，实现 90 度旋转。
        例如：[x1, x2, x3, x4] -> [-x2, x1, -x4, x3]
        这是 RoPE 的核心操作，用于实现复数乘法的实部形式。
        
        参数:
        - x: 形状为 [..., d_k] 的张量
        
        返回:
        - 旋转后的张量，形状与输入相同
        """
        # 将最后一维按每2个元素分组：(..., d_k) -> (..., d_k//2, 2)
        x = einops.rearrange(x, "... (d r) -> ... d r", r=2)
        
        # 拆分为两部分：x1 是偶数索引部分，x2 是奇数索引部分
        x1, x2 = x.unbind(dim=-1)  # x1: [..., d_k//2], x2: [..., d_k//2]
        
        # 构造旋转后的形式：(-x2, x1)
        x = torch.stack((-x2, x1), dim=-1)  # [..., d_k//2, 2]
        
        # 重新合并为 [..., d_k]
        return einops.rearrange(x, "... d r -> ... (d r)")

    def forward(self, x: torch.Tensor, token_positions: torch.Tensor) -> torch.Tensor:
        """
        应用旋转位置编码到输入张量。
        
        RoPE 的数学本质是：将 token 嵌入视为复数，然后乘以 e^{i * m * θ}（m 是位置，θ 是频率），
        等价于：x * cos(mθ) + rotate(x) * sin(mθ)
        
        参数:
        - x: 输入张量，形状 [..., seq_len, d_k]
        - token_positions: 位置索引张量，形状 [..., seq_len]，
          通常为 [0, 1, 2, ..., seq_len-1]，支持动态位置（如滑动窗口）
        
        返回:
        - 应用 RoPE 后的张量，形状与 x 相同
        """
        seq_len = x.size(-2)  # 获取序列长度
        
        # 如果未提供 token_positions，则自动生成从 0 到 seq_len-1 的位置索引
        if token_positions is None:
            token_positions = torch.arange(seq_len, device=x.device)
            # 扩展为与 batch 维度对齐：(seq_len,) -> (B, seq_len)
            token_positions = token_positions.unsqueeze(0).expand(x.size(0), -1)

        # 计算每个位置和每个频率维度的旋转角度 θ = m * inv_freq
        # token_positions: [..., seq_len]
        # inv_freq: [d_k // 2]
        # 结果 theta: [..., seq_len, d_k // 2]
        theta = torch.einsum("... n, d -> ... n d", token_positions, self.inv_freq)
        
        # 计算 cos(θ) 和 sin(θ)，并重复每个值两次以匹配 d_k 维度
        # 例如：[cos0, cos1] -> [cos0, cos0, cos1, cos1]
        cos = theta.cos().repeat_interleave(2, dim=-1)  # [..., seq_len, d_k]
        sin = theta.sin().repeat_interleave(2, dim=-1)  # [..., seq_len, d_k]
        
        # 应用 RoPE 公式：x_rot = x * cos + rotate_half(x) * sin
        # 这等价于复数乘法：(x_real + i*x_imag) * (cos + i*sin)
        x = x * cos + self._rotate_half(x) * sin

        return x

def softmax(x: torch.Tensor, dim: int = -1) -> torch.Tensor:
    x = x - torch.max(x, dim=dim, keepdim=True).values
    x = torch.exp(x)
    x = x / torch.sum(x, dim=dim, keepdim=True)
    return x

def scaled_dot_product_attention(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, mask: torch.Tensor | None = None) -> torch.Tensor:
    """
    q: (B, S_q, D)
    k: (B, S_k, D)
    v: (B, S_v, D)
    mask: (B, S_q, S_k) or None
    """
    
    d_k = k.size(-1)
    
    scores = torch.matmul(q, k.transpose(-2, -1)) / (d_k ** 0.5)  # (B, S_q, S_k)
    
    if mask is not None:
        scores = scores.masked_fill(mask == 0, float('-inf'))

    attn_weights = softmax(scores, dim=-1)  # (B, S_q, S_k)
    
    output = torch.matmul(attn_weights, v)  # (B, S_q, D)
    
    return output

class MultiHeadAttention(nn.Module):
    def __init__(
        self,
        d_model: int,                  # 模型的隐藏层维度（例如 512）
        num_heads: int,                # 注意力头的数量（例如 8）
        use_rope: bool = False,        # 是否使用旋转位置编码（RoPE）
        max_seq_len: int | None = None, # 最大序列长度（用于 RoPE 缓存）
        theta: float | None = None,     # RoPE 的基底参数（通常为 10000.0）
        token_positions: torch.Tensor | None = None,  # 显式指定每个 token 的位置（用于 RoPE）
    ):
        super().__init__()
        
        self.d_model = d_model
        self.num_heads = num_heads
        
        # 检查 d_model 是否能被 num_heads 整除：每个头必须分到整数维度
        if d_model % num_heads != 0:
            raise ValueError("d_model must be divisible by num_heads")
        
        self.head_dim = d_model // num_heads  # 每个注意力头的维度

        # RoPE 相关配置
        self.use_rope = use_rope
        self.max_seq_len = max_seq_len
        self.theta = theta
        self.token_positions = token_positions

        # 如果启用 RoPE 且提供了必要参数，则初始化 RoPE 模块
        if use_rope and (max_seq_len is not None and theta is not None):
            # 注意：RoPE 通常作用于每个头的 head_dim 维度
            self.rope = RotaryPositionalEmbedding(theta, self.head_dim, max_seq_len)

        # 线性投影层：将输入分别映射为 Query、Key、Value
        self.q_proj = nn.Linear(d_model, d_model)  # 生成 Q
        self.k_proj = nn.Linear(d_model, d_model)  # 生成 K
        self.v_proj = nn.Linear(d_model, d_model)  # 生成 V
        
        # 输出投影层：将多头拼接后的结果映射回 d_model 维度
        self.out_proj = nn.Linear(d_model, d_model)
    
    def _causal_mask(self, seq_len: int) -> torch.Tensor:
        """
        生成因果掩码（Causal Mask），用于自回归生成（如语言模型）。
        只允许每个位置关注它自己及之前的位置，不能看未来。
        返回形状: (1, 1, seq_len, seq_len)
        """
        # torch.tril 保留下三角（包括对角线），上三角为 0
        mask = torch.tril(torch.ones(seq_len, seq_len)).bool()  # (S, S)
        # 增加 batch 和 head 维度以适配 attention 输入：(1, 1, S, S)
        mask = mask.unsqueeze(0).unsqueeze(0)
        return mask

    def forward(self, in_features: torch.Tensor):
        """
        前向传播：实现多头注意力计算。
        
        参数:
            in_features: 输入特征，形状 (B, S, D)
                B = batch size
                S = sequence length
                D = d_model（模型维度）
        
        返回:
            output: 注意力输出，形状 (B, S, D)
        """
        B, S, D = in_features.size()
        
        # 1. 线性投影 + 拆分为多头
        # 先投影到 d_model，再 reshape 成 (B, S, num_heads, head_dim)
        # 然后转置为 (B, num_heads, S, head_dim)，便于后续矩阵运算
        q = self.q_proj(in_features).view(B, S, self.num_heads, self.head_dim).transpose(1, 2)  # (B, H, S, D/H)
        k = self.k_proj(in_features).view(B, S, self.num_heads, self.head_dim).transpose(1, 2)  # (B, H, S, D/H)
        v = self.v_proj(in_features).view(B, S, self.num_heads, self.head_dim).transpose(1, 2)  # (B, H, S, D/H)

        # 2. （可选）应用旋转位置编码（RoPE）为了保留词汇的相对位置关系
        if self.use_rope:
            # RoPE 将位置信息通过旋转矩阵融入 q 和 k，保留相对位置关系
            q = self.rope(q, self.token_positions)  # token_positions 可为 None，此时默认使用 0,1,2,...
            k = self.rope(k, self.token_positions)
        
        # 3. 生成因果掩码（用于解码器或自回归任务）防止模型看到未来的数据
        mask = self._causal_mask(S) 
        mask = mask.to(q.device)  # 确保 mask 与数据在同一设备（CPU/GPU）
        
        # 4. 调用缩放点积注意力（注意：此处假设 scaled_dot_product_attention 支持 (B, H, S_q, S_k) 输入）
        # 输出形状: (B, H, S, D/H)
        attn_output = scaled_dot_product_attention(q, k, v, mask)
        
        # 5. 合并多头：转置回 (B, S, H, D/H) → reshape 为 (B, S, D)
        attn_output = attn_output.transpose(1, 2).contiguous().view(B, S, D)
        
        # 6. 最终线性投影
        output = self.out_proj(attn_output)  # (B, S, D)
        return output