"""render a mosaic video.  python3 render.py stills NAME t1,t2,...  |  python3 render.py part NAME i n  |  python3 render.py join NAME n"""
import sys, os, subprocess, numpy as np
from PIL import Image
from vis_common import *
import engines
VID = ROOT + 'mosaic/video/'; os.makedirs(VID, exist_ok=True)
def make(name):
    c = Ctx(name); return c, engines.ENGINES[name](c)
def stills(name, ts):
    c, e = make(name); got = []; dt = 1 / FPS
    for T in sorted(ts):
        e = engines.ENGINES[name](c)
        for f in range(max(0, int((T - 2.5) * FPS)), int(T * FPS) + 1): img = e.frame(f / FPS, dt)
        got.append(finish(img, T, c.dur))
    ims = [Image.fromarray((g * 255).astype(np.uint8)).resize((640, 360)) for g in got]
    S = Image.new('RGB', (1920, 360 * ((len(ims) + 2) // 3)))
    for k, im in enumerate(ims): S.paste(im, ((k % 3) * 640, (k // 3) * 360))
    S.save(f'/tmp/claude-0/st_{name}.jpg')
def part(name, i, n, warm=4.0):
    c, e = make(name); nf = int(c.dur * FPS); a = nf * i // n; b = nf * (i + 1) // n
    out = VID + f'{name}_p{i}.mp4'; tmp = out.replace('.mp4', '.tmp.mp4')
    p = subprocess.Popen(['ffmpeg', '-v', 'error', '-y', '-f', 'rawvideo', '-pix_fmt', 'rgb24', '-s', f'{W}x{H}', '-r', str(FPS), '-i', '-',
                          '-c:v', 'libx264', '-preset', 'veryfast', '-crf', '18', '-pix_fmt', 'yuv420p', tmp], stdin=subprocess.PIPE)
    start = max(0, a - int(warm * FPS)); dt = 1 / FPS
    for f in range(start, b):
        t = f / FPS; img = e.frame(t, dt)
        if f >= a: p.stdin.write((finish(img, t, c.dur) * 255).astype(np.uint8).tobytes())
        if f % 240 == 0: print(name, i, round(t, 1), flush=True)
    p.stdin.close(); p.wait()
    if p.returncode == 0: os.replace(tmp, out)
def join(name, n):
    lst = VID + f'{name}_list.txt'; open(lst, 'w').write(''.join(f"file '{VID}{name}_p{i}.mp4'\n" for i in range(n)))
    subprocess.run(['ffmpeg', '-v', 'error', '-y', '-f', 'concat', '-safe', '0', '-i', lst, '-c', 'copy', VID + f'{name}_v.mp4'], check=True)
    subprocess.run(['ffmpeg', '-v', 'error', '-y', '-i', VID + f'{name}_v.mp4', '-i', (AUD + name + '_sd.wav') if os.path.exists(AUD + name + '_sd.wav') else (AUD + name + '.wav'), '-map', '0:v', '-map', '1:a', '-c:v', 'copy',
                    '-af', 'loudnorm=I=-14:TP=-1.0', '-c:a', 'aac', '-b:a', '192k', '-shortest', '-movflags', '+faststart', VID + f'afterlight-{name}.mp4'], check=True)
if __name__ == '__main__':
    m = sys.argv[1]
    if m == 'stills': stills(sys.argv[2], [float(x) for x in sys.argv[3].split(',')])
    elif m == 'part': part(sys.argv[2], int(sys.argv[3]), int(sys.argv[4]))
    elif m == 'join': join(sys.argv[2], int(sys.argv[3]))