"""Train TinyLM with GRPO from the SFT checkpoint and plot what happens.

    python run.py                          # the default recipe, about 40-60 s on a laptop
    python run.py --G 2 --prompts 64       # same samples per step, tiny groups
    python run.py --beta 0                 # no KL leash
    python run.py --no-std                 # mean-centred advantages only
    python run.py --inner-steps 8 --clip 1000   # reuse each batch 8 times, effectively no clip

Saves grpo_tinylm.png (or --out) and the trained weights tinylm_grpo.pt next to this
file, and prints a one-line summary. Then try:  python passk.py
"""

import argparse
import importlib
import os

import matplotlib

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

grpo = importlib.import_module("solutions.grpo" if os.environ.get("RL_SOLUTIONS") else "grpo")

HERE = os.path.dirname(os.path.abspath(__file__))


def smooth(x, k=10):
    x = np.asarray(x, dtype=float)
    if len(x) < k:
        return x
    return np.convolve(x, np.ones(k) / k, mode="valid")


def main():
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--iters", type=int, default=300)
    ap.add_argument("--G", type=int, default=8, help="group size")
    ap.add_argument("--prompts", type=int, default=16, help="prompts per iteration")
    ap.add_argument("--inner-steps", type=int, default=2, help="gradient steps per batch")
    ap.add_argument("--lr", type=float, default=3e-4)
    ap.add_argument("--clip", type=float, default=0.2, help="clip range epsilon")
    ap.add_argument("--beta", type=float, default=0.04, help="KL coefficient")
    ap.add_argument("--temperature", type=float, default=1.0)
    ap.add_argument("--no-std", action="store_true", help="do not divide advantages by the group std")
    ap.add_argument("--seed", type=int, default=0)
    ap.add_argument("--out", default="grpo_tinylm.png")
    a = ap.parse_args()

    cfg = grpo.GRPOConfig(
        iters=a.iters, G=a.G, prompts_per_iter=a.prompts, inner_steps=a.inner_steps, lr=a.lr,
        clip_eps=a.clip, beta=a.beta, temperature=a.temperature, scale_std=not a.no_std, seed=a.seed,
    )
    print(f"GRPO on TinyLM: {cfg}")
    policy, h = grpo.train_grpo(cfg, sft_path=os.path.join(HERE, "tinylm_sft.pt"))
    torch.save(policy.state_dict(), os.path.join(HERE, "tinylm_grpo.pt"))  # used by passk.py

    it = np.arange(1, len(h["reward"]) + 1)
    k = 10
    its = it[k - 1:] if len(it) >= k else it
    fig, ax = plt.subplots(2, 2, figsize=(10, 7))

    ax[0, 0].plot(h["eval_iter"], h["acc_sampled"], "o-", label="sampled (T = 1), what GRPO trains")
    ax[0, 0].plot(h["eval_iter"], h["acc_greedy"], "s--", label="greedy")
    ax[0, 0].set(title="Accuracy on 500 held-out problems", xlabel="GRPO iteration", ylim=(0, 1))
    ax[0, 0].legend(fontsize=8)

    ax[0, 1].plot(its, smooth(h["reward"], k), label="mean reward of the batch")
    ax[0, 1].plot(its, smooth(h["frac_zero_signal"], k), label="fraction of zero-signal groups")
    ax[0, 1].set(title=f"Reward and wasted groups (G = {cfg.G}, {k}-iter average)", xlabel="GRPO iteration", ylim=(0, 1))
    ax[0, 1].legend(fontsize=8)

    ax[1, 0].plot(it, h["kl"], alpha=0.35, color="C2")
    ax[1, 0].plot(its, smooth(h["kl"], k), color="C2")
    ax[1, 0].set(title=f"KL to reference (mean k3 per token), beta = {cfg.beta}", xlabel="GRPO iteration")

    ax[1, 1].plot(it, h["entropy"], alpha=0.35, color="C3")
    ax[1, 1].plot(its, smooth(h["entropy"], k), color="C3")
    ax[1, 1].set(title="Policy entropy per response token (nats)", xlabel="GRPO iteration")

    fig.tight_layout()
    path = os.path.join(HERE, a.out)
    fig.savefig(path, dpi=120)
    print(f"sampled accuracy {h['acc_sampled'][0]:.3f} -> {h['acc_sampled'][-1]:.3f}, "
          f"greedy {h['acc_greedy'][0]:.3f} -> {h['acc_greedy'][-1]:.3f}, "
          f"final KL {np.mean(h['kl'][-10:]):.3f}, entropy {h['entropy'][0]:.3f} -> {np.mean(h['entropy'][-10:]):.3f}; "
          f"saved {path}")


if __name__ == "__main__":
    main()