"""第 14 章：控制变量比较初始化与学习率；CPU 上运行。"""

from copy import deepcopy
import numpy as np
import torch
from torch import nn
from sklearn.datasets import make_moons
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler


def make_data():
    x, y = make_moons(n_samples=600, noise=0.18, random_state=2026)
    x_dev, x_test, y_dev, y_test = train_test_split(
        x, y, test_size=0.2, stratify=y, random_state=2026)
    x_train, x_val, y_train, y_val = train_test_split(
        x_dev, y_dev, test_size=0.25, stratify=y_dev, random_state=2027)
    scaler = StandardScaler().fit(x_train)
    datasets = []
    for features, labels in [(x_train, y_train), (x_val, y_val), (x_test, y_test)]:
        datasets.append((torch.tensor(scaler.transform(features), dtype=torch.float64),
                         torch.tensor(labels, dtype=torch.long)))
    return datasets, scaler


def make_model(zero=False):
    model = nn.Sequential(nn.Linear(2, 32), nn.ReLU(),
                          nn.Linear(32, 32), nn.ReLU(), nn.Linear(32, 2)).double()
    layers = [module for module in model if isinstance(module, nn.Linear)]
    for index, layer in enumerate(layers):
        if zero:
            nn.init.zeros_(layer.weight)
        elif index == len(layers)-1:
            nn.init.xavier_normal_(layer.weight)
        else:
            nn.init.kaiming_normal_(layer.weight, nonlinearity="relu")
        nn.init.zeros_(layer.bias)
    return model


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


def train(train_data, validation, rate, zero=False, epochs=60):
    torch.manual_seed(2026)
    model = make_model(zero)
    optimizer = torch.optim.SGD(model.parameters(), lr=rate, momentum=0.9,
                                weight_decay=1e-4)
    scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=30, gamma=0.3)
    generator = torch.Generator().manual_seed(2027)
    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(32):
            optimizer.zero_grad(set_to_none=True)
            loss = nn.functional.cross_entropy(model(x[batch]), y[batch])
            if not torch.isfinite(loss):
                raise RuntimeError("训练损失不是有限数，请检查输入、尺度和学习率")
            loss.backward()
            optimizer.step()
        train_loss, _ = metrics(model, train_data)
        val_loss, val_accuracy = metrics(model, validation)
        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())
        scheduler.step()
    model.load_state_dict(best_state)
    return model, best_epoch, history


def run():
    torch.set_num_threads(1)
    torch.use_deterministic_algorithms(True)
    datasets, scaler = make_data()
    configurations = [("步长很小", 0.0001, False),
                      ("常用设置", 0.05, False),
                      ("全零初始化", 0.05, True)]
    results = []
    for name, rate, zero in configurations:
        model, epoch, history = train(datasets[0], datasets[1], rate, zero)
        val_loss, val_accuracy = metrics(model, datasets[1])
        results.append((name, model, epoch, history, val_loss, val_accuracy))
    return datasets, scaler, results


def main():
    datasets, _, results = run()
    print("训练/验证/测试：360 / 120 / 120")
    for name, _, epoch, _, loss, accuracy in results:
        print(f"{name}：最佳轮次={epoch}，验证交叉熵={loss:.3f}，准确率={accuracy:.3f}")
    name, model, _, _, _, _ = min(results, key=lambda row: row[4])
    loss, accuracy = metrics(model, datasets[2])
    print(f"固定方案：{name}；测试交叉熵={loss:.3f}，准确率={accuracy:.3f}")


if __name__ == "__main__":
    main()