import numpy as np
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import matplotlib.patches as mpatches
from matplotlib.patches import FancyBboxPatch
import matplotlib.patheffects as pe
from matplotlib.colors import LinearSegmentedColormap
import subprocess, os, shutil

# ── Data ──────────────────────────────────────────────────────────────────────
casp_data = [
    ("CASP 1",  "1994", 39),
    ("CASP 2",  "1996", 47),
    ("CASP 3",  "1998", 49),
    ("CASP 4",  "2000", 45),
    ("CASP 5",  "2002", 57),
    ("CASP 6",  "2004", 60),
    ("CASP 7",  "2006", 65),
    ("CASP 8",  "2008", 66),
    ("CASP 9",  "2010", 62),
    ("CASP 10", "2012", 55),
    ("CASP 11", "2014", 51),
    ("CASP 12", "2016", 56),
    ("CASP 13", "2018", 74),
    ("CASP 14", "2020", 90),
]

# ── New colour theme  (dark teal bg, cyan→orange bars, purple highlight) ──────
BG_COLOR      = "#0a1628"          # deep navy-teal background
CHART_BG      = "#0d1f35"          # slightly lighter chart area
AXIS_COLOR    = "#4dd9c0"          # cyan-teal axis lines / tick labels
TEXT_COLOR    = "#e8f4f8"          # near-white text
YLABEL_COLOR  = "#a0d8ef"          # soft blue for y-label
HIGHLIGHT_BG  = "#2d0a4e"          # deep purple zone above 90
HIGHLIGHT_FG  = "#cc44ff"          # bright purple line at 90
BAR_GRAD_BOT  = "#ff6a00"          # orange bottom of bar gradient
BAR_GRAD_TOP  = "#00d4ff"          # cyan top
GLOW_COLOR    = "#ff6a0088"        # semi-transparent orange glow
LOGO_COLOR    = "#4dd9c0"

OUT_DIR   = "/home/daytona/workspace/video-recolor/render_frames"
OUT_VIDEO = "/home/daytona/workspace/video-recolor/casp_recolored.mp4"
FPS       = 30
W, H      = 1914, 1080

os.makedirs(OUT_DIR, exist_ok=True)

# ── Animation plan ─────────────────────────────────────────────────────────────
# Phase 1 (0-3 s): bars 0..12 all grow simultaneously, CASP14 hidden
# Phase 2 (3-8 s): CASP14 bar shoots up from 0→90, glow/highlight appears
TOTAL_FRAMES = int(8.475 * FPS)   # match original duration

P1_START = 0
P1_END   = int(3.0 * FPS)
P2_START = P1_END
P2_END   = TOTAL_FRAMES

def ease_in_out(t):
    """Smooth cubic easing."""
    return t * t * (3 - 2 * t)

def draw_gradient_bar(ax, x, width, height, color_bot, color_top, alpha=1.0, glow=False):
    """Draw a bar with vertical gradient using imshow."""
    if height <= 0:
        return
    n = max(2, int(height * 5))
    grad = np.linspace(0, 1, n).reshape(n, 1)
    # colour blend bot→top
    rb, gb, bb = int(color_bot[1:3],16)/255, int(color_bot[3:5],16)/255, int(color_bot[5:7],16)/255
    rt, gt, bt = int(color_top[1:3],16)/255, int(color_top[3:5],16)/255, int(color_top[5:7],16)/255
    rgba = np.zeros((n, 1, 4))
    rgba[:,:,0] = rb + (rt-rb)*grad
    rgba[:,:,1] = gb + (gt-gb)*grad
    rgba[:,:,2] = bb + (bt-bb)*grad
    rgba[:,:,3] = alpha
    ax.imshow(rgba, extent=[x - width/2, x + width/2, 0, height],
              aspect='auto', origin='lower', zorder=3)
    if glow:
        # orange glow rectangle behind bar
        for pad in [0.12, 0.25, 0.4]:
            glow_rect = plt.Rectangle(
                (x - width/2 - pad, -0.5), width + pad*2, height + 1,
                color=BAR_GRAD_BOT, alpha=0.07, zorder=2, linewidth=0)
            ax.add_patch(glow_rect)

def render_frame(frame_idx):
    fig = plt.figure(figsize=(W/100, H/100), dpi=100, facecolor=BG_COLOR)
    ax = fig.add_axes([0.10, 0.15, 0.87, 0.72], facecolor=CHART_BG)

    ax.set_xlim(-0.5, len(casp_data) - 0.5)
    ax.set_ylim(0, 105)

    # ── purple highlight zone (90+) ────────────────────────────────────────────
    highlight_rect = plt.Rectangle((-0.5, 90), len(casp_data), 15,
                                   color=HIGHLIGHT_BG, zorder=1)
    ax.add_patch(highlight_rect)
    ax.axhline(90, color=HIGHLIGHT_FG, linewidth=1.5, linestyle='--',
               alpha=0.85, zorder=4)

    # ── grid lines ────────────────────────────────────────────────────────────
    for y in range(0, 101, 10):
        ax.axhline(y, color=AXIS_COLOR, linewidth=0.3, alpha=0.2, zorder=1)

    # ── Compute bar heights for this frame ────────────────────────────────────
    t_p1 = np.clip((frame_idx - P1_START) / (P1_END - P1_START), 0, 1)
    progress_p1 = ease_in_out(t_p1)

    t_p2 = np.clip((frame_idx - P2_START) / (P2_END - P2_START), 0, 1)
    progress_p2 = ease_in_out(t_p2)

    bar_width = 0.6

    for i, (name, year, val) in enumerate(casp_data):
        if i < 13:
            h = val * progress_p1
            alpha = min(1.0, progress_p1 * 1.5)
            draw_gradient_bar(ax, i, bar_width, h, BAR_GRAD_BOT, BAR_GRAD_TOP,
                              alpha=alpha, glow=False)
        else:
            # CASP 14 — phase 2
            h = val * progress_p2
            alpha = min(1.0, progress_p2 * 2.0)
            glow = (progress_p2 > 0.05)
            draw_gradient_bar(ax, i, bar_width, h, BAR_GRAD_BOT, BAR_GRAD_TOP,
                              alpha=alpha, glow=glow)

    # ── x-axis labels ─────────────────────────────────────────────────────────
    for i, (name, year, val) in enumerate(casp_data):
        vis = progress_p1 if i < 13 else progress_p2
        if vis < 0.02:
            continue
        a = min(1.0, vis * 2)
        ax.text(i, -2, name, ha='center', va='top', fontsize=7.5,
                color=AXIS_COLOR, alpha=a * 0.7, transform=ax.transData,
                fontfamily='monospace')
        ax.text(i, -5.5, year, ha='center', va='top',
                fontsize=10, fontweight='bold',
                color=TEXT_COLOR, alpha=a,
                transform=ax.transData)

    # ── y-axis ────────────────────────────────────────────────────────────────
    ax.set_yticks(range(0, 101, 10))
    ax.tick_params(axis='y', colors=AXIS_COLOR, labelsize=9)
    ax.tick_params(axis='x', bottom=False, labelbottom=False)
    for spine in ax.spines.values():
        spine.set_edgecolor(AXIS_COLOR)
        spine.set_linewidth(0.8)
        spine.set_alpha(0.4)
    ax.spines['top'].set_visible(False)
    ax.spines['right'].set_visible(False)

    # ── y-axis label (rotated) ─────────────────────────────────────────────────
    ax.set_ylabel("Global Distance Test", color=YLABEL_COLOR,
                  fontsize=11, labelpad=10)

    # ── "Ve" logo watermark (bottom-right) ────────────────────────────────────
    ax.text(0.995, 0.012, "Ve", transform=ax.transAxes,
            ha='right', va='bottom', fontsize=7, color=LOGO_COLOR, alpha=0.5,
            bbox=dict(boxstyle='square,pad=0.3', facecolor='#1a2a3a',
                      edgecolor=LOGO_COLOR, linewidth=0.5, alpha=0.6))

    # ── glow on CASP 14 bar highlight text ─────────────────────────────────────
    if progress_p2 > 0.8:
        glow_alpha = (progress_p2 - 0.8) / 0.2
        ax.text(13, 92, "AlphaFold", ha='center', va='bottom',
                fontsize=9, color='#cc44ff', alpha=glow_alpha * 0.9,
                fontweight='bold')

    plt.tight_layout(pad=0)
    path = os.path.join(OUT_DIR, f"frame_{frame_idx:05d}.jpg")
    plt.savefig(path, dpi=100, bbox_inches='tight',
                facecolor=BG_COLOR, format='jpg', quality=92)
    plt.close(fig)

# ── Render all frames ──────────────────────────────────────────────────────────
print(f"Rendering {TOTAL_FRAMES} frames...")
for f in range(TOTAL_FRAMES):
    render_frame(f)
    if f % 30 == 0:
        print(f"  frame {f}/{TOTAL_FRAMES}")

print("Encoding video...")
cmd = [
    "ffmpeg", "-y",
    "-framerate", str(FPS),
    "-i", os.path.join(OUT_DIR, "frame_%05d.jpg"),
    "-c:v", "libx264",
    "-preset", "fast",
    "-crf", "18",
    "-pix_fmt", "yuv420p",
    "-vf", f"scale={W}:{H}:flags=lanczos",
    OUT_VIDEO
]
result = subprocess.run(cmd, capture_output=True, text=True)
print(result.stdout[-500:] if result.stdout else "")
print(result.stderr[-500:] if result.stderr else "")
print("Done!", OUT_VIDEO)