# Exporta una corrida (o un conjunto) al visor 3D web: terreno, bloques, mapa de calor de profundidad/velocidad/presión,
# isócronas de llegada (veriles de propagación), probabilidad de alcance (si hay varias corridas) y fichas de explicación.
# Uso: python exportar_visor.py <corrida_principal.npz> [otras corridas para probabilidad...]
import sys, json, numpy as np
from scipy.interpolate import LinearNDInterpolator, NearestNDInterpolator
from scipy import ndimage as nd
from PIL import Image
from pyproj import Transformer

T = np.load('terreno.npz'); X0, Y0, DX, NX, NY = float(T['X0']), float(T['Y0']), float(T['DX']), int(T['NX']), int(T['NY'])
R = 4.0                                                     # resolución de las capas (m)
gx = X0 + (np.arange(NX) + 0.5) * DX; gy = Y0 + (np.arange(NY) + 0.5) * DX
GX, GY = np.meshgrid(gx, gy)                                # fila 0 = sur

def campo(r, nombre):
    xy = np.c_[r['cx'], r['cy']]; v = r[nombre].astype(np.float64)
    if nombre == 'tarr': v = np.where(np.isnan(v), -1.0, v)
    f = NearestNDInterpolator(xy, v); return f(GX, GY)

principal = sys.argv[1]; r = np.load(principal); p = json.loads(str(r['params']))
H = campo(r, 'hmax'); V = campo(r, 'vmax'); P = campo(r, 'pdin'); TA = campo(r, 'tarr')
wet = H > 0.05
dist = nd.distance_transform_edt(~wet)                       # recorta el relleno del vecino más cercano fuera de la malla mojada
H[~wet] = 0; V[~wet] = 0; P[~wet] = 0; TA[~wet] = -1
def png16(a, vmax, nombre):
    q = np.clip(a / vmax, 0, 1) * 65535; q = q.astype(np.uint32)[::-1]   # fila 0 = norte para la imagen
    Image.fromarray(np.stack([(q >> 8) & 255, q & 255, (a[::-1] > 0).astype(np.uint8) * 255], -1).astype(np.uint8)).save(f'web/{nombre}.png', optimize=True)
png16(H, 3.0, 'h'); png16(V, 8.0, 'v'); png16(P, 60000.0, 'p'); png16(np.where(TA > 0, TA, 0), 3600.0, 't')
# probabilidad de alcance (fracción de corridas con h > 5 cm) si se pasan varias
if len(sys.argv) > 2:
    acc = np.zeros_like(H)
    for f in sys.argv[1:]: acc += campo(np.load(f), 'hmax') > 0.05
    png16(acc / (len(sys.argv) - 1), 1.0, 'prob')
# terreno (DTM) para el relieve 3D, cada 8 m
step = 2; Z = T['DTM'][::step, ::step].astype(np.float64); zmin, zmax = float(Z.min()), float(Z.max())
q = (np.clip((Z - zmin) / (zmax - zmin), 0, 1) * 65535).astype(np.uint32)[::-1]
Image.fromarray(np.stack([(q >> 8) & 255, q & 255, np.zeros_like(q)], -1).astype(np.uint8)).save('web/dtm.png', optimize=True)
# bloques con estadísticas de impacto
tf = Transformer.from_crs('EPSG:4326', 'EPSG:32719', always_xy=True)
B = json.load(open('bloques.json')); bl = []
for b in B:
    xs, ys = tf.transform([q[1] for q in b['ring']], [q[0] for q in b['ring']])
    i0, i1 = np.clip(((np.array(xs) - X0) / DX).astype(int), 0, NX - 1), np.clip(((np.array(ys) - Y0) / DX).astype(int), 0, NY - 1)
    sl = (slice(max(0, i1.min() - 3), i1.max() + 4), slice(max(0, i0.min() - 3), i0.max() + 4))   # entorno de 12 m alrededor de la planta
    hb, vb, pb = float(H[sl].max()), float(V[sl].max()), float(P[sl].max())
    bl.append({'r': [[round(x - X0, 1), round(y - Y0, 1)] for x, y in zip(xs, ys)], 'h': b['h'], 'zb': round(float(T['DTM'][int(i1.mean()), int(i0.mean())]), 1),
               'hm': round(b['dsm_dtm_medio'], 1), 'hx': round(b['dsm_dtm_p95'], 1), 'f': b['fuente'][:3], 'n': b.get('name') or '',
               'imp': {'h': round(hb, 2), 'v': round(vb, 2), 'p': round(pb / 1000, 1)}})
o = json.load(open('osm.json'))
calles = [{'n': c['name'], 'p': [[round(x - X0, 1), round(y - Y0, 1)] for x, y in zip(*tf.transform([q[1] for q in c['pts']], [q[0] for q in c['pts']]))]}
          for c in o['calles'] if c['k'] in ('primary', 'secondary', 'tertiary', 'residential')]
meta = {'X0': X0, 'Y0': Y0, 'W': NX * DX, 'D': NY * DX, 'zmin': zmin, 'zmax': zmax, 'escalas': {'h': 3.0, 'v': 8.0, 'p': 60000.0, 't': 3600.0},
        'corrida': principal, 'params': p, 'area_ha': round(float(np.sum(wet)) * DX * DX / 1e4, 1), 'h_max': round(float(H.max()), 2), 'v_max': round(float(V.max()), 2),
        'n_prob': len(sys.argv) - 1}
json.dump({'meta': meta, 'bloques': bl, 'calles': calles}, open('web/escena.json', 'w'), separators=(',', ':'))
print('exportado', meta['area_ha'], 'ha · bloques', len(bl), '· calles', len(calles))