"""Train PPO on CartPole four ways and plot what the clip buys you.

    python run.py            # ~2 minutes on a laptop CPU; saves ppo_cartpole.png

Variants (same seed, same network, same 40k environment steps):
  1. PPO                    eps = 0.2, K = 10 epochs, lr 1e-3   (the default PPOConfig)
  2. PPO, aggressive        eps = 0.2, K = 40 epochs, lr 3e-3   (squeeze each batch hard)
  3. no clip, aggressive    eps = infinity, K = 40, lr 3e-3     (the unclipped surrogate)
  4. one epoch              eps = 0.2, K = 1                     (use each batch once)
"""

import importlib
import os
import time

import matplotlib

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

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

TOTAL_STEPS = 40_960
SEED = 0
VARIANTS = {
    "PPO (K=10)": PPOConfig(),
    "PPO, aggressive (K=40, lr 3e-3)": PPOConfig(epochs=40, lr=3e-3),
    "no clip, aggressive (K=40, lr 3e-3)": PPOConfig(epochs=40, lr=3e-3, clip_eps=1e9),
    "one epoch (K=1)": PPOConfig(epochs=1),
}


def main():
    results = {}
    t0 = time.time()
    for name, cfg in VARIANTS.items():
        h = ppo.train_ppo(cfg, total_steps=TOTAL_STEPS, seed=SEED)
        results[name] = h
        print(f"{name:38s} final mean return {h['mean_return'][-1]:6.1f}   "
              f"mean approx KL {np.mean(h['approx_kl']):.4f}   ({time.time() - t0:.0f}s)")

    fig, axes = plt.subplots(2, 2, figsize=(11, 7.5))
    panels = [
        ("mean_return", "mean return (last 20 episodes)", False),
        ("approx_kl", "approx KL(pi_old || pi_new) per update", True),
        ("clip_frac", "clip fraction", False),
        ("entropy", "policy entropy (nats)", False),
    ]
    for ax, (key, title, logy) in zip(axes.flat, panels):
        for name, h in results.items():
            y = np.maximum(h[key], 1e-6) if logy else h[key]
            ax.plot(h["steps"], y, label=name, lw=1.6)
        ax.set_title(title)
        ax.set_xlabel("environment steps")
        if logy:
            ax.set_yscale("log")
        ax.grid(alpha=0.3)
    axes[0, 0].axhline(500, color="gray", ls=":", lw=1)
    axes[0, 0].legend(fontsize=8)
    fig.suptitle("PPO on CartPole: what the clip buys you")
    fig.tight_layout()
    out = os.path.join(os.path.dirname(os.path.abspath(__file__)), "ppo_cartpole.png")
    fig.savefig(out, dpi=120)
    best = max(results, key=lambda n: results[n]["mean_return"][-1])
    print(f"saved {out}; best final return: {best} ({results[best]['mean_return'][-1]:.0f})")


if __name__ == "__main__":
    main()