# 把 10 個數字倒過來：小型 Transformer（自注意力）與 LSTM 比較，並畫出 Transformer 學到的注意力
import json
import math
from pathlib import Path

import torch
from torch import nn

L, V, D = 10, 10, 32                                # 序列長 10、數字 0～9、向量 32 維

def batch(n, g):
    x = torch.randint(0, V, (n, L), generator=g)
    return x, x.flip(1)                             # 目標：倒過來的序列

def pos_encoding(L, d):
    """正弦位置編碼：PE[p, 2i] = sin(p / 10000^(2i/d))，PE[p, 2i+1] = cos(…)"""
    p = torch.arange(L).unsqueeze(1)
    w = torch.exp(-math.log(10000.0) * torch.arange(0, d, 2) / d)
    pe = torch.zeros(L, d)
    pe[:, 0::2], pe[:, 1::2] = torch.sin(p * w), torch.cos(p * w)
    return pe

class TinyTransformer(nn.Module):
    def __init__(self):
        super().__init__()
        self.emb = nn.Embedding(V, D)
        self.register_buffer("pe", pos_encoding(L, D))
        self.att = nn.MultiheadAttention(D, num_heads=1, batch_first=True)
        self.ff = nn.Sequential(nn.Linear(D, 64), nn.ReLU(), nn.Linear(64, D))
        self.n1, self.n2 = nn.LayerNorm(D), nn.LayerNorm(D)
        self.out = nn.Linear(D, V)
    def forward(self, x, need_att=False):
        h = self.emb(x) + self.pe                    # 數字的向量＋位置
        a, w = self.att(h, h, h, need_weights=need_att)
        h = self.n1(h + a)                           # 殘差＋層正規化
        h = self.n2(h + self.ff(h))
        return (self.out(h), w) if need_att else self.out(h)

class TinyLSTM(nn.Module):
    def __init__(self):
        super().__init__()
        self.emb = nn.Embedding(V, D)
        self.rnn = nn.LSTM(D, 48, batch_first=True, bidirectional=True)   # 雙向：每個位置看得到整個序列
        self.out = nn.Linear(96, V)
    def forward(self, x):
        return self.out(self.rnn(self.emb(x))[0])

g = torch.Generator().manual_seed(0)
xt, yt = batch(2000, torch.Generator().manual_seed(99))                # 測試資料
result = {}
for name, Net in [("Transformer", TinyTransformer), ("雙向 LSTM", TinyLSTM)]:
    torch.manual_seed(0)
    net = Net()
    opt = torch.optim.Adam(net.parameters(), lr=3e-3)
    for step in range(1500):
        x, y = batch(64, g)
        loss = nn.functional.cross_entropy(net(x).reshape(-1, V), y.reshape(-1))
        opt.zero_grad(); loss.backward(); opt.step()
    with torch.no_grad():
        pred = net(xt).argmax(-1)
        tok, seq = (pred == yt).float().mean().item(), (pred == yt).all(1).float().mean().item()
    n_par = sum(p.numel() for p in net.parameters())
    print(f"{name:11s}：參數 {n_par:5d}，每個數字的正確率 {tok:.1%}，整串全對 {seq:.1%}")
    if name == "Transformer":
        with torch.no_grad():
            _, w = net(xt[:1], need_att=True)
        result["att"] = w[0].tolist()                                  # 第一筆測試資料的注意力權重（10×10）
        result["x"] = xt[0].tolist()
Path(__file__).with_name("reverse_seq.json").write_text(json.dumps(result), encoding="utf-8")