"""
01_attention_from_scratch.py
================================
用最少的程式（只靠 numpy）把 Transformer 的心臟——「自注意力 self-attention」——拆開給你看。
對照講義 docs/01_LLM_原理.md 的「注意力」一節。

核心公式（縮放點積注意力, scaled dot-product attention）：

    Attention(Q, K, V) = softmax( Q · Kᵀ / sqrt(d_k) ) · V

直覺：
  - Query (Q) = 「我在找什麼」
  - Key   (K) = 「我有什麼」
  - Value (V) = 「我能提供的內容」
  每個 token 用自己的 Q 去和所有 token 的 K 做點積（算相關度），
  正規化成權重後，對所有 token 的 V 做加權平均，得到「融合了上下文」的新表示。

執行：python3 01_attention_from_scratch.py
（不需要 GPU，不需要安裝任何東西，只要 numpy）
"""

import numpy as np

np.random.seed(0)


def softmax(x, axis=-1):
    """數值穩定版 softmax：先減去最大值再取 exp，避免溢位。"""
    x = x - np.max(x, axis=axis, keepdims=True)
    e = np.exp(x)
    return e / np.sum(e, axis=axis, keepdims=True)


def scaled_dot_product_attention(Q, K, V, mask=None):
    """
    Q, K, V 形狀都是 (T, d)：T 個 token、每個是 d 維向量。
    回傳 (輸出 (T, d), 注意力權重 (T, T))。
    """
    d_k = Q.shape[-1]
    scores = Q @ K.T / np.sqrt(d_k)        # (T, T)：每個 token 對每個 token 的相關度
    if mask is not None:
        # 因果遮罩：把「未來」的位置設成 -infinity，softmax 後權重會變 0，
        # 確保第 i 個字只能看到第 1..i 個字（這就是 GPT 一個字一個字往下生成的關鍵）。
        scores = np.where(mask == 0, -1e9, scores)
    weights = softmax(scores, axis=-1)     # (T, T)：每一列加總為 1
    output = weights @ V                   # (T, d)：用權重對 Value 做加權平均
    return output, weights


def multi_head_attention(x, n_heads, Wq, Wk, Wv, Wo, causal=True):
    """
    多頭注意力：把 d 維切成 n_heads 份，每份各做一次注意力（各看一種關係），最後拼回來再線性轉換。
    x: (T, d_model)。Wq/Wk/Wv: (d_model, d_model)。Wo: (d_model, d_model)。
    """
    T, d_model = x.shape
    d_head = d_model // n_heads

    # 1) 把輸入投影成 Q, K, V（這三個投影矩陣就是「學來的參數」）
    Q = x @ Wq      # (T, d_model)
    K = x @ Wk
    V = x @ Wv

    # 2) 切成多個 head
    def split_heads(m):
        return m.reshape(T, n_heads, d_head).transpose(1, 0, 2)  # (n_heads, T, d_head)

    Qh, Kh, Vh = split_heads(Q), split_heads(K), split_heads(V)

    # 因果遮罩（下三角為 1）
    mask = np.tril(np.ones((T, T))) if causal else None

    # 3) 每個 head 各做一次注意力
    head_outs = []
    head_weights = []
    for h in range(n_heads):
        out, w = scaled_dot_product_attention(Qh[h], Kh[h], Vh[h], mask)
        head_outs.append(out)
        head_weights.append(w)

    # 4) 把所有 head 拼回 (T, d_model)，再經過輸出投影 Wo
    concat = np.concatenate(head_outs, axis=-1)  # (T, d_model)
    output = concat @ Wo
    return output, np.stack(head_weights)        # 權重 shape: (n_heads, T, T)


def demo():
    # 一個小例子：4 個 token，模型維度 d_model = 8，2 個 head
    T, d_model, n_heads = 4, 8, 2
    tokens = ["the", "cat", "sat", "down"]

    # 假裝這是斷詞 + embedding 之後的結果：每個 token 一個 8 維向量（這裡用亂數代表）
    x = np.random.randn(T, d_model)

    # 學來的權重（這裡用亂數代表；真實模型是訓練出來的）
    Wq = np.random.randn(d_model, d_model) * 0.3
    Wk = np.random.randn(d_model, d_model) * 0.3
    Wv = np.random.randn(d_model, d_model) * 0.3
    Wo = np.random.randn(d_model, d_model) * 0.3

    out, weights = multi_head_attention(x, n_heads, Wq, Wk, Wv, Wo, causal=True)

    print("=" * 60)
    print("輸入序列：", tokens)
    print("輸入 x 形狀 (T, d_model):", x.shape)
    print("輸出 形狀 (T, d_model):", out.shape)
    print("=" * 60)
    print("第 1 個 head 的注意力權重矩陣 (T x T)：")
    print("（第 i 列 = 第 i 個字分配給每個字的注意力；因為有因果遮罩，右上角為 0）\n")
    header = "          " + "".join(f"{t:>8}" for t in tokens)
    print(header)
    for i, t in enumerate(tokens):
        row = "".join(f"{weights[0, i, j]:8.3f}" for j in range(T))
        print(f"{t:>8}  {row}")
    print("\n觀察：'down'（最後一個字）可以看到前面全部的字，")
    print("      'the'（第一個字）只能看到自己 → 權重 1.000。")
    print("這就是『因果』自注意力：預測下一個字時，只准看過去。")


if __name__ == "__main__":
    demo()

"""
02_tiny_gpt.py
================================
一個「迷你 GPT」：用 PyTorch 從頭把 docs/01_LLM_原理.md 講的東西實作出來——
embedding → 多層 Transformer（含因果自注意力）→ 預測下一個字 → 取樣生成。
這就是 Claude / GPT 的完整機制，只是縮到能在你筆電上跑。

需要先安裝 PyTorch（這台沙盒沒網路，所以請在你自己的電腦執行）：
    pip install torch
然後：
    python3 02_tiny_gpt.py
它會讀 data/input.txt（沒有的話用內建文字），訓練幾分鐘（CPU 也可以），最後生成一段文字。
有 GPU 會自動使用。想更快可把下面的 max_iters 調小。

對照：
  - token/char embedding ...... docs/01 的「斷詞」「embedding」
  - CausalSelfAttention ........ docs/01 的「注意力」公式 softmax(QKᵀ/√d)V（和 01_attention_from_scratch.py 同一套）
  - 訓練 loss = cross_entropy .. docs/01 的「預訓練：預測下一個 token」
  - generate() 取樣 ............ docs/01 的「推理」
"""

import os
import math
import torch
import torch.nn as nn
from torch.nn import functional as F

# ---------------- 超參數（刻意設小，方便在 CPU 上跑）----------------
block_size = 64        # 一次最多看幾個字（context window）
n_embd = 128           # 每個字的向量維度
n_head = 4             # 注意力的頭數
n_layer = 4            # 疊幾層 Transformer
dropout = 0.1
batch_size = 32
max_iters = 2000       # 訓練步數（想更快可改 500）
eval_interval = 200
eval_iters = 50
learning_rate = 3e-4
device = "cuda" if torch.cuda.is_available() else "cpu"
torch.manual_seed(1337)

# ---------------- 讀資料、建立字元詞表 ----------------
HERE = os.path.dirname(os.path.abspath(__file__))
data_path = os.path.join(HERE, "data", "input.txt")
if os.path.exists(data_path):
    with open(data_path, "r", encoding="utf-8") as f:
        text = f.read()
else:
    text = ("hello world. this is a tiny language model learning to predict "
            "the next character. the more it reads, the better it writes. ") * 200

chars = sorted(list(set(text)))
vocab_size = len(chars)
stoi = {ch: i for i, ch in enumerate(chars)}      # 字 -> 整數 ID
itos = {i: ch for i, ch in enumerate(chars)}      # 整數 ID -> 字
encode = lambda s: [stoi[c] for c in s]
decode = lambda l: "".join(itos[i] for i in l)

data = torch.tensor(encode(text), dtype=torch.long)
n = int(0.9 * len(data))
train_data, val_data = data[:n], data[n:]


def get_batch(split):
    """隨機取一批 (輸入 x, 目標 y)：y 是把 x 往右移一格（也就是『下一個字』）。"""
    d = train_data if split == "train" else val_data
    ix = torch.randint(len(d) - block_size, (batch_size,))
    x = torch.stack([d[i:i + block_size] for i in ix])
    y = torch.stack([d[i + 1:i + 1 + block_size] for i in ix])
    return x.to(device), y.to(device)


# ---------------- 模型零件 ----------------
class CausalSelfAttention(nn.Module):
    """因果（單向）多頭自注意力：實作 softmax(QKᵀ/√d)·V，並遮住未來。"""

    def __init__(self):
        super().__init__()
        assert n_embd % n_head == 0
        self.c_attn = nn.Linear(n_embd, 3 * n_embd)   # 一次算出 Q、K、V
        self.c_proj = nn.Linear(n_embd, n_embd)       # 輸出投影
        self.attn_drop = nn.Dropout(dropout)
        self.resid_drop = nn.Dropout(dropout)
        # 下三角遮罩（buffer 不是參數，不會被訓練）
        self.register_buffer("mask", torch.tril(torch.ones(block_size, block_size))
                             .view(1, 1, block_size, block_size))

    def forward(self, x):
        B, T, C = x.shape
        q, k, v = self.c_attn(x).split(n_embd, dim=2)   # 各 (B, T, C)
        hd = C // n_head
        # 切成多頭：(B, n_head, T, head_dim)
        q = q.view(B, T, n_head, hd).transpose(1, 2)
        k = k.view(B, T, n_head, hd).transpose(1, 2)
        v = v.view(B, T, n_head, hd).transpose(1, 2)
        # 注意力分數 = Q·Kᵀ / sqrt(head_dim)
        att = (q @ k.transpose(-2, -1)) / math.sqrt(hd)        # (B, n_head, T, T)
        att = att.masked_fill(self.mask[:, :, :T, :T] == 0, float("-inf"))  # 遮住未來
        att = F.softmax(att, dim=-1)
        att = self.attn_drop(att)
        y = att @ v                                            # 用權重對 V 加權平均
        y = y.transpose(1, 2).contiguous().view(B, T, C)       # 拼回 (B, T, C)
        return self.resid_drop(self.c_proj(y))


class MLP(nn.Module):
    """每個位置各自做的前饋網路（很多知識存在這裡）。"""

    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(n_embd, 4 * n_embd),
            nn.GELU(),
            nn.Linear(4 * n_embd, n_embd),
            nn.Dropout(dropout),
        )

    def forward(self, x):
        return self.net(x)


class Block(nn.Module):
    """一層 Transformer：LayerNorm → 注意力 → 殘差；LayerNorm → MLP → 殘差。"""

    def __init__(self):
        super().__init__()
        self.ln1 = nn.LayerNorm(n_embd)
        self.attn = CausalSelfAttention()
        self.ln2 = nn.LayerNorm(n_embd)
        self.mlp = MLP()

    def forward(self, x):
        x = x + self.attn(self.ln1(x))   # 殘差連接：讓梯度好傳、訓練更穩
        x = x + self.mlp(self.ln2(x))
        return x


class TinyGPT(nn.Module):
    def __init__(self):
        super().__init__()
        self.tok_emb = nn.Embedding(vocab_size, n_embd)     # 字 -> 向量
        self.pos_emb = nn.Embedding(block_size, n_embd)     # 位置 -> 向量
        self.drop = nn.Dropout(dropout)
        self.blocks = nn.Sequential(*[Block() for _ in range(n_layer)])
        self.ln_f = nn.LayerNorm(n_embd)
        self.head = nn.Linear(n_embd, vocab_size)           # 向量 -> 每個字的分數

    def forward(self, idx, targets=None):
        B, T = idx.shape
        pos = torch.arange(T, device=idx.device)
        x = self.drop(self.tok_emb(idx) + self.pos_emb(pos))   # 加總字向量與位置向量
        x = self.blocks(x)
        x = self.ln_f(x)
        logits = self.head(x)                                  # (B, T, vocab_size)
        loss = None
        if targets is not None:
            # 交叉熵：比較「模型猜的下一字分佈」和「真正的下一字」
            loss = F.cross_entropy(logits.view(-1, vocab_size), targets.view(-1))
        return logits, loss

    @torch.no_grad()
    def generate(self, idx, max_new_tokens, temperature=1.0, top_k=None):
        """從目前的字往後一個一個生成（自迴歸）。"""
        for _ in range(max_new_tokens):
            idx_cond = idx[:, -block_size:]            # 只保留最近 block_size 個字
            logits, _ = self(idx_cond)
            logits = logits[:, -1, :] / temperature    # 只看最後一個位置的預測
            if top_k is not None:
                v, _ = torch.topk(logits, top_k)
                logits[logits < v[:, [-1]]] = float("-inf")
            probs = F.softmax(logits, dim=-1)
            next_id = torch.multinomial(probs, num_samples=1)  # 依機率抽一個字
            idx = torch.cat([idx, next_id], dim=1)
        return idx


@torch.no_grad()
def estimate_loss(model):
    model.eval()
    out = {}
    for split in ["train", "val"]:
        losses = torch.zeros(eval_iters)
        for k in range(eval_iters):
            x, y = get_batch(split)
            _, loss = model(x, y)
            losses[k] = loss.item()
        out[split] = losses.mean().item()
    model.train()
    return out


def main():
    print(f"裝置：{device}　字元數(詞表大小)：{vocab_size}　資料長度：{len(data)} 字")
    model = TinyGPT().to(device)
    n_params = sum(p.numel() for p in model.parameters())
    print(f"模型參數量：{n_params/1e3:.1f}K\n開始訓練（目標：讓 loss 下降）...")

    optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate)
    for it in range(max_iters + 1):
        if it % eval_interval == 0:
            losses = estimate_loss(model)
            print(f"step {it:5d} | train loss {losses['train']:.3f} | val loss {losses['val']:.3f}")
        x, y = get_batch("train")
        _, loss = model(x, y)
        optimizer.zero_grad(set_to_none=True)
        loss.backward()              # 反向傳播：算梯度
        optimizer.step()             # 梯度下降：更新參數

    print("\n===== 生成結果（訓練後）=====")
    context = torch.zeros((1, 1), dtype=torch.long, device=device)
    out = model.generate(context, max_new_tokens=500, temperature=0.8, top_k=20)
    print(decode(out[0].tolist()))


if __name__ == "__main__":
    main()

"""
02b_tiny_gpt_exercise.py  ──「填空練習版」
================================
這是 02_tiny_gpt.py 的練習版：核心三段被挖空成 TODO，由你親手補上。
目的：把「讀得懂」升級成「寫得出來」——這才是路線圖第 3 站真正的里程碑。
（與其一直切出去看影片跟著別人手寫，不如在這裡自己寫、再對答案。）

玩法：
  1. 找到三個「✏️ TODO」，照上面的提示把程式碼補完（每個約 1～5 行）。
     （三個都要補；執行時不一定照 1→2→3 的順序報錯，把三個都寫完再跑即可。）
  2. 執行：python3 02b_tiny_gpt_exercise.py
     - 沒寫完會在該處 raise，告訴你還差哪裡。
     - 寫對了會開始訓練、loss 下降、最後生成文字。
  3. 卡住了，對照解答：02_tiny_gpt.py（同一支的完整版）。

需要 PyTorch：pip install torch
"""

import os
import math
import torch
import torch.nn as nn
from torch.nn import functional as F

# ---------------- 超參數 ----------------
block_size = 64
n_embd = 128
n_head = 4
n_layer = 4
dropout = 0.1
batch_size = 32
max_iters = 2000
eval_interval = 200
eval_iters = 50
learning_rate = 3e-4
device = "cuda" if torch.cuda.is_available() else "cpu"
torch.manual_seed(1337)

# ---------------- 讀資料、建立字元詞表 ----------------
HERE = os.path.dirname(os.path.abspath(__file__))
data_path = os.path.join(HERE, "data", "input.txt")
if os.path.exists(data_path):
    with open(data_path, "r", encoding="utf-8") as f:
        text = f.read()
else:
    text = ("hello world. this is a tiny language model learning to predict "
            "the next character. the more it reads, the better it writes. ") * 200

chars = sorted(list(set(text)))
vocab_size = len(chars)
stoi = {ch: i for i, ch in enumerate(chars)}
itos = {i: ch for i, ch in enumerate(chars)}
encode = lambda s: [stoi[c] for c in s]
decode = lambda l: "".join(itos[i] for i in l)

data = torch.tensor(encode(text), dtype=torch.long)
n = int(0.9 * len(data))
train_data, val_data = data[:n], data[n:]


def get_batch(split):
    d = train_data if split == "train" else val_data
    ix = torch.randint(len(d) - block_size, (batch_size,))
    x = torch.stack([d[i:i + block_size] for i in ix])
    y = torch.stack([d[i + 1:i + 1 + block_size] for i in ix])
    return x.to(device), y.to(device)


class CausalSelfAttention(nn.Module):
    def __init__(self):
        super().__init__()
        assert n_embd % n_head == 0
        self.c_attn = nn.Linear(n_embd, 3 * n_embd)
        self.c_proj = nn.Linear(n_embd, n_embd)
        self.attn_drop = nn.Dropout(dropout)
        self.resid_drop = nn.Dropout(dropout)
        self.register_buffer("mask", torch.tril(torch.ones(block_size, block_size))
                             .view(1, 1, block_size, block_size))

    def forward(self, x):
        B, T, C = x.shape
        q, k, v = self.c_attn(x).split(n_embd, dim=2)
        hd = C // n_head
        q = q.view(B, T, n_head, hd).transpose(1, 2)
        k = k.view(B, T, n_head, hd).transpose(1, 2)
        v = v.view(B, T, n_head, hd).transpose(1, 2)
        # ─────────────────────────────────────────────────────────
        # ✏️ TODO 1（心臟）：實作 softmax(Q·Kᵀ / sqrt(head_dim)) · V，並遮住未來。
        #   步驟提示：
        #     a) att = q @ k 的轉置(最後兩維互換) / math.sqrt(hd)        # (B, n_head, T, T)
        #     b) 用 self.mask[:, :, :T, :T] == 0 把「未來」位置 masked_fill 成 float("-inf")
        #     c) att = F.softmax(att, dim=-1)，再過 self.attn_drop
        #     d) y = att @ v，得到加權平均後的 (B, n_head, T, hd)
        #   （對答案見 02_tiny_gpt.py 的 CausalSelfAttention.forward）
        raise NotImplementedError("✏️ 請完成 TODO 1：縮放點積注意力 + 因果遮罩")
        # ─────────────────────────────────────────────────────────
        y = y.transpose(1, 2).contiguous().view(B, T, C)
        return self.resid_drop(self.c_proj(y))


class MLP(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(n_embd, 4 * n_embd),
            nn.GELU(),
            nn.Linear(4 * n_embd, n_embd),
            nn.Dropout(dropout),
        )

    def forward(self, x):
        return self.net(x)


class Block(nn.Module):
    def __init__(self):
        super().__init__()
        self.ln1 = nn.LayerNorm(n_embd)
        self.attn = CausalSelfAttention()
        self.ln2 = nn.LayerNorm(n_embd)
        self.mlp = MLP()

    def forward(self, x):
        # ✏️ TODO 2（殘差連接）：補上兩條殘差。
        #   提示：x = x + 注意力(先過 ln1 的 x)；再 x = x + mlp(先過 ln2 的 x)。
        #   為什麼是「x + ...」而不是直接覆蓋？殘差讓深層網路梯度好傳、訓練穩定。
        raise NotImplementedError("✏️ 請完成 TODO 2：兩條殘差連接")
        return x


class TinyGPT(nn.Module):
    def __init__(self):
        super().__init__()
        self.tok_emb = nn.Embedding(vocab_size, n_embd)
        self.pos_emb = nn.Embedding(block_size, n_embd)
        self.drop = nn.Dropout(dropout)
        self.blocks = nn.Sequential(*[Block() for _ in range(n_layer)])
        self.ln_f = nn.LayerNorm(n_embd)
        self.head = nn.Linear(n_embd, vocab_size)

    def forward(self, idx, targets=None):
        B, T = idx.shape
        pos = torch.arange(T, device=idx.device)
        x = self.drop(self.tok_emb(idx) + self.pos_emb(pos))
        x = self.blocks(x)
        x = self.ln_f(x)
        logits = self.head(x)
        loss = None
        if targets is not None:
            # ✏️ TODO 3（訓練目標）：用交叉熵比較「模型猜的下一字分佈」與「真正的下一字」。
            #   提示：F.cross_entropy(logits.view(-1, vocab_size), targets.view(-1))
            #   為什麼要 view 攤平？cross_entropy 想要 (N, C) 的 logits 與 (N,) 的目標。
            raise NotImplementedError("✏️ 請完成 TODO 3：交叉熵損失")
        return logits, loss

    @torch.no_grad()
    def generate(self, idx, max_new_tokens, temperature=1.0, top_k=None):
        for _ in range(max_new_tokens):
            idx_cond = idx[:, -block_size:]
            logits, _ = self(idx_cond)
            logits = logits[:, -1, :] / temperature
            if top_k is not None:
                v, _ = torch.topk(logits, top_k)
                logits[logits < v[:, [-1]]] = float("-inf")
            probs = F.softmax(logits, dim=-1)
            next_id = torch.multinomial(probs, num_samples=1)
            idx = torch.cat([idx, next_id], dim=1)
        return idx


@torch.no_grad()
def estimate_loss(model):
    model.eval()
    out = {}
    for split in ["train", "val"]:
        losses = torch.zeros(eval_iters)
        for k in range(eval_iters):
            x, y = get_batch(split)
            _, loss = model(x, y)
            losses[k] = loss.item()
        out[split] = losses.mean().item()
    model.train()
    return out


def main():
    print(f"裝置：{device}　詞表大小：{vocab_size}　資料長度：{len(data)} 字")
    model = TinyGPT().to(device)
    n_params = sum(p.numel() for p in model.parameters())
    print(f"模型參數量：{n_params/1e3:.1f}K\n開始訓練（目標：讓 loss 下降）...")

    optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate)
    for it in range(max_iters + 1):
        if it % eval_interval == 0:
            losses = estimate_loss(model)
            print(f"step {it:5d} | train loss {losses['train']:.3f} | val loss {losses['val']:.3f}")
        x, y = get_batch("train")
        _, loss = model(x, y)
        optimizer.zero_grad(set_to_none=True)
        loss.backward()
        optimizer.step()

    print("\n===== 生成結果（訓練後）=====")
    context = torch.zeros((1, 1), dtype=torch.long, device=device)
    out = model.generate(context, max_new_tokens=500, temperature=0.8, top_k=20)
    print(decode(out[0].tolist()))


if __name__ == "__main__":
    main()

"""
05_bpe_tokenizer.py
================================
從零做一個 BPE（Byte-Pair Encoding）斷詞器——docs/01 講的「斷詞」這一步。
真實的 GPT/Claude 都用 BPE 把文字切成 token。

做法（位元組級，能處理任何語言，包含中文）：
  1. 先把文字用 UTF-8 編成位元組（0~255 的整數）。
  2. 重複：找出「出現最多次的相鄰一對」，把它合併成一個新 token（新 ID 從 256 開始），記下這條合併規則。
  3. 做指定次數後，就得到一張詞表與一組合併規則。
  4. encode：照合併規則把文字壓成 token；decode：反過來還原。

執行：python3 05_bpe_tokenizer.py   （純 Python，不需任何套件）
參考：Andrej Karpathy 的 minbpe。
"""


def get_stats(ids):
    """統計每個相鄰 pair 出現幾次。"""
    counts = {}
    for a, b in zip(ids, ids[1:]):
        counts[(a, b)] = counts.get((a, b), 0) + 1
    return counts


def merge(ids, pair, new_id):
    """把序列中所有的 pair 換成 new_id。"""
    out, i = [], 0
    while i < len(ids):
        if i < len(ids) - 1 and ids[i] == pair[0] and ids[i + 1] == pair[1]:
            out.append(new_id)
            i += 2
        else:
            out.append(ids[i])
            i += 1
    return out


def train_bpe(text, num_merges):
    """訓練：回傳 merges（合併規則）與 vocab（id → bytes）。"""
    ids = list(text.encode("utf-8"))
    merges = {}                                  # (a,b) -> new_id
    vocab = {i: bytes([i]) for i in range(256)}  # 前 256 個是原始位元組
    for k in range(num_merges):
        stats = get_stats(ids)
        if not stats:
            break
        pair = max(stats, key=stats.get)         # 出現最多次的相鄰對
        new_id = 256 + k
        ids = merge(ids, pair, new_id)
        merges[pair] = new_id
        vocab[new_id] = vocab[pair[0]] + vocab[pair[1]]
    return merges, vocab


def encode(text, merges):
    """把文字壓成 token id 序列：依合併規則的順序，能合就合。"""
    ids = list(text.encode("utf-8"))
    while len(ids) >= 2:
        stats = get_stats(ids)
        # 找出「最早被學到」的那個可合併的 pair（min merge index）
        pair = min(stats, key=lambda p: merges.get(p, float("inf")))
        if pair not in merges:
            break                                # 沒有可再合併的了
        ids = merge(ids, pair, merges[pair])
    return ids


def decode(ids, vocab):
    """把 token id 還原成文字。"""
    data = b"".join(vocab[i] for i in ids)
    return data.decode("utf-8", errors="replace")


def main():
    text = (
        "機器學習很有趣，機器學習能讓電腦從資料中學習。"
        "深度學習是機器學習的一種，深度學習用神經網路。"
        "tokenization tokenization is the first step. "
        "the model reads tokens, the model predicts tokens. "
    ) * 3

    print(f"原始文字長度：{len(text)} 個字")
    raw_bytes = len(text.encode("utf-8"))
    print(f"UTF-8 位元組數：{raw_bytes}")

    num_merges = 60
    merges, vocab = train_bpe(text, num_merges)
    print(f"\n學了 {len(merges)} 條合併規則，詞表大小 = {len(vocab)}")

    print("\n前 8 條學到的合併規則（被合併的兩段 → 合併後的字串）：")
    for i, (pair, nid) in enumerate(list(merges.items())[:8]):
        piece = vocab[nid].decode("utf-8", errors="replace")
        print(f"  {i+1}. {pair} → id {nid}  代表「{piece}」")

    sample = "機器學習很有趣"
    ids = encode(sample, merges)
    print(f"\n編碼範例：「{sample}」")
    print(f"  → token ids: {ids}")
    print(f"  → token 數：{len(ids)}（原本 {len(sample.encode('utf-8'))} 個位元組，被壓縮了）")
    print(f"  → 還原 decode：「{decode(ids, merges and vocab)}」（應與原句相同）")

    # 整段壓縮率
    all_ids = encode(text, merges)
    print(f"\n整段壓縮：{raw_bytes} 位元組 → {len(all_ids)} 個 token"
          f"（壓縮率約 {raw_bytes/len(all_ids):.2f}×）")
    print("\n重點：BPE 學會把『常一起出現的片段』併成一個 token，")
    print("      所以常見字詞用更少 token 表示——這就是 GPT 處理文字的第一步。")


if __name__ == "__main__":
    main()