"""Chapter 20 lab: run all three experiments and save rl_systems.png (a few seconds).

    python run.py                    # uses your memory.py, scheduler.py, tis.py
    RL_SOLUTIONS=1 python run.py     # uses the reference solutions
"""

import importlib
import os

import matplotlib

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

SOL = os.environ.get("RL_SOLUTIONS")
memory = importlib.import_module("solutions.memory" if SOL else "memory")
scheduler = importlib.import_module("solutions.scheduler" if SOL else "scheduler")
tis = importlib.import_module("solutions.tis" if SOL else "tis")

from scheduler_core import lognormal_lengths  # noqa: E402
from toy import estimator_moments, make_toy  # noqa: E402

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


def part1_memory():
    print("Part 1: model-state memory for a 7B policy (GB; activations and KV cache NOT included)")
    rows = [
        ("PPO (policy, critic, ref, RM all 7B)", memory.rl_memory_gb("ppo", 7e9)),
        ("GRPO + reference + verifier", memory.rl_memory_gb("grpo", 7e9)),
        ("GRPO, beta = 0 (no reference)", memory.rl_memory_gb("grpo", 7e9, use_ref=False)),
        ("GRPO + LoRA (40M adapter params)", memory.rl_memory_gb("grpo", 7e9, lora_params=40e6)),
    ]
    print(f"  {'setup':40s} {'policy':>8s} {'ref':>6s} {'critic':>7s} {'RM':>6s} {'total':>7s} {'H100s':>6s}")
    for name, d in rows:
        n_gpu = int(np.ceil(d["total"] / 80))
        print(f"  {name:40s} {d['policy']:8.1f} {d['reference']:6.1f} {d['critic']:7.1f} "
              f"{d['reward_model']:6.1f} {d['total']:7.1f} {n_gpu:6d}")
    kv = memory.kv_cache_bytes(28, 4, 128, 1)
    print(f"  KV cache, Qwen2.5-7B shape: {kv:,} bytes/token; 64 sequences x 16k tokens = "
          f"{memory.kv_cache_bytes(28, 4, 128, 64 * 16384) / 1e9:.1f} GB")


def part2_scheduling(ax_speed, ax_lag):
    print("Part 2: sync vs async rollouts (32 lanes, batch 32, train time 1, 60 steps)")
    ks = list(range(0, 7))
    for sigma, color in [(0.25, "C0"), (0.75, "C1"), (1.25, "C3")]:
        L = lognormal_lengths(40000, sigma, seed=1)
        base = scheduler.simulate(L, 32, 32, 0, 1.0, 60)
        speed, lag = [], []
        for k in ks:
            r = scheduler.simulate(L, 32, 32, k, 1.0, 60)
            speed.append(r["throughput"] / base["throughput"])
            lag.append(r["mean_lag"])
        ax_speed.plot(ks, speed, "o-", color=color, label=f"tail σ = {sigma}")
        ax_lag.plot(ks, lag, "o-", color=color, label=f"tail σ = {sigma}")
        print(f"  sigma={sigma}: sync lane utilization {base['lane_util']:.0%}; "
              f"k=1 speedup {speed[1]:.2f}x (mean lag {lag[1]:.2f}); k=4 speedup {speed[4]:.2f}x (mean lag {lag[4]:.2f})")
    ax_speed.set(xlabel="max staleness k (k = 0 is synchronous)", ylabel="throughput / synchronous",
                 title="Async buys throughput...")
    ax_lag.set(xlabel="max staleness k", ylabel="mean policy lag of trained samples",
               title="...and pays in staleness")
    for ax in (ax_speed, ax_lag):
        ax.grid(alpha=0.3)
        ax.legend(fontsize=8)


def part3_tis(ax_bias, ax_rmse):
    print("Part 3: gradient error under train-inference mismatch (toy LM, T=6, V=4)")
    mismatches = np.linspace(0, 1.0, 11)
    styles = {"none": ("k", "naive (no correction)"), "token": ("C0", "per-token IS"),
              "token_tis": ("C2", "per-token TIS, C=2"), "seq": ("C3", "sequence IS"),
              "seq_tis": ("C1", "sequence TIS, C=2")}
    table = {m: ([], []) for m in tis.MODES}
    for s in mismatches:
        toy = make_toy(mismatch=s)
        for m in tis.MODES:
            w = tis.importance_weights(toy["logp_train"], toy["logp_sampler"], m, 2.0)
            mo = estimator_moments(toy, tis.pg_estimate(toy["R"], toy["score"], w))
            table[m][0].append(mo["rel_bias"])
            table[m][1].append(mo["rel_rmse"](256))
    for m, (col, lab) in styles.items():
        ax_bias.plot(mismatches, table[m][0], "o-", ms=3, color=col, label=lab)
        ax_rmse.plot(mismatches, table[m][1], "o-", ms=3, color=col, label=lab)
    ax_bias.set(xlabel="sampler/trainer mismatch scale", ylabel="|bias| / |true gradient|",
                title="Bias: does NOT shrink with batch size")
    ax_rmse.set(xlabel="sampler/trainer mismatch scale", ylabel="relative RMSE, batch of 256",
                title="Total error = bias² + variance/N", ylim=(0, 2.5))
    for ax in (ax_bias, ax_rmse):
        ax.grid(alpha=0.3)
        ax.legend(fontsize=7)
    i = 3  # mismatch 0.3
    print("  at mismatch 0.3, relative bias / RMSE(N=256): " + ", ".join(
        f"{m} {table[m][0][i]:.2f}/{table[m][1][i]:.2f}" for m in tis.MODES))


def main():
    part1_memory()
    fig, axes = plt.subplots(2, 2, figsize=(11, 8.5))
    part2_scheduling(axes[0, 0], axes[0, 1])
    part3_tis(axes[1, 0], axes[1, 1])
    fig.tight_layout()
    out = os.path.join(HERE, "rl_systems.png")
    fig.savefig(out, dpi=120)
    print(f"Saved {out}")


if __name__ == "__main__":
    main()