# -*- coding: utf-8 -*-
"""熱帯低気圧の経路（IBTrACS v04r01, 1980–2025, main track）。v4 の表から導いた最初の一枚。
線×観測の時間 → 線の規範（階級で間引き、太さに階級）＋ 積層（画素あたりの回数）。"""
import sys, os, json
import numpy as np, pandas as pd
sys.path.insert(0, '/home/claude/forest'); sys.path.insert(0, '/home/claude/quakes')
import forest1 as F, sylvania as S, q1 as Q
from PIL import Image, ImageDraw
from scipy import ndimage
OUT = '/home/claude/cyclones/out'; os.makedirs(OUT, exist_ok=True)
PAPER = np.array([255,255,255], np.float32); LANDC = np.array([238,238,235], np.float32)
INK = np.array([78, 64, 170], np.float32)     # 暫定: 嵐＝青紫（色相の規則は要確認）
BBOX = (-180, -52, 180, 62); SIZE = (1350, int(round(1350*114/360))); SS = 3; LON0 = 180.0   # 太平洋を中央に（題材が及ぶ範囲で中心を決める）
df = pd.read_pickle('/home/claude/cyclones/tracks.pkl'); df['LON'] = ((df.LON + 180) % 360) - 180   # IBTrACS は 180 を超える経度を含む
CLS = [(0, 34, 0.3, 0.45), (34, 64, 0.55, 0.8), (64, 96, 0.9, 0.95), (96, 999, 1.3, 1.0)]   # (kt lo, hi, width px, alpha): TD / TS / Cat1-2 / Cat3-5

def scene(ss=SS):
    sc = F.Scene(BBOX, (SIZE[0]*ss, SIZE[1]*ss), proj='pc', pad=0.0, lon0=LON0)
    lon, lat, ok = sc.lonlat_grid()
    land = F.sample(Q.land_grid_all().astype(np.float32), lon, lat, ok) > 0.5
    return sc, lon, lat, ok, land

def segments():
    """経路を区間（p0, p1, wind）に分解。日付変更線をまたぐ区間は捨てる。"""
    out = []
    for sid, g in df.sort_values(['SID','t']).groupby('SID', sort=False):
        lon = g.LON.values.astype(float); lat = g.LAT.values.astype(float); w = g.wind.values.astype(float)
        for i in range(len(g)-1):
            d = ((lon[i+1]-lon[i]+180) % 360) - 180
            if abs(d) > 30: continue
            lon[i+1] = lon[i] + d
            ww = np.nanmax([w[i], w[i+1]]) if not (np.isnan(w[i]) and np.isnan(w[i+1])) else np.nan
            out.append((lon[i], lat[i], lon[i+1], lat[i+1], ww))
    a = np.array(out, np.float32); print('segments', len(a), 'wind nan', int(np.isnan(a[:,4]).sum())); return a

def to_px(sc, lon, lat):
    lon = ((np.asarray(lon) - LON0 + 180) % 360) - 180
    return (lon*111320.0 - sc.cx)/sc.s + sc.W/2, sc.H/2 - (lat*111320.0 - sc.cy)/sc.s

_SEG = None
def lines_img(name='T1', drop_td=False, min_wind=0, cls=None, alpha_scale=1.0):
    global _SEG
    sc, lon, lat, ok, land = scene(); W, H = sc.W, sc.H
    if _SEG is None: _SEG = segments()
    seg = _SEG; wind = np.nan_to_num(seg[:,4], nan=30.0); cls = cls or CLS
    x0, y0 = to_px(sc, seg[:,0], seg[:,1]); x1, y1 = to_px(sc, seg[:,2], seg[:,3])
    out = np.zeros((H, W, 3), np.float32); out[:] = PAPER; out[land] = LANDC
    for lo, hi, wpx, al in cls:
        if drop_td and lo == 0: continue
        if hi <= min_wind: continue
        al = al*alpha_scale
        m = (wind >= max(lo, min_wind)) & (wind < hi)
        cv = Image.new('L', (W, H), 0); d = ImageDraw.Draw(cv)
        for i in np.nonzero(m)[0]:
            if abs(x1[i]-x0[i]) > W/2: continue
            d.line([(x0[i], y0[i]), (x1[i], y1[i])], fill=255, width=max(1, int(round(wpx*SS))))
        a = (np.array(cv).astype(np.float32)/255 * al)[..., None]
        out = out*(1-a) + INK*a
    im = Image.fromarray(np.clip(out,0,255).astype(np.uint8)).resize(SIZE, Image.LANCZOS)
    im.save(f'{OUT}/{name}.png'); print(name, SIZE, flush=True); return im

def stack_img(name='T2', cell=0.5, bands=8):
    """画素あたりの回数（経路が通った嵐の数）。log1p で尺度を固定。"""
    sc, lon, lat, ok, land = scene()
    w, h = int(360/cell), int(180/cell); cnt = np.zeros((h, w), np.float32)
    for sid, g in df.groupby('SID'):
        hit = np.zeros((h, w), bool)
        ix = np.clip(((g.LON.values+180)/cell).astype(int), 0, w-1); iy = np.clip(((90-g.LAT.values)/cell).astype(int), 0, h-1)
        hit[iy, ix] = True; cnt += hit
    cs = ndimage.gaussian_filter(cnt, 1.0, mode=('nearest','wrap')); ref = float(np.quantile(cs[cs>0], 0.995))
    v = np.clip(np.log1p(cs)/np.log1p(ref), 0, 1).astype(np.float32)
    # 継ぎ目を 0° に移してから標本化（180° の継ぎ目が太平洋の真ん中に出るため）
    fr = F.sample_bilinear(np.roll(v, -w//2, axis=1), (lon % 360) - 180, lat, ok)
    img = np.zeros((sc.H, sc.W, 3), np.uint8); img[:] = PAPER.astype(np.uint8); img[land] = LANDC.astype(np.uint8)
    a, b = S.to_lab(PAPER), S.to_lab(INK); cols = [S.from_lab(a + (b-a)*(k+0.5)/bands).astype(np.uint8) for k in range(bands)]
    on = fr > 1.0/bands/2
    for k in range(bands):
        m = on & (fr > k/bands) & (fr <= (k+1)/bands + (1 if k == bands-1 else 0)); img[m] = cols[k]
    im = Image.fromarray(img).resize(SIZE, Image.LANCZOS); im.save(f'{OUT}/{name}.png')
    json.dump(dict(cell=cell, sigma=1.0, ref=ref, max=float(cnt.max())), open(f'{OUT}/{name}.json','w')); print(name, ref, cnt.max(), flush=True); return im

if __name__ == '__main__':
    lines_img('T1')                                   # 全区間、太さ＝風速の階級
    lines_img('T3', min_wind=64)                      # Cat1 以上に間引く
    lines_img('T4', drop_td=True, cls=[(34,64,0.35,0.30),(64,96,0.6,0.5),(96,999,0.9,0.7)])   # 細く淡く、密度を濃淡に
    lines_img('T5', min_wind=96, cls=[(96,999,1.0,0.9)])   # Cat3 以上だけ
    stack_img('T2')