#!/usr/bin/env python3
"""remeasure.py — re-measure self-report faithfulness under the JACOBIAN lens, and compare it to
the logit lens.

The first white-box number we published (Yi-34B: mean 0.103, median 0.000, concealment 20/20) was
measured with a LOGIT lens, which is not the instrument nemo's brief specifies and whose raw top-k is
dominated by attention-sink tokens. This runs the harness's own validator against the JACOBIAN
readouts so the headline is re-measured with the right instrument before it is quoted again.

For each model: build the GEO results CSV from the raw sweep (mapping the raw run_label to the lens
run_label — they differ, and the join is exact), run criteria_layer.validation, and report.

  python3 remeasure.py --lens-root lens_jac --raw raw/geo_all.jsonl --harness <path to jlens_lab_harness>
"""
from __future__ import annotations
import argparse, csv, glob, json, os, re, statistics as st, subprocess, sys, tempfile

# Below this many graded questions a row is marked provisional: a mean over 1-4 questions is not
# comparable to a mean over 20, and in a markdown table both render as a bare float.
MIN_N = 5

# lens label -> the run_label the raw sweep used (they were named independently; the join is exact)
RAW_LABEL = {
    "yi-1.5-34b": "yi-1.5-34b-chat",
    "llama-3.2-1b": "llama-3.2-1b-instruct",
    "llama-3.2-3b": "llama-3.2-3b-instruct",
    "seed-oss-36b": "seed-oss-36b-instruct",
    "qwen2.5-7b": "qwen2.5-7b-instruct",
    "qwen2.5-14b": "qwen2.5-14b-instruct",
    "qwen2.5-32b": "qwen2.5-32b-instruct",
    "phi-4": "phi-4",
    # lens labels that differ from the raw sweep's label (the join is exact, so map explicitly)
    "nemotron-nano-8b": "llama-3.1-nemotron-nano-8b-v1",
    "nemotron-3-nano-4b": "nvidia-nemotron-3-nano-4b-bf16",
    "deepseek-r1-distill-32b": "deepseek-r1-distill-qwen-32b",
    "deepseek-r1-qwen3-8b": "deepseek-r1-0528-qwen3-8b",
    "yi-coder-9b": "yi-coder-9b-chat",
    "internlm3-8b": "internlm3-8b-instruct",
    "hunyuan-7b": "hunyuan-7b-instruct",
    "baichuan-m2-32b": "baichuan-m2-32b",
    "hunyuan-4b": "hunyuan-4b-instruct",
    # Google lineage: the sweep labels carry the -it suffix, the lens labels do not.
    "gemma-2-2b": "gemma-2-2b-it",
    "gemma-2-9b": "gemma-2-9b-it",
    "gemma-2-27b": "gemma-2-27b-it",
    "gemma-1.1-2b": "gemma-1.1-2b-it",
    "gemma-1.1-7b": "gemma-1.1-7b-it",
    "medgemma-27b": "medgemma-27b-it",
    # Models lensed on the sparks. The probe runner derives run_label from the HF id
    # (m.split("/")[-1].lower()), NOT from the --out filename, so these must be spelled out or the
    # join skips the model with only a printed line to show for it.
    "mistral-7b": "mistral-7b-instruct-v0.3",
    "mistral-small-24b": "mistral-small-24b-instruct-2501",
    "granite-3.3-8b": "granite-3.3-8b-instruct",
    "aya-expanse-8b": "aya-expanse-8b",
    "qwen3-32b": "qwen3-32b",
    "llama-3.1-8b": "llama-3.1-8b-instruct",
    "olmo-2-7b": "olmo-2-1124-7b-instruct",
    "olmo-2-13b": "olmo-2-1124-13b-instruct",
    "olmo-2-32b": "olmo-2-0325-32b-instruct",
    "ornith-9b": "ornith-1.0-9b",
    "ornith-35b": "ornith-1.0-35b",
    "falcon3-7b": "falcon3-7b-instruct",
    # Phase 3.5 layer ablation: same raw rows, alternate lens layer selection.
    "olmo-2-7b-alt": "olmo-2-1124-7b-instruct",
    "olmo-2-13b-alt": "olmo-2-1124-13b-instruct",
    "gemma-2b": "gemma-2b-it",
    # C0 replicates: same model probed again with --label-suffix; graded against the SAME lens
    # dir (aliased on abzu) so only the probe run varies, which is exactly what C0 measures.
    "qwen3-4b-r2": "qwen3-4b-r2",
    "qwen3-4b-r3": "qwen3-4b-r3",
    "qwen3-8b-r2": "qwen3-8b-r2",
    "qwen3-8b-r3": "qwen3-8b-r3",
    "exaone-4.0-32b-r2": "exaone-4.0-32b-r2",
    "qwen2.5-14b-r2": "qwen2.5-14b-instruct-r2",
    "qwen3.8-27b": "qwen3.8-27b",
    "qwen2.5-3b": "qwen2.5-3b-instruct",
    "yi-1.5-6b-chat": "yi-1.5-6b-chat",
    "exaone-4.0-1.2b-r2": "exaone-4.0-1.2b-r2",
    "qwen3-30b-a3b": "qwen3-30b-a3b",
    "mistral-small-24b-r2": "mistral-small-24b-instruct-2501-r2",
    "falcon3-10b": "falcon3-10b-instruct",
    "glm-4.5-air": "glm-4.5-air",
    "qwen3-14b-spark": "qwen3-14b-r2",
    "exaone-4.0-32b-r3": "exaone-4.0-32b-r3",
    "granite-4.2-30b": "granite-4.2-30b",
    "qwen3-8b": "qwen3-8b",
    "phi-3.5-mini": "phi-3.5-mini-instruct",
    "internlm2.5-7b": "internlm2_5-7b-chat",
    # Qwen3 scale ladder, all six already on the NFS mount. Suffixed -abzu where a spark ran the
    # same model, so the two lens dirs cannot collide.
    "qwen3-0.6b": "qwen3-0.6b",
    "qwen3-1.7b": "qwen3-1.7b",
    "qwen3-4b": "qwen3-4b",
    "qwen3-8b-abzu": "qwen3-8b",
    "qwen3-14b": "qwen3-14b",
    "qwen3-32b-abzu": "qwen3-32b",
}

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--lens-root", default="lens_jac")
    ap.add_argument("--raw", default="raw/geo_all.jsonl")
    ap.add_argument("--harness", required=True)
    ap.add_argument("--out", default="JACOBIAN-REMEASURE.md")
    a = ap.parse_args()

    raw = [json.loads(l) for l in open(a.raw) if l.strip()]
    by = {}
    for r in raw:
        by.setdefault(r["run_label"], {})[r["query_id"]] = r

    rows_out = []
    for d in sorted(glob.glob(os.path.join(a.lens_root, "*/"))):
        lens_label = os.path.basename(d.rstrip("/"))
        qids = [os.path.basename(f)[:-5] for f in glob.glob(d + "Q*.json")]
        if len(qids) < 5:
            continue
        # A readout FILE is not a readout. medgemma-27b wrote 20 well-formed JSON files with an
        # EMPTY token_mass in every one; the sweep guard counts files, so it reported "done: 20
        # readouts", and the validator dutifully graded 20 rows against nothing and returned 0.000
        # on every metric with 0/20 concealed. That is not a model that conceals nothing — it is a
        # model that was never measured, and it would have been the striking result in the table.
        nonempty = 0
        for f in glob.glob(d + "Q*.json"):
            try:
                if json.load(open(f)).get("token_mass"):
                    nonempty += 1
            except Exception:
                pass
        if nonempty == 0:
            print(f"  [{lens_label}] {len(qids)} readout files, ALL WITH EMPTY token_mass — the "
                  f"lens produced nothing. Reported as failed, NOT as a score of zero.",
                  file=sys.stderr)
            rows_out.append({"model": lens_label, "n": 0, "graded": 0, "no_lens": None,
                             "post_mean": None, "post_median": None, "cold_mean": None,
                             "concealed": None, "confabulated": None, "unparseable": 0,
                             "empty_lens": len(qids)})
            continue
        if nonempty < len(qids):
            print(f"  [{lens_label}] WARNING: only {nonempty}/{len(qids)} readouts carry "
                  f"token_mass", file=sys.stderr)

        raw_label = RAW_LABEL.get(lens_label, lens_label)
        src = by.get(raw_label)
        if not src:
            print(f"  [{lens_label}] no raw rows under '{raw_label}' — skipped", file=sys.stderr)
            continue
        cols = ["query_id", "experiment_type", "condition", "run_label",
                "criteria", "criteria_cold", "answer_text"]
        csv_path = tempfile.mktemp(suffix=f"_{lens_label}.csv")
        n = unparseable = 0
        with open(csv_path, "w", newline="") as f:
            w = csv.DictWriter(f, fieldnames=cols); w.writeheader()
            for q in qids:
                r = src.get(q)
                if not r:
                    continue
                # The runner records an unparseable self-report as {"_parse_error": ..., "_raw": ...}
                # rather than a criteria LIST. The harness's canonicalize() assumes a list of dicts
                # and dies on it. That is not a grade of zero — the model emitted no structured
                # self-report at all, which is itself the finding for the small Llamas. Count these
                # separately and keep them out of the graded rows so they cannot be read as a score.
                if not (isinstance(r["criteria_posthoc"], list) and isinstance(r["criteria_cold"], list)):
                    unparseable += 1
                    continue
                w.writerow({"query_id": q, "experiment_type": "baseline", "condition": "baseline",
                            "run_label": lens_label,
                            "criteria": json.dumps(r["criteria_posthoc"]),
                            "criteria_cold": json.dumps(r["criteria_cold"]),
                            "answer_text": r["answer_text"]})
                n += 1
        if not n:
            print(f"  [{lens_label}] {unparseable}/{len(qids)} self-reports UNPARSEABLE — "
                  f"nothing gradable; reported as such, not as zero", file=sys.stderr)
            rows_out.append({"model": lens_label, "n": 0, "graded": 0, "no_lens": None,
                             "post_mean": None, "post_median": None, "cold_mean": None,
                             "concealed": None, "confabulated": None, "unparseable": unparseable})
            continue
        res = subprocess.run([sys.executable, "-m", "criteria_layer.validation",
                              "--results", csv_path, "--lens-dir", os.path.abspath(d)],
                             cwd=a.harness, capture_output=True, text=True)
        # A CRASHED validator is not a measurement of zero. This swallowed the validator's exit
        # code and stderr, so a missing `pandas` on the host came back as `graded=0` for every
        # model — a table of ten rows of None that looked like a finished run and would have
        # overwritten seven real results. Fail loudly instead: the run stops, and the reason is
        # the actual traceback rather than a plausible zero.
        if res.returncode != 0:
            print(f"\n!!! [{lens_label}] VALIDATOR FAILED (rc={res.returncode}) — this is a "
                  f"FAILURE, not a score of zero.\n{(res.stderr or '').strip()[-900:]}\n",
                  file=sys.stderr)
            sys.exit(4)
        graded = re.search(r"Rows graded:\s*(\d+)", res.stdout or "")
        nolens = re.search(r"no lens readout:\s*(\d+)", res.stdout or "")
        pj, cj, conc, conf = [], [], 0, 0
        try:
            for r in csv.DictReader(open(csv_path)):
                v = r.get("criteria_faithfulness_jaccard", "")
                if v not in ("", None): pj.append(float(v))
                v = r.get("criteria_cold_faithfulness_jaccard", "")
                if v not in ("", None): cj.append(float(v))
                if r.get("criteria_concealed", "[]") not in ("", "[]"): conc += 1
                if r.get("criteria_confabulated", "[]") not in ("", "[]"): conf += 1
        except Exception as e:
            print(f"  [{lens_label}] could not read graded CSV: {e}", file=sys.stderr)
        rows_out.append({
            "model": lens_label, "n": n,
            "graded": int(graded.group(1)) if graded else 0,
            "no_lens": int(nolens.group(1)) if nolens else None,
            "post_mean": round(st.mean(pj), 3) if pj else None,
            "post_median": round(st.median(pj), 3) if pj else None,
            "cold_mean": round(st.mean(cj), 3) if cj else None,
            "concealed": conc, "confabulated": conf, "unparseable": unparseable,
        })
        print(f"  [{lens_label}] graded={rows_out[-1]['graded']} "
              f"post_mean={rows_out[-1]['post_mean']} conc={conc}/{n}", file=sys.stderr, flush=True)

    with open(a.out, "w") as f:
        f.write("# Jacobian re-measurement — faithfulness under the lens nemo specified\n\n")
        f.write("The first white-box number (Yi-34B: mean 0.103 / median 0.000, concealment 20/20) was\n"
                "measured with a **logit** lens. This table re-measures with the **Jacobian** lens\n"
                "(`softmax(W_U · norm(J·h))`, JVP, eager attention) using the harness's own validator.\n\n")
        f.write("| model | n | graded | unparseable | posthoc mean | posthoc median | cold mean | concealed | confabulated |\n")
        f.write("|---|--:|--:|--:|--:|--:|--:|--:|--:|\n")
        for r in rows_out:
            # An ungraded model gets em-dashes, never "0/20". A zero in a concealment column reads
            # as "we looked and found no concealment"; these rows mean "there was nothing to look
            # at". Same glyph, opposite claims.
            if not r["graded"]:
                f.write(f"| {r['model']} | — | 0 | {r.get('unparseable', 0)} | — | — | — | — | — |\n")
            else:
                # A mean over one or two questions renders identically to a mean over twenty and
                # invites the reader to compare them. Mark the thin rows rather than trusting the
                # `n` column to be read — it sits three columns away from the number it qualifies.
                mark = " ⚠PROVISIONAL" if r["graded"] < MIN_N else ""
                f.write(f"| {r['model']}{mark} | {r['n']} | {r['graded']} | {r.get('unparseable', 0)} | "
                        f"{r['post_mean']} | {r['post_median']} | {r['cold_mean']} | "
                        f"{r['concealed']}/{r['n']} | {r['confabulated']}/{r['n']} |\n")
        f.write("\n**Preliminary and uncertified.** Open weights only; no claim about API models.\n")
        f.write(f"\n\u26a0PROVISIONAL marks a model graded on fewer than {MIN_N} questions — "
                "its mean is not comparable to the full rows and should not be ranked against "
                "them.\n")
        f.write("\n`unparseable` counts questions where the model emitted no parseable JSON "
                "self-report at all. Those rows are excluded from the grade rather than scored "
                "zero — a model that cannot produce a structured self-report has not been measured "
                "as perfectly unfaithful, it has not been measured.\n")
    print(f"\nwrote {a.out} — {len(rows_out)} models", file=sys.stderr)

if __name__ == "__main__":
    main()