"""第 25 章：相同矩阵乘法预算下，比模型宽度、数据量与颜色一致性。"""

import argparse
import json
from math import cos, exp, pi
from pathlib import Path
from statistics import mean, stdev
import torch
from torch import nn
from compute_budget import BUDGET, matrix_flops, parameters
from tiny_transformer import TinyTransformer
from small_language_model import PAD, corpus, vocabulary, make_windows, evaluate, corrupt_colors

SEEDS = [11, 22, 33]


def train(texts, validation, vocab, width, seed):
    torch.manual_seed(seed)
    model = TinyTransformer(len(vocab), width=width).double()
    assert sum(p.numel() for p in model.parameters()) == parameters(width)
    optimizer = torch.optim.AdamW(model.parameters(), lr=0.003, weight_decay=0.01)
    generator = torch.Generator().manual_seed(seed+1000)
    training = make_windows(texts, vocab, 32)
    x, y, _ = training
    cost = matrix_flops(width)
    steps = BUDGET//cost
    checkpoints = {max(1, round(steps*fraction)) for fraction in [0, .25, .5, .75, 1]}
    history, seen = [], set()
    for step in range(1, steps+1):
        model.train()
        batch = torch.randint(len(x), (16,), generator=generator)
        seen.update(batch.tolist())
        progress = step/steps
        multiplier = progress/.1 if progress <= .1 else .1+.9*(1+cos(pi*(progress-.1)/.9))/2
        for group in optimizer.param_groups:
            group["lr"] = .003*multiplier
        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()
        norm = nn.utils.clip_grad_norm_(model.parameters(), 1.0, error_if_nonfinite=True)
        optimizer.step()
        if step in checkpoints:
            val_loss, accuracy = evaluate(model, validation)
            train_loss, _ = evaluate(model, training)
            history.append(dict(step=step, compute=step*cost, train_loss=train_loss,
                                val_loss=val_loss, accuracy=accuracy, grad_norm=float(norm)))
    # 预先规定用预算结束时的参数；不在每个种子中再挑选最好的一步。
    val_loss, accuracy = evaluate(model, validation)
    return model, dict(seed=seed, val_loss=val_loss, accuracy=accuracy, history=history,
                       unique_seen=len(seen), steps=steps, compute=steps*cost)


def run():
    torch.set_num_threads(1)
    torch.use_deterministic_algorithms(True)
    train_texts, validation, test = corpus()
    vocab = vocabulary(train_texts)
    assert vocabulary(train_texts[:48]) == vocab
    val_data, test_data = [make_windows(texts, vocab, 32) for texts in [validation, test]]
    configurations = [("48 条一致", 32, train_texts[:48]),
                      ("48 条加宽", 64, train_texts[:48]),
                      ("192 条一致", 32, train_texts),
                      ("192 条错配", 32, corrupt_colors(train_texts)),
                      ("48 条错配", 32, corrupt_colors(train_texts[:48]))]
    results, models = [], []
    for name, width, texts in configurations:
        trials, candidates = [], []
        for seed in SEEDS:
            model, trial = train(texts, val_data, vocab, width, seed)
            candidates.append(model)
            trials.append(trial)
        results.append(dict(name=name, width=width, documents=len(texts), trials=trials,
                            mean_val=mean(t["val_loss"] for t in trials)))
        models.append(candidates)
    selected = min(range(len(results)), key=lambda index: results[index]["mean_val"])
    # 只有选定配置后才访问测试评分；不根据测试结果调整配置或挑种子。
    test_scores = [evaluate(model, test_data) for model in models[selected]]
    return dict(budget=BUDGET, seeds=SEEDS, results=results, selected=selected,
                test_scores=test_scores, vocab_size=len(vocab))


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--out", type=Path, help="可选：保存全部实验指标 JSON")
    args = parser.parse_args()
    report = run()
    print(f"预算上限={BUDGET:,} 矩阵 FLOPs；种子={SEEDS}；验证/测试各 32 条")
    for result in report["results"]:
        trials = result["trials"]
        losses, accuracies = [[trial[key] for trial in trials] for key in ["val_loss", "accuracy"]]
        print(f"{result['name']}：参数={parameters(result['width']):,}，更新={trials[0]['steps']}，"
              f"验证损失={mean(losses):.3f}±{stdev(losses):.3f}，颜色={mean(accuracies):.1%}")
        print("  各种子验证损失：" + ", ".join(f"{loss:.3f}" for loss in losses))
    base = report["results"][0]["trials"]
    wider = report["results"][1]["trials"]
    clean = report["results"][2]["trials"]
    noisy = report["results"][3]["trials"]
    small_noisy = report["results"][4]["trials"]
    for label, left, right in [("加宽−基线", wider, base), ("增量−基线", clean, base),
                               ("192 条：错配−一致", noisy, clean),
                               ("48 条：错配−一致", small_noisy, base)]:
        differences = [a["val_loss"]-b["val_loss"] for a, b in zip(left, right)]
        print(label + "：" + ", ".join(f"{value:+.3f}" for value in differences))
    selected = report["results"][report["selected"]]
    losses, accuracies = zip(*report["test_scores"])
    print(f"按平均验证损失选择：{selected['name']}")
    print(f"测试损失={mean(losses):.3f}±{stdev(losses):.3f}；"
          f"各种子困惑度=" + ", ".join(f"{exp(loss):.3f}" for loss in losses))
    print("测试颜色正确数：" + ", ".join(f"{round(value*32)}/32" for value in accuracies))
    if args.out:
        args.out.parent.mkdir(parents=True, exist_ok=True)
        args.out.write_text(json.dumps(report, ensure_ascii=False, indent=2))
        print(f"实验记录：{args.out}")


if __name__ == "__main__":
    main()