"""Lab 19 experiments. Each part saves a PNG next to this file.

    python run.py overopt    # Part 1: best-of-n overoptimization curves + Gao fit (a few seconds)
    python run.py hack       # Part 2: TinyLM learns to hack a proxy; the KL leash stops it (~2 min)
    python run.py entropy    # Part 3: entropy under symmetric clip, clip-higher, no upper clip (~2-3 min)
    python run.py            # all three (~5 min)

Uses your overopt.py / failures.py, or the reference solutions with RL_SOLUTIONS=1.
"""

import importlib
import os
import sys
import time

import matplotlib

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

HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, HERE)
SOL = bool(os.environ.get("RL_SOLUTIONS"))
overopt = importlib.import_module("solutions.overopt" if SOL else "overopt")
impl = importlib.import_module("solutions.failures" if SOL else "failures")

from common import RLConfig, gold_reward, make_pool, proxy_reward, rl_train  # noqa: E402


def smooth(x, k=10):
    """Trailing moving average over the last k steps, ignoring NaNs (steps with no data)."""
    x = np.asarray(x, dtype=float)
    ok = ~np.isnan(x)
    s = np.cumsum(np.insert(np.where(ok, x, 0.0), 0, 0.0))
    c = np.cumsum(np.insert(ok.astype(float), 0, 0.0))
    lo = np.maximum(np.arange(1, len(x) + 1) - k, 0)
    hi = np.arange(1, len(x) + 1)
    n = c[hi] - c[lo]
    with np.errstate(invalid="ignore", divide="ignore"):
        return np.where(n > 0, (s[hi] - s[lo]) / n, np.nan)


# ---------------------------------------------------------------- Part 1
def part_overopt():
    print("Part 1: best-of-n against a proxy, measured with the gold reward")
    ns = np.unique(np.round(np.logspace(0, np.log10(3000), 40)).astype(int))
    d = np.sqrt(overopt.bon_kl(ns))
    settings = [
        ("length bias b=1, no noise", 1.0, 0.0),
        ("noise only b=0, σ=1", 0.0, 1.0),
        ("bias and noise b=1, σ=1", 1.0, 1.0),
    ]
    fig, axes = plt.subplots(1, 3, figsize=(13, 3.8), sharey=True)
    for ax, (name, b, sigma) in zip(axes, settings):
        proxy, gold = make_pool(200_000, b=b, sigma=sigma, seed=0)
        stats = np.array([overopt.best_of_n_stats(proxy, gold, n) for n in ns])
        pg, gg = stats[:, 0] - proxy.mean(), stats[:, 1] - gold.mean()
        alpha, beta = overopt.fit_gao_bon(d, gg)
        i = int(np.argmax(gg))
        ax.plot(d, pg, "--", color="C1", label="proxy gain")
        ax.plot(d, gg, "o", ms=3.5, color="C0", label="gold gain (exact, from pool)")
        dd = np.linspace(0, d.max(), 200)
        ax.plot(dd, alpha * dd - beta * dd ** 2, color="C0", alpha=0.7,
                label=f"fit  {alpha:.2f}·d {'−' if beta >= 0 else '+'} {abs(beta):.2f}·d²")
        ax.axhline(0, color="gray", lw=0.8)
        ax.set_title(name)
        ax.set_xlabel("d = √KL(best-of-n ‖ base)")
        ax.set_ylim(-1.5, 4)
        ax.legend(fontsize=8, loc="upper left")
        print(f"  {name:28s} fit alpha={alpha:.3f} beta={beta:.3f}   gold peaks at n={ns[i]} "
              f"(KL {d[i]**2:.2f} nats, gain {gg[i]:+.3f});  at n={ns[-1]}: proxy {pg[-1]:+.2f}, gold {gg[-1]:+.2f}")
    axes[0].set_ylabel("gain over the base policy")
    fig.tight_layout()
    path = os.path.join(HERE, "overopt.png")
    fig.savefig(path, dpi=120)
    print(f"  saved {path}")


# ---------------------------------------------------------------- Part 2
def load_sft():
    from tinylm import load_or_train_sft
    return load_or_train_sft(os.path.join(HERE, "sft_tinylm.pt"))


def part_hack(steps=200):
    print("Part 2: RL on TinyLM against the hackable proxy_reward (read it in common.py)")
    runs = [
        ("proxy reward, no KL", proxy_reward, 0.0),
        ("proxy reward, KL β=0.05", proxy_reward, 0.05),
        ("gold reward (exact match)", gold_reward, 0.0),
    ]
    hists = {}
    for name, fn, beta in runs:
        t0 = time.time()
        print(f"  {name}")
        hists[name] = rl_train(load_sft(), fn, RLConfig(steps=steps, beta=beta), impl, log_every=50)
        print(f"    ({time.time() - t0:.0f}s)")
    fig, axes = plt.subplots(1, 4, figsize=(15, 3.6))
    for (name, _, _), color in zip(runs, ["C3", "C2", "C0"]):
        h = hists[name]
        for ax, key in zip(axes, ["proxy", "gold", "length", "kl"]):
            ax.plot(smooth(h[key]), color=color, label=name)
    for ax, title in zip(axes, ["proxy reward (what RL sees)", "gold: exact-match accuracy",
                                "response length (chars)", "KL(π ‖ π_ref), nats"]):
        ax.set_title(title)
        ax.set_xlabel("RL step")
    axes[0].legend(fontsize=8)
    fig.tight_layout()
    path = os.path.join(HERE, "hack.png")
    fig.savefig(path, dpi=120)
    for name, _, _ in runs:
        h = hists[name]
        print(f"  {name:28s} gold {np.mean(h['gold'][:10]):.2f} -> {np.mean(h['gold'][-10:]):.2f}   "
              f"proxy {np.mean(h['proxy'][:10]):.2f} -> {np.mean(h['proxy'][-10:]):.2f}   "
              f"length {np.mean(h['length'][-10:]):.2f}")
    print(f"  saved {path}")


# ---------------------------------------------------------------- Part 3
def part_entropy(steps=120):
    print("Part 3: policy entropy during GRPO (gold reward, 4 gradient epochs per batch so the clip binds)")
    runs = [("symmetric clip ε = 0.2 / 0.2", 0.2), ("clip-higher ε = 0.2 / 0.28", 0.28),
            ("no upper clip ε = 0.2 / ∞", 1e9)]
    hists = {}
    for name, eh in runs:
        t0 = time.time()
        print(f"  {name}")
        cfg = RLConfig(steps=steps, epochs=4, eps_low=0.2, eps_high=eh)
        hists[name] = rl_train(load_sft(), gold_reward, cfg, impl, log_every=40)
        print(f"    ({time.time() - t0:.0f}s)")
    fig, axes = plt.subplots(1, 3, figsize=(14, 3.6))
    for (name, _), color in zip(runs, ["C0", "C1", "C3"]):
        h = hists[name]
        axes[0].plot(smooth(h["entropy"]), color=color, label=name)
        axes[1].plot(smooth(h["gold"]), color=color, label=name)
        if name.startswith("symmetric"):
            axes[2].plot(smooth(h["clip_rare"]), color=color, label="rare tokens (π_old < 0.2)")
            axes[2].plot(smooth(h["clip_common"]), color=color, ls="--", label="common tokens (π_old > 0.8)")
    axes[0].set_title("mean token entropy (nats)")
    axes[1].set_title("gold: exact-match accuracy")
    axes[2].set_title("symmetric clip: positive-advantage tokens\nswitched off by the upper clip")
    for ax in axes:
        ax.set_xlabel("RL step")
    axes[0].legend(fontsize=8)
    axes[2].legend(fontsize=8)
    fig.tight_layout()
    path = os.path.join(HERE, "entropy.png")
    fig.savefig(path, dpi=120)
    for name, _ in runs:
        h = hists[name]
        print(f"  {name:28s} entropy {np.mean(h['entropy'][:10]):.3f} -> {np.mean(h['entropy'][-10:]):.3f}   "
              f"gold {np.mean(h['gold'][:10]):.2f} -> {np.mean(h['gold'][-10:]):.2f}   "
              f"upper-clipped: rare {np.nanmean(h['clip_rare']):.2f}, common {np.nanmean(h['clip_common']):.3f}")
    print(f"  saved {path}")


if __name__ == "__main__":
    which = sys.argv[1] if len(sys.argv) > 1 else "all"
    if which in ("overopt", "all"):
        part_overopt()
    if which in ("hack", "all"):
        part_hack()
    if which in ("entropy", "all"):
        part_entropy()