"""第 45 章：真实训练、独立校准、保存及重新加载。CPU，约数秒至几十秒。"""

from argparse import ArgumentParser
import copy
import hashlib
import json
import math
from pathlib import Path
import platform
import time

import numpy as np
import scipy
from scipy.optimize import minimize_scalar
from scipy.special import logsumexp
import sklearn
from sklearn.datasets import load_digits
from sklearn.model_selection import train_test_split
import torch
from torch import nn
from torch.nn import functional as F

HERE = Path(__file__).resolve().parent
DEFAULT_CONFIG = HERE / "data/calibration-config.json"


def digest(value):
    return hashlib.sha256(json.dumps(value, sort_keys=True, separators=(",", ":"),
                                     allow_nan=False).encode()).hexdigest()


def file_hash(path):
    return hashlib.sha256(Path(path).read_bytes()).hexdigest()


def state_hash(state):
    h = hashlib.sha256()
    for name, tensor in sorted(state.items()):
        a = tensor.detach().cpu().contiguous().numpy()
        h.update(json.dumps([name, list(a.shape), a.dtype.str]).encode())
        h.update(a.tobytes())
    return h.hexdigest()


def configure():
    torch.set_num_threads(1)
    torch.use_deterministic_algorithms(True)


def load_data(config):
    x, y = load_digits(return_X_y=True)
    raw_hash = digest({"pixels": x.astype(int).tolist(), "labels": y.tolist()})
    pool = np.arange(len(y))
    splits = {}
    for role, seed in zip(["test", "calibration", "selection"], config["split_seeds"]):
        pool, splits[role] = train_test_split(
            pool, test_size=config[role + "_size"],
            stratify=y[pool], random_state=seed)
    splits["train"] = pool
    # 16 是数据定义给出的像素上限，不从全部样本估计归一化统计量。
    return (torch.tensor(x / 16, dtype=torch.float32),
            torch.tensor(y, dtype=torch.long), splits, raw_hash)


def make_model(config):
    a, b = config["hidden"]
    return nn.Sequential(nn.Linear(64, a), nn.ReLU(), nn.Linear(a, b),
                         nn.ReLU(), nn.Linear(b, 10))


def checked(logits, labels, temperature=1.0):
    z, y = np.asarray(logits, dtype=np.float64), np.asarray(labels)
    if (z.ndim != 2 or len(z) == 0 or z.shape[1] < 2 or
            y.shape != (len(z),) or not np.issubdtype(y.dtype, np.integer) or
            not np.isfinite(z).all() or np.any(y < 0) or np.any(y >= z.shape[1]) or
            not math.isfinite(temperature) or temperature <= 0):
        raise ValueError("需要有限分数、合法整数标签和正温度")
    return (z - z.max(axis=1, keepdims=True)) / temperature, y


def nll(logits, labels, temperature=1.0):
    z, y = checked(logits, labels, temperature)
    return float(np.mean(logsumexp(z, axis=1) - z[np.arange(len(y)), y]))


def metrics(logits, labels, temperature=1.0, bins=15):
    z, y = checked(logits, labels, temperature)
    p = np.exp(z - logsumexp(z, axis=1, keepdims=True))
    prediction, confidence = z.argmax(axis=1), p.max(axis=1)
    correct = prediction == y
    bucket = np.clip(np.ceil(confidence * bins).astype(int) - 1, 0, bins - 1)
    groups = []
    for m in range(bins):
        use = bucket == m
        if use.any():
            groups.append({"bin": m, "count": int(use.sum()),
                           "confidence": float(confidence[use].mean()),
                           "accuracy": float(correct[use].mean())})
    return {"nll": nll(logits, labels, temperature),
            "accuracy": float(correct.mean()), "confidence": float(confidence.mean()),
            "ece": sum(g["count"] / len(y) * abs(g["accuracy"] - g["confidence"])
                       for g in groups), "bins": groups}


def fit_temperature(logits, labels, bounds):
    low, high = bounds
    if not (0 < low < high and math.isfinite(high)):
        raise ValueError("温度区间必须有限且为正")
    objective = lambda s: nll(logits, labels, math.exp(s))
    result = minimize_scalar(objective, bounds=(math.log(low), math.log(high)),
                             method="bounded", options={"xatol": 1e-10})
    if not result.success:
        raise RuntimeError(result.message)
    # 显式比较端点；也保留区间内的 T=1，避免数值误差输给“不校准”。
    candidates = [low, high, math.exp(result.x)]
    if low <= 1 <= high:
        candidates.append(1.0)
    return float(min(candidates, key=lambda t: nll(logits, labels, t)))


def train_one(config, seed, x_train, y_train, x_selection, y_selection):
    """接口只接触训练和选模数据，校准与测试数据不进入训练。"""
    torch.manual_seed(seed)
    generator = torch.Generator().manual_seed(seed + 1000)
    model = make_model(config)
    initial_hash = state_hash(model.state_dict())
    optimizer = torch.optim.Adam(model.parameters(), lr=config["learning_rate"],
                                 weight_decay=0)
    curves, best, best_loss, first_order = [], None, math.inf, None
    for epoch in range(1, config["epochs"] + 1):
        model.train()
        order = torch.randperm(len(y_train), generator=generator)
        if first_order is None:
            first_order = digest(order.tolist())
        for ids in order.split(config["batch_size"]):
            optimizer.zero_grad(set_to_none=True)
            F.cross_entropy(model(x_train[ids]), y_train[ids]).backward()
            optimizer.step()
        model.eval()
        with torch.no_grad():
            train_loss = F.cross_entropy(model(x_train), y_train).item()
            selection_loss = F.cross_entropy(model(x_selection), y_selection).item()
        curves.append({"epoch": epoch, "train_nll": train_loss,
                       "selection_nll": selection_loss})
        if selection_loss < best_loss:
            best_loss, best_epoch = selection_loss, epoch
            best = copy.deepcopy(model.state_dict())
    model.load_state_dict(best)
    return model, {"seed": seed, "selected_epoch": best_epoch, "curves": curves,
                   "initial_sha256": initial_hash, "first_order_sha256": first_order,
                   "weights_sha256": state_hash(best)}


def environment():
    return {"python": platform.python_version(), "platform": platform.platform(),
            "machine": platform.machine(), "numpy": np.__version__,
            "scipy": scipy.__version__, "sklearn": sklearn.__version__,
            "torch": torch.__version__, "torch_build": torch.__config__.show(),
            "device": "cpu", "dtype": "float32", "threads": torch.get_num_threads(),
            "deterministic": torch.are_deterministic_algorithms_enabled()}


def run(config, out=None):
    configure()
    if out is not None:
        out.mkdir(parents=True, exist_ok=False)
        (out / "config.json").write_text(json.dumps(config, indent=2) + "\n")
    start = time.perf_counter()
    x, y, splits, raw_hash = load_data(config)
    tr, sel = splits["train"], splits["selection"]
    results, artifacts = [], {}
    for seed in config["seeds"]:
        model, row = train_one(config, seed, x[tr], y[tr], x[sel], y[sel])
        row["logits"] = {}
        model.eval()
        with torch.no_grad():
            for role in ["selection", "calibration", "test"]:
                row["logits"][role] = model(x[splits[role]]).tolist()
        cal = row["logits"]["calibration"]
        cy = y[splits["calibration"]].tolist()
        row["temperature"] = fit_temperature(cal, cy, config["temperature_bounds"])
        row["metrics"] = {}
        for role in ["calibration", "test"]:
            labels = y[splits[role]].tolist()
            row["metrics"][role] = {
                stage: metrics(row["logits"][role], labels, t, config["ece_bins"])
                for stage, t in [("before", 1.0), ("after", row["temperature"])]}
        if out is not None:
            name = f"seed-{seed}.pt"
            torch.save(model.state_dict(), out / name)
            artifacts[name] = file_hash(out / name)
        results.append(row)
    numeric = {"splits": {k: v.tolist() for k, v in splits.items()},
               "labels": {k: y[v].tolist() for k, v in splits.items()},
               "results": results}
    report = {"config": config, "config_sha256": digest(config),
              "source_sha256": file_hash(__file__), "data_sha256": raw_hash,
              "environment": environment(), "parameters": sum(p.numel() for p in model.parameters()),
              "numeric_sha256": digest(numeric), "numeric": numeric,
              "artifacts": artifacts, "seconds": time.perf_counter() - start}
    if out is not None:
        (out / "report.json").write_text(json.dumps(report, indent=2) + "\n")
    return report


def verify(folder):
    """只加载本例生成的可信 state_dict；重新计算分数，拒绝静默不一致。"""
    configure()
    report = json.loads((folder / "report.json").read_text())
    config = json.loads((folder / "config.json").read_text())
    if config != report["config"] or digest(config) != report["config_sha256"]:
        raise ValueError("配置不一致")
    if file_hash(__file__) != report["source_sha256"]:
        raise ValueError("源代码已变化，请保留原版本再核对")
    if digest(report["numeric"]) != report["numeric_sha256"]:
        raise ValueError("数值记录已变化")
    x, y, splits, raw_hash = load_data(config)
    if raw_hash != report["data_sha256"]:
        raise ValueError("数据版本不一致")
    for role, ids in splits.items():
        if (ids.tolist() != report["numeric"]["splits"][role] or
                y[ids].tolist() != report["numeric"]["labels"][role]):
            raise ValueError("数据划分或标签不一致")
    for row in report["numeric"]["results"]:
        name = f"seed-{row['seed']}.pt"
        if file_hash(folder / name) != report["artifacts"][name]:
            raise ValueError("模型文件已变化")
        state = torch.load(folder / name, map_location="cpu", weights_only=True)
        if state_hash(state) != row["weights_sha256"]:
            raise ValueError("权重内容不一致")
        model = make_model(config)
        model.load_state_dict(state)
        model.eval()
        with torch.no_grad():
            for role in ["selection", "calibration", "test"]:
                if model(x[splits[role]]).tolist() != row["logits"][role]:
                    raise ValueError("重新推理的分数不完全一致；请检查环境")
        cal = row["logits"]["calibration"]
        t = fit_temperature(cal, y[splits["calibration"]].tolist(),
                            config["temperature_bounds"])
        if t != row["temperature"]:
            raise ValueError("重新拟合的温度不一致")
        for role in ["calibration", "test"]:
            for stage, value in [("before", 1.0), ("after", t)]:
                actual = metrics(row["logits"][role], y[splits[role]].tolist(),
                                 value, config["ece_bins"])
                if actual != row["metrics"][role][stage]:
                    raise ValueError("重新评分不一致")
    return report


def show(report):
    n = report["numeric"]
    print("样本数：" + " / ".join(f"{k}={len(n['splits'][k])}"
                              for k in ["train", "selection", "calibration", "test"]))
    print(f"MLP 参数量：{report['parameters']}；每个种子训练 {report['config']['epochs']} 轮")
    for row in n["results"]:
        a, b = row["metrics"]["test"]["before"], row["metrics"]["test"]["after"]
        print(f"seed={row['seed']}：选第 {row['selected_epoch']} 轮，T={row['temperature']:.3f}；"
              f"准确率 {a['accuracy']:.3f} → {b['accuracy']:.3f}；"
              f"NLL {a['nll']:.4f} → {b['nll']:.4f}；ECE {a['ece']:.4f} → {b['ece']:.4f}")
    print("温度只使用校准标签；三个种子全部保留。")


def main():
    parser = ArgumentParser(description=__doc__)
    parser.add_argument("--config", type=Path, default=DEFAULT_CONFIG)
    group = parser.add_mutually_exclusive_group()
    group.add_argument("--out", type=Path, help="保存到尚不存在的目录")
    group.add_argument("--verify", type=Path, help="从本例输出目录加载并核对")
    args = parser.parse_args()
    if args.verify:
        report = verify(args.verify)
        print("核对通过：数据、划分、权重、逐条分数、温度和指标一致。")
    else:
        config = json.loads(args.config.read_text())
        report = run(config, args.out)
    show(report)


if __name__ == "__main__":
    main()