"""Did the adapter move the voice TOWARD the held-out author? A seat-free relative measure.

Written for Yarros, since used on Brontë and Hemingway. The author is now a REQUIRED
argument rather than a hardcoded string — see the note on `--author` in main().

NOT the frozen adjudication. That needs a control-author panel (to place an absolute
band and a hard-negative sister), a seed-to-seed spread, and — for BEAT INCUMBENT — the
gen seat, none of which are available here. This answers the smaller, honest question
the operator can act on: of the arms generated on ONE harness, which sits closest to
the real held-out author, and does the adapter beat the base control?

Instrument: Burrows's Delta over CHARACTER BIGRAMS (hence delta_cb). Char bigrams are
dominated by function-word morphology and rhythm, not proper nouns, so the rename does
not move them. Reference profile is the HELD-OUT (val) split — text no arm was trained
on. Each arm's pooled generations are scored against it; lower = closer to the author.

Discipline: this is a RELATIVE reading (arms vs each other, same harness), never an
absolute-band claim. The A-vs-A floor below is the only thing that makes a between-arm
gap meaningful — half-vs-half of the held-out reference gives the distance the metric
returns for two samples of the SAME author, so a between-arm gap smaller than that floor
is not a finding.
"""
from __future__ import annotations
import argparse, json, re, sys, statistics as st
from collections import Counter
from pathlib import Path


def bigrams(text: str) -> Counter:
    t = re.sub(r"\s+", " ", text.lower())
    return Counter(t[i:i+2] for i in range(len(t) - 1))


def profile(text: str, keys: list[str]) -> dict:
    c = bigrams(text); n = sum(c.values()) or 1
    return {k: c.get(k, 0) / n for k in keys}


def delta(arm_text: str, ref_prof: dict, mu: dict, sd: dict, keys: list[str]) -> float:
    ap = profile(arm_text, keys)
    # Burrows's Delta = mean |z(arm) - z(ref)| over the shared feature set
    return st.mean(abs((ap[k] - mu[k]) / sd[k] - (ref_prof[k] - mu[k]) / sd[k]) for k in keys)


def main() -> int:
    # ⚠ --author IS REQUIRED, and that is the fix for a defect this script shipped with.
    # The reference label was hardcoded "Yarros". Run against Brontë it printed
    # "reference: held-out Yarros" over Brontë's numbers, and that output is now sitting
    # in a committed artifact saying the wrong author. A default would have kept the
    # silent-wrong-label failure and only moved it; naming the author is one word at the
    # call site and the label can no longer disagree with the data.
    ap = argparse.ArgumentParser()
    ap.add_argument("corpus", help="renamed corpus dir containing copies/ with split=val records")
    ap.add_argument("evaldir", help="dir of voice.<arm>.jsonl; the control arm's name must "
                                    "contain the substring `unadapted`")
    ap.add_argument("--author", required=True,
                    help="reference author label, e.g. Hemingway. Required: see above.")
    a = ap.parse_args()
    corpus = Path(a.corpus)
    evaldir = Path(a.evaldir)
    author = a.author
    # reference = held-out val text
    val = []
    for p in sorted((corpus / "copies").glob("*.jsonl")):
        for l in p.read_text(encoding="utf-8").splitlines():
            r = json.loads(l)
            if r.get("split") == "val":
                val.append(r["text"])
    # dedup identical val chapters across copies (renaming aside, the same chapter recurs)
    ref_text = "\n".join(dict.fromkeys(val))
    # feature set: the most frequent bigrams in the reference (stable, high-signal)
    keys = [k for k, _ in bigrams(ref_text).most_common(400)]
    # mu/sd across the val text split into chunks, for z-scoring
    words = ref_text.split()
    chunks = [" ".join(words[i:i+800]) for i in range(0, len(words), 800) if len(words[i:i+800]) > 200]
    profs = [profile(c, keys) for c in chunks]
    mu = {k: st.mean(p[k] for p in profs) for k in keys}
    sd = {k: (st.pstdev(p[k] for p in profs) or 1e-9) for k in keys}
    ref_prof = profile(ref_text, keys)

    # SAME-AUTHOR REFERENCE (the target, not a significance threshold): two halves
    # of the held-out author. A perfect mimic scores about this; you cannot get closer
    # to the author than the author gets to itself at this sample size.
    half = len(words) // 2
    same_author = delta(" ".join(words[:half]), profile(" ".join(words[half:]), keys), mu, sd, keys)

    print(f"reference: held-out {author}, {len(words):,} words, {len(chunks)} chunks, {len(keys)} char-bigram features")
    print(f"same-author target (held-out {author} vs itself): delta_cb = {same_author:.3f}")
    print(f"  -> the floor of what any arm could reach; lower is more {author}-like, this is the best possible\n")

    def arm_texts(f):
        return [json.loads(l) for l in f.read_text(encoding="utf-8").splitlines()]

    rows = []
    for f in sorted(evaldir.glob("voice.*.jsonl")):
        arm = f.stem.replace("voice.", "")
        recs = arm_texts(f)
        allt = "\n".join(r["continuation"] for r in recs)
        d = delta(allt, ref_prof, mu, sd, keys)
        # within-arm sampling spread = the REAL noise floor for a between-arm gap:
        # split by seed and score each subset; the range is this metric's variance
        # at this sample size, measured rather than assumed.
        by_seed = {}
        for r in recs:
            by_seed.setdefault(r["seed"], []).append(r["continuation"])
        seed_ds = [delta("\n".join(v), ref_prof, mu, sd, keys) for v in by_seed.values() if len(v) > 2]
        spread = (max(seed_ds) - min(seed_ds)) if len(seed_ds) > 1 else float("nan")
        rows.append((arm, d, len(allt.split()), seed_ds, spread))

    print("  arm                 delta_cb   per-seed        [words]")
    for arm, d, w, sd_, spread in sorted(rows, key=lambda x: x[1]):
        seeds = " ".join(f"{x:.3f}" for x in sd_)
        print(f"  {arm:20s} {d:.3f}    ({seeds})   [{w}]  spread {spread:.3f}")

    # ⚠⚠ THE FLOOR IS PAIRWISE, and that is a RULE CHANGE made because the all-arms
    # rule decided lv-bronte. Measured there:
    #     base-unadapted  spread 0.062
    #     ckpt475         spread 0.092   <- the candidate that shipped
    #     ckpt925         spread 0.251   <- set the floor, on ONE outlier seed
    # ckpt475's +0.193 was failed by a floor contributed entirely by a THIRD arm nobody
    # was shipping. Run as base-vs-ckpt475 the floor is 0.092 and the same gap clears at
    # 2.1x. A candidate's verdict must not depend on which other arms you happened to
    # generate, so the comparison's floor is the larger of the TWO arms being compared.
    # The all-arms number is still printed, because lv-bronte's record used it and a
    # reader comparing the two runs needs both.
    floors = [r[4] for r in rows if r[4] == r[4]]
    noise_all = max(floors) if floors else float("nan")
    spread_of = {r[0]: r[4] for r in rows}

    base_row = next(((a, d) for a, d, _, _, _ in rows if "unadapted" in a), None)
    if base_row is None:
        print(f"\n  all-arms noise floor (largest within-arm seed spread): {noise_all:.3f}")
        print("  ⚠ no arm name contains `unadapted` -- no control identified, no verdict\n")
        return 0
    base_arm, base = base_row
    print(f"\n  all-arms noise floor (largest within-arm seed spread, lv-bronte's rule): {noise_all:.3f}")
    print(f"  PAIRWISE floor is the verdict: max(spread(candidate), spread({base_arm}) = "
          f"{spread_of[base_arm]:.3f})\n")

    print(f"  vs {base_arm} control (positive gap = moved toward {author}):")
    for arm, d, _, _, _ in sorted(rows, key=lambda x: x[1]):
        if arm == base_arm:
            continue
        gap = base - d
        pair_floor = max(spread_of[arm], spread_of[base_arm])
        verdict = (f"MOVED toward {author} ({gap / pair_floor:.1f}x the pairwise floor "
                   f"{pair_floor:.3f})" if gap > pair_floor
                   else f"within the pairwise floor {pair_floor:.3f} -- NOT a finding")
        flag = "" if (gap > noise_all) == (gap > pair_floor) else "   <- the two rules DISAGREE"
        print(f"    {arm:20s} {gap:+.3f}  ({verdict}){flag}")

    ordered = [a for a, *_ in sorted(rows, key=lambda x: x[1])]
    print(f"\n  ordering: {' < '.join(ordered)}  (lower = more {author}-like)")
    # ⚠ This used to assert "one seed-pair per arm" and claim corroboration from a
    # "Base < Instruct" held-out loss ordering. Both were Yarros-run facts hardcoded
    # as if they were properties of the instrument: by lv-bronte every arm carried
    # four seeds, and no Base-vs-Instruct comparison was in the run at all. Report
    # what this run actually has instead of a remembered one.
    nseeds = sorted({len(r[3]) for r in rows})
    print(f"  ⚠ RELATIVE reading on one harness: {nseeds if len(nseeds) > 1 else nseeds[0]} "
          f"seed group(s) per arm, scored against this corpus's own held-out split. "
          f"It is not an absolute-band claim and corroborates nothing on its own.")
    return 0


if __name__ == "__main__":
    sys.exit(main())