# -*- coding: utf-8 -*-
"""河川を、森林（N3）・気温と同じ図法・同じ枠で描く。

図法: 正距円筒。枠: 森林 N3 と同一（lon -180..180 / lat -58..83、1350×1080、pad 0.02）。
これで 3 作品が同じ座標系・同じ余白に載る。
"""
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/out'; os.makedirs(OUT, exist_ok=True)
PAPER = np.array([255, 255, 255], np.float32)
INK = np.array([86, 120, 196], np.float32)
Z = np.load('/home/claude/rivers/rivers.npz')
PTS, CNT, CLS = Z['pts'], Z['cnt'], Z['cls']
OFF = np.concatenate([[0], np.cumsum(CNT)])
BBOX = (-180, -58, 180, 83); SIZE = (1350, 1080)      # 森林 N3 と同じ

def build(name, groups, widths, ss=3, alphas=None, size=SIZE, ink=INK):
    sc = F.Scene(BBOX, (size[0]*ss, size[1]*ss), proj='pc')   # 同じ Scene で座標を決める
    W, H = sc.W, sc.H
    img = np.zeros((size[1], size[0], 3), np.float32); img[:] = PAPER
    if alphas is None: alphas = [1.0]*len(groups)
    for g, wpx, al in zip(groups, widths, alphas):
        im = Image.new('L', (W, H), 0); d = ImageDraw.Draw(im)
        sel = np.where(np.isin(CLS, g))[0]; t0 = time.time()
        for i in sel:
            xy = PTS[OFF[i]:OFF[i+1]]
            x = (xy[:, 0]*111320.0 - sc.cx) / sc.s + W/2
            y = H/2 - (xy[:, 1]*111320.0 - sc.cy) / sc.s
            if x.max() - x.min() > W*0.5: continue
            d.line(list(zip(x.tolist(), y.tolist())), fill=255, width=max(1, int(round(wpx*ss))))
        a = np.asarray(im.resize(size, Image.LANCZOS), np.float32) / 255 * al
        img = img * (1 - a[..., None]) + ink * a[..., None]
        print('  ', g[:2], len(sel), round(time.time()-t0, 1), 's', flush=True)
    Image.fromarray(np.clip(img, 0, 255).astype(np.uint8)).save(f'{OUT}/{name}.png')
    print(name, size, flush=True)

if __name__ == '__main__':
    build('R_width', [[1,2],[3,4],[5],[6]], [1.0, 0.7, 0.45, 0.3], alphas=[1, 1, .95, .75])  # 階級で太さ（E2 と同じ設計）
    build('R_even', [list(range(1, 6))], [0.4])                                              # 一律の細線（E1 と同じ設計）