"""SAM: Sharpness-Aware Minimization Animation
Series 19, Post 2 — @fminxyz
"""
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation, FFMpegWriter
import warnings
warnings.filterwarnings('ignore')

fig, ax = plt.subplots(1, 1, figsize=(6, 6), facecolor='#0d1117')
ax.set_facecolor('#0d1117')
ax.set_xlim(0, 10)
ax.set_ylim(0, 10)
ax.axis('off')
fig.tight_layout(pad=0)

BLUE   = '#58a6ff'
GREEN  = '#3fb950'
ORANGE = '#f0883e'
RED    = '#f85149'
GRAY   = '#8b949e'
WHITE  = '#e6edf3'
PURPLE = '#bc8cff'
YELLOW = '#e3b341'

n_frames = 300
fps = 12

def clear_ax():
    ax.clear()
    ax.set_facecolor('#0d1117')
    ax.set_xlim(0, 10)
    ax.set_ylim(0, 10)
    ax.axis('off')

# Loss landscape curves
def sharp_minimum(x, center=3.0):
    return 3.0 + 8.0 * (x - center)**2

def flat_minimum(x, center=7.0):
    return 3.0 + 0.6 * (x - center)**2

def combined_loss(x):
    return np.minimum(sharp_minimum(x), flat_minimum(x) + 0.5 * np.sin(x * 4) * 0.3)

# Phase 1: Sharp vs flat minima problem (0-74)
def phase1_landscape(t):
    clear_ax()
    p = min(t / 65, 1.0)

    ax.text(5, 9.4, 'SAM: Острые vs Плоские Минимумы', ha='center',
            fontsize=11, color=WHITE, fontweight='bold', fontfamily='monospace')

    x = np.linspace(1.0, 9.0, 500)
    # Sharp minimum on left
    y_sharp = 1.5 + 6.0 * (x - 3.0)**2
    y_sharp_clipped = np.clip(y_sharp, 1.0, 8.5)

    # Flat minimum on right
    y_flat = 1.5 + 0.7 * (x - 7.0)**2
    y_flat_clipped = np.clip(y_flat, 1.0, 8.5)

    if p > 0.1:
        a1 = min((p - 0.1) / 0.35, 1.0)
        ax.plot(x[:250], y_sharp_clipped[:250], color=RED, linewidth=2.0, alpha=a1)
        ax.plot(3.0, 1.5, 'o', color=RED, markersize=7, alpha=a1)
        ax.text(3.0, 0.9, 'Острый\nminimum', ha='center', fontsize=8, color=RED, alpha=a1)

    if p > 0.4:
        a2 = min((p - 0.4) / 0.35, 1.0)
        ax.plot(x[250:], y_flat_clipped[250:], color=GREEN, linewidth=2.0, alpha=a2)
        ax.plot(7.0, 1.5, 'o', color=GREEN, markersize=7, alpha=a2)
        ax.text(7.0, 0.9, 'Плоский\nminimum', ha='center', fontsize=8, color=GREEN, alpha=a2)

    # Key question
    if p > 0.7:
        a3 = min((p - 0.7) / 0.3, 1.0)
        ax.text(5, 7.0, 'Оба минимума: train loss ≈ 0', ha='center',
                fontsize=9, color=GRAY, alpha=a3)
        ax.text(5, 6.3, 'Но острый → плохо generalize', ha='center',
                fontsize=9, color=RED, alpha=a3)
        ax.text(5, 5.6, 'Плоский → хорошо generalize', ha='center',
                fontsize=9, color=GREEN, alpha=a3)
        # PAC-Bayes intuition
        ax.text(5, 4.7, 'Интуиция: малый сдвиг θ → большой рост loss', ha='center',
                fontsize=8, color=GRAY, alpha=a3 * 0.8)
        ax.text(5, 4.1, 'в острой точке (чувствит. к dist. shift)', ha='center',
                fontsize=8, color=GRAY, alpha=a3 * 0.8)


# Phase 2: SAM objective (75-149)
def phase2_objective(t):
    clear_ax()
    p = min(t / 65, 1.0)

    ax.text(5, 9.4, 'SAM: Minimax Objective', ha='center',
            fontsize=11, color=WHITE, fontweight='bold', fontfamily='monospace')

    # Standard min
    if p > 0.0:
        a = min(p / 0.2, 1.0)
        ax.text(5, 8.4, 'Стандарт:', ha='center', fontsize=9, color=GRAY, alpha=a)
        ax.text(5, 7.8, 'min_θ  L(θ)', ha='center', fontsize=12, color=BLUE,
                fontweight='bold', alpha=a)

    # SAM min-max
    if p > 0.25:
        a = min((p - 0.25) / 0.25, 1.0)
        ax.text(5, 6.7, 'SAM:', ha='center', fontsize=9, color=ORANGE, alpha=a)
        ax.text(5, 6.0, 'min_θ  max_{||ε||≤ρ}  L(θ+ε)', ha='center',
                fontsize=12, color=ORANGE, fontweight='bold', alpha=a)
        ax.text(5, 5.3, '↑ найти наиболее острую окрестность',
                ha='center', fontsize=8, color=GRAY, alpha=a * 0.8)

    # Two-step explanation
    if p > 0.5:
        a = min((p - 0.5) / 0.3, 1.0)
        # Step 1
        rect1 = plt.Rectangle((0.5, 3.5), 4.2, 1.2, linewidth=1.2,
                               edgecolor=RED, facecolor='#2d1a0f', alpha=a)
        ax.add_patch(rect1)
        ax.text(2.6, 4.55, 'Шаг 1: worst case', ha='center', fontsize=8.5,
                color=RED, fontweight='bold', alpha=a)
        ax.text(2.6, 4.0, 'ε̂ = ρ·∇L(θ)/||∇L(θ)||', ha='center',
                fontsize=8.5, color=WHITE, alpha=a)

        # Step 2
        rect2 = plt.Rectangle((5.3, 3.5), 4.2, 1.2, linewidth=1.2,
                               edgecolor=GREEN, facecolor='#0f2d1a', alpha=a)
        ax.add_patch(rect2)
        ax.text(7.4, 4.55, 'Шаг 2: descent', ha='center', fontsize=8.5,
                color=GREEN, fontweight='bold', alpha=a)
        ax.text(7.4, 4.0, 'θ -= η·∇L(θ+ε̂)', ha='center',
                fontsize=8.5, color=WHITE, alpha=a)

    # Cost note
    if p > 0.8:
        a = min((p - 0.8) / 0.2, 1.0)
        ax.text(5, 2.8, 'Цена: 2 forward pass вместо 1', ha='center',
                fontsize=9, color=YELLOW, alpha=a)
        ax.text(5, 2.2, '≈ 2× медленнее → но generalization лучше', ha='center',
                fontsize=8.5, color=GRAY, alpha=a * 0.8)


# Phase 3: SAM step visualization (150-224)
def phase3_visualization(t):
    clear_ax()
    p = min(t / 65, 1.0)

    ax.text(5, 9.4, 'SAM шаг: визуализация', ha='center',
            fontsize=11, color=WHITE, fontweight='bold', fontfamily='monospace')

    x = np.linspace(1.5, 8.5, 400)
    y = 2.0 + 4.5 * (x - 5.0)**2 / 9 + 1.5 * np.sin((x - 3.0) * 1.5)**2
    y_norm = (y - y.min()) / (y.max() - y.min()) * 5.0 + 1.5
    ax.plot(x, y_norm, color=BLUE, linewidth=2.0, alpha=0.6)

    theta = 3.5
    idx_theta = np.argmin(np.abs(x - theta))
    y_theta = y_norm[idx_theta]

    if p > 0.1:
        a = min((p - 0.1) / 0.2, 1.0)
        ax.plot(theta, y_theta, 'o', color=WHITE, markersize=8, alpha=a)
        ax.text(theta, y_theta + 0.4, 'θ (текущий)', ha='center',
                fontsize=8.5, color=WHITE, alpha=a)

    # Perturbation
    eps = 0.9
    theta_eps = theta + eps
    idx_eps = np.argmin(np.abs(x - theta_eps))
    y_eps = y_norm[idx_eps]

    if p > 0.35:
        a = min((p - 0.35) / 0.25, 1.0)
        ax.annotate('', xy=(theta_eps, y_eps + 0.1), xytext=(theta + 0.1, y_theta),
                    arrowprops=dict(arrowstyle='->', color=RED, lw=1.5), alpha=a)
        ax.plot(theta_eps, y_eps, 's', color=RED, markersize=7, alpha=a)
        ax.text(theta_eps, y_eps - 0.6, 'θ+ε̂\n(worst case)', ha='center',
                fontsize=8, color=RED, alpha=a)
        ax.text(4.2, 4.5, 'ε̂ = ρ·g/||g||', fontsize=8, color=RED, alpha=a)

    # Gradient at perturbed point
    theta_new = theta - 0.7
    idx_new = np.argmin(np.abs(x - theta_new))
    y_new = y_norm[idx_new]

    if p > 0.6:
        a = min((p - 0.6) / 0.25, 1.0)
        ax.annotate('', xy=(theta_new, y_new + 0.1), xytext=(theta - 0.05, y_theta - 0.05),
                    arrowprops=dict(arrowstyle='->', color=GREEN, lw=2.0), alpha=a)
        ax.plot(theta_new, y_new, '*', color=GREEN, markersize=10, alpha=a)
        ax.text(theta_new, y_new - 0.6, 'θ_new\n(SAM шаг)', ha='center',
                fontsize=8, color=GREEN, alpha=a)
        ax.text(4.2, 3.8, 'g_SAM = ∇L(θ+ε̂)', fontsize=8, color=GREEN, alpha=a)

    # Note about flat region
    if p > 0.85:
        a = min((p - 0.85) / 0.15, 1.0)
        ax.text(5, 1.2, 'SAM шагает туда где подъём минимален → плоская зона',
                ha='center', fontsize=8, color=ORANGE, alpha=a)


# Phase 4: results table (225-299)
def phase4_results(t):
    clear_ax()
    p = min(t / 65, 1.0)

    ax.text(5, 9.4, 'SAM: результаты', ha='center',
            fontsize=11, color=WHITE, fontweight='bold', fontfamily='monospace')

    rows = [
        ('ResNet-50 CIFAR-10', 'SGD', '93.2%', 'Adam', '93.8%', 'SAM', '94.4%'),
        ('ViT ImageNet', 'Adam', '77.9%', '—', '—', 'SAM', '+0.5%'),
        ('BERT Fine-tune', 'AdamW', '85.1%', '—', '—', 'SAM', '86.3%'),
    ]

    if p > 0.15:
        a = min((p - 0.15) / 0.25, 1.0)
        ax.text(5, 8.4, 'Accuracy (test set)', ha='center', fontsize=9,
                color=YELLOW, fontweight='bold', alpha=a)

    for i, (task, opt1, acc1, opt2, acc2, opt3, acc3) in enumerate(rows):
        y_base = 7.3 - i * 1.4
        a_row = min(max((p - 0.3 - i * 0.15) / 0.25, 0.0), 1.0)
        ax.text(5, y_base + 0.5, task, ha='center', fontsize=8.5,
                color=GRAY, alpha=a_row)
        ax.text(2.2, y_base, f'{opt1}: {acc1}', ha='center', fontsize=8,
                color=BLUE, alpha=a_row)
        ax.text(5.0, y_base, f'{opt2}: {acc2}', ha='center', fontsize=8,
                color=ORANGE, alpha=a_row)
        ax.text(7.8, y_base, f'{opt3}: {acc3}', ha='center', fontsize=8,
                color=GREEN, fontweight='bold', alpha=a_row)

    if p > 0.7:
        a = min((p - 0.7) / 0.3, 1.0)
        ax.axhline(2.8, xmin=0.08, xmax=0.92, color=GRAY, alpha=a * 0.3, linewidth=0.8)
        ax.text(5, 2.4, 'Вывод: SAM стабильно даёт +0.4–1.2% top-1',
                ha='center', fontsize=9, color=WHITE, alpha=a)
        ax.text(5, 1.8, 'Цена: 2× медленнее. Рецепт: SAM + AdamW = best default',
                ha='center', fontsize=8, color=GRAY, alpha=a)

    if p > 0.9:
        a = min((p - 0.9) / 0.1, 1.0)
        rect = plt.Rectangle((0.5, 0.6), 9, 0.9, linewidth=1.5,
                              edgecolor=PURPLE, facecolor='#1a0d3c', alpha=a)
        ax.add_patch(rect)
        ax.text(5, 1.1, 'Завтра: Muon — геометрия пространства весов (2024)',
                ha='center', fontsize=9, color=PURPLE, fontweight='bold', alpha=a)


def animate(frame):
    if frame < 75:
        phase1_landscape(frame)
    elif frame < 150:
        phase2_objective(frame - 75)
    elif frame < 225:
        phase3_visualization(frame - 150)
    else:
        phase4_results(frame - 225)


anim = FuncAnimation(fig, animate, frames=n_frames, interval=1000/fps)
writer = FFMpegWriter(fps=fps, bitrate=1800,
                      extra_args=['-vf', 'scale=1080:1080',
                                  '-c:v', 'libx264', '-pix_fmt', 'yuv420p'])
out_path = '/root/Strategy/content/drafts/sam_animation.mp4'
anim.save(out_path, writer=writer, dpi=180)
print(f'Saved: {out_path}')