"""
Loss Landscape Visualization for @fminxyz Telegram Channel
Generates a 3D animated visualization of a neural network loss surface
showing SGD trajectories navigating sharp vs flat minima.
"""

import numpy as np
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import matplotlib.animation as animation
from matplotlib.colors import LightSource
from mpl_toolkits.mplot3d import Axes3D
import warnings
warnings.filterwarnings('ignore')

# ─── Configuration ────────────────────────────────────────────────────────────
OUTPUT_VIDEO = "/root/Strategy/content/drafts/loss_landscape_animation.mp4"
OUTPUT_THUMB = "/root/Strategy/content/drafts/loss_landscape_thumbnail.png"
FIG_SIZE_IN = (10.8, 10.8)   # 1080x1080 @ 100 dpi
DPI = 100
FPS = 30
DURATION_SEC = 25
N_FRAMES = FPS * DURATION_SEC

GRID_N = 300       # resolution of the landscape
TRAJ_STEPS = 500   # optimizer steps to pre-compute


# ─── Loss Landscape Construction ──────────────────────────────────────────────
def build_landscape(n=GRID_N):
    """
    Construct a realistic-looking loss surface that includes:
      - A broad, flat (good-generalizing) global minimum region
      - Several sharp narrow minima (poor generalizers)
      - Saddle points
      - A rough high-loss plateau (initial random parameters live here)
    """
    rng = np.random.default_rng(42)

    x = np.linspace(-4, 4, n)
    y = np.linspace(-4, 4, n)
    X, Y = np.meshgrid(x, y)

    # ── Base landscape: gentle bowl (flat minimum region) ──────────────────
    Z = 0.10 * (X**2 + Y**2)

    # ── Broad flat minimum (the "good" generalizing basin) ─────────────────
    flat_cx, flat_cy = -1.2, -0.8
    flat_r = 1.4
    dist_flat = np.sqrt((X - flat_cx)**2 + (Y - flat_cy)**2)
    Z -= 2.8 * np.exp(-dist_flat**2 / (2 * flat_r**2))

    # ── Sharp narrow minima (poor generalizers) ─────────────────────────────
    sharp_minima = [
        ( 2.0,  1.5, -2.2, 0.12),
        (-0.5,  2.8, -1.8, 0.10),
        ( 2.8, -1.8, -2.0, 0.09),
        (-3.0,  0.5, -1.6, 0.11),
        ( 0.8, -2.8, -1.9, 0.13),
        (-2.5, -2.0, -1.5, 0.10),
    ]
    for cx, cy, depth, width in sharp_minima:
        dist = np.sqrt((X - cx)**2 + (Y - cy)**2)
        Z += depth * np.exp(-dist**2 / (2 * width**2))

    # ── Saddle points (added as hyperbolic-ish features) ────────────────────
    saddle_pts = [
        ( 0.5,  1.8, 0.5, 0.6),
        (-1.8,  1.2, 0.4, 0.5),
        ( 1.5, -0.8, 0.3, 0.7),
        (-0.8, -2.2, 0.45, 0.55),
    ]
    for cx, cy, amp, scale in saddle_pts:
        dx, dy = X - cx, Y - cy
        saddle = amp * (dx**2 - dy**2) * np.exp(-(dx**2 + dy**2) / (2 * scale**2))
        Z += saddle

    # ── High-frequency texture (rough local structure everywhere) ───────────
    # Simulate the noisy rugged terrain seen in real nets
    roughness_phases_x = rng.uniform(0, 2*np.pi, 18)
    roughness_phases_y = rng.uniform(0, 2*np.pi, 18)
    roughness_freqs = rng.uniform(1.5, 4.0, 18)
    roughness_amps = rng.uniform(0.04, 0.18, 18)
    for i in range(18):
        f = roughness_freqs[i]
        Z += roughness_amps[i] * np.sin(f * X + roughness_phases_x[i]) \
                                * np.cos(f * Y + roughness_phases_y[i])

    # ── Soft clip to avoid extreme spikes ───────────────────────────────────
    Z = np.tanh(Z / 3.0) * 3.0 + 2.5   # shift so min is ~0

    # Normalize to [0, 1] range for cleaner coloring
    Z_min, Z_max = Z.min(), Z.max()
    Z_norm = (Z - Z_min) / (Z_max - Z_min)

    return X, Y, Z_norm, x, y


# ─── Optimizer Trajectory Simulation ─────────────────────────────────────────
def loss_fn_scalar(px, py, X, Y, Z_norm):
    """Bilinear interpolation of loss value at (px, py)."""
    xi = np.interp(px, X[0, :], np.arange(X.shape[1]))
    yi = np.interp(py, Y[:, 0], np.arange(Y.shape[0]))
    xi_c = int(np.clip(xi, 0, X.shape[1]-2))
    yi_c = int(np.clip(yi, 0, Y.shape[0]-2))
    fx = xi - xi_c
    fy = yi - yi_c
    z00 = Z_norm[yi_c,   xi_c]
    z10 = Z_norm[yi_c,   xi_c+1]
    z01 = Z_norm[yi_c+1, xi_c]
    z11 = Z_norm[yi_c+1, xi_c+1]
    return (1-fx)*(1-fy)*z00 + fx*(1-fy)*z10 + (1-fx)*fy*z01 + fx*fy*z11


def numerical_grad(px, py, X, Y, Z_norm, h=0.02):
    gx = (loss_fn_scalar(px+h, py, X, Y, Z_norm)
        - loss_fn_scalar(px-h, py, X, Y, Z_norm)) / (2*h)
    gy = (loss_fn_scalar(px, py+h, X, Y, Z_norm)
        - loss_fn_scalar(px, py-h, X, Y, Z_norm)) / (2*h)
    return gx, gy


def simulate_trajectory(X, Y, Z_norm, start=(3.5, 3.0),
                         lr=0.06, momentum=0.85, noise=0.04,
                         steps=TRAJ_STEPS, seed=7):
    """SGD with momentum + Gaussian noise (mimics SGD on mini-batches)."""
    rng = np.random.default_rng(seed)
    px, py = start
    vx, vy = 0.0, 0.0
    path_x, path_y, path_z = [px], [py], [loss_fn_scalar(px, py, X, Y, Z_norm)]

    for _ in range(steps):
        gx, gy = numerical_grad(px, py, X, Y, Z_norm)
        # Gradient clipping
        gnorm = np.sqrt(gx**2 + gy**2) + 1e-8
        clip = 3.0
        if gnorm > clip:
            gx, gy = gx / gnorm * clip, gy / gnorm * clip
        # Momentum update
        vx = momentum * vx - lr * gx + rng.normal(0, noise)
        vy = momentum * vy - lr * gy + rng.normal(0, noise)
        px = np.clip(px + vx, X[0,0], X[0,-1])
        py = np.clip(py + vy, Y[0,0], Y[-1,0])
        path_x.append(px)
        path_y.append(py)
        path_z.append(loss_fn_scalar(px, py, X, Y, Z_norm) + 0.002)  # tiny lift

    return (np.array(path_x), np.array(path_y), np.array(path_z))


# ─── Render & Animation ───────────────────────────────────────────────────────
def make_frame(fig, ax, X, Y, Z_norm, traj, frame_idx, n_frames,
               surf_handle_container):
    """Render a single animation frame."""
    total_traj = len(traj[0])
    # How many trajectory points to reveal by this frame
    reveal = int((frame_idx / n_frames) * total_traj)
    reveal = max(2, reveal)

    # ── Camera rotation ─────────────────────────────────────────────────────
    # Full 360° rotation over the video
    azim_start = -50
    azim_end   = azim_start + 340
    azim = azim_start + (azim_end - azim_start) * (frame_idx / n_frames)
    elev = 28 + 6 * np.sin(frame_idx / n_frames * 2 * np.pi)  # gentle bob

    ax.view_init(elev=elev, azim=azim)

    # ── Trajectory line ─────────────────────────────────────────────────────
    # Remove old trajectory artists (all lines + scatter)
    while ax.lines:
        ax.lines[0].remove()
    while ax.collections and len(ax.collections) > 1:
        # Keep the first collection (the surface)
        ax.collections[-1].remove()

    tx = traj[0][:reveal]
    ty = traj[1][:reveal]
    tz = traj[2][:reveal] + 0.01

    if reveal > 2:
        ax.plot(tx, ty, tz, color='#FF4444', linewidth=1.5,
                alpha=0.85, zorder=5)

    # Current position marker
    ax.scatter([tx[-1]], [ty[-1]], [tz[-1]],
               color='#FFFF00', s=60, zorder=6, edgecolors='white',
               linewidths=0.8)

    # Trail: color-coded by loss height (red=high, green=low)
    if reveal > 10:
        seg_n = min(reveal, 80)  # last N points as colored trail
        tx_t = traj[0][reveal-seg_n:reveal]
        ty_t = traj[1][reveal-seg_n:reveal]
        tz_t = traj[2][reveal-seg_n:reveal] + 0.01
        colors = plt.cm.RdYlGn_r(np.linspace(0, 1, seg_n))
        for i in range(seg_n-1):
            ax.plot(tx_t[i:i+2], ty_t[i:i+2], tz_t[i:i+2],
                    color=colors[i], linewidth=2.2, alpha=0.7, zorder=5)


def create_static_surface(ax, X, Y, Z_norm):
    """Draw the surface once; return the handle."""
    ls = LightSource(azdeg=225, altdeg=45)
    cmap = plt.cm.plasma

    # Shade the surface with LightSource for 3D depth feel
    rgb = ls.shade(Z_norm, cmap=cmap, vert_exag=1.5, blend_mode='soft')

    surf = ax.plot_surface(
        X, Y, Z_norm,
        facecolors=rgb,
        rstride=2, cstride=2,
        linewidth=0,
        antialiased=True,
        alpha=0.92,
        shade=False,
    )
    return surf


def build_animation():
    print("Building landscape...")
    X, Y, Z_norm, x, y = build_landscape()

    print("Simulating optimizer trajectory...")
    traj = simulate_trajectory(X, Y, Z_norm,
                                start=(3.5, 3.2),
                                lr=0.055, momentum=0.87,
                                noise=0.035, steps=TRAJ_STEPS)

    print("Setting up figure...")
    fig = plt.figure(figsize=FIG_SIZE_IN, dpi=DPI)
    fig.patch.set_facecolor('#0d0d1a')

    ax = fig.add_subplot(111, projection='3d')
    ax.set_facecolor('#0d0d1a')

    # Draw static surface
    surf = create_static_surface(ax, X, Y, Z_norm)

    # ── Axes styling ────────────────────────────────────────────────────────
    ax.set_xlabel('θ₁', color='#aaaacc', fontsize=11, labelpad=8)
    ax.set_ylabel('θ₂', color='#aaaacc', fontsize=11, labelpad=8)
    ax.set_zlabel('Loss', color='#aaaacc', fontsize=11, labelpad=8)
    ax.tick_params(colors='#555577', labelsize=7)
    for pane in [ax.xaxis.pane, ax.yaxis.pane, ax.zaxis.pane]:
        pane.fill = False
        pane.set_edgecolor('#1a1a2e')
    ax.grid(True, color='#1a1a2e', linewidth=0.4, alpha=0.5)

    # ── Title & annotations ─────────────────────────────────────────────────
    fig.text(0.5, 0.97, 'Loss Landscape of a Neural Network',
             ha='center', va='top', color='white',
             fontsize=15, fontweight='bold', alpha=0.95)
    fig.text(0.5, 0.93,
             'SGD trajectory: escaping sharp minima → finding the flat basin',
             ha='center', va='top', color='#aaaacc', fontsize=9, alpha=0.85)
    fig.text(0.5, 0.02, '@fminxyz', ha='center', va='bottom',
             color='#6666aa', fontsize=10, alpha=0.7, style='italic')

    # Legend
    from matplotlib.lines import Line2D
    legend_elements = [
        Line2D([0], [0], color='#FF4444', linewidth=2, label='SGD path'),
        Line2D([0], [0], marker='o', color='w', markerfacecolor='#FFFF00',
               markersize=8, label='Current position', linewidth=0),
    ]
    legend = ax.legend(handles=legend_elements, loc='upper left',
                       framealpha=0.3, facecolor='#0d0d1a',
                       edgecolor='#444466', labelcolor='#ccccee',
                       fontsize=8)

    # Colorbar
    mappable = plt.cm.ScalarMappable(cmap='plasma')
    mappable.set_array(Z_norm)
    mappable.set_clim(0, 1)
    cbar = fig.colorbar(mappable, ax=ax, shrink=0.4, aspect=15, pad=0.05)
    cbar.set_label('Loss value', color='#aaaacc', fontsize=8)
    cbar.ax.yaxis.set_tick_params(color='#555577')
    plt.setp(plt.getp(cbar.ax.axes, 'yticklabels'), color='#aaaacc', fontsize=7)

    surf_container = [surf]

    def animate(frame):
        if frame % 30 == 0:
            print(f"  Frame {frame}/{N_FRAMES}")
        make_frame(fig, ax, X, Y, Z_norm, traj, frame, N_FRAMES, surf_container)
        return []

    print(f"Rendering {N_FRAMES} frames at {FPS}fps ({DURATION_SEC}s)...")
    anim = animation.FuncAnimation(
        fig, animate,
        frames=N_FRAMES,
        interval=1000 // FPS,
        blit=False,
    )

    Writer = animation.FFMpegWriter
    writer = Writer(fps=FPS, bitrate=4000,
                    extra_args=['-vcodec', 'libx264', '-pix_fmt', 'yuv420p',
                                '-preset', 'fast', '-crf', '18'])
    print(f"Writing video to {OUTPUT_VIDEO}...")
    anim.save(OUTPUT_VIDEO, writer=writer, dpi=DPI,
              savefig_kwargs={'facecolor': '#0d0d1a'})
    print("Video saved.")
    plt.close(fig)
    return X, Y, Z_norm, traj


def create_thumbnail(X, Y, Z_norm, traj):
    """High-quality static thumbnail (end of trajectory)."""
    print("Creating thumbnail...")
    fig = plt.figure(figsize=FIG_SIZE_IN, dpi=DPI)
    fig.patch.set_facecolor('#0d0d1a')

    ax = fig.add_subplot(111, projection='3d')
    ax.set_facecolor('#0d0d1a')

    create_static_surface(ax, X, Y, Z_norm)

    # Full trajectory
    ax.plot(traj[0], traj[1], traj[2]+0.01,
            color='#FF4444', linewidth=1.8, alpha=0.8, zorder=5)

    # Color-coded trail (last 150 steps)
    seg_n = 150
    tx_t = traj[0][-seg_n:]
    ty_t = traj[1][-seg_n:]
    tz_t = traj[2][-seg_n:] + 0.01
    colors = plt.cm.RdYlGn_r(np.linspace(0, 1, seg_n))
    for i in range(seg_n-1):
        ax.plot(tx_t[i:i+2], ty_t[i:i+2], tz_t[i:i+2],
                color=colors[i], linewidth=2.5, alpha=0.8, zorder=5)

    # Final position
    ax.scatter([traj[0][-1]], [traj[1][-1]], [traj[2][-1]+0.02],
               color='#FFFF00', s=120, zorder=6,
               edgecolors='white', linewidths=1.2)

    # Start marker
    ax.scatter([traj[0][0]], [traj[1][0]], [traj[2][0]+0.02],
               color='#FF8800', s=80, zorder=6, marker='^',
               edgecolors='white', linewidths=0.8)

    ax.view_init(elev=30, azim=-40)

    ax.set_xlabel('θ₁', color='#aaaacc', fontsize=11, labelpad=8)
    ax.set_ylabel('θ₂', color='#aaaacc', fontsize=11, labelpad=8)
    ax.set_zlabel('Loss', color='#aaaacc', fontsize=11, labelpad=8)
    ax.tick_params(colors='#555577', labelsize=7)
    for pane in [ax.xaxis.pane, ax.yaxis.pane, ax.zaxis.pane]:
        pane.fill = False
        pane.set_edgecolor('#1a1a2e')
    ax.grid(True, color='#1a1a2e', linewidth=0.4, alpha=0.5)

    fig.text(0.5, 0.97, 'Loss Landscape of a Neural Network',
             ha='center', va='top', color='white',
             fontsize=15, fontweight='bold')
    fig.text(0.5, 0.93,
             'SGD trajectory: sharp minima traps → flat generalizing basin',
             ha='center', va='top', color='#aaaacc', fontsize=9)
    fig.text(0.5, 0.02, '@fminxyz', ha='center', va='bottom',
             color='#6666aa', fontsize=10, style='italic')

    from matplotlib.lines import Line2D
    from matplotlib.patches import Patch
    legend_elements = [
        Line2D([0], [0], color='#FF4444', linewidth=2, label='SGD path'),
        Line2D([0], [0], marker='^', color='w', markerfacecolor='#FF8800',
               markersize=10, label='Start (random init)', linewidth=0),
        Line2D([0], [0], marker='o', color='w', markerfacecolor='#FFFF00',
               markersize=10, label='End (flat minimum)', linewidth=0),
    ]
    ax.legend(handles=legend_elements, loc='upper left',
              framealpha=0.3, facecolor='#0d0d1a',
              edgecolor='#444466', labelcolor='#ccccee', fontsize=9)

    mappable = plt.cm.ScalarMappable(cmap='plasma')
    mappable.set_array(Z_norm)
    mappable.set_clim(0, 1)
    cbar = fig.colorbar(mappable, ax=ax, shrink=0.4, aspect=15, pad=0.05)
    cbar.set_label('Loss value', color='#aaaacc', fontsize=8)
    cbar.ax.yaxis.set_tick_params(color='#555577')
    plt.setp(plt.getp(cbar.ax.axes, 'yticklabels'), color='#aaaacc', fontsize=7)

    plt.tight_layout()
    fig.savefig(OUTPUT_THUMB, dpi=DPI, facecolor='#0d0d1a',
                bbox_inches='tight', pad_inches=0.1)
    print(f"Thumbnail saved to {OUTPUT_THUMB}")
    plt.close(fig)


# ─── Main ─────────────────────────────────────────────────────────────────────
if __name__ == '__main__':
    import sys
    print("=" * 60)
    print("Loss Landscape Animator — @fminxyz")
    print("=" * 60)
    X, Y, Z_norm, traj = build_animation()
    create_thumbnail(X, Y, Z_norm, traj)
    print("=" * 60)
    print("Done!")
    print(f"  Video:     {OUTPUT_VIDEO}")
    print(f"  Thumbnail: {OUTPUT_THUMB}")
    print("=" * 60)