"""线性头测试（T-G4）：纯 Python 推理与黄金向量一致 / 工件缺失 ArtifactMissing。"""
import json
import math
import tempfile
from pathlib import Path

import pytest

from gateway.sense.classifier import LinearHead
from gateway.sense.errors import ArtifactMissing

W = [[0.5, -0.5, 0.0], [0.0, 0.5, -0.5], [-0.5, 0.0, 0.5]]
B = [0.1, 0.0, -0.1]


def _head(tmp_path):
    p = tmp_path / "head.json"
    p.write_text(json.dumps({"version": "test", "dim": 3, "labels": ["t1", "t2", "t3"],
                             "weights": W, "bias": B}), encoding="utf-8")
    return p


def test_predict_matches_reference_softmax(tmp_path):
    """纯 Python 点积 + softmax 与手工参考实现一致（黄金向量）。"""
    head = LinearHead.load(_head(tmp_path))
    vec = [1.0, 2.0, 3.0]
    logits = [sum(wi * vi for wi, vi in zip(w, vec)) + b for w, b in zip(W, B)]
    m = max(logits)
    exps = [math.exp(z - m) for z in logits]
    total = sum(exps)
    expected = {"t1": exps[0] / total, "t2": exps[1] / total, "t3": exps[2] / total}
    probs = head.predict(vec)
    assert sum(probs.values()) == pytest.approx(1.0, abs=1e-9)
    for k, v in expected.items():
        assert probs[k] == pytest.approx(v, rel=1e-9)


def test_predict_argmax_consistent_with_logit_order(tmp_path):
    head = LinearHead.load(_head(tmp_path))
    # [3,1,2] 点积最大维度是 0 -> t1 概率最大
    probs = head.predict([3.0, 1.0, 2.0])
    assert max(probs, key=probs.get) == "t1"


def test_load_missing_raises_artifact_missing(tmp_path):
    with pytest.raises(ArtifactMissing):
        LinearHead.load(tmp_path / "ghost.json")


def test_load_corrupted_raises_artifact_missing(tmp_path):
    p = tmp_path / "bad.json"
    p.write_text('{"weights": [], "bias": []}', encoding="utf-8")
    with pytest.raises(ArtifactMissing):
        LinearHead.load(p)


def test_train_synthetic_roundtrip(tmp_path):
    """T-G4 验收：合成集训练 macro-F1 > 0.9，head.json 可被 LinearHead 推理。"""
    import importlib.util
    import numpy as np

    script = Path(__file__).resolve().parent.parent / "scripts" / "train_tier_head.py"
    spec = importlib.util.spec_from_file_location("train_tier_head", script)
    mod = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(mod)

    rng = np.random.default_rng(7)
    n, d = 150, 32
    centers = [rng.normal(0, 1, d), rng.normal(5, 1, d), rng.normal(-5, 1, d)]
    X = np.vstack([c + rng.normal(0, 0.7, (n, d)) for c in centers])
    y = np.repeat([0, 1, 2], n)
    W, b = mod.train(X, y, epochs=150)
    m = mod.metrics(X, y, W, b)
    assert m["macro_f1"] > 0.9, m

    out = tmp_path / "head.json"
    out.write_text(json.dumps({"version": "synthetic", "dim": d,
                               "labels": ["t1", "t2", "t3"],
                               "weights": W.T.tolist(), "bias": b.tolist()}),
                   encoding="utf-8")
    head = LinearHead.load(out)
    probs = head.predict(X[0].tolist())
    assert sum(probs.values()) == pytest.approx(1.0, abs=1e-9)