"""
Post-run analysis for one GRPO arm. Produces logs/experiments/<name>.md plus
figures, from outputs/<name>/reward_history.json.

Run this after EVERY training run, before launching the next one. The point is
to catch a dead or hacked reward channel while there is still time to change
something, rather than discovering it in the final report -- which is exactly
how the prior run wasted 300 steps.
"""
from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path

import numpy as np

ROOT = Path(__file__).resolve().parent.parent


def smooth(x, w=15):
    x = np.asarray(x, dtype=np.float64)
    if len(x) < w:
        return x
    k = np.ones(w) / w
    return np.convolve(x, k, mode="valid")


def trend(x, frac=0.25):
    """(early mean, late mean, delta) over the first/last `frac` of the run."""
    x = np.asarray(x, dtype=np.float64)
    n = max(1, int(len(x) * frac))
    a, b = float(x[:n].mean()), float(x[-n:].mean())
    return a, b, b - a


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--name", required=True)
    ap.add_argument("--title", default=None)
    args = ap.parse_args()

    import logbook

    run_dir = ROOT / "outputs" / args.name
    hist = json.load(open(run_dir / "reward_history.json"))
    if not hist:
        print("empty history"); return 1

    steps = np.arange(len(hist))
    series = {k: np.array([h[k] for h in hist], dtype=np.float64) for k in
              ("gate_pass", "ends_cleanly", "mean_quality_passing", "mean_novelty",
               "frac_above_tau", "mean_deviation", "mean_logdet", "mean_words",
               "frac_groups_degenerate")}

    tr = {k: trend(v) for k, v in series.items()}

    # --- reward-hacking diagnostic ---------------------------------------
    dev_a, dev_b, dev_d = tr["mean_deviation"]
    q_a, q_b, q_d = tr["mean_quality_passing"]
    g_a, g_b, g_d = tr["gate_pass"]
    hack = (dev_d > 0.02 and q_d < -0.5) or (g_b < 0.6 and g_d < -0.2)
    verdict = ("SUSPECTED REWARD HACKING" if hack else
               "healthy" if (g_b > 0.7 and q_d > -0.5) else "degraded but not hacking")

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

    # TRL-side metrics (entropy / KL) live in a separate history
    trl_path = run_dir / "trl_log_history.json"
    trl = json.load(open(trl_path)) if trl_path.exists() else []
    ent = np.array([h["entropy"] for h in trl if "entropy" in h], dtype=np.float64)
    kl = np.array([h["kl"] for h in trl if "kl" in h], dtype=np.float64)

    fig, ax = plt.subplots(3, 3, figsize=(16, 11.5))
    panels = [
        ("gate_pass", "Gate pass rate", "#c0392b", (0, 1.05)),
        ("ends_cleanly", "Ends cleanly", "#27ae60", (0, 1.05)),
        ("mean_quality_passing", "Judge quality (gate-passers)", "#2980b9", None),
        ("mean_deviation", "Mean pairwise deviation d_i", "#8e44ad", None),
        ("mean_logdet", "Group log-det volume", "#d35400", None),
        ("frac_above_tau", "Frac above tau (diversity-eligible)", "#16a085", (0, 1.05)),
        ("mean_words", "Mean story length (words)", "#7f8c8d", None),
        ("frac_groups_degenerate", "Groups with no valid sample", "#c0392b", (0, 1.05)),
    ]
    axes = ax.ravel()
    for a, (key, title, c, ylim) in zip(axes, panels):
        v = series[key]
        a.plot(steps, v, alpha=0.25, color=c, lw=0.8)
        sm = smooth(v)
        a.plot(steps[len(steps) - len(sm):], sm, color=c, lw=2)
        a.set_title(title, fontsize=10)
        a.set_xlabel("reward-batch"); a.grid(alpha=.3)
        if ylim:
            a.set_ylim(*ylim)

    # entropy panel: the creativity-death detector. Falling entropy means the
    # policy is becoming deterministic. Rising entropy is permissive, NOT proof
    # of creativity -- a policy can get noisier inside a single narrative mode --
    # so this is read alongside log-det, never as a substitute for it.
    a = axes[8]
    if ent.size:
        x = np.arange(ent.size)
        a.plot(x, ent, alpha=0.25, color="#e67e22", lw=0.8)
        sm = smooth(ent)
        a.plot(x[len(x) - len(sm):], sm, color="#e67e22", lw=2, label="entropy")
        e0, e1, ed = trend(ent)
        a.axhline(e0, ls=":", c="gray", lw=1)
        a.set_title(f"Per-token policy entropy  ({e0:.3f} → {e1:.3f}, Δ{ed:+.3f})",
                    fontsize=10)
        if kl.size:
            a2 = a.twinx()
            a2.plot(np.arange(kl.size), kl, color="#95a5a6", lw=1, alpha=.7)
            a2.set_ylabel("KL to ref", color="#95a5a6", fontsize=8)
    else:
        a.text(.5, .5, "no entropy logged\n(liger path?)", ha="center",
               va="center", transform=a.transAxes, color="crimson")
        a.set_title("Per-token policy entropy — MISSING", fontsize=10)
    a.set_xlabel("optimizer step"); a.grid(alpha=.3)
    fig.suptitle(f"{args.title or args.name} — training diagnostics ({verdict})",
                 fontsize=12)
    plt.tight_layout()
    figp = logbook.FIGS / f"{args.name}_diagnostics.png"
    figp.parent.mkdir(parents=True, exist_ok=True)
    plt.savefig(figp, dpi=140)
    plt.close()

    # --- gate failure census ---------------------------------------------
    from collections import Counter
    cnt = Counter()
    for h in hist:
        for k, v in (h.get("reasons") or {}).items():
            cnt[k] += v
    total_stories = sum(h["n"] for h in hist)

    rows = [{"metric": k, "early": round(a, 4), "late": round(b, 4),
             "delta": round(d, 4)} for k, (a, b, d) in tr.items()]
    if ent.size:
        a_, b_, d_ = trend(ent)
        rows.append({"metric": "policy_entropy", "early": round(a_, 4),
                     "late": round(b_, 4), "delta": round(d_, 4)})
        entropy_note = (
            f"Per-token policy entropy moved {a_:.4f} → {b_:.4f} "
            f"(Δ {d_:+.4f}, {100*d_/max(a_,1e-9):+.1f}%). "
            + ("**Entropy collapse** — the policy is becoming deterministic; "
               "treat any diversity gain reported below with suspicion."
               if d_ < -0.15 * a_ else
               "No entropy collapse. Note that entropy holding up is a "
               "necessary but not sufficient condition for diversity: it "
               "permits varied output without demonstrating it.")
        )
    else:
        entropy_note = ("Entropy was NOT logged for this run. TRL only emits it "
                        "on the non-liger loss path; check `use_liger_kernel`.")

    cost = {}
    cp = run_dir / "judge_cost.json"
    if cp.exists():
        cost = json.load(open(cp))

    body = f"""# {args.title or args.name} — training analysis

**Verdict: {verdict}**

Reward batches: {len(hist)} | stories scored: {total_stories}

## Trend (first 25% vs last 25% of the run)

{logbook.table(rows)}

## Gate failures (count over the whole run, {total_stories} stories)

{logbook.table([{"reason": k, "count": v, "pct_of_stories": round(100*v/max(1,total_stories), 2)}
                for k, v in cnt.most_common()]) if cnt else "_No gate failures._"}

## Policy entropy

{entropy_note}

Entropy is read as an *asymmetric* signal here. A large fall is strong evidence
that creativity is dying — the policy is collapsing toward deterministic output.
A rise is only permissive: a policy can raise per-token entropy while staying
inside one narrative mode (noisier word choice, same story). The prior run
demonstrated exactly that dissociation — surface variation up, semantic
diversity down. So entropy is never optimized, and diversity claims rest on
log-det / effective rank.

## Reward-hacking check

The signature we watch for is **diversity up while quality or validity goes
down** — the policy discovering it can farm the diversity channel by emitting
text that is different because it is worse.

- mean deviation: {dev_a:.4f} → {dev_b:.4f} (Δ {dev_d:+.4f})
- judge quality:  {q_a:.3f} → {q_b:.3f} (Δ {q_d:+.3f})
- gate pass:      {g_a:.3f} → {g_b:.3f} (Δ {g_d:+.3f})

Trip conditions: `Δdeviation > +0.02 AND Δquality < -0.5`, or
`gate_pass < 0.60 AND Δgate_pass < -0.20`. → **{"TRIPPED" if hack else "not tripped"}**

## Judge cost

```json
{json.dumps(cost, indent=1) if cost else "{}"}
```

![diagnostics](../figures/{args.name}_diagnostics.png)
"""
    p = logbook.write_report(args.name, body)
    print(f"verdict: {verdict}")
    for r in rows:
        print(f"  {r['metric']:26} {r['early']:>9.4f} -> {r['late']:>9.4f}  ({r['delta']:+.4f})")
    print("report ->", p)
    print("figure ->", figp)
    logbook.note(f"analysis: {args.name}", f"verdict={verdict}; report {p}")
    return 0


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