"""
Heavy Tails Animation for @fminxyz Telegram Channel
1080x1080 px, 12 fps, 25 seconds (= 300 frames)

Three acts:
  Act 1 (0-8s):   PDF comparison: Gaussian vs Cauchy
  Act 2 (8-16s):  Sampling animation with outliers
  Act 3 (16-25s): Log-log power law plot
"""

import os
import sys
import tempfile
import subprocess
import numpy as np
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
from matplotlib.lines import Line2D
from scipy import stats
import warnings
warnings.filterwarnings('ignore')

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

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

rng = np.random.default_rng(42)

x_pdf       = np.linspace(-8, 8, 400)
gauss_pdf   = stats.norm.pdf(x_pdf, 0, 1)
cauchy_pdf  = stats.cauchy.pdf(x_pdf, 0, 1)
x_tail      = np.linspace(2.5, 8, 150)
gauss_tail  = stats.norm.pdf(x_tail, 0, 1)
cauchy_tail = stats.cauchy.pdf(x_tail, 0, 1)

N_S      = 140
gauss_s  = rng.standard_normal(N_S)
cauchy_s = rng.standard_cauchy(N_S)

x_ll        = np.logspace(0, 4, 100)
pareto_ccdf = x_ll ** (-1.5)
exp_ccdf    = np.exp(-x_ll / 150)

A1S, A1E = 0,        8 * FPS    # 0..96
A2S, A2E = 8 * FPS,  16 * FPS   # 96..192
A3S, A3E = 16 * FPS, NFRAM      # 192..300
BLEND     = 5


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 make_frame(frame, fig, axes):
    ax1, ax1t, ax2g, ax2h, ax3, axT, axB = axes
    fig.patch.set_facecolor(BG)

    def hide():
        for a in [ax1, ax1t, ax2g, ax2h, ax3]:
            a.set_visible(False); a.cla()
        for a in [axT, axB]:
            a.cla(); a.axis('off')
            a.set_xlim(0, 1); a.set_ylim(0, 1)

    def sty(ax, xl='', yl='', ti=''):
        ax.set_facecolor(PAN)
        ax.tick_params(colors=TXT, labelsize=10)
        for sp in ax.spines.values(): sp.set_color(GRD)
        if xl: ax.set_xlabel(xl, fontsize=11, color=TXT)
        if yl: ax.set_ylabel(yl, fontsize=11, color=TXT)
        if ti: ax.set_title(ti, fontsize=13, color=TXT, pad=5)
        ax.grid(True, color=GRD, alpha=0.35, linewidth=0.5)

    # ── ACT 1 ──────────────────────────────────────────────────────────────────
    def act1(f):
        lf = f - A1S
        dur = A1E - A1S
        hide()
        ax1.set_visible(True)
        sty(ax1, xl='x', yl='PDF(x)')
        ax1.set_xlim(-6.5, 6.5)
        ax1.set_ylim(-0.01, 0.42)

        fg = fp(lf, 0.00, 0.38, dur)
        ng = max(2, int(fg * len(x_pdf)))
        ax1.plot(x_pdf[:ng], gauss_pdf[:ng], color=CYAN, lw=2.5, zorder=3)

        fc = fp(lf, 0.22, 0.60, dur)
        nc = max(2, int(fc * len(x_pdf)))
        ax1.plot(x_pdf[:nc], cauchy_pdf[:nc], color=PINK, lw=2.5, zorder=3)

        ft = fp(lf, 0.52, 0.82, dur)
        if ft > 0:
            mask = x_pdf > 2.5
            ax1.fill_between(x_pdf[mask], 0, gauss_pdf[mask],
                             color=CYAN, alpha=ft * 0.20)
            ax1.fill_between(x_pdf[mask], 0, cauchy_pdf[mask],
                             color=PINK, alpha=ft * 0.30)

        fa = fp(lf, 0.68, 1.00, dur)
        if fa > 0.05:
            ax1.annotate("Тяжёлый хвост",
                xy=(4.0, stats.cauchy.pdf(4.0)),
                xytext=(3.1, 0.05 + 0.13 * fa),
                fontsize=12, color=PINK, fontweight='bold', alpha=fa,
                arrowprops=dict(arrowstyle='->', color=PINK, lw=1.5,
                                connectionstyle='arc3,rad=-0.2'))
            ax1.annotate("Лёгкий хвост",
                xy=(3.0, stats.norm.pdf(3.0)),
                xytext=(2.0, 0.12 + 0.12 * fa),
                fontsize=12, color=CYAN, alpha=fa,
                arrowprops=dict(arrowstyle='->', color=CYAN, lw=1.5,
                                connectionstyle='arc3,rad=0.2'))

        fz = fp(lf, 0.75, 1.00, dur)
        ax1t.set_visible(fz > 0.05)
        if fz > 0.05:
            sty(ax1t, ti='Хвост (x > 2.5)')
            ax1t.set_facecolor('#0d0d1f')
            ax1t.set_xlim(2.5, 7.5)
            ax1t.set_ylim(-0.003, 0.075)
            for sp in ax1t.spines.values():
                sp.set_color(PINK); sp.set_linewidth(1.8)
            nt = max(2, int(fz * len(x_tail)))
            ax1t.plot(x_tail[:nt], gauss_tail[:nt], color=CYAN, lw=2, alpha=0.9)
            ax1t.plot(x_tail[:nt], cauchy_tail[:nt], color=PINK, lw=2, alpha=0.9)
            ax1t.fill_between(x_tail[:nt], 0, cauchy_tail[:nt], color=PINK, alpha=0.25)

        if fc > 0.3:
            legs = [Line2D([0],[0], color=CYAN, lw=2.5, label='Гауссовское N(0,1)'),
                    Line2D([0],[0], color=PINK, lw=2.5, label='Коши (тяжёлый хвост)')]
            ax1.legend(handles=legs, loc='upper left', fontsize=11,
                       facecolor='#0d0d20', edgecolor=GRD, labelcolor=TXT)

        fa0 = fp(lf, 0.00, 0.18, dur)
        axT.text(0.5, 0.55, 'Тяжёлые хвосты', ha='center', va='center',
                 fontsize=30, fontweight='bold', color=TXT, alpha=fa0)
        axT.text(0.5, 0.12, 'Heavy Tails', ha='center', va='center',
                 fontsize=16, color=ACC, alpha=fa0 * 0.8)
        axB.text(0.5, 0.6, 'Функция плотности вероятности',
                 ha='center', va='center', fontsize=14,
                 color=ACC, alpha=fp(lf, 0.08, 0.28, dur))

    # ── ACT 2 ──────────────────────────────────────────────────────────────────
    def act2(f):
        lf = f - A2S
        dur = A2E - A2S
        hide()
        ax2g.set_visible(True); ax2h.set_visible(True)
        sty(ax2g, xl='Сэмпл №', yl='Значение', ti='Гауссовское N(0,1)')
        sty(ax2h, xl='Сэмпл №', yl='', ti='Коши (тяжёлый хвост)')
        YR = 22
        for ax in [ax2g, ax2h]:
            ax.set_xlim(0, N_S); ax.set_ylim(-YR, YR)
            ax.axhline(0, color=GRD, lw=0.8)

        n = min(N_S, max(0, int(fp(lf, 0.00, 0.88, dur) * N_S)))

        if n > 0:
            xs = np.arange(n)
            gv = gauss_s[:n]
            cv = np.clip(cauchy_s[:n], -YR*1.5, YR*1.5)
            ax2g.scatter(xs, gv, c=CYAN, s=14, alpha=0.7, zorder=3)
            cc = [YELL if abs(cauchy_s[i]) > 8 else PINK for i in range(n)]
            ax2h.scatter(xs, cv, c=cc, s=14, alpha=0.7, zorder=3)
            if n > 1:
                ax2g.plot(xs, gv, color=CYAN, alpha=0.18, lw=0.6)
                ax2h.plot(xs, np.clip(cauchy_s[:n], -YR, YR),
                          color=PINK, alpha=0.18, lw=0.6)

        n_out = int(np.sum(np.abs(cauchy_s[:n]) > 8)) if n > 0 else 0
        fa_o = fp(lf, 0.35, 0.58, dur)
        if n_out > 0 and fa_o > 0:
            ax2h.text(0.97, 0.96, f'Выбросов: {n_out}',
                      transform=ax2h.transAxes, ha='right', va='top',
                      fontsize=14, color=YELL, fontweight='bold', alpha=fa_o)

        if n > 15:
            fa_s = fp(lf, 0.45, 0.68, dur)
            if fa_s > 0:
                ax2g.text(0.97, 0.96, f'σ ≈ {np.std(gauss_s[:n]):.2f}',
                          transform=ax2g.transAxes, ha='right', va='top',
                          fontsize=14, color=CYAN, fontweight='bold', alpha=fa_s)

        ab = fp(lf, 0.28, 0.55, dur) * 0.12
        ax2g.axhspan(-2, 2, color=CYAN, alpha=ab)
        fa_b = fp(lf, 0.42, 0.62, dur)
        if fa_b > 0:
            ax2g.text(1.5, 2.5, '±2σ = 95%', fontsize=10, color=CYAN, alpha=fa_b)

        fa0 = fp(lf, 0.00, 0.15, dur)
        axT.text(0.5, 0.55, 'Сэмплирование', ha='center', va='center',
                 fontsize=30, fontweight='bold', color=TXT, alpha=fa0)
        axB.text(0.5, 0.6, 'Тяжёлые хвосты = экстремальные выбросы',
                 ha='center', va='center', fontsize=14, color=PINK, alpha=fa0)

    # ── ACT 3 ──────────────────────────────────────────────────────────────────
    def act3(f):
        lf = f - A3S
        dur = A3E - A3S
        hide()
        ax3.set_visible(True)
        sty(ax3, xl='x  (log scale)', yl='P(X > x)  (log scale)',
            ti='Степенной закон = прямая линия в log-log')
        ax3.set_xscale('log'); ax3.set_yscale('log')
        ax3.set_xlim(1, 15000); ax3.set_ylim(5e-6, 2.0)

        p1 = fp(lf, 0.00, 0.42, dur)
        n1 = max(2, int(p1 * len(x_ll)))
        ax3.plot(x_ll[:n1], exp_ccdf[:n1], color=CYAN, lw=2.5)

        p2 = fp(lf, 0.25, 0.70, dur)
        n2 = max(2, int(p2 * len(x_ll)))
        ax3.plot(x_ll[:n2], pareto_ccdf[:n2], color=PINK, lw=3.0)

        p3 = fp(lf, 0.60, 0.88, dur)
        if p3 > 0.01:
            ax3.text(0.60, 0.72, 'Наклон = −α = −1.5',
                     transform=ax3.transAxes,
                     fontsize=14, color=PINK, fontweight='bold', alpha=p3,
                     bbox=dict(boxstyle='round,pad=0.4', facecolor='#1a0020',
                               edgecolor=PINK, alpha=p3 * 0.7))

        anns = [(200,  "Частоты слов\n(Закон Ципфа)", 0.55),
                (1500, "Распределение\nбогатства",    0.68),
                (8000, "Scaling Laws\nнейросетей",    0.80)]
        for xi, label, sf in anns:
            pa = fp(lf, sf, sf + 0.18, dur)
            if pa > 0.01:
                yi = xi ** (-1.5)
                ax3.annotate(label,
                    xy=(xi, yi), xytext=(xi * 1.8, yi * 5),
                    fontsize=11, color=ACC, fontweight='bold', alpha=pa,
                    arrowprops=dict(arrowstyle='->', color=ACC, lw=1.2,
                                    connectionstyle='arc3,rad=0.15'),
                    bbox=dict(boxstyle='round,pad=0.3', facecolor='#0d0d25',
                              edgecolor=ACC, alpha=pa * 0.8))

        if p2 > 0.3:
            fa = fp(lf, 0.45, 0.65, dur)
            legs = [Line2D([0],[0], color=CYAN, lw=2.5,
                           label='Экспоненциальный (лёгкий)'),
                    Line2D([0],[0], color=PINK, lw=3.0,
                           label='Степенной закон (тяжёлый)')]
            ax3.legend(handles=legs, loc='lower left', fontsize=11,
                       facecolor='#0d0d20', edgecolor=GRD,
                       labelcolor=TXT, framealpha=max(0, fa))

        fa0 = fp(lf, 0.00, 0.16, dur)
        axT.text(0.5, 0.55, 'Реальный мир', ha='center', va='center',
                 fontsize=30, fontweight='bold', color=TXT, alpha=fa0)
        axB.text(0.5, 0.6, 'Тяжёлые хвосты повсюду в ML и природе',
                 ha='center', va='center', fontsize=14, color=ACC, alpha=fa0)

    # ── dispatch ───────────────────────────────────────────────────────────────
    if frame < A1E - BLEND:
        act1(frame)
    elif frame < A2S + BLEND:
        act1(frame) if (frame - (A1E - BLEND)) < BLEND else act2(frame)
    elif frame < A2E - BLEND:
        act2(frame)
    elif frame < A3S + BLEND:
        act2(frame) if (frame - (A2E - BLEND)) < BLEND else act3(frame)
    else:
        act3(frame)


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

    fig  = plt.figure(figsize=(FW, FH), dpi=DPI, facecolor=BG)
    ax1  = fig.add_axes([0.09, 0.17, 0.86, 0.64])
    ax1t = fig.add_axes([0.60, 0.38, 0.31, 0.31])
    ax2g = fig.add_axes([0.07, 0.17, 0.41, 0.64])
    ax2h = fig.add_axes([0.54, 0.17, 0.41, 0.64])
    ax3  = fig.add_axes([0.10, 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.15], facecolor='none')
    axB.set_xlim(0,1); axB.set_ylim(0,1); axB.axis('off')
    axes = (ax1, ax1t, ax2g, ax2h, ax3, axT, axB)

    for i in range(NFRAM):
        make_frame(i, fig, axes)
        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)")


if __name__ == '__main__':
    main()
