#!/usr/bin/env python3
"""#000057 — control-arm characterization sweep (fox: "we both do not
know which framing is right, measure all benchmarks and bring the
results forward for review"; "we have qwen with and without reasoning
to use"; "make the n huge huge and check in every ~7 turns").

This is NOT a verdict generator. It is the wide, parallel, incremental
sweep the bench-maxing discipline demands when a config/framing
decision is genuinely open: sweep model × question-framing at HUGE N,
log every item the instant it lands, present the SAME judge verdicts
under BOTH narrative framings, and hand the table to review.

WHY this sweep (the scientific question fox is probing):
  Is "Hermes-solo confidently fabricates the current officeholder" a
  real capability gap, or an artefact of (a) an 8B model being weak
  and (b) an unfair *truth* framing (penalising a model for knowing
  post-corpus reality)?  The control arm answered across
    {Hermes-3-8B, Qwen3.6-27B reasoning, Qwen3.6-27B no-reasoning}
  × {plain, source-relative, as-of-corpus-era} questions, judged vs
  the FIXED corpus-vintage gold, decides it empirically:
    * if Qwen-27B-reasoning honestly ABSTAINS where Hermes-8B
      fabricates → the gap is largely a weak-small-model artefact;
    * if even Qwen-27B-reasoning confidently asserts post-corpus →
      the gap is real and scale-independent;
    * `as_of_corpus` ("As of 2010, who was…") separates "can't recall
      the corpus era" from "won't constrain to a source" — the exact
      ambiguity fox flagged. (Corpus ~2010-2011, verified from the
      artefact — see mine_questions.py — NOT the 2003 dump CLAUDE.md
      names.)

DESIGN for huge N:
  * the huge N goes on the CONTROL (solo × 3 models × 3 framings) —
    that is where fox's open statistical question lives. The Arborist
    treatment arm is a FIXED smaller A/B reference (--arborist-n,
    default 40): re-measuring Arborist 386× does not buy power on the
    control question, and it is the heaviest call (retrieval over
    ~40GB shards + proof path). Power belongs where the question is.
  * parallel: independent hermetic judge calls + independent solo HTTP
    calls fanned across a bounded worker pool (CLAUDE.md bench-maxing:
    "serial-by-caution is halting in disguise"). Work unit = one
    fixture item (all its variants/models) so each Arborist call gets
    a per-unit qa_db — zero SQLite write contention; shards are
    read-only and safe under concurrent readers.
  * incremental: every record is appended+flushed the instant it
    lands. `--report-only PATH.jsonl` aggregates a PARTIAL file with
    NO LLM/judge spend — that is how the loop checks in every ~7
    turns while the sweep is still running.

FRAMINGS, measured not chosen (all three from the AskUserQuestion):
  accuracy / grounding = two deterministic relabellings of the SAME
  judge verdicts (zero extra spend). faithfulness-ablation = a
  different instrument (judge vs the context each arm was given,
  Arborist-only); ticket §4b scopes it v2/ext-ii — NOT run here, so
  the omission is explicit not silent. Arborist×Qwen needs proof-path
  guided_json+extra_body surgery — the coupled follow-up, also not
  done in a measurement harness.

Gate: judge-self-test 4/4 in-script, once, before any cell (aborts on
fail). `--report-only` skips it (read-only aggregation, no judge).
"""
from __future__ import annotations

import argparse
import json
import os
import re
import sys
import threading
import time
from collections import Counter
from concurrent.futures import ThreadPoolExecutor, as_completed
from pathlib import Path

ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT))
from bench.control_ab import _SOLO_SYS, _descaffold, _gold  # noqa: E402
# Judge selection is now run-time via --judge {code,opus} so a sweep can
# collect data WITHOUT burning Opus quota (fox 2026-05-19). Both judges
# share the Verdict shape & verdict vocabulary so swapping is a pointer
# reassignment; nothing downstream changes. Default=code.
import bench.judge as _judge_opus  # noqa: E402
import bench.judge_code as _judge_code  # noqa: E402

_JUDGES = {
    "code": (_judge_code.judge, _judge_code.self_test,
             _judge_code.JUDGE_MODEL),
    "opus": (_judge_opus.judge, _judge_opus.self_test,
             _judge_opus.JUDGE_MODEL),
}

MODELS: dict[str, dict] = {
    "hermes": dict(
        endpoint="https://hermes.ai.unturf.com/v1",
        model="adamo1139/Hermes-3-Llama-3.1-8B-FP8-Dynamic",
        extra=None, max_tokens=512, reasoning=False),
    "qwen-think": dict(
        endpoint="https://qwen.ai.unturf.com/v1",
        model="Qwen3.6-27B-UD-Q4_K_XL.gguf",
        extra=None, max_tokens=1024, reasoning=True),
    "qwen-nothink": dict(
        endpoint="https://qwen.ai.unturf.com/v1",
        model="Qwen3.6-27B-UD-Q4_K_XL.gguf",
        extra={"chat_template_kwargs": {"enable_thinking": False}},
        max_tokens=512, reasoning=False),
}


def _lower_first(s: str) -> str:
    return s[:1].lower() + s[1:] if s else s


def _v_plain(q: str) -> str:
    return q


def _v_source_relative(q: str) -> str:
    return "According to the reference knowledge base, " + _lower_first(q)


def _v_as_of_corpus(q: str) -> str:
    q2 = re.sub(r"^who is\b", "As of 2010, who was", q, flags=re.I)
    return q2 if q2 != q else f"As of 2010: {q}"


VARIANTS = {
    "plain": _v_plain,
    "source_relative": _v_source_relative,
    "as_of_corpus": _v_as_of_corpus,
}


# ---------------------------------------------------------------- work

def _process_item(idx: int, it: dict, variants: list[str],
                  models: list[str], shards_dir: Path,
                  arborist_on: bool, arborist_ref: str,
                  ts: str, judge_fn=None,
                  skip_solo: bool = False,
                  arborist_answer_mode: str = "claim_lattice") -> list[dict]:
    """All variants × models for ONE fixture item. Self-contained:
    its own qa_db, its own clients — safe to run concurrently.

    ``judge_fn`` defaults to the code judge for back-compat with any
    direct callers that pre-date the --judge switch; main() passes
    the user's choice explicitly. ``skip_solo`` lets a caller run
    ONLY the arborist arm — useful when the solo data for the same
    model + fixture already exists in a prior sweep and the new run
    is just adding the arborist arm (fox 2026-05-19: 'we don't need
    to redo anything')."""
    if judge_fn is None:
        judge_fn = _judge_code.judge
    from arborist.qa.client import OpenAICompatibleClient
    from arborist.qa.query import query
    from bench.stock_v1 import assert_not_drifted as _assert_stock
    from bench.stock_v1 import policy_for as _stock_policy_for

    out: list[dict] = []
    q0 = it["question"]
    gold = _gold(shards_dir, it.get("shard", ""), it["target_root"])
    if not gold:
        return [{"arm": "skip", "i": idx, "question_orig": q0,
                 "reason": "no gold"}]

    clients: dict[str, OpenAICompatibleClient] = {}

    def cl(mkey: str) -> OpenAICompatibleClient:
        if mkey not in clients:
            clients[mkey] = OpenAICompatibleClient(
                base_url=MODELS[mkey]["endpoint"],
                api_key=os.environ.get("ARBORIST_LLM_API_KEY"))
        return clients[mkey]

    qa_db = Path("/tmp") / f"control_sweep_{ts}_{idx}.db"
    try:
        for variant in variants:
            q_asked = VARIANTS[variant](q0)
            # Solo arm is skipped when --skip-solo is set OR when the
            # caller passed no models. Either way the arborist arm
            # below still runs if arborist_on.
            for mkey in [] if skip_solo else models:
                cfg = MODELS[mkey]
                try:
                    ans = cl(mkey).chat_completion(
                        [{"role": "system", "content": _SOLO_SYS},
                         {"role": "user", "content": q_asked}],
                        model=cfg["model"],
                        max_tokens=cfg["max_tokens"],
                        extra_body=cfg["extra"])
                except Exception as e:  # noqa: BLE001
                    ans = f"[solo-error: {type(e).__name__}: {e}]"
                v = judge_fn(q_asked, _descaffold(ans), gold)
                out.append({"arm": "solo", "model": mkey,
                            "variant": variant, "i": idx,
                            "question_orig": q0,
                            "question_asked": q_asked,
                            "answer": (ans or "")[:1500],
                            "verdict": v.label,
                            "verdict_eff": v.label,
                            "arb_audit_mode": None,
                            "judge_raw": (v.raw or "")[:700]})
            if arborist_on:
                try:
                    # Pass through MODELS[arborist_ref]["extra"] so
                    # per-model knobs (e.g. Qwen's enable_thinking
                    # toggle under chat_template_kwargs) reach the
                    # synthesis call in the proof-path. Without this,
                    # an arborist-ref=qwen-nothink sweep would still
                    # emit reasoning tokens because the toggle never
                    # crosses query()'s boundary. JSON-schema extras
                    # are added inside query() via
                    # claim_lattice_structured_output_extras() and
                    # merge with this per-model extras dict.
                    reasoning_ref = bool(MODELS[arborist_ref].get("reasoning"))
                    # STOCK V.1 frozen substrate-ON policy for this answer
                    # mode (bench.stock_v1). Non-reasoning runs are
                    # drift-guarded — a mid-campaign DEFAULT_QUERY_POLICY
                    # edit halts the sweep rather than silently redefining
                    # substrate-ON. Reasoning refs (phase 3) layer the
                    # documented overrides below and diverge from the pin
                    # by design, so they skip the assert.
                    if not reasoning_ref:
                        _assert_stock(arborist_answer_mode)
                    arb_policy = _stock_policy_for(arborist_answer_mode)
                    # Reasoning models emit a multi-line trace before
                    # the JSON; the claim_lattice "\n\n" runaway-guard
                    # stop sequence (tuned for single-line Hermes JSON)
                    # truncates that trace to an EMPTY answer (measured
                    # 2026-05-20: arborist+qwen-think produced 100%
                    # empty → ABSTAINED). Clear the stop for reasoning
                    # refs so the JSON actually lands. (With json-schema
                    # grammar enforcement the reasoning trace is itself
                    # suppressed, so output is clean single-line JSON —
                    # but the stop must still be cleared or the first
                    # structural newline truncates it.)
                    if reasoning_ref:
                        arb_policy["claim_lattice_json_stop_sequences"] = []
                        # Reasoning burns 1300-3300 completion tokens on
                        # the (internal) reasoning trace BEFORE the JSON
                        # (measured 2026-05-20). DEFAULT_QUERY_POLICY's
                        # 512-token budget — fine for non-reasoning
                        # single-line JSON — guarantees finish_reason=
                        # 'length' mid-reasoning → EMPTY completion.
                        # A/B at large context: max_tokens=1024 → 4/4
                        # empty (all finish='length'); 4096 → 0/4 (used
                        # 1339-3295). 8192 = generous headroom. The JSON
                        # answer itself is tiny; the budget is entirely
                        # for the reasoning trace — which is exactly why
                        # reasoning costs ~20-50x qwen-nothink's tokens
                        # for this workload.
                        arb_policy["max_tokens"] = 8192
                    # Empty-output self-heal for reasoning refs: even with
                    # an adequate token budget, a reasoning model can
                    # occasionally still emit an empty completion under
                    # the json-schema grammar — a model-side artefact,
                    # NOT a real abstention. Retry (burning the cached
                    # empty) up to 3 attempts so an artefact-empty doesn't
                    # masquerade as ABSTAINED in the scorecard. Non-
                    # reasoning refs don't exhibit this (qwen-nothink
                    # phase 3 had 0
                    # spurious empties) so they take a single pass. We
                    # never fabricate — an answer that is still empty
                    # after retries is recorded as the empty it is.
                    max_attempts = 3 if reasoning_ref else 1
                    arb_raw, arb_mode = "", None
                    for attempt in range(max_attempts):
                        r = query(question=q_asked, qa_db=qa_db,
                                  chat_client=cl(arborist_ref),
                                  model_id=MODELS[arborist_ref]["model"],
                                  shards_dir=shards_dir,
                                  extra_body=MODELS[arborist_ref]["extra"],
                                  policy=arb_policy,
                                  burn_existing=(attempt > 0))
                        arb_raw = (r.get("raw_answer")
                                   or r.get("answer_text") or "")
                        arb_mode = r.get("audit_mode")
                        if arb_raw.strip():
                            break
                except Exception as e:  # noqa: BLE001
                    arb_raw = f"[arborist-error: {type(e).__name__}: {e}]"
                    arb_mode = "ERROR"
                va = judge_fn(q_asked, _descaffold(arb_raw), gold)
                eff = ("ABSTAINED" if arb_mode == "UNGROUNDED"
                       and va.label in ("WRONG", "FABRICATED")
                       else va.label)
                out.append({"arm": "arborist", "model": arborist_ref,
                            "variant": variant, "i": idx,
                            "question_orig": q0,
                            "question_asked": q_asked,
                            "answer": (arb_raw or "")[:1500],
                            "verdict": va.label, "verdict_eff": eff,
                            "arb_audit_mode": arb_mode,
                            "judge_raw": (va.raw or "")[:700]})
    finally:
        for c in clients.values():
            try:
                c.close()
            except Exception:  # noqa: BLE001
                pass
        try:
            qa_db.unlink(missing_ok=True)
        except Exception:  # noqa: BLE001
            pass
    return out


# -------------------------------------------------------------- report

def _bucket(c: Counter) -> str:
    tot = sum(c.values()) or 1
    return (f"CG={c['CORRECT_GROUNDED']} W={c['WRONG']} "
            f"F={c['FABRICATED']} A={c['ABSTAINED']} "
            f"JE={c['JUDGE_ERROR']}  (n={tot})")


def _load_recs(jsonl_path: Path) -> list[dict]:
    """Tolerant: a killed-mid-write run can leave a truncated final
    line — skip it rather than crash the aggregator/resume scan."""
    recs = []
    for ln in jsonl_path.read_text().splitlines():
        if not ln.strip():
            continue
        try:
            recs.append(json.loads(ln))
        except json.JSONDecodeError:
            continue  # truncated trailing line from a kill
    return recs


def _aggregate(jsonl_path: Path):
    # Dedupe by (i, arm, model, variant), last-wins: a resume re-runs
    # any item that was incomplete when the prior run was killed, so
    # the same cell can appear twice — the later (complete-run) record
    # is authoritative. Without this, a restart double-counts.
    latest: dict[tuple, dict] = {}
    for r in _load_recs(jsonl_path):
        if r.get("arm") == "skip":
            continue
        k = (r["i"], r["arm"], r["model"], r["variant"])
        latest[k] = r
    tally: dict[tuple, Counter] = {}
    seen_items: set = set()
    for r in latest.values():
        seen_items.add(r["i"])
        key = (r["arm"], r["model"], r["variant"])
        tally.setdefault(key, Counter())[r["verdict_eff"]] += 1
    return tally, len(seen_items), len(latest)


def _report(tally, n_items, n_recs, args, jsonl_path, ts,
            done: bool) -> str:
    L = [f"# #000057 control sweep — {ts}"
         f"{'' if done else '  (INTERIM — sweep still running)'}",
         "",
         f"Fixture `{args.fixture}` · items scored so far={n_items} · "
         f"records={n_recs}",
         f"Raw per-item JSONL: `{jsonl_path}`",
         "",
         "**Review input, not a verdict.** fox: measure all "
         "benchmarks, bring forward for review.",
         "",
         "Codes: CG=CORRECT_GROUNDED W=WRONG F=FABRICATED "
         "A=ABSTAINED JE=JUDGE_ERROR. Gold = fixed corpus-vintage "
         "article (~2010-2011, verified from the artefact, NOT the "
         "2003 dump CLAUDE.md names). "
         f"Judge = {getattr(args, 'judge', 'opus')} "
         "(see bench/judge_code.py for the code judge / "
         "bench/judge.py for the gated Opus judge).",
         "",
         "## Raw verdicts (framing-neutral)",
         "",
         "| arm | model | variant | verdicts |",
         "|-----|-------|---------|----------|"]
    for (arm, mkey, variant), c in sorted(tally.items()):
        L.append(f"| {arm} | {mkey} | {variant} | {_bucket(c)} |")
    # Framing B headline numbers, computed deterministically.
    L += ["",
          "## Framing B — grounding-fidelity (the defensible read)",
          "",
          "A = honest about ungroundedness (GOOD); W+F = ungrounded "
          "confident assertion; CG = matched the designated source. "
          "Solo abstention-rate per (model,variant):",
          ""]
    for (arm, mkey, variant), c in sorted(tally.items()):
        if arm != "solo":
            continue
        tot = sum(c.values()) or 1
        a = c["ABSTAINED"]
        wf = c["WRONG"] + c["FABRICATED"]
        L.append(f"- `{mkey}/{variant}`: abstain "
                 f"{a}/{tot} ({a/tot:.0%}) · ungrounded-assert "
                 f"{wf}/{tot} ({wf/tot:.0%})")
    L += ["",
          "## Framing A — accuracy (shown only so review sees why it "
          "misleads: it penalises a model for knowing post-corpus "
          "reality; \"Macron\" is world-true).",
          "",
          "## Not run here (explicit, not silent)",
          "- faithfulness-ablation (judge vs each arm's own context) "
          "— different instrument, ticket §4b v2/ext-ii;",
          "- Arborist×Qwen — proof-path guided_json+extra_body merge, "
          "coupled follow-up.",
          ""]
    return "\n".join(L) + "\n"


def _emit_report(jsonl_path: Path, args, ts, done: bool) -> Path:
    tally, n_items, n_recs = _aggregate(jsonl_path)
    md = _report(tally, n_items, n_recs, args, jsonl_path, ts, done)
    mdp = jsonl_path.with_suffix(".md")
    mdp.write_text(md)
    print(md)
    return mdp


# ----------------------------------------------------------------- main

def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--fixture",
                    default="bench/qa_questions_stale_map.json")
    ap.add_argument("--n", type=int, default=386,
                    help="solo (control) item count — the HUGE axis")
    ap.add_argument("--arborist-n", type=int, default=40,
                    help="Arborist A/B reference item count (fixed "
                         "small — power belongs on the control)")
    ap.add_argument("--max-workers", type=int, default=6)
    ap.add_argument("--shards-dir",
                    default=str(Path.home() / ".arborist" / "shards"))
    ap.add_argument("--models", default="hermes,qwen-think,qwen-nothink")
    ap.add_argument("--variants",
                    default="plain,source_relative,as_of_corpus")
    ap.add_argument("--arborist-ref", default="hermes")
    ap.add_argument("--arborist-answer-mode",
                    choices=["quote", "claim_lattice"],
                    default="claim_lattice",
                    help="STOCK V.1 substrate-ON answer shape for the "
                         "treatment arm (the two-cell config family). "
                         "Frozen bench.stock_v1 policy per mode.")
    ap.add_argument("--out-dir", default="bench/qa_results")
    ap.add_argument("--report-only", default="",
                    help="aggregate a (partial) JSONL with NO spend "
                         "and exit — the interim check-in path")
    ap.add_argument("--resume", default="",
                    help="append to an existing JSONL, skipping items "
                         "already COMPLETE in it (re-runs partial "
                         "items; aggregator dedupes last-wins) — for "
                         "restarting at higher --max-workers without "
                         "losing finished units")
    ap.add_argument("--skip-self-test", action="store_true")
    ap.add_argument("--judge", choices=sorted(_JUDGES.keys()),
                    default="code",
                    help="which judge to use. 'code' (default, 2026-05-19) "
                         "is deterministic / no LLM / no quota; 'opus' is "
                         "the original headless Opus judge (requires "
                         "ARBORIST_JUDGE_ENABLE=1 — gated to prevent "
                         "accidental quota burn).")
    ap.add_argument("--skip-solo", action="store_true",
                    help="skip the solo arm; run ONLY the arborist arm. "
                         "Use when the model's solo data already exists "
                         "from a prior sweep on the same fixture — avoids "
                         "redundant LLM spend.")
    a = ap.parse_args()
    judge_fn, judge_self_test, judge_model_id = _JUDGES[a.judge]

    if a.report_only:
        jp = Path(a.report_only)
        ts = jp.stem.replace("control_sweep_", "")
        done = not any(  # crude: running if a sweep proc still alive
            True for _ in [])
        _emit_report(jp, a, ts, done=False)
        return 0

    models = [m for m in a.models.split(",") if m in MODELS]
    variants = [v for v in a.variants.split(",") if v in VARIANTS]
    items = json.loads(Path(a.fixture).read_text())[:a.n]

    n_solo = 0 if a.skip_solo else len(models) * len(variants) * len(items)
    arb_units = min(a.arborist_n, len(items))
    n_arb = len(variants) * arb_units
    print(f"#000057 control sweep — fixture={a.fixture}")
    print(f"  solo N={len(items)}  arborist-ref N={arb_units}  "
          f"models={models}  variants={variants}  "
          f"workers={a.max_workers}")
    print(f"  judge={a.judge} ({judge_model_id})")
    judge_cost_note = ("0 LLM calls" if a.judge == "code"
                       else f"~{n_solo + n_arb} claude -p judge")
    print(f"  bounded spend: {n_solo} solo + {n_arb} Arborist + "
          f"{judge_cost_note} + 4 self-test")

    if a.skip_self_test:
        print("  !! self-test SKIPPED — verdicts UNTRUSTED")
    else:
        print(f"  judge self-test ({a.judge}; instrument gate) …")
        if judge_self_test() != 0:
            print(f"ABORT: judge ({a.judge}) unreliable")
            return 1

    skip_items: set[int] = set()
    if a.resume:
        outp = Path(a.resume)
        ts = outp.stem.replace("control_sweep_", "")
        exp_solo = len(models) * len(variants)
        per_item: dict[int, Counter] = {}
        for r in _load_recs(outp):
            if r.get("arm") == "skip":
                continue
            per_item.setdefault(r["i"], Counter())[r["arm"]] += 1
        for i in range(1, len(items) + 1):
            c = per_item.get(i)
            if not c:
                continue
            need_arb = len(variants) if i <= arb_units else 0
            if c["solo"] >= exp_solo and c["arborist"] >= need_arb:
                skip_items.add(i)
        open_mode = "a"
        print(f"  RESUME {outp}: {len(skip_items)} complete items "
              f"skipped, {len(items)-len(skip_items)} to run "
              f"(dedupe is last-wins at aggregation)")
    else:
        ts = time.strftime("%Y-%m-%dT%H-%M-%SZ", time.gmtime())
        outp = Path(a.out_dir) / f"control_sweep_{ts}.jsonl"
        open_mode = "w"
    outp.parent.mkdir(parents=True, exist_ok=True)
    shards_dir = Path(a.shards_dir)
    lock = threading.Lock()
    t0 = time.time()
    done_units = 0

    with open(outp, open_mode) as log, ThreadPoolExecutor(
            max_workers=a.max_workers) as ex:
        futs = {
            ex.submit(_process_item, i, it, variants, models,
                      shards_dir, i <= arb_units, a.arborist_ref,
                      ts, judge_fn, a.skip_solo,
                      a.arborist_answer_mode): i
            for i, it in enumerate(items, 1)
            if i not in skip_items
        }
        for fut in as_completed(futs):
            i = futs[fut]
            try:
                recs = fut.result()
            except Exception as e:  # noqa: BLE001
                recs = [{"arm": "skip", "i": i,
                         "reason": f"unit-crash {type(e).__name__}: {e}"}]
            with lock:
                for rec in recs:
                    log.write(json.dumps(rec, ensure_ascii=False) + "\n")
                log.flush()
                done_units += 1
            el = time.time() - t0
            to_run = len(items) - len(skip_items)
            rate = done_units / el if el else 0
            eta = (to_run - done_units) / rate / 60 if rate else 0
            print(f"  unit {done_units}/{to_run} (item {i}) "
                  f"· {el/60:.1f}m elapsed · ETA {eta:.0f}m")

    print(f"\nsweep complete in {(time.time()-t0)/60:.1f}m")
    mdp = _emit_report(outp, a, ts, done=True)
    print(f"review deliverables:\n  {outp}\n  {mdp}")
    return 0


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