"""Lab 12 experiment: the whole RLHF pipeline on TinyLM, then the KL leash on and off.

    python run.py                 # about 2 minutes on a laptop CPU; saves rlhf_runs.png
    RL_SOLUTIONS=1 python run.py  # the same, with the reference solution

Steps:
  1. SFT: load (or train and cache) a TinyLM that does addition about half the time.
  2. Preferences: for 600 prompts, sample K = 4 answers; the hidden labeler ranks them.
  3. Reward model: fit a DigitRM with your bt_loss on the K-choose-2 pairs of each ranking.
  4. PPO: four RLHF runs against that RM, with different KL coefficients beta (and one with PPO-ptx).
"""

import importlib
import os
import time

import matplotlib

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

rlhf = importlib.import_module("solutions.rlhf" if os.environ.get("RL_SOLUTIONS") else "rlhf")
from prefs import collect_rankings, digit_counts  # noqa: E402
from tinylm import decode_responses, load_or_train_sft, sample  # noqa: E402

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

RUNS = [
    ("beta = 0 (no leash)", dict(beta=0.0)),
    ("beta = 0.1", dict(beta=0.1)),
    ("beta = 1", dict(beta=1.0)),
    ("beta = 0.1 + PPO-ptx", dict(beta=0.1, ptx_coef=1.0)),
]


def smooth(xs, k=8):
    out, acc = [], 0.0
    for i, x in enumerate(xs):
        acc = x if i == 0 else acc + (x - acc) / min(i + 1, k)
        out.append(acc)
    return out


def main():
    t0 = time.time()
    sft = load_or_train_sft(os.path.join(HERE, "sft_tinylm.pt"))

    # ---- preferences and reward model
    train = collect_rankings(sft, 600, K=4, seed=1)
    test = collect_rankings(sft, 200, K=4, seed=2)
    n_pairs = sum(len(rlhf.ranking_to_pairs(d["ranking"])) for d in train)
    rm = rlhf.train_reward_model(train)
    acc, ceiling = rlhf.rm_accuracy(rm, test)
    with torch.no_grad():
        offset = torch.cat([rm(digit_counts(d["out"])) for d in train]).mean().item()
    print(f"reward model: {n_pairs} training pairs, held-out accuracy {acc:.3f} (ceiling {ceiling:.3f})")
    print("  digit weights w[0..9]:", " ".join(f"{d}:{w:+.2f}" for d, w in enumerate(rm.w.tolist())))
    print(f"  offset (mean RM score of SFT answers, subtracted during RL): {offset:+.2f}")

    # ---- RLHF runs
    probe = ["37+48=", "12+35=", "9+8=", "50+46="]
    results = []
    for name, kw in RUNS:
        t1 = time.time()
        policy, hist = rlhf.train_rlhf(sft, rm, offset, rlhf.RLHFConfig(**kw))
        torch.manual_seed(0)
        answers = decode_responses(sample(policy, probe, greedy=True))
        last = {k: sum(v[-20:]) / 20 for k, v in hist.items()}
        results.append((name, hist))
        print(f"{name:22s} RM {last['rm']:+.2f}  gold {last['gold']:+.2f}  KL {last['kl']:5.2f}  "
              f"accuracy {last['acc']:.2f}  malformed {last['malformed']:.2f}  "
              f"greedy answers {dict(zip(probe, answers))}  ({time.time() - t1:.0f}s)")

    # ---- plot
    fig, axes = plt.subplots(2, 2, figsize=(11, 7.5))
    panels = [
        ("rm", "Reward-model score (what PPO optimizes)"),
        ("gold", "Hidden labeler's score (what we wanted)"),
        ("kl", "KL to reference, nats per response"),
        ("acc", "Arithmetic accuracy (never rewarded)"),
    ]
    for ax, (key, title) in zip(axes.flat, panels):
        for name, hist in results:
            ax.plot(smooth(hist[key]), label=name)
        ax.set_title(title)
        ax.set_xlabel("PPO iteration")
        ax.grid(alpha=0.3)
    axes[1, 0].set_yscale("symlog", linthresh=1.0)
    axes[0, 0].legend(fontsize=8)
    fig.suptitle("RLHF on TinyLM: the reward model is hackable; the KL leash and PPO-ptx limit the damage")
    fig.tight_layout()
    path = os.path.join(HERE, "rlhf_runs.png")
    fig.savefig(path, dpi=120)
    print(f"saved {path}  (total {time.time() - t0:.0f}s)")


if __name__ == "__main__":
    main()