# -*- coding: utf-8 -*-
"""C: 大河が支流によってつくられていく。

A（海までの距離の順）と B（流量の階級で太さを決める）を重ねる。
描く順序は海までの距離の大きい順なので、ある川の上流の細い区間が先に現れ、
その下流の太い区間は必ず後から現れる。つまり本流は、支流が出そろってから太くなる。
伸びていく先端だけを一瞬濃く描き、0.72 倍ずつ減衰させて「育つ向き」を見せる。
"""
import os, sys, time
import numpy as np
sys.path.insert(0, '/home/claude/forest')
import forest1 as F
from PIL import Image, ImageDraw
OUT = '/home/claude/rivers/video/C'; os.makedirs(f'{OUT}/frames', exist_ok=True)
PAPER = np.array([255, 255, 255], np.float32)
INK = np.array([86, 120, 196], np.float32)
sys.path.insert(0, '/home/claude/forest')
import sylvania as S   # Lab 変換のみ利用
FRONT = S.from_lab(S.to_lab(INK) + np.array([-16.0, 0, 0])).astype(np.float32)   # 先端は一段濃い青
Z = np.load('/home/claude/rivers/rivers.npz')
PTS, CNT, CLS = Z['pts'], Z['cnt'], Z['cls']
DIST = np.load('/home/claude/rivers/dist_dn.npy')
OFF = np.concatenate([[0], np.cumsum(CNT)])
BBOX = (-180, -58, 180, 83); SIZE = (1350, 1080); SS = 3
GROUPS = [[1, 2], [3, 4], [5], [6]]; WIDTHS = [1.0, 0.7, 0.45, 0.3]; ALPHAS = [1.0, 1.0, 0.95, 0.75]
sc = F.Scene(BBOX, (SIZE[0]*SS, SIZE[1]*SS), proj='pc'); BW, BH = sc.W, sc.H
IDX = np.where(np.isin(CLS, [1, 2, 3, 4, 5, 6]))[0]
GI = np.full(len(CLS), -1, np.int8)
for gi, g in enumerate(GROUPS): GI[np.isin(CLS, g)] = gi

def draw(ds, ids, wscale=1.0):
    for i in ids:
        g = GI[i]
        if g < 0: continue
        xy = PTS[OFF[i]:OFF[i+1]]
        x = (xy[:, 0]*111320.0 - sc.cx) / sc.s + BW/2
        y = BH/2 - (xy[:, 1]*111320.0 - sc.cy) / sc.s
        if x.max() - x.min() > BW*0.5: continue
        ds[g].line(list(zip(x.tolist(), y.tolist())), fill=255, width=max(1, int(round(WIDTHS[g]*SS*wscale))))

FRAMES, HOLD0, HOLD1 = 480, 45, 105
order = IDX[np.argsort(-DIST[IDX])]                       # 源流側から河口側へ
t = np.linspace(0, 1, FRAMES + 1) ** 1.45                 # 終盤（本流が着く所）をゆっくり
edges = (t * len(order)).astype(int)
base = [Image.new('L', (BW, BH), 0) for _ in GROUPS]; bd = [ImageDraw.Draw(im) for im in base]
front_acc = np.zeros((SIZE[1], SIZE[0]), np.float32)
k = 0; t0 = time.time()

def render():
    img = np.zeros((SIZE[1], SIZE[0], 3), np.float32); img[:] = PAPER
    for gi, im in enumerate(base):
        a = np.asarray(im.resize(SIZE, Image.LANCZOS), np.float32) / 255 * ALPHAS[gi]
        img = img * (1 - a[..., None]) + INK * a[..., None]
    a = np.clip(front_acc, 0, 1)[..., None] * 0.85
    return Image.fromarray(np.clip(img * (1 - a) + FRONT * a, 0, 255).astype(np.uint8))

first = render(); first.save(f'{OUT}/frames/{k:04d}.png'); k += 1
for _ in range(HOLD0 - 1):
    os.link(f'{OUT}/frames/0000.png', f'{OUT}/frames/{k:04d}.png'); k += 1
for f in range(FRAMES):
    ids = order[edges[f]:edges[f+1]]
    draw(bd, ids)
    ft = [Image.new('L', (BW, BH), 0) for _ in GROUPS]; fd = [ImageDraw.Draw(im) for im in ft]
    draw(fd, ids, wscale=1.6)                              # 先端は少し太く
    add = np.zeros((SIZE[1], SIZE[0]), np.float32)
    for im in ft: add = np.maximum(add, np.asarray(im.resize(SIZE, Image.LANCZOS), np.float32) / 255)
    front_acc = np.maximum(front_acc * 0.72, add)
    render().save(f'{OUT}/frames/{k:04d}.png'); k += 1
    if f % 80 == 0: print('C', f, len(ids), round(time.time()-t0, 1), flush=True)
last = f'{OUT}/frames/{k-1:04d}.png'
for _ in range(HOLD1):
    os.link(last, f'{OUT}/frames/{k:04d}.png'); k += 1
print('C frames', k, round(k/30, 1), 's', flush=True)
os.system(f'cd {OUT} && ffmpeg -y -loglevel error -framerate 30 -i frames/%04d.png -vf "pad=ceil(iw/2)*2:ceil(ih/2)*2:0:0:white" -c:v libx264 -pix_fmt yuv420p -crf 18 ../rivers_tributaries.mp4')
