"""Project 4: MaxRL vs. the mean baseline vs. GRPO's std normalization.

All three use the same samples and differ only in how a group's rewards become advantages:
  mean   A_i = r_i - p_hat                    (RLOO / Dr. GRPO style; population weight w(p) = 1)
  std    A_i = (r_i - p_hat) / std            (GRPO;                  w(p) ~ 1 / sqrt(p(1-p)))
  maxrl  A_i = (r_i - p_hat) / p_hat          (MaxRL, your Task 5;    w(p) -> 1 / p, maximum likelihood)

Hypothesis (Tajwar et al. 2026): up-weighting hard prompts (small p) improves pass@k at
large k more than pass@1. Saves project4_maxrl.png.
"""
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=200)
ap.add_argument("--seed", type=int, default=0)
ap.add_argument("--sft-steps", type=int, default=300)
args = ap.parse_args()

ARMS = ["mean", "std", "maxrl"]
hists = []
for adv in ARMS:
    m = get_student(args.sft_steps)
    hists.append(rl_train(m, args.steps, adv=adv, seed=args.seed, tag=adv))
fig, axes = plt.subplots(1, 3, figsize=(14, 4))
for ax, key in zip(axes, ["pass@1", "pass@4", "pass@16"]):
    plot_histories(ax, hists, key, ARMS)
fig.tight_layout()
fig.savefig(f"{HERE}/project4_maxrl.png", dpi=110)
for key in ["pass@1", "pass@16"]:
    print(f"final {key}: " + ", ".join(f"{a} {h[-1][key]:.3f}" for a, h in zip(ARMS, hists)))