"""Project 5: does RL sharpen or expand?  pass@k curves before and after RL.

Yue et al. (2025) found that RL-trained models win at pass@1 but their base models catch up
or win at large k. Here: train the student with RL, then compare pass@k for k = 1..64 on
(a) in-distribution 2-digit problems and (b) 3-digit + 1-digit problems never seen in SFT.

Hypothesis: big gains at k = 1, shrinking gains as k grows (sharpening). Whether the curves
cross at large k is the open question. Saves project5_sharpen_or_expand.png.
"""
import argparse
import copy
import random

import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt  # noqa: E402
import numpy as np  # noqa: E402
import torch  # noqa: E402

from common import HERE, bucket_problem, get_student, rl_train, tools  # noqa: E402
from tinylm import decode_responses, make_problem, sample  # noqa: E402

ap = argparse.ArgumentParser()
ap.add_argument("--steps", type=int, default=300)
ap.add_argument("--seed", type=int, default=0)
args = ap.parse_args()
N, KS = 64, [1, 2, 4, 8, 16, 32, 64]


@torch.no_grad()
def passk_curve(model, bucket, n_prompts=60, seed=5):
    """bucket=None: the training distribution (both operands 0..99); else a common.bucket_problem bucket."""
    rng = random.Random(seed)
    pa = [make_problem(rng) if bucket is None else bucket_problem(rng, bucket) for _ in range(n_prompts)]
    P = [p for p, _ in pa for _ in range(N)]
    A = [a for _, a in pa for _ in range(N)]
    ok = np.array([y == a for y, a in zip(decode_responses(sample(model, P)), A)]).reshape(n_prompts, N).sum(1)
    return [float(np.mean([tools.pass_at_k(N, int(c), k) for c in ok])) for k in KS]


base = get_student(300)
rl = copy.deepcopy(base)
rl_train(rl, args.steps, adv="mean", seed=args.seed, eval_every=100, tag="rl")
fig, axes = plt.subplots(1, 2, figsize=(10, 4))
for ax, (b, name) in zip(axes, [(None, "2-digit (in distribution)"), (4, "3-digit + 1-digit (new)")]):
    cb, cr = passk_curve(base, b), passk_curve(rl, b)
    ax.plot(KS, cb, marker="o", label="before RL")
    ax.plot(KS, cr, marker="o", label="after RL")
    ax.set_xscale("log", base=2)
    ax.set_xlabel("k")
    ax.set_ylabel("pass@k")
    ax.set_title(name)
    ax.grid(alpha=0.3)
    ax.legend()
    print(f"{name:26s} pass@k for k={KS}:\n   before {np.round(cb, 3).tolist()}\n   after  {np.round(cr, 3).tolist()}")
fig.tight_layout()
fig.savefig(f"{HERE}/project5_sharpen_or_expand.png", dpi=110)