#!/usr/bin/env python3
"""exp15/Q4a — omission characterisation probe (sanity-gate follow-up). The first sanity pro-draft
OMITTED the opening ~13 source paragraphs (aperture/44%/丙 exposition); a fresh pro call was COMPLETE
-> stochastic omission at temp 0.3, NOT a capture artifact (translation is in content, not reasoning).
This probe samples K drafts of a chunk from BOTH models and measures, per meaning-trap, whether the
trap's phenomenon sentence is RENDERED (locatable) vs OMITTED — to decide seed count + whether omission
is a Q4a signal. $ persisted to the gen ledger (D30.10)."""
from __future__ import annotations
import argparse
import json
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "exp14"))
sys.path.insert(0, str(Path(__file__).resolve().parent))
import exp14_common as C
import exp14_prompts as P
import q4a_traps as Q

RECS = {(r["chapter"], r["chunk_idx"]): r
        for r in json.load(open("/home/ubuntu/books/gu-zhenren/rerun/records.json"))}
INJ = json.load(open("/home/ubuntu/books/gu-zhenren/exp14/injection_blocks.json"))
LEDGER = Path("/home/ubuntu/books/gu-zhenren/exp15/q4a_gen_costs.jsonl")
OUTF = Path("/home/ubuntu/books/gu-zhenren/exp15/q4a_omission_probe.json")


def run(chunks, k):
    traps = Q.traps()
    results = {}
    for cid in chunks:
        ch, ck = cid.split("."); src = RECS[(int(ch), int(ck))]["source"]
        gloss = INJ[f"{ch}/{ck}"]["bilingual_glossary"]
        msgs = C.messages_with_injection(P.translator_system(), gloss, P.translator_user(src))
        chunk_traps = [t for t in traps if t["chunk"] == cid]
        results[cid] = {"traps": [t["id"] for t in chunk_traps], "samples": {}}
        for model in ("deepseek-v4-flash", "deepseek-v4-pro"):
            rows = []
            for s in range(k):
                spec = C.spec(model, temp=0.3, max_tokens=8000)
                text, err, usage = C.call(spec, msgs)
                with LEDGER.open("a", encoding="utf-8") as f:
                    f.write(json.dumps({"cell": f"PROBE_{model}", "chunk": cid, "seed": s, "model": model,
                                        "usage": usage, "cost": C.cost_of(model, usage), "err": err},
                                       ensure_ascii=False) + "\n")
                if err:
                    rows.append({"seed": s, "err": err}); continue
                # per-trap: rendered (fix/fail) vs omitted ('?')
                verds = {t["id"]: Q.det_score(text, t)[0] for t in chunk_traps}
                rendered = sum(1 for v in verds.values() if v in ("fix", "fail"))
                rows.append({"seed": s, "chars": len(text), "finish": usage.get("finish_reason"),
                             "rendered": rendered, "n_traps": len(chunk_traps), "verds": verds})
                print(f"  {model:<18} {cid} s{s}: {len(text)}c finish={usage.get('finish_reason')} "
                      f"rendered {rendered}/{len(chunk_traps)}  {verds}")
            results[cid]["samples"][model] = rows
    OUTF.write_text(json.dumps(results, ensure_ascii=False, indent=2), encoding="utf-8")
    tot = sum(json.loads(l).get("cost", 0) for l in LEDGER.read_text().splitlines())
    print(f"\nprobe done -> {OUTF}; gen ledger total ${tot:.4f}")


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--chunks", nargs="+", default=["6.0", "7.0"])
    ap.add_argument("--k", type=int, default=5)
    a = ap.parse_args()
    run(a.chunks, a.k)


if __name__ == "__main__":
    main()