"""Project 2: rubric rewards, and how a rubric gets hacked.   (Chapter 21, open project)

Rubric rewards (Rubrics as Rewards, Gunjal et al. 2025; Chapter 14) score several criteria
and combine them with weights. The gold goal here is an exactly correct sum. Three rewards:

  binary    : 1 if exactly right, else 0 (a verifier)
  partial   : partial credit, (right length + right last digit + right first digit
              + 2 x exactly right) / 5
  stylish   : (2 x exactly right + w x "looks thorough") / (2 + w), where "looks thorough"
              means the answer has 3 or more digits (think: a judge that likes long,
              detailed answers). Default w = 2.

Hypotheses: (1) partial credit gives a denser signal and learns at least as fast as binary;
(2) once a criterion that is NOT the goal carries enough weight, the policy optimizes it
at the expense of the goal (Goodhart, Chapter 19): the proxy reward rises while gold accuracy
falls, and the policy starts writing 3-digit answers to problems whose sum is below 100.

Saves project2_rubric.png.
"""
import argparse
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  # noqa: E402
from tinylm import decode_responses, sample  # noqa: E402

ap = argparse.ArgumentParser()
ap.add_argument("--steps", type=int, default=200)
ap.add_argument("--seed", type=int, default=0)
ap.add_argument("--w", type=float, default=2.0, help="weight of the 'looks thorough' criterion")
args = ap.parse_args()


def binary(y, a):
    return float(y == a)


def partial(y, a):
    return (1.0 * (len(y) == len(a)) + 1.0 * (len(y) > 0 and y[-1] == a[-1])
            + 1.0 * (len(y) > 0 and y[0] == a[0]) + 2.0 * (y == a)) / 5.0


def stylish(y, a):
    return (2.0 * (y == a) + args.w * (len(y) >= 3)) / (2.0 + args.w)


@torch.no_grad()
def frac_long_on_small_sums(model, n=200):
    """Fraction of 3+-digit answers to problems whose true sum is below 100 (should be ~0)."""
    rng = random.Random(9)
    pa = [bucket_problem(rng, 1) for _ in range(n)]
    ys = decode_responses(sample(model, [p for p, _ in pa]))
    return sum(len(y) >= 3 for y in ys) / n


ARMS = {"binary": binary, "partial credit": partial, f"stylish (w={args.w:g})": stylish}
hists, proxies, longs = [], [], []
for name, fn in ARMS.items():
    m = get_student(300)
    proxy = []
    hists.append(rl_train(m, args.steps, reward_fn=fn, seed=args.seed, tag=name[:10],
                          on_batch=lambda infos, r, proxy=proxy: proxy.append(float(r.mean()))))
    proxies.append(proxy)
    longs.append(frac_long_on_small_sums(m))

fig, axes = plt.subplots(1, 3, figsize=(14, 4))
for h, name in zip(hists, ARMS):
    axes[0].plot([e["step"] for e in h], [e["pass@1"] for e in h], marker="o", ms=3, label=name)
for p, name in zip(proxies, ARMS):
    axes[1].plot(np.convolve(p, np.ones(10) / 10, mode="valid"), label=name)
axes[2].bar(list(ARMS), longs, color=["C0", "C1", "C2"])
axes[0].set_title("gold: exact-match accuracy (sampled)")
axes[1].set_title("proxy: mean training reward (10-step avg)")
axes[2].set_title("3-digit answers to sums < 100")
for ax in axes[:2]:
    ax.set_xlabel("RL step")
    ax.grid(alpha=0.3)
    ax.legend(fontsize=8)
fig.tight_layout()
fig.savefig(f"{HERE}/project2_rubric.png", dpi=110)
for (name, h, p, lf) in zip(ARMS, hists, proxies, longs):
    print(f"{name:16s} gold pass@1 {h[0]['pass@1']:.2f} -> {h[-1]['pass@1']:.2f}   "
          f"proxy (last 20 steps) {np.mean(p[-20:]):.2f}   3-digit answers to small sums {lf:.2f}")