"""Project 1: does clip-higher slow entropy collapse here?   (Chapter 21, open project)

Hypothesis (DAPO, Yu et al. 2025): with a symmetric PPO clip [1-0.2, 1+0.2], tokens the
policy finds unlikely can only be raised a little per batch, while likely tokens get
sharpened freely, so entropy falls fast. Raising only the upper bound (eps_high = 0.28)
should keep entropy higher for longer and help pass@k at large k.

What this script does: three runs from the same weak student, identical except for the
clip range, with 4 gradient epochs per batch (so the ratio leaves 1 and the clip acts).
It plots entropy, pass@1 and pass@16 and saves project1_clip_higher.png.

Your job (see README): run >= 3 seeds (--seed), decide whether the entropy gap is larger
than the seed-to-seed spread, and then try the extension in the README.
"""
import argparse

import matplotlib

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

from common import HERE, get_student, plot_histories, rl_train  # noqa: E402

ap = argparse.ArgumentParser()
ap.add_argument("--steps", type=int, default=150)
ap.add_argument("--seed", type=int, default=0)
args = ap.parse_args()

ARMS = {"symmetric 0.2/0.2": (0.2, 0.2), "clip-higher 0.2/0.28": (0.2, 0.28), "no clip": (10.0, 10.0)}
hists = []
for name, (lo, hi) in ARMS.items():
    m = get_student(300)
    hists.append(rl_train(m, args.steps, eps_low=lo, eps_high=hi, epochs=4, lr=3e-4, seed=args.seed,
                          eval_every=25, tag=name[:10]))
fig, axes = plt.subplots(1, 3, figsize=(14, 4))
for ax, key in zip(axes, ["entropy", "pass@1", "pass@16"]):
    plot_histories(ax, hists, key, list(ARMS))
fig.tight_layout()
fig.savefig(f"{HERE}/project1_clip_higher.png", dpi=110)
print("final entropy: " + ", ".join(f"{n} {h[-1]['entropy']:.3f}" for n, h in zip(ARMS, hists)))
print("final pass@1:  " + ", ".join(f"{n} {h[-1]['pass@1']:.3f}" for n, h in zip(ARMS, hists)))