"""``python -m deg.evaluate calibrate``: fits temperatures on val + jev dev and writes a
calibrated copy of the checkpoint without touching the source (tiny random Qwen3.5, CPU)."""

from __future__ import annotations

import hashlib
import json
from pathlib import Path

import pytest

torch = pytest.importorskip("torch")
pytest.importorskip("transformers")

from test_model import CONTEXT, SNAPSHOT, _sub, tiny_lm  # noqa: E402

from deg import evaluate  # noqa: E402
from deg.model import KINDS, load_model, load_tokenizer  # noqa: E402


def _digest(path):
    return {p.name: hashlib.sha256(p.read_bytes()).hexdigest() for p in sorted(path.iterdir())}


@pytest.mark.skipif(SNAPSHOT is None, reason="Qwen3.5-0.8B tokenizer not cached")
def test_calibrate_writes_calibrated_copy_and_keeps_source(tmp_path, monkeypatch):
    from safetensors.torch import load_file

    tok = load_tokenizer(SNAPSHOT)
    model = load_model(None, device="cpu", dtype=torch.float32, lm=tiny_lm(), tokenizer=tok)
    torch.manual_seed(1)
    with torch.no_grad():
        for p in model.heads.parameters():
            p.normal_(0.0, 0.5)
    src = tmp_path / "step-0000010"
    model.save(src, {"temperature": dict.fromkeys(KINDS, 1.0)})
    before = _digest(src)

    a = CONTEXT.index("AB-1234")
    rec = {"uid": "v", "split": "val", "schema": _sub(["department", "urgent", "order_id"]),
           "context": CONTEXT,
           "value": {"department": "billing", "urgent": True, "order_id": "AB-1234"},
           "spans": {"/order_id": [a, a + 7]}, "span_alternatives": {}, "span_probs": {},
           "decide_probs": {"/department": [0.8, 0.15, 0.05]}}
    records = tmp_path / "records.jsonl"
    records.write_text("".join(json.dumps({**rec, "uid": f"v{i}"}) + "\n" for i in range(12))
                       + json.dumps({**rec, "uid": "t", "split": "train"}) + "\n")

    loaded = []

    def fake_load(path, device="cuda"):
        m = load_model(path, device="cpu", dtype=torch.float32, lm=tiny_lm(), tokenizer=tok)
        loaded.append(m)
        return m

    monkeypatch.setattr(evaluate, "load_model", fake_load)
    out = tmp_path / "final"
    evaluate.main(["calibrate", "--ckpt", str(src), "--out", str(out), "--records",
                   str(records), "--jev-dev", str(tmp_path / "missing.jsonl"), "--n-jev", "0"])

    assert _digest(src) == before  # source untouched
    meta = json.loads((out / "deg.json").read_text())
    cal = meta["calibration"]
    assert cal["source"] == str(src) and cal["val_records"] == 12 and "jev_records" not in cal
    assert cal["fit_source"] == "datagen-val" and cal["fit_on"] == ["datagen_val"]
    temps = cal["temperature"]
    assert temps["choice"] != 1.0 or temps["noul"] != 1.0 or temps["extract"] != 1.0
    assert temps["count"] == 1.0  # never fit
    rep = cal["report"]["datagen_val"]
    assert rep["choice"]["n"] == 12 and rep["choice"]["t_after"] == temps["choice"]
    assert rep["decide"]["ece_after"] is not None and "jev_dev" not in cal["report"]
    assert meta["temperature"] == pytest.approx({k: temps.get(k, 1.0) for k in KINDS})
    # every tensor but the temperatures is byte-identical to the source
    backbone = "backbone.safetensors"
    assert (out / backbone).read_bytes() == (src / backbone).read_bytes()
    h_src, h_out = (load_file(str(p / "heads.safetensors")) for p in (src, out))
    assert set(h_src) == set(h_out)
    for k in h_src:
        if k != "temperature":
            assert torch.equal(h_src[k], h_out[k]), k
    re = fake_load(out)
    assert re.heads.temperature.tolist() == pytest.approx([meta["temperature"][k] for k in KINDS])
    assert re.array_mode == loaded[0].array_mode

    with pytest.raises(ValueError):
        evaluate.write_calibrated(loaded[0], src, src, {})


@pytest.mark.skipif(SNAPSHOT is None, reason="Qwen3.5-0.8B tokenizer not cached")
def test_calibrate_fits_datagen_val_only_and_reports_jev_separately(tmp_path):
    tok = load_tokenizer(SNAPSHOT)

    def fresh():
        m = load_model(None, device="cpu", dtype=torch.float32, lm=tiny_lm(), tokenizer=tok)
        torch.manual_seed(1)
        with torch.no_grad():
            for p in m.heads.parameters():
                p.normal_(0.0, 0.5)
        return m

    a = CONTEXT.index("AB-1234")
    rec = {"uid": "v", "split": "val", "schema": _sub(["department", "urgent", "order_id"]),
           "context": CONTEXT,
           "value": {"department": "billing", "urgent": True, "order_id": "AB-1234"},
           "spans": {"/order_id": [a, a + 7]}, "span_alternatives": {}, "span_probs": {},
           "decide_probs": {}}
    records = tmp_path / "records.jsonl"
    records.write_text("".join(json.dumps({**rec, "uid": f"v{i}"}) + "\n" for i in range(12)))
    jev = tmp_path / "dev.jsonl"
    q_choice = {"type": "choice", "instructions": "Which team?",
                "criteria": {"billing": "money", "technical": "bugs", "sales": "buying"},
                "label": {"probs": [0.05, 0.9, 0.05]}}
    q_noul = {"type": "noul", "instructions": "Is it about shipping?",
              "criteria": {"true": "shipping", "false": "other"}, "label": {"hard": 0}}
    jev.write_text("".join(json.dumps({"uid": f"j{i}", "state": CONTEXT,
                                       "questions": [q_choice, q_noul]}) + "\n"
                           for i in range(15)))

    info_val: dict = {}
    t_val = evaluate.calibrate(fresh(), [records], jev, 15, info=info_val)
    info_rep: dict = {}
    t_rep = evaluate.calibrate(fresh(), [records], jev, 15, info=info_rep, report_jev=True)
    assert t_rep == t_val  # reporting jev dev never changes the fit
    assert info_val["fit_on"] == ["datagen_val"] and "jev_dev" not in info_val["report"]
    assert set(info_rep["report"]) == {"datagen_val", "jev_dev"}
    jr = info_rep["report"]["jev_dev"]
    assert jr["choice"]["n"] == 15 and jr["noul"]["n"] == 15
    assert jr["choice"]["t_after"] == t_val["choice"]  # reported at the val-fit temperature
    assert info_rep["rows"]["choice"] == 12  # the fit pool: datagen val rows only

    info_both: dict = {}
    evaluate.calibrate(fresh(), [records], jev, 15, info=info_both,
                       source="datagen-val+jev-dev")
    assert info_both["fit_on"] == ["datagen_val", "jev_dev"]
    assert info_both["rows"]["choice"] == 27
    with pytest.raises(ValueError):
        evaluate.calibrate(fresh(), [records], jev, 15, source="jev")


def test_hard_target_nll_scores_the_filled_value():
    from types import SimpleNamespace

    logits = torch.tensor([2.0, 0.5, -1.0])
    row = (logits, SimpleNamespace(target=[0.6, 0.3, 0.1]))  # soft; argmax = filled value 0
    logp = torch.log_softmax(logits / 1.5, -1)
    assert evaluate._nll("choice", [row], 1.5, hard=True) == pytest.approx(-float(logp[0]))
    soft = -float((torch.tensor([0.6, 0.3, 0.1]) * logp).sum())
    assert evaluate._nll("choice", [row], 1.5) == pytest.approx(soft)
    # a correct, under-confident head: the hard fit sharpens (T < 1), the soft one does not
    best_hard = min(evaluate.TEMPS, key=lambda t: evaluate._nll("choice", [row], t, True))
    best_soft = min(evaluate.TEMPS, key=lambda t: evaluate._nll("choice", [row], t))
    assert best_hard < 1.0 <= best_soft
    with pytest.raises(ValueError):
        evaluate.calibrate(None, [], Path("x"), 0, target="medium")