"""Project 3: a self-play-style proposer that picks problems at the learner's frontier.

Absolute Zero (Zhao et al. 2025) trains a *proposer* with the learnability reward
    r_propose = 0 if the solver's success rate p_hat = 0,  else 1 - p_hat,
so problems that are always solved (p_hat = 1) or never solved (p_hat = 0) are worth nothing.
Here the proposer is a small bandit (Chapter 2) over five buckets of problems
(see common.BUCKET_NAMES). The student's success rates differ a lot across them: about
2% on 1d+1d (rare in its SFT data), 30-50% on the 2-digit buckets, ~0% on 3d+1d (never seen).

Three arms, same number of RL steps:
  target-only   : train only on the hardest bucket (3d+1d)
  uniform       : pick a bucket uniformly at random
  learnability  : softmax over each bucket's running learnability reward (+10% uniform)

Hypothesis: training only on a target the student never solves gives (almost) no gradient;
the learnability proposer stops spending samples on unsolvable and solved buckets and so
learns at least as fast as uniform sampling on the learnable ones. Saves project3_learnability.png and prints how often each bucket was chosen.
"""
import argparse
import math

import matplotlib

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

from common import BUCKET_NAMES, HERE, bucket_problem, evaluate, get_student, learnability_reward, rl_train  # noqa: E402

ap = argparse.ArgumentParser()
ap.add_argument("--steps", type=int, default=200)
ap.add_argument("--seed", type=int, default=0)
args = ap.parse_args()
NB = len(BUCKET_NAMES)


class Proposer:
    def __init__(self, mode, tau=0.1, eps=0.1, decay=0.8):
        self.mode, self.tau, self.eps, self.decay = mode, tau, eps, decay
        self.score = np.full(NB, 0.5)          # optimistic start: every bucket looks learnable
        self.counts = np.zeros(NB, dtype=int)

    def probs(self):
        if self.mode == "uniform":
            return np.full(NB, 1 / NB)
        if self.mode == "target-only":
            return np.eye(NB)[NB - 1]
        z = np.exp((self.score - self.score.max()) / self.tau)
        return (1 - self.eps) * z / z.sum() + self.eps / NB

    def problem(self, rng):
        b = int(rng.choices(range(NB), weights=self.probs())[0])
        self.counts[b] += 1
        p, a = bucket_problem(rng, b)
        return p, a, b

    def update(self, buckets, rewards):       # rewards: (P, G) array
        for b, r in zip(buckets, rewards):
            lr = learnability_reward(float(r.mean()))
            self.score[b] = self.decay * self.score[b] + (1 - self.decay) * lr


def eval_buckets(model, buckets=(0, 2, 3, 4)):
    return {b: evaluate(model, n_prompts=60, n_samples=8, ks=(1,), problem_fn=lambda rng, b=b: bucket_problem(rng, b))["pass@1"]
            for b in buckets}


results = {}
for mode in ["target-only", "uniform", "learnability"]:
    m = get_student(300)
    prop = Proposer(mode)
    curve = []
    for chunk in range(4):
        rl_train(m, args.steps // 4, problem_fn=prop.problem, on_batch=prop.update, seed=args.seed + chunk,
                 eval_every=10**9, log=None)
        curve.append(eval_buckets(m))
        print(f"{mode:12s} after {(chunk + 1) * (args.steps // 4):4d} steps: " +
              "  ".join(f"{BUCKET_NAMES[b]} {v:.2f}" for b, v in curve[-1].items()), flush=True)
    results[mode] = (curve, prop.counts / prop.counts.sum())
    print(f"{mode:12s} bucket usage: " + "  ".join(f"{BUCKET_NAMES[b]} {f:.2f}" for b, f in enumerate(results[mode][1])))

fig, axes = plt.subplots(1, 4, figsize=(16, 4))
for ax, b in zip(axes, (0, 2, 3, 4)):
    for mode, (curve, _) in results.items():
        ax.plot([(i + 1) * (args.steps // 4) for i in range(len(curve))], [c[b] for c in curve], marker="o", label=mode)
    ax.set_title(f"pass@1 on {BUCKET_NAMES[b]}")
    ax.set_xlabel("RL step")
    ax.grid(alpha=0.3)
axes[0].legend(fontsize=8)
fig.tight_layout()
fig.savefig(f"{HERE}/project3_learnability.png", dpi=110)