"""Calibration and a small eval for deg checkpoints (SPEC.md §11, §12).

``calibrate`` fits one temperature per head kind (choice, score, noul, extract) by NLL on
held-out data: datagen ``val`` records only by default (``--source datagen-val``; the
/v1/fill heads are calibrated for /v1/fill data). jev dev can be reported next to it
(``--report-jev``, fit pool unchanged) or fit on (``--source jev-dev`` or
``datagen-val+jev-dev``, the run-1 pool). ``evaluate`` runs the full
:func:`deg.infer.fill` path on datagen ``val`` and ``test`` (held-out domains) records and
reports, per split: valid-response rate, decide accuracy + ECE, extract exact match
(positional, and ``extract_em_aligned`` with array items matched order-insensitively), array
count accuracy, and generated tokens vs the tokens of the gold JSON written out in full
(the full-JSON constrained-decoding baseline emits at least that many). jev dev decisions
are scored against the teacher argmax.

    python -m deg.evaluate --ckpt /data/deg/runs/smoke/final
    python -m deg.evaluate calibrate --ckpt /data/deg/runs/run-1/step-0020000 \
        --out /data/deg/runs/run-1/final-v2 --report-jev

``calibrate`` writes a calibrated copy of a checkpoint (same backbone and head weights, the
fitted temperatures in ``heads.safetensors`` and ``deg.json``); the source is not modified.
"""

from __future__ import annotations

import argparse
import json
import random
import shutil
import time
from pathlib import Path
from typing import Any

import torch

from deg.infer import fill
from deg.model import KINDS, DegModel, load_model
from deg.request import FillOptions
from deg.response import FillError
from deg.schema import compile_schema

TEMPS = [round(0.4 + 0.1 * i, 2) for i in range(37)]  # 0.4 .. 4.0
CALIBRATION_SOURCES = ("datagen-val", "jev-dev", "datagen-val+jev-dev")
CALIBRATION_TARGETS = ("hard", "soft")
DECIDE_KINDS = ("choice", "score", "noul")


def _flatten(value: Any, prefix: str = "") -> dict[str, Any]:
    if isinstance(value, dict):
        out: dict[str, Any] = {}
        for k, v in value.items():
            out |= _flatten(v, f"{prefix}/{k.replace('~', '~0').replace('/', '~1')}")
        return out
    if isinstance(value, list):
        out = {f"{prefix}#len": len(value)}
        for i, v in enumerate(value):
            out |= _flatten(v, f"{prefix}/{i}")
        return out
    return {prefix: value}


def align_arrays(gold: Any, pred: Any, eq: Any = None) -> Any:
    """``pred`` with its arrays reordered (recursively) so item i is the predicted item
    that best matches gold item i (greedy by the fraction of matching non-null leaves,
    ties by index); unmatched gold items become ``None`` and unmatched predicted items
    are dropped. Order-insensitive scoring: anchored arrays come out in document order,
    which need not be the gold's order."""
    eq = eq or (lambda g, p: g == p)
    if isinstance(gold, dict) and isinstance(pred, dict):
        return {k: (align_arrays(gold[k], v, eq) if k in gold else v) for k, v in pred.items()}
    if isinstance(gold, list) and isinstance(pred, list):
        pairs = []
        for i, g in enumerate(gold):
            fg = {k: v for k, v in _flatten(g).items() if not k.endswith("#len") and v is not None}
            for j, q in enumerate(pred):
                fq = _flatten(q)
                hit = sum(1 for k, v in fg.items() if k in fq and eq(v, fq[k]))
                pairs.append((-hit / max(1, len(fg)), i, j))
        pairs.sort()
        used_g: set[int] = set()
        used_p: set[int] = set()
        match: dict[int, int] = {}
        for score, i, j in pairs:
            if score >= 0 or i in used_g or j in used_p:
                continue
            used_g.add(i)
            used_p.add(j)
            match[i] = j
        free = [j for j in range(len(pred)) if j not in used_p]  # unmatched: keep in order
        out = []
        for i, g in enumerate(gold):
            if i in match:
                out.append(align_arrays(g, pred[match[i]], eq))
            elif free:
                out.append(align_arrays(g, pred[free.pop(0)], eq))
        return out + [pred[j] for j in free]
    return pred


def _abstract(ptr: str) -> str:
    return "/".join("*" if p.isdigit() else p for p in ptr.split("/"))


def ece(pairs: list[tuple[float, bool]], bins: int = 10) -> float | None:
    if not pairs:
        return None
    total = 0.0
    for b in range(bins):
        lo, hi = b / bins, (b + 1) / bins
        sel = [(c, ok) for c, ok in pairs if lo < c <= hi or (b == 0 and c == 0)]
        if sel:
            conf = sum(c for c, _ in sel) / len(sel)
            acc = sum(ok for _, ok in sel) / len(sel)
            total += len(sel) / len(pairs) * abs(conf - acc)
    return round(total, 4)


def _collect_logits(model: DegModel, examples: list[Any]) -> dict[str, list]:
    """kind -> [(logits/start-end, target)] on the inference path."""
    got: dict[str, list] = {k: [] for k in KINDS}
    for ex in examples:
        items = [it for it in ex.items if it.query is not None]
        if not items:
            continue
        hidden, state = model.trunk([ex.ctx])
        outs = model.answer_batched(state, hidden, ex.ctx, [it.query for it in items])
        for it, out in zip(items, outs):
            if it.kind == "extract":
                got["extract"].append(((out.start.float(), out.end.float()), it))
            elif it.kind in ("choice", "score", "noul"):
                got[it.kind].append((out.logits.float(), it))
    return got


def _nll(kind: str, rows: list, t: float, hard: bool = False) -> float:
    """Mean NLL at temperature ``t``; ``hard``: decide targets become one-hot of their
    argmax (the filled value / the teacher's top outcome), the label accuracy and ECE are
    scored against."""
    total = 0.0
    for logits, it in rows:
        if kind == "extract":
            start, end = logits
            for lg, pos in ((start, it.starts), (end, it.ends)):
                logp = torch.log_softmax(lg / t, -1)
                idx = torch.tensor([lg.shape[0] - 1 if p < 0 else p for p in pos],
                                   device=lg.device)
                total -= float(torch.logsumexp(logp[idx.unique()], 0))
        else:
            target = torch.tensor(it.target, device=logits.device).clamp_min(0)
            if hard:
                target = torch.nn.functional.one_hot(target.argmax(), len(target)).float()
            total -= float((target / target.sum() * torch.log_softmax(logits / t, -1)).sum())
    return total / max(1, len(rows))


def _top1_pairs(kind: str, rows: list, t: float) -> list[tuple[float, bool]]:
    """(confidence, correct) of each decide row's argmax at temperature ``t``; correct =
    the target's argmax (the filled value, or the teacher's top outcome)."""
    out = []
    for logits, it in rows:
        probs = torch.softmax(logits / t, -1)
        best = max(range(len(it.target)), key=it.target.__getitem__)
        out.append((float(probs.max()), int(probs.argmax()) == best))
    return out


def fit_report(got: dict[str, list], temps: dict[str, float], before: dict[str, float],
               hard: bool = True) -> dict[str, Any]:
    """Per head kind of one pool: rows, NLL and (decide kinds) top-1 accuracy + ECE at the
    ``before`` and fitted temperatures, plus all decide kinds pooled."""
    out: dict[str, Any] = {}
    pooled: dict[str, list] = {"before": [], "after": []}
    for kind in KINDS:
        rows = got.get(kind) or []
        if not rows:
            continue
        m: dict[str, Any] = {"n": len(rows), "t_before": before[kind], "t_after": temps[kind],
                             "nll_before": round(_nll(kind, rows, before[kind], hard), 4),
                             "nll_after": round(_nll(kind, rows, temps[kind], hard), 4)}
        if kind in DECIDE_KINDS:
            for tag, t in (("before", before[kind]), ("after", temps[kind])):
                pairs = _top1_pairs(kind, rows, t)
                pooled[tag] += pairs
                m[f"ece_{tag}"] = ece(pairs)
            m["acc"] = round(sum(ok for _, ok in pairs) / len(pairs), 4)
        out[kind] = m
    if pooled["after"]:
        out["decide"] = {"n": len(pooled["after"]), "ece_before": ece(pooled["before"]),
                         "ece_after": ece(pooled["after"])}
    return out


def _as_paths(records: Path | list[Path] | tuple[Path, ...]) -> list[Path]:
    return [Path(records)] if isinstance(records, (str, Path)) else [Path(r) for r in records]


@torch.no_grad()
def calibrate(model: DegModel, records: Path | list[Path], jev_dev: Path, n_jev: int = 300,
              seed: int = 0, info: dict[str, Any] | None = None, *,
              source: str = "datagen-val", report_jev: bool = False,
              max_context_tokens: int | None = None, target: str = "hard"
              ) -> dict[str, float]:
    """Fit and set ``model``'s temperatures on ``source``: datagen ``val`` records of every
    ``records`` file (default), jev dev, or both pooled. ``info`` (if given) receives the
    sample sizes and, per pool (``datagen_val``, ``jev_dev`` when fit on or
    ``report_jev``), NLL/ECE before (the model's current temperatures) and after.
    ``max_context_tokens`` leaves longer val contexts out (the end-of-training call).
    ``target="hard"`` (default) fits decide heads to the one-hot filled value, the label
    accuracy/ECE are judged against; ``"soft"`` fits the soft training targets (Clef
    ``decide_probs``, jev teacher probs), which are less confident than the gold and pull
    the temperatures above 1 (calibration-v2.md)."""
    from deg.train import deg_example, jev_example, read_jsonl

    if target not in CALIBRATION_TARGETS:
        raise ValueError(f"calibration target {target!r} not in {CALIBRATION_TARGETS}")
    hard = target == "hard"
    if source not in CALIBRATION_SOURCES:
        raise ValueError(f"calibration source {source!r} not in {CALIBRATION_SOURCES}")
    fit_val = source in ("datagen-val", "datagen-val+jev-dev")
    fit_jev = source in ("jev-dev", "datagen-val+jev-dev")
    rng = random.Random(seed)
    pools: dict[str, dict[str, list]] = {}
    sizes: dict[str, int] = {}
    if fit_val:
        examples = []
        for path in _as_paths(records):
            for r in read_jsonl(path, "val"):
                ex = deg_example(model, r, rng, shuffle_p=0.0)
                if max_context_tokens is None or len(ex.ctx.ids) <= max_context_tokens:
                    examples.append(ex)
        sizes["val_records"] = len(examples)
        pools["datagen_val"] = _collect_logits(model, examples)
    if (fit_jev or report_jev) and jev_dev.exists() and n_jev:
        examples = [e for r in read_jsonl(jev_dev, None, n_jev)
                    if (e := jev_example(model, r, rng, shuffle_p=0.0)) is not None]
        sizes["jev_records"] = len(examples)
        pools["jev_dev"] = _collect_logits(model, examples)
    fit_on = [k for k, on in (("datagen_val", fit_val), ("jev_dev", fit_jev)) if on]
    got = {k: [row for name in fit_on for row in pools.get(name, {}).get(k, [])]
           for k in KINDS}
    before = dict(zip(KINDS, (float(t) for t in model.heads.temperature.tolist())))
    temps = {}
    for i, kind in enumerate(KINDS):
        rows = got.get(kind) or []
        if len(rows) < 10:
            temps[kind] = before[kind]  # too few rows: unchanged
            continue
        temps[kind] = min(TEMPS, key=lambda t: _nll(kind, rows, t, hard))
        model.heads.temperature[i] = temps[kind]
    if info is not None:
        info.update({"fit_source": source, "fit_target": target, "fit_on": fit_on, **sizes,
                     "rows": {k: len(v) for k, v in got.items()},
                     "report": {name: fit_report(pool, temps, before, hard)
                                for name, pool in pools.items()}})
    return temps


@torch.no_grad()
def evaluate(model: DegModel, records: Path, jev_dev: Path, n_records: int = 60,
             n_jev: int = 300) -> dict[str, Any]:
    from deg.train import jev_example, read_jsonl

    tok = model.renderer.tokenizer
    report: dict[str, Any] = {}
    for split in ("val", "test"):
        recs = read_jsonl(records, split, n_records)
        m: dict[str, Any] = {"records": len(recs), "valid": 0, "errors": {}}
        dec_pairs: list[tuple[float, bool]] = []
        ext_ok = ext_n = cnt_ok = cnt_n = gen_fields = 0
        ext_ok_al = ext_n_al = 0
        gen_tokens = baseline_tokens = 0
        t0 = time.time()
        for rec in recs:
            plan = compile_schema(rec["schema"])
            baseline_tokens += len(tok(json.dumps(rec["value"], ensure_ascii=False),
                                       add_special_tokens=False)["input_ids"])
            try:
                resp = fill(model, rec["context"], rec["schema"], FillOptions(), plan=plan)
            except FillError as exc:
                key = str(exc)[:80]
                m["errors"][key] = m["errors"].get(key, 0) + 1
                continue
            m["valid"] += 1
            gen_tokens += resp["usage"]["generated_tokens"]
            gold, pred = _flatten(rec["value"]), _flatten(resp["value"])
            modes = {}
            for f in plan.fields:
                modes[f.pointer] = f.mode
            for ptr, gv in gold.items():
                if ptr.endswith("#len"):
                    cnt_n += 1
                    cnt_ok += pred.get(ptr) == gv
                    continue
                mode = modes.get(_abstract(ptr), "")
                if mode.startswith("decide."):
                    meta = resp["meta"].get(ptr)
                    if meta is None:
                        continue
                    p = meta["probability"]
                    conf = max(p, 1 - p) if mode == "decide.noul" else p
                    dec_pairs.append((conf, pred.get(ptr) == gv and ptr in pred))
                elif mode.startswith("extract."):
                    ext_n += 1
                    ext_ok += ptr in pred and pred[ptr] == gv
                elif mode == "generate":
                    gen_fields += 1
            # extract fields the gold omits: correct iff the prediction omits/nulls them too
            for ptr, pv in pred.items():
                mode = modes.get(_abstract(ptr), "")
                if mode.startswith("extract.") and ptr not in gold:
                    ext_n += 1
                    ext_ok += pv is None
            # order-insensitive: array items matched to their best gold item
            aligned = _flatten(align_arrays(rec["value"], resp["value"]))
            for ptr in set(gold) | set(aligned):
                if not modes.get(_abstract(ptr), "").startswith("extract."):
                    continue
                if ptr in gold:
                    ext_n_al += 1
                    ext_ok_al += ptr in aligned and aligned[ptr] == gold[ptr]
                else:
                    ext_n_al += 1
                    ext_ok_al += aligned[ptr] is None
        m.update({
            "valid_rate": round(m["valid"] / max(1, len(recs)), 4),
            "decide_n": len(dec_pairs),
            "decide_acc": round(sum(ok for _, ok in dec_pairs) / max(1, len(dec_pairs)), 4),
            "decide_ece": ece(dec_pairs),
            "extract_n": ext_n, "extract_em": round(ext_ok / max(1, ext_n), 4),
            "extract_em_aligned": round(ext_ok_al / max(1, ext_n_al), 4),
            "count_n": cnt_n, "count_acc": round(cnt_ok / max(1, cnt_n), 4),
            "generate_fields": gen_fields,
            "generated_tokens": gen_tokens, "full_json_tokens": baseline_tokens,
            "generated_vs_full_json": round(gen_tokens / max(1, baseline_tokens), 4),
            "s_per_record": round((time.time() - t0) / max(1, len(recs)), 3),
        })
        report[split] = m
    if jev_dev.exists() and n_jev:
        rng = random.Random(0)
        examples = [e for r in read_jsonl(jev_dev, None, n_jev)
                    if (e := jev_example(model, r, rng, shuffle_p=0.0)) is not None]
        got = _collect_logits(model, examples)
        out: dict[str, Any] = {}
        pairs_all = []
        for kind in ("choice", "score", "noul"):
            pairs = []
            for logits, it in got[kind]:
                probs = torch.softmax(logits / model.heads.temp(kind), -1)
                pairs.append((float(probs.max()),
                              int(probs.argmax()) == max(range(len(it.target)),
                                                         key=it.target.__getitem__)))
            pairs_all += pairs
            out[kind] = {"n": len(pairs),
                         "acc": round(sum(ok for _, ok in pairs) / max(1, len(pairs)), 4),
                         "ece": ece(pairs)}
        out["all"] = {"n": len(pairs_all),
                      "acc": round(sum(ok for _, ok in pairs_all) / max(1, len(pairs_all)), 4),
                      "ece": ece(pairs_all)}
        report["jev_dev"] = out
    report["temperature"] = dict(zip(KINDS, model.heads.temperature.tolist()))
    return report


def write_calibrated(model: DegModel, src: Path, out: Path,
                     calibration: dict[str, Any]) -> None:
    """Copy checkpoint ``src`` to ``out`` with ``model``'s fitted temperatures. Every other
    tensor is copied from ``src`` byte for byte (a run-1 checkpoint has no anchor heads, so
    saving ``model`` itself would add freshly initialized ones and flip it to anchor mode)."""
    from safetensors.torch import load_file, save_file

    src, out = Path(src), Path(out)
    if out.resolve() == src.resolve():
        raise ValueError("calibrate --out must differ from --ckpt (the source is kept as is)")
    out.mkdir(parents=True, exist_ok=True)
    shutil.copyfile(src / "backbone.safetensors", out / "backbone.safetensors")
    heads = load_file(str(src / "heads.safetensors"))
    heads["temperature"] = model.heads.temperature.detach().float().cpu().contiguous()
    save_file(heads, str(out / "heads.safetensors"))
    meta = json.loads((src / "deg.json").read_text())
    meta["temperature"] = dict(zip(KINDS, model.heads.temperature.tolist()))
    meta["calibration"] = {"source": str(src), **calibration}
    (out / "deg.json").write_text(json.dumps(meta, indent=2))


def calibrate_main(argv: list[str]) -> None:
    from deg.train import DEG_RECORDS, JEV_DEV

    ap = argparse.ArgumentParser(description="fit head temperatures; write a calibrated copy")
    ap.add_argument("--ckpt", type=Path, required=True)
    ap.add_argument("--out", type=Path, required=True)
    ap.add_argument("--records", type=Path, action="append", default=None,
                    help=f"datagen records file(s) whose val split is fit on "
                         f"(repeatable; default {DEG_RECORDS})")
    ap.add_argument("--source", choices=CALIBRATION_SOURCES, default="datagen-val",
                    help="fit pool (default: datagen val only)")
    ap.add_argument("--report-jev", action="store_true",
                    help="also report jev dev NLL/ECE before/after (separately; not fit on "
                         "unless --source includes jev-dev)")
    ap.add_argument("--target", choices=CALIBRATION_TARGETS, default="hard",
                    help="decide NLL target: hard = one-hot filled value (default), soft = "
                         "the soft training targets")
    ap.add_argument("--jev-dev", type=Path, default=JEV_DEV)
    ap.add_argument("--n-jev", type=int, default=1000)
    args = ap.parse_args(argv)
    records = args.records or [DEG_RECORDS]
    model = load_model(args.ckpt, device="cuda")
    t0 = time.time()
    info: dict[str, Any] = {}
    temps = calibrate(model, records, args.jev_dev, args.n_jev, info=info,
                      source=args.source, report_jev=args.report_jev, target=args.target)
    calibration = {"temperature": temps, "records": [str(r) for r in records],
                   "jev_dev": str(args.jev_dev), **info, "array_mode": model.array_mode,
                   "seconds": round(time.time() - t0, 1)}
    write_calibrated(model, args.ckpt, args.out, calibration)
    print(json.dumps({"out": str(args.out), **calibration}))


def main(argv: list[str] | None = None) -> None:
    import sys

    from deg.train import DEG_RECORDS, JEV_DEV

    argv = sys.argv[1:] if argv is None else argv
    if argv[:1] == ["calibrate"]:
        calibrate_main(argv[1:])
        return

    ap = argparse.ArgumentParser(description="deg eval")
    ap.add_argument("--ckpt", type=Path, default=None, help="None = untrained heads")
    ap.add_argument("--records", type=Path, default=DEG_RECORDS)
    ap.add_argument("--n-records", type=int, default=60)
    ap.add_argument("--n-jev", type=int, default=300)
    ap.add_argument("--out", type=Path, default=None)
    ap.add_argument("--set", choices=("datagen", "extractbench"), default="datagen",
                    help="extractbench: human-gold ExtractBench records (deg.extractbench "
                         "prepare); EVAL ONLY")
    ap.add_argument("--extractbench-records", type=Path, default=None)
    ap.add_argument("--predictions", type=Path, default=None,
                    help="extractbench: write one {uid, gold, pred, error} line per document")
    args = ap.parse_args(argv)
    model = load_model(args.ckpt, device="cuda")
    if args.set == "extractbench":
        from deg.extractbench import RECORDS, evaluate_records

        report = evaluate_records(model, args.extractbench_records or RECORDS,
                                  predictions=args.predictions)
    else:
        report = evaluate(model, args.records, JEV_DEV, args.n_records, args.n_jev)
    text = json.dumps(report, indent=2)
    print(text)
    if args.out:
        args.out.write_text(text)


if __name__ == "__main__":
    main()