"""
generate_fig4_hazardmap.py
==========================
Generates Figure 4 of SERRA_Paper:
  Six-panel seismic hazard map for Sumatra (475-yr and 2475-yr PGA,
  hazard curves, disaggregation, city comparison, amplification ratio).

Output: fig4_hazard_map.png  (300 dpi, no suptitle, (A)-(F) corner labels)

Requires: numpy, scipy, pandas, matplotlib
Run from directory containing lgcp_parameters.csv and declustered_gk.csv
"""

import numpy as np
import pandas as pd
import matplotlib
import os
import cartopy
import cartopy.crs as ccrs
import cartopy.feature as cfeature
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import matplotlib.colors as mcolors
from mpl_toolkits.axes_grid1.inset_locator import inset_axes
from scipy.stats import norm as sp_norm
import warnings
warnings.filterwarnings('ignore')
from matplotlib.patheffects import withStroke as _ws

# ── Configuration ─────────────────────────────────────────────────────────────
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)
OUTPUT_FILE  = 'fig4_hazard_map.png'
DPI          = 300
LAT_MIN, LAT_MAX = -6., 6.
LON_MIN, LON_MAX = 95., 108.
lat0, lon0       = 0., 101.5
KM_LON = 111.320 * np.cos(np.radians(lat0))
KM_LAT = 111.320
NQ_SRC   = 40          # PSHA source grid
NQ_SITE  = 28          # hazard evaluation site grid
T_OBS    = 26.0
N_MAIN   = 693
B_VAL    = 1.05
MC       = 4.5
MMAX     = 9.0
MAG_BINS = np.arange(4.75, 9.01, 0.5)
PGA_LEVELS = np.logspace(np.log10(0.001), np.log10(3.0), 60)
LAM_475  = 1. / 475
LAM_2475 = 1. / 2475

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'),
}

# ── Tectonic geometry ─────────────────────────────────────────────────────────
TRENCH_LATS = 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])
TRENCH_LONS = 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])
COUPLING    = 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])
FAULT_LATS  = np.array([
     5.912,  5.737,  5.021,  4.414,  4.293,  4.117,  3.560,
     3.218,  2.180,  2.017,  1.844,  1.507,  0.704,  0.520,
     0.251,  0.205,  0.080, -0.026, -0.072, -0.555, -0.728,
    -0.903, -1.148, -1.596, -1.797, -2.193, -2.334, -2.731,
    -2.777, -3.262, -3.441, -3.577, -4.304, -4.350, -5.177,
    -5.313, -5.834, -6.154, -6.604, -6.697
])
FAULT_LONS  = np.array([
     95.326,  95.094,  95.951,  96.593,  96.813,  97.064,  97.849,
     97.987,  98.893,  98.985,  99.115,  99.331, 100.008, 100.133,
    100.054, 100.155, 100.247, 100.155, 100.247, 100.561, 100.607,
    100.727, 100.852, 101.168, 101.214, 101.606, 101.560, 101.866,
    102.034, 102.370, 102.511, 102.557, 103.210, 103.421, 104.293,
    104.339, 104.764, 104.883, 104.837, 104.764
])

def city_km(lat, lon):
    return (lon - lon0)*KM_LON, (lat - lat0)*KM_LAT

def coupling_idw(x, y):
    tx = (TRENCH_LONS - lon0)*KM_LON; ty = (TRENCH_LATS - lat0)*KM_LAT
    w  = 1./((x[:,None]-tx)**2 + (y[:,None]-ty)**2 + 50**2)
    return (w*COUPLING).sum(1)/w.sum(1)

def dist_fault(x, y):
    fx = (FAULT_LONS-lon0)*KM_LON; fy = (FAULT_LATS-lat0)*KM_LAT
    md = np.full(len(x), np.inf)
    for k in range(len(fx)-1):
        x1,y1,x2,y2 = fx[k],fy[k],fx[k+1],fy[k+1]
        L2 = (x2-x1)**2+(y2-y1)**2
        t  = np.clip(((x-x1)*(x2-x1)+(y-y1)*(y2-y1))/L2, 0, 1)
        md = np.minimum(md, np.sqrt((x-x1-t*(x2-x1))**2+(y-y1-t*(y2-y1))**2))
    return md

# ── Load LGCP parameters ──────────────────────────────────────────────────────
p      = pd.read_csv('lgcp_parameters.csv').set_index('Parameter')['Value']
b0_lgcp= p['beta0']
b1_lgcp= p['beta1_coupling']
b2_lgcp= 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']

# ── Build source grid ─────────────────────────────────────────────────────────
print("Building source grid …")
slat_s = np.linspace(LAT_MIN, LAT_MAX, NQ_SRC)
slon_s = np.linspace(LON_MIN, LON_MAX, NQ_SRC)
SSLO, SSLA = np.meshgrid(slon_s, slat_s)
sx = (SSLO.ravel()-lon0)*KM_LON; sy = (SSLA.ravel()-lat0)*KM_LAT
phi = coupling_idw(sx, sy); dfl = dist_fault(sx, sy)
eta = b0_lgcp + b1_lgcp*(phi-phi_mu)/phi_sd + b2_lgcp*(dfl-df_mu)/df_sd
lam = np.exp(np.clip(eta,-20,20))
cell_area = ((LAT_MAX-LAT_MIN)*KM_LAT/NQ_SRC)*((LON_MAX-LON_MIN)*KM_LON/NQ_SRC)
scale = (N_MAIN/T_OBS)/(lam*cell_area).sum()
n_dot = lam*scale*cell_area

# ── Build site grid ───────────────────────────────────────────────────────────
slat_e = np.linspace(LAT_MIN, LAT_MAX, NQ_SITE)
slon_e = np.linspace(LON_MIN, LON_MAX, NQ_SITE)
ESLO, ESLA = np.meshgrid(slon_e, slat_e)
site_x = (ESLO.ravel()-lon0)*KM_LON
site_y = (ESLA.ravel()-lat0)*KM_LAT
N_src, N_site = len(sx), len(site_x)

# ── GMPE ─────────────────────────────────────────────────────────────────────
def youngs1997(M, R_hypo):
    c1,c2,c3,c4,c5,c7 = 0.2418,1.414,-2.552,-1.7818,0.554,0.00607
    return c1+c2*M+c3*np.log(R_hypo+np.exp(c4+c5*M))+c7*25.

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

# ── PSHA computation ──────────────────────────────────────────────────────────
print("Running PSHA …")
hazard = np.zeros((N_site, len(PGA_LEVELS)))
dm = 0.5
for m in MAG_BINS:
    fm_dm = gr_density(m)*dm
    sig   = float(np.clip(1.45-0.1*m if m<8. else 0.65, 0.55, 1.0))
    for si in range(N_src):
        r_km   = np.sqrt((site_x-sx[si])**2+(site_y-sy[si])**2)
        r_hypo = np.sqrt(np.maximum(r_km,30.)**2+25.**2)
        ln_pga = youngs1997(m, r_hypo)
        for pi, im in enumerate(PGA_LEVELS):
            hazard[:,pi] += n_dot[si]*fm_dm*sp_norm.sf((np.log(im)-ln_pga)/sig)

def pga_at_rate(haz_row, target):
    if haz_row.max() < target: return PGA_LEVELS[-1]
    idx = np.clip(np.searchsorted(-haz_row,-target),1,len(haz_row)-1)
    x0,x1 = np.log(PGA_LEVELS[idx-1]), np.log(PGA_LEVELS[idx])
    y0,y1 = np.log(haz_row[idx-1]+1e-30), np.log(haz_row[idx]+1e-30)
    return np.exp(x0+(np.log(target)-y0)*(x1-x0)/(y1-y0)) if y1!=y0 else PGA_LEVELS[idx]

# ── City hazard curves ────────────────────────────────────────────────────────
print("Computing city hazard curves …")
city_hazard = {}; city_pga_475 = {}; city_pga_2475 = {}
for cname,(clat,clon,col) in CITIES.items():
    cx,cy = city_km(clat,clon)
    si = int(np.argmin((site_x-cx)**2+(site_y-cy)**2))
    city_hazard[cname]  = hazard[si]
    city_pga_475[cname] = pga_at_rate(hazard[si], LAM_475)
    city_pga_2475[cname]= pga_at_rate(hazard[si], LAM_2475)

city_names = list(CITIES.keys())
print("\nCity PGA summary:")
for c in city_names:
    print(f"  {c}: 475-yr = {city_pga_475[c]:.3f} g, 2475-yr = {city_pga_2475[c]:.3f} g")

# ── PGA map grids ─────────────────────────────────────────────────────────────
pga_475_map  = np.array([pga_at_rate(hazard[i], LAM_475)  for i in range(N_site)]).reshape(NQ_SITE,NQ_SITE)
pga_2475_map = np.array([pga_at_rate(hazard[i], LAM_2475) for i in range(N_site)]).reshape(NQ_SITE,NQ_SITE)
ratio_map    = pga_2475_map / np.maximum(pga_475_map, 0.001)

# ── Magnitude disaggregation (Padang, PGA > 0.1 g) ───────────────────────────
print("Disaggregation …")
cx_p,cy_p = city_km(-0.95,100.35)
si_p = int(np.argmin((site_x-cx_p)**2+(site_y-cy_p)**2))
IM_THRESH = 0.1
disagg = np.zeros(len(MAG_BINS))
for bi,m in enumerate(MAG_BINS):
    fm_dm = gr_density(m)*dm
    sig   = float(np.clip(1.45-0.1*m if m<8. else 0.65,0.55,1.0))
    for si in range(N_src):
        r_km  = np.sqrt((site_x[si_p]-sx[si])**2+(site_y[si_p]-sy[si])**2)
        r_hypo= np.sqrt(max(r_km,30.)**2+25.**2)
        ln_p  = float(youngs1997(m,np.array([r_hypo]))[0])
        disagg[bi] += n_dot[si]*fm_dm*sp_norm.sf((np.log(IM_THRESH)-ln_p)/sig)
disagg /= disagg.sum() if disagg.sum() > 0 else 1.

print("Drawing figures (split A/B/C and D/E/F) …")
EXTENT = [LON_MIN, LON_MAX, LAT_MIN, LAT_MAX]

def add_base_map(ax, left_labels=True, 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 add_fault(ax):
    ax.plot(FAULT_LONS, FAULT_LATS, 'w--', lw=1.2, alpha=0.9,
            transform=_PC, zorder=9,
            path_effects=[_ws(linewidth=2.0, foreground='black')])

def add_cities(ax, fontsize=7.5):
    for cname,(clat,clon,col) in CITIES.items():
        ax.scatter(clon, clat, s=40, c='white', edgecolors='k', lw=0.8,
                   transform=_PC, zorder=10)
        ax.text(clon+0.18, clat+0.1, cname, fontsize=fontsize,
                color='white', fontweight='normal',
                path_effects=[_ws(linewidth=1.5, foreground='black')],
                transform=_PC, zorder=11)

def corner_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.8, pad=2)
    ax.text(0.97, 0.97, txt, **kw)

# ──────────────────────────────────────────────────────────────────────────────
#  MAIN-TEXT FIGURE 4 :  (A) 475-yr map  (B) 2475-yr map  (C) hazard curves
# ──────────────────────────────────────────────────────────────────────────────
fig = plt.figure(figsize=(16, 5.5))
gs = fig.add_gridspec(1, 3)
axes = [fig.add_subplot(gs[0], projection=_PC),
        fig.add_subplot(gs[1], projection=_PC),
        fig.add_subplot(gs[2])]
plt.subplots_adjust(wspace=0.38, left=0.06, right=0.96, top=0.95, bottom=0.13)

# Panel A — 475-yr PGA
ax = axes[0]
im = ax.imshow(pga_475_map, 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)
cax = inset_axes(ax, width='4%', height='100%', loc='lower left',
                 bbox_to_anchor=(1.02,0,1,1), bbox_transform=ax.transAxes, borderpad=0)
plt.colorbar(im, cax=cax, label='PGA (g)')
corner_label(ax, '(A)')

# Panel B — 2475-yr PGA
ax = axes[1]
im2 = ax.imshow(pga_2475_map, extent=EXTENT, origin='lower',
                cmap='YlOrRd', vmin=0, vmax=2.5, aspect='auto',
                transform=_PC)
add_base_map(ax, left_labels=False, bottom_labels=True)
add_fault(ax); add_cities(ax)
cax2 = inset_axes(ax, width='4%', height='100%', loc='lower left',
                  bbox_to_anchor=(1.02,0,1,1), bbox_transform=ax.transAxes, borderpad=0)
plt.colorbar(im2, cax=cax2, label='PGA (g)')
corner_label(ax, '(B)')

# Panel C — Hazard curves
ax = axes[2]
for cname,(clat,clon,col) in CITIES.items():
    ax.semilogy(PGA_LEVELS, city_hazard[cname], lw=2, color=col, label=cname)
ax.axhline(LAM_475,  ls='--', color='grey', lw=0.9)
ax.axhline(LAM_2475, ls='--', color='grey', lw=0.9)
ax.set_xlim(0, 2.5); ax.set_ylim(1e-5, 0.1)
ax.set_xlabel('PGA (g)', fontsize=9); ax.set_ylabel('Annual exceedance rate', fontsize=9)
ax.legend(fontsize=7.5, framealpha=0.8)
corner_label(ax, '(C)')

fig.savefig('fig4_hazard_map.png', dpi=DPI, bbox_inches='tight')
plt.close(fig)
print("  saved → fig4_hazard_map.png  (A, B, C)")

# ──────────────────────────────────────────────────────────────────────────────
#  SUPPLEMENT S9 FIGURE :  (A) disaggregation  (B) city PGA  (C) amplification
#  (relabelled A/B/C for the standalone supplement figure; these were D/E/F)
# ──────────────────────────────────────────────────────────────────────────────
figs = plt.figure(figsize=(16, 5.0))
gss = figs.add_gridspec(1, 3)
axes = [figs.add_subplot(gss[0]), figs.add_subplot(gss[1]),
        figs.add_subplot(gss[2], projection=_PC)]
plt.subplots_adjust(wspace=0.30, left=0.06, right=0.96, top=0.95, bottom=0.20)

# Panel A (was D) — Disaggregation
ax = axes[0]
colors_d = plt.cm.YlOrRd(np.linspace(0.35, 0.85, len(MAG_BINS)))
ax.bar(MAG_BINS, disagg*100, width=0.4, color=colors_d, edgecolor='white', lw=0.5)
ax.axvline(4.75, ls='--', color='grey', lw=0.8, label=r'Peak $M_w=4.75$')
ax.set_xlabel('Moment magnitude $M_w$', fontsize=9)
ax.set_ylabel('Contribution to hazard (%)', fontsize=9)
ax.set_xticks(MAG_BINS); ax.set_xticklabels([f'{m:.2f}' for m in MAG_BINS], rotation=45, fontsize=7)
corner_label(ax, '(A)')

# Panel B (was E) — City PGA comparison
ax = axes[1]
x = np.arange(len(city_names)); w = 0.35
pga_475_vals  = [city_pga_475[c]  for c in city_names]
pga_2475_vals = [city_pga_2475[c] for c in city_names]
ax.bar(x-w/2, pga_475_vals,  w, label='475-yr',  color='#f39c12', edgecolor='white')
ax.bar(x+w/2, pga_2475_vals, w, label='2475-yr', color='#c0392b', edgecolor='white')
for i,(p4,p24) in enumerate(zip(pga_475_vals, pga_2475_vals)):
    ax.text(i-w/2, p4+0.03,  f'{p4:.2f}',  ha='center', fontsize=7.5)
    ax.text(i+w/2, p24+0.03, f'{p24:.2f}', ha='center', fontsize=7.5)
ax.set_xticks(x); ax.set_xticklabels(city_names, fontsize=8.5, rotation=12)
ax.set_ylabel('PGA (g)', fontsize=9); ax.legend(fontsize=8.5)
corner_label(ax, '(B)')

# Panel C (was F) — Amplification ratio
ax = axes[2]
mean_ratio = float(np.nanmean(ratio_map))
im3 = ax.imshow(ratio_map, extent=EXTENT, origin='lower',
                cmap='Blues', vmin=1.0, vmax=3.5, aspect='auto',
                transform=_PC)
add_base_map(ax, left_labels=True, bottom_labels=True)
add_fault(ax)
for cname,(clat,clon,col) in CITIES.items():
    ax.scatter(clon, clat, s=40, c='orange', edgecolors='k', lw=0.8,
               transform=_PC, zorder=10)
cax3 = inset_axes(ax, width='4%', height='100%', loc='lower left',
                  bbox_to_anchor=(1.02,0,1,1), bbox_transform=ax.transAxes, borderpad=0)
plt.colorbar(im3, cax=cax3, label='PGA$_{2475}$ / PGA$_{475}$')
ax.text(0.05, 0.05, f'Mean ratio: {mean_ratio:.2f}×',
        transform=ax.transAxes, fontsize=8.5, color='white', va='bottom',
        bbox=dict(facecolor='#2c3e50', alpha=0.7, pad=3))
corner_label(ax, '(C)', bg='white')

figs.savefig('figS9_hazard_disagg.png', dpi=DPI, bbox_inches='tight')
plt.close(figs)
print("  saved → figS9_hazard_disagg.png  (was D, E, F)")
print("\nDone.")
