# Figura de la nota técnica: profundidad máxima para H1, H2 y H4 sobre satélite, con Quebrada Honda, la embocadura del colector
# en Fco. Bulnes Correa y los sitios afectados. Uso: python figura_colector.py res/d04.npz res/e03.npz res/e01.npz
import sys, os, json, numpy as np, matplotlib
matplotlib.use('Agg'); import matplotlib.pyplot as plt
from matplotlib.colors import LogNorm
from pyproj import Transformer
from PIL import Image
plt.rcParams.update({'font.family': 'serif', 'font.size': 8})
tf = Transformer.from_crs('EPSG:4326', 'EPSG:32719', always_xy=True)
T = np.load('terreno.npz'); X0, Y0, DX, NX, NY = [float(T[k]) for k in ('X0', 'Y0', 'DX', 'NX', 'NY')]
sat = Image.open('web/satelite.jpg'); o = json.load(open('osm.json'))
ext = (X0 + 1200, X0 + NX * DX - 150, Y0 + 300, Y0 + 3000)
ROT = {'H1': 'H1: salida solo en el abanico', 'H2': 'H2: salida solo en la embocadura', 'H4': 'H4: abanico + embocadura'}
fig, axs = plt.subplots(1, 3, figsize=(7.2, 2.9), sharey=True)
SAL = sys.argv[4] if len(sys.argv) > 4 else 'paper/nota-colector/figs/comparacion.pdf'
for ax, arg in zip(axs, sys.argv[1:4]):
    f, _, tit = arg.partition(':'); r = np.load(f); p = json.loads(str(r['params'])); m = r['hmax'] > 0.05
    ax.imshow(sat, extent=(X0, X0 + NX * DX, Y0, Y0 + NY * DX), alpha=0.65)
    sc = ax.scatter(r['cx'][m], r['cy'][m], c=r['hmax'][m], s=0.15, cmap='turbo', norm=LogNorm(0.05, 2.5), lw=0, rasterized=True)
    for c in o['calles']:
        if c['name'] == 'Quebrada Honda':
            x, y = tf.transform([q[1] for q in c['pts']], [q[0] for q in c['pts']]); ax.plot(x, y, c='white', lw=1.6, ls=(0, (2, 1)))
    ex, ey = tf.transform(-70.51629, -33.40596); ax.plot(ex, ey, 'o', ms=6, mfc='none', mec='red', mew=1.5)
    ax.annotate('embocadura', (ex, ey), (ex - 900, ey + 650), color='white', fontsize=7, arrowprops=dict(arrowstyle='-', color='white', lw=0.6))
    for la, lo in ((-33.40014, -70.50874), (-33.40727, -70.50373), (-33.40063, -70.50540), (-33.40424, -70.50660)):
        x, y = tf.transform(lo, la); ax.plot(x, y, 's', ms=3, mfc='white', mec='k', mew=0.5)
    ax.set_xlim(ext[0], ext[1]); ax.set_ylim(ext[2], ext[3]); ax.set_aspect('equal'); ax.set_xticks([]); ax.set_yticks([])
    ax.set_title(tit.replace('|', '\n') if tit else p['origen'], fontsize=7.5)
cb = fig.colorbar(sc, ax=axs, fraction=0.02, pad=0.01); cb.set_label('h máx (m)')
os.makedirs(os.path.dirname(SAL), exist_ok=True); fig.savefig(SAL, bbox_inches='tight', dpi=250); print('ok')