"""Chapter 7 lab: make the two plots that sum up the chapter.

Left:  spread of single-sample REINFORCE estimates on the 3-armed bandit as you
       slide a constant baseline b (rewards shifted by +10, as in the chapter).
Right: REINFORCE on CartPole with three per-step weights (whole return,
       reward-to-go, reward-to-go minus a learned baseline), 4 seeds each.

    python run.py                    # uses your pg.py
    RL_SOLUTIONS=1 python run.py     # uses the reference solution

Saves reinforce.png next to this file. Takes about a minute on a laptop CPU.
"""

import importlib
import os
import time

import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np

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

HERE = os.path.dirname(os.path.abspath(__file__))
SEEDS = [0, 1, 2, 3]
N_ITERS = 50


def smooth(x, k=20):
    x = np.asarray(x, dtype=float)
    c = np.cumsum(np.insert(x, 0, 0.0))
    out = np.empty_like(x)
    for i in range(len(x)):
        lo = max(0, i - k + 1)
        out[i] = (c[i + 1] - c[lo]) / (i + 1 - lo)
    return out


def main():
    t0 = time.time()
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(11, 4.2))

    # ---- Left: variance of the bandit estimator as a function of the baseline
    q = np.array([1.0, 2.0, 4.0]) + 10.0
    theta = np.zeros(3)
    rng = np.random.default_rng(0)
    a, r = pg.sample_bandit(theta, q, 20_000, rng, noise=1.0)
    bs = np.linspace(0, 20, 81)
    var = [pg.total_variance(pg.reinforce_estimates(theta, a, r, baseline=b)) for b in bs]
    J = pg.softmax(theta) @ q
    ax1.semilogy(bs, var, color="C0")
    ax1.axvline(J, color="C1", ls="--", label=f"b = J = {J:.2f}")
    ax1.set_xlabel("baseline b")
    ax1.set_ylabel("total variance of one estimate")
    ax1.set_title("Bandit (rewards 11, 12, 14): spread vs baseline")
    ax1.legend()
    best_b = bs[int(np.argmin(var))]

    # ---- Right: CartPole learning curves
    labels = {"total": "whole return G_0", "to_go": "reward-to-go G_t", "baseline": "G_t - V(s_t)"}
    finals = {}
    for i, variant in enumerate(["total", "to_go", "baseline"]):
        curves = [pg.train_reinforce(variant, n_iters=N_ITERS, seed=s) for s in SEEDS]
        finals[variant] = np.mean([c[-40:].mean() for c in curves])
        n = min(len(c) for c in curves)
        for c in curves:
            ax2.plot(smooth(c[:n]), color=f"C{i}", alpha=0.25, lw=0.8)
        ax2.plot(np.mean([smooth(c[:n]) for c in curves], axis=0), color=f"C{i}", lw=2, label=labels[variant])
    ax2.set_xlabel("episode")
    ax2.set_ylabel("return (smoothed)")
    ax2.set_title(f"REINFORCE on CartPole ({len(SEEDS)} seeds each)")
    ax2.legend(loc="upper left")

    fig.tight_layout()
    out = os.path.join(HERE, "reinforce.png")
    fig.savefig(out, dpi=120)
    print(f"best constant baseline ~ {best_b:.2f} (J = {J:.2f}); final mean return: "
          + ", ".join(f"{k} {v:.0f}" for k, v in finals.items())
          + f"; saved {os.path.basename(out)} in {time.time() - t0:.0f}s")


if __name__ == "__main__":
    main()