"""Extract null margin (options.null_margin): request parsing, the decode rule, server
default, determinism / batch invariance on a tiny model, and the calibration tooling
(deg.bench.null_margin)."""

from __future__ import annotations

import dataclasses
import json
import math
from types import SimpleNamespace

import pytest

from deg.bench import null_margin as nm
from deg.request import FillOptions, RequestError, parse_fill_request
from deg.schema import compile_schema

torch = pytest.importorskip("torch")

BASE = {"context": "Order AB-1 shipped.", "schema": {
    "type": "object", "properties": {"x": {"type": ["string", "null"]}}, "required": ["x"]}}


def test_option_parses_and_defaults_to_none():
    assert parse_fill_request(BASE).options.null_margin is None
    for v in (0, 1.5, -2, 0.0):
        assert parse_fill_request({**BASE, "options": {"null_margin": v}}).options.null_margin == v
    for bad in (True, "1", float("inf"), [1]):
        with pytest.raises(RequestError):
            parse_fill_request({**BASE, "options": {"null_margin": bad}})


def test_decision_api_accepts_null_margin():
    from deg.api import REQUEST_OPTION_KEYS

    assert "null_margin" in REQUEST_OPTION_KEYS


# --------------------------------------------------------------------- decode rule

CTX = "alpha beta gamma delta"
OFFSETS = [(0, 5), (5, 10), (10, 16), (16, 22)]  # " beta" etc.: trim_span strips spaces


def _logits(span_lp: float, null_lp: float) -> tuple[torch.Tensor, torch.Tensor]:
    """start/end logits [T+1] whose best span (token 1..1 = "beta") has log-prob ~span_lp
    and the null outcome ~null_lp (the rest share what is left)."""
    def row(a: float, b: float) -> torch.Tensor:
        p = torch.full((5,), 1e-6, dtype=torch.float64)
        p[1], p[4] = math.exp(a), math.exp(b)
        p[0] = p[2] = p[3] = max(1e-6, (1 - p[1] - p[4]) / 3)
        return p.log()
    return row(span_lp / 2, null_lp / 2), row(span_lp / 2, null_lp / 2)


def _extract(schema_type: object, span_lp: float, null_lp: float, margin: float | None,
             required: bool = True):
    from deg.infer import extract_result

    schema = {"type": "object", "properties": {"x": {"type": schema_type}},
              "required": ["x"] if required else []}
    fld = compile_schema(schema).fields[0]
    start, end = _logits(span_lp, null_lp)
    return extract_result(fld, CTX, SimpleNamespace(offsets=OFFSETS), start, end,
                          torch.tensor(1.0, dtype=torch.float64),
                          FillOptions(null_margin=margin))


def test_margin_shifts_the_null_decision():
    # span log-prob -1.0 > null -1.6: margin 0 / None take the span
    assert _extract(["string", "null"], -1.0, -1.6, None).value == "beta"
    assert _extract(["string", "null"], -1.0, -1.6, 0.0).value == "beta"
    assert _extract(["string", "null"], -1.0, -1.6, 0.5).value == "beta"
    # null wins once span < null + margin
    assert _extract(["string", "null"], -1.0, -1.6, 0.7).value is None
    # a negative margin keeps a span null would beat at 0
    assert _extract(["string", "null"], -1.6, -1.0, 0.0).value is None
    assert _extract(["string", "null"], -1.6, -1.0, -0.7).value == "beta"


def test_margin_omits_optional_and_never_touches_forbid():
    omitted = _extract("string", -1.0, -1.6, 2.0, required=False)
    assert omitted.value is None and omitted.omit
    assert _extract("string", -1.0, -1.6, 50.0, required=True).value == "beta"  # forbid


def test_server_default_applies_only_when_unset():
    from deg.server import with_null_margin

    assert with_null_margin(FillOptions(), 1.5).null_margin == 1.5
    assert with_null_margin(FillOptions(null_margin=0.0), 1.5).null_margin == 0.0
    assert with_null_margin(FillOptions(null_margin=-1), 1.5).null_margin == -1


def test_bind_model_passes_default(monkeypatch):
    from deg import infer, server

    seen = []
    monkeypatch.setattr(infer, "fill", lambda m, c, s, o: seen.append(o.null_margin) or {})
    monkeypatch.setattr(infer, "fill_many",
                        lambda m, jobs: [seen.append(o.null_margin) for _, _, o in jobs])
    fill, many = server.bind_model(object(), null_margin=2.0)
    fill("c", {}, FillOptions())
    fill("c", {}, FillOptions(null_margin=0.5))
    many([("c", {}, FillOptions()), ("c", {}, FillOptions(null_margin=0.0))])
    assert seen == [2.0, 0.5, 2.0, 0.0]


# ------------------------------------------------------------ tiny model (end to end)


@pytest.fixture(scope="module")
def model():
    pytest.importorskip("transformers")
    from test_model import SNAPSHOT, tiny_lm

    from deg.model import load_model, load_tokenizer

    if SNAPSHOT is None:
        pytest.skip("Qwen3.5-0.8B tokenizer not cached")
    m = load_model(None, device="cpu", dtype=torch.float32, lm=tiny_lm(),
                   tokenizer=load_tokenizer(SNAPSHOT))
    torch.manual_seed(3)
    with torch.no_grad():
        for p in m.heads.parameters():
            p.normal_(0.0, 0.5)
    return m


QSCHEMA = {"type": "object", "properties": {
    f"q{i}": {"type": ["string", "null"], "description": d} for i, d in enumerate(
        ["Who signed?", "Where?", "When was it paid?", "What is the color of the car?",
         "Which team won?", "What was the amount?"])},
    "required": [f"q{i}" for i in range(6)]}
QCTX = ("The contract was signed by Jane Roe in Paris on 3 May 2024. The fee of 400 euros "
        "was paid by Globex. The Red team beat the Blue team in the final.")


def test_fill_margin_monotone_deterministic_and_batch_invariant(model):
    from deg.infer import fill, fill_many

    outs = {m: fill(model, QCTX, QSCHEMA, FillOptions(null_margin=m))["value"]
            for m in (-50.0, -1.0, 0.0, 1.0, 50.0)}
    nulls = [sum(v is None for v in outs[m].values()) for m in sorted(outs)]
    assert nulls == sorted(nulls)  # larger margin -> never fewer nulls
    assert nulls[-1] == 6 and nulls[0] < 6
    assert fill(model, QCTX, QSCHEMA)["value"] == outs[0.0]  # None = 0
    assert fill(model, QCTX, QSCHEMA, FillOptions(null_margin=1.0))["value"] == outs[1.0]
    batched = fill_many(model, [(QCTX, QSCHEMA, FillOptions(null_margin=m)) for m in sorted(outs)])
    assert [b["value"] for b in batched] == [outs[m] for m in sorted(outs)]


def test_leaf_probe_reproduces_margin0(model):
    from deg.infer import fill

    with nm.LeafProbe() as probe:
        out = fill(model, QCTX, QSCHEMA, FillOptions())["value"]
    assert [r["ptr"] for r in probe.rows] == [f"/q{i}" for i in range(6)]
    leaves = [{**r, "gold": "ans", "alts": []} for r in probe.rows]
    assert nm.check_margin0(leaves) == 0
    for r in probe.rows:
        assert r["value"] == out[r["ptr"][1:]]
        # the sweep rule at any margin = the real decode at that margin
    for m in (-1.0, 1.0):
        real = fill(model, QCTX, QSCHEMA, FillOptions(null_margin=m))["value"]
        assert {r["ptr"][1:]: (None if nm.is_null(r, m) else r["span_value"])
                for r in probe.rows} == real


# --------------------------------------------------------------- calibration tooling


def _leaf(gold: str, span_lp: float, null_lp: float, value: object = "x",
          alts: list | None = None) -> dict:
    return {"ptr": "/q", "gold": gold, "alts": alts if alts is not None else (
        ["x"] if gold == "ans" else []), "span_value": value, "span_lp": span_lp,
        "null_lp": null_lp, "value": value if span_lp >= null_lp else None}


def test_point_rates_and_cost():
    leaves = [_leaf("ans", -0.1, -3.0), _leaf("ans", -1.0, -1.5, "y"),
              _leaf("unans", -0.5, -1.0), _leaf("unans", -2.0, -1.0)]
    p = nm.point(leaves, 0.0)
    assert p["made_up_rate"] == 0.5 and p["omission_rate"] == 0.0
    assert p["answerable_correct"] == 0.5 and p["answerable_wrong_span"] == 0.5
    assert p["cost"] == pytest.approx(1.0)
    q = nm.point(leaves, 1.0)  # drops the -0.5/-1.0 made-up one and the -1.0/-1.5 answer
    assert q["made_up_rate"] == 0.0 and q["omission_rate"] == 0.5 and q["cost"] == 0.5


def test_choose_refuses_always_null_collapse():
    def pt(m, cost, omit):
        return {"margin": m, "cost": cost, "sets": {"a": {"omission_rate": omit}}}
    curve = [pt(0.0, 0.8, 0.2), pt(1.0, 0.6, 0.3), pt(9.0, 0.5, 0.9)]
    assert nm.choose(curve)["margin"] == 1.0
    assert nm.choose([pt(-1.0, 0.6, 0.1), pt(1.0, 0.6, 0.1), pt(0.5, 0.6, 0.1)])["margin"] == 0.5


def test_squad_records_exclude_suite_contexts():
    rows = [{"context": f"Paragraph {i} text.", "id": f"{i}-{j}", "question": f"Q{j}?",
             "title": "T", "answers": ["text"] if j % 2 else []}
            for i in range(5) for j in range(4)]
    recs, skipped = nm.squad_records(rows, {"Paragraph 2 text."}, None, 10)
    assert len(recs) == 4 and skipped["suite_context"] == 1
    assert all(r["context"] != "Paragraph 2 text." for r in recs)
    r = recs[0]
    assert set(r["schema"]["properties"]) == {"q1", "q2", "q3", "q4"}
    assert sum(v is None for v in r["value"].values()) == 2
    assert nm.leaf_gold(r, "/q1")[0] == ("ans" if r["alternatives"]["/q1"] else "unans")


def test_leaf_gold_nested_and_regold(tmp_path):
    rec = {"uid": "r", "value": {"a": {"b": "x", "c": None}, "d": "y"}}
    assert nm.leaf_gold(rec, "/a/b") == ("ans", ["x"])
    assert nm.leaf_gold(rec, "/a/c") == ("unans", [])
    assert nm.leaf_gold(rec, "/a/zz") == ("unans", [])
    assert nm.leaf_gold(rec, "/d") == ("ans", ["y"])
    nm.write_jsonl(tmp_path / "h.jsonl", [rec])
    rows = [{"uid": "r", "error": None, "leaves": [{"ptr": "/a/b", "gold": "unans"}]},
            {"uid": "r", "error": "boom", "leaves": [{"ptr": "/d"}]}]
    assert nm.regold(rows, tmp_path / "h.jsonl") == [{"ptr": "/a/b", "gold": "ans",
                                                      "alts": ["x"]}]


def test_extract_null_rates():
    from deg.bench.score import Prediction

    recs = [{"uid": "a", "schema": {"type": "object", "properties": {
        "p": {"type": ["string", "null"]}, "q": {"type": ["string", "null"]},
        "r": {"type": "string"}, "s": {"type": "string"}}, "required": ["p", "q", "r"]},
        "value": {"p": "x", "q": None, "r": "z"}}]
    preds = {"a": Prediction("a", value={"p": None, "q": "made up", "r": "z", "s": "w"})}
    out = nm.extract_null_rates(recs, preds)
    # p answerable->null (wrong null); q, s(optional, absent gold) unanswerable->answered;
    # r is forbid (required, non-nullable): not counted
    assert out == {"unanswerable_n": 2, "answerable_n": 1, "made_up_rate": 1.0,
                   "wrong_null_rate": 1.0}


def test_suite_label_gets_margin_suffix(monkeypatch):
    from deg.bench.systems import deg as ds

    class M:
        generate_with, meta = "base", {}

    monkeypatch.setattr("deg.model.load_model", lambda *a, **k: M())
    monkeypatch.setattr(ds.DegSystem, "_warm", lambda self: None)
    import pathlib
    monkeypatch.setattr(pathlib.Path, "exists", lambda self: True)
    s = ds.DegSystem("/x/run/step-1", null_margin=1.5)
    assert s.name.endswith("-basegen-nm1.5") and s.options.null_margin == 1.5
    assert s.info["null_margin"] == 1.5
    assert ds.DegSystem("/x/run/step-1").options == FillOptions()


def test_dataclass_replace_keeps_other_options():
    o = FillOptions(array_decode="hybrid", seed=3)
    r = dataclasses.replace(o, null_margin=1.0)
    assert (r.array_decode, r.seed, r.null_margin) == ("hybrid", 3, 1.0)
    json.dumps(dataclasses.asdict(r))