"""第 44 章：把温度缩放公式变成可检查的计算；数据全为手工构造。"""

from argparse import ArgumentParser
from hashlib import sha256
import json
from math import ceil, exp, isfinite, log, sqrt
from pathlib import Path
import platform


def scaled_logits(logits, temperature):
    if not isfinite(temperature) or temperature <= 0:
        raise ValueError("温度必须为有限正数")
    if len(logits) < 2 or not all(isfinite(z) for z in logits):
        raise ValueError("至少需要两个有限类别分数")
    # 先减最大值，避免所有分数同时很大时的指数溢出。
    maximum = max(logits)
    return [(z - maximum) / temperature for z in logits]


def probabilities(logits, temperature=1.0):
    values = [exp(z) for z in scaled_logits(logits, temperature)]
    total = sum(values)
    return [v / total for v in values]


def nll(logits, labels, temperature=1.0):
    if not logits or len(logits) != len(labels):
        raise ValueError("分数和标签必须非空、长度相同")
    losses = []
    for row, label in zip(logits, labels):
        if not isinstance(label, int) or not 0 <= label < len(row):
            raise ValueError("标签必须是有效类别编号")
        z = scaled_logits(row, temperature)
        losses.append(log(sum(exp(v) for v in z)) - z[label])
    return sum(losses) / len(losses)


def metrics(logits, labels, temperature=1.0, bins=15):
    loss = nll(logits, labels, temperature)
    if not isinstance(bins, int) or bins < 1:
        raise ValueError("分箱数必须为正整数")
    rows = []
    for row, label in zip(logits, labels):
        p = probabilities(row, temperature)
        prediction = max(range(len(row)), key=row.__getitem__)
        confidence = p[prediction]
        # 第 m 格为 (m/M, (m+1)/M]，m 从 0 开始。
        bucket = min(bins - 1, max(0, ceil(confidence * bins) - 1))
        rows.append({"prediction": prediction, "label": label,
                     "confidence": confidence, "correct": prediction == label,
                     "bucket": bucket})
    grouped = []
    for bucket in range(bins):
        selected = [r for r in rows if r["bucket"] == bucket]
        if selected:
            grouped.append({"bin": bucket, "count": len(selected),
                            "confidence": sum(r["confidence"] for r in selected) / len(selected),
                            "accuracy": sum(r["correct"] for r in selected) / len(selected)})
    return {"nll": loss, "accuracy": sum(r["correct"] for r in rows) / len(rows),
            "confidence": sum(r["confidence"] for r in rows) / len(rows),
            "ece": sum(r["count"] * abs(r["accuracy"] - r["confidence"])
                       for r in grouped) / len(rows),
            "bins": grouped, "rows": rows}


def fit_temperature(logits, labels, bounds=(0.05, 20.0)):
    """在预先给定的正温度区间，用黄金分割搜索平均 NLL 的最低点。"""
    lower, upper = bounds
    if not (isfinite(lower) and isfinite(upper) and 0 < lower < upper):
        raise ValueError("温度边界必须为递增的有限正数")
    left, right = log(lower), log(upper)
    ratio = (sqrt(5) - 1) / 2
    a = right - ratio * (right - left)
    b = left + ratio * (right - left)
    fa, fb = nll(logits, labels, exp(a)), nll(logits, labels, exp(b))
    for _ in range(80):
        if fa < fb:
            right, b, fb = b, a, fa
            a = right - ratio * (right - left)
            fa = nll(logits, labels, exp(a))
        else:
            left, a, fa = a, b, fb
            b = left + ratio * (right - left)
            fb = nll(logits, labels, exp(b))
    candidates = [lower, exp((left + right) / 2), upper]
    return min(candidates, key=lambda t: nll(logits, labels, t))


def run():
    logits = [[log(9), 0.0] for _ in range(10)]
    calibration_labels = [0] * 8 + [1] * 2
    counter_labels = [0] * 9 + [1]
    temperature = fit_temperature(logits, calibration_labels)
    analytic = log(9) / log(4)
    assert abs(temperature - analytic) < 1e-6
    temperatures = [0.5 + i * 0.025 for i in range(141)]
    return {
        "source_sha256": sha256(Path(__file__).read_bytes()).hexdigest(),
        "python": platform.python_version(),
        "data_kind": "hand_constructed_logits_and_labels",
        "logits": logits,
        "calibration_ids": [f"cal-{i}" for i in range(10)],
        "counterexample_ids": [f"counter-{i}" for i in range(10)],
        "calibration_labels": calibration_labels, "counterexample_labels": counter_labels,
        "temperature": temperature, "analytic_temperature": analytic,
        "bounds": [0.05, 20.0], "bins": 15,
        "calibration": {"before": metrics(logits, calibration_labels),
                        "after": metrics(logits, calibration_labels, temperature)},
        "counterexample": {"before": metrics(logits, counter_labels),
                           "after": metrics(logits, counter_labels, temperature)},
        "three_class": [{"temperature": t, "probabilities": probabilities([2, 1, 0], t)}
                        for t in [0.5, 1.0, 2.0]],
        "curve": [{"temperature": t,
                   "calibration_nll": nll(logits, calibration_labels, t),
                   "counterexample_nll": nll(logits, counter_labels, t)}
                  for t in temperatures],
    }


def show(report):
    for row in report["three_class"]:
        formatted = ", ".join(f"{p:.3f}" for p in row["probabilities"])
        print(f"T={row['temperature']:.1f}：三类概率=[{formatted}]；最高分类别=0")
    print(f"十条手工校准样本：拟合 T={report['temperature']:.6f}")
    for key, name in [("calibration", "校准组"), ("counterexample", "另构造的反例组")]:
        before, after = report[key]["before"], report[key]["after"]
        print(f"{name}：准确率 {before['accuracy']:.3f} → {after['accuracy']:.3f}；"
              f"NLL {before['nll']:.4f} → {after['nll']:.4f}；"
              f"ECE {before['ece']:.3f} → {after['ece']:.3f}")
    print("这是手工分数的公式实验，没有训练分类器，也没有复现论文的基准成绩。")


if __name__ == "__main__":
    parser = ArgumentParser(description=__doc__)
    parser.add_argument("--out", type=Path)
    args = parser.parse_args()
    result = run()
    show(result)
    if args.out:
        args.out.parent.mkdir(parents=True, exist_ok=True)
        args.out.write_text(json.dumps(result, ensure_ascii=False, indent=2) + "\n")