"""
AlphaFold Animation для @fminxyz — Series 8, Пост 1
Тема: AlphaFold решил 50-летнюю задачу биологии
4 фазы:
  1. Проблема: последовательность → структура (25× → 1 конфигурация)
  2. Evoformer: попарные представления + MSA
  3. Structure Module: итеративное уточнение 3D
  4. Результат: pLDDT confidence + 200M структур

Формат: 1080×1080 MP4, 25 секунд, без звука
"""

import numpy as np
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import matplotlib.patches as mpatches
from matplotlib.patches import FancyArrowPatch, FancyBboxPatch
import matplotlib.patheffects as pe
from matplotlib.animation import FuncAnimation, FFMpegWriter

# ── Цветовая схема ──────────────────────────────────────────────
BG      = '#0D1117'
BLUE    = '#58A6FF'
GREEN   = '#3FB950'
ORANGE  = '#F78166'
PURPLE  = '#BC8CFF'
YELLOW  = '#E3B341'
WHITE   = '#E6EDF3'
GRAY    = '#8B949E'
CYAN    = '#39D3F5'
RED     = '#FF7B72'

fig = plt.figure(figsize=(10.8, 10.8), facecolor=BG)
ax  = fig.add_axes([0, 0, 1, 1])
ax.set_xlim(0, 10.8)
ax.set_ylim(0, 10.8)
ax.set_aspect('equal')
ax.axis('off')
ax.set_facecolor(BG)

# ── Параметры анимации ───────────────────────────────────────────
FPS      = 24
DURATION = 25
N_FRAMES = FPS * DURATION

# фазы (секунды):  0-6  |  6-12  |  12-19  |  19-25
PHASE_FRAMES = [
    (0,           6 * FPS),   # Фаза 1: проблема сворачивания
    (6 * FPS,    12 * FPS),   # Фаза 2: Evoformer
    (12 * FPS,   19 * FPS),   # Фаза 3: Structure module
    (19 * FPS,   N_FRAMES),   # Фаза 4: результат
]

def t_local(frame, phase_idx):
    """Локальное время фазы [0, 1]"""
    s, e = PHASE_FRAMES[phase_idx]
    return np.clip((frame - s) / max(e - s - 1, 1), 0, 1)

def easing(t, k=4):
    """Smooth easing: медленнее в начале и конце"""
    return t**k / (t**k + (1 - t)**k)

def clear():
    ax.cla()
    ax.set_xlim(0, 10.8)
    ax.set_ylim(0, 10.8)
    ax.set_aspect('equal')
    ax.axis('off')
    ax.set_facecolor(BG)

def draw_label(x, y, text, color=WHITE, size=22, weight='bold', alpha=1.0, ha='center'):
    ax.text(x, y, text, color=color, fontsize=size, fontweight=weight,
            ha=ha, va='center', alpha=alpha,
            path_effects=[pe.withStroke(linewidth=3, foreground=BG)])

def draw_box(x, y, w, h, color, alpha=0.85, lw=2.5, corner=0.3):
    fancy = FancyBboxPatch((x - w/2, y - h/2), w, h,
                           boxstyle=f"round,pad={corner}",
                           facecolor=color, edgecolor=color,
                           alpha=alpha, linewidth=lw)
    ax.add_patch(fancy)

def draw_arrow(x1, y1, x2, y2, color=WHITE, lw=2.5, alpha=0.85):
    ax.annotate('', xy=(x2, y2), xytext=(x1, y1),
                arrowprops=dict(arrowstyle='->', color=color,
                                lw=lw, mutation_scale=20),
                alpha=alpha)

# ═══════════════════════════════════════════════════════════════════
# ФАЗА 1: Проблема сворачивания (0–6 сек)
# ═══════════════════════════════════════════════════════════════════

def draw_phase1(frame):
    clear()
    t = easing(t_local(frame, 0))

    # Заголовок
    draw_label(5.4, 10.2, 'Задача сворачивания белка', color=YELLOW, size=28)

    # Аминокислотная цепочка (слева)
    aa_colors = [BLUE, GREEN, ORANGE, PURPLE, CYAN, RED, BLUE, GREEN, ORANGE, PURPLE]
    aa_labels = ['M','A','G','L','K','D','V','T','P','E']
    cx = 1.2
    for i, (col, lbl) in enumerate(zip(aa_colors, aa_labels)):
        y_pos = 9.0 - i * 0.72
        alpha_aa = min(1.0, t * len(aa_labels) / (i + 0.5))
        circle = plt.Circle((cx, y_pos), 0.28, color=col, alpha=alpha_aa * 0.85, zorder=3)
        ax.add_patch(circle)
        ax.text(cx, y_pos, lbl, color=WHITE, fontsize=13, ha='center', va='center',
                fontweight='bold', zorder=4, alpha=alpha_aa)
        if i < len(aa_labels) - 1:
            ax.plot([cx, cx], [y_pos - 0.28, y_pos - 0.44], color=GRAY, lw=2, alpha=alpha_aa * 0.7)

    ax.text(cx, 2.4, 'Последовательность\n(N=10)', color=GRAY, fontsize=13,
            ha='center', va='center', alpha=t)

    # Стрелка вопроса
    q_alpha = min(1.0, max(0, (t - 0.3) * 5))
    ax.text(3.5, 5.5, '?', color=YELLOW, fontsize=72, ha='center', va='center',
            fontweight='bold', alpha=q_alpha,
            path_effects=[pe.withStroke(linewidth=5, foreground=BG)])
    draw_arrow(2.0, 5.5, 4.5, 5.5, color=YELLOW, lw=3, alpha=q_alpha)

    # 3D структура (спираль) справа
    if t > 0.4:
        s_alpha = easing(min(1.0, (t - 0.4) * 2.5))
        theta = np.linspace(0, 4 * np.pi, 80)
        x_h = 7.0 + 0.8 * np.cos(theta)
        y_h = 5.5 + 2.5 * (theta / (4 * np.pi) - 0.5) + 0.3 * np.sin(theta * 2)
        ax.plot(x_h, y_h, color=BLUE, lw=3.5, alpha=s_alpha * 0.9)
        # добавить случайные боковые цепи
        np.random.seed(42)
        for i in range(0, 80, 10):
            dx = np.random.uniform(-0.6, 0.6)
            dy = np.random.uniform(-0.4, 0.4)
            col = aa_colors[i % len(aa_colors)]
            ax.plot([x_h[i], x_h[i] + dx], [y_h[i], y_h[i] + dy],
                    color=col, lw=2, alpha=s_alpha * 0.7)
            circle = plt.Circle((x_h[i] + dx, y_h[i] + dy), 0.12,
                                 color=col, alpha=s_alpha * 0.8, zorder=3)
            ax.add_patch(circle)

        ax.text(7.0, 2.4, '3D структура\n(1 из 3¹⁰⁰ вариантов)', color=GRAY, fontsize=13,
                ha='center', va='center', alpha=s_alpha)

    # Снизу: факт
    fact_alpha = easing(min(1.0, max(0, (t - 0.7) * 5)))
    draw_box(5.4, 1.3, 8.5, 1.1, '#1C2128', alpha=fact_alpha * 0.9)
    ax.text(5.4, 1.5, '50 лет биологи не могли решить эту задачу надёжно',
            color=WHITE, fontsize=16, ha='center', va='center', alpha=fact_alpha, fontweight='bold')
    ax.text(5.4, 1.0, '1972 — Christian Anfinsen: "структура закодирована в последовательности"',
            color=GRAY, fontsize=13, ha='center', va='center', alpha=fact_alpha * 0.8)


# ═══════════════════════════════════════════════════════════════════
# ФАЗА 2: Evoformer (6–12 сек)
# ═══════════════════════════════════════════════════════════════════

def draw_phase2(frame):
    clear()
    t = easing(t_local(frame, 1))

    draw_label(5.4, 10.2, 'Evoformer: сердце AlphaFold2', color=CYAN, size=28)

    # MSA (Multiple Sequence Alignment) блок — слева
    msa_alpha = min(1.0, t * 4)
    draw_box(2.2, 7.5, 3.5, 2.0, '#1A2E45', alpha=msa_alpha * 0.9)
    ax.text(2.2, 8.6, 'MSA матрица', color=CYAN, fontsize=16,
            ha='center', va='center', fontweight='bold', alpha=msa_alpha)

    # Сетка MSA
    if t > 0.15:
        seqs = ['MAGLKDVTPE', 'MAGMKDVTPE', 'MAGLKDVQPE', 'MGGLKDVTPE', 'MAGLKDVTPA']
        colors_map = {
            'M': BLUE, 'A': GREEN, 'G': ORANGE, 'L': PURPLE,
            'K': CYAN, 'D': RED, 'V': YELLOW, 'T': GREEN,
            'P': BLUE, 'E': ORANGE, 'Q': PURPLE
        }
        for j, seq in enumerate(seqs):
            for k, aa in enumerate(seq[:8]):
                cx_g = 0.65 + k * 0.35
                cy_g = 8.2 - j * 0.28
                col = colors_map.get(aa, WHITE)
                cell_alpha = min(1.0, (t - 0.1) * 6) * 0.8
                rect = plt.Rectangle((cx_g - 0.14, cy_g - 0.11), 0.28, 0.22,
                                     color=col, alpha=cell_alpha, zorder=3)
                ax.add_patch(rect)

    # Попарная матрица — справа
    pair_alpha = easing(min(1.0, max(0, (t - 0.2) * 3)))
    draw_box(8.2, 7.5, 3.5, 2.0, '#2D1A45', alpha=pair_alpha * 0.9)
    ax.text(8.2, 8.6, 'Pair representation', color=PURPLE, fontsize=16,
            ha='center', va='center', fontweight='bold', alpha=pair_alpha)

    n_res = 8
    if t > 0.3:
        for i_r in range(n_res):
            for j_r in range(n_res):
                val = np.exp(-abs(i_r - j_r) * 0.5) + 0.1 * np.random.RandomState(i_r * 10 + j_r).rand()
                val = min(val, 1.0)
                cx_p = 6.65 + j_r * 0.28
                cy_p = 8.2 - i_r * 0.28
                col = plt.cm.plasma(val)
                cell_alpha = pair_alpha * 0.85
                rect = plt.Rectangle((cx_p - 0.12, cy_p - 0.11), 0.24, 0.22,
                                     color=col, alpha=cell_alpha, zorder=3)
                ax.add_patch(rect)

    # Стрелки обновления
    update_alpha = easing(min(1.0, max(0, (t - 0.45) * 4)))
    if update_alpha > 0.01:
        draw_arrow(3.9, 7.5, 6.5, 7.5, color=WHITE, lw=2.5, alpha=update_alpha)
        draw_arrow(6.5, 7.2, 3.9, 7.2, color=GRAY, lw=2, alpha=update_alpha * 0.7)
        ax.text(5.2, 7.85, 'Row attention', color=WHITE, fontsize=12,
                ha='center', va='center', alpha=update_alpha)
        ax.text(5.2, 7.05, 'Col attention', color=GRAY, fontsize=12,
                ha='center', va='center', alpha=update_alpha * 0.7)

    # Evoformer блоки (48 повторений)
    block_alpha = easing(min(1.0, max(0, (t - 0.55) * 4)))
    draw_label(5.4, 6.3, '× 48 блоков Evoformer', color=YELLOW, size=18, alpha=block_alpha)

    # Итог: дистанционная матрица + торсионные углы
    out_alpha = easing(min(1.0, max(0, (t - 0.7) * 5)))
    draw_box(2.8, 4.5, 4.5, 1.5, '#1A3020', alpha=out_alpha * 0.9)
    ax.text(2.8, 4.9, '→ Матрица расстояний', color=GREEN, fontsize=15,
            ha='center', va='center', fontweight='bold', alpha=out_alpha)
    ax.text(2.8, 4.3, 'P(dᵢⱼ < threshold)', color=GRAY, fontsize=12,
            ha='center', va='center', alpha=out_alpha * 0.8)

    draw_box(8.0, 4.5, 4.5, 1.5, '#3A1A20', alpha=out_alpha * 0.9)
    ax.text(8.0, 4.9, '→ Торсионные углы', color=ORANGE, fontsize=15,
            ha='center', va='center', fontweight='bold', alpha=out_alpha)
    ax.text(8.0, 4.3, 'φ, ψ ∈ Ramachandran', color=GRAY, fontsize=12,
            ha='center', va='center', alpha=out_alpha * 0.8)

    # Ключевой инсайт
    insight_alpha = easing(min(1.0, max(0, (t - 0.82) * 8)))
    draw_box(5.4, 2.8, 9.5, 1.4, '#1C2128', alpha=insight_alpha * 0.9)
    ax.text(5.4, 3.1, 'Ключ: MSA → эволюционные ограничения',
            color=WHITE, fontsize=16, ha='center', va='center',
            fontweight='bold', alpha=insight_alpha)
    ax.text(5.4, 2.55, 'Если два остатка мутируют вместе → они пространственно близки',
            color=GRAY, fontsize=13, ha='center', va='center', alpha=insight_alpha * 0.85)

    ax.text(5.4, 1.9, 'Нобелевская премия 2024 — Demis Hassabis, John Jumper',
            color=YELLOW, fontsize=14, ha='center', va='center',
            fontweight='bold', alpha=insight_alpha,
            path_effects=[pe.withStroke(linewidth=3, foreground=BG)])


# ═══════════════════════════════════════════════════════════════════
# ФАЗА 3: Structure Module — итерационное уточнение (12–19 сек)
# ═══════════════════════════════════════════════════════════════════

def draw_phase3(frame):
    clear()
    t = easing(t_local(frame, 2))

    draw_label(5.4, 10.2, 'Structure Module: 3D из математики', color=GREEN, size=28)

    # Итерация уточнения
    n_steps = 8
    step_now = int(t * n_steps * 0.9)
    step_alpha = min(1.0, t * 5)

    ax.text(5.4, 9.3, f'Итерация: {step_now + 1} / {n_steps}  (Invariant Point Attention)',
            color=GRAY, fontsize=15, ha='center', va='center', alpha=step_alpha)

    # Белковая спираль с нарастающим качеством
    np.random.seed(7)
    n_atoms = 40
    theta_base = np.linspace(0, 6 * np.pi, n_atoms)

    # Начальная шумная конфигурация
    noise_level = 1.0 - t * 0.85
    theta = theta_base + noise_level * np.random.randn(n_atoms) * 0.4
    x_prot = 5.4 + 1.5 * np.cos(theta) + noise_level * np.random.randn(n_atoms) * 0.5
    y_prot = 5.4 + 2.8 * (theta_base / (6 * np.pi) - 0.5) + noise_level * np.random.randn(n_atoms) * 0.5

    # Красим по pLDDT (confidence)
    plddt_vals = 30 + 70 * (1 - noise_level) + np.random.RandomState(42).rand(n_atoms) * 10
    plddt_vals = np.clip(plddt_vals, 0, 100)

    colors_prot = [plt.cm.RdYlBu(v / 100) for v in plddt_vals]

    # Рисуем backbone
    for i in range(n_atoms - 1):
        ax.plot([x_prot[i], x_prot[i+1]], [y_prot[i], y_prot[i+1]],
                color=WHITE, lw=2.5, alpha=0.6, zorder=2)

    # Рисуем атомы (Cα)
    for i, (xp, yp, col) in enumerate(zip(x_prot, y_prot, colors_prot)):
        circle = plt.Circle((xp, yp), 0.18, color=col, alpha=0.9, zorder=3)
        ax.add_patch(circle)

    # Легенда pLDDT
    legend_alpha = min(1.0, t * 3)
    ax.text(9.5, 8.5, 'pLDDT\n(confidence)', color=WHITE, fontsize=14,
            ha='center', va='center', fontweight='bold', alpha=legend_alpha)
    for i, (label, col_v) in enumerate([('90–100', 0.95), ('70–90', 0.7),
                                          ('50–70', 0.45), ('<50', 0.15)]):
        col_b = plt.cm.RdYlBu(col_v)
        circle = plt.Circle((8.7, 7.8 - i * 0.55), 0.18, color=col_b, alpha=0.9, zorder=3)
        ax.add_patch(circle)
        ax.text(9.1, 7.8 - i * 0.55, label, color=WHITE, fontsize=12,
                ha='left', va='center', alpha=legend_alpha)

    # Выравнивание рамок (SE3-equivariant)
    frame_alpha = easing(min(1.0, max(0, (t - 0.6) * 4)))
    if frame_alpha > 0.05:
        # Рисуем несколько систем координат на атомах
        for i in [5, 15, 25, 35]:
            if i < n_atoms:
                scale = 0.4
                angle = theta[i] * 0.3
                ax.annotate('', xy=(x_prot[i] + scale * np.cos(angle),
                                    y_prot[i] + scale * np.sin(angle)),
                            xytext=(x_prot[i], y_prot[i]),
                            arrowprops=dict(arrowstyle='->', color=RED,
                                           lw=1.5, mutation_scale=12),
                            alpha=frame_alpha)
                ax.annotate('', xy=(x_prot[i] - scale * np.sin(angle),
                                    y_prot[i] + scale * np.cos(angle)),
                            xytext=(x_prot[i], y_prot[i]),
                            arrowprops=dict(arrowstyle='->', color=GREEN,
                                           lw=1.5, mutation_scale=12),
                            alpha=frame_alpha)

    ax.text(1.8, 8.2, 'SE(3)-\nequivariant', color=RED, fontsize=12,
            ha='center', va='center', alpha=frame_alpha)

    # Метрика качества
    quality = t * 0.88
    qual_alpha = min(1.0, t * 2)
    draw_box(5.4, 1.8, 8.5, 1.4, '#0D1117', alpha=qual_alpha * 0.95, lw=0)
    bar_w = quality * 8.0
    bar = plt.Rectangle((1.4, 1.3), bar_w, 0.6,
                         color=GREEN if quality > 0.7 else YELLOW if quality > 0.4 else RED,
                         alpha=qual_alpha * 0.9, zorder=3)
    ax.add_patch(bar)
    ax.text(5.4, 2.2, f'TM-score: {quality:.2f}  (>0.5 = правильная укладка)',
            color=WHITE, fontsize=14, ha='center', va='center',
            fontweight='bold', alpha=qual_alpha)


# ═══════════════════════════════════════════════════════════════════
# ФАЗА 4: Результат и масштаб (19–25 сек)
# ═══════════════════════════════════════════════════════════════════

def draw_phase4(frame):
    clear()
    t = easing(t_local(frame, 3))

    draw_label(5.4, 10.2, 'AlphaFold: масштаб', color=YELLOW, size=30)

    # Цифры
    stats = [
        ('200 млн', 'структур белков\n(весь UniProt)', GREEN, 2.7, 7.5),
        ('2 мин', 'на предсказание\nvs 2 года X-ray', BLUE, 8.1, 7.5),
        ('AlphaFold3', 'РНК, ДНК,\nмалые молекулы', PURPLE, 2.7, 4.5),
        ('Nobel 2024', 'Hassabis + Jumper\n(Chemistry)', YELLOW, 8.1, 4.5),
    ]

    for stat_val, stat_label, col, sx, sy in stats:
        appear_t = easing(min(1.0, max(0, t * 4 - stats.index((stat_val, stat_label, col, sx, sy)) * 0.3)))
        if appear_t > 0.01:
            draw_box(sx, sy, 4.2, 2.2, col, alpha=appear_t * 0.15)
            ax.text(sx, sy + 0.55, stat_val, color=col, fontsize=28,
                    ha='center', va='center', fontweight='bold', alpha=appear_t,
                    path_effects=[pe.withStroke(linewidth=3, foreground=BG)])
            ax.text(sx, sy - 0.35, stat_label, color=WHITE, fontsize=14,
                    ha='center', va='center', alpha=appear_t * 0.85)

    # Применения
    apps_alpha = easing(min(1.0, max(0, (t - 0.55) * 4)))
    draw_box(5.4, 2.1, 9.8, 1.8, '#1C2128', alpha=apps_alpha * 0.9)
    ax.text(5.4, 2.6, 'Применения: Insilico Medicine (первый AI-designed drug в Phase 2),',
            color=WHITE, fontsize=13, ha='center', va='center',
            fontweight='bold', alpha=apps_alpha)
    ax.text(5.4, 2.15, 'CASP14: GDT ≈ 92 (Human Expert: ~87)  |  ESMFold (Meta), RoseTTAFold (UW)',
            color=GRAY, fontsize=12, ha='center', va='center', alpha=apps_alpha * 0.85)

    # Финальный факт
    final_alpha = easing(min(1.0, max(0, (t - 0.78) * 8)))
    draw_label(5.4, 1.1, 'AI4Science: это наше направление 🔬',
               color=CYAN, size=18, alpha=final_alpha)


# ═══════════════════════════════════════════════════════════════════
# ГЛАВНАЯ ФУНКЦИЯ КАДРА
# ═══════════════════════════════════════════════════════════════════

def draw_frame(frame):
    if   frame < PHASE_FRAMES[0][1]: draw_phase1(frame)
    elif frame < PHASE_FRAMES[1][1]: draw_phase2(frame)
    elif frame < PHASE_FRAMES[2][1]: draw_phase3(frame)
    else:                              draw_phase4(frame)

    # Постоянный прогресс-бар
    progress = frame / (N_FRAMES - 1)
    ax.add_patch(plt.Rectangle((0, 0), 10.8 * progress, 0.05, color=CYAN, alpha=0.7))


# ═══════════════════════════════════════════════════════════════════
# РЕНДЕР
# ═══════════════════════════════════════════════════════════════════

anim = FuncAnimation(fig, draw_frame, frames=N_FRAMES, interval=1000/FPS)

writer = FFMpegWriter(fps=FPS, bitrate=2000,
                      extra_args=['-vcodec', 'libx264', '-pix_fmt', 'yuv420p',
                                  '-preset', 'fast', '-crf', '23'])

out_path = '/root/Strategy/content/drafts/alphafold_animation.mp4'
anim.save(out_path, writer=writer, dpi=100)
print(f"Saved: {out_path}")

import os
size_kb = os.path.getsize(out_path) // 1024
print(f"File size: {size_kb} KB")
