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

# ── Data: PDB depositions per year (from RCSB PDB historical data) ──────────
years = list(range(1976, 2024))
values = [
    1, 2, 4, 7, 14, 21, 32, 44, 65, 92,      # 1976-1985
    141, 177, 231, 369, 521, 737, 968, 1340,   # 1986-1993
    1820, 2530, 3358, 4208, 5194, 6461, 7893,  # 1994-2000 (cumulative-ish)
    9527, 11490, 13589, 16216, 19636, 24027,   # 2001-2006
    29174, 35619, 42855, 50715, 59962, 72981,  # 2007-2012 (approx cumulative)
    83897, 96818, 109979, 119677, 129464,      # 2013-2017
    138720, 151057, 162846, 173526,            # 2018-2021
    185688, 199662, 213433,                    # 2022-2024 (approx)
]
# Trim to match
n = min(len(years), len(values))
years = years[:n]
values = values[:n]

# ── Color scheme: vibrant gradient per bar height ──────────────────────────
# Background: deep dark navy
BG_COLOR    = '#0a0e1a'
CHART_BG    = '#0d1220'
AXIS_COLOR  = '#3a4a6a'
TEXT_COLOR  = '#c8d8f0'
LABEL_COLOR = '#7aa0d4'

def bar_color(fraction):
    """Return an RGBA color based on normalized bar height (0-1).
    Low bars -> cyan/teal, high bars -> orange/crimson."""
    # Interpolate through: teal -> green -> yellow -> orange -> red
    cmap = mcolors.LinearSegmentedColormap.from_list(
        'custom',
        ['#00f5d4', '#00bbf9', '#fee440', '#f15bb5', '#ff006e'],
        N=256
    )
    return cmap(fraction)

def bar_glow_color(fraction):
    """Lighter version of bar color for glow effect."""
    c = list(bar_color(fraction))
    c[3] = 0.18  # alpha for glow
    return tuple(c)

# ── Figure setup ─────────────────────────────────────────────────────────────
FPS = 30
DURATION = 17.3  # seconds
TOTAL_FRAMES = int(FPS * DURATION)

fig = plt.figure(figsize=(1914/100, 1080/100), dpi=100, facecolor=BG_COLOR)
ax = fig.add_axes([0.14, 0.13, 0.82, 0.78], facecolor=CHART_BG)

# Style axes
for spine in ax.spines.values():
    spine.set_color(AXIS_COLOR)
ax.spines['top'].set_visible(False)
ax.spines['right'].set_visible(False)
ax.tick_params(colors=LABEL_COLOR, which='both', length=0)
ax.xaxis.label.set_color(TEXT_COLOR)
ax.yaxis.label.set_color(TEXT_COLOR)

# Rounded chart background patch
fancy = FancyBboxPatch((0.01, 0.01), 0.98, 0.98,
                        boxstyle="round,pad=0.01",
                        transform=ax.transAxes,
                        linewidth=1.5,
                        edgecolor='#1e2d50',
                        facecolor=CHART_BG,
                        zorder=0)
ax.add_patch(fancy)

max_val = max(values)

def get_display_count(frame):
    """How many bars to show based on frame number."""
    # First ~2s: blank/intro, then bars appear one by one over ~13s, last 2s hold
    intro_frames = int(FPS * 1.5)
    hold_frames  = int(FPS * 2.0)
    anim_frames  = TOTAL_FRAMES - intro_frames - hold_frames
    if frame < intro_frames:
        return 0
    if frame >= TOTAL_FRAMES - hold_frames:
        return n
    progress = (frame - intro_frames) / anim_frames
    # Use easing: ease-in-out
    t = progress
    eased = t * t * (3 - 2 * t)
    return max(1, int(eased * n))

def get_bar_height_fraction(bar_idx, frame, display_count):
    """For the newest bar, animate it growing up; older bars are full height."""
    if bar_idx >= display_count:
        return 0.0
    if bar_idx < display_count - 1:
        return 1.0
    # Last bar: compute how far into its "slot" we are
    intro_frames = int(FPS * 1.5)
    hold_frames  = int(FPS * 2.0)
    anim_frames  = TOTAL_FRAMES - intro_frames - hold_frames
    if frame >= TOTAL_FRAMES - hold_frames:
        return 1.0
    progress = (frame - intro_frames) / anim_frames
    t = progress
    eased = t * t * (3 - 2 * t)
    # Fraction within current bar slot
    slot_size = 1.0 / n
    slot_start = bar_idx * slot_size
    within = (eased - slot_start) / slot_size
    return min(1.0, max(0.0, within))

def format_yval(v):
    if v >= 1_000_000:
        return f'{v/1_000_000:.0f}M'
    if v >= 1_000:
        return f'{v/1_000:.0f},000' if v % 1000 == 0 else f'{v:,.0f}'
    return f'{int(v):,}'

# Pre-compute y-axis max per frame (smoothly scales up)
def get_ymax(display_count):
    if display_count == 0:
        return 500
    visible_vals = [values[i] for i in range(display_count)]
    cur_max = max(visible_vals)
    # Round up to nice number
    mag = 10 ** int(np.log10(cur_max + 1))
    nice = np.ceil(cur_max / mag) * mag
    return max(nice, 100)

def animate(frame):
    ax.clear()
    ax.set_facecolor(CHART_BG)
    for spine in ax.spines.values():
        spine.set_color(AXIS_COLOR)
    ax.spines['top'].set_visible(False)
    ax.spines['right'].set_visible(False)
    ax.tick_params(colors=LABEL_COLOR, which='both', length=0, labelsize=9)

    display_count = get_display_count(frame)
    ymax = get_ymax(display_count) if display_count > 0 else 500

    # Grid lines
    ax.set_ylim(0, ymax * 1.12)
    yticks = ax.yaxis.get_major_locator()

    # Draw subtle grid
    for yval in np.linspace(0, ymax, 7):
        ax.axhline(yval, color='#1a2540', linewidth=0.6, zorder=1)

    bar_width = 0.72
    bar_positions = list(range(n))

    # Draw bars
    for i in range(display_count):
        frac = get_bar_height_fraction(i, frame, display_count)
        val = values[i] * frac
        color_frac = values[i] / max_val  # color based on final value
        color = bar_color(color_frac)
        gcolor = bar_glow_color(color_frac)

        # Glow shadow (wider, semi-transparent)
        if val > 0:
            ax.bar(i, val, width=bar_width + 0.6, color=gcolor,
                   zorder=2, align='center')

        # Main bar with gradient simulation via stacked thin bars
        if val > 0:
            segments = 40
            seg_h = val / segments
            for s in range(segments):
                seg_bottom = s * seg_h
                # Gradient: darker at bottom, brighter at top
                t = s / segments
                base_c = list(color)
                brightness = 0.55 + 0.45 * t
                seg_c = (
                    min(1.0, base_c[0] * brightness),
                    min(1.0, base_c[1] * brightness),
                    min(1.0, base_c[2] * brightness),
                    base_c[3]
                )
                ax.bar(i, seg_h, bottom=seg_bottom,
                       width=bar_width, color=seg_c,
                       zorder=3, align='center')

            # Top highlight (bright line at top of bar)
            hi_c = (
                min(1.0, color[0] * 1.6),
                min(1.0, color[1] * 1.6),
                min(1.0, color[2] * 1.6),
                0.9
            )
            ax.bar(i, val * 0.025, bottom=val * 0.975,
                   width=bar_width * 0.85, color=hi_c,
                   zorder=4, align='center')

    # X-axis labels (show every 5 years)
    tick_positions = [i for i, y in enumerate(years) if y % 5 == 0]
    tick_labels    = [str(years[i]) for i in tick_positions]
    ax.set_xticks(tick_positions)
    ax.set_xticklabels(tick_labels, color=LABEL_COLOR, fontsize=10,
                       fontfamily='DejaVu Sans')
    ax.set_xlim(-1, n)

    # Y-axis
    ax.set_ylim(0, ymax * 1.12)
    ax.yaxis.set_major_formatter(
        matplotlib.ticker.FuncFormatter(lambda x, _: f'{int(x):,}')
    )
    ax.tick_params(axis='y', colors=LABEL_COLOR, labelsize=9)

    # Axis labels
    ax.set_xlabel('Year', color=TEXT_COLOR, fontsize=16,
                  fontweight='bold', labelpad=12)
    ax.set_ylabel('Protein Structures', color=TEXT_COLOR, fontsize=13,
                  fontweight='bold', labelpad=12)

    # Title
    fig.text(0.5, 0.96, 'Protein Structures Deposited in PDB',
             ha='center', va='top', color=TEXT_COLOR,
             fontsize=15, fontweight='bold', transform=fig.transFigure)

    # Watermark "Ve"
    fig.text(0.975, 0.025, 'Ve', ha='right', va='bottom',
             color='#3a4a6a', fontsize=9, alpha=0.7,
             transform=fig.transFigure)

    # Show current year label on the newest visible bar
    if display_count > 0:
        idx = display_count - 1
        frac = get_bar_height_fraction(idx, frame, display_count)
        val = values[idx] * frac
        yr_label = str(years[idx])
        ax.text(idx, val + ymax * 0.02, yr_label,
                ha='center', va='bottom',
                color=TEXT_COLOR, fontsize=8, fontweight='bold', zorder=5)

print("Script loaded. Starting animation render...")

anim = FuncAnimation(fig, animate, frames=TOTAL_FRAMES,
                     interval=1000 / FPS, blit=False)

writer = FFMpegWriter(fps=FPS, codec='libx264',
                      extra_args=['-preset', 'fast', '-crf', '18',
                                  '-pix_fmt', 'yuv420p'])

out_path = '/home/daytona/workspace/video-color/protein_structures_colored.mp4'
print(f"Rendering {TOTAL_FRAMES} frames to {out_path} ...")
anim.save(out_path, writer=writer)
print("Done!")

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

# ── Colors ──────────────────────────────────────────────────────────────────
BG          = '#f0f0f0'
TEAL        = '#2db5a3'   # node fill
TEAL_DARK   = '#1a8c7c'   # node edge
ORANGE      = '#e07b54'   # active/highlight node
ORANGE_EDGE = '#c05a30'
BLUE_TITLE  = '#1a3a7c'   # main title
DARK_TEXT   = '#222222'
BULLET_TEXT = '#333333'
ARROW_BLUE  = '#4a90c4'
ARROW_ORANGE= '#e07b54'
EDGE_GRAY   = '#aaaaaa'
EDGE_ORANGE = '#e07b54'

# ── Canvas ───────────────────────────────────────────────────────────────────
fig = plt.figure(figsize=(19.2, 10.8), dpi=200, facecolor=BG)
ax  = fig.add_axes([0, 0, 1, 1], facecolor=BG)
ax.set_xlim(0, 1)
ax.set_ylim(0, 1)
ax.axis('off')

# ─────────────────────────────────────────────────────────────────────────────
# Helper: draw a filled circle node
# ─────────────────────────────────────────────────────────────────────────────
def node(cx, cy, r=0.018, color=TEAL, edge=TEAL_DARK, zorder=10, alpha=1.0):
    c = Circle((cx, cy), r, color=color, ec=edge, linewidth=1.5,
               transform=ax.transAxes, zorder=zorder)
    ax.add_patch(c)

def diamond(cx, cy, r=0.018, color=TEAL, edge=TEAL_DARK, zorder=10):
    # Diamond shape via RegularPolygon with 4 sides rotated 45°
    p = RegularPolygon((cx, cy), 4, radius=r*1.1,
                       orientation=np.pi/4,
                       color=color, ec=edge, linewidth=1.5,
                       transform=ax.transAxes, zorder=zorder)
    ax.add_patch(p)

def square_node(cx, cy, w=0.038, color=TEAL, edge=TEAL_DARK, zorder=10):
    rect = FancyBboxPatch((cx - w/2, cy - w/2), w, w,
                          boxstyle="round,pad=0.003",
                          facecolor=color, edgecolor=edge, linewidth=1.5,
                          transform=ax.transAxes, zorder=zorder)
    ax.add_patch(rect)

def line(x0, y0, x1, y1, color=EDGE_GRAY, lw=1.5, zorder=5):
    ax.plot([x0, x1], [y0, y1], color=color, lw=lw,
            transform=ax.transAxes, zorder=zorder)

def arrow(x0, y0, x1, y1, color=ARROW_BLUE, lw=1.8, zorder=6):
    ax.annotate('', xy=(x1, y1), xytext=(x0, y0),
                xycoords='axes fraction', textcoords='axes fraction',
                arrowprops=dict(arrowstyle='->', color=color,
                                lw=lw, mutation_scale=14))

def bullet_text(x, y, items, size=13.5, spacing=0.038):
    for i, item in enumerate(items):
        ax.text(x, y - i*spacing, f'\u2022  {item}',
                transform=ax.transAxes,
                fontsize=size, color=BULLET_TEXT,
                va='top', ha='left',
                fontfamily='sans-serif')

# ─────────────────────────────────────────────────────────────────────────────
# TITLE
# ─────────────────────────────────────────────────────────────────────────────
ax.text(0.035, 0.945, 'Inductive Bias for Deep Learning Models',
        transform=ax.transAxes,
        fontsize=32, fontweight='bold', color=BLUE_TITLE,
        va='top', ha='left', fontfamily='sans-serif')

# ─────────────────────────────────────────────────────────────────────────────
# ══ TOP-LEFT: Convolutional Networks ══
# ─────────────────────────────────────────────────────────────────────────────
# Two columns of 5 nodes, orange node on right col row 3
# Left col nodes
lx, rx = 0.075, 0.165
ys = [0.77, 0.70, 0.63, 0.56, 0.49]

for y in ys:
    node(lx, y)
for i, y in enumerate(ys):
    color = ORANGE if i == 2 else TEAL
    edge  = ORANGE_EDGE if i == 2 else TEAL_DARK
    node(rx, y, color=color, edge=edge)

# Lines from right col row3 (orange) to all left col nodes - orange lines
for y in ys:
    line(lx, y, rx, ys[2], color=EDGE_ORANGE, lw=1.8)

# Light lines between same-level left-right
for i, y in enumerate(ys):
    if i != 2:
        line(lx, y, rx, y, color=EDGE_GRAY, lw=1.0)

# Label
ax.text(0.20, 0.785, 'Convolutional Networks',
        transform=ax.transAxes,
        fontsize=16, fontweight='bold', color=DARK_TEXT,
        va='top', ha='left')
ax.text(0.20, 0.748, '(e.g. computer vision)',
        transform=ax.transAxes,
        fontsize=15, color=DARK_TEXT, va='top', ha='left', style='italic')

bullet_text(0.20, 0.700,
    ['data in regular grid',
     'information flow to local neighbours',
     'AlphaFold 1'],
    size=13.5)

# ─────────────────────────────────────────────────────────────────────────────
# ══ TOP-RIGHT: Recurrent Networks ══
# ─────────────────────────────────────────────────────────────────────────────
# 4 square hidden nodes in a row (middle row), 4 teal circle inputs below,
# 4 teal circle outputs above, 1 orange output above col 3
sx0 = 0.535
sq_xs = [sx0 + i*0.048 for i in range(4)]
sq_y  = 0.63

# Input nodes (bottom row)
in_y = 0.49
out_y = 0.77

for i, x in enumerate(sq_xs):
    # Square hidden node (teal/dark)
    color = TEAL if i != 2 else '#8b9b6b'  # slightly different for active
    square_node(x, sq_y, color='#4a7080', edge='#2a5060')
    # Input circle
    node(x, in_y)
    # Output circle (orange for col 3)
    if i == 2:
        node(x, out_y, color=ORANGE, edge=ORANGE_EDGE)
    else:
        node(x, out_y)
    # Vertical arrows up from input to square
    arrow(x, in_y + 0.018, x, sq_y - 0.022, color=ARROW_BLUE)
    # Vertical arrows up from square to output
    oc = ARROW_ORANGE if i == 2 else ARROW_BLUE
    arrow(x, sq_y + 0.022, x, out_y - 0.018, color=oc)

# Horizontal arrows between squares
for i in range(3):
    arrow(sq_xs[i] + 0.022, sq_y, sq_xs[i+1] - 0.022, sq_y,
          color=ARROW_BLUE)

# Label
ax.text(0.685, 0.785, 'Recurrent Networks',
        transform=ax.transAxes,
        fontsize=16, fontweight='bold', color=DARK_TEXT,
        va='top', ha='left')
ax.text(0.685, 0.748, '(e.g. language)',
        transform=ax.transAxes,
        fontsize=15, color=DARK_TEXT, va='top', ha='left', style='italic')

bullet_text(0.685, 0.700,
    ['data in ordered sequence',
     'information flow sequentially'],
    size=13.5)

# ─────────────────────────────────────────────────────────────────────────────
# ══ BOTTOM-LEFT: Graph Networks ══
# ─────────────────────────────────────────────────────────────────────────────
# Irregular graph with orange node in middle-bottom
gx, gy = 0.075, 0.26   # orange hub
# Surrounding nodes
gnodes = [
    (0.058, 0.42), (0.14, 0.44),   # top
    (0.065, 0.33),                  # mid-left
    (0.155, 0.35),                  # mid-right
    (0.04,  0.22), (0.11, 0.20),    # bottom
    (0.17,  0.22),                  # bottom-right
]
# Draw edges first (gray)
edge_pairs = [(0,1),(0,2),(1,3),(2,3),(2,4),(3,6),(4,5),(5,6)]
for a, b in edge_pairs:
    line(gnodes[a][0], gnodes[a][1], gnodes[b][0], gnodes[b][1],
         color=EDGE_GRAY, lw=1.5)
# Orange edges from hub
for nx, ny in gnodes[2:]:
    line(gx, gy, nx, ny, color=EDGE_ORANGE, lw=2.2)
# Also connect some top nodes to hub
line(gnodes[0][0], gnodes[0][1], gx, gy, color=EDGE_ORANGE, lw=2.2)
line(gnodes[1][0], gnodes[1][1], gx, gy, color=EDGE_ORANGE, lw=2.2)

# Draw nodes
for i, (nx, ny) in enumerate(gnodes):
    node(nx, ny)
node(gx, gy, color=ORANGE, edge=ORANGE_EDGE)  # hub

# Label
ax.text(0.20, 0.455, 'Graph Networks (e.g. recommender',
        transform=ax.transAxes,
        fontsize=16, fontweight='bold', color=DARK_TEXT,
        va='top', ha='left')
ax.text(0.20, 0.415, 'systems or molecules)',
        transform=ax.transAxes,
        fontsize=16, fontweight='bold', color=DARK_TEXT,
        va='top', ha='left')

bullet_text(0.20, 0.368,
    ['data in fixed graph structure',
     'information flow along fixed edges'],
    size=13.5)

# ─────────────────────────────────────────────────────────────────────────────
# ══ BOTTOM-RIGHT: Attention Module ══
# ─────────────────────────────────────────────────────────────────────────────
# Left: 5 rows of diamonds (input set)
# Right: 5 rows of diamonds (keys/queries)
# Orange diamond on right, orange hub node far right
dlx  = 0.535
drx  = 0.625
d_ys = [0.44, 0.375, 0.31, 0.245, 0.18]
hub_x = 0.665
hub_y = 0.31  # same level as middle row

# Draw left diamonds
for y in d_ys:
    diamond(dlx, y)

# Draw right diamonds + orange for middle
for i, y in enumerate(d_ys):
    if i == 2:
        diamond(drx, y, color=ORANGE, edge=ORANGE_EDGE)
    else:
        diamond(drx, y)

# Orange hub node
node(hub_x, hub_y, r=0.020, color=ORANGE, edge=ORANGE_EDGE)

# Lines from hub to all right diamonds (orange)
for y in d_ys:
    lc = EDGE_ORANGE
    lw = 2.0 if y == d_ys[2] else 1.6
    line(drx, y, hub_x, hub_y, color=lc, lw=lw)

# Lighter lines connecting left to right diamonds (same row)
for y in d_ys:
    line(dlx, y, drx, y, color=EDGE_GRAY, lw=1.2)

# Label
ax.text(0.695, 0.455, 'Attention Module (e.g. language)',
        transform=ax.transAxes,
        fontsize=16, fontweight='bold', color=BLUE_TITLE,
        va='top', ha='left')

bullet_text(0.695, 0.400,
    ['data in unordered set',
     'information flow dynamically controlled',
     'by the network (via keys and queries)'],
    size=13.5, spacing=0.038)

# ─────────────────────────────────────────────────────────────────────────────
# Divider lines (horizontal center)
# ─────────────────────────────────────────────────────────────────────────────
ax.axhline(0.465, xmin=0.03, xmax=0.97,
           color='#cccccc', lw=1.0, transform=ax.transAxes, zorder=1)
ax.axvline(0.5, ymin=0.04, ymax=0.88,
           color='#cccccc', lw=1.0, transform=ax.transAxes, zorder=1)

# ─────────────────────────────────────────────────────────────────────────────
# Logo watermark (bottom-right) — simple "S" swirl placeholder
# ─────────────────────────────────────────────────────────────────────────────
logo_circle = Circle((0.955, 0.055), 0.028, color='#1a5fbf',
                      transform=ax.transAxes, zorder=20)
ax.add_patch(logo_circle)
ax.text(0.955, 0.055, 'S', transform=ax.transAxes,
        fontsize=16, fontweight='bold', color='white',
        ha='center', va='center', zorder=21)

ax.text(0.965, 0.022, '17', transform=ax.transAxes,
        fontsize=14, color='#555555',
        ha='right', va='bottom', zorder=21)

# ─────────────────────────────────────────────────────────────────────────────
# Save
# ─────────────────────────────────────────────────────────────────────────────
out = '/home/daytona/workspace/hires-diagram/inductive_bias_hires.png'
fig.savefig(out, dpi=200, bbox_inches='tight',
            facecolor=BG, format='png')
print(f"Saved: {out}")