"""
RLHF Pipeline Animation: SFT → Reward Model → PPO
@fminxyz Series 3, Post 1 — 11 марта 2026
1080x1080 px, 2 fps, 50 frames (25 sec)
"""

import numpy as np
import matplotlib.pyplot as plt
import matplotlib.patches as mpatches
import matplotlib.animation as animation
from matplotlib.patches import FancyBboxPatch

OUTPUT = "/root/Strategy/content/drafts/rlhf_animation.mp4"
FPS = 2
N_FRAMES = 50

bg = '#0a0a14'

fig, ax = plt.subplots(figsize=(10, 10), dpi=108, facecolor=bg)

# Simulated policy "quality" curve across PPO iterations
ppo_iters = 12
quality = np.array([0.3, 0.38, 0.48, 0.56, 0.63, 0.69, 0.73, 0.76, 0.78, 0.80, 0.81, 0.82])
kl_div = np.array([0.0, 0.02, 0.05, 0.09, 0.14, 0.18, 0.21, 0.23, 0.24, 0.25, 0.25, 0.26])

def draw_frame(i):
    ax.clear()
    ax.set_facecolor(bg)
    ax.set_xlim(0, 10)
    ax.set_ylim(0, 10)
    ax.axis('off')

    fade = min(1.0, (i + 1) / 4)

    # Title
    ax.text(5, 9.6, 'RLHF: Reinforcement Learning from Human Feedback',
            ha='center', fontsize=13, fontweight='bold', color='white',
            fontfamily='monospace', alpha=fade)
    ax.text(5, 9.25, 'Как ChatGPT научился слушаться',
            ha='center', fontsize=11, color='#8899cc',
            fontfamily='monospace', alpha=fade, fontstyle='italic')

    # ─── 3 PIPELINE BOXES ─────────────────────────────────────────────────────
    # Box 1: SFT (always visible from frame 2)
    b1fade = min(1.0, max(0, (i - 1) / 4))
    b1 = FancyBboxPatch((0.3, 7.0), 2.8, 1.6,
        boxstyle='round,pad=0.12',
        facecolor='#0d1a2d', edgecolor='#4fc3f7', linewidth=2, alpha=b1fade)
    ax.add_patch(b1)
    ax.text(1.7, 8.35, '① SFT', ha='center', fontsize=13, color='#4fc3f7',
            fontfamily='monospace', alpha=b1fade, fontweight='bold')
    ax.text(1.7, 8.05, 'Supervised', ha='center', fontsize=10,
            color='#7bccee', fontfamily='monospace', alpha=b1fade)
    ax.text(1.7, 7.75, 'Fine-Tuning', ha='center', fontsize=10,
            color='#7bccee', fontfamily='monospace', alpha=b1fade)
    ax.text(1.7, 7.2, 'GPT-3 + демо\n= базовая модель', ha='center', fontsize=9,
            color='#446688', fontfamily='monospace', alpha=b1fade)

    # Arrow 1→2
    if i >= 6:
        a1fade = min(1.0, (i - 6) / 4)
        ax.annotate('', xy=(3.8, 7.8), xytext=(3.15, 7.8),
            arrowprops=dict(arrowstyle='->', color='#6677aa', lw=2, alpha=a1fade))

    # Box 2: Reward Model (visible from frame 8)
    b2fade = min(1.0, max(0, (i - 7) / 4))
    b2 = FancyBboxPatch((3.6, 7.0), 2.8, 1.6,
        boxstyle='round,pad=0.12',
        facecolor='#1a0d2d', edgecolor='#ce93d8', linewidth=2, alpha=b2fade)
    ax.add_patch(b2)
    ax.text(5.0, 8.35, '② Reward Model', ha='center', fontsize=12, color='#ce93d8',
            fontfamily='monospace', alpha=b2fade, fontweight='bold')
    ax.text(5.0, 8.0, 'y_w ≻ y_l', ha='center', fontsize=11,
            color='#bb88dd', fontfamily='monospace', alpha=b2fade)
    ax.text(5.0, 7.65, 'Human ranking', ha='center', fontsize=10,
            color='#9966bb', fontfamily='monospace', alpha=b2fade)
    ax.text(5.0, 7.2, 'RM(y) → reward\nскор ответа', ha='center', fontsize=9,
            color='#664488', fontfamily='monospace', alpha=b2fade)

    # Arrow 2→3
    if i >= 14:
        a2fade = min(1.0, (i - 14) / 4)
        ax.annotate('', xy=(7.1, 7.8), xytext=(6.45, 7.8),
            arrowprops=dict(arrowstyle='->', color='#6677aa', lw=2, alpha=a2fade))

    # Box 3: PPO (visible from frame 16)
    b3fade = min(1.0, max(0, (i - 15) / 4))
    b3 = FancyBboxPatch((6.9, 7.0), 2.8, 1.6,
        boxstyle='round,pad=0.12',
        facecolor='#0d1a0d', edgecolor='#81c784', linewidth=2, alpha=b3fade)
    ax.add_patch(b3)
    ax.text(8.3, 8.35, '③ PPO', ha='center', fontsize=13, color='#81c784',
            fontfamily='monospace', alpha=b3fade, fontweight='bold')
    ax.text(8.3, 8.0, 'Policy Gradient', ha='center', fontsize=10,
            color='#66bb66', fontfamily='monospace', alpha=b3fade)
    ax.text(8.3, 7.65, 'Optimization', ha='center', fontsize=10,
            color='#66bb66', fontfamily='monospace', alpha=b3fade)
    ax.text(8.3, 7.2, 'max E[R(y)]\n−β·KL(π||π_ref)', ha='center', fontsize=9,
            color='#447744', fontfamily='monospace', alpha=b3fade)

    # ─── PPO OBJECTIVE FORMULA (frame 22+) ─────────────────────────────────────
    if i >= 22:
        ffade = min(1.0, (i - 22) / 5)
        form_box = FancyBboxPatch((0.5, 5.6), 9.0, 1.1,
            boxstyle='round,pad=0.12',
            facecolor='#0d0d22', edgecolor='#4466ff', linewidth=1.5, alpha=ffade)
        ax.add_patch(form_box)
        ax.text(5, 6.45, 'Objective: max E[R(y)] − β · KL(π || π_ref)',
                ha='center', fontsize=12, color='#aabbff',
                fontfamily='monospace', alpha=ffade, fontweight='bold')
        ax.text(5, 6.05, 'R(y) = reward;    β·KL = penalty за отклонение от SFT',
                ha='center', fontsize=10, color='#667799',
                fontfamily='monospace', alpha=ffade)

    # ─── KL PENALTY NOTE (frame 28+) ─────────────────────────────────────────
    if i >= 28:
        kfade = min(1.0, (i - 28) / 4)
        ax.text(5, 5.3, 'KL-penalty = не уйти слишком далеко от SFT',
                ha='center', fontsize=11, color='#ffaa44',
                fontfamily='monospace', alpha=kfade)

    # ─── PPO QUALITY CURVE (frame 32+) ────────────────────────────────────────
    if i >= 32:
        gfade = min(1.0, (i - 32) / 5)

        # Draw axes
        ax.plot([0.8, 9.2], [1.5, 1.5], '-', color='#334455', lw=1.5, alpha=gfade)
        ax.plot([0.8, 0.8], [1.5, 4.8], '-', color='#334455', lw=1.5, alpha=gfade)
        ax.text(5, 1.2, 'PPO iterations ->', ha='center', fontsize=10,
                color='#445566', fontfamily='monospace', alpha=gfade)
        ax.text(0.4, 3.2, 'quality', ha='center', fontsize=9,
                color='#445566', fontfamily='monospace', alpha=gfade, rotation=90)

        # How many curve points to show
        curve_pts = max(0, min(ppo_iters, i - 31))
        if curve_pts > 0:
            xs = np.linspace(0.8, 9.0, ppo_iters)[:curve_pts]
            ys = 1.5 + quality[:curve_pts] * 3.3
            ax.plot(xs, ys, '-o', color='#81c784', lw=2.5,
                    markersize=5, alpha=gfade, zorder=5)

            # KL divergence (warning zone)
            ys_kl = 1.5 + kl_div[:curve_pts] * 3.3 * 0.8
            ax.plot(xs, ys_kl, '--', color='#ffaa44', lw=1.5, alpha=gfade * 0.6)

        ax.text(7.5, 4.7, '<- policy', ha='left', fontsize=9,
                color='#81c784', fontfamily='monospace', alpha=gfade)
        ax.text(7.5, 3.3, '<- KL drift', ha='left', fontsize=9,
                color='#ffaa44', fontfamily='monospace', alpha=gfade)

    # ─── FINAL RESULT (frame 46+) ──────────────────────────────────────────────
    if i >= 46:
        rfade = min(1.0, (i - 46) / 3)
        res_box = FancyBboxPatch((2.0, 0.2), 6.0, 0.9,
            boxstyle='round,pad=0.1',
            facecolor='#0a1a0a', edgecolor='#44ff88', linewidth=2, alpha=rfade)
        ax.add_patch(res_box)
        ax.text(5, 0.85, 'GPT-3 -> ChatGPT: helpful, harmless, honest',
                ha='center', fontsize=11, color='#44ff88',
                fontfamily='monospace', alpha=rfade, fontweight='bold')

    # Bottom tag
    if i >= 5:
        ax.text(9.8, 0.1, '@fminxyz', ha='right', fontsize=9,
                color='#223344', fontfamily='monospace',
                alpha=min(1.0, (i - 5) / 5))

anim = animation.FuncAnimation(fig, draw_frame, frames=N_FRAMES, interval=1000 // FPS)
anim.save(OUTPUT, writer='ffmpeg', fps=FPS, dpi=108,
          extra_args=['-vcodec', 'libx264', '-pix_fmt', 'yuv420p',
                      '-crf', '22', '-preset', 'fast'])
plt.close()
print(f"Saved: {OUTPUT}")