"""
Peano / Hilbert Curve Animation for @fminxyz Telegram Channel
1080x1080 px, 12 fps, 25 seconds (= 300 frames)

Four acts:
  Act 1 (0-7s):   Hilbert curve iterations 1 → 5  (progressive draw)
  Act 2 (7-14s):  Level-6 curve filling the square (rainbow gradient)
  Act 3 (14-20s): Locality property demo
  Act 4 (20-25s): Applications text reveal
"""

import os
import sys
import tempfile
import subprocess
import numpy as np
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
from matplotlib.collections import LineCollection
import warnings
warnings.filterwarnings('ignore')

OUTPUT = "/root/Strategy/content/drafts/peano_curve_animation.mp4"
FPS    = 12
DUR    = 25
NFRAM  = FPS * DUR   # 300
DPI    = 100
FW, FH = 10.8, 10.8

BG   = "#1a1a2e"
PAN  = "#16213e"
CYAN = "#00d4ff"
GREE = "#00ff99"
PINK = "#ff6b9d"
TXT  = "#e8e8f0"
ACC  = "#a8a8ff"
GRD  = "#2a2a4a"
YELL = "#ffe066"
ORG  = "#ff9933"

# ── Hilbert curve generation (iterative, integer grid) ──────────────────────
def hilbert_xy(n):
    """Return (x, y) arrays of Hilbert curve points for order n.
    Grid size = 2^n, so 4^n points total."""
    N = 1 << n          # 2^n
    total = N * N
    xs = np.zeros(total, dtype=float)
    ys = np.zeros(total, dtype=float)
    for i in range(total):
        x = y = 0
        s = 1
        t = i
        while s < N:
            rx = 1 if (t & 2) else 0
            ry = 1 if (t & 1) ^ rx else 0
            # rotate
            if ry == 0:
                if rx == 1:
                    x = s - 1 - x
                    y = s - 1 - y
                x, y = y, x
            x += s * rx
            y += s * ry
            t >>= 2
            s <<= 1
        xs[i] = x
        ys[i] = y
    # normalize to [0, 1]
    xs = xs / (N - 1)
    ys = ys / (N - 1)
    return xs, ys


# Precompute curves for levels 1..6
print("Precomputing Hilbert curves ...")
curves = {}
for level in range(1, 7):
    curves[level] = hilbert_xy(level)
    print(f"  level {level}: {len(curves[level][0])} points")

# ── Frame timing ─────────────────────────────────────────────────────────────
A1S, A1E = 0,         7 * FPS     # 0..84
A2S, A2E = 7 * FPS,   14 * FPS    # 84..168
A3S, A3E = 14 * FPS,  20 * FPS    # 168..240
A4S, A4E = 20 * FPS,  NFRAM       # 240..300
BLEND     = 4


def cl(v, lo=0.0, hi=1.0):
    return float(max(lo, min(hi, v)))


def fp(local, s, e, dur):
    a = s * dur
    b = e * dur
    return cl((local - a) / max(1, b - a))


def ease(t):
    """Smooth ease in-out."""
    return t * t * (3 - 2 * t)


def sty_ax(ax):
    ax.set_facecolor(PAN)
    ax.set_xlim(-0.06, 1.06)
    ax.set_ylim(-0.06, 1.06)
    ax.set_aspect('equal')
    ax.tick_params(left=False, bottom=False,
                   labelleft=False, labelbottom=False)
    for sp in ax.spines.values():
        sp.set_color(GRD)
        sp.set_linewidth(1.2)


# ── ACT 1: Hilbert iterations 1→5 ────────────────────────────────────────────
def act1(f, ax, axT, axB):
    lf = f - A1S
    dur = A1E - A1S

    # Show levels sequentially, each occupying ~1.2s (14 frames)
    # Levels 1-5 across the 7s window with overlap
    # Level k visible in [k-1, k+0.4] / 5 fraction
    level_times = {
        1: (0.00, 0.20),
        2: (0.18, 0.38),
        3: (0.35, 0.56),
        4: (0.53, 0.74),
        5: (0.71, 1.00),
    }

    # Determine which level to show
    current_level = 1
    for lv in [5, 4, 3, 2, 1]:
        ts, te = level_times[lv]
        if lf / dur >= ts:
            current_level = lv
            break

    ts, te = level_times[current_level]
    progress_in_level = fp(lf, ts, te, dur)
    draw_frac = ease(progress_in_level)

    xs, ys = curves[current_level]
    n_pts = len(xs)
    n_draw = max(2, int(draw_frac * n_pts))

    # Draw curve up to n_draw
    ax.plot(xs[:n_draw], ys[:n_draw],
            color=CYAN, lw=max(0.5, 3.0 - 0.4 * current_level),
            alpha=0.92, solid_capstyle='round')

    # Level label
    ax.text(0.5, -0.04, f'Итерация {current_level}  ({n_pts} точек)',
            ha='center', va='top', fontsize=14, color=ACC,
            transform=ax.transAxes)

    # Title
    fa0 = fp(lf, 0.00, 0.12, dur)
    axT.text(0.5, 0.55, 'Кривая Гильберта', ha='center', va='center',
             fontsize=30, fontweight='bold', color=TXT, alpha=fa0)
    axT.text(0.5, 0.10, 'Hilbert Curve', ha='center', va='center',
             fontsize=16, color=ACC, alpha=fa0 * 0.8)


# ── ACT 2: Space-filling at level 6 ──────────────────────────────────────────
def act2(f, ax, axT, axB):
    lf = f - A2S
    dur = A2E - A2S

    xs, ys = curves[6]
    n_pts = len(xs)

    draw_frac = ease(fp(lf, 0.00, 0.80, dur))
    n_draw = max(2, int(draw_frac * n_pts))

    # Segments with viridis colormap
    t_vals = np.linspace(0, 1, n_draw)
    segs = np.stack([np.c_[xs[:n_draw - 1], ys[:n_draw - 1]],
                     np.c_[xs[1:n_draw],    ys[1:n_draw]]], axis=1)
    lc = LineCollection(segs, cmap='plasma',
                        norm=plt.Normalize(0, 1),
                        linewidths=0.9, alpha=0.88)
    lc.set_array(t_vals[:-1])
    ax.add_collection(lc)

    # Fill percentage text
    fill_pct = draw_frac * 100
    fa_pct = fp(lf, 0.25, 0.50, dur)
    ax.text(0.5, -0.04, f'Заполнение: {fill_pct:.0f}%',
            ha='center', va='top', fontsize=14, color=GREE,
            transform=ax.transAxes, alpha=fa_pct)

    fa0 = fp(lf, 0.00, 0.12, dur)
    axT.text(0.5, 0.55, 'Заполнение квадрата', ha='center', va='center',
             fontsize=28, fontweight='bold', color=TXT, alpha=fa0)
    axT.text(0.5, 0.10, '4096 × 4096 = 16 млн точек', ha='center',
             va='center', fontsize=14, color=ACC, alpha=fa0 * 0.7)

    # Colorbar label
    fa_cb = fp(lf, 0.50, 0.80, dur)
    axB.text(0.5, 0.65, 'Цвет = позиция вдоль кривой',
             ha='center', va='center', fontsize=14, color=ACC, alpha=fa_cb)
    axB.text(0.5, 0.25,
             'непрерывная биективная функция [0,1] → [0,1]²',
             ha='center', va='center', fontsize=12,
             color=YELL, alpha=fa_cb * 0.85)


# ── ACT 3: Locality property ──────────────────────────────────────────────────
def act3(f, ax, axT, axB):
    lf = f - A3S
    dur = A3E - A3S

    # Full level-5 curve as backdrop
    xs5, ys5 = curves[5]
    ax.plot(xs5, ys5, color='#334466', lw=1.2, alpha=0.5, zorder=1)

    # Two nearby points on the curve and their 2D proximity
    # Pick indices ~5% apart on level-5 curve
    n5 = len(xs5)
    idx_a = int(0.35 * n5)
    idx_b = int(0.40 * n5)   # close on curve
    idx_c = int(0.80 * n5)   # far on curve, far in space

    xa, ya = xs5[idx_a], ys5[idx_a]
    xb, yb = xs5[idx_b], ys5[idx_b]
    xc, yc = xs5[idx_c], ys5[idx_c]

    # Phase 1: show two close points
    f1 = fp(lf, 0.00, 0.28, dur)
    if f1 > 0.01:
        ax.scatter([xa, xb], [ya, yb], color=CYAN, s=120, zorder=5, alpha=f1)

    # Phase 2: draw arrow between them in 2D + text
    f2 = fp(lf, 0.22, 0.50, dur)
    if f2 > 0.01:
        ax.annotate('', xy=(xb, yb), xytext=(xa, ya),
                    arrowprops=dict(arrowstyle='<->', color=CYAN,
                                    lw=2.0, alpha=f2))
        dist2d = np.sqrt((xa - xb)**2 + (ya - yb)**2)
        mx, my = (xa + xb) / 2, (ya + yb) / 2
        ax.text(mx + 0.04, my + 0.04,
                f'd₂D = {dist2d:.3f}', fontsize=11, color=CYAN, alpha=f2)

    # Phase 3: show position on 1D curve (highlight segment)
    f3 = fp(lf, 0.40, 0.65, dur)
    if f3 > 0.01:
        seg_len = idx_b - idx_a
        ax.plot(xs5[idx_a:idx_b], ys5[idx_a:idx_b],
                color=YELL, lw=3.5, alpha=f3 * 0.9, zorder=3)
        frac_1d = seg_len / n5
        ax.text(0.5, -0.04, f'Близко на кривой: {frac_1d * 100:.1f}% пути',
                ha='center', va='top', fontsize=13, color=YELL,
                transform=ax.transAxes, alpha=f3)

    # Phase 4: locality text
    f4 = fp(lf, 0.60, 0.90, dur)
    if f4 > 0.01:
        ax.text(0.5, 0.5, 'Locality\nPreserving',
                ha='center', va='center', fontsize=22,
                fontweight='bold', color=GREE, alpha=f4 * 0.9,
                transform=ax.transAxes,
                bbox=dict(boxstyle='round,pad=0.5', facecolor='#001a0d',
                          edgecolor=GREE, alpha=f4 * 0.7))

    fa0 = fp(lf, 0.00, 0.14, dur)
    axT.text(0.5, 0.55, 'Свойство локальности', ha='center', va='center',
             fontsize=28, fontweight='bold', color=TXT, alpha=fa0)
    axT.text(0.5, 0.10, 'Близкие точки на кривой → близкие в 2D',
             ha='center', va='center', fontsize=14, color=CYAN, alpha=fa0)


# ── ACT 4: Applications ───────────────────────────────────────────────────────
def act4(f, ax, axT, axB):
    lf = f - A4S
    dur = A4E - A4S

    # Faded curve in background
    xs5, ys5 = curves[5]
    ax.plot(xs5, ys5, color='#223355', lw=1.0, alpha=0.3, zorder=1)

    apps = [
        (0.08, 0.42, "LSH\n(Locality-Sensitive Hashing)", CYAN,   0.10, 0.48),
        (0.30, 0.42, "Database indexing\n(Z-order / Morton)", GREE,  0.30, 0.65),
        (0.55, 0.42, "Neural architecture\n(ViT patch ordering)", ACC,   0.52, 0.82),
        (0.75, 0.42, "Google S2\ngeo-indexing",                 YELL,  0.72, 1.00),
    ]

    # Draw each app as floating box appearing sequentially
    for (ts, te, label, color, xs_pos, ys_pos) in apps:
        fa = fp(lf, ts, te, dur)
        if fa > 0.01:
            ax.text(xs_pos, ys_pos, label,
                    ha='center', va='center',
                    fontsize=13, fontweight='bold',
                    color=color, alpha=fa,
                    transform=ax.transAxes,
                    bbox=dict(boxstyle='round,pad=0.5',
                              facecolor=BG,
                              edgecolor=color,
                              linewidth=1.5,
                              alpha=fa * 0.9))

    # Connection lines from curve to boxes (appear with boxes)
    # Just a visual connector dot on the curve
    connector_pts = [
        (int(0.15 * len(xs5)), CYAN),
        (int(0.35 * len(xs5)), GREE),
        (int(0.60 * len(xs5)), ACC),
        (int(0.80 * len(xs5)), YELL),
    ]
    for (ci, (ts, te, *_)), (idx, color) in zip(
            enumerate(apps), connector_pts):
        fa = fp(lf, ts, te, dur)
        if fa > 0.01:
            ax.scatter([xs5[idx]], [ys5[idx]],
                       color=color, s=80, zorder=4, alpha=fa)

    fa0 = fp(lf, 0.00, 0.14, dur)
    axT.text(0.5, 0.55, 'Применения', ha='center', va='center',
             fontsize=32, fontweight='bold', color=TXT, alpha=fa0)
    axT.text(0.5, 0.10, 'Кривая Гильберта в современных системах',
             ha='center', va='center', fontsize=14, color=ACC, alpha=fa0 * 0.8)


# ── Main render loop ──────────────────────────────────────────────────────────
def make_frame(frame, fig, ax, axT, axB):
    fig.patch.set_facecolor(BG)
    ax.cla()
    axT.cla(); axT.axis('off')
    axB.cla(); axB.axis('off')
    axT.set_xlim(0, 1); axT.set_ylim(0, 1)
    axB.set_xlim(0, 1); axB.set_ylim(0, 1)

    sty_ax(ax)

    if frame < A1E - BLEND:
        act1(frame, ax, axT, axB)
    elif frame < A2S + BLEND:
        p = (frame - (A1E - BLEND)) / (2 * BLEND)
        if p < 0.5:
            act1(frame, ax, axT, axB)
        else:
            act2(frame, ax, axT, axB)
    elif frame < A2E - BLEND:
        act2(frame, ax, axT, axB)
    elif frame < A3S + BLEND:
        p = (frame - (A2E - BLEND)) / (2 * BLEND)
        if p < 0.5:
            act2(frame, ax, axT, axB)
        else:
            act3(frame, ax, axT, axB)
    elif frame < A3E - BLEND:
        act3(frame, ax, axT, axB)
    elif frame < A4S + BLEND:
        p = (frame - (A3E - BLEND)) / (2 * BLEND)
        if p < 0.5:
            act3(frame, ax, axT, axB)
        else:
            act4(frame, ax, axT, axB)
    else:
        act4(frame, ax, axT, axB)


def main():
    tmpdir = tempfile.mkdtemp(prefix='peano_')
    print(f"Rendering {NFRAM} frames ({DUR}s @ {FPS}fps) → {tmpdir}")

    fig = plt.figure(figsize=(FW, FH), dpi=DPI, facecolor=BG)
    ax  = fig.add_axes([0.08, 0.14, 0.84, 0.70])
    axT = fig.add_axes([0.0, 0.86, 1.0, 0.12], facecolor='none')
    axT.set_xlim(0, 1); axT.set_ylim(0, 1); axT.axis('off')
    axB = fig.add_axes([0.0, 0.00, 1.0, 0.13], facecolor='none')
    axB.set_xlim(0, 1); axB.set_ylim(0, 1); axB.axis('off')

    for i in range(NFRAM):
        make_frame(i, fig, ax, axT, axB)
        fig.savefig(os.path.join(tmpdir, f"f{i:04d}.png"),
                    facecolor=BG, dpi=DPI)
        if i % 30 == 0:
            print(f"  {i}/{NFRAM}", flush=True)

    plt.close(fig)
    print("Encoding with ffmpeg ...")

    cmd = ['ffmpeg', '-y',
           '-framerate', str(FPS),
           '-i', os.path.join(tmpdir, 'f%04d.png'),
           '-c:v', 'libx264', '-preset', 'fast', '-crf', '20',
           '-pix_fmt', 'yuv420p', '-movflags', '+faststart',
           OUTPUT]
    r = subprocess.run(cmd, capture_output=True, text=True)
    if r.returncode != 0:
        print("ffmpeg error:\n", r.stderr[-2000:])
        sys.exit(1)

    import shutil; shutil.rmtree(tmpdir)
    size_mb = os.path.getsize(OUTPUT) / 1e6
    print(f"Done! {OUTPUT} ({size_mb:.1f} MB)")

    # Verify with ffprobe
    probe = subprocess.run(
        ['ffprobe', '-v', 'quiet', '-print_format', 'json',
         '-show_streams', OUTPUT],
        capture_output=True, text=True)
    print("ffprobe output:", probe.stdout[:500])


if __name__ == '__main__':
    main()
