"""pass@k before and after GRPO: the pass@k debate in miniature.

    python run.py      # first: trains GRPO and saves tinylm_grpo.pt
    python passk.py    # then: about 15 s; saves passk.png

For each of 300 held-out problems, draws n = 64 samples (temperature 1) from the SFT
model and from the GRPO model, counts the correct ones c, and computes the unbiased
pass@k estimate  1 - C(n - c, k) / C(n, k)  (Chen et al. 2021) for k = 1, 2, 4, ..., 64.
It also counts the problems each model never solves in 64 tries.
"""

import math
import os
import random

import matplotlib

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

import tinylm  # noqa: E402

HERE = os.path.dirname(os.path.abspath(__file__))
N_PROBLEMS, N_SAMPLES = 300, 64
KS = [1, 2, 4, 8, 16, 32, 64]


def pass_at_k(n, c, k):
    """Unbiased estimate of P(at least one of k samples is correct) from c correct out of n."""
    if n - c < k:
        return 1.0
    return 1.0 - math.comb(n - c, k) / math.comb(n, k)


@torch.no_grad()
def correct_counts(model, prompts, answers, n=N_SAMPLES, chunk=50):
    counts = []
    for i in range(0, len(prompts), chunk):
        ps, ans = prompts[i:i + chunk], answers[i:i + chunk]
        rep_p = [p for p in ps for _ in range(n)]
        rep_a = [a for a in ans for _ in range(n)]
        resp = tinylm.decode_responses(tinylm.sample(model, rep_p))
        ok = torch.tensor([r == a for r, a in zip(resp, rep_a)], dtype=torch.float32).view(len(ps), n)
        counts += ok.sum(1).long().tolist()
    return counts


def main():
    sft_path, grpo_path = os.path.join(HERE, "tinylm_sft.pt"), os.path.join(HERE, "tinylm_grpo.pt")
    if not os.path.exists(grpo_path):
        raise SystemExit("tinylm_grpo.pt not found: run `python run.py` first.")
    sft = tinylm.load_or_train_sft(sft_path)
    rl = tinylm.TinyLM()
    rl.load_state_dict(torch.load(grpo_path))
    prompts, answers = tinylm.make_batch(N_PROBLEMS, random.Random(2024))
    torch.manual_seed(0)
    curves, never = {}, {}
    for name, model in [("SFT (before RL)", sft), ("after GRPO", rl)]:
        c = correct_counts(model, prompts, answers)
        curves[name] = [sum(pass_at_k(N_SAMPLES, ci, k) for ci in c) / len(c) for k in KS]
        never[name] = sum(ci == 0 for ci in c)
    fig, ax = plt.subplots(figsize=(6, 4))
    for name, ys in curves.items():
        ax.plot(KS, ys, "o-", label=f"{name}: never solved {never[name]}/{N_PROBLEMS}")
    ax.set(xscale="log", xlabel="k (samples per problem)", ylabel="pass@k", title="pass@k on held-out addition problems")
    ax.set_xticks(KS, [str(k) for k in KS])
    ax.legend(fontsize=8)
    fig.tight_layout()
    path = os.path.join(HERE, "passk.png")
    fig.savefig(path, dpi=120)
    fmt = lambda ys: " ".join(f"{y:.3f}" for y in ys)  # noqa: E731
    print(f"k = {KS}")
    for name in curves:
        print(f"{name:16s} pass@k: {fmt(curves[name])}   never solved: {never[name]}")
    print(f"saved {path}")


if __name__ == "__main__":
    main()