"""Extract null-margin calibration (EVAL ONLY): choose ``options.null_margin`` on held-out data.

A nullable/optional extract leaf is null when ``best_span_logprob < null_logprob + margin``
(:func:`deg.infer.extract_result`). The user's rule: a made-up answer costs twice an
omission, so the operating point minimises

    cost(m) = 2 * made_up_rate(unanswerable leaves) + omission_rate(answerable leaves)

on HELD-OUT data (never the deg-suite v1 test items), and must not collapse to
always-null. Steps:

    # held-out sets: SQuAD v2 dev paragraphs NOT in suite-v1 (suite contexts excluded and
    # SuiteGuard-checked) + datagen val-split records (full-off, extract-real-v1, v3-long)
    python -m deg.bench.null_margin build --out results/deg/null-margin-v1
    # one model pass per set: per top-level nullable/optional extract leaf, the best valid
    # span (value, log-prob) and the null log-prob
    CUDA_VISIBLE_DEVICES=<uuid> python -m deg.bench.null_margin dump --ckpt <ckpt> \\
        --label deg2b-1-step-0004000 --out results/deg/null-margin-v1
    # offline sweep -> curve + chosen margin per label
    python -m deg.bench.null_margin sweep --out results/deg/null-margin-v1
    # made-up / wrong-null rates of suite-v1 prediction dirs (before/after)
    python -m deg.bench.null_margin rates --system <suite label> [...]

Only top-level leaves are dumped (their pointer is unique per record); the margin itself
applies to every nullable/optional extract leaf, also inside array items. The dump probes
each leaf at margins -inf and +inf beside the real (margin 0) call, which decodes the same
head outputs, so the sweep is exact: at margin m a leaf is the best valid span when its
log-prob >= null + m, else null (checked against the margin-0 output for every leaf).
"""

from __future__ import annotations

import argparse
import dataclasses
import hashlib
import json
import math
from collections.abc import Iterable
from pathlib import Path
from typing import Any

OUT = Path("/shared/rwkv/results/deg/null-margin-v1")
SQUAD_DEV = Path("/data/deg/external/suite-v1/squad_v2/validation.parquet")
DATAGEN = Path("/data/deg/datagen")
DATAGEN_SOURCES = ("full-off", "extract-real-v1", "v3-long")
SQUAD_Q_PER_PARA = 4
HELDOUT_SQUAD_PARAS = 300
BIG = 1e30  # a probe margin: never null (-BIG) / always null (+BIG)
GRID = tuple(round(-3 + 0.1 * i, 2) for i in range(131))  # -3.0 .. 10.0 nats
HALLUCINATION_COST = 2.0
MAX_OMISSION = 0.5  # an operating point that drops half the answerable leaves is a collapse


def sha1(s: str) -> str:
    return hashlib.sha1(s.encode()).hexdigest()


def write_jsonl(path: Path, rows: Iterable[dict[str, Any]]) -> int:
    path.parent.mkdir(parents=True, exist_ok=True)
    n = 0
    with path.open("w", encoding="utf-8") as fh:
        for r in rows:
            fh.write(json.dumps(r, ensure_ascii=False) + "\n")
            n += 1
    return n


def read_jsonl(path: Path) -> list[dict[str, Any]]:
    return [json.loads(x) for x in path.read_text(encoding="utf-8").split("\n") if x.strip()]


# ------------------------------------------------------------------------ held-out sets


def squad_records(rows: Iterable[dict[str, Any]], exclude_contexts: set[str], guard: Any,
                  limit: int) -> tuple[list[dict[str, Any]], dict[str, int]]:
    """SQuAD v2 rows (``context, id, question, title, answers``) -> suite-format records
    (one per paragraph, up to 4 questions as nullable string fields, sha1 order), skipping
    suite paragraphs and anything the suite guard flags."""
    paras: dict[str, list[dict[str, Any]]] = {}
    for row in rows:
        paras.setdefault(row["context"], []).append(row)
    out, skipped = [], {"suite_context": 0, "guard": 0, "few_questions": 0}
    for ctx in sorted(paras, key=sha1):
        if ctx in exclude_contexts:
            skipped["suite_context"] += 1
            continue
        if guard is not None and guard.match(ctx):
            skipped["guard"] += 1
            continue
        qs = sorted(paras[ctx], key=lambda q: sha1(q["id"]))[:SQUAD_Q_PER_PARA]
        if len(qs) < 3:
            skipped["few_questions"] += 1
            continue
        props, value, alts = {}, {}, {}
        for i, q in enumerate(qs):
            key = f"q{i + 1}"
            props[key] = {"type": ["string", "null"],
                          "description": q["question"] + " (copy the answer span; null if the "
                                                         "passage does not answer it)"}
            value[key] = q["answers"][0] if q["answers"] else None
            alts[f"/{key}"] = list(q["answers"])
        out.append({"uid": "squad2-heldout:" + sha1(ctx)[:16], "context": ctx,
                    "schema": {"type": "object", "properties": props, "required": list(props),
                               "additionalProperties": False},
                    "value": value, "alternatives": alts, "domain": f"squad2/{qs[0]['title']}",
                    "question_ids": [q["id"] for q in qs]})
        if len(out) == limit:
            break
    return out, skipped


def load_squad_dev(path: Path = SQUAD_DEV) -> list[dict[str, Any]]:
    import pandas as pd

    df = pd.read_parquet(path)
    return [{"context": r.context, "id": r.id, "question": r.question, "title": r.title,
             "answers": sorted({str(a) for a in r.answers["text"]},
                               key=list(r.answers["text"]).index)}
            for r in df.itertuples()]


def nullable_top_leaves(schema: Any) -> list[str]:
    """Pointers of extract leaves outside arrays (top-level or in nested objects) that may be
    null or omitted (on_null != forbid)."""
    from deg.schema import compile_schema

    return [f.pointer for f in compile_schema(schema).fields
            if f.mode.startswith("extract.") and "*" not in f.pointer and f.on_null != "forbid"]


def datagen_records(paths: dict[str, Path], guard: Any, jev_guard: Any,
                    exclude_uids: set[str]) -> tuple[list[dict[str, Any]], dict[str, int]]:
    """``split == "val"`` records with >= 1 nullable/optional top-level extract leaf, as
    bench records (``x-generate`` fields dropped: extract outputs do not depend on them)."""
    from deg.bench.score import drop_generate

    out, skipped = [], {"guard": 0, "jev_guard": 0, "suite_uid": 0, "no_nullable_leaf": 0}
    for name, path in paths.items():
        for line in path.open(encoding="utf-8"):
            r = json.loads(line)
            if r.get("split") != "val":
                continue
            if r["uid"] in exclude_uids:
                skipped["suite_uid"] += 1
                continue
            if guard is not None and guard.match(r["context"]):
                skipped["guard"] += 1
                continue
            if jev_guard is not None and jev_guard.match(r["context"]):
                skipped["jev_guard"] += 1
                continue
            schema = drop_generate(r["schema"])
            if not nullable_top_leaves(schema):
                skipped["no_nullable_leaf"] += 1
                continue
            out.append({"uid": f"{name}:{r['uid']}", "context": r["context"], "schema": schema,
                        "value": r["value"], "domain": r.get("domain")})
    return out, skipped


def suite_uids() -> set[str]:
    from deg.bench.suite import load_source

    return {r["uid"] for r in load_source("datagen-test")}


def build(out: Path, squad_paras: int = HELDOUT_SQUAD_PARAS) -> dict[str, Any]:
    from deg.bench.jevbench import ContaminationGuard
    from deg.bench.suite import SuiteGuard, load_source

    guard, jev = SuiteGuard(), ContaminationGuard()
    suite_ctx = {r["context"] for r in load_source("squad2")}
    suite_qids = {q for r in load_source("squad2") for q in r["question_ids"]}
    squad, sq_skip = squad_records(load_squad_dev(), suite_ctx, guard, squad_paras)
    assert not {q for r in squad for q in r["question_ids"]} & suite_qids
    dg, dg_skip = datagen_records({n: DATAGEN / n / "records.jsonl" for n in DATAGEN_SOURCES},
                                  guard, jev, suite_uids())
    manifest = {"squad2-heldout": {"records": write_jsonl(out / "heldout" / "squad2.jsonl", squad),
                                   "skipped": sq_skip, "suite_guard_texts": guard.loaded,
                                   "origin": "SQuAD v2 dev paragraphs not in deg-suite v1"},
                "datagen-val": {"records": write_jsonl(out / "heldout" / "datagen-val.jsonl", dg),
                                "skipped": dg_skip,
                                "origin": "split=val records of " + ", ".join(DATAGEN_SOURCES)}}
    for name in ("squad2", "datagen-val"):
        p = out / "heldout" / f"{name}.jsonl"
        manifest["squad2-heldout" if name == "squad2" else name]["sha256"] = hashlib.sha256(
            p.read_bytes()).hexdigest()
    (out / "heldout" / "manifest.json").write_text(json.dumps(manifest, indent=1) + "\n")
    return manifest


# ------------------------------------------------------------------------------- dump


class LeafProbe:
    """Wraps :func:`deg.infer.extract_result`: beside the real call, records for each
    top-level nullable/optional leaf the best valid span (margin -inf) and the null
    outcome (margin +inf)."""

    def __init__(self) -> None:
        self.rows: list[dict[str, Any]] = []
        self._orig: Any = None

    def __enter__(self) -> LeafProbe:
        from deg import infer

        self._orig = orig = infer.extract_result

        def probed(fld: Any, ctx_text: str, ctx: Any, start: Any, end: Any, temperature: Any,
                   options: Any) -> Any:
            res = orig(fld, ctx_text, ctx, start, end, temperature, options)
            if fld.on_null != "forbid" and "*" not in fld.pointer:
                best = orig(fld, ctx_text, ctx, start, end, temperature,
                            dataclasses.replace(options, null_margin=-BIG))
                null = orig(fld, ctx_text, ctx, start, end, temperature,
                            dataclasses.replace(options, null_margin=BIG))
                span_lp = (_log(best.meta["confidence"]) if best.value is not None
                           else -math.inf)
                self.rows.append({"ptr": fld.pointer, "value": res.value,
                                  "span_value": best.value, "span_lp": span_lp,
                                  "null_lp": _log(null.meta["confidence"])})
            return res

        infer.extract_result = probed
        return self

    def __exit__(self, *exc: Any) -> None:
        from deg import infer

        infer.extract_result = self._orig


def _log(p: float) -> float:
    return math.log(p) if p > 0 else -math.inf


def leaf_gold(rec: dict[str, Any], ptr: str) -> tuple[str, list[Any]]:
    """("unans", []) when the gold is null/absent, else ("ans", accepted answers)."""
    if "alternatives" in rec and ptr in rec["alternatives"]:
        alts = rec["alternatives"][ptr]
        return ("ans", alts) if alts else ("unans", [])
    gold: Any = rec["value"]
    for part in ptr.split("/")[1:]:  # nested object leaves (no arrays: top-level dumps)
        key = part.replace("~1", "/").replace("~0", "~")
        gold = gold.get(key) if isinstance(gold, dict) else None
    return ("unans", []) if gold is None else ("ans", [gold])


def dump(ckpt: str, label: str, out: Path, sets: list[str]) -> None:
    import torch

    from deg.infer import fill
    from deg.model import load_model
    from deg.request import FillOptions

    model = load_model(Path(ckpt), device="cuda")
    for name in sets:
        recs = read_jsonl(out / "heldout" / f"{name}.jsonl")
        path = out / "dumps" / label / f"{name}.jsonl"
        path.parent.mkdir(parents=True, exist_ok=True)
        done = {r["uid"] for r in read_jsonl(path)} if path.exists() else set()
        with path.open("a", encoding="utf-8") as fh, LeafProbe() as probe:
            for i, rec in enumerate(recs):
                if rec["uid"] in done:
                    continue
                probe.rows = []
                err = None
                try:
                    fill(model, rec["context"], rec["schema"], FillOptions())
                except Exception as exc:  # noqa: BLE001 - recorded, excluded from the sweep
                    err = f"{type(exc).__name__}: {str(exc)[:160]}"
                    torch.cuda.empty_cache()
                leaves = []
                for row in probe.rows:
                    kind, alts = leaf_gold(rec, row["ptr"])
                    leaves.append({**row, "gold": kind, "alts": alts})
                fh.write(json.dumps({"uid": rec["uid"], "error": err, "leaves": leaves},
                                    ensure_ascii=False) + "\n")
                fh.flush()
                if (i + 1) % 50 == 0:
                    print(json.dumps({"label": label, "set": name, "done": i + 1,
                                      "of": len(recs)}), flush=True)
    print(json.dumps({"label": label, "dumped": sets, "ckpt": ckpt}), flush=True)


# ------------------------------------------------------------------------------ sweep


def is_null(leaf: dict[str, Any], m: float) -> bool:
    return leaf["span_value"] is None or leaf["span_lp"] < leaf["null_lp"] + m


def matches(value: Any, alts: list[Any]) -> bool:
    from deg.bench.score import normalize

    return value is not None and normalize(value) in {normalize(a) for a in alts}


def point(leaves: list[dict[str, Any]], m: float) -> dict[str, Any]:
    ans = [x for x in leaves if x["gold"] == "ans"]
    un = [x for x in leaves if x["gold"] == "unans"]
    omit = sum(is_null(x, m) for x in ans)
    right = sum(not is_null(x, m) and matches(x["span_value"], x["alts"]) for x in ans)
    made_up = sum(not is_null(x, m) for x in un)
    mu = made_up / len(un) if un else 0.0
    om = omit / len(ans) if ans else 0.0
    return {"margin": m, "answerable_n": len(ans), "unanswerable_n": len(un),
            "made_up_rate": mu, "omission_rate": om,
            "answerable_correct": right / len(ans) if ans else None,
            "answerable_wrong_span": (len(ans) - omit - right) / len(ans) if ans else None,
            "unanswerable_null": 1 - mu,
            # suite-style balanced score (squad2 component definition)
            "balanced": ((right / len(ans) if ans else 0.0) + (1 - mu)) / 2,
            "cost": HALLUCINATION_COST * mu + om}


def check_margin0(leaves: list[dict[str, Any]]) -> int:
    """Leaves whose offline margin-0 decision differs from the real margin-0 output (0 =
    the sweep reproduces inference exactly)."""
    return sum((x["value"] is None) != is_null(x, 0.0)
               or (x["value"] is not None and x["value"] != x["span_value"]) for x in leaves)


def choose(curve: list[dict[str, Any]]) -> dict[str, Any]:
    """Min cost (the mean of the per-set costs) among non-collapsed points (omission <=
    MAX_OMISSION on every set); ties -> the margin closest to 0."""
    ok = [p for p in curve if all(s["omission_rate"] <= MAX_OMISSION for s in p["sets"].values())]
    return min(ok, key=lambda p: (round(p["cost"], 9), abs(p["margin"])))


def regold(rows: list[dict[str, Any]], heldout: Path) -> list[dict[str, Any]]:
    """The leaves of error-free dump rows, gold re-derived from the held-out records (the
    dump's own gold fields are not trusted: the sweep always uses :func:`leaf_gold`)."""
    recs = {r["uid"]: r for r in read_jsonl(heldout)}
    out = []
    for r in rows:
        if r["error"] is not None:
            continue
        for leaf in r["leaves"]:
            kind, alts = leaf_gold(recs[r["uid"]], leaf["ptr"])
            out.append({**leaf, "gold": kind, "alts": alts})
    return out


def sweep(out: Path, labels: list[str] | None = None) -> dict[str, Any]:
    result: dict[str, Any] = {"rule": f"cost = {HALLUCINATION_COST:g}*made_up_rate(unanswerable)"
                                      " + omission_rate(answerable), averaged over the held-out"
                                      " sets; omission <= "
                                      f"{MAX_OMISSION:g} on every set (no always-null collapse)",
                              "models": {}}
    for d in sorted((out / "dumps").iterdir()):
        if labels and d.name not in labels:
            continue
        sets = {p.stem: regold(read_jsonl(p), out / "heldout" / p.name)
                for p in sorted(d.glob("*.jsonl"))}
        curve = []
        for m in GRID:
            per = {k: point(v, m) for k, v in sets.items()}
            curve.append({"margin": m, "cost": sum(s["cost"] for s in per.values()) / len(per),
                          "sets": per})
        best = choose(curve)
        zero = next(p for p in curve if p["margin"] == 0)
        result["models"][d.name] = {
            "chosen_margin": best["margin"], "chosen": best, "margin0": zero,
            "margin0_mismatch": {k: check_margin0(v) for k, v in sets.items()},
            "errors": {p.stem: sum(r["error"] is not None for r in read_jsonl(p))
                       for p in sorted(d.glob("*.jsonl"))},
            "curve": curve}
    (out / "sweep.json").write_text(json.dumps(result, indent=1) + "\n")
    (out / "curve.md").write_text(curve_md(result))
    return result


def curve_md(result: dict[str, Any]) -> str:
    lines = ["# null-margin curve (held-out)", "", result["rule"], ""]
    for label, m in result["models"].items():
        names = list(m["chosen"]["sets"])
        lines += [f"## {label}", "", f"chosen margin **{m['chosen_margin']:g}** nats "
                  f"(margin-0 reproduction mismatches: {m['margin0_mismatch']}; "
                  f"errors: {m['errors']})", "",
                  "| margin | cost | " + " | ".join(
                      f"{n} made-up | {n} omit | {n} ans-correct | {n} balanced"
                      for n in names) + " |",
                  "|---|---|" + "---|" * (4 * len(names))]
        for p in m["curve"]:
            if round(p["margin"] * 10) % 5 and p["margin"] != m["chosen_margin"]:
                continue  # every 0.5 nats + the chosen point
            mark = " **<-**" if p["margin"] == m["chosen_margin"] else ""
            lines.append(f"| {p['margin']:g}{mark} | {p['cost']:.3f} | " + " | ".join(
                f"{s['made_up_rate']:.3f} | {s['omission_rate']:.3f} | "
                f"{s['answerable_correct']:.3f} | {s['balanced']:.3f}"
                for s in p["sets"].values()) + " |")
        lines.append("")
    return "\n".join(lines)


# ------------------------------------------------------------------------ suite rates


def extract_null_rates(recs: list[dict[str, Any]], preds: dict[str, Any]) -> dict[str, Any]:
    """Made-up (gold null/absent, prediction non-null) and wrong-null (gold non-null,
    prediction null/absent) rates over the extract leaves of suite records. A record with an
    error counts every leaf as null (its value is missing)."""
    from deg.evaluate import _abstract, _flatten
    from deg.schema import compile_schema

    ans = un = made_up = wrong_null = 0
    for rec in recs:
        p = preds.get(rec["uid"])
        got = _flatten(p.value) if p is not None and p.value is not None else {}
        fields = {f.pointer: f for f in compile_schema(rec["schema"]).fields}
        gold = _flatten(rec["value"])
        if "alternatives" in rec:
            gold = {ptr: (alts[0] if alts else None) for ptr, alts in rec["alternatives"].items()}
        ptrs = set(gold) | {f.pointer for f in fields.values()
                            if "*" not in f.pointer and f.mode.startswith("extract.")}
        for ptr in ptrs:
            if ptr.endswith("#len"):
                continue
            f = fields.get(_abstract(ptr))
            if f is None or not f.mode.startswith("extract.") or f.on_null == "forbid":
                continue
            gv, pv = gold.get(ptr), got.get(ptr)
            if gv is None:
                un += 1
                made_up += pv is not None
            else:
                ans += 1
                wrong_null += pv is None
    return {"unanswerable_n": un, "answerable_n": ans,
            "made_up_rate": made_up / un if un else None,
            "wrong_null_rate": wrong_null / ans if ans else None}


def suite_rates(system: str, results: Path | None = None,
                sources: tuple[str, ...] = ("datagen-test", "squad2", "mixed")) -> dict[str, Any]:
    from deg.bench.run import load_preds
    from deg.bench.suite import RESULTS, load_source, system_dir

    sdir = system_dir(results or RESULTS, system)
    out = {}
    for src in sources:
        path = sdir / f"{src}.preds.jsonl"
        if path.exists():
            out[src] = extract_null_rates(load_source(src), load_preds(path, keep_errors=True))
    return out


def main(argv: list[str] | None = None) -> None:
    ap = argparse.ArgumentParser(description=__doc__.split("\n")[0])
    sub = ap.add_subparsers(dest="cmd", required=True)
    b = sub.add_parser("build")
    b.add_argument("--out", type=Path, default=OUT)
    b.add_argument("--squad-paras", type=int, default=HELDOUT_SQUAD_PARAS)
    d = sub.add_parser("dump")
    d.add_argument("--ckpt", required=True)
    d.add_argument("--label", required=True)
    d.add_argument("--out", type=Path, default=OUT)
    d.add_argument("--sets", nargs="*", default=["squad2", "datagen-val"])
    d.add_argument("--max-gpu-gb", type=float, default=None)
    s = sub.add_parser("sweep")
    s.add_argument("--out", type=Path, default=OUT)
    s.add_argument("--labels", nargs="*", default=None)
    r = sub.add_parser("rates")
    r.add_argument("--system", nargs="+", required=True)
    r.add_argument("--results", type=Path, default=None)
    args = ap.parse_args(argv)
    if args.cmd == "build":
        print(json.dumps(build(args.out, args.squad_paras), indent=1))
    elif args.cmd == "dump":
        if args.max_gpu_gb:
            import torch
            total = torch.cuda.get_device_properties(0).total_memory
            torch.cuda.set_per_process_memory_fraction(min(1.0, args.max_gpu_gb * 2**30 / total))
        dump(args.ckpt, args.label, args.out, args.sets)
    elif args.cmd == "sweep":
        res = sweep(args.out, args.labels)
        print(json.dumps({k: {"chosen_margin": v["chosen_margin"],
                              "cost": v["chosen"]["cost"], "cost0": v["margin0"]["cost"],
                              "mismatch": v["margin0_mismatch"]}
                          for k, v in res["models"].items()}, indent=1))
    else:
        print(json.dumps({s: suite_rates(s, args.results) for s in args.system}, indent=1))


if __name__ == "__main__":
    main()