# -*- coding: utf-8 -*-
"""地図の地の候補。Q3d の印（赤、大きさ＝M、薄さ＝深さ、M6/M7 で不透明度↑）を共通に載せる。"""
import sys, os; sys.path.insert(0,'/home/claude/quakes'); sys.path.insert(0,'/home/claude/forest')
import q1 as Q, numpy as np, forest1 as F, sylvania as S, tifffile
from nyt import solid_grat
from PIL import Image, ImageDraw
from scipy import ndimage
Q.INK = np.array([196,34,28],np.float32)
sc, lon, lat, ok, land = Q.scene(); W,H=sc.W,sc.H
SR = tifffile.imread('/home/claude/quakes/ne/GRAY_50M_SR.tif').astype(np.float32)
OB = tifffile.imread('/home/claude/quakes/ne/GRAY_50M_SR_OB.tif').astype(np.float32)
sr = F.sample_bilinear(SR, lon, lat, ok); ob = F.sample_bilinear(OB, lon, lat, ok)

def lab_img(Lmap, a=0.0, b=0.0):
    """L のマップから RGB（a,b 一定）。"""
    out = np.zeros((H,W,3),np.float32)
    # 近似: 無彩色なら L→sRGB は 1 次元なので表引き
    Ls = np.linspace(0,100,1001); tab = np.array([S.from_lab(np.array([L,a,b])) for L in Ls])
    idx = np.clip((Lmap*10).astype(int),0,1000); return tab[idx]

def base(v):
    img = np.zeros((H,W,3),np.float32)
    grat = True
    if v == 'B0': img[:] = Q.PAPER; img[land] = Q.LANDC
    elif v == 'B1': img[:] = Q.PAPER; img[land] = (226,224,219); grat = False
    elif v == 'B2':
        img[:] = Q.PAPER; img[land] = Q.LANDC
    elif v == 'B3':   # 陸に起伏。相対的な起伏 (SR 143..251) を L 96→88 に。海は白
        rel = np.clip((sr-143)/(251-143),0,1)
        L = np.where(land, 95.5 - 8.0*rel, 100.0); img = lab_img(L, 0.3, 1.2)
    elif v == 'B4':   # 陸の起伏＋海底（海嶺が見える）。海は L 100→96.5
        rel = np.clip((sr-143)/(251-143),0,1); bath = np.clip((ob-80)/(147-80),0,1)   # 深い=0, 浅い=1
        L = np.where(land, 95.5 - 8.0*rel, 97.0 + 3.0*bath); img = lab_img(L, 0.3, 1.2)
    elif v == 'B5':   # 反転: 海が薄い灰の紙、陸が白
        img[:] = (244,243,240); img[land] = (255,255,255)
    elif v == 'B6':   # B4 の海底だけ（陸は平坦な灰）
        bath = np.clip((ob-80)/(147-80),0,1)
        L = np.where(land, 93.5, 97.0 + 3.0*bath); img = lab_img(L, 0.3, 1.2)
    return img, grat

def coast_mask():
    e = land ^ ndimage.binary_erosion(land, iterations=Q.SS)
    return e.astype(np.float32)

x=(Q.LON*111320.0-sc.cx)/sc.s+W/2; y=H/2-(Q.LAT*111320.0-sc.cy)/sc.s
r=0.9*2.4**(Q.MAG-5.0)*Q.SS
AMAX=0.55; al=np.where(Q.MAG<6,0.30,np.where(Q.MAG<7,0.40,AMAX))
edges=[0,35,70,150,300,800]; depth_L=[0,0.15,0.35,0.55,0.72]
a_ink,b_pap=S.to_lab(Q.INK),S.to_lab(Q.PAPER)
layers=[]
for k in range(4,-1,-1):
    m=(Q.DEP>=edges[k])&(Q.DEP<edges[k+1]); idx=np.nonzero(m)[0]; idx=idx[np.argsort(Q.MAG[idx])]
    cv=Image.new('L',(W,H),0); dr=ImageDraw.Draw(cv)
    for i in idx: dr.ellipse([x[i]-r[i],y[i]-r[i],x[i]+r[i],y[i]+r[i]],fill=int(255*al[i]/AMAX))
    layers.append((np.array(cv).astype(np.float32)/255*AMAX, S.from_lab(a_ink+(b_pap-a_ink)*depth_L[k]).astype(np.float32)))

def render(v):
    img, grat = base(v); out = img.copy()
    im = Image.fromarray(np.clip(out,0,255).astype(np.uint8)); F.NAVY=Q.INK; F.WHITE=Q.PAPER
    if v == 'B2': im = F.comp_rel(im, coast_mask(), 0.14, 0.10)
    out = np.array(im).astype(np.float32)
    for a, col in layers: out = out*(1-a[...,None]) + col*a[...,None]
    im = Image.fromarray(np.clip(out,0,255).astype(np.uint8))
    if grat: im = F.comp_rel(im, F.aa_mask((W,H), solid_grat(sc)), 0.14, 0.10)
    im = im.resize(Q.SIZE, Image.LANCZOS); im.save(f'/home/claude/quakes/out/{v}.png'); print(v, flush=True); return im

if __name__ == '__main__':
    vs = sys.argv[1:] or ['B1','B2','B3','B4','B5','B6']
    ims = [render(v) for v in vs]
    s = Image.new('RGB', (1350, (675+10)*len(ims)), (200,200,200))
    for i, im in enumerate(ims): s.paste(im, (0, i*685))
    s.save('/home/claude/quakes/out/Bsheet.png')