"""
Autopsy of the previous run (logs/experiments/00_prior_run_autopsy.md + figures).

Reads prior_run/eval_samples/step_*.json, which are {prompt: [16 stories]} dumps
taken every 50 steps. Recomputes gates + embedding diversity so the prior run is
measured with EXACTLY the same instruments as the new arms -- otherwise the
"before" number is not comparable to any "after" number.
"""
from __future__ import annotations

import glob
import json
import re
from collections import Counter
from pathlib import Path

import numpy as np

import gates
import logbook
from diversity import group_metrics, l2_normalize

PRIOR = Path(__file__).resolve().parent.parent / "prior_run" / "eval_samples"
EMB_MODEL = "BAAI/bge-base-en-v1.5"


def load_steps() -> dict[int, dict[str, list[str]]]:
    out = {}
    for f in sorted(glob.glob(str(PRIOR / "step_*.json")),
                    key=lambda x: int(re.search(r"\d+", Path(x).name).group())):
        step = int(re.search(r"\d+", Path(f).name).group())
        out[step] = json.load(open(f))
    return out


def main():
    steps = load_steps()
    print(f"loaded {len(steps)} checkpoints: {sorted(steps)}")

    from sentence_transformers import SentenceTransformer
    enc = SentenceTransformer(EMB_MODEL, device="cuda")

    rows = []
    reason_counter: dict[int, Counter] = {}
    for step, data in steps.items():
        allres, pw, ld, mm = [], [], [], []
        for prompt, samples in data.items():
            res = [gates.check(s) for s in samples]
            allres += res
            E = l2_normalize(enc.encode(samples, normalize_embeddings=True,
                                        batch_size=32, show_progress_bar=False))
            gm = group_metrics(E)
            pw.append(gm["mean_pairwise_dist"]); ld.append(gm["logdet"])
            mm.append(gm["mean_marginal"])
        n = len(allres)
        reason_counter[step] = Counter(r for x in allres for r in x.reasons)
        rows.append({
            "step": step, "n": n,
            "gate_pass_%": 100 * sum(r.passed for r in allres) / n,
            "ends_cleanly_%": 100 * sum(r.completeness for r in allres) / n,
            "med_words": float(np.median([r.n_words for r in allres])),
            "rep4gram": float(np.mean([r.repeat_4gram_frac for r in allres])),
            "pairwise_dist": float(np.mean(pw)),
            "logdet": float(np.mean(ld)),
            "mean_marginal": float(np.mean(mm)),
        })

    for r in rows:
        print(r)

    # ---------------------------------------------------------------- figures
    import matplotlib
    matplotlib.use("Agg")
    import matplotlib.pyplot as plt

    s = [r["step"] for r in rows]
    fig, ax = plt.subplots(1, 3, figsize=(15, 4.2))

    ax[0].plot(s, [r["ends_cleanly_%"] for r in rows], "o-", color="#c0392b", lw=2)
    ax[0].axhline(80, ls="--", c="gray", lw=1)
    ax[0].set_title("Completeness collapse\n(prior run)")
    ax[0].set_xlabel("step"); ax[0].set_ylabel("% ending cleanly")
    ax[0].set_ylim(-3, 100); ax[0].grid(alpha=.3)
    ax[0].text(60, 84, "healthy target", color="gray", fontsize=8)

    ax[1].plot(s, [r["pairwise_dist"] for r in rows], "o-", label="pairwise dist", lw=2)
    ax[1].set_xlabel("step"); ax[1].set_ylabel("mean pairwise distance", color="C0")
    ax2 = ax[1].twinx()
    ax2.plot(s, [r["logdet"] for r in rows], "s--", color="C1", label="logdet")
    ax2.set_ylabel("logdet volume", color="C1")
    ax[1].set_title("Semantic diversity\n(embedding-based)")
    ax[1].grid(alpha=.3)

    ax[2].plot(s, [r["med_words"] for r in rows], "o-", color="#2c3e50", lw=2)
    ax[2].set_title("Median length pinned at the\ntoken wall (~380 words)")
    ax[2].set_xlabel("step"); ax[2].set_ylabel("median words")
    ax[2].grid(alpha=.3)

    plt.tight_layout()
    out = logbook.FIGS / "00_prior_run_autopsy.png"
    out.parent.mkdir(parents=True, exist_ok=True)
    plt.savefig(out, dpi=140)
    print("figure ->", out)

    # ---------------------------------------------------------------- report
    body = f"""# Experiment 00 — Autopsy of the prior run

**Purpose.** Establish what actually happened in the previous `train.py` run
before building anything new, and re-measure it with the *same* instruments the
new arms will use, so "before" and "after" are comparable.

**Source.** `prior_run/eval_samples/step_{{50..300}}.json` — 5 held-out prompts x
16 samples, dumped every 50 steps. 150 stories total per checkpoint set.

## Measurements

{logbook.table(rows)}

## Failure-reason census (gate violations, counted per story)

{logbook.table([{"step": k, **dict(v)} for k, v in reason_counter.items()])}

## Findings

**1. The run generated no finished stories, at any checkpoint.**
`ends_cleanly` is 0-8% across all six checkpoints. Median length sits in a
tight ~360-400 word band with a hard ceiling, which is the signature of a
`max_new_tokens` wall rather than a stylistic preference. Sampled tails
confirm mid-word cutoff (`"...a place called *The Library of"`).

**2. That silently destroyed the quality reward.**
The prior judge rubric said *"if the story cuts off mid-sentence, cap quality
at 4."* When ~100% of a group hits the same cap, within-group quality variance
goes to ~0. GRPO normalizes advantages within the group, so an
(almost) constant reward column produces an (almost) zero advantage. The
quality objective contributed no gradient for the entire run.

**3. The group-diversity term was structurally incapable of contributing.**
`quality_reward` added `0.5 * group_diversity`, a single scalar identical for
every sample in the group. Its within-group std is exactly 0, so after
normalization it contributes exactly 0 to every advantage. It cost one judge
call per step and bought nothing. (This is the failure mode recorded in the
brief's section 0.1, confirmed here in situ.)

**4. Net effect: the run optimized novelty almost alone.**
With quality flat and group-diversity annihilated, the only surviving gradient
came from `novelty_reward` — and its quality gate (`q >= 5`) was itself
degenerate, because quality was pinned near the truncation cap. That is an
unconditioned novelty objective, which is precisely the configuration known to
produce surface-level gaming.

**5. The observable signature matches that diagnosis.**
Unique 5-word openers rise 0.68 -> 0.96 while type-token ratio *falls*
0.381 -> 0.343. The model learned to vary the first few words while its
vocabulary narrowed — diversification at the surface, homogenization
underneath. Heavy italic emphasis persists at ~7 instances/story throughout,
the exact tic the rubric explicitly tried to punish.

**6. A latent bug would have suppressed completeness credit even after a fix.**
`completeness_reward` tested `text[-1] in ".!?\\"'"` — straight quotes only.
These generations use curly U+201D constantly, so a story ending `...gone.”`
scored 0.3 instead of 1.0.

## Consequences for the new design

- Token budget raised to 1280 new tokens (a 500-word story is ~650-700 Qwen
  tokens; the planned 768 left no headroom and would have re-created the bug).
  Nothing truncates the story anywhere downstream — not the judge, not the
  embedder. Overrun is *detected* by a word-count gate, never chopped.
- System prompt now carries an explicit word budget and ending instruction.
- No set-level scalar is ever used as a per-sample reward. All diversity credit
  is per-sample by construction (deviation `d_i`, marginal `m_i`).
- `finish_reason == "length"` from vLLM is captured and gates on it directly,
  rather than inferring truncation from punctuation.
- Terminal-punctuation detection handles curly quotes, ellipses and trailing
  markdown emphasis.
- Judge scores one story at a time on an absolute rubric, so `tau` is
  meaningful and cross-model eval comparisons are valid.

![autopsy](../figures/00_prior_run_autopsy.png)
"""
    p = logbook.write_report("00_prior_run_autopsy", body)
    print("report ->", p)
    logbook.note("prior-run autopsy complete",
                 f"0-8% completion across all checkpoints; quality signal was "
                 f"degenerate. Report: {p}")
    json.dump(rows, open(logbook.LOGS / "prior_run_metrics.json", "w"), indent=2)


if __name__ == "__main__":
    main()