"""Compare zoo presets with YOUR pg_loss (or the reference with RL_SOLUTIONS=1).

Two experiments, one figure (zoo_comparison.png, saved next to this file):

  1. TinyLM addition ("37+48="), starting from the SFT checkpoint (~50% sampled
     accuracy). Several presets, same seed, same budget. Top row: sampled
     accuracy and mean response length.  (~2 minutes on a laptop CPU)
  2. LengthToy (common.py), where filler length is a free choice that never
     changes the reward. Bottom row: mean length of WRONG and of RIGHT responses
     during training. Seq-mean aggregation (GRPO, GSPO) should make wrong
     answers long and right answers short; the others should not. (seconds)
     (GSPO is left out of the default: on-policy, with one update per batch, its
     sequence ratio is exactly 1 and it is the same estimator as GRPO. Try it.)

Usage:
    python run.py                                   # default presets
    python run.py --presets grpo dapo cispo gspo    # any keys of common.PRESETS
    python run.py --steps 60                        # quicker, noisier
"""

import argparse
import importlib
import math
import os
import time

import matplotlib

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

from common import PRESETS, train, train_length_toy  # noqa: E402

zoo = importlib.import_module("solutions.zoo" if os.environ.get("RL_SOLUTIONS") else "zoo")
HERE = os.path.dirname(os.path.abspath(__file__))


def smooth(xs, k=15):
    out, buf = [], []
    for x in xs:
        if x == x:  # skip NaN (no wrong / right response in that batch)
            buf.append(x)
        buf = buf[-k:]
        out.append(sum(buf) / len(buf) if buf else math.nan)
    return out


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--presets", nargs="+", default=["reinforce", "rloo", "grpo", "dr_grpo"])
    ap.add_argument("--toy-presets", nargs="+", default=["grpo", "dapo", "dr_grpo"])
    ap.add_argument("--steps", type=int, default=120)
    ap.add_argument("--lr", type=float, default=5e-4)
    ap.add_argument("--seed", type=int, default=0)
    args = ap.parse_args()

    t0 = time.time()
    sft_path = os.path.join(HERE, "sft_tinylm.pt")
    tiny = {}
    for name in args.presets:
        print(f"[TinyLM] {name}")
        tiny[name] = train(zoo.pg_loss, PRESETS[name], steps=args.steps, lr=args.lr, seed=args.seed,
                           sft_path=sft_path, eval_every=max(1, args.steps // 6), eval_n=500)
    toy = {}
    for name in args.toy_presets:
        toy[name] = train_length_toy(zoo.pg_loss, PRESETS[name], steps=300, seed=args.seed)

    fig, ax = plt.subplots(2, 2, figsize=(11, 8))
    for name, h in tiny.items():
        ax[0, 0].plot(h["eval_step"], h["eval_acc"], marker="o", label=name)
        ax[0, 1].plot(h["step"], smooth(h["len"]), label=name)
    ax[0, 0].set(title="TinyLM: sampled accuracy (500 held-out prompts)", xlabel="step", ylabel="accuracy")
    ax[0, 1].set(title="TinyLM: mean response length (tokens, incl. <eos>)", xlabel="step", ylabel="tokens",
                 ylim=(1, 6.5))
    for name, h in toy.items():
        line, = ax[1, 0].plot(smooth(h["len_wrong"]), label=name)
        ax[1, 1].plot(smooth(h["len_right"]), label=name, color=line.get_color())
    ax[1, 0].set(title="LengthToy: mean length of WRONG responses", xlabel="step", ylabel="tokens")
    ax[1, 1].set(title="LengthToy: mean length of RIGHT responses", xlabel="step", ylabel="tokens")
    for a in ax.flat:
        a.grid(alpha=0.3)
        a.legend(fontsize=8)
    fig.tight_layout()
    out = os.path.join(HERE, "zoo_comparison.png")
    fig.savefig(out, dpi=110)

    acc = ", ".join(f"{n} {h['eval_acc'][0]:.2f}->{h['eval_acc'][-1]:.2f}" for n, h in tiny.items())
    drift = ", ".join(
        f"{n} {smooth(h['len_wrong'])[19]:.1f}->{smooth(h['len_wrong'])[-1]:.1f}" for n, h in toy.items())
    print(f"TinyLM accuracy: {acc} | LengthToy wrong-answer length: {drift} | saved {out} ({time.time() - t0:.0f}s)")


if __name__ == "__main__":
    main()