"""
@fminxyz Series 5 Post 5: Universality + будущее MI
Феанор, 2026-02-26 10:00 MSK

Концепция: одни и те же circuits в разных моделях
- Фаза 1: вопрос — случайность или закономерность?
- Фаза 2: side-by-side Gabor-like features в 4 моделях
- Фаза 3: timeline прогресса MI 2020-2025
- Фаза 4: что дальше — "физические законы нейросетей"

25 секунд, 1080x1080, 2fps
"""

import numpy as np
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import matplotlib.patches as patches
import matplotlib.gridspec as gridspec
from matplotlib.patches import FancyArrowPatch, Rectangle, FancyBboxPatch
from matplotlib.lines import Line2D
import matplotlib.patheffects as pe
from mpl_toolkits.axes_grid1 import make_axes_locatable
import os

# --- CONFIG ---
FPS = 2
DURATION = 25
N_FRAMES = FPS * DURATION
W, H = 1080, 1080
DPI = 72
FIG_SIZE = W / DPI

# Colors
BG = '#0d1117'
FG = '#e6edf3'
ACCENT1 = '#58a6ff'   # blue (ResNet)
ACCENT2 = '#3fb950'   # green (ViT)
ACCENT3 = '#ff7b72'   # red (ConvNet)
ACCENT4 = '#d2a8ff'   # purple (CLIP)
GOLD = '#f0b429'
GRAY = '#8b949e'
GRID_C = '#21262d'

# Phases
PHASE_FRAMES = [0, 12, 28, 36, 42]  # cumulative
# Phase 0 (0-11):   question + setup
# Phase 1 (12-27):  side-by-side model features
# Phase 2 (28-35):  timeline 2020-2025
# Phase 3 (36-41):  conclusion / "physical laws"
# Phase 4 (42-49):  roundup teaser


def gabor_feature(size=32, angle=0, freq=4.0, sigma=5.0):
    """Generate Gabor-like feature pattern."""
    x = np.linspace(-size//2, size//2, size)
    y = np.linspace(-size//2, size//2, size)
    X, Y = np.meshgrid(x, y)
    Xr = X * np.cos(angle) + Y * np.sin(angle)
    Yr = -X * np.sin(angle) + Y * np.cos(angle)
    gauss = np.exp(-(Xr**2 + Yr**2) / (2 * sigma**2))
    sinusoid = np.cos(2 * np.pi * freq / size * Xr)
    return gauss * sinusoid


def curve_detector(size=32, angle=0, curvature=0.3):
    """Simplified curve detector feature."""
    x = np.linspace(-1, 1, size)
    y = np.linspace(-1, 1, size)
    X, Y = np.meshgrid(x, y)
    # Rotate
    Xr = X * np.cos(angle) + Y * np.sin(angle)
    Yr = -X * np.sin(angle) + Y * np.cos(angle)
    # Curved line
    curve = Yr - curvature * Xr**2
    feature = np.exp(-curve**2 * 20) * np.exp(-(Xr)**2 * 2)
    return feature


def draw_phase0(ax, alpha):
    """Question: is it random or pattern?"""
    ax.set_facecolor(BG)
    ax.set_xlim(0, 10)
    ax.set_ylim(0, 10)
    ax.axis('off')

    # Title
    title = ax.text(5, 8.5, "Случайность\nили закономерность?",
                    ha='center', va='center', fontsize=28, fontweight='bold',
                    color=FG, alpha=alpha)
    title.set_path_effects([pe.withStroke(linewidth=3, foreground=BG)])

    # Two models showing similar features
    # ResNet feature (left)
    ax_inset1 = ax.inset_axes([0.1, 0.3, 0.35, 0.35])
    feature1 = gabor_feature(32, angle=np.pi/4, freq=4)
    ax_inset1.imshow(feature1, cmap='RdBu', vmin=-1, vmax=1)
    ax_inset1.axis('off')
    ax.text(2.85, 3.2, "ResNet", ha='center', fontsize=11, color=ACCENT1, alpha=alpha)

    # ViT feature (right) — almost same pattern
    ax_inset2 = ax.inset_axes([0.55, 0.3, 0.35, 0.35])
    feature2 = gabor_feature(32, angle=np.pi/4 + 0.05, freq=4)  # nearly identical
    ax_inset2.imshow(feature2, cmap='RdBu', vmin=-1, vmax=1)
    ax_inset2.axis('off')
    ax.text(7.3, 3.2, "ViT", ha='center', fontsize=11, color=ACCENT2, alpha=alpha)

    # Question mark between them
    ax.text(5, 5.2, "≈ ?", ha='center', va='center', fontsize=40,
            color=GOLD, alpha=alpha, fontweight='bold')

    ax.text(5, 1.8, "Одинаковые фичи в разных архитектурах?",
            ha='center', fontsize=13, color=GRAY, alpha=alpha*0.8)


def draw_phase1(ax, t):
    """Side-by-side model features: same Gabor patterns."""
    ax.set_facecolor(BG)
    ax.set_xlim(0, 10)
    ax.set_ylim(0, 10)
    ax.axis('off')

    t_norm = min(t, 1.0)

    ax.text(5, 9.3, "Те же самые circuits в разных моделях",
            ha='center', fontsize=18, fontweight='bold', color=FG, alpha=t_norm)

    # 4 models, 4 similar features
    models = [
        ("ResNet", ACCENT1, 0, np.pi/6),
        ("ViT", ACCENT2, np.pi/5, np.pi/6 + 0.1),
        ("ConvNet", ACCENT3, np.pi/4, np.pi/6 + 0.05),
        ("CLIP", ACCENT4, np.pi/3, np.pi/6 - 0.08),
    ]

    positions = [(0.03, 0.52), (0.27, 0.52), (0.52, 0.52), (0.76, 0.52)]
    feature_show = [
        gabor_feature(32, angle=np.pi/6, freq=4),
        gabor_feature(32, angle=np.pi/6 + 0.1, freq=4),
        gabor_feature(32, angle=np.pi/6 + 0.05, freq=3.9),
        gabor_feature(32, angle=np.pi/6 - 0.08, freq=4.1),
    ]

    reveal_times = [0.0, 0.25, 0.5, 0.75]

    for i, ((name, color, _, _), pos, feat, rt) in enumerate(
            zip(models, positions, feature_show, reveal_times)):
        a = min(max((t_norm - rt) / 0.25, 0), 1)
        if a > 0:
            inset = ax.inset_axes([pos[0], pos[1], 0.20, 0.30])
            inset.imshow(feat, cmap='RdBu', vmin=-1, vmax=1, alpha=a)
            inset.axis('off')
            ax.text(pos[0]*10 + 1.0, pos[1]*10 - 0.3, name,
                    ha='center', fontsize=12, color=color, alpha=a, fontweight='bold')

    # Arrow indicating similarity
    if t_norm > 0.8:
        a2 = (t_norm - 0.8) / 0.2
        ax.text(5, 4.8, "Gabor-подобные фичи",
                ha='center', fontsize=14, color=GOLD, alpha=a2)
        ax.text(5, 4.2, "в 8 из 9 исследованных архитектур",
                ha='center', fontsize=12, color=GRAY, alpha=a2 * 0.8)

    # Curve detectors row
    if t_norm > 0.4:
        a3 = min((t_norm - 0.4) / 0.3, 1)
        ax.text(5, 2.7, "Детекторы кривых (45°, 90°, 135°...)",
                ha='center', fontsize=12, color=FG, alpha=a3)

        angles = [0, np.pi/4, np.pi/2, 3*np.pi/4]
        for j, ang in enumerate(angles):
            xpos = 0.1 + j * 0.22
            inset2 = ax.inset_axes([xpos, 0.05, 0.18, 0.22])
            feat2 = curve_detector(32, angle=ang)
            inset2.imshow(feat2, cmap='viridis', alpha=a3)
            inset2.axis('off')


def draw_phase2(ax, t):
    """Timeline of MI progress 2020-2025."""
    ax.set_facecolor(BG)
    ax.set_xlim(0, 10)
    ax.set_ylim(0, 10)
    ax.axis('off')

    t_norm = min(t, 1.0)

    ax.text(5, 9.3, "Прогресс Mechanistic Interpretability",
            ha='center', fontsize=17, fontweight='bold', color=FG, alpha=t_norm)

    # Timeline entries
    timeline = [
        (2020, "Circuits hypothesis", "Olah et al.", ACCENT1, "Детекторы кривых в InceptionV1"),
        (2021, "Universality", "Elhage et al.", ACCENT2, "Те же circuits в 9 разных сетях"),
        (2022, "Superposition", "Elhage et al.", ACCENT3, "100 нейронов → 300 концептов"),
        (2023, "Induction heads", "Olsson et al.", ACCENT4, "In-context learning = 2 heads"),
        (2024, "Scaling SAE", "Anthropic", GOLD, "34M интерпретируемых фич в Claude 3"),
        (2025, "→ Safety apps", "Anthropic", GRAY, "Steering vectors, backdoor detection"),
    ]

    t_step = 1.0 / len(timeline)
    for i, (year, name, author, color, desc) in enumerate(timeline):
        reveal = min(max((t_norm - i * t_step) / t_step, 0), 1)
        if reveal > 0:
            y = 8.0 - i * 1.25
            # Year
            ax.text(0.5, y, str(year), ha='left', fontsize=14, color=color,
                    alpha=reveal, fontweight='bold')
            # Name
            ax.text(1.6, y, name, ha='left', fontsize=12, color=FG, alpha=reveal,
                    fontweight='bold')
            # Author
            ax.text(1.6, y - 0.4, f"  {author}", ha='left', fontsize=10,
                    color=GRAY, alpha=reveal * 0.8)
            # Line
            ax.plot([1.3, 1.3], [y - 0.6, y + 0.4], color=color, alpha=reveal * 0.5,
                    linewidth=2)
            # Dot
            ax.plot(1.3, y, 'o', color=color, alpha=reveal, markersize=8)


def draw_phase3(ax, t):
    """Conclusion: physical laws of neural networks."""
    ax.set_facecolor(BG)
    ax.set_xlim(0, 10)
    ax.set_ylim(0, 10)
    ax.axis('off')

    t_norm = min(t, 1.0)

    # Central quote
    ax.text(5, 7.5, '"Возможно, нейросети открывают',
            ha='center', fontsize=17, color=FG, alpha=t_norm, style='italic')
    ax.text(5, 6.8, 'одни и те же алгоритмы снова и снова —',
            ha='center', fontsize=17, color=FG, alpha=t_norm, style='italic')
    ax.text(5, 6.1, 'как физические законы"',
            ha='center', fontsize=17, color=GOLD, alpha=t_norm, style='italic',
            fontweight='bold')
    ax.text(5, 5.4, '— Chris Olah, 2020',
            ha='center', fontsize=12, color=GRAY, alpha=t_norm * 0.8)

    if t_norm > 0.4:
        a = min((t_norm - 0.4) / 0.3, 1)
        # Three implications
        items = [
            (ACCENT1, "🔬 Safety", "Найти опасные circuits → исправить"),
            (ACCENT2, "⚡ Efficiency", "Убрать лишние нейроны → speedup"),
            (ACCENT3, "🔮 Capability", "Понять → улучшить → управлять"),
        ]
        for j, (color, title, desc) in enumerate(items):
            y = 4.0 - j * 1.2
            ax.text(1.0, y, title, ha='left', fontsize=13, color=color, alpha=a,
                    fontweight='bold')
            ax.text(1.0, y - 0.45, desc, ha='left', fontsize=11, color=GRAY, alpha=a * 0.9)


def draw_frame(frame_idx):
    fig, ax = plt.subplots(1, 1, figsize=(FIG_SIZE, FIG_SIZE), facecolor=BG)
    ax.set_facecolor(BG)

    # Phase logic
    total_frames = N_FRAMES
    p0_end = 10
    p1_end = 28
    p2_end = 38
    p3_end = total_frames

    if frame_idx < p0_end:
        t = frame_idx / p0_end
        draw_phase0(ax, alpha=min(t * 3, 1.0))

    elif frame_idx < p1_end:
        t = (frame_idx - p0_end) / (p1_end - p0_end)
        draw_phase1(ax, t)

    elif frame_idx < p2_end:
        t = (frame_idx - p1_end) / (p2_end - p1_end)
        draw_phase2(ax, t)

    else:
        t = (frame_idx - p2_end) / (p3_end - p2_end)
        draw_phase3(ax, t)

    # Watermark
    ax.text(9.8, 0.2, "@fminxyz", ha='right', fontsize=10, color=GRAY,
            alpha=0.5, transform=ax.transData)

    plt.tight_layout(pad=0)
    return fig


def main():
    out_dir = "/root/Strategy/content/drafts"
    frames_dir = os.path.join(out_dir, "_univ_frames")
    os.makedirs(frames_dir, exist_ok=True)

    print(f"Rendering {N_FRAMES} frames...")
    for i in range(N_FRAMES):
        fig = draw_frame(i)
        fig.savefig(os.path.join(frames_dir, f"frame_{i:04d}.png"),
                    dpi=DPI, facecolor=BG, bbox_inches='tight',
                    pad_inches=0)
        plt.close(fig)
        if i % 5 == 0:
            print(f"  frame {i}/{N_FRAMES}")

    print("Encoding video...")
    mp4_path = os.path.join(out_dir, "universality_animation.mp4")
    os.system(
        f"ffmpeg -y -r {FPS} -i {frames_dir}/frame_%04d.png "
        f"-vf 'scale=1080:1080' "
        f"-c:v libx264 -pix_fmt yuv420p -preset fast -crf 23 "
        f"{mp4_path} 2>/dev/null"
    )

    # Cleanup frames
    import shutil
    shutil.rmtree(frames_dir)
    print(f"Done: {mp4_path}")


if __name__ == "__main__":
    main()