"""Project 0 (fully worked): on-policy distillation vs. off-policy distillation vs. RL.

A strong teacher TinyLM (~100% on 2-digit addition) and a weak student (~30% sampled
accuracy). Four ways to improve the student, each for the same number of steps and
the same number of sequences per step (64):

  opd-sampled   student samples; per-token reward log pi_T - log pi_theta (your Tasks 1-2)
  opd-full      student samples; exact full-vocabulary reverse KL at each visited state (Task 3)
  sft-teacher   *teacher* samples; plain next-token cross-entropy on them (off-policy distillation)
  rl            student samples; 0/1 exact-match reward; G = 4 per prompt, mean baseline

Prints a summary line and saves opd_comparison.png. About 1.5-2 minutes on a laptop CPU
(+20 s the first time, to train and cache the teacher).

Usage:  RL_SOLUTIONS=1 python run.py      (or plain `python run.py` once Tasks 1-5 are done)
"""

import copy
import importlib
import os
import random
import time

import matplotlib

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

from common import HERE, evaluate, get_student, get_teacher, response_logits  # noqa: E402
from tinylm import decode_responses, make_batch, reward, sample, sequence_logprobs  # noqa: E402

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

STEPS, EVAL_EVERY, SEQS, LR = 200, 25, 64, 3e-4


@torch.no_grad()
def reverse_kl_to_teacher(student, teacher, n=256, seed=7):
    """Per-token KL(student || teacher) on the student's own samples (lower = closer)."""
    prompts, _ = make_batch(n, random.Random(seed))
    out = sample(student, prompts)
    return opd.full_vocab_reverse_kl(response_logits(student, out), response_logits(teacher, out), out["mask"]).item()


def train(arm, student0, teacher, seed=0):
    m = copy.deepcopy(student0)
    opt = torch.optim.AdamW(m.parameters(), lr=LR, weight_decay=0.0)
    rng = random.Random(seed)
    torch.manual_seed(seed)
    hist = []
    for step in range(STEPS + 1):
        if step % EVAL_EVERY == 0:
            ev = evaluate(m, n_prompts=100, n_samples=8, ks=(1, 8))
            ev["kl"] = reverse_kl_to_teacher(m, teacher)
            ev["step"] = step
            hist.append(ev)
        if step == STEPS:
            break
        if arm == "rl":
            G = 4
            p, a = make_batch(SEQS // G, rng)
            P, A = [x for x in p for _ in range(G)], [x for x in a for _ in range(G)]
            out = sample(m, P)
            r = torch.tensor([reward(y, t) for y, t in zip(decode_responses(out), A)]).view(-1, G)
            adv = (r - r.mean(1, keepdim=True)).reshape(-1, 1)
            logp, mask = sequence_logprobs(m, out), out["mask"]
            loss = -(adv * logp * mask).sum() / mask.sum()
        elif arm == "sft-teacher":
            p, _ = make_batch(SEQS, rng)
            out = sample(teacher, p)                       # the TEACHER's states and tokens
            logp, mask = sequence_logprobs(m, out), out["mask"]
            loss = -(logp * mask).sum() / mask.sum()
        elif arm == "opd-sampled":
            p, _ = make_batch(SEQS, rng)
            out = sample(m, p)                             # the STUDENT's states and tokens
            logp, mask = sequence_logprobs(m, out), out["mask"]
            with torch.no_grad():
                t_logp = sequence_logprobs(teacher, out)
            r = opd.reverse_kl_rewards(logp, t_logp, mask)
            loss = opd.opd_pg_loss(logp, r, mask)
        elif arm == "opd-full":
            p, _ = make_batch(SEQS, rng)
            out = sample(m, p)
            with torch.no_grad():
                t_logits = response_logits(teacher, out)
            loss = opd.full_vocab_reverse_kl(response_logits(m, out), t_logits, out["mask"])
        else:
            raise ValueError(arm)
        opt.zero_grad()
        loss.backward()
        opt.step()
    return hist


def main():
    t0 = time.time()
    teacher, student = get_teacher(), get_student(300)
    teacher.eval()
    arms = ["opd-full", "opd-sampled", "sft-teacher", "rl"]
    hists = {}
    for arm in arms:
        hists[arm] = train(arm, student, teacher)
        h = hists[arm][-1]
        print(f"{arm:12s} pass@1 {hists[arm][0]['pass@1']:.2f} -> {h['pass@1']:.2f}   "
              f"KL to teacher {hists[arm][0]['kl']:.3f} -> {h['kl']:.3f}   ({time.time() - t0:.0f}s)", flush=True)

    fig, axes = plt.subplots(1, 3, figsize=(14, 4))
    for arm in arms:
        s = [e["step"] for e in hists[arm]]
        axes[0].plot(s, [e["pass@1"] for e in hists[arm]], marker="o", ms=3, label=arm)
        axes[1].plot(s, [e["pass@8"] for e in hists[arm]], marker="o", ms=3, label=arm)
        axes[2].plot(s, [e["kl"] for e in hists[arm]], marker="o", ms=3, label=arm)
    for ax, t in zip(axes, ["sampled accuracy (pass@1)", "pass@8", "per-token KL(student || teacher)"]):
        ax.set_title(t)
        ax.set_xlabel("step (64 sequences each)")
        ax.grid(alpha=0.3)
    axes[2].set_yscale("log")
    axes[0].legend(fontsize=8)
    fig.tight_layout()
    path = os.path.join(HERE, "opd_comparison.png")
    fig.savefig(path, dpi=110)
    best = max(arms, key=lambda a: hists[a][-1]["pass@1"])
    print(f"saved {path}; best final pass@1: {best} ({hists[best][-1]['pass@1']:.2f}); "
          f"total {time.time() - t0:.0f}s")


if __name__ == "__main__":
    main()