# -*- coding: utf-8 -*-
"""第 2 世代。(1) 実物と同じ濃さの配色にする (2) 5°の格子を表示用に滑らかにする。

平滑化は氷の輪郭と同じ規範で行う: 表示用にだけ滑らかにし、統計は元格子（5°）のまま。
欠測は平滑化に混ぜない（観測のある格子だけで重み付き平均を取り、欠測は欠測のまま残す）。
"""
import os, sys, json
import numpy as np, netCDF4 as nc
sys.path.insert(0, '/home/claude/forest')
import forest1 as F, sylvania as S
from PIL import Image
from scipy import ndimage
OUT = '/home/claude/stripes/out'; os.makedirs(OUT, exist_ok=True)
PAPER = np.array([255, 255, 255], np.float32)

def hx(s): return np.array([int(s[i:i+2], 16) for i in (1, 3, 5)], np.float32)
# 実物の温暖化ストライプと同じ ColorBrewer 8 段（淡→濃）
COOL_H = [hx(c) for c in ('#deebf7', '#c6dbef', '#9ecae1', '#6baed6', '#4292c6', '#2171b5', '#08519c', '#08306b')]
WARM_H = [hx(c) for c in ('#fee0d2', '#fcbba1', '#fc9272', '#fb6a4a', '#ef3b2c', '#cb181d', '#a50f15', '#67000d')]
# SORAH のインクで同じ濃度設計（淡い端も紙より確実に濃く、深い端まで伸ばす）
def ink_ramp(deep, light_L=92, deep_L=None, n=8):
    lab = S.to_lab(deep); L, a, b = lab
    deep_L = L if deep_L is None else deep_L
    out = []
    for k in range(n):
        t = (k + 1) / n
        LL = light_L + (deep_L - light_L) * t
        out.append(S.from_lab(np.array([LL, a * (0.35 + 0.65 * t), b * (0.35 + 0.65 * t)])).astype(np.float32))
    return out
WARM_O = ink_ramp(np.array(F.r.col(46, 50, 70), np.float32), deep_L=32)
COOL_O = ink_ramp(np.array(F.r.col(30, 270, 40), np.float32), deep_L=24)

_y = None
def annual():
    global _y
    if _y is None:
        d = nc.Dataset('/home/claude/stripes/grid.nc')
        a = np.ma.masked_invalid(d.variables['tas_mean'][:])
        n = a.shape[0] // 12 * 12
        _y = (np.arange(1850, 1850 + n // 12), a[:n].reshape(-1, 12, 36, 72).mean(1))
    return _y

def smooth_field(g, up=10, sigma_deg=3.5):
    """5°格子を 0.5°に展開し、観測のある所だけで重み付き平滑（表示用）。欠測は欠測のまま返す。"""
    v = g.filled(np.nan)
    m = np.isfinite(v).astype(np.float32)
    v0 = np.where(m > 0, v, 0).astype(np.float32)
    V = np.repeat(np.repeat(v0, up, 0), up, 1)
    M = np.repeat(np.repeat(m, up, 0), up, 1)
    s = sigma_deg * up / 5.0
    num = ndimage.gaussian_filter(V * M, s, mode=('nearest', 'wrap'))
    den = ndimage.gaussian_filter(M, s, mode=('nearest', 'wrap'))
    out = np.where(den > 0.25, num / np.maximum(den, 1e-6), np.nan)
    out = np.where(M > 0.5, out, np.nan)      # 観測のない格子は埋めない
    return out

def fill_field(g, up=10, sigma_deg=3.5):
    """欠測を周囲の形から補間しつつ、どこが補間かを別に返す（表示用）。

    観測のある格子だけで重み付き平滑をかけ、足りない所は sigma を段階的に広げて埋める。
    元の曲面の曲がりを引き継ぐので、格子の縁が立たない。"""
    v = g.filled(np.nan)
    m = np.isfinite(v).astype(np.float32)
    V = np.repeat(np.repeat(np.where(m > 0, v, 0).astype(np.float32), up, 0), up, 1)
    M = np.repeat(np.repeat(m, up, 0), up, 1)
    md = ('nearest', 'wrap')
    base = sigma_deg * up / 5.0
    num = ndimage.gaussian_filter(V * M, base, mode=md); den = ndimage.gaussian_filter(M, base, mode=md)
    out = np.where(den > 0.25, num / np.maximum(den, 1e-6), np.nan)
    for mult in (2, 4, 8, 16, 32):     # 埋まらない所だけ、順に広い平滑で埋める
        if not np.isnan(out).any(): break
        s2 = base * mult
        n2 = ndimage.gaussian_filter(V * M, s2, mode=md); d2 = ndimage.gaussian_filter(M, s2, mode=md)
        f2 = np.where(d2 > 1e-4, n2 / np.maximum(d2, 1e-9), np.nan)
        out = np.where(np.isnan(out), f2, out)
    # 観測の有無も同じ幅でぼかし、灰への移りを段にしない
    conf = ndimage.gaussian_filter(M, base * 0.8, mode=md)
    conf = np.clip((conf - 0.15) / 0.55, 0, 1)
    return out, conf

def bands(v, lim, warm, cool, n=8):
    k = np.clip((np.abs(v) / lim * n).astype(int), 0, n - 1)
    return k

def grey_out(c, k=0.62, lift=6):
    """補間した区画の色。彩度を落とし、わずかに明るく（帯の形は残す）。"""
    L, a, b = S.to_lab(np.array(c, np.float32))
    return S.from_lab(np.array([min(97, L + lift), a * (1 - k), b * (1 - k)])).astype(np.float32)

def extrema(fld, nwarm=4, ncool=2, sep_deg=28, floor=0.35):
    """場の山と谷を拾う。近すぎるものは強い方を残す（0.5°格子の添字で返す）。"""
    from scipy import ndimage as nd
    out = []
    for sign, count in ((1, nwarm), (-1, ncool)):
        f = np.nan_to_num(fld * sign, nan=-9e9)
        mx = nd.maximum_filter(f, size=21, mode=('nearest', 'wrap'))
        cand = np.argwhere((f == mx) & (f > floor))
        cand = sorted(cand.tolist(), key=lambda p: -f[p[0], p[1]])
        keep = []
        for r, c in cand:
            if len(keep) >= count: break
            if all(min(abs(c - c2), fld.shape[1] - abs(c - c2)) * 0.1 > sep_deg or abs(r - r2) * 0.1 > sep_deg for r2, c2 in keep):
                keep.append((r, c))
        out += [(r, c, float(fld[r, c])) for r, c in keep]
    return out

def label_extrema(im, sc, fld, items, size=15):
    """山と谷の中心に数値を置く。下地の明暗でトーンを決め、PIL の描画（アンチエイリアス付き）で 1 回に描く。"""
    import forest1 as FF
    from PIL import ImageDraw
    h, w = fld.shape
    arr = np.array(im).astype(np.float32); L = FF.lstar(arr)
    d = ImageDraw.Draw(im); f = FF.g1.font('Regular', size)
    for r, c, v in items:
        lat = -90 + (r + 0.5) * 180 / h; lon = -180 + (c + 0.5) * 360 / w
        x, y = sc.to_px(lon, lat)
        if not (78 < x < sc.W - 78 and 20 < y < sc.H - 20): continue   # 端で切れる位置には置かない
        y0 = int(max(0, y - 14)); y1 = int(min(sc.H, y + 14)); x0 = int(max(0, x - 46)); x1 = int(min(sc.W, x + 46))
        # 北極海氷と同じ規範: 固定色ではなく下地からの相対シフト（明るい所は沈め、暗い所は持ち上げる）
        base = arr[y0:y1, x0:x1].reshape(-1, 3).mean(0)
        lab = S.to_lab(base.astype(np.float32))
        lab[0] = lab[0] - 34 if lab[0] > 55 else lab[0] + 40
        lab[1] *= 0.55; lab[2] *= 0.55
        tone = tuple(int(v) for v in S.from_lab(lab))
        txt = ('+' if v >= 0 else '\u2212') + f'{abs(v):.1f}\u00b0C'
        d.text((x, y), txt, font=f, fill=tone, anchor='mm')
    return im

def world(name, warm, cool, year=2025, size=(1350, 675), lim=2.0, smooth=3.5, land=True, fill=True, ss=3, marks=True):
    size = (size[0]*ss, size[1]*ss)   # 3 倍で描いてから縮小（帯の境目と海岸線の階段を消す）
    yrs, y = annual(); g = y[int(np.where(yrs == year)[0][0])]
    if fill:
        fld, obs = fill_field(g, sigma_deg=smooth)
    else:
        fld = smooth_field(g, sigma_deg=smooth) if smooth else np.repeat(np.repeat(g.filled(np.nan), 10, 0), 10, 1)
        obs = np.isfinite(fld).astype(np.float32)
    sc = F.Scene((-180, -90, 180, 90), size, proj='pc', pad=0.0)
    lon, lat, ok = sc.lonlat_grid()
    h, w = fld.shape
    # 双線形で場を読む（最近傍だと帯の境目が格子の形に折れる）
    fx = np.clip((lon + 180) / 360 * w - 0.5, 0, w - 1.001); fy = np.clip((lat + 90) / 180 * h - 0.5, 0, h - 1.001)
    x0 = fx.astype(int); y0 = fy.astype(int); tx = fx - x0; ty = fy - y0
    def bil(a):
        a = np.nan_to_num(a, nan=0.0)
        return (a[y0, x0]*(1-tx)*(1-ty) + a[y0, x0+1]*tx*(1-ty) + a[y0+1, x0]*(1-tx)*ty + a[y0+1, x0+1]*tx*ty)
    v = np.where(ok, bil(fld), np.nan)
    ob = np.where(ok, bil(obs.astype(np.float32)), 0.0)
    fin = np.isfinite(v)
    img = np.zeros((sc.H, sc.W, 3), np.uint8); img[:] = PAPER.astype(np.uint8)
    k = bands(np.nan_to_num(v), lim, warm, cool)
    gw = [grey_out(c) for c in warm]; gc = [grey_out(c) for c in cool]
    full = np.zeros((sc.H, sc.W, 3), np.float32); grey = np.zeros_like(full)
    for b in range(8):
        mw = fin & (v >= 0) & (k == b); mc = fin & (v < 0) & (k == b)
        full[mw] = warm[b]; grey[mw] = gw[b]
        full[mc] = cool[b]; grey[mc] = gc[b]
    c3 = ob[..., None]
    mix = np.clip(grey * (1 - c3) + full * c3, 0, 255)
    img[fin] = mix[fin].astype(np.uint8)
    # 補間した区画の外周を灰色の線で囲う（どこが観測でどこが補間かを明示する）
    interp = fin & (ob < 0.5)
    if interp.any():
        wpx = max(1, int(round(ss)))          # 出力で約 1px
        ring = interp & ~ndimage.binary_erosion(interp, iterations=wpx)
        b0 = np.array(img).astype(np.float32)
        b0[ring] = b0[ring] * 0.68 + np.array([120, 120, 124], np.float32) * 0.32   # 透過させて馴染ませる
        img = np.clip(b0, 0, 255).astype(np.uint8)
    im = Image.fromarray(img)
    if land:
        lm = F.land_mask(sc) > 0.5
        edge = lm & ~ndimage.binary_erosion(lm)
        b_ = np.array(im).astype(np.float32)
        b_[edge] = b_[edge] * 0.62
        im = Image.fromarray(np.clip(b_, 0, 255).astype(np.uint8))
    if ss > 1: im = im.resize((size[0]//ss, size[1]//ss), Image.LANCZOS)
    if marks:
        sc2 = F.Scene((-180, -90, 180, 90), (size[0]//ss, size[1]//ss), proj='pc', pad=0.0)
        im = label_extrema(im, sc2, fld, extrema(fld))
    im.save(f'{OUT}/{name}.png')
    obs = int(np.isfinite(g.filled(np.nan)).sum())
    json.dump(dict(year=year, lim=lim, sigma_deg=smooth, observed_cells=obs, cells=36*72),
              open(f'{OUT}/{name}.json', 'w'), indent=1)
    print(name, size, 'observed', obs, flush=True)

def stripes(name, warm, cool, size=(1350, 1080), lim=1.2):
    import csv
    rows = list(csv.DictReader(open('/home/claude/stripes/global_annual.csv')))
    vals = [float(r['Anomaly (deg C)']) for r in rows]
    W, H = size; img = np.zeros((H, W, 3), np.uint8); n = len(vals)
    for i, v in enumerate(vals):
        x0 = int(i * W / n); x1 = int((i + 1) * W / n)
        k = min(7, int(abs(v) / lim * 8))
        img[:, x0:x1] = (warm[k] if v >= 0 else cool[k]).astype(np.uint8)
    Image.fromarray(img).save(f'{OUT}/{name}.png'); print(name, size, flush=True)

if __name__ == '__main__':
    world('V1', WARM_H, COOL_H, 1880)                 # 補間して灰に落とす
    world('V2', WARM_H, COOL_H, 1850)
    world('V3', WARM_H, COOL_H, 1900)
    world('V4', WARM_H, COOL_H, 2025)