#!/usr/bin/env python3
from PIL import Image, ImageEnhance, ImageFilter, ImageDraw
import numpy as np, random, math, pathlib, subprocess, os

SRC = pathlib.Path('/home/ubuntu/.hermes/image_cache/img_184c12b449fc.jpg')
OUT_DIR = pathlib.Path('/home/ubuntu/lobbytracker/public/media/larreta-impact-frames')
OUT_DIR.mkdir(parents=True, exist_ok=True)
MP4 = pathlib.Path('/home/ubuntu/lobbytracker/public/media/larreta-impact-short.mp4')
GIF = pathlib.Path('/home/ubuntu/lobbytracker/public/media/larreta-impact-short.gif')

W,H = 720,1280
FPS = 15
DUR = 4.2
N = int(FPS*DUR)
random.seed(7)

src = Image.open(SRC).convert('RGB')
# crop/fit source poster into a tall shorts canvas
# attached image is 4:5; keep it huge, nearly full width, centered vertically with room for shake.
base_card_w = 690
base_card_h = int(base_card_w * src.height / src.width)
card_base = src.resize((base_card_w, base_card_h), Image.Resampling.LANCZOS)

# background: enlarged blurred copy of same poster, dark/red graded
bg_src = src.resize((W, int(W*src.height/src.width)), Image.Resampling.LANCZOS)
if bg_src.height < H:
    bg_src = src.resize((int(H*src.width/src.height), H), Image.Resampling.LANCZOS)
left = (bg_src.width-W)//2
upper = (bg_src.height-H)//2
bg_src = bg_src.crop((left, upper, left+W, upper+H)).filter(ImageFilter.GaussianBlur(16))

def ease_out_cubic(x): return 1 - (1-x)**3

def rgb_split(im, amt):
    if amt <= 0: return im
    arr = np.array(im)
    out = np.zeros_like(arr)
    a = int(amt)
    out[:,:,0] = np.roll(arr[:,:,0], a, axis=1)
    out[:,:,1] = arr[:,:,1]
    out[:,:,2] = np.roll(arr[:,:,2], -a, axis=1)
    return Image.fromarray(out)

def add_scanlines(im, alpha=28):
    overlay = Image.new('RGBA', im.size, (0,0,0,0))
    d = ImageDraw.Draw(overlay)
    for y in range(0, im.height, 4):
        d.line([(0,y),(im.width,y)], fill=(0,0,0,alpha), width=1)
    return Image.alpha_composite(im.convert('RGBA'), overlay)

def add_vignette(im):
    arr = np.array(im).astype(np.float32)
    yy,xx = np.mgrid[0:H,0:W]
    cx,cy = W/2,H/2
    dist = np.sqrt(((xx-cx)/(W/2))**2 + ((yy-cy)/(H/2))**2)
    mask = np.clip(1 - (dist-0.25)*0.72, 0.38, 1.0)
    arr[:,:,:3] *= mask[:,:,None]
    return Image.fromarray(np.clip(arr,0,255).astype(np.uint8), 'RGBA')

def red_grade(im, strength=1.0):
    arr = np.array(im).astype(np.float32)
    arr[:,:,0] *= 1.12 + 0.10*strength
    arr[:,:,1] *= 0.86
    arr[:,:,2] *= 0.82
    return Image.fromarray(np.clip(arr,0,255).astype(np.uint8))

beats = [0.55, 1.35, 2.45, 3.35]
flash_frames = {int(b*FPS) for b in beats} | {int(b*FPS)+1 for b in beats}
glitch_frames = set()
for b in beats:
    f=int(b*FPS)
    glitch_frames.update([f-1,f,f+1])

for i in range(N):
    t = i/(N-1)
    sec = i/FPS
    # slow push in with beat punches
    zoom = 1.0 + 0.10*t
    for b in beats:
        pulse = max(0, 1 - abs(sec-b)/0.13)
        zoom += 0.055 * pulse
    shake = 0
    for b in beats:
        if abs(sec-b) < 0.11: shake = 10
    shake += 2
    dx = random.randint(-shake, shake)
    dy = random.randint(-shake, shake)

    bg = bg_src.copy()
    bg = red_grade(bg, 1.0)
    bg = ImageEnhance.Contrast(bg).enhance(1.28)
    bg = ImageEnhance.Brightness(bg).enhance(0.55)
    canvas = Image.new('RGB', (W,H), (12,0,0))
    canvas.paste(bg, (0,0))

    # red radial rays / lightning-ish strokes behind the card
    overlay = Image.new('RGBA', (W,H), (0,0,0,0))
    d = ImageDraw.Draw(overlay)
    for k in range(10):
        ang = (k/10)*math.tau + 0.4*math.sin(sec*2)
        x2 = W/2 + math.cos(ang)*780
        y2 = H/2 + math.sin(ang)*780
        d.line([(W/2,H/2),(x2,y2)], fill=(255,0,0,22), width=8)
    for k in range(3):
        if (i+k*7) % 19 < 6:
            x=random.randint(50,W-50)
            points=[]
            y=0
            while y<H:
                points.append((x+random.randint(-25,25), y))
                y += random.randint(90,180)
            d.line(points, fill=(255,35,35,80), width=2)
    canvas = Image.alpha_composite(canvas.convert('RGBA'), overlay)

    cw = int(base_card_w*zoom)
    ch = int(base_card_h*zoom)
    card = card_base.resize((cw,ch), Image.Resampling.LANCZOS)
    card = ImageEnhance.Contrast(card).enhance(1.05)
    card = ImageEnhance.Sharpness(card).enhance(1.12)
    # enter from bottom in first second
    enter = ease_out_cubic(min(1, sec/0.85))
    x = (W-cw)//2 + dx
    y_final = 135 - int(65*t)
    y = int(H + 80 - (H + 80 - y_final) * enter) + dy

    # drop shadow / red glow
    shadow = Image.new('RGBA', (cw+50,ch+50), (0,0,0,0))
    sd = ImageDraw.Draw(shadow)
    sd.rounded_rectangle([20,20,cw+30,ch+30], radius=18, fill=(255,0,0,70))
    shadow = shadow.filter(ImageFilter.GaussianBlur(20))
    canvas.alpha_composite(shadow, (x-25,y-25))
    canvas.paste(card.convert('RGBA'), (x,y))

    frame = canvas
    # impact dark/red flashes
    if i in flash_frames:
        red = Image.new('RGBA', (W,H), (255,0,0,82))
        frame = Image.alpha_composite(frame, red)
    if i in {f+2 for f in flash_frames}:
        black = Image.new('RGBA', (W,H), (0,0,0,120))
        frame = Image.alpha_composite(frame, black)

    # horizontal glitch bars and RGB split
    if i in glitch_frames:
        frame = rgb_split(frame.convert('RGB'), random.randint(6,16)).convert('RGBA')
        arr = np.array(frame)
        for _ in range(8):
            yy=random.randint(0,H-28); hh=random.randint(8,28); off=random.randint(-45,45)
            arr[yy:yy+hh,:,:] = np.roll(arr[yy:yy+hh,:,:], off, axis=1)
        frame = Image.fromarray(arr, 'RGBA')

    # grain
    arr = np.array(frame).astype(np.int16)
    noise = np.random.normal(0, 7, arr.shape[:2]).astype(np.int16)
    arr[:,:,:3] = np.clip(arr[:,:,:3] + noise[:,:,None], 0, 255)
    frame = Image.fromarray(arr.astype(np.uint8), 'RGBA')
    frame = add_scanlines(frame, 22)
    frame = add_vignette(frame)

    frame.convert('RGB').save(OUT_DIR / f'frame_{i:03d}.png', quality=95)

# MP4
subprocess.run([
    'ffmpeg','-y','-framerate',str(FPS),'-i',str(OUT_DIR/'frame_%03d.png'),
    '-vf','format=yuv420p','-movflags','+faststart','-crf','18',str(MP4)
], check=True, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
# GIF smaller using ffmpeg palette
palette='/tmp/larreta_palette.png'
subprocess.run(['ffmpeg','-y','-i',str(MP4),'-vf','fps=12,scale=480:-1:flags=lanczos,palettegen',palette], check=True, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
subprocess.run(['ffmpeg','-y','-i',str(MP4),'-i',palette,'-filter_complex','fps=12,scale=480:-1:flags=lanczos[x];[x][1:v]paletteuse=dither=bayer:bayer_scale=3',str(GIF)], check=True, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
print(MP4)
print(GIF)