"""第三篇实践：本地手写数字分类、学习曲线与受控比较；CPU 实验。"""

from copy import deepcopy
import numpy as np
import torch
from torch import nn
from sklearn.datasets import load_digits
from sklearn.model_selection import train_test_split


def make_data():
    digits = load_digits()
    indices = np.arange(len(digits.target))
    development, test = train_test_split(
        indices, test_size=0.2, stratify=digits.target, random_state=2026)
    train, validation = train_test_split(
        development, test_size=0.25, stratify=digits.target[development],
        random_state=2027)
    datasets = []
    for ids in [train, validation, test]:
        datasets.append((torch.tensor(digits.images[ids, None], dtype=torch.float64),
                         torch.tensor(digits.target[ids], dtype=torch.long)))
    return datasets, [train, validation, test]


def make_model(kind):
    if kind == "mlp":
        return nn.Sequential(nn.Flatten(), nn.Linear(64, 32), nn.ReLU(),
                             nn.Linear(32, 10)).double()
    return nn.Sequential(nn.Conv2d(1, 8, 3, padding=1), nn.ReLU(),
                         nn.MaxPool2d(2), nn.Flatten(), nn.Linear(128, 32),
                         nn.ReLU(), nn.Linear(32, 10)).double()


def metrics(model, dataset, scale):
    model.eval()
    with torch.no_grad():
        x, y = dataset
        logits = model(x/scale)
        loss = nn.functional.cross_entropy(logits, y).item()
        prediction = logits.argmax(1)
        accuracy = (prediction == y).double().mean().item()
    return loss, accuracy, prediction


def train(kind, scale, train_data, validation, seed=2026, epochs=25):
    torch.manual_seed(seed)
    model = make_model(kind)
    optimizer = torch.optim.AdamW(model.parameters(), lr=0.001, weight_decay=0.001)
    generator = torch.Generator().manual_seed(2030)
    best_loss, best_epoch, best_state = float("inf"), 0, None
    history = []
    x, y = train_data
    for epoch in range(1, epochs+1):
        model.train()
        for batch in torch.randperm(len(x), generator=generator).split(64):
            optimizer.zero_grad(set_to_none=True)
            loss = nn.functional.cross_entropy(model(x[batch]/scale), y[batch])
            if not torch.isfinite(loss):
                raise RuntimeError("损失出现非有限值")
            loss.backward()
            optimizer.step()
        train_loss, _, _ = metrics(model, train_data, scale)
        val_loss, val_accuracy, _ = metrics(model, validation, scale)
        history.append((train_loss, val_loss, val_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():
    torch.set_num_threads(1)
    torch.use_deterministic_algorithms(True)
    datasets, ids = make_data()
    configurations = [("MLP / 16", "mlp", 16.0),
                      ("CNN / 16", "cnn", 16.0),
                      ("CNN 原像素", "cnn", 1.0)]
    results = []
    for name, kind, scale in configurations:
        model, epoch, history = train(kind, scale, datasets[0], datasets[1])
        val_loss, accuracy, _ = metrics(model, datasets[1], scale)
        results.append((name, scale, model, epoch, history, val_loss, accuracy))
    selected = min(results, key=lambda row: row[5])
    return datasets, ids, results, selected


def main():
    datasets, ids, results, selected = run()
    print("训练/验证/测试：" + " / ".join(str(len(x)) for x, _ in datasets))
    majority = torch.bincount(datasets[0][1]).argmax()
    baseline = (datasets[1][1] == majority).double().mean().item()
    print(f"多数类基线验证准确率={baseline:.3f}")
    for name, _, model, epoch, _, loss, accuracy in results:
        count = sum(p.numel() for p in model.parameters())
        print(f"{name}：参数={count}，最佳轮次={epoch}，"
              f"验证交叉熵={loss:.3f}，准确率={accuracy:.3f}")
    name, scale, model, _, _, _, _ = selected
    loss, accuracy, prediction = metrics(model, datasets[2], scale)
    print(f"固定方案：{name}；测试交叉熵={loss:.3f}，准确率={accuracy:.3f}")
    target = datasets[2][1]
    wrong = torch.where(prediction != target)[0]
    print(f"测试错误数：{len(wrong)} / {len(target)}")
    for index in wrong[:3]:
        print(f"样本 ID={ids[2][index]}：真实={target[index].item()}，"
              f"预测={prediction[index].item()}")


if __name__ == "__main__":
    main()