"""
Generate a 2-panel figure comparing variance reduction techniques across trajectory horizons T.
Estimators compared:
  1. Vanilla REINFORCE: sum_t X_t * R(tau)
  2. Reward-to-Go (RTG): sum_t X_t * G_t
  3. RTG + State Baseline: sum_t X_t * (G_t - V(s_t))
  4. 1-Step TD Advantage: sum_t X_t * (r_t + V(s_{t+1}) - V(s_t))
"""

import numpy as np
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from scipy.special import softmax

BG_COLOR_FIG = "#1E222A"
BG_COLOR_AX = "#252931"
COLOR_TEXT = "#ECEFF4"
COLOR_MUTED = "#D8DEE9"
GRID_COLOR = "#3B4252"
SPINE_COLOR = "#4C566A"

COLOR_VANILLA = "#BF616A"
COLOR_RTG = "#EBCB8B"
COLOR_BASE = "#81A1C1"
COLOR_AC = "#A3BE8C"

def simulate_estimators(n_samples, T, n_feat, n_actions, seed=42):
    rng = np.random.default_rng(seed)
    theta_rng = np.random.default_rng(0)
    theta = theta_rng.normal(0, 0.01, (n_actions, n_feat))
    d = n_actions * n_feat

    scores_list, rewards_list, v_true_list = [], [], []

    for t in range(T):
        states = rng.uniform(-1, 1, (n_samples, n_feat))
        logits = states @ theta.T
        probs = softmax(logits, axis=1)
        cum_probs = np.cumsum(probs, axis=1)
        u = rng.uniform(0, 1, n_samples)
        actions = np.sum(cum_probs < u[:, None], axis=1)
        actions = np.clip(actions, 0, n_actions - 1)

        e_a = np.zeros((n_samples, n_actions))
        e_a[np.arange(n_samples), actions] = 1.0
        diff = e_a - probs

        score_vecs = (diff[:, :, None] * states[:, None, :]).reshape(n_samples, d)
        scores_list.append(score_vecs)

        theta_a = theta[actions]
        r = 1.0 + 0.1 * np.sum(theta_a * states, axis=1)
        rewards_list.append(r)

        exp_r_s = 1.0 + 0.1 * np.sum(probs * (states @ theta.T), axis=1)
        v_true_list.append(exp_r_s + (T - 1 - t) * 1.0)

    scores = np.stack(scores_list, axis=1)
    rewards = np.stack(rewards_list, axis=1)
    v_true = np.stack(v_true_list, axis=1)

    # 1. Vanilla REINFORCE
    g_vanilla = np.sum(scores * np.sum(rewards, axis=1, keepdims=True)[:, :, None], axis=1)

    # 2. Reward-to-Go
    G = np.cumsum(rewards[:, ::-1], axis=1)[:, ::-1]
    g_rtg = np.sum(scores * G[:, :, None], axis=1)

    # 3. RTG + Optimal State Baseline
    g_base = np.sum(scores * (G - v_true)[:, :, None], axis=1)

    # 4. 1-Step TD Advantage
    td_adv = np.zeros((n_samples, T))
    for t in range(T):
        if t < T - 1:
            td_adv[:, t] = rewards[:, t] + (T - 2 - t) * 1.0 - v_true[:, t]
        else:
            td_adv[:, t] = rewards[:, t] - v_true[:, t]
    g_ac = np.sum(scores * td_adv[:, :, None], axis=1)

    return (np.var(g_vanilla, axis=0).sum(), np.var(g_rtg, axis=0).sum(),
            np.var(g_base, axis=0).sum(), np.var(g_ac, axis=0).sum())

T_list = np.array([5, 10, 20, 40, 80, 160, 320, 640])
vars_vanilla, vars_rtg, vars_base, vars_ac = [], [], [], []
for T in T_list:
    vv, vr, vb, vac = simulate_estimators(30000, T, 4, 4)
    vars_vanilla.append(vv); vars_rtg.append(vr)
    vars_base.append(vb); vars_ac.append(vac)

vars_vanilla, vars_rtg = np.array(vars_vanilla), np.array(vars_rtg)
vars_base, vars_ac = np.array(vars_base), np.array(vars_ac)
s_van = np.polyfit(np.log(T_list), np.log(vars_vanilla), 1)[0]
s_rtg = np.polyfit(np.log(T_list), np.log(vars_rtg), 1)[0]
s_base = np.polyfit(np.log(T_list), np.log(vars_base), 1)[0]
s_ac = np.polyfit(np.log(T_list), np.log(vars_ac), 1)[0]

# --- 1. Figure: Reward-to-Go vs Vanilla ---
fig1, (ax1a, ax1b) = plt.subplots(1, 2, figsize=(9.0, 3.8))
for ax in (ax1a, ax1b):
    ax.set_facecolor("none")
    ax.spines["top"].set_visible(False); ax.spines["right"].set_visible(False)
    ax.spines["left"].set_color(SPINE_COLOR); ax.spines["bottom"].set_color(SPINE_COLOR)
    ax.tick_params(which='both', colors=COLOR_MUTED, labelsize=9)
    ax.grid(True, linestyle="--", alpha=0.4, color=GRID_COLOR)
fig1.patch.set_facecolor("none")

ax1a.loglog(T_list, vars_vanilla, "o-", color=COLOR_VANILLA, lw=2, ms=5, label=f"Vanilla REINFORCE (empirical slope {s_van:.2f})")
ax1a.loglog(T_list, vars_rtg, "s-", color=COLOR_RTG, lw=2, ms=5, label=f"Reward-to-Go (empirical slope {s_rtg:.2f})")
ax1a.set_xlabel("Trajectory length $T$", fontsize=10, color=COLOR_TEXT)
ax1a.set_ylabel(r"Estimator variance $\mathrm{Var}(\hat{g})$", fontsize=10, color=COLOR_TEXT)
ax1a.set_title(r"Variance Scaling with Horizon $T$", fontsize=11, fontweight="bold", color=COLOR_TEXT, pad=8)
leg1a = ax1a.legend(fontsize=8, loc="upper left", frameon=False)
plt.setp(leg1a.get_texts(), color=COLOR_TEXT)

ax1b.loglog(T_list, vars_vanilla / vars_rtg, "s-", color=COLOR_RTG, lw=2, ms=5, label="Reward-to-Go vs Vanilla")
ax1b.axhline(3.0, color=COLOR_MUTED, linestyle=":", lw=1.5, label="Asymptotic limit (3.0)")
ax1b.set_xlabel("Trajectory length $T$", fontsize=10, color=COLOR_TEXT)
ax1b.set_ylabel(r"Variance Reduction Factor", fontsize=10, color=COLOR_TEXT)
ax1b.set_title(r"Variance Reduction Ratio (Higher is Better)", fontsize=11, fontweight="bold", color=COLOR_TEXT, pad=8)
leg1b = ax1b.legend(fontsize=8, loc="upper left", frameon=False)
plt.setp(leg1b.get_texts(), color=COLOR_TEXT)
fig1.tight_layout()
fig1.savefig("/Users/pedregosa/dev/webpage/content/images/2026/policy_gradient_rtg_variance.png", dpi=200, bbox_inches="tight", transparent=True)
plt.close(fig1)

# --- 2. Figure: RTG + State Baseline ---
fig2, (ax2a, ax2b) = plt.subplots(1, 2, figsize=(9.0, 3.8))
for ax in (ax2a, ax2b):
    ax.set_facecolor("none")
    ax.spines["top"].set_visible(False); ax.spines["right"].set_visible(False)
    ax.spines["left"].set_color(SPINE_COLOR); ax.spines["bottom"].set_color(SPINE_COLOR)
    ax.tick_params(which='both', colors=COLOR_MUTED, labelsize=9)
    ax.grid(True, linestyle="--", alpha=0.4, color=GRID_COLOR)
fig2.patch.set_facecolor("none")

ax2a.loglog(T_list, vars_vanilla, "o-", color=COLOR_VANILLA, lw=2, ms=5, label=f"Vanilla REINFORCE (empirical slope {s_van:.2f})")
ax2a.loglog(T_list, vars_rtg, "s-", color=COLOR_RTG, lw=2, ms=5, label=f"Reward-to-Go (empirical slope {s_rtg:.2f})")
ax2a.loglog(T_list, vars_base, "^-", color=COLOR_BASE, lw=2, ms=5, label=f"Reward-to-Go + State Baseline (empirical slope {s_base:.2f})")
ax2a.set_xlabel("Trajectory length $T$", fontsize=10, color=COLOR_TEXT)
ax2a.set_ylabel(r"Estimator variance $\mathrm{Var}(\hat{g})$", fontsize=10, color=COLOR_TEXT)
ax2a.set_title(r"Variance Scaling with Horizon $T$", fontsize=11, fontweight="bold", color=COLOR_TEXT, pad=8)
leg2a = ax2a.legend(fontsize=8, loc="upper left", frameon=False)
plt.setp(leg2a.get_texts(), color=COLOR_TEXT)

ax2b.loglog(T_list, vars_vanilla / vars_rtg, "s-", color=COLOR_RTG, lw=2, ms=5, label="Reward-to-Go vs Vanilla")
ax2b.loglog(T_list, vars_vanilla / vars_base, "^-", color=COLOR_BASE, lw=2, ms=5, label="Reward-to-Go + Baseline vs Vanilla")
ax2b.set_xlabel("Trajectory length $T$", fontsize=10, color=COLOR_TEXT)
ax2b.set_ylabel(r"Variance Reduction Factor", fontsize=10, color=COLOR_TEXT)
ax2b.set_title(r"Variance Reduction Ratio (Higher is Better)", fontsize=11, fontweight="bold", color=COLOR_TEXT, pad=8)
leg2b = ax2b.legend(fontsize=8, loc="upper left", frameon=False)
plt.setp(leg2b.get_texts(), color=COLOR_TEXT)
fig2.tight_layout()
fig2.savefig("/Users/pedregosa/dev/webpage/content/images/2026/policy_gradient_baseline_variance.png", dpi=200, bbox_inches="tight", transparent=True)
fig2.savefig("/Users/pedregosa/dev/webpage/content/images/2026/policy_gradient_baseline_variance.svg", bbox_inches="tight", transparent=True)
plt.close(fig2)

# --- 3. Figure: All Estimators (including 1-Step TD) ---
fig3, (ax3a, ax3b) = plt.subplots(1, 2, figsize=(9.0, 3.8))
for ax in (ax3a, ax3b):
    ax.set_facecolor("none")
    ax.spines["top"].set_visible(False); ax.spines["right"].set_visible(False)
    ax.spines["left"].set_color(SPINE_COLOR); ax.spines["bottom"].set_color(SPINE_COLOR)
    ax.tick_params(which='both', colors=COLOR_MUTED, labelsize=9)
    ax.grid(True, linestyle="--", alpha=0.4, color=GRID_COLOR)
fig3.patch.set_facecolor("none")

ax3a.loglog(T_list, vars_vanilla, "o-", color=COLOR_VANILLA, lw=2, ms=5, label=f"Vanilla REINFORCE (empirical slope {s_van:.2f})")
ax3a.loglog(T_list, vars_rtg, "s-", color=COLOR_RTG, lw=2, ms=5, label=f"Reward-to-Go (empirical slope {s_rtg:.2f})")
ax3a.loglog(T_list, vars_base, "^-", color=COLOR_BASE, lw=2, ms=5, label=f"Reward-to-Go + State Baseline (empirical slope {s_base:.2f})")
ax3a.loglog(T_list, vars_ac, "d-", color=COLOR_AC, lw=2, ms=5, label=f"1-Step TD Advantage (empirical slope {s_ac:.2f})")
ax3a.set_xlabel("Trajectory length $T$", fontsize=10, color=COLOR_TEXT)
ax3a.set_ylabel(r"Estimator variance $\mathrm{Var}(\hat{g})$", fontsize=10, color=COLOR_TEXT)
ax3a.set_title(r"Variance Scaling with Horizon $T$", fontsize=11, fontweight="bold", color=COLOR_TEXT, pad=8)
leg3a = ax3a.legend(fontsize=7.5, loc="upper left", frameon=False)
plt.setp(leg3a.get_texts(), color=COLOR_TEXT)

ax3b.loglog(T_list, vars_vanilla / vars_rtg, "s-", color=COLOR_RTG, lw=2, ms=5, label="Reward-to-Go vs Vanilla")
ax3b.loglog(T_list, vars_vanilla / vars_base, "^-", color=COLOR_BASE, lw=2, ms=5, label="Reward-to-Go + Baseline vs Vanilla")
ax3b.loglog(T_list, vars_vanilla / vars_ac, "d-", color=COLOR_AC, lw=2, ms=5, label="1-Step TD Advantage vs Vanilla")
ax3b.set_xlabel("Trajectory length $T$", fontsize=10, color=COLOR_TEXT)
ax3b.set_ylabel(r"Variance Reduction Factor", fontsize=10, color=COLOR_TEXT)
ax3b.set_title(r"Variance Reduction Ratio (Higher is Better)", fontsize=11, fontweight="bold", color=COLOR_TEXT, pad=8)
leg3b = ax3b.legend(fontsize=7.5, loc="upper left", frameon=False)
plt.setp(leg3b.get_texts(), color=COLOR_TEXT)
fig3.tight_layout()
fig3.savefig("/Users/pedregosa/dev/webpage/content/images/2026/policy_gradient_variance_reduction.png", dpi=200, bbox_inches="tight", transparent=True)
plt.close(fig3)

print(f"Slopes: Vanilla={s_van:.3f}, RTG={s_rtg:.3f}, Base={s_base:.3f}, AC={s_ac:.3f}")
print("All three variance reduction plots successfully saved!")