#!/usr/bin/env python3
"""
verify_holdout_relm.py
======================
Reproduces (1) the temporal hold-out spatial validation of main-text
Sect. 3.5 (L-Test, R-Test, N-Test on the 2016-2024 test window), (2) the
segment-level RELM evaluation of Online Supplement S7 (20 Great Sumatran
Fault segments, +/-30 km buffers), and (3) the Padang site hazard table
(main-text Table 5 / Supplement Table S9.1).

Covariates, grid, projection, and Poisson log-linear MLE are identical to
sumatra_lgcp.py. Segment endpoints follow Sieh & Natawidjaja (2000):
latitude bounds from their Table 1; longitudes obtained by interpolating
the GEM/EOS Southeast Asia fault-trace polyline used throughout this
package, so that the segment geometry is fully reproducible from the
files distributed here. Batee (splay) and Barumun (NE strand) are
represented by straight-line endpoint pairs (see SEGMENTS below).

Run:  python verify_holdout_relm.py
"""

import numpy as np
import pandas as pd

RNG_HOLDOUT = np.random.default_rng(42)
RNG_RELM    = np.random.default_rng(42)

# ── Domain & projection (identical to sumatra_lgcp.py) ──────────────────────
GK_CSV = "declustered_gk.csv"
LAT_MIN, LAT_MAX = -6.0, 6.0
LON_MIN, LON_MAX = 95.0, 108.0
lat0, lon0 = 0.0, 101.5
KM_LON = 111.320 * np.cos(np.radians(lat0))
KM_LAT = 111.320
X_MIN = (LON_MIN - lon0) * KM_LON
X_MAX = (LON_MAX - lon0) * KM_LON
Y_MIN = (LAT_MIN - lat0) * KM_LAT
Y_MAX = (LAT_MAX - lat0) * KM_LAT
W_AREA = (X_MAX - X_MIN) * (Y_MAX - Y_MIN)
MC_LGCP = 4.5
NQ = 60
MAX_ITER, STEP_DAMP = 40, 0.5

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 to_km(lats, lons):
    return (np.asarray(lons) - lon0) * KM_LON, (np.asarray(lats) - lat0) * KM_LAT

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

def dist_polyline(x_km, y_km, fx, fy):
    md = np.full(len(x_km), 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_km - x1) * (x2 - x1) + (y_km - y1) * (y2 - y1)) / L2, 0, 1)
        md = np.minimum(md, np.sqrt((x_km - x1 - t * (x2 - x1)) ** 2 +
                                    (y_km - y1 - t * (y2 - y1)) ** 2))
    return md

def fault_lon_at_lat(lat):
    """Longitude of the fault trace polyline at a given latitude
    (trace is monotone in latitude from NW to SE)."""
    return np.interp(lat, FAULT_LATS[::-1], FAULT_LONS[::-1])

# ── Load catalogue, build covariates (identical pipeline) ───────────────────
df_raw = pd.read_csv(GK_CSV)
mask_sp = ((df_raw['Latitude'] >= LAT_MIN) & (df_raw['Latitude'] <= LAT_MAX) &
           (df_raw['Longitude'] >= LON_MIN) & (df_raw['Longitude'] <= LON_MAX))
mask_ms = (df_raw['is_mainshock'] == 1) & (df_raw['Magnitude'] >= MC_LGCP)
ev = df_raw[mask_sp & mask_ms].copy().reset_index(drop=True)
ev['x'], ev['y'] = to_km(ev['Latitude'].values, ev['Longitude'].values)
n_ev = len(ev)
print(f"Mainshocks M>={MC_LGCP}: {n_ev}")

TRENCH_X, TRENCH_Y = to_km(TRENCH_LATS, TRENCH_LONS)
FAULT_X, FAULT_Y = to_km(FAULT_LATS, FAULT_LONS)

xq_1d = np.linspace(X_MIN, X_MAX, NQ)
yq_1d = np.linspace(Y_MIN, Y_MAX, NQ)
XQ, YQ = np.meshgrid(xq_1d, yq_1d)
xq_f, yq_f = XQ.ravel(), YQ.ravel()
cell_area = W_AREA / (NQ * NQ)

phi_q    = coupling_idw(xq_f, yq_f, TRENCH_X, TRENCH_Y, COUPLING)
dfault_q = dist_polyline(xq_f, yq_f, FAULT_X, FAULT_Y)
phi_mu, phi_sd = phi_q.mean(), phi_q.std()
df_mu, df_sd   = dfault_q.mean(), dfault_q.std()
X_q = np.column_stack([np.ones(NQ * NQ),
                       (phi_q - phi_mu) / phi_sd,
                       (dfault_q - df_mu) / df_sd])

def event_design(x, y):
    phi = coupling_idw(x, y, TRENCH_X, TRENCH_Y, COUPLING)
    dfl = dist_polyline(x, y, FAULT_X, FAULT_Y)
    return np.column_stack([np.ones(len(x)),
                            (phi - phi_mu) / phi_sd,
                            (dfl - df_mu) / df_sd])

def fit_poisson(X_ev_local):
    beta = np.array([np.log(len(X_ev_local) / W_AREA), 0.0, 0.0])
    for _ in range(MAX_ITER):
        lam_q = np.exp(np.clip(X_q @ beta, -20, 20))
        score = X_ev_local.sum(0) - cell_area * (X_q * lam_q[:, None]).sum(0)
        FI = (X_q.T * (lam_q * cell_area)[None, :]) @ X_q
        step = np.linalg.solve(FI + 1e-8 * np.eye(3), score)
        beta = beta + STEP_DAMP * step
        if np.abs(step).max() < 1e-7:
            break
    return beta

# ════════════════════════════════════════════════════════════════════════════
# PART 1 — TEMPORAL HOLD-OUT (main text Sect. 3.5)
# ════════════════════════════════════════════════════════════════════════════
print("\n=== PART 1: temporal hold-out (train 1998-2015, test 2016-2024) ===")
tr = ev[(ev['Year'] >= 1998) & (ev['Year'] <= 2015)]
te = ev[(ev['Year'] >= 2016) & (ev['Year'] <= 2024)]
print(f"n_train = {len(tr)},  n_test = {len(te)}")

X_tr = event_design(tr['x'].values, tr['y'].values)
beta_tr = fit_poisson(X_tr)
print(f"training betas: b0={beta_tr[0]:.3f}, b1={beta_tr[1]:+.3f}, b2={beta_tr[2]:+.3f}")

lam_q_tr = np.exp(np.clip(X_q @ beta_tr, -20, 20))      # intensity, events/km^2
Lam_tr = (lam_q_tr * cell_area).sum()                    # integrates to ~n_train

# R-Test: per-event log-density ratio of LGCP forecast vs uniform null
X_te = event_design(te['x'].values, te['y'].values)
lam_te = np.exp(np.clip(X_te @ beta_tr, -20, 20))
dLL_holdout = np.sum(np.log(lam_te / Lam_tr) - np.log(1.0 / W_AREA))
print(f"R-Test  Delta-LL (LGCP vs uniform, {len(te)} test events): {dLL_holdout:+.1f}")

# L-Test: forecast normalised to the observed test count, gridded Poisson LL
n_test = len(te)
lam_cell = lam_q_tr * cell_area * (n_test / Lam_tr)      # expected counts/cell
ix = np.clip(np.searchsorted(xq_1d, te['x'].values) - 1, 0, NQ - 1)
iy = np.clip(np.searchsorted(yq_1d, te['y'].values) - 1, 0, NQ - 1)
obs_cell = np.zeros(NQ * NQ)
np.add.at(obs_cell, iy * NQ + ix, 1)

from scipy.special import gammaln
def pois_ll(n, lam):
    return np.sum(n * np.log(lam + 1e-300) - lam - gammaln(n + 1))

LL_obs = pois_ll(obs_cell, lam_cell)
N_SIM = 5000
sims = RNG_HOLDOUT.poisson(lam_cell, size=(N_SIM, NQ * NQ))
LL_sim = (sims * np.log(lam_cell + 1e-300)).sum(1) - lam_cell.sum() \
         - gammaln(sims + 1).sum(1)
gamma_q = float((LL_sim <= LL_obs).mean())
print(f"L-Test  quantile gamma (N_sim={N_SIM}): {gamma_q:.2f}")

# N-Test: rate forecast scaled by time
T_train, T_test = 18.0, 9.0
n_forecast = Lam_tr / T_train * T_test
print(f"N-Test  forecast count {n_forecast:.0f} vs observed {n_test} "
      f"-> over-prediction x{n_forecast / n_test:.1f}; "
      f"test-window rate {n_test / T_test:.1f} events/yr")

# ════════════════════════════════════════════════════════════════════════════
# PART 2 — SEGMENT-LEVEL RELM (Online Supplement S7)
# ════════════════════════════════════════════════════════════════════════════
print("\n=== PART 2: RELM, 20 GSF segments, +/-30 km buffers ===")
# Latitude bounds from Sieh & Natawidjaja (2000, Table 1); longitudes from
# the package fault-trace polyline at those latitudes. Batee (splay) and
# Barumun (NE strand) use explicit endpoint pairs.
SEG_LAT_BOUNDS = {
    'Seulimeum': (5.0, 5.9), 'Aceh': (4.4, 5.4), 'Tripa': (3.4, 4.4),
    'Renun': (2.0, 3.55), 'Toru': (1.2, 2.0), 'Angkola': (0.3, 1.8),
    'Sumpur': (0.1, 0.3), 'Sianok': (-0.7, 0.1), 'Sumani': (-1.0, -0.5),
    'Suliti': (-1.75, -1.0), 'Siulak': (-2.25, -1.7), 'Dikit': (-2.75, -2.3),
    'Ketaun': (-3.35, -2.75), 'Musi': (-3.65, -3.25), 'Manna': (-4.35, -3.8),
    'Kumering': (-5.3, -4.35), 'Semangko': (-5.9, -5.25), 'Sunda': (-6.75, -5.9),
}
SEGMENTS = {}
for name, (la1, la2) in SEG_LAT_BOUNDS.items():
    SEGMENTS[name] = ((la1, fault_lon_at_lat(la1)), (la2, fault_lon_at_lat(la2)))
# Barumun: NE strand parallel to Angkola (offset ~0.35 deg lon to the NE)
SEGMENTS['Barumun'] = ((0.3, fault_lon_at_lat(0.3) + 0.35),
                       (1.2, fault_lon_at_lat(1.2) + 0.35))
# Batee: splay diverging SSE from the main trace near 4.35N to the coast
SEGMENTS['Batee'] = ((4.35, fault_lon_at_lat(4.35)), (3.0, 97.4))

BUF = 30.0  # km
lam_q_full = np.exp(np.clip(X_q @ fit_poisson(event_design(ev['x'].values,
                                                           ev['y'].values)), -20, 20))
lam_bar = n_ev / W_AREA

rows = []
for name, ((la1, lo1), (la2, lo2)) in SEGMENTS.items():
    (x1,), (y1,) = to_km([la1], [lo1])
    (x2,), (y2,) = to_km([la2], [lo2])
    fx, fy = np.array([x1, x2]), np.array([y1, y2])
    d_cells = dist_polyline(xq_f, yq_f, fx, fy)
    in_buf = d_cells <= BUF
    A_j = in_buf.sum() * cell_area
    lam_L = (lam_q_full[in_buf] * cell_area).sum()
    lam_U = lam_bar * A_j
    d_ev = dist_polyline(ev['x'].values, ev['y'].values, fx, fy)
    w_j = int((d_ev <= BUF).sum())
    rows.append((name, A_j, lam_L, lam_U, w_j))

seg = pd.DataFrame(rows, columns=['segment', 'A_km2', 'lam_lgcp', 'lam_unif', 'omega'])
seg['OE_unif'] = seg['omega'] / seg['lam_unif']

def cls(r):
    return ('Highly Active' if r >= 3.0 else 'Active' if r >= 1.0
            else 'Background' if r >= 0.5 else 'Quiescent')
seg['class'] = seg['OE_unif'].apply(cls)
print(seg[['segment', 'lam_lgcp', 'lam_unif', 'omega', 'OE_unif', 'class']]
      .to_string(index=False, float_format=lambda v: f'{v:8.2f}'))

lamL = seg['lam_lgcp'].values
lamU = seg['lam_unif'].values
w    = seg['omega'].values.astype(float)

dLL_relm = (pois_ll(w, lamL) - pois_ll(w, lamU))
N_SIM_RELM = 2000
w_sim = RNG_RELM.poisson(lamU, size=(N_SIM_RELM, len(lamU))).astype(float)
dLL_sim = ((w_sim * np.log(lamL + 1e-300)).sum(1) - lamL.sum()
           - (w_sim * np.log(lamU + 1e-300)).sum(1) + lamU.sum())
alpha_R = float((dLL_sim <= dLL_relm).mean())

LL_obs_relm = pois_ll(w, lamL)
wL_sim = RNG_RELM.poisson(lamL, size=(N_SIM_RELM, len(lamL))).astype(float)
LL_sim_relm = ((wL_sim * np.log(lamL + 1e-300)).sum(1) - lamL.sum()
               - gammaln(wL_sim + 1).sum(1))
gamma_relm = float((LL_sim_relm <= LL_obs_relm).mean())
N_tot_sim = RNG_RELM.poisson(lamL.sum(), size=N_SIM_RELM)
delta_relm = float((N_tot_sim <= w.sum()).mean())

print(f"\nRELM R-Test:  Delta-LL = {dLL_relm:+.1f},  alpha_R = {alpha_R:.3f}")
print(f"RELM L-Test:  gamma = {gamma_relm:.2f}   (N_sim = {N_SIM_RELM})")
print(f"RELM N-Test:  delta = {delta_relm:.2f}   "
      f"(forecast {lamL.sum():.0f} vs observed {int(w.sum())} in buffers)")

# ════════════════════════════════════════════════════════════════════════════
# PART 3 — PADANG HAZARD TABLE (main text Table 5 / Supplement Table S9.1)
# ════════════════════════════════════════════════════════════════════════════
print("\n=== PART 3: Padang site hazard (same pipeline as generate_fig4_hazard.py) ===")
from scipy.stats import norm as sp_norm
NQ_SRC, NQ_SITE = 40, 28
T_OBS, N_MAIN, B_VAL, MC, MMAX = 26.0, 693, 1.05, 4.5, 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)

p = pd.read_csv('lgcp_parameters.csv').set_index('Parameter')['Value']
slat = np.linspace(LAT_MIN, LAT_MAX, NQ_SRC)
slon = np.linspace(LON_MIN, LON_MAX, NQ_SRC)
SSLO, SSLA = np.meshgrid(slon, slat)
sx = (SSLO.ravel() - lon0) * KM_LON
sy = (SSLA.ravel() - lat0) * KM_LAT
phi_s = coupling_idw(sx, sy, TRENCH_X, TRENCH_Y, COUPLING)
dfl_s = dist_polyline(sx, sy, FAULT_X, FAULT_Y)
eta_s = (p['beta0'] + p['beta1_coupling'] * (phi_s - p['phi_mean']) / p['phi_std']
         + p['beta2_fault_dist'] * (dfl_s - p['dfault_mean_km']) / p['dfault_std_km'])
lam_s = np.exp(np.clip(eta_s, -20, 20))
ca_src = ((LAT_MAX - LAT_MIN) * KM_LAT / NQ_SRC) * ((LON_MAX - LON_MIN) * KM_LON / NQ_SRC)
n_dot = lam_s * ((N_MAIN / T_OBS) / (lam_s * ca_src).sum()) * ca_src

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
cx, cy = (100.35 - lon0) * KM_LON, (-0.95 - lat0) * KM_LAT
si = int(np.argmin((site_x - cx) ** 2 + (site_y - cy) ** 2))

def youngs1997(M, R):
    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 + np.exp(c4 + c5 * M)) + c7 * 25.0

def gr_density(m, b=B_VAL):
    beta = b * np.log(10.0)
    return beta * np.exp(-beta * (m - MC)) / (1.0 - np.exp(-beta * (MMAX - MC)))

haz = np.zeros(len(PGA_LEVELS))
for m in MAG_BINS:
    fm_dm = gr_density(m) * 0.5
    sig = float(np.clip(1.45 - 0.1 * m if m < 8.0 else 0.65, 0.55, 1.0))
    r = np.sqrt((site_x[si] - sx) ** 2 + (site_y[si] - sy) ** 2)
    r_hypo = np.sqrt(np.maximum(r, 30.0) ** 2 + 25.0 ** 2)
    ln_pga = youngs1997(m, r_hypo)
    for pi, im in enumerate(PGA_LEVELS):
        haz[pi] += (n_dot * fm_dm * sp_norm.sf((np.log(im) - ln_pga) / sig)).sum()

def pga_at_rate(hrow, target):
    idx = np.clip(np.searchsorted(-hrow, -target), 1, len(hrow) - 1)
    x0, x1 = np.log(PGA_LEVELS[idx - 1]), np.log(PGA_LEVELS[idx])
    y0, y1 = np.log(hrow[idx - 1] + 1e-30), np.log(hrow[idx] + 1e-30)
    return float(np.exp(x0 + (np.log(target) - y0) * (x1 - x0) / (y1 - y0)))

print("T_R (yr)   PGA (g)   Low_ep   High_ep")
for TR in [43, 72, 224, 475, 975, 1500, 2475, 4975]:
    g = pga_at_rate(haz, 1.0 / TR)
    print(f"{TR:7d}   {g:7.3f}   {0.7 * g:6.3f}   {1.3 * g:7.3f}")
