# -*- coding: utf-8 -*-
"""船と航空の動画 v2。最初から両方が動く。航空＝路線の太い順に弧を伸ばす（前と同じ）。
船＝さざなみの粒がじりじり動く（標本化の位置を、ゆっくり回る滑らかな乱れ場でずらす。表示の効果であって観測ではない）。"""
import sys, os, shutil, subprocess
import numpy as np
sys.path.insert(0, '/home/claude/transport'); sys.path.insert(0, '/home/claude/forest')
import combo as C, flights as FL, forest1 as F
from scipy import ndimage
from PIL import Image, ImageDraw
SS = 2; C.SS = SS
W, H = C.SIZE[0]*SS, C.SIZE[1]*SS
C.SHIP = ndimage.maximum_filter(C.SHIP, size=3, mode=('nearest','wrap'))
sc, lon, lat, ok, land = C.scene(0.0)
base = np.zeros((H, W, 3), np.float32); base[:] = C.NAVY; base[land] = C.LANDC
# 船の場（log 尺度）を一度だけ作り、毎コマ位置をずらして最近傍で拾う
v = np.where(C.SHIP >= 1000, C.SHIP, 0); v = np.clip(np.log1p(v)/np.log1p(10**7.5), 0, 1).astype(np.float32)
vv = np.roll(v, -v.shape[1]//2, axis=1); lon2 = (lon % 360) - 180
rng = np.random.default_rng(7)
def smooth_noise(sig=6.0):
    n = rng.standard_normal((H, W)).astype(np.float32); n = ndimage.gaussian_filter(n, sig); return n / n.std()
N1x, N2x, N1y, N2y = smooth_noise(), smooth_noise(), smooth_noise(), smooth_noise()
AMP = 0.035          # 度。粒が 1〜2 セル分だけ動く
PERIOD = 300         # コマ。10 秒で一巡（航空より遅い）
def ships(f):
    w = 2*np.pi*f/PERIOD
    dx = AMP*(N1x*np.cos(w) + N2x*np.sin(w)); dy = AMP*(N1y*np.cos(w) + N2y*np.sin(w))
    fr = F.sample(vv, lon2 + dx, lat + dy, ok); fr = np.where(land, 0, fr)
    t = np.where(fr > 0, 0.6 + 0.4*fr, 0); a = np.where(fr > 0, np.clip(fr*1.6, 0, 1), 0)
    return LUT[np.clip((t*1000).astype(int), 0, 1000)], a
LUT = C.ramp(np.linspace(0, 1, 1001)).astype(np.float32)     # 階調の表引き（毎コマ Lab 変換しない）
CLS = [(1,2,0.35,0.2),(2,4,0.5,0.3),(4,8,0.7,0.45),(8,999,1.0,0.6)]; T_CLS = [0.35, 0.42, 0.5, 0.6]
COLS = [C.ramp(np.array([t]))[0] for t in T_CLS]; WHITE = C.ramp(np.array([1.0]))[0]
def px(la, lo): return (lo*111320.0 - sc.cx)/sc.s + W/2, H/2 - (la*111320.0 - sc.cy)/sc.s
arcs = []
P = FL.pairs.sort_values('n', ascending=False).reset_index(drop=True)
for a, b, n in zip(P.a, P.b, P.n):
    g = FL.gc(FL.pos.at[a,'lat'], FL.pos.at[a,'lon'], FL.pos.at[b,'lat'], FL.pos.at[b,'lon'])
    if g is None: continue
    k = next(i for i, (lo_, hi_, _, _) in enumerate(CLS) if lo_ <= n < hi_)
    x, y = px(g[0], g[1]); arcs.append((k, list(zip(x, y)), CLS[k][2]))
NA = len(arcs)
HOLD0, N_AIR, HOLD1 = 30, 600, 90; GROW = 24
starts = np.linspace(0, N_AIR - GROW, NA)
canv = [Image.new('L', (W, H), 0) for _ in CLS]; drs = [ImageDraw.Draw(c) for c in canv]
flash = np.zeros((H, W), np.float32)
FD = '/home/claude/transport/video/frames2'; shutil.rmtree(FD, ignore_errors=True); os.makedirs(FD)
def compose(f):
    out = base.copy()
    for k in range(len(CLS)):
        a = (np.array(canv[k]).astype(np.float32)/255 * CLS[k][3])[..., None]; out = out*(1-a) + COLS[k]*a
    col, a = ships(f); a = a[..., None]; out = out*(1-a) + col*a
    a = (flash*0.9)[..., None]; out = out*(1-a) + WHITE*a
    return Image.fromarray(np.clip(out,0,255).astype(np.uint8)).resize(C.SIZE, Image.LANCZOS)
n = 0
def emit(im):
    global n
    im.save(f'{FD}/{n:04d}.png'); n += 1
for f in range(HOLD0): emit(compose(f))          # 船だけが動いている
prog = np.zeros(NA, int)
for f in range(N_AIR):
    tips = Image.new('L', (W, H), 0); dt = ImageDraw.Draw(tips)
    act = np.nonzero((starts <= f) & (starts + GROW > f))[0]
    for i in act:
        k, pts, wpx = arcs[i]; nseg = len(pts) - 1
        u = (f - starts[i] + 1)/GROW; e = 1 - (1-u)**2; target = min(nseg, int(round(e*nseg))); j0 = prog[i]
        if target > j0:
            seg = pts[j0:target+1]; run = [seg[0]]
            for p0, p1 in zip(seg[:-1], seg[1:]):
                if abs(p1[0]-p0[0]) > W/2:
                    if len(run) > 1: drs[k].line(run, fill=255, width=max(1, int(round(wpx*SS))))
                    run = [p1]
                else: run.append(p1)
            if len(run) > 1: drs[k].line(run, fill=255, width=max(1, int(round(wpx*SS))))
            prog[i] = target; tx, ty = pts[target]; r = 1.6*SS; dt.ellipse([tx-r, ty-r, tx+r, ty+r], fill=255)
    flash = np.maximum(flash*0.7, np.array(tips).astype(np.float32)/255)
    emit(compose(HOLD0 + f))
    if f % 100 == 0: print('A', f, len(act), flush=True)
for j in range(HOLD1):
    flash *= 0.7; emit(compose(HOLD0 + N_AIR + j))        # 航空は止まり、船は動き続ける
subprocess.run(['ffmpeg','-y','-loglevel','error','-framerate','30','-i',f'{FD}/%04d.png','-c:v','libx264','-pix_fmt','yuv420p','-crf','18',
                '-vf','pad=ceil(iw/2)*2:ceil(ih/2)*2:0:0:black','/home/claude/transport/video/transport_live.mp4'], check=True)
print('frames', n, n/30, 's')