# Figuras y tablas del paper (paper/figs/*.pdf, paper/tablas/*.tex) a partir de res/.
# Uso: python figuras.py <mejor_corrida.npz> [corridas para envolvente...]
import sys, json, csv, os, numpy as np, matplotlib
matplotlib.use('Agg'); import matplotlib.pyplot as plt
from matplotlib.tri import Triangulation
from matplotlib.colors import LogNorm
plt.rcParams.update({'font.family': 'serif', 'font.size': 9, 'axes.titlesize': 9, 'figure.dpi': 200})
os.makedirs('paper/figs', exist_ok=True); os.makedirs('paper/tablas', exist_ok=True)
T = np.load('terreno.npz'); X0, Y0 = float(T['X0']), float(T['Y0'])
best = sys.argv[1]; r = np.load(best); p = json.loads(str(r['params']))
xs, ys = r['cx'] - X0, r['cy'] - Y0
o = json.load(open('osm.json'))
from pyproj import Transformer
tf = Transformer.from_crs('EPSG:4326', 'EPSG:32719', always_xy=True)
def calles(ax):
    for c in o['calles']:
        if c['k'] not in ('primary', 'secondary', 'tertiary'): continue
        x, y = tf.transform([q[1] for q in c['pts']], [q[0] for q in c['pts']]); ax.plot(np.array(x) - X0, np.array(y) - Y0, lw=0.4, c='0.55', zorder=1)
ext = (3000, 5900, 600, 2900)
def mapa(ax, v, titulo, cmap, norm=None, umbral=0.05, cb=''):
    m = r['hmax'] > umbral; ax.set_facecolor('0.96'); calles(ax)
    sc = ax.scatter(xs[m], ys[m], c=v[m], s=0.6, cmap=cmap, norm=norm, linewidths=0, zorder=2)
    ax.set_xlim(ext[0], ext[1]); ax.set_ylim(ext[2], ext[3]); ax.set_aspect('equal'); ax.set_title(titulo); ax.set_xticks([]); ax.set_yticks([])
    plt.colorbar(sc, ax=ax, fraction=0.035, pad=0.01, label=cb)
fig, axs = plt.subplots(2, 2, figsize=(7.2, 5.4))
mapa(axs[0, 0], r['hmax'], '(a) Profundidad máxima', 'viridis', LogNorm(0.05, 3), cb='m')
mapa(axs[0, 1], r['vmax'], '(b) Velocidad máxima', 'magma', cb='m/s')
mapa(axs[1, 0], r['pdin'] / 1000, '(c) Presión dinámica ½ρU²', 'inferno', cb='kPa')
ta = np.where(np.isnan(r['tarr']), np.nan, r['tarr'] / 60)
mapa(axs[1, 1], ta, '(d) Llegada del frente', 'turbo_r', cb='min')
fig.tight_layout(); fig.savefig('paper/figs/mapas.pdf'); plt.close(fig)
# envolvente de probabilidad (fracción de corridas aceptables con h > 5 cm), sobre la malla de la mejor corrida
if len(sys.argv) > 2:
    from scipy.spatial import cKDTree
    tree = cKDTree(np.c_[r['cx'], r['cy']]); acc = np.zeros(len(xs))
    for f in sys.argv[2:]:
        q = np.load(f); d, i = cKDTree(np.c_[q['cx'], q['cy']]).query(np.c_[r['cx'], r['cy']]); acc += q['hmax'][i] > 0.05
    acc /= len(sys.argv) - 2
    fig, ax = plt.subplots(figsize=(7.2, 3.6)); ax.set_facecolor('0.96'); calles(ax); mm = acc > 0
    sc = ax.scatter(xs[mm], ys[mm], c=acc[mm], s=0.6, cmap='YlOrRd', vmin=0, vmax=1, linewidths=0, zorder=2)
    ax.set_xlim(ext[0], ext[1]); ax.set_ylim(ext[2], ext[3]); ax.set_aspect('equal'); ax.set_xticks([]); ax.set_yticks([])
    plt.colorbar(sc, ax=ax, fraction=0.03, label='fracción de corridas aceptables'); fig.tight_layout(); fig.savefig('paper/figs/probabilidad.pdf'); plt.close(fig)
# convergencia y parada
vc = json.load(open('res/verif_parada.json')); v2 = json.load(open('res/verif_conv2.json'))
cfl = [1.0, 0.5, 0.25] + [0.1, 0.05]; fr = [vc['convergencia_CFL'][str(c)]['frente_final_m'] for c in (1.0, 0.5, 0.25)] + [v2['DE1 CFL 0.1']['frente_m'], v2['DE1 CFL 0.05']['frente_m']]
fig, ax = plt.subplots(figsize=(3.4, 2.4)); ax.plot(cfl, fr, 'o-', label='DE1'); ax.axhline(v2['DE0 CFL 0.1']['frente_m'], ls='--', c='0.4', label='DE0 (todas)')
ax.set_xscale('log'); ax.set_xlabel('número de Courant'); ax.set_ylabel('alcance del frente (m)'); ax.legend(frameon=False); fig.tight_layout(); fig.savefig('paper/figs/convergencia.pdf'); plt.close(fig)
# balance de volumen
b = r['balance']; fig, ax = plt.subplots(figsize=(3.4, 2.4)); ax.plot(b[:, 0] / 60, b[:, 1] / 1000, label='entrado'); ax.plot(b[:, 0] / 60, b[:, 2] / 1000, '--', label='en la malla')
ax.set_xlabel('tiempo (min)'); ax.set_ylabel('volumen (mil m³)'); ax.legend(frameon=False); fig.tight_layout(); fig.savefig('paper/figs/balance.pdf'); plt.close(fig)
# tabla de corridas (h* = 5 cm)
filas = [f for f in csv.DictReader(open('res/corridas.csv')) if f['hstar'] == '0.05']
with open('paper/tablas/corridas.tex', 'w') as t:
    t.write('\\begin{tabular}{llrrlrrrrrr}\\toprule\nCorrida & Origen & $V$ (mil m$^3$) & $Q_p$ (m$^3$/s) & Muestra & $C_v$ & Área (ha) & $R$ & $R_\\mathrm{pres}$ & $P$ & $F$\\\\\\midrule\n')
    for f in sorted(filas, key=lambda f: -float(f['F'])):
        t.write(f"{f['corrida'].replace('.npz','')} & {f['origen']} & {float(f['V'])/1000:.0f} & {float(f['Qp']):.0f} & {f['suelo']} & {float(f['Cv']):.2f} & {float(f['area_ha']):.1f} & {float(f['R']):.2f} & {float(f['R_presencia']):.2f} & {float(f['P']):.2f} & {float(f['F']):.2f}\\\\\n")
    t.write('\\bottomrule\\end{tabular}\n')
print('figuras y tablas listas')