from pathlib import Path
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation, PillowWriter

# ==============================================================================
# 1. COMPUTATIONAL SETUP & OPTIMIZATION TRAJECTORIES
# ==============================================================================

# Objective function definition: J(theta) = 7.0 - 0.5 * theta^2
theta_grid = np.linspace(-3.8, 3.8, 200)
J_grid = 7.0 - 0.5 * (theta_grid ** 2)

n_steps = 40
alpha = 0.065  # Learning rate
theta_init = 2.2
np.random.seed(123)

V_init_base = 7.0 - 0.5 * (theta_init ** 2)
J_base_grid = J_grid - V_init_base

theta_unc = [theta_init]
theta_cen = [theta_init]
actions_unc, rewards_unc = [], []
actions_cen, rewards_cen = [], []

# Pre-compute optimization trajectory
for t in range(n_steps):
    z = np.random.normal(0, 1.0)

    # 1. Uncentered Step (without baseline)
    th_u = theta_unc[-1]
    a_u = th_u + z
    r_u = 7.5 - 0.5 * (a_u ** 2)
    g_u = (a_u - th_u) * r_u
    next_u = np.clip(th_u + alpha * g_u, -2.8, 2.8)
    theta_unc.append(next_u)
    actions_unc.append(a_u)
    rewards_unc.append(r_u)

    # 2. Centered Step (with baseline V(s_t) = J(theta_t))
    th_c = theta_cen[-1]
    a_c = th_c + z
    r_c = 7.5 - 0.5 * (a_c ** 2)
    v_c = 7.0 - 0.5 * (th_c ** 2)
    g_c = (a_c - th_c) * (r_c - v_c)
    next_c = np.clip(th_c + alpha * g_c, -2.8, 2.8)
    theta_cen.append(next_c)
    actions_cen.append(a_c)
    rewards_cen.append(r_c)

# Pre-compute surrogate curves for each step
surrogates_unc = []
surrogates_cen = []
for t in range(n_steps):
    th_u, a_u, r_u = theta_unc[t], actions_unc[t], rewards_unc[t]
    j_u = 7.0 - 0.5 * (th_u ** 2)
    L_u = j_u - 0.5 * ((theta_grid - a_u)**2 - (th_u - a_u)**2) * r_u
    surrogates_unc.append(L_u)

    th_c, a_c, r_c = theta_cen[t], actions_cen[t], rewards_cen[t]
    v_c = 7.0 - 0.5 * (th_c ** 2)
    L_c = (v_c - V_init_base) - 0.5 * ((theta_grid - a_c)**2 - (th_c - a_c)**2) * (r_c - v_c)
    surrogates_cen.append(L_c)

j_unc_history = [7.0 - 0.5 * (th ** 2) for th in theta_unc]
j_cen_history = [(7.0 - 0.5 * (th ** 2)) - V_init_base for th in theta_cen]

# ==============================================================================
# 2. PLOTTING & ANIMATION
# ==============================================================================

BG_COLOR = '#151515'
TEXT_COLOR = '#dddddd'
GRID_COLOR = '#333333'

fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(14, 6), dpi=150, facecolor=BG_COLOR)

for ax in (ax1, ax2):
    ax.set_facecolor(BG_COLOR)
    ax.tick_params(colors=TEXT_COLOR, labelsize=12)
    ax.xaxis.label.set_color(TEXT_COLOR)
    ax.yaxis.label.set_color(TEXT_COLOR)
    ax.title.set_color(TEXT_COLOR)
    ax.grid(True, color=GRID_COLOR, linestyle='--', alpha=0.5)
    for spine in ax.spines.values():
        spine.set_color('#555555')

# Left Plot: Uncentered Trajectory
ax1.plot(theta_grid, J_grid, color='#ffff00', linewidth=3, label='Expected Return $J(\\theta)$', zorder=5)
ax1.plot(0, 7.0, '*', color='#ffff00', markersize=15, label='Optimum $\\theta^* = 0$', zorder=7)
line_surr_u, = ax1.plot([], [], color='#ff4444', linewidth=2.4, label='Current Surrogate $L(\\theta; a_t)$', zorder=4)
trail_u, = ax1.plot([], [], 'o', color='#00ff88', markersize=5, alpha=0.6, label='Step History $\\theta_t$', zorder=6)
dot_u, = ax1.plot([], [], 'o', color='#00ff88', markersize=10, label='Current Estimate', zorder=8)
lines_past_u = [ax1.plot([], [], color='#ff4444', alpha=0.25, linewidth=1.2, zorder=3)[0] for _ in range(n_steps)]

ax1.set_title('Without Baseline', fontsize=14, fontweight='bold', pad=15)
ax1.set_xlabel('Policy Parameter $\\theta$', fontsize=14)
ax1.set_ylabel('Expected Return $J(\\theta)$', fontsize=14)
ax1.set_ylim(0.0, 9.2)
ax1.set_xlim(-3.6, 3.6)
ax1.text(-3.45, 8.1, 'High Variance:\nNoisy gradients cause unstable optimization',
         fontsize=11, color=TEXT_COLOR, bbox=dict(facecolor='#222222', alpha=0.85, edgecolor='#ff4444', boxstyle='round,pad=0.4'))

# Right Plot: Centered Trajectory
ax2.plot(theta_grid, J_base_grid, color='#ffff00', linewidth=3, label='Modified Objective $J_{\\text{base}}(\\theta)$', zorder=5)
ax2.plot(0, 7.0 - V_init_base, '*', color='#ffff00', markersize=15, label='Optimum $\\theta^* = 0$', zorder=7)
line_surr_c, = ax2.plot([], [], color='#00d2ff', linewidth=2.4, label='Current Surrogate $L_{\\text{base}}(\\theta; a_t)$', zorder=4)
trail_c, = ax2.plot([], [], 'o', color='#00ff88', markersize=5, alpha=0.6, label='Step History $\\theta_t$', zorder=6)
dot_c, = ax2.plot([], [], 'o', color='#00ff88', markersize=10, label='Current Estimate', zorder=8)
lines_past_c = [ax2.plot([], [], color='#00d2ff', alpha=0.30, linewidth=1.2, zorder=3)[0] for _ in range(n_steps)]

ax2.set_title('With Baseline', fontsize=14, fontweight='bold', pad=15)
ax2.set_xlabel('Policy Parameter $\\theta$', fontsize=14)
ax2.set_ylabel('Centered Objective $J_{\\text{base}}(\\theta)$', fontsize=14)
ax2.set_ylim(-4.8, 4.4)
ax2.set_xlim(-3.6, 3.6)
ax2.text(-3.45, 3.3, 'Low Variance:\nBaseline leads to more stable dynamics',
         fontsize=11, color=TEXT_COLOR, bbox=dict(facecolor='#222222', alpha=0.85, edgecolor='#00d2ff', boxstyle='round,pad=0.4'))

for ax in (ax1, ax2):
    leg = ax.legend(loc='lower left', fontsize=9.5, facecolor='#222222', edgecolor='#555555')
    plt.setp(leg.get_texts(), color=TEXT_COLOR)

plt.tight_layout()

def update_anim(frame):
    if frame == 0:
        for line in lines_past_u + lines_past_c:
            line.set_data([], [])

    if frame > 0:
        lines_past_u[frame - 1].set_data(theta_grid, surrogates_unc[frame - 1])
        lines_past_c[frame - 1].set_data(theta_grid, surrogates_cen[frame - 1])

    line_surr_u.set_data(theta_grid, surrogates_unc[frame])
    trail_u.set_data(theta_unc[:frame + 1], j_unc_history[:frame + 1])
    dot_u.set_data([theta_unc[frame]], [j_unc_history[frame]])

    line_surr_c.set_data(theta_grid, surrogates_cen[frame])
    trail_c.set_data(theta_cen[:frame + 1], j_cen_history[:frame + 1])
    dot_c.set_data([theta_cen[frame]], [j_cen_history[frame]])

    return [line_surr_u, trail_u, dot_u, line_surr_c, trail_c, dot_c] + lines_past_u + lines_past_c

ani = FuncAnimation(fig, update_anim, frames=n_steps, blit=False)

output_dir = Path(__file__).resolve().parents[2] / "images" / "2026"
output_dir.mkdir(parents=True, exist_ok=True)
gif_path = output_dir / "objective_landscape_1d.gif"

ani.save(gif_path, writer=PillowWriter(fps=3))
plt.close()
print(f"Saved animated GIF to {gif_path}")