"""
generate_fig_sni_comparison_b1.05.py  (self-contained, fitted-grid)
====================================================================
Figure 5: LGCP PSHA vs SNI 1726:2019 for Sumatra at b = 1.05 (Aki-Utsu MLE).
  (A) LGCP 475-yr PGA   (B) SNI 475-yr PGA   (C) ratio LGCP/SNI

The LGCP 475-yr surface is computed here on a 60x60 grid from the FITTED
intensity field: log-lambda = beta0 + beta1*phi~ + beta2*Dfault~, using the
30-node Chlieh (2008) coupling field and GEM/EOS fault polyline of
sumatra_lgcp.py and the fitted coefficients in lgcp_parameters.csv. No /tmp
inputs. Run from a directory containing lgcp_parameters.csv.
"""
import numpy as np, pandas as pd, matplotlib
import os
import cartopy
import cartopy.crs as ccrs
import cartopy.feature as cfeature
matplotlib.use('Agg')
import matplotlib.pyplot as plt
from matplotlib.colors import TwoSlopeNorm
from mpl_toolkits.axes_grid1.inset_locator import inset_axes
from matplotlib.patheffects import withStroke
from scipy.stats import norm as sp_norm
from scipy.interpolate import RBFInterpolator
import warnings; warnings.filterwarnings('ignore')

# ── Geometry / constants ──────────────────────────────────────────────────────
cartopy.config['data_dir'] = os.environ.get('CARTOPY_DATA_DIR',
                                            '/private/tmp/cartopy_data')
_PC = ccrs.PlateCarree()
_COAST_10M = cfeature.NaturalEarthFeature(
    'physical', 'coastline', '10m',
    edgecolor='#111111', facecolor='none', linewidth=1.15)
_BORDER_10M = cfeature.NaturalEarthFeature(
    'cultural', 'admin_0_boundary_lines_land', '10m',
    edgecolor='#444444', facecolor='none', linewidth=0.35, alpha=0.65)
LAT_MIN, LAT_MAX, LON_MIN, LON_MAX = -6., 6., 95., 108.
lat0, lon0 = 0., 101.5
KM_LON = 111.320 * np.cos(np.radians(lat0)); KM_LAT = 111.320
N_MAIN, T_OBS, MMAX = 693, 26.0, 9.0
B_VAL, MC = 1.05, 4.5
NQ_SRC, NG = 40, 60          # source quadrature / map grid
PGA = np.logspace(np.log10(0.001), np.log10(3.0), 45)

# ── Fitted covariate field (sumatra_lgcp.py): 30-node Chlieh coupling ─────────
TR_LAT = np.array([13.5,12.5,11.5,10.5,9.5,8.5,8.0,6.5,5.5,4.5,4.0,3.0,2.0,1.0,0.0,
                   -0.5,-1.0,-1.5,-2.0,-2.5,-2.8,-3.2,-3.5,-3.8,-4.2,-4.5,-5.0,-5.5,-6.0,-6.5])
TR_LON = np.array([92.2,92.3,92.4,92.5,92.6,92.7,92.8,93.0,93.2,93.4,93.5,93.6,93.8,94.2,95.0,
                   95.5,96.0,96.5,97.2,98.0,98.5,99.0,99.5,100.0,100.5,101.0,101.5,102.0,102.5,103.0])
COUP   = np.array([0.08,0.09,0.10,0.11,0.12,0.13,0.14,0.17,0.20,0.25,0.28,0.32,0.38,0.50,0.65,
                   0.70,0.75,0.80,0.85,0.88,0.87,0.85,0.80,0.75,0.68,0.62,0.60,0.57,0.55,0.54])
FL_LAT = np.array([5.5,4.5,3.5,2.5,1.5,0.5,-0.5,-1.5,-2.5,-3.5,-4.5,-5.5])
FL_LON = np.array([95.8,96.4,97.2,98,98.8,99.5,100.2,100.9,101.5,102.2,102.9,104.2])

def idw(x, y, tx, ty, phi, reg=50.0**2):
    w = 1.0 / ((x[:,None]-tx)**2 + (y[:,None]-ty)**2 + reg)
    return (w*phi).sum(1) / w.sum(1)

def dist_polyline(x, y, fx, fy):
    md = np.full(len(x), np.inf)
    for k in range(len(fx)-1):
        ax, ay, bx, by = fx[k], fy[k], fx[k+1], fy[k+1]
        ab2 = (bx-ax)**2 + (by-ay)**2
        t = np.clip(((x-ax)*(bx-ax) + (y-ay)*(by-ay))/ab2, 0, 1)
        md = np.minimum(md, np.sqrt((ax+t*(bx-ax)-x)**2 + (ay+t*(by-ay)-y)**2))
    return md

def youngs1997(M, R_hypo):
    return 0.2418 + 1.414*M - 2.552*np.log(R_hypo + np.exp(-1.7818+0.554*M)) + 0.00607*25.

def gr_density(m, b=B_VAL):
    bl = b*np.log(10); return bl*np.exp(-bl*(m-MC)) / (1 - np.exp(-bl*(MMAX-MC)))

# ── Fitted intensity -> source rate n_dot ─────────────────────────────────────
p = pd.read_csv('lgcp_parameters.csv').set_index('Parameter')['Value']
b0, b1, b2 = p['beta0'], p['beta1_coupling'], p['beta2_fault_dist']
phi_mu, phi_sd = p['phi_mean'], p['phi_std']
df_mu, df_sd   = p['dfault_mean_km'], p['dfault_std_km']
trx, tryy = (TR_LON-lon0)*KM_LON, (TR_LAT-lat0)*KM_LAT
flx, fly  = (FL_LON-lon0)*KM_LON, (FL_LAT-lat0)*KM_LAT

slat = np.linspace(LAT_MIN, LAT_MAX, NQ_SRC); slon = np.linspace(LON_MIN, LON_MAX, NQ_SRC)
SLO, SLA = np.meshgrid(slon, slat)
sx, sy = (SLO.ravel()-lon0)*KM_LON, (SLA.ravel()-lat0)*KM_LAT
phi_q = idw(sx, sy, trx, tryy, COUP); dfl_q = dist_polyline(sx, sy, flx, fly)
lam = np.exp(np.clip(b0 + b1*(phi_q-phi_mu)/phi_sd + b2*(dfl_q-df_mu)/df_sd, -20, 20))
cell = ((LAT_MAX-LAT_MIN)*KM_LAT/NQ_SRC) * ((LON_MAX-LON_MIN)*KM_LON/NQ_SRC)
n_dot = lam * ((N_MAIN/T_OBS)/(lam*cell).sum()) * cell

# ── LGCP 475-yr PGA on the map grid ───────────────────────────────────────────
print(f"Computing fitted-grid LGCP 475-yr surface ({NG}x{NG}) at b={B_VAL} ...")
glat = np.linspace(LAT_MIN, LAT_MAX, NG); glon = np.linspace(LON_MIN, LON_MAX, NG)
GLO, GLA = np.meshgrid(glon, glat)
ex, ey = (GLO.ravel()-lon0)*KM_LON, (GLA.ravel()-lat0)*KM_LAT
H = np.zeros((len(ex), len(PGA)))
lnPGA = np.log(PGA)
for m in np.arange(4.75, 9.01, 0.5):
    fm = gr_density(m)*0.5
    sig = float(np.clip(1.45-0.1*m if m < 8 else 0.65, 0.55, 1.0))
    Rk = np.sqrt(np.maximum(np.sqrt((ex[:,None]-sx[None,:])**2 + (ey[:,None]-sy[None,:])**2), 30.)**2 + 25.**2)
    mu = youngs1997(m, Rk)
    for pi in range(len(PGA)):
        H[:,pi] += (n_dot[None,:]*fm*sp_norm.sf((lnPGA[pi]-mu)/sig)).sum(axis=1)

def pga475(row):
    t = 1/475.
    if row.max() < t: return PGA[-1]
    i = np.clip(np.searchsorted(-row, -t), 1, len(row)-1)
    x0, x1 = np.log(PGA[i-1]), np.log(PGA[i]); y0, y1 = np.log(row[i-1]+1e-30), np.log(row[i]+1e-30)
    return np.exp(x0 + (np.log(t)-y0)*(x1-x0)/(y1-y0)) if y1 != y0 else PGA[i]
lgcp_grid = np.array([pga475(H[i]) for i in range(len(ex))]).reshape(NG, NG)

# ── SNI surface (RBF from published city anchors) ─────────────────────────────
CITIES = {'Banda Aceh':(5.55,95.32,'#e74c3c'), 'Padang':(-0.95,100.35,'#3498db'),
          'Medan':(3.59,98.67,'#2ecc71'), 'Pekanbaru':(0.51,101.45,'#9b59b6')}
SNI_CITY = {'Banda Aceh':0.85, 'Padang':0.70, 'Medan':0.60, 'Pekanbaru':0.50}
pts = np.array([[la,lo] for (la,lo,_) in CITIES.values()])
vals = np.array([SNI_CITY[c] for c in CITIES])
sni_grid = np.clip(RBFInterpolator(pts, vals, kernel='linear', smoothing=0.0)(
                   np.c_[GLA.ravel(), GLO.ravel()]).reshape(NG, NG), 0.3, 0.95)

ratio_grid = np.clip(lgcp_grid / np.maximum(sni_grid, 0.05), 0.3, 3.0)

# ── Figure ────────────────────────────────────────────────────────────────────
EXTENT = [LON_MIN, LON_MAX, LAT_MIN, LAT_MAX]
halo = [withStroke(linewidth=2.5, foreground='white')]
fig, axes = plt.subplots(1, 3, figsize=(17, 6.5),
                         subplot_kw={'projection': _PC})
plt.subplots_adjust(wspace=0.32, left=0.06, right=0.95, top=0.98, bottom=0.10)

def add_base_map(ax, left_labels=False, bottom_labels=True):
    ax.set_extent(EXTENT, crs=_PC)
    ax.set_aspect('auto')
    ax.add_feature(_COAST_10M, zorder=8)
    ax.add_feature(_BORDER_10M, zorder=8)
    gl = ax.gridlines(crs=_PC, draw_labels=True, linewidth=0.3,
                      color='gray', alpha=0.45, x_inline=False,
                      y_inline=False)
    gl.top_labels = False
    gl.right_labels = False
    gl.left_labels = left_labels
    gl.bottom_labels = bottom_labels
    gl.xlabel_style = {'size': 8}
    gl.ylabel_style = {'size': 8}

def label(ax, txt, bg=None):
    kw = dict(fontsize=13, fontweight='bold', ha='right', va='top', transform=ax.transAxes, zorder=20)
    if bg: kw['bbox'] = dict(facecolor=bg, edgecolor='none', alpha=0.75, pad=2)
    ax.text(0.97, 0.97, txt, **kw)
def add_cities(ax):
    for c,(la,lo,col) in CITIES.items():
        ax.scatter(lo, la, s=55, c='white', edgecolors='k', lw=1.0,
                   transform=_PC, zorder=10)
        ax.text(lo+0.15, la+0.12, c, fontsize=7, color='black',
                path_effects=halo, transform=_PC, zorder=11)
def add_fault(ax):
    ax.plot(FL_LON, FL_LAT, 'k--', lw=1.2, alpha=0.8,
            path_effects=halo, transform=_PC, zorder=9)
def cbar(ax, im, lab, ticks=None):
    cax = inset_axes(ax, width='4%', height='100%', loc='lower left',
                     bbox_to_anchor=(1.03,0.,1,1), bbox_transform=ax.transAxes, borderpad=0)
    cb = plt.colorbar(im, cax=cax, label=lab)
    if ticks is not None: cb.set_ticks(ticks)
    return cb

# Panel A — LGCP
ax = axes[0]
imA = ax.imshow(lgcp_grid, extent=EXTENT, origin='lower', cmap='YlOrRd',
                vmin=0, vmax=1.2, aspect='auto', transform=_PC)
add_base_map(ax, left_labels=True, bottom_labels=True)
add_fault(ax); add_cities(ax); ax.set_ylabel('Latitude (°N)', fontsize=9)
cbar(ax, imA, 'PGA (g)'); label(ax, '(A)')
# Panel B — SNI
ax = axes[1]
imB = ax.imshow(sni_grid, extent=EXTENT, origin='lower', cmap='YlOrRd',
                vmin=0, vmax=1.2, aspect='auto', transform=_PC)
add_base_map(ax, left_labels=False, bottom_labels=True)
add_fault(ax); add_cities(ax); cbar(ax, imB, 'PGA (g)'); label(ax, '(B)')
# Panel C — ratio
ax = axes[2]
imC = ax.imshow(ratio_grid, extent=EXTENT, origin='lower', cmap='RdBu_r',
                norm=TwoSlopeNorm(vmin=0.3, vcenter=1.0, vmax=3.0),
                aspect='auto', transform=_PC)
add_base_map(ax, left_labels=False, bottom_labels=True)
add_fault(ax); add_cities(ax)
cs = ax.contour(GLO, GLA, ratio_grid, levels=[1.0], colors='black',
                linewidths=1.2, linestyles='--', transform=_PC, zorder=6)
ax.clabel(cs, fmt='1.0', fontsize=8, inline=True)
cbar(ax, imC, 'LGCP / SNI ratio', ticks=[0.5,1.0,1.5,2.0,2.5,3.0]); label(ax, '(C)', bg='white')

print("City-level LGCP/SNI (fitted grid):")
for c,(la,lo,_) in CITIES.items():
    ii, jj = np.unravel_index(np.argmin((GLA-la)**2 + (GLO-lo)**2), GLA.shape)
    print(f"  {c:12s} LGCP={lgcp_grid[ii,jj]:.2f}g  SNI={SNI_CITY[c]:.2f}g  ratio={lgcp_grid[ii,jj]/SNI_CITY[c]:.2f}")

OUT = 'fig_sni_comparison.png'
fig.savefig(OUT, dpi=300, bbox_inches='tight'); plt.close(fig)
print(f"Figure saved -> {OUT}")
