"""
Publication-quality figure: QNG vs GQNG on curved vs flat manifolds.
Surface shape matches the Gemini reference image — a dramatic hill/landscape
with steep drop-offs, visible contour lines, and clear geometric elements.
"""
import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d.art3d import Poly3DCollection
from matplotlib.patches import FancyArrowPatch
from mpl_toolkits.mplot3d import proj3d
import matplotlib.patheffects as pe
from matplotlib.lines import Line2D

# ── 3D Arrow utility ────────────────────────────────────────────────────
class Arrow3D(FancyArrowPatch):
    def __init__(self, xs, ys, zs, *args, **kwargs):
        super().__init__((0, 0), (0, 0), *args, **kwargs)
        self._verts3d = xs, ys, zs
    def do_3d_projection(self, renderer=None):
        xs3d, ys3d, zs3d = self._verts3d
        xs, ys, zs = proj3d.proj_transform(xs3d, ys3d, zs3d, self.axes.M)
        self.set_positions((xs[0], ys[0]), (xs[1], ys[1]))
        return min(zs)

# ── Style ───────────────────────────────────────────────────────────────
plt.rcParams.update({
    'font.family': 'serif',
    'font.serif': ['DejaVu Serif', 'Computer Modern Roman'],
    'font.size': 9,
    'mathtext.fontset': 'cm',
})

BLUE     = '#1a56db'
RED      = '#c0392b'
DARK_RED = '#8b0000'
PURPLE   = '#4a2d8a'
BLACK    = '#1a1a1a'
ERR_COL  = '#d63031'
ws  = pe.withStroke(linewidth=3.5, foreground='white')
ws2 = pe.withStroke(linewidth=2.5, foreground='white')

# ── Surface: Gaussian hill + modulation ─────────────────────────────────
# Matches the Gemini original: a dramatic hill with visible contour lines
# and steep drop-offs. NOT a shallow saddle.
def surface(x, y):
    # Main hill (tall Gaussian — gives dramatic curvature)
    hill = 2.8 * np.exp(-0.12*(x + 0.3)**2 - 0.08*(y - 0.2)**2)
    # Secondary bump for contour interest
    bump = 0.6 * np.exp(-0.25*(x - 1.5)**2 - 0.20*(y + 0.8)**2)
    # Subtle modulation for visual texture
    mod  = 0.15 * np.sin(1.4*x) * np.cos(1.1*y)
    return hill + bump + mod

def surf_grad(x, y, eps=1e-5):
    dzdx = (surface(x+eps, y) - surface(x-eps, y)) / (2*eps)
    dzdy = (surface(x, y+eps) - surface(x, y-eps)) / (2*eps)
    return dzdx, dzdy

def bez(p0, p1, p2, n=100):
    t = np.linspace(0, 1, n)
    return ((1-t)**2)[:,None]*p0 + 2*((1-t)*t)[:,None]*p1 + (t**2)[:,None]*p2


# ═══════════════════════════════════════════════════════════════════════
fig = plt.figure(figsize=(15, 7))

# ════════════════════════════════════════════════════════════════════════
#  (a) CURVED MANIFOLD
# ════════════════════════════════════════════════════════════════════════
ax = fig.add_subplot(121, projection='3d', computed_zorder=False)

# ── Render surface ──────────────────────────────────────────────────────
N = 160
xs = np.linspace(-3.0, 3.5, N)
ys = np.linspace(-2.8, 3.0, N)
X, Y = np.meshgrid(xs, ys)
Z = surface(X, Y)

# Surface with topographic coloring
ax.plot_surface(X, Y, Z, cmap='YlGnBu_r', alpha=0.55,
                linewidth=0, antialiased=True, zorder=1,
                rstride=2, cstride=2)

# Contour lines ON the surface (gives the topographic look)
ax.contour(X, Y, Z, levels=15, colors='#2d5f2d', alpha=0.30,
           linewidths=0.4, zorder=2)

# Projected contours at floor
ax.contour(X, Y, Z, levels=12, cmap='Greens_r', alpha=0.20,
           linewidths=0.3, offset=Z.min() - 0.2, zorder=0)

# ── Geometry ────────────────────────────────────────────────────────────
# θ: on the upper-left slope of the hill (high elevation)
th_xy = np.array([-1.5, 0.5])
th_z = surface(*th_xy)
th = np.array([*th_xy, th_z])

# θ_next: where the geodesic arrives (further along the surface, lower)
tn_xy = np.array([1.0, 1.8])
tn_z = surface(*tn_xy)
tn = np.array([*tn_xy, tn_z])

# θ': where QNG retraction lands (CLEARLY separated from θ_next)
tp_xy = np.array([1.5, 0.2])
tp_z = surface(*tp_xy)
tp = np.array([*tp_xy, tp_z])

# ── Tangent plane at θ ──────────────────────────────────────────────────
dzdx, dzdy = surf_grad(*th_xy)
# Tangent plane basis vectors in 3D
e1 = np.array([1, 0, dzdx]); e1 /= np.linalg.norm(e1)
e2 = np.array([0, 1, dzdy]); e2 /= np.linalg.norm(e2)

# Rectangle corners (extending mainly to the right from θ)
corners = [
    th - 0.4*e1 - 0.7*e2,
    th + 2.4*e1 - 0.7*e2,
    th + 2.4*e1 + 0.9*e2,
    th - 0.4*e1 + 0.9*e2,
]
ax.add_collection3d(Poly3DCollection(
    [corners], alpha=0.32, facecolor='#d8d8d8',
    edgecolor='#777', linewidth=0.8, zorder=5))

# ── Tangent vector v = -g̃⁻¹∇L (arrow on tangent plane) ─────────────
# Direction: from θ toward tn projected onto tangent plane
disp = tn_xy - th_xy
v3 = np.array([disp[0], disp[1], dzdx*disp[0] + dzdy*disp[1]])
v3_unit = v3 / np.linalg.norm(v3)
arrow_length = 2.5

P = th
Pp = th + arrow_length * v3_unit

ax.add_artist(Arrow3D(
    [P[0], Pp[0]], [P[1], Pp[1]], [P[2], Pp[2]],
    mutation_scale=18, lw=2.5, arrowstyle='-|>',
    color=DARK_RED, zorder=14))

# ── Blue geodesic: θ → θ_next ──────────────────────────────────────────
# Wide curve that follows the manifold curvature
geo_mid = 0.5*(th_xy + tn_xy) + np.array([0.5, -0.7])  # big lateral bend
geo_xy = bez(th_xy, geo_mid, tn_xy)
geo_z = surface(geo_xy[:,0], geo_xy[:,1])
ax.plot(geo_xy[:,0], geo_xy[:,1], geo_z,
        color=BLUE, lw=3.0, zorder=9, solid_capstyle='round')

# ── Red QNG retraction: θ → θ' ─────────────────────────────────────────
# Straighter path, clearly diverging from the geodesic
ret_mid = 0.5*(th_xy + tp_xy) + np.array([0.05, 0.0])
ret_xy = bez(th_xy, ret_mid, tp_xy)
ret_z = surface(ret_xy[:,0], ret_xy[:,1])
ax.plot(ret_xy[:,0], ret_xy[:,1], ret_z,
        color=RED, lw=2.6, zorder=8, solid_capstyle='round')

# ── Retraction error (dashed line θ_next ↔ θ') ─────────────────────────
ax.plot([tn[0], tp[0]], [tn[1], tp[1]], [tn[2], tp[2]],
        color=ERR_COL, lw=2.0, ls='--', zorder=10, alpha=0.85)
# Double-headed arrow
ax.add_artist(Arrow3D(
    [tn[0], tp[0]], [tn[1], tp[1]], [tn[2], tp[2]],
    mutation_scale=11, lw=1.4, arrowstyle='<|-|>',
    color=ERR_COL, zorder=11))

# ── Points ──────────────────────────────────────────────────────────────
S = 70
ax.scatter(*th, color=BLACK, s=S, zorder=16, depthshade=False)
ax.scatter(*tn, color=BLUE, s=S, zorder=16, depthshade=False,
           edgecolors='white', linewidth=0.7)
ax.scatter(*tp, color=RED, s=S, zorder=16, depthshade=False,
           edgecolors='white', linewidth=0.7)

# ── Labels ──────────────────────────────────────────────────────────────
# ONLY essential labels — no clutter
pt_lbl = dict(fontsize=13, fontweight='bold', path_effects=[ws])

ax.text(th[0]-0.5, th[1]+0.15, th[2]+0.12,
        r'$\theta$', color=BLACK, **pt_lbl, zorder=20)
ax.text(tn[0]+0.20, tn[1]+0.30, tn[2]+0.20,
        r'$\theta_{\mathrm{next}}$', color=BLUE, **pt_lbl, zorder=20)
ax.text(tp[0]+0.15, tp[1]-0.35, tp[2]+0.05,
        r"$\theta'$", color=RED, **pt_lbl, zorder=20)

# v label — well above the arrow, clear of P'
vmid = 0.15*P + 0.85*Pp
ax.text(vmid[0]+0.1, vmid[1]+0.55, vmid[2]+0.28,
        r'$\mathbf{v}=-\tilde{g}^{-1}\nabla_\theta\mathcal{L}$',
        color=DARK_RED, fontsize=10, fontweight='bold',
        path_effects=[ws], zorder=20)

# P and P' on tangent plane
ax.text(P[0]-0.20, P[1]+0.22, P[2]+0.10,
        r'$P$', color=BLACK, fontsize=10, path_effects=[ws2], zorder=20)
ax.text(Pp[0]+0.15, Pp[1]-0.10, Pp[2]+0.10,
        r"$P'$", color=DARK_RED, fontsize=10, path_effects=[ws2], zorder=20)

# Retraction error label — clearly beside the dashed line
me = 0.5*(tn + tp)
ax.text(me[0]+0.50, me[1]-0.10, me[2]+0.22,
        'Retraction\nError', color=ERR_COL, fontsize=9.5,
        fontstyle='italic', ha='left', path_effects=[ws], zorder=20)

# Tangent plane label
ax.text(corners[1][0]-0.3, corners[1][1]-0.15, corners[1][2]-0.08,
        r'Tangent $\mathcal{T}_\theta\mathcal{M}$',
        color='#666', fontsize=9, fontstyle='italic',
        path_effects=[ws2], zorder=20)

# Manifold label (far from action, bottom-left)
ax.text(-2.8, -2.2, surface(-2.8, -2.2)+0.15,
        r'Curved Riemannian' '\n' r'Manifold $\mathcal{M}_{\mathrm{Curved}}$',
        color='#2d5f2d', fontsize=9, path_effects=[ws], zorder=20)

# g̃ label
ax.text(-2.8, -1.4, surface(-2.8, -1.4)+0.15,
        r'$\tilde{g}$', color='#2d5f2d', fontsize=12,
        path_effects=[ws], zorder=20)

# ── View ────────────────────────────────────────────────────────────────
ax.set_xlim(-3.0, 3.5)
ax.set_ylim(-2.8, 3.0)
ax.set_zlim(Z.min()-0.3, Z.max()+0.5)
ax.view_init(elev=33, azim=-55)
ax.set_axis_off()

# ── Text annotations outside the 3D view (using fig coords) ────────────
# Path descriptions at bottom of left panel
ax.text2D(0.02, 0.08,
          r'$\mathbf{Blue}$: True Geodesic Path, $R^E(\mathbf{v})=\mathrm{Exp}_\theta(\mathbf{v})$ (Ideal)',
          transform=ax.transAxes, fontsize=8.5, color=BLUE,
          fontweight='bold')
ax.text2D(0.02, 0.02,
          r'$\mathbf{Red}$: Standard QNG Retraction Path (deviates due to geometry infidelity)',
          transform=ax.transAxes, fontsize=8.5, color=RED,
          fontweight='bold')

ax.set_title('Case 1: Optimization on a Curved\nRiemannian Manifold (Standard QNG)',
             fontsize=11, fontweight='bold', pad=8, y=1.0)


# ════════════════════════════════════════════════════════════════════════
#  (b) FLAT MANIFOLD (DFS-QNG)
# ════════════════════════════════════════════════════════════════════════
ax2 = fig.add_subplot(122, projection='3d', computed_zorder=False)

# Light green flat plane
xf = np.linspace(-2.8, 2.8, 35)
XF, YF = np.meshgrid(xf, xf)
ax2.plot_surface(XF, YF, np.zeros_like(XF), color='#ddeedd', alpha=0.40,
                 linewidth=0, antialiased=True, zorder=1)

# Subtle grid
for val in np.linspace(-2.8, 2.8, 12):
    ax2.plot([val, val], [-2.8, 2.8], [0, 0], color='#aaccaa', lw=0.25, zorder=2)
    ax2.plot([-2.8, 2.8], [val, val], [0, 0], color='#aaccaa', lw=0.25, zorder=2)

# Points
ft = np.array([-2.0, -1.2, 0])
fn = np.array([1.8, 1.2, 0])

# Single path (geodesic ≡ retraction ≡ straight line)
ax2.plot([ft[0], fn[0]], [ft[1], fn[1]], [0, 0],
         color=PURPLE, lw=3.5, zorder=10, solid_capstyle='round')

# Arrow overlay
ax2.add_artist(Arrow3D(
    [ft[0], fn[0]], [ft[1], fn[1]], [0.01, 0.01],
    mutation_scale=18, lw=2.5, arrowstyle='-|>',
    color=DARK_RED, zorder=12))

# Points
ax2.scatter(*ft, color=BLACK, s=S, zorder=16, depthshade=False)
ax2.scatter(*fn, color=PURPLE, s=S, zorder=16, depthshade=False,
            edgecolors='white', linewidth=0.7)

# ── Labels ──────────────────────────────────────────────────────────────
ax2.text(ft[0]-0.55, ft[1]-0.45, 0.05,
         r'$\theta$', color=BLACK, **pt_lbl, zorder=20)
ax2.text(fn[0]+0.15, fn[1]+0.25, 0.05,
         r"$\theta_{\mathrm{next}}=\theta'$", color=PURPLE,
         **pt_lbl, zorder=20)

# v label above line
mf = 0.5*(ft + fn)
ax2.text(mf[0]+0.1, mf[1]+0.65, 0.10,
         r'$\mathbf{v}=-\tilde{g}^{-1}\nabla_\theta\mathcal{L}$',
         color=DARK_RED, fontsize=10, fontweight='bold',
         path_effects=[ws], zorder=20)

# Path equivalence below line
ax2.text(mf[0]-0.8, mf[1]-0.75, -0.02,
         r'Path $\equiv$ Geodesic $\equiv$ Retraction',
         color=PURPLE, fontsize=9, fontstyle='italic',
         path_effects=[ws], zorder=20)
ax2.text(mf[0]-0.3, mf[1]-1.15, -0.03,
         r'$R^E(\mathbf{v}) \equiv R^I(\mathbf{v}) \equiv$ Straight Line',
         color=PURPLE, fontsize=9, fontstyle='italic',
         path_effects=[ws], zorder=20)

# Properties (top-right corner, away from line)
ax2.text(1.5, 2.5, 0.12,
         r'$\tilde{g},\ \Gamma = 0$',
         color='#444', fontsize=11, path_effects=[ws], zorder=20)
ax2.text(-0.5, 2.5, 0.08,
         r'Manifold $\equiv$ Tangent Space',
         color='#444', fontsize=9, path_effects=[ws], zorder=20)
ax2.text(-0.8, 2.1, 0.06,
         r'$(\mathcal{M}_{\mathrm{Flat}}\equiv\mathcal{T}_\theta\mathcal{M}_{\mathrm{Flat}})$',
         color='#444', fontsize=8.5, path_effects=[ws], zorder=20)

# Summary (bottom)
ax2.text2D(0.08, 0.06,
           'All points and paths coincide perfectly.\n'
           'No geometry error or retraction deviation.\n'
           'Updates are exact geodesics on flattened manifolds.',
           transform=ax2.transAxes, fontsize=8, color='#444')

# Flat manifold label
ax2.text(2.0, -2.3, 0.02,
         'Flat Manifold\n(DFS-QNG)', color='#2d5f2d', fontsize=9,
         ha='right', path_effects=[ws], zorder=20)

ax2.set_xlim(-2.8, 2.8)
ax2.set_ylim(-2.8, 2.8)
ax2.set_zlim(-0.6, 0.7)
ax2.view_init(elev=28, azim=-52)
ax2.set_axis_off()
ax2.set_title('Case 2: DFS-QNG\n(Dually Flat Stratified Manifold)',
              fontsize=11, fontweight='bold', pad=8, y=1.0)


# ── Shared legend ───────────────────────────────────────────────────────
legend = [
    Line2D([0],[0], color=BLUE, lw=3.0,
           label=r'True Geodesic (GQNG)'),
    Line2D([0],[0], color=RED, lw=2.6,
           label='Standard QNG Retraction'),
    Line2D([0],[0], color=DARK_RED, lw=2.5,
           label=r'Tangent vector $\mathbf{v}$'),
    Line2D([0],[0], color=ERR_COL, lw=1.8, ls='--',
           label='Retraction Error'),
]
fig.legend(handles=legend, loc='lower center', ncol=4, fontsize=9,
           frameon=True, fancybox=True, framealpha=0.92,
           edgecolor='#ccc', bbox_to_anchor=(0.5, -0.01),
           columnspacing=2)

plt.subplots_adjust(wspace=0.02, bottom=0.08, top=0.90)
plt.show()

# ── Save ────────────────────────────────────────────────────────────────
# for fmt in ['pdf', 'png', 'svg']:
#     fig.savefig(f'/home/claude/manifold_fig.{fmt}', format=fmt,
#                 bbox_inches='tight', pad_inches=0.15,
#                 **(dict(dpi=400) if fmt == 'png' else {}))
# print("✓ Saved manifold_fig.{pdf,png,svg}")