#!/usr/bin/env python3
"""NLI shadow GRID sweep — #000049 §7 #22 follow-up: characterize the
whole {model × candidate-clause-cap × aggregation × θc × θe} space at
once, to learn whether *any* config of the lexical-candidate approach
clears the §7 #12 gate (catch the recombinations, ~0 false-demote on
already-STRICT answers) before deciding the Phase-3 verifier-matched-
clause hook is the only path.

Inputs (both shapes from `nli_shadow_sweep.py` are accepted per record:
`{answer_text|claim, context|source}`; `_meta` lines skipped):
  --should-demote   records where demoting IS correct (false claims /
                    recombination fixtures). Default: the in-repo
                    `bench/fixtures/5f/falsification-hard-v1.jsonl`
                    (12; the 2 recombination fixtures + 10 other false
                    near-misses). For the richer 28-case synthetic set
                    point this at `arborist-nli-bench/eval/recombination.jsonl`.
  --should-not-demote  records where demoting is a FALSE POSITIVE — a
                    bench-qa JSONL; rows are filtered to `audit_mode ==
                    STRICT` (the confidently-grounded answers) unless
                    --no-strict-filter. (Add `arborist-nli-bench/eval/
                    legit_summary.jsonl` too if you have it.)

How it stays cheap: NLI inference runs ONCE per (model, record) over
the record's top-`KMAX` candidate clauses; the k/aggregation/θ grid is
then pure arithmetic on the cached per-clause scores. On a GPU the
whole grid is seconds. Models come from `arborist/qa/nli/manifest.json`
(`nli_model_version`+`hf_repo`+`pinned_revision`, plus `alternates`).

SHADOW ONLY — writes nothing but `bench/results/nli-shadow-grid.json`.

Usage:
  python3 bench/scripts/nli_shadow_grid.py --should-not-demote bench/qa_results/<...>.jsonl
  python3 bench/scripts/nli_shadow_grid.py --should-not-demote a.jsonl --should-demote ~/git/arborist-nli-bench/eval/recombination.jsonl --models all
  python3 bench/scripts/nli_shadow_grid.py --should-not-demote a.jsonl --models nli-shadow-v1-minilm2-l6-h768,nli-shadow-v1-bart-large-mnli
"""
from __future__ import annotations

import argparse
import json
import sys
import time
from pathlib import Path

REPO = Path(__file__).resolve().parents[2]
KMAX = 12  # NLI is run on the top-KMAX candidate clauses; the grid sub-selects k <= KMAX.
K_GRID = (1, 2, 3, 4, 6, 8, 12)
# how to combine the per-clause contradiction scores over the top-k clauses:
#   max / mean — over all k; top2 / top3 — the 2nd / 3rd highest (needs agreement);
#   meanTop2 / meanTop3 — mean of the 2 / 3 highest; margin — max over clauses of
#   (p_contra - p_entail) of THAT clause (folds the entailment guard into the score).
AGGS = ("max", "mean", "top2", "top3", "meanTop2", "meanTop3", "margin")
# the entailment guard variant: max_entail = max p_entail over the top-k clauses;
# paired_entail = p_entail of the SAME clause that produced the aggregated contra
# (the principled version — "the contradicting clause itself doesn't entail").
GUARDS = ("max_entail", "paired_entail")
THETA_C = (0.30, 0.40, 0.50, 0.60, 0.70, 0.80, 0.85, 0.90, 0.92, 0.94, 0.95, 0.96,
           0.97, 0.975, 0.98, 0.985, 0.99, 0.995, 0.998, 0.999)
THETA_E = (0.30, 0.50, 0.70, 0.90, 0.95, 1.01)   # 1.01 == "no guard"


def _records(path: Path):
    for ln in path.read_text().splitlines():
        ln = ln.strip()
        if not ln:
            continue
        o = json.loads(ln)
        if "_meta" in o:
            continue
        claim = o.get("answer_text") or o.get("claim")
        source = o.get("context") or o.get("source")
        if not claim or not source:
            continue
        yield {"id": o.get("id", path.stem), "claim": claim, "source": source,
               "audit_mode": o.get("audit_mode")}


def _models_from_manifest(which: str | None):
    sys.path.insert(0, str(REPO))
    from arborist.qa.nli.shadow import load_manifest
    m = load_manifest()
    entries = [{"nli_model_version": m["nli_model_version"], "hf_repo": m["hf_repo"],
                "pinned_revision": m.get("pinned_revision")}]
    entries += [{"nli_model_version": a["nli_model_version"], "hf_repo": a["hf_repo"],
                 "pinned_revision": a.get("pinned_revision")} for a in m.get("alternates", [])]
    if which in (None, ""):
        return entries[:1]
    if which == "all":
        return entries
    want = {w.strip() for w in which.split(",")}
    sel = [e for e in entries if e["nli_model_version"] in want or e["hf_repo"] in want]
    return sel or entries[:1]


def _aggregate(top: list[tuple], how: str) -> tuple[float, int]:
    """Return (aggregated_contradiction_score, idx_of_the_clause_it_came_from)
    over `top` = [(overlap, pe, pn, pc), ...] (top-k clauses, overlap-ranked).
    For pooled aggs (mean/topN/meanTopN) the 'idx' is the highest-pc clause in
    the pool — used for the paired-entail guard."""
    if not top:
        return 0.0, -1
    pcs = [t[3] for t in top]
    order = sorted(range(len(top)), key=lambda i: pcs[i], reverse=True)
    if how == "max":
        return pcs[order[0]], order[0]
    if how == "mean":
        return sum(pcs) / len(pcs), order[0]
    if how == "top2":
        return (pcs[order[1]], order[1]) if len(order) >= 2 else (0.0, order[0])
    if how == "top3":
        return (pcs[order[2]], order[2]) if len(order) >= 3 else (0.0, order[-1])
    if how == "meanTop2":
        sel = order[:2]; return sum(pcs[i] for i in sel) / len(sel), order[0]
    if how == "meanTop3":
        sel = order[:3]; return sum(pcs[i] for i in sel) / len(sel), order[0]
    if how == "margin":  # max over clauses of (pc - pe); guard folded in
        margins = [t[3] - t[1] for t in top]
        bi = max(range(len(top)), key=lambda i: margins[i])
        return margins[bi], bi
    raise ValueError(how)


def _cache_scores(nli, recs: list[dict]) -> list[dict]:
    """For each record: segment, take top-KMAX candidate clauses, NLI each.
    Returns [{id, audit_mode, n_clauses, clause_scores=[(overlap,pe,pn,pc),...]}]."""
    from arborist.qa.nli.shadow import clauses as _clauses, candidate_clauses as _cand
    out = []
    for r in recs:
        n_clauses = len(_clauses(r["source"]))
        cand = _cand(r["claim"], r["source"], KMAX)  # list[(clause, overlap)]
        triples = nli._nli_batch([(cl, r["claim"]) for cl, _ in cand]) if cand else []
        clause_scores = [(ov, pe, pn, pc) for (_, ov), (pe, pn, pc) in zip(cand, triples)]
        out.append({"id": r["id"], "audit_mode": r.get("audit_mode"),
                    "n_clauses": n_clauses, "clause_scores": clause_scores})
    return out


def _rate(cached: list[dict], k: int, agg: str, guard: str, tc: float, te: float) -> float:
    """fraction of records that would_demote under (k, agg, guard, θc, θe).
    For agg='margin' the guard is folded into the score (margin = pc-pe), so
    only θc applies; the explicit guard loop is skipped for it (te is ignored)."""
    n = len(cached)
    if not n:
        return 0.0
    fired = 0
    is_margin = (agg == "margin")
    for c in cached:
        top = c["clause_scores"][:k]
        if not top:
            continue
        agg_c, src_i = _aggregate(top, agg)
        if agg_c < tc:
            continue
        if is_margin:
            fired += 1
            continue
        if guard == "max_entail":
            guard_e = max(t[1] for t in top)
        else:  # paired_entail — the entail of the clause that produced agg_c
            guard_e = top[src_i][1] if 0 <= src_i < len(top) else max(t[1] for t in top)
        if guard_e < te:
            fired += 1
    return fired / n


def _pareto(pts: list[dict]) -> list[dict]:
    """Pareto frontier: sorted by fp asc, keep a point only if its catch beats
    the best catch seen at a lower-or-equal fp."""
    out, best = [], -1.0
    for p in sorted(pts, key=lambda p: (p["fp"], -p["catch"])):
        if p["catch"] > best:
            out.append(p); best = p["catch"]
    return out


def main(argv=None) -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--should-demote", type=Path,
                    default=REPO / "bench" / "fixtures" / "5f" / "falsification-hard-v1.jsonl")
    ap.add_argument("--should-not-demote", type=Path, required=True,
                    help="a bench-qa JSONL (rows filtered to audit_mode==STRICT)")
    ap.add_argument("--no-strict-filter", action="store_true",
                    help="don't filter the should-not-demote set to STRICT rows")
    ap.add_argument("--models", default=None, help="manifest model id(s) comma-sep, or 'all' (default: main only)")
    ap.add_argument("--extra-models", default=None, help="comma-sep HF repo ids to also test (no manifest pin needed) — for 'no stone unturned' breadth")
    ap.add_argument("--out", type=Path, default=REPO / "bench" / "results" / "nli-shadow-grid.json")
    args = ap.parse_args(argv)

    try:
        import torch  # noqa
        sys.path.insert(0, str(REPO))
        from arborist.qa.nli.shadow import ShadowNLI
    except ImportError as e:
        print(f"[nli-grid] missing dependency: {e}  (pip install 'arborist[nli]')", file=sys.stderr)
        return 2

    sd = list(_records(args.should_demote))
    snd = list(_records(args.should_not_demote))
    if not args.no_strict_filter:
        snd_all = len(snd)
        snd = [r for r in snd if r.get("audit_mode") == "STRICT"]
        print(f"[nli-grid] should-not-demote: {len(snd)}/{snd_all} rows are audit_mode==STRICT", flush=True)
    print(f"[nli-grid] should-demote: {len(sd)} records ({args.should_demote.name})", flush=True)
    if not sd or not snd:
        print("[nli-grid] need non-empty inputs on both sides", file=sys.stderr)
        return 2

    models = _models_from_manifest(args.models)
    if args.extra_models:
        for repo in [r.strip() for r in args.extra_models.split(",") if r.strip()]:
            mv = "extra:" + repo.split("/")[-1]
            if not any(m["hf_repo"] == repo for m in models):
                models.append({"nli_model_version": mv, "hf_repo": repo, "pinned_revision": None})
    print(f"[nli-grid] models ({len(models)}): {[m['nli_model_version'] for m in models]}", flush=True)

    report = {"generated_at": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
              "should_demote_file": str(args.should_demote), "n_should_demote": len(sd),
              "should_not_demote_file": str(args.should_not_demote), "n_should_not_demote": len(snd),
              "k_grid": list(K_GRID), "aggs": list(AGGS), "guards": list(GUARDS),
              "theta_c_grid": list(THETA_C), "theta_e_grid": list(THETA_E),
              "kmax": KMAX, "models": {}}

    for me in models:
        mv = me["nli_model_version"]
        print(f"[nli-grid] === {mv}  ({me['hf_repo']}) ===", flush=True)
        manifest_override = {"nli_model_version": mv, "hf_repo": me["hf_repo"],
                             "pinned_revision": me.get("pinned_revision"),
                             "thresholds": {"contradiction_veto": 0.5, "entailment_block_veto": 0.9},
                             "max_length": 256, "max_candidate_clauses": KMAX}
        nli = ShadowNLI(manifest=manifest_override)
        nli._ensure_loaded()
        if not nli.available:
            print(f"[nli-grid]   skip — {nli._reason}", file=sys.stderr)
            continue
        t0 = time.time()
        cache_sd = _cache_scores(nli, sd)
        cache_snd = _cache_scores(nli, snd)
        infer_s = time.time() - t0
        n_pairs = sum(len(c["clause_scores"]) for c in cache_sd + cache_snd)
        print(f"[nli-grid]   backend={nli.backend} device={nli.device} · {n_pairs} NLI pairs in {infer_s:.2f}s", flush=True)

        configs = []           # best fp=0 (else fp<=.05 / knee) point per (k, agg, guard)
        all_pts = []            # every (model-internal) point, for the global pareto + global best
        for k in K_GRID:
            for agg in AGGS:
                guard_list = ("max_entail",) if agg == "margin" else GUARDS
                te_list = (1.01,) if agg == "margin" else THETA_E
                for guard in guard_list:
                    pts = []
                    for te in te_list:
                        for tc in THETA_C:
                            catch = _rate(cache_sd, k, agg, guard, tc, te)
                            fp = _rate(cache_snd, k, agg, guard, tc, te)
                            p = {"k": k, "agg": agg, "guard": guard, "theta_c": tc,
                                 "theta_e": te, "catch": round(catch, 4), "fp": round(fp, 4)}
                            pts.append(p); all_pts.append(p)
                    zero = [p for p in pts if p["fp"] == 0.0]
                    low = [p for p in pts if p["fp"] <= 0.05]
                    if zero:
                        best = max(zero, key=lambda p: p["catch"]); best_kind = "fp=0"
                    elif low:
                        best = max(low, key=lambda p: p["catch"]); best_kind = "fp<=0.05"
                    else:
                        best = max(pts, key=lambda p: p["catch"] - p["fp"]); best_kind = "knee"
                    configs.append({"k": k, "agg": agg, "guard": guard, "best": best, "best_kind": best_kind})
        def _ck(c):
            b, kind = c["best"], c["best_kind"]
            return (0 if kind == "fp=0" else 1 if kind == "fp<=0.05" else 2, -b["catch"], b["fp"])
        configs.sort(key=_ck)
        # global best fp=0 point across ALL (k,agg,guard,θc,θe), then the pareto frontier
        zero_all = [p for p in all_pts if p["fp"] == 0.0]
        global_best = max(zero_all, key=lambda p: p["catch"]) if zero_all else max(all_pts, key=lambda p: p["catch"] - p["fp"])
        pareto = _pareto(all_pts)
        report["models"][mv] = {"hf_repo": me["hf_repo"], "backend": nli.backend, "device": nli.device,
                                "n_nli_pairs": n_pairs, "infer_seconds": round(infer_s, 2),
                                "global_best": global_best, "pareto": pareto, "configs_ranked": configs}
        gb = global_best
        print(f"[nli-grid]   GLOBAL best fp=0: k={gb['k']} agg={gb['agg']} guard={gb['guard']} θc={gb['theta_c']} θe={gb['theta_e']} "
              f"→ catch={gb['catch']} fp={gb['fp']}", flush=True)

    args.out.parent.mkdir(parents=True, exist_ok=True)
    args.out.write_text(json.dumps(report, indent=2))
    print(f"[nli-grid] wrote {args.out}")

    # summary table
    print("\n" + "=" * 100)
    print(f"NLI SHADOW GRID  —  global-best fp=0 config per model  ({len(sd)} should-demote vs {len(snd)} should-not-demote)")
    print("=" * 100)
    print(f"{'model':<32} {'hf_repo':<44} {'k':>3} {'agg':>9} {'guard':>14} {'θc':>6} {'θe':>5} {'catch':>7} {'fp':>6}")
    print("-" * 100)
    ranked = sorted(report["models"].items(), key=lambda kv: (-(kv[1]["global_best"]["catch"] if kv[1]["global_best"]["fp"] == 0 else -1), kv[1]["global_best"]["fp"]))
    for mv, md in ranked:
        g = md["global_best"]
        print(f"{mv:<32} {md['hf_repo']:<44} {g['k']:>3} {g['agg']:>9} {g['guard']:>14} {g['theta_c']:>6} {g['theta_e']:>5} {g['catch']:>7.3f} {g['fp']:>6.3f}")
    print("-" * 100)
    print("catch = fraction of should-demote (recombination falsehoods) the veto fires on  (higher better)")
    print("fp    = fraction of should-not-demote (confidently-grounded STRICT answers) the veto WRONGLY fires on  (≈0 clears §7 #12)")
    print("'global_best' = the highest-catch point with fp==0 across the FULL {k × agg × guard × θc × θe} grid for that model.")
    print("Each model's 'pareto' (in the JSON) is the catch-vs-fp frontier — read it to pick an operating point above fp=0 if a small false-demote rate is acceptable.")
    print("If no model reaches a useful catch at fp=0, the lexical-candidate approach is precision-limited — Phase-3 verifier-matched-clause hook is the path.")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())
