# -*- coding: utf-8 -*-
"""地震の動画 v2。各円は震源から広がる（半径 0→r を K コマで ease-out、K = 4 + 3·(M−5)）。
広がり切ったら恒久の印（Q3d と同じ薄さ・不透明度）に沈み、先端の濃いインクは 0.72/コマで減衰。
月ごとに 1 コマ。終端は静止画に一致する。"""
import sys, os, subprocess, shutil, datetime
import numpy as np
sys.path.insert(0, '/home/claude/quakes'); sys.path.insert(0, '/home/claude/forest')
import q1 as Q
Q.SS = 2
import forest1 as F, sylvania as S
from nyt import solid_grat
from PIL import Image, ImageDraw
Q.INK = np.array([196, 34, 28], np.float32)
SS = Q.SS; W, H = Q.SIZE[0]*SS, Q.SIZE[1]*SS
sc, lon, lat, ok, land = Q.scene()
BG = sys.argv[1] if len(sys.argv) > 1 else 'plain'
Y0, Y1 = (int(sys.argv[2]), int(sys.argv[3])) if len(sys.argv) > 3 else (1973, 2025)
OUTMP4 = sys.argv[4] if len(sys.argv) > 4 else f'/home/claude/quakes/video/quakes_{BG}_{Y0}_{Y1}.mp4'
# 地
GPARAMS = {'G1': dict(ocean_amp=0.10), 'G2': dict(ocean_amp=0.18), 'G3': dict(ocean_amp=0.18, ocean_depth=6.0), 'G4': dict(ocean_amp=0.28, vex_sea=14.0), 'F5': dict()}
if BG in GPARAMS:
    import relief
    base = relief.f5_base(sc, lon, lat, ok, land, **GPARAMS[BG])   # 経緯線なし
else:
    base = np.zeros((H, W, 3), np.float32); base[:] = Q.PAPER; base[land] = Q.LANDC
    grat = F.aa_mask((W, H), solid_grat(sc)); F.NAVY = Q.INK; F.WHITE = Q.PAPER
    base = np.array(F.comp_rel(Image.fromarray(base.astype(np.uint8)), grat, 0.14, 0.10)).astype(np.float32)
# 事象
x = (Q.LON*111320.0 - sc.cx)/sc.s + W/2; y = H/2 - (Q.LAT*111320.0 - sc.cy)/sc.s
r = 0.9 * 2.4 ** (Q.MAG - 5.0) * SS
AMAX = 0.55; al = np.where(Q.MAG < 6, 0.30, np.where(Q.MAG < 7, 0.40, AMAX))
K = (4 + 3*(Q.MAG - 5.0)).astype(int)
edges = [0, 35, 70, 150, 300, 800]; depth_L = [0.0, 0.15, 0.35, 0.55, 0.72]
a_ink, b_pap = S.to_lab(Q.INK), S.to_lab(Q.PAPER)
cols = [S.from_lab(a_ink + (b_pap - a_ink) * t).astype(np.float32) for t in depth_L]
dcls = np.digitize(Q.DEP, edges[1:-1])
dt = [datetime.datetime.utcfromtimestamp(int(s)) for s in Q.T]
mon = np.array([(d.year - Y0) * 12 + d.month - 1 for d in dt])
NM = (Y1 - Y0 + 1) * 12
perm = [np.zeros((H, W), np.float32) for _ in range(5)]   # 恒久（fill = al/AMAX）
flash = np.zeros((H, W), np.float32)
FD = f'/home/claude/quakes/video/frames_{BG}_{Y0}_{Y1}'; shutil.rmtree(FD, ignore_errors=True); os.makedirs(FD)
def compose():
    out = base.copy()
    for k in range(4, -1, -1):
        a = (perm[k] * AMAX)[..., None]; out = out * (1 - a) + cols[k] * a
    a = (flash * 0.85)[..., None]; out = out * (1 - a) + Q.INK * a
    return Image.fromarray(np.clip(out, 0, 255).astype(np.uint8)).resize(Q.SIZE, Image.LANCZOS)
n = 0
def emit(im, k=1):
    global n
    p = f'{FD}/{n:04d}.png'; im.save(p); n += 1
    for _ in range(k - 1): os.link(p, f'{FD}/{n:04d}.png'); n += 1
emit(compose(), 60)
active = []   # (i, start_month)
order = np.argsort(mon)
for m in range(NM):
    idx = np.nonzero(mon == m)[0]
    for i in idx: active.append((i, m))
    cv = Image.new('L', (W, H), 0); dr = ImageDraw.Draw(cv)
    layer = [None]*5; still = []
    for i, m0 in active:
        t = (m - m0 + 1) / K[i]
        if t >= 1.0:
            k = dcls[i]
            if layer[k] is None: layer[k] = (Image.new('L', (W, H), 0)); 
            ImageDraw.Draw(layer[k]).ellipse([x[i]-r[i], y[i]-r[i], x[i]+r[i], y[i]+r[i]], fill=int(255*al[i]/AMAX))
            continue
        e = 1 - (1 - t) ** 2; rr = r[i] * e
        dr.ellipse([x[i]-rr, y[i]-rr, x[i]+rr, y[i]+rr], fill=255); still.append((i, m0))
    for k in range(5):
        if layer[k] is not None: perm[k] = np.maximum(perm[k], np.array(layer[k]).astype(np.float32)/255)
    active = still
    flash = np.maximum(flash * 0.72, np.array(cv).astype(np.float32)/255)
    emit(compose())
    if m % 60 == 0: print('M', m, n, len(active), flush=True)
# 残った広がりを畳む
for _ in range(20):
    layer = [None]*5
    for i, m0 in active:
        k = dcls[i]
        if layer[k] is None: layer[k] = Image.new('L', (W, H), 0)
        ImageDraw.Draw(layer[k]).ellipse([x[i]-r[i], y[i]-r[i], x[i]+r[i], y[i]+r[i]], fill=int(255*al[i]/AMAX))
    for k in range(5):
        if layer[k] is not None: perm[k] = np.maximum(perm[k], np.array(layer[k]).astype(np.float32)/255)
    active = []; flash *= 0.72; emit(compose())
    if flash.max() < 0.02: break
flash[:] = 0; last = compose(); emit(last, 90); last.save(f'/home/claude/quakes/out/end_{BG}_{Y0}_{Y1}.png')
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:white', OUTMP4], check=True)
print('frames', n, n/30, 's', OUTMP4)