"""
Characterize the base policy from the generation pool.

The pool is 1000 prompts x 16 samples = 16,000 stories, which is 20x the
held-out eval set. It is the most statistically solid picture of baseline mode
collapse in the whole study, so it gets its own report rather than being used
only as DPO feed.

Answers:
  - How collapsed is the base model, per prompt? (deviation, log-det, eff. rank)
  - Is collapse uniform, or are some prompts far worse than others?
  - Does the judge's quality correlate with diversity? (i.e. is there really a
    quality-diversity tension to trade off, or are they independent?)
  - What does the quality distribution look like -- does tau=5 / rho=6 bite?
"""
from __future__ import annotations

import argparse
import json
import sys
from collections import defaultdict
from pathlib import Path

import numpy as np

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


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--tag", default="4b")
    ap.add_argument("--split", default="train")
    args = ap.parse_args()

    import logbook
    from diversity import effective_rank, logdet_volume, pairwise_deviation

    pdir = ROOT / "outputs" / f"pool_{args.tag}"
    rows = [json.loads(l) for l in open(pdir / f"pool_{args.split}.jsonl") if l.strip()]
    E = np.load(pdir / f"emb_{args.split}.npy").astype(np.float64)
    summary = json.load(open(pdir / f"summary_{args.split}.json"))

    by = defaultdict(list)
    for i, r in enumerate(rows):
        by[r["prompt_id"]].append(i)

    per = []
    for pid, ids in by.items():
        sub = E[ids]
        q = np.array([rows[i]["quality"] for i in ids])
        gp = np.array([rows[i]["gate_passed"] for i in ids])
        per.append({
            "prompt_id": pid,
            "dev": float(pairwise_deviation(sub).mean()),
            "logdet": float(logdet_volume(sub)),
            "eff_rank": float(effective_rank(sub)),
            "quality": float(q[gp].mean()) if gp.any() else 0.0,
            "gate_pass": float(gp.mean()),
            "n": len(ids),
        })

    dev = np.array([p["dev"] for p in per])
    ld = np.array([p["logdet"] for p in per])
    er = np.array([p["eff_rank"] for p in per])
    ql = np.array([p["quality"] for p in per])
    qual_all = np.array([r["quality"] for r in rows if r["gate_passed"]])
    N = per[0]["n"]

    # quality-diversity correlation across prompts
    def corr(a, b):
        if a.std() < 1e-9 or b.std() < 1e-9:
            return 0.0
        return float(np.corrcoef(a, b)[0, 1])

    # --- how much signal do the E1 / E2 channels actually carry? -----------
    # GDPO z-scores each reward channel WITHIN its group, so a channel with tiny
    # within-group spread gets amplified to unit variance regardless. If the
    # within-group ordering of d_i is not meaningful, E1 is largely learning
    # from amplified noise. Compare within-group spread against between-group
    # spread: a ratio far below 1 means the channel mostly encodes "which prompt
    # is this", which per-group normalization deliberately removes.
    dev_within, marg_within = [], []
    for pid, ids in by.items():
        sub = E[ids]
        d = pairwise_deviation(sub)
        from diversity import marginal_contributions
        m = marginal_contributions(sub)
        dev_within.append(d.std())
        marg_within.append(m.std())
    dev_within = np.array(dev_within); marg_within = np.array(marg_within)

    stats = {
        "n_prompts": len(per), "n_per_prompt": N, "n_stories": len(rows),
        "dev_within_group_sd": float(dev_within.mean()),
        "dev_between_group_sd": float(dev.std()),
        "dev_within_over_between": float(dev_within.mean() / max(dev.std(), 1e-9)),
        "marginal_within_group_sd": float(marg_within.mean()),
        "gate_pass_rate": summary["gate_pass_rate"],
        "ends_cleanly_rate": summary["ends_cleanly_rate"],
        "quality_mean": float(qual_all.mean()), "quality_sd": float(qual_all.std()),
        "quality_p10": float(np.percentile(qual_all, 10)),
        "quality_median": float(np.median(qual_all)),
        "quality_p90": float(np.percentile(qual_all, 90)),
        "frac_quality_ge_5(tau)": float((qual_all >= 5).mean()),
        "frac_quality_ge_6": float((qual_all >= 6).mean()),
        "frac_quality_ge_7(rho)": float((qual_all >= 7).mean()),
        "deviation_mean": float(dev.mean()), "deviation_sd": float(dev.std()),
        "deviation_p10": float(np.percentile(dev, 10)),
        "deviation_p90": float(np.percentile(dev, 90)),
        "logdet_mean": float(ld.mean()), "logdet_sd": float(ld.std()),
        "eff_rank_mean": float(er.mean()), "eff_rank_sd": float(er.std()),
        "eff_rank_p10": float(np.percentile(er, 10)),
        "eff_rank_p90": float(np.percentile(er, 90)),
        "eff_rank_ceiling": N,
        "corr(quality, deviation)": corr(ql, dev),
        "corr(quality, eff_rank)": corr(ql, er),
        "corr(deviation, eff_rank)": corr(dev, er),
    }

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

    fig, ax = plt.subplots(1, 4, figsize=(19, 4.3))
    ax[0].hist(qual_all, bins=np.arange(-0.25, 10.75, 0.5), color="#2980b9",
               edgecolor="white")
    ax[0].axvline(5, c="crimson", ls="--", label="tau=5 (diversity gate)")
    ax[0].axvline(7, c="darkorange", ls="--", label="rho=7 (DivPO, swept)")
    ax[0].set_title("Judge quality, all gate-passing stories")
    ax[0].set_xlabel("quality"); ax[0].legend(fontsize=7); ax[0].grid(alpha=.3)

    ax[1].hist(er, bins=30, color="#8e44ad", edgecolor="white")
    ax[1].axvline(N, c="green", ls="--", label=f"ceiling = {N}")
    ax[1].set_title(f"Effective rank per prompt\n(1 = total collapse, {N} = orthogonal)")
    ax[1].set_xlabel("effective rank"); ax[1].legend(fontsize=7); ax[1].grid(alpha=.3)

    ax[2].hist(dev, bins=30, color="#16a085", edgecolor="white")
    ax[2].set_title("Mean pairwise distance per prompt")
    ax[2].set_xlabel("1 - cos"); ax[2].grid(alpha=.3)

    ax[3].scatter(er, ql, s=6, alpha=.35, color="#c0392b")
    ax[3].set_xlabel("effective rank"); ax[3].set_ylabel("mean judge quality")
    ax[3].set_title(f"Quality vs diversity across prompts\nr = {stats['corr(quality, eff_rank)']:+.3f}")
    ax[3].grid(alpha=.3)

    plt.tight_layout()
    figp = logbook.FIGS / f"03_pool_{args.tag}_baseline.png"
    plt.savefig(figp, dpi=140)
    plt.close()

    worst = sorted(per, key=lambda p: p["eff_rank"])[:5]
    best = sorted(per, key=lambda p: -p["eff_rank"])[:5]

    r_qd = stats["corr(quality, eff_rank)"]
    tension = ("a genuine quality-diversity **tension**" if r_qd < -0.15 else
               "quality and diversity are **largely independent**" if abs(r_qd) <= 0.15 else
               "quality and diversity are **positively** related")

    body = f"""# Baseline characterization — {args.tag} pool ({len(rows):,} stories)

The pool is {stats['n_prompts']} prompts x {N} samples from the **base policy**,
20x the held-out eval set. This is the most statistically solid picture of
baseline mode collapse in the study, so it is reported in its own right rather
than treated only as DPO feed.

## Summary

{logbook.table([{"metric": k, "value": v} for k, v in stats.items()])}

## Findings

**Baseline collapse is severe.** Mean effective rank is
**{stats['eff_rank_mean']:.2f} out of a ceiling of {N}** — the {N} samples for a
given prompt span only ~{stats['eff_rank_mean']:.1f} effective directions. Mean
pairwise distance is {stats['deviation_mean']:.3f}, i.e. same-prompt stories sit
at ~{1-stats['deviation_mean']:.2f} cosine similarity. This is the thing every
arm is trying to move.

**Collapse is not uniform across prompts.** Effective rank runs from
{stats['eff_rank_p10']:.2f} (p10) to {stats['eff_rank_p90']:.2f} (p90), so some
prompts admit far more variation than others. Per-group normalization (GDPO)
handles this correctly: each prompt's advantage is computed within its own
group, so an intrinsically constrained prompt does not drag the update.

**Quality vs diversity: r = {r_qd:+.3f}** across prompts, so {tension}. This
matters for reading the frontier: if the correlation is near zero, then a method
that raises diversity without lowering quality is not defying a tradeoff, it is
exploiting slack that was already there.

**How much signal does E1's channel carry?** Within-group SD of `d_i` is
{stats['dev_within_group_sd']:.4f} against a between-group SD of
{stats['dev_between_group_sd']:.4f} — a ratio of
**{stats['dev_within_over_between']:.2f}**. This matters because GDPO z-scores
each channel *within* its group, so whatever within-group spread exists is
amplified to unit variance. A low ratio means most of `d_i`'s variation encodes
*which prompt this is* rather than *which sample is the odd one out* — and
per-group normalization deliberately discards exactly the former. Read E1's
result with this number in mind: if E1 underperforms, weak within-group
resolution is the first hypothesis, not a refutation of pairwise diversity as an
idea. The marginal channel's within-group SD is
{stats['marginal_within_group_sd']:.4f} on the raw log scale (it is z-scored
before use, so only its ordering matters).

**Threshold placement.** {100*stats['frac_quality_ge_5(tau)']:.1f}% of
gate-passing stories score >= tau=5 and
{100*stats['frac_quality_ge_7(rho)']:.1f}% score >= rho=7 (swept up from 6 because only {100*(1-stats['frac_quality_ge_6']):.1f}% fall below 6, starving DivPO of negatives on 311/1000 prompts). tau therefore acts as
a *floor* that only bites when a story is genuinely bad, which is its intent
(anti-gaming, not selection). rho is the more selective threshold and determines
DivPO's skip rate.

## Most collapsed prompts

{logbook.table(worst, ["prompt_id", "eff_rank", "dev", "logdet", "quality"])}

## Most diverse prompts

{logbook.table(best, ["prompt_id", "eff_rank", "dev", "logdet", "quality"])}

![pool baseline](../figures/03_pool_{args.tag}_baseline.png)
"""
    p = logbook.write_report(f"03_pool_{args.tag}_baseline", body)
    print(json.dumps(stats, indent=1))
    print("report ->", p, "\nfigure ->", figp)
    json.dump({"stats": stats, "per_prompt": per},
              open(logbook.LOGS / f"pool_{args.tag}_analysis.json", "w"), indent=1)
    logbook.note(f"pool analysis ({args.tag})",
                 f"eff_rank {stats['eff_rank_mean']:.2f}/{N}, "
                 f"dev {stats['deviation_mean']:.3f}, "
                 f"corr(q,div) {r_qd:+.3f}")
    return 0


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