# -*- coding: utf-8 -*-
"""エアロゾルの動画。周期の時間: 12 か月の気候値（各暦月の 2006–2011 平均）を一巡してループ。
細かい粒（密度＝AOD、大きさ＝由来）が、月ごとに増減しながらじりじり動く。各月で 4 コマ止まり、残りは値の補間。"""
import sys, os, shutil, subprocess, glob
import numpy as np, netCDF4 as nc
sys.path.insert(0, '/home/claude/aerosol'); sys.path.insert(0, '/home/claude/forest'); sys.path.insert(0, '/home/claude/quakes'); sys.path.insert(0, '/home/claude/stripes')
import aer as A, forest1 as F, sylvania as S, q1 as Q, temp2 as T
from PIL import Image, ImageDraw
from scipy import ndimage
SS = 2; W, H = A.SIZE[0]*SS, A.SIZE[1]*SS
sc = F.Scene(A.BBOX, (W, H), proj='pc', pad=0.0); lon, lat, ok = sc.lonlat_grid()
land = F.sample(Q.land_grid_all().astype(np.float32), lon, lat, ok) > 0.5
coast = (land ^ ndimage.binary_erosion(land, iterations=SS)).astype(np.float32)
base = np.zeros((H, W, 3), np.float32); base[:] = A.PAPER; base[land] = A.LANDC
F.NAVY = A.INK; F.WHITE = A.PAPER
base = np.array(F.comp_rel(Image.fromarray(base.astype(np.uint8)), coast, 0.22, 0.16)).astype(np.float32)
# 月ごとの気候値（AOD と Ångström 指数）
fs = sorted(glob.glob('/home/claude/aerosol/data/DBAER_*.nc'))
t5 = np.array([np.ma.filled(nc.Dataset(f).variables['aerosol_optical_thickness_550_land_ocean_mean'][:], np.nan) for f in fs])[:, ::-1]
t6 = np.array([np.ma.filled(nc.Dataset(f).variables['aerosol_optical_thickness_630_land_ocean_mean'][:], np.nan) for f in fs])[:, ::-1]
cnt = A.CNT
FR, AL = [], []
for m in range(12):
    idx = [i for i in range(len(fs)) if int(A.MONTHS[i][-2:]) == m+1]
    g = A.field(idx); v, conf = T.fill_field(g, up=10, sigma_deg=5.0)
    fr = F.sample_bilinear(np.nan_to_num(v), lon, lat, ok); cf = F.sample_bilinear(conf, lon, lat, ok)
    FR.append(np.where(cf > 0.12, fr, 0).astype(np.float32))
    a5 = np.nanmean(t5[idx], 0); a6 = np.nanmean(t6[idx], 0)
    al = -np.log(np.maximum(a6, 1e-3)/np.maximum(a5, 1e-3))/np.log(630/550)
    av, _ = T.fill_field(np.ma.masked_invalid(al), up=10, sigma_deg=5.0)
    AL.append(F.sample_bilinear(np.nan_to_num(av, nan=1.0), lon, lat, ok).astype(np.float32))
    print('month', m+1, flush=True)
# 粒の候補点（格子＋ずらし）と、点ごとの閾値
rng = np.random.default_rng(11); STEP = 2
ys, xs = np.mgrid[0:H:STEP, 0:W:STEP].astype(np.float32)
ys = (ys + rng.uniform(-1, 1, ys.shape)).ravel(); xs = (xs + rng.uniform(-1, 1, xs.shape)).ravel()
U = rng.uniform(0, 1, xs.shape).astype(np.float32)
# ゆっくり回る滑らかな乱れ場（粒の動き。表示の効果）
def smooth_noise(sig=8.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 = 2.2*SS; PERIOD = 360
GAMMA, GAIN = 1.6, 0.75
FPM, HOLD = 30, 4
LOOPS = 2
FD = '/home/claude/aerosol/video/frames'; shutil.rmtree(FD, ignore_errors=True); os.makedirs(FD)
def frame(f, fr, al):
    w = 2*np.pi*f/PERIOD
    yi = np.clip(ys.astype(int), 0, H-1); xi = np.clip(xs.astype(int), 0, W-1)
    dx = AMP*(N1x[yi, xi]*np.cos(w) + N2x[yi, xi]*np.sin(w)); dy = AMP*(N1y[yi, xi]*np.cos(w) + N2y[yi, xi]*np.sin(w))
    px = xs + dx; py = ys + dy
    pi = np.clip(py.astype(int), 0, H-1); pj = np.clip(px.astype(int), 0, W-1)
    p = np.clip(fr[pi, pj]/A.LIM, 0, 1)**GAMMA * GAIN; keep = U < p
    r = (0.35 + 0.75*np.clip((1.6 - al[pi, pj])/1.4, 0, 1)) * SS/2      # 細かい粒: 0.35〜1.1 px（出力）
    cv = Image.new('L', (W, H), 0); d = ImageDraw.Draw(cv)
    for x, y, rr in zip(px[keep], py[keep], r[keep]):
        if rr < 0.55: d.point((x, y), fill=200)
        else: d.ellipse([x-rr, y-rr, x+rr, y+rr], fill=255)
    a = (np.array(cv).astype(np.float32)/255*0.85)[..., None]
    out = base*(1-a) + A.INK*a
    return Image.fromarray(np.clip(out,0,255).astype(np.uint8)).resize(A.SIZE, Image.LANCZOS)
n = 0
for loop in range(LOOPS):
    for m in range(12):
        f0, f1 = FR[m], FR[(m+1) % 12]; a0, a1 = AL[m], AL[(m+1) % 12]
        for k in range(FPM):
            if k < HOLD: t = 0.0
            else:
                u = (k - HOLD + 1)/(FPM - HOLD); t = 3*u*u - 2*u*u*u
            im = frame(n, f0*(1-t) + f1*t, a0*(1-t) + a1*t); im.save(f'{FD}/{n:04d}.png'); n += 1
        print('M', loop, m+1, n, flush=True)
subprocess.run(['ffmpeg','-y','-loglevel','error','-framerate','30','-i',f'{FD}/%04d.png','-c:v','libx264','-pix_fmt','yuv420p','-crf','20',
                '-vf','pad=ceil(iw/2)*2:ceil(ih/2)*2:0:0:white','/home/claude/aerosol/video/aerosol_seasons.mp4'], check=True)
print('frames', n, n/30, 's')