"""第 21 章：训练、比较、保存和生成短文本；与 tiny_transformer.py 放在同一目录。"""

import argparse
from collections import Counter
from copy import deepcopy
from hashlib import sha256
from itertools import product
from math import exp
from pathlib import Path
from random import Random
from tempfile import TemporaryDirectory
import torch
from torch import nn
from tiny_transformer import TinyTransformer

PAD, BOS, EOS = 0, 1, 2
COLOR_INDEX = 21  # 原句中第二次颜色的位置，输入加 BOS 后仍是相应目标的索引。


def corpus():
    records = [
        f"{name}选了{color}色{thing}，放到{place}。最后{name}拿走了{color}色{thing}。"
        for name, color, thing, place in product(
            ["小林", "小周", "小陈", "小李"], ["红", "蓝", "绿", "白"],
            ["杯子", "盒子", "本子", "袋子"], ["桌上", "架上", "柜里", "门边"])
    ]
    assert all(len(text) == 26 for text in records)
    Random(2026).shuffle(records)
    return records[:192], records[192:224], records[224:]


def corrupt_colors(texts):
    # 打乱第二次颜色的对应关系，保持文本长度和整个训练集的字符计数不变。
    colors = [text[COLOR_INDEX] for text in texts]
    Random(2028).shuffle(colors)
    return [text[:COLOR_INDEX]+color+text[COLOR_INDEX+1:]
            for text, color in zip(texts, colors)]


def vocabulary(texts):
    return ["<PAD>", "<BOS>", "<EOS>"]+sorted(set("".join(texts)))


def make_windows(texts, vocab, block_size):
    ids = {token: index for index, token in enumerate(vocab)}
    rows = []
    for text in texts:
        sequence = [BOS]+[ids[char] for char in text]+[EOS]
        for start in range(0, len(sequence)-1, block_size):
            x = sequence[start:start+block_size]
            y = sequence[start+1:start+block_size+1]
            # 最后一块的输入长度必须与可预测目标数量一致。
            x = x[:len(y)]
            marks = [start+i == COLOR_INDEX for i in range(len(y))]
            padding = block_size-len(y)
            rows.append((x+[PAD]*padding, y+[PAD]*padding, marks+[False]*padding))
    return tuple(torch.tensor([row[i] for row in rows], dtype=dtype)
                 for i, dtype in enumerate([torch.long, torch.long, torch.bool]))


@torch.no_grad()
def evaluate(model, data):
    model.eval()
    x, y, marks = data
    loss_sum, target_count, color_correct, color_count = 0.0, 0, 0, 0
    for start in range(0, len(x), 64):
        batch = slice(start, start+64)
        logits = model(x[batch])
        loss_sum += nn.functional.cross_entropy(
            logits.flatten(0, 1), y[batch].flatten(), ignore_index=PAD, reduction="sum").item()
        target_count += (y[batch] != PAD).sum().item()
        color_correct += ((logits.argmax(-1) == y[batch]) & marks[batch]).sum().item()
        color_count += marks[batch].sum().item()
    return loss_sum/target_count, color_correct/color_count


def train(texts, validation, vocab, width, block_size, epochs):
    torch.manual_seed(2026)
    model = TinyTransformer(len(vocab), max_context=32, width=width, heads=4, layers=2).double()
    data = make_windows(texts, vocab, block_size)
    val_data = make_windows(validation, vocab, block_size)
    optimizer = torch.optim.AdamW(model.parameters(), lr=0.003, weight_decay=0.01)
    generator = torch.Generator().manual_seed(2027)
    batch_size = 16 if block_size == 32 else 64
    best_loss, best_epoch, best_state = float("inf"), 0, None
    history = []
    x, y, _ = data
    for epoch in range(1, epochs+1):
        model.train()
        for batch in torch.randperm(len(x), generator=generator).split(batch_size):
            optimizer.zero_grad(set_to_none=True)
            logits = model(x[batch])
            loss = nn.functional.cross_entropy(logits.flatten(0, 1), y[batch].flatten(),
                                               ignore_index=PAD)
            if not torch.isfinite(loss):
                raise RuntimeError("损失不是有限数，请检查数据、掩码与学习率")
            loss.backward()
            nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
            optimizer.step()
        train_loss, _ = evaluate(model, data)
        val_loss, accuracy = evaluate(model, val_data)
        history.append((train_loss, val_loss, accuracy))
        if val_loss < best_loss:
            best_loss, best_epoch = val_loss, epoch
            best_state = deepcopy(model.state_dict())
    model.load_state_dict(best_state)
    return model, best_epoch, history


def run(epochs=40):
    torch.set_num_threads(1)
    torch.use_deterministic_algorithms(True)
    train_texts, validation, test = corpus()
    vocab = vocabulary(train_texts)
    configs = [
        ("基础模型", 32, 32, train_texts),
        ("加宽模型", 64, 32, train_texts),
        ("短上下文", 32, 8, train_texts),
        ("颜色错配", 32, 32, corrupt_colors(train_texts)),
    ]
    results = []
    for name, width, block_size, texts in configs:
        model, epoch, history = train(texts, validation, vocab, width, block_size, epochs)
        val_loss, accuracy = evaluate(model, make_windows(validation, vocab, block_size))
        results.append((name, block_size, model, epoch, history, val_loss, accuracy))
    selected = min(results, key=lambda row: row[5])
    return (train_texts, validation, test), vocab, results, selected


def choose(logits, generator, temperature=0.8, top_k=None, top_p=1.0):
    if temperature < 0 or not 0 < top_p <= 1:
        raise ValueError("temperature 不小于零，top_p 在 (0, 1] 内")
    logits = logits.clone()
    logits[[PAD, BOS]] = float("-inf")
    if temperature == 0:
        return int(logits.argmax())
    logits /= temperature
    if top_k is not None:
        if not 1 <= top_k <= len(logits):
            raise ValueError("top_k 必须在 1 与词表大小之间")
        indices = logits.topk(top_k).indices
        keep = torch.zeros_like(logits, dtype=torch.bool)
        keep[indices] = True
        logits[~keep] = float("-inf")
    sorted_logits, indices = logits.sort(descending=True)
    probabilities = sorted_logits.softmax(-1)
    # 保留使累计概率首次达到阈值的那个词元，至少保留第一项。
    remove = probabilities.cumsum(-1)-probabilities >= top_p
    sorted_logits[remove] = float("-inf")
    chosen = torch.multinomial(sorted_logits.softmax(-1), 1, generator=generator)
    return int(indices[chosen].item())


@torch.no_grad()
def generate(model, vocab, block_size, prompt="", temperature=0.8, top_k=8, top_p=0.9):
    model.eval()
    ids = {token: index for index, token in enumerate(vocab)}
    if any(char not in ids for char in prompt):
        raise ValueError("提示中有词表未收录的字符")
    sequence = [BOS]+[ids[char] for char in prompt]
    generator = torch.Generator().manual_seed(2029)
    for _ in range(60):
        inputs = torch.tensor([sequence[-block_size:]])
        token = choose(model(inputs)[0, -1], generator, temperature, top_k, top_p)
        if token == EOS:
            return "".join(vocab[index] for index in sequence[1:]), "EOS"
        sequence.append(token)
    return "".join(vocab[index] for index in sequence[1:]), "长度上限"


def save_model(path, model, vocab, block_size, name, epochs, best_epoch, train_texts):
    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)
    torch.save({"config": model.config, "state": model.state_dict(), "vocab": vocab,
                "block_size": block_size, "name": name, "epochs": epochs,
                "best_epoch": best_epoch,
                "torch_version": str(torch.__version__),
                "training_sha256": sha256("\n".join(train_texts).encode()).hexdigest()}, path)


def load_model(path):
    checkpoint = torch.load(path, map_location="cpu", weights_only=True)
    model = TinyTransformer(**checkpoint["config"]).double()
    model.load_state_dict(checkpoint["state"])
    model.eval()
    return model, checkpoint["vocab"], checkpoint["block_size"]


def show_generation(model, vocab, block_size):
    prompt = "小林选了红色杯子，放到桌上。最后小林拿走了"
    for label, prefix, temperature in [("贪心续写", prompt, 0.0), ("采样生成", "", 0.8)]:
        text, stop = generate(model, vocab, block_size, prefix, temperature)
        print(f"{label}：{text}（{stop}）")


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--epochs", type=int, default=40)
    actions = parser.add_mutually_exclusive_group()
    actions.add_argument("--save", type=Path)
    actions.add_argument("--load", type=Path)
    args = parser.parse_args()
    torch.set_num_threads(1)
    if args.load:
        show_generation(*load_model(args.load))
        return
    if args.epochs < 1:
        parser.error("--epochs 必须大于零")
    datasets, vocab, results, selected = run(args.epochs)
    counts = Counter(char for text in datasets[0] for char in text)
    counts["<EOS>"] = len(datasets[0])
    probabilities = torch.tensor([counts[token]+0.1 for token in vocab], dtype=torch.float64)
    probabilities /= probabilities.sum()
    indices = {token: i for i, token in enumerate(vocab)}
    targets = [indices[token] for text in datasets[1] for token in list(text)+["<EOS>"]]
    baseline_loss = -probabilities[targets].log().mean().item()
    print(f"训练/验证/测试文本：{len(datasets[0])} / {len(datasets[1])} / {len(datasets[2])}")
    print(f"字符词表大小：{len(vocab)}；每轮训练目标数={len(datasets[0])*27}")
    print(f"一元基线验证交叉熵={baseline_loss:.3f}，困惑度={exp(baseline_loss):.3f}")
    corrupt = corrupt_colors(datasets[0])
    wrong = sum(a[COLOR_INDEX] != b[COLOR_INDEX] for a, b in zip(datasets[0], corrupt))
    print(f"颜色错配训练文本：{wrong} / {len(corrupt)}")
    for name, block_size, model, epoch, _, loss, accuracy in results:
        count = sum(p.numel() for p in model.parameters())
        print(f"{name}：参数={count}，块长={block_size}，最佳轮次={epoch}，"
              f"验证交叉熵={loss:.3f}，困惑度={exp(loss):.3f}，末次颜色准确率={accuracy:.3f}")
    name, block_size, model, best_epoch, _, _, _ = selected
    loss, accuracy = evaluate(model, make_windows(datasets[2], vocab, block_size))
    print(f"固定方案：{name}；测试交叉熵={loss:.3f}，"
          f"困惑度={exp(loss):.3f}，末次颜色准确率={accuracy:.3f}")
    prompt = "小林选了红色杯子，放到桌上。最后小林拿走了"
    for other_name, other_block, other_model, _, _, _, _ in results[2:]:
        text, stop = generate(other_model, vocab, other_block, prompt, temperature=0)
        print(f"{other_name}续写：{text}（{stop}）")
    with TemporaryDirectory() as folder:
        path = args.save or Path(folder)/"tiny-lm.pt"
        selected_training = corrupt if name == "颜色错配" else datasets[0]
        save_model(path, model, vocab, block_size, name, args.epochs, best_epoch, selected_training)
        restored, restored_vocab, restored_block = load_model(path)
        probe = torch.tensor([[BOS, vocab.index("小"), vocab.index("林")]])
        with torch.no_grad():
            equal = torch.equal(model(probe), restored(probe))
        assert restored_vocab == vocab and restored_block == block_size
        print(f"保存再加载后输出一致：{equal}")
        if args.save:
            print(f"已保存：{args.save}")
        show_generation(restored, vocab, restored_block)


if __name__ == "__main__":
    main()