import numpy as np
import pywt
from scipy.interpolate import griddata
from scipy.stats import pearsonr, spearmanr
import matplotlib.pyplot as plt

rng = np.random.default_rng(0)
NX = NY = 64; N = NX*NY
BETA = 2.5                      # potential-field truth: power spectrum ~ k^-BETA (before attenuation)
Z_SRC = 1.5                     # source depth (cells): upward-continuation attenuation exp(-2*pi*k*Z_SRC)
N_MODES = 12                    # Fourier-sparse truth: number of modes (low-frequency biased)
STRIDE = 3                      # regular sparse stride -> delta ~ 0.118
WAVELET, LEVELS, MODE = 'sym8', 2, 'periodization'   # 2 = pywt.dwt_max_level(64, 'sym8')
SYM_SPINS = 16                  # cycle-spin shifts for the symlet reconstruction
NBR, NBA, RMAX, BLK = 45, 36, 90.0, 4   # BLK=4 ~ paper's patch size A/N_phi (2.9 px, non-integer)
SA_ITERS = 4000
POCS_ITERS, POCS_LMIN = 800, 1e-4
LOO_M = 60                      # leave-one-out calibration points
ENS_K = 128                     # ensemble validation realizations

# ---------------- ground truths (two tiers) ----------------
def powerlaw_field(beta, z_src=Z_SRC):
    # random-phase power law with source-depth attenuation: without the exp(-2*pi*k*z) factor the
    # field is rough at pixel scale, which real (upward-continued) potential-field data never is
    kx = np.fft.fftfreq(NX)[:, None]; ky = np.fft.fftfreq(NY)[None, :]
    k = np.hypot(kx, ky); k[0, 0] = 1.0
    amp = k**(-beta/2)*np.exp(-2*np.pi*k*z_src); amp[0, 0] = 0.0
    f = np.real(np.fft.ifft2(amp*np.exp(2j*np.pi*rng.random((NX, NY)))))
    return (f - f.mean())/f.std()

def fourier_sparse_field(n_modes):
    # visually structured sparse signal: amplitude ~ |k|^-1/2 so low frequencies dominate the look,
    # with half the modes beyond the regular grid's Nyquist 1/(2*STRIDE) so that grid aliases
    k_alias = 1/(2*STRIDE)
    C = np.zeros((NX, NY), complex); picked = set()
    m = 0
    while m < n_modes:
        lo, hi = ((k_alias, 0.30) if m < n_modes//2 else (0.03, k_alias))
        kmag = np.exp(rng.uniform(np.log(lo), np.log(hi)))
        ang = rng.uniform(0, 2*np.pi)
        i = int(round(kmag*np.cos(ang)*NX)) % NX
        j = int(round(kmag*np.sin(ang)*NY)) % NY
        if (i, j) in picked or (i, j) == (0, 0):
            continue
        picked.add((i, j))
        C[i, j] = kmag**-0.5*np.exp(2j*np.pi*rng.random())
        m += 1
    f = np.fft.ifft2(C).real
    return (f - f.mean())/f.std()

truth_pl = powerlaw_field(BETA)           # tier 2: hard, not l1-sparse in any basis
truth_sp = fourier_sparse_field(N_MODES)  # tier 1: exactly sparse, the paper's Fig. A.1 setting

# ---------------- sampling patterns ----------------
mask_reg = np.zeros((NX, NY), bool); mask_reg[::STRIDE, ::STRIDE] = True
NS = int(mask_reg.sum()); DELTA = NS/N

def random_mask():
    m = np.zeros(N, bool); m[rng.choice(N, NS, replace=False)] = True
    return m.reshape(NX, NY)

# ---------------- ISA properties (Eq. 4 ingredients) ----------------
def H(p):
    p = np.abs(p); nz = p > 0
    return -np.sum(p[nz]*np.log(p[nz]))

def dense_pair_hists(nx, ny):
    # all-pairs interval/angle histograms of the dense grid via closed-form pair counts
    DX, DY = np.meshgrid(np.arange(-(nx-1), nx), np.arange(0, ny), indexing='ij')
    keep = (DY > 0) | ((DY == 0) & (DX > 0))
    DX, DY = DX[keep], DY[keep]
    w = (nx - np.abs(DX))*(ny - np.abs(DY))
    hr, _ = np.histogram(np.hypot(DX, DY), NBR, (0, RMAX), weights=w)
    # angles folded to [0, pi): the all-pairs rose diagram is 180-degree symmetric
    ha, _ = np.histogram(np.arctan2(DY, DX) % np.pi, NBA, (0, np.pi), weights=w)
    return hr/hr.max(), ha/ha.max()

def sparse_pair_hists(mask):
    c = np.argwhere(mask).astype(float)
    d = c[:, None, :] - c[None, :, :]
    iu = np.triu_indices(len(c), 1)
    dx, dy = d[..., 0][iu], d[..., 1][iu]
    hr, _ = np.histogram(np.hypot(dx, dy), NBR, (0, RMAX))
    ha, _ = np.histogram(np.arctan2(dy, dx) % np.pi, NBA, (0, np.pi))
    return hr/hr.max(), ha/ha.max()

def block_density(mask):
    # per-block occupancy fraction gamma_phi; compared against delta*gamma_theta = DELTA (Eq. 2)
    return mask.reshape(NX//BLK, BLK, NY//BLK, BLK).sum((1, 3))/BLK**2

def srf(mask):
    P = np.abs(np.fft.fft2(mask.astype(float))); P /= P[0, 0]
    xi = P.ravel()[1:]
    return xi.max(), xi

HR0, HA0 = dense_pair_hists(NX, NY)

def objective(mask):
    # Eq. 4 with the paper's 2D weights w_alpha = w_beta = w_gamma = mu
    hr, ha = sparse_pair_hists(mask)
    mu, xi = srf(mask)
    return mu*(H(hr - HR0) + H(ha - HA0) + H(block_density(mask) - DELTA) + H(xi))

def ergodic_mask(iters=SA_ITERS, T0=0.05):
    m = random_mask(); f = objective(m); best = (f, m.copy())
    for it in range(iters):
        T = T0*(1 - it/iters)
        on = np.flatnonzero(m.ravel()); off = np.flatnonzero(~m.ravel())
        i, j = rng.choice(on), rng.choice(off)
        m.flat[i] = False; m.flat[j] = True
        f2 = objective(m)
        if f2 < f or rng.random() < np.exp(-(f2-f)/max(T, 1e-9)):
            f = f2
            if f < best[0]: best = (f, m.copy())
        else:
            m.flat[i] = True; m.flat[j] = False
    return best[1], best[0]

mask_erg, f_erg = ergodic_mask()
f_rand = np.mean([objective(random_mask()) for _ in range(20)])
print(f"N_samples={NS} (delta={DELTA:.3f})  obj: regular={objective(mask_reg):.3f}  "
      f"random(mean)={f_rand:.3f}  ergodic={f_erg:.3f}")

# ---------------- transforms ----------------
_, SLICES = pywt.coeffs_to_array(pywt.wavedec2(np.zeros((NX, NY)), WAVELET, mode=MODE, level=LEVELS))
def Wsym(x):  return pywt.coeffs_to_array(pywt.wavedec2(x, WAVELET, mode=MODE, level=LEVELS))[0]
def Wtsym(c): return pywt.waverec2(pywt.array_to_coeffs(c, SLICES, output_format='wavedec2'), WAVELET, mode=MODE)
SYM_APP = np.zeros((NX, NY), bool); SYM_APP[SLICES[0]] = True    # approximation subband
def Wfft(x):  return np.fft.fft2(x)
def Wtfft(c): return np.fft.ifft2(c).real

# ---------------- CS reconstruction: iterative thresholding + data reinsertion (POCS) ----------------
def cs_reconstruct(y_img, mask, W=Wfft, Wt=Wtfft, n_iter=POCS_ITERS, lam_min_frac=POCS_LMIN):
    x = y_img.copy()
    lam_max = np.abs(W(y_img)).max()
    for lam in np.geomspace(lam_max, lam_max*lam_min_frac, n_iter):
        c = W(x)
        c *= np.maximum(1 - lam/np.maximum(np.abs(c), 1e-300), 0)   # soft threshold on magnitude
        x = Wt(c)
        x[mask] = y_img[mask]
    return x

# symlet variant (the paper's transform), tuned to reproduce the paper's accuracy in its regime:
# never threshold the approximation subband, and average reconstructions over circular shifts
# (cycle spinning) to remove the decimated transform's shift-variance
def cs_reconstruct_sym(y_img, mask, n_iter=POCS_ITERS, lam_min_frac=POCS_LMIN, n_shift=SYM_SPINS):
    def single(y2, m2):
        x = y2.copy()
        lam_max = np.abs(np.where(SYM_APP, 0, Wsym(y2))).max()
        for lam in np.geomspace(lam_max, lam_max*lam_min_frac, n_iter):
            c = Wsym(x)
            t = np.sign(c)*np.maximum(np.abs(c) - lam, 0)
            x = Wtsym(np.where(SYM_APP, c, t))
            x[m2] = y2[m2]
        return x
    recs = []
    for _ in range(n_shift):
        s = (int(rng.integers(NX)), int(rng.integers(NY)))
        xs = single(np.roll(y_img, s, (0, 1)), np.roll(mask, s, (0, 1)))
        recs.append(np.roll(xs, (-s[0], -s[1]), (0, 1)))
    x = np.mean(recs, 0)
    x[mask] = y_img[mask]
    return x

# ---------------- kriging / GP: Bayes-optimal for the power-law field class ----------------
KXF = np.fft.fftfreq(NX)[:, None]; KYF = np.fft.fftfreq(NY)[None, :]
KRAD = np.hypot(KXF, KYF)
GRID = np.argwhere(np.ones((NX, NY), bool))

def circulant_cov(beta, z_src):
    # power spectrum k^-beta * exp(-4*pi*k*z): the amplitude-attenuated generator's power spectrum
    S = np.where(KRAD > 0, KRAD, 1.0)**(-beta)*np.exp(-4*np.pi*KRAD*z_src); S[0, 0] = 0.0
    c = np.fft.ifft2(S).real
    return c/c[0, 0]

def gp_matrices(mask, beta, z_src):
    on = np.argwhere(mask)
    cov = circulant_cov(beta, z_src)
    Kss = cov[(on[:, None, 0]-on[None, :, 0]) % NX, (on[:, None, 1]-on[None, :, 1]) % NY]
    Kxs = cov[(GRID[:, None, 0]-on[None, :, 0]) % NX, (GRID[:, None, 1]-on[None, :, 1]) % NY]
    return Kss + 1e-6*np.eye(len(on)), Kxs

def gp_fit_params(y_img, mask):
    # truth-free ML over spectral slope and source depth, variance profiled out
    yv = y_img[mask]; best = None
    for beta in np.arange(1.5, 4.01, 0.5):
        for z_src in (0.0, 0.5, 1.0, 1.5, 2.0, 3.0):
            Kss, _ = gp_matrices(mask, beta, z_src)
            try:
                L = np.linalg.cholesky(Kss)
            except np.linalg.LinAlgError:
                continue
            a = np.linalg.solve(L, yv)
            ll = -0.5*len(yv)*np.log(a @ a/len(yv)) - np.log(np.diag(L)).sum()
            if best is None or ll > best[0]: best = (ll, beta, z_src)
    return best[1], best[2]

def gp_reconstruct(y_img, mask, beta, z_src):
    yv = y_img[mask]
    Kss, Kxs = gp_matrices(mask, beta, z_src)
    w = np.linalg.solve(Kss, yv)
    xh = (Kxs @ w).reshape(NX, NY)
    A = np.linalg.solve(Kss, Kxs.T)
    var = np.maximum(1.0 - np.einsum('ij,ji->i', Kxs, A), 0)
    s2 = yv @ np.linalg.solve(Kss, yv)/len(yv)          # profiled variance scale
    return xh, np.sqrt(var*s2).reshape(NX, NY)

# ---------------- acquire + reconstruct ----------------
def leakage(x, truth): return np.corrcoef((x - truth).ravel(), truth.ravel())[0, 1]
def rms(x, truth):     return np.sqrt(np.mean((x - truth)**2))

def spline_interp(truth, mask):
    pts = np.argwhere(mask); X, Y = np.mgrid[0:NX, 0:NY]
    xr = griddata(pts, truth[mask], (X, Y), method='cubic')
    nn = np.isnan(xr)
    xr[nn] = griddata(pts, truth[mask], (X, Y), method='nearest')[nn]
    return xr

def fourier_oracle(truth, k=NS):
    C = np.fft.fft2(truth).ravel()
    C[np.argsort(np.abs(C))[:-k]] = 0
    return np.fft.ifft2(C.reshape(NX, NY)).real

# tier 1: Fourier-sparse truth: the CS claim of the paper (Fig. A.1 / Figs. 15, 22)
print("\n--- tier 1: Fourier-sparse truth (paper's CS regime) ---")
x_sp_spl = spline_interp(truth_sp, mask_reg)
x_sp_erg_spl = spline_interp(truth_sp, mask_erg)
x_sp_reg = cs_reconstruct(truth_sp*mask_reg, mask_reg)
x_sp_erg = cs_reconstruct(truth_sp*mask_erg, mask_erg)
x_sp_sym = cs_reconstruct_sym(truth_sp*mask_erg, mask_erg)   # paper's literal transform
# kriging with the power-law covariance is misspecified for a 12-mode signal; shown as the
# symmetric counterpart to CS-on-the-random-field so every process appears on both truths
GP_SP = gp_fit_params(truth_sp*mask_erg, mask_erg)
x_sp_gp, sig_sp_raw = gp_reconstruct(truth_sp*mask_erg, mask_erg, *GP_SP)
x_sp_reg_gp = gp_reconstruct(truth_sp*mask_reg, mask_reg, *GP_SP)[0]
for name, x in [('regular+spline', x_sp_spl), ('regular+CS', x_sp_reg),
                ('regular+kriging', x_sp_reg_gp), ('ergodic+spline', x_sp_erg_spl),
                ('ergodic+CS', x_sp_erg), ('ergodic+CS-symlet', x_sp_sym),
                ('ergodic+kriging', x_sp_gp)]:
    print(f"{name:18s} leakage={leakage(x, truth_sp):+.3f}  rms={rms(x, truth_sp):.3f}")

# tier 2: potential-field truth: not l1-sparse; report against the oracle floor
print(f"\n--- tier 2: potential-field truth (beta={BETA}, source depth {Z_SRC}) ---")
print(f"oracle {NS}-term Fourier floor: rms={rms(fourier_oracle(truth_pl), truth_pl):.3f}")
BETA_HAT, Z_HAT = gp_fit_params(truth_pl*mask_erg, mask_erg)
print(f"GP ML estimates: beta_hat={BETA_HAT:.1f} z_hat={Z_HAT:.1f} "
      f"(generator beta={BETA}, z={Z_SRC})")
y_erg = truth_pl*mask_erg
x_pl_spl = spline_interp(truth_pl, mask_reg)
x_pl_erg_spl = spline_interp(truth_pl, mask_erg)
x_pl_reg = cs_reconstruct(truth_pl*mask_reg, mask_reg)
x_pl_reg_gp = gp_reconstruct(truth_pl*mask_reg, mask_reg, BETA_HAT, Z_HAT)[0]
x_erg = cs_reconstruct(y_erg, mask_erg)                      # primary reconstruction (paper's method)
x_erg_sym = cs_reconstruct_sym(y_erg, mask_erg)              # symlet-wavelet comparison
x_gp, sig_gp_raw = gp_reconstruct(y_erg, mask_erg, BETA_HAT, Z_HAT)
print("method            leakage  rms    -rms/std(truth)   (leakage ~ -rms/std when error is unrecovered signal)")
for name, x in [('regular+spline', x_pl_spl),
                ('regular+CS', x_pl_reg),
                ('regular+kriging', x_pl_reg_gp),
                ('ergodic+spline', x_pl_erg_spl),
                ('ergodic+CS-Fourier', x_erg),
                ('ergodic+CS-symlet', x_erg_sym),
                ('ergodic+kriging', x_gp)]:
    print(f"{name:18s} {leakage(x, truth_pl):+.3f}  {rms(x, truth_pl):.3f}   {-rms(x, truth_pl)/truth_pl.std():+.3f}")

# ---------------- predicted sigma (truth-free: from y + mask only) ----------------
# LOO calibration: reconstruct without one sample, compare prediction at that point to its value
def loo_rms(y_img, mask, rec_fn, m=LOO_M):
    on = np.argwhere(mask)
    rs = []
    for i in rng.choice(len(on), m, replace=False):
        m2 = mask.copy(); m2[tuple(on[i])] = False
        rs.append(rec_fn(y_img*m2, m2)[tuple(on[i])] - y_img[tuple(on[i])])
    return np.sqrt(np.mean(np.square(rs)))

def recalibrate(sig, loo):
    return sig*loo/np.sqrt(np.mean(sig**2))

# a single realization |err| = sigma*|z| caps the achievable correlation even for a perfect map:
def pearson_ceiling(s):
    return np.sqrt((2/np.pi)*s.var()/((s**2).mean() - (2/np.pi)*s.mean()**2))

def spearman_ceiling(s, ndraw=50):
    return np.mean([spearmanr(s.ravel()*np.abs(rng.standard_normal(s.size)), s.ravel())[0]
                    for _ in range(ndraw)])

OFF = ~mask_erg      # evaluate z-scores at unsampled pixels only (err and sigma are both 0 at samples)

def reliability(s, a, nb=10, reduce=np.mean):
    q = np.quantile(s, np.linspace(0, 1, nb+1)); q[-1] += 1e-12
    m, r = [], []
    for lo, hi in zip(q[:-1], q[1:]):
        sel = (s >= lo) & (s < hi)
        if sel.any():
            m.append(s[sel].mean()); r.append(reduce(a[sel]))
    return np.array(m), np.array(r)

def rms_reduce(a): return np.sqrt(np.mean(a**2))

# ---------------- plots (no abbreviations in any figure text) ----------------
C_ENS, C_ONE = '#1f77b4', '#ff7f0e'   # blue = ensemble comparison, orange = single realization
OBJ_REG = objective(mask_reg)

# --- figure 1: the three sampling patterns (paper Figures 3/14 analog) ---
fig, axes = plt.subplots(1, 3, figsize=(13.5, 4.8), constrained_layout=True)
for a, m, t in zip(axes,
                   [np.ones_like(mask_reg), mask_reg, mask_erg],
                   [f'Dense reference grid\n{N} samples, objective = 0 by definition',
                    f'Regular sparse subset\n{NS} samples, objective = {OBJ_REG:.1f}',
                    f'Ergodic subset, optimized by equation 4\n{NS} samples, objective = {f_erg:.1f}']):
    a.scatter(*np.argwhere(m).T, s=3, c='k')
    a.set_aspect(1); a.set_xlim(-1, NX); a.set_ylim(-1, NY)
    a.set_title(t); a.set_xlabel('grid x (cells)')
axes[0].set_ylabel('grid y (cells)')
fig.suptitle(f'Sampling patterns on the {NX}×{NY} grid: each sparse pattern keeps {DELTA:.1%} of the '
             f'samples (mean objective of 20 purely random patterns: {f_rand:.1f}; lower is better)')
plt.savefig('./plt1.png', dpi=120)

# --- figures 2a/2b: one figure per ground truth: reconstructions (top) and errors (bottom) ---
COLS = ['Ground truth',
        'Regular + spline interpolation\n(naive control)',
        'Regular + compressive sensing\n(Fourier)',
        'Regular + kriging\n(Gaussian process)',
        'Ergodic + spline interpolation',
        'Ergodic + compressive sensing\n(Fourier)',
        'Ergodic + compressive sensing\n(symlet wavelet, as in the paper)',
        'Ergodic + kriging\n(Gaussian process)']
COL_MASKS = [np.ones_like(mask_reg), mask_reg, mask_reg, mask_reg,
             mask_erg, mask_erg, mask_erg, mask_erg]
for fname, gt, imgs, note in [
        ('plt2a', truth_sp,
         [truth_sp, x_sp_spl, x_sp_reg, x_sp_reg_gp, x_sp_erg_spl, x_sp_erg, x_sp_sym, x_sp_gp],
         f'Fourier-sparse truth ({N_MODES} modes); compressive sensing with the ergodic pattern '
         'recovers the signal exactly'),
        ('plt2b', truth_pl,
         [truth_pl, x_pl_spl, x_pl_reg, x_pl_reg_gp, x_pl_erg_spl, x_erg, x_erg_sym, x_gp],
         f'potential-field truth (spectral exponent {BETA}, source depth {Z_SRC} cells); '
         'kriging is the best reconstruction for this field class')]:
    fig, axes = plt.subplots(3, 8, figsize=(32, 12.4), constrained_layout=True)
    v = np.abs(gt).max()
    for c, (a, x, t) in enumerate(zip(axes[0], imgs, COLS)):
        im = a.imshow(x.T, origin='lower', vmin=-v, vmax=v, cmap='RdBu_r')
        sub = '' if c == 0 else (f'\nleakage {leakage(x, gt):+.2f} · '
                                 f'root-mean-square error {rms(x, gt):.3f}')
        a.set_title(t + sub, fontsize=9)
    axes[0, 0].set_ylabel('Reconstruction\n\ngrid y (cells)')
    fig.colorbar(im, ax=axes[0], shrink=0.9,
                 label='field value\n(units of the truth standard deviation)')
    errors = [gt - x for x in imgs[1:]]
    ve = max(np.abs(e).max() for e in errors)
    axes[1, 0].axis('off')
    for a, e in zip(axes[1, 1:], errors):
        im2 = a.imshow(e.T, origin='lower', vmin=-ve, vmax=ve, cmap='RdBu_r')
    axes[1, 1].set_ylabel('Error: ground truth − reconstruction\n\ngrid y (cells)')
    fig.colorbar(im2, ax=axes[1, 1:], shrink=0.9,
                 label='error\n(units of the truth standard deviation)')
    for a, mm in zip(axes[2], COL_MASKS):
        a.scatter(*np.argwhere(mm).T, s=2, c='k')
        a.set_aspect(1); a.set_xlim(-1, NX); a.set_ylim(-1, NY)
        a.set_title(f'{int(mm.sum())} of {N} samples', fontsize=10)
        a.set_xlabel('grid x (cells)')
    axes[2, 0].set_ylabel('Sampling locations\n\ngrid y (cells)')
    fig.suptitle(f'Reconstruction from {DELTA:.1%} of the samples: {note}')
    plt.savefig(f'./{fname}.png', dpi=120)

# --- figures 3a/3b: predicted standard deviation versus actual error, one per ground truth ---
def sigma_study(fname, desc, make_truth, truth, x_rec, sig_raw):
    y_img = truth*mask_erg
    loo = loo_rms(y_img, mask_erg, cs_reconstruct)
    sig = recalibrate(sig_raw, loo)    # GP shape transfers to the l1 solver; LOO sets the level
    err = x_rec - truth; aerr = np.abs(err)
    print(f"\n--- sigma study ({fname}): {desc} ---")
    print(f"LOO error rms (truth-free calibration level): {loo:.3f}   "
          f"actual rms: {np.sqrt((err**2).mean()):.3f}")
    pr = pearsonr(aerr.ravel(), sig.ravel())[0]; sr = spearmanr(aerr.ravel(), sig.ravel())[0]
    pc, sc = pearson_ceiling(sig), spearman_ceiling(sig)
    print(f"|err| vs sigma_GP: pearson={pr:+.3f} (ceiling {pc:.3f}, ratio {pr/pc:+.2f})  "
          f"spearman={sr:+.3f} (ceiling {sc:.3f}, ratio {sr/sc:+.2f})  "
          f"z-std={np.std(err[OFF]/sig[OFF]):.2f}")
    ens_errs = np.array([cs_reconstruct(t*mask_erg, mask_erg) - t
                         for t in (make_truth() for _ in range(ENS_K))])
    sig_emp = np.sqrt((ens_errs**2).mean(0))
    pr2 = pearsonr(sig_emp.ravel(), sig.ravel())[0]; sr2 = spearmanr(sig_emp.ravel(), sig.ravel())[0]
    print(f"ensemble (K={ENS_K} truths, same mask; sigma_emp has ~{1/np.sqrt(2*ENS_K):.0%} noise): "
          f"pearson={pr2:+.3f} spearman={sr2:+.3f} "
          f"pooled z-std={np.std(ens_errs[:, OFF]/sig[OFF]):.2f}")

    fig, axes = plt.subplots(3, 3, figsize=(14.5, 13.8), constrained_layout=True)
    vs = max(sig.max(), sig_emp.max(), aerr.max())
    v = np.abs(truth).max()
    im0 = axes[0, 0].imshow(truth.T, origin='lower', vmin=-v, vmax=v, cmap='RdBu_r')
    axes[0, 0].set_title('Ground truth', fontsize=10)
    axes[0, 1].imshow(x_rec.T, origin='lower', vmin=-v, vmax=v, cmap='RdBu_r')
    axes[0, 1].set_title('Reconstruction: ergodic + compressive sensing (Fourier)\n'
                         f'leakage {leakage(x_rec, truth):+.2f} · root-mean-square error '
                         f'{rms(x_rec, truth):.3f}', fontsize=10)
    fig.colorbar(im0, ax=axes[0, :2], shrink=0.9,
                 label='field value\n(units of the truth standard deviation)')
    imE = axes[0, 2].imshow(err.T, origin='lower', vmin=-vs, vmax=vs, cmap='RdBu_r')
    axes[0, 2].set_title('Signed error:\nreconstruction − truth', fontsize=10)
    fig.colorbar(imE, ax=axes[0, 2], shrink=0.9,
                 label='error\n(units of the truth standard deviation)')
    for a, x, t in zip(axes[1],
                       [sig, sig_emp, aerr],
                       ['Predicted standard deviation:\nGaussian process posterior (truth-free)',
                        f'Empirical standard deviation\n({ENS_K} independent truths, '
                        'same sampling pattern)',
                        'Actual absolute error, this realization:\n|reconstruction − truth|']):
        im1 = a.imshow(x.T, origin='lower', vmin=0, vmax=vs, cmap='magma')
        a.set_title(t, fontsize=10)
        a.set_xlabel('grid x (cells)')
    fig.colorbar(im1, ax=axes[1], shrink=0.9,
                 label='standard deviation or absolute error\n(units of the truth standard deviation)')
    for r in (0, 1):
        axes[r, 0].set_ylabel('grid y (cells)')

    a = axes[2, 0]   # the sampling pattern that produced everything above
    a.scatter(*np.argwhere(mask_erg).T, s=2, c='k')
    a.set_aspect(1); a.set_xlim(-1, NX); a.set_ylim(-1, NY)
    a.set_title(f'Ergodic sampling locations\n({NS} of {N} samples)', fontsize=10)
    a.set_xlabel('grid x (cells)'); a.set_ylabel('grid y (cells)')

    a = axes[2, 1]   # scatter with decile-binned means: does the prediction rank the error?
    a.scatter(sig[OFF], aerr[OFF], s=2, alpha=.2, color='#9ecae1')
    bm, bv = reliability(sig[OFF], aerr[OFF])
    a.plot(bm, bv, 'o-', color='k', lw=1.6)
    xs = np.linspace(0, sig[OFF].max(), 50)
    a.plot(xs, np.sqrt(2/np.pi)*xs, '--', color='k', lw=1.2)
    a.set_xlabel('predicted standard deviation'); a.set_ylabel('absolute error, this realization')
    a.ticklabel_format(style='sci', scilimits=(-2, 3))
    a.set_title('Prediction versus error at unsampled cells\n(single-realization Pearson '
                f'correlation is capped\nat about {pc:.2f} even for a perfect prediction)',
                fontsize=10)
    a.legend(handles=[plt.Line2D([], [], marker='o', ls='none', ms=4, color='#9ecae1',
                                 label='unsampled cells'),
                      plt.Line2D([], [], marker='o', ls='-', ms=5, color='k',
                                 label='mean absolute error\nwithin prediction decile'),
                      plt.Line2D([], [], ls='--', color='k',
                                 label='expected mean absolute error\nfor a perfect prediction')],
             loc='upper left', fontsize=8)

    a = axes[2, 2]   # merged calibration panel: against one realization and against the ensemble
    for target, lab, cc in [(aerr, 'against this single realization', C_ONE),
                            (sig_emp, f'against the {ENS_K}-truth ensemble', C_ENS)]:
        mm, rr = reliability(sig, target, reduce=rms_reduce)
        a.plot(mm, rr, 'o-', color=cc, label=lab)
    cal_lim = max(a.get_xlim()[1], a.get_ylim()[1])
    a.plot([0, cal_lim], [0, cal_lim], 'k--', lw=.8, label='perfect calibration')
    a.set_xlim(0, cal_lim); a.set_ylim(0, cal_lim); a.set_aspect(1)
    a.set_xlabel('mean predicted standard deviation within decile')
    a.set_ylabel('root-mean-square error or\nempirical standard deviation within decile')
    a.ticklabel_format(style='sci', scilimits=(-2, 3))
    a.set_title('Calibration of the Gaussian process prediction\n(cells binned into deciles of '
                'the prediction)', fontsize=10)
    a.legend(loc='upper left', fontsize=8)
    fig.suptitle('Truth-free predicted standard deviation versus actual reconstruction error\n'
                 f'({desc}; ergodic sampling; compressive-sensing reconstruction)\n'
                 'Uniform ergodic coverage intentionally flattens the prediction; its narrow '
                 'range caps single-realization correlation;\nthe ensemble comparison is the '
                 'decisive test.', fontsize=11)
    plt.savefig(f'./{fname}.png', dpi=120)

sigma_study('plt3a', f'Fourier-sparse truth ({N_MODES} modes)',
            lambda: fourier_sparse_field(N_MODES), truth_sp, x_sp_erg, sig_sp_raw)
sigma_study('plt3b', f'potential-field truth (spectral exponent {BETA}, source depth {Z_SRC} cells)',
            lambda: powerlaw_field(BETA), truth_pl, x_erg, sig_gp_raw)

# ---------------- Monte-Carlo RMS table over truth classes ----------------
N_TRIALS = 25

def white_noise_field():
    # incompressible control: no method can beat predicting zero at unsampled cells,
    # whose rms is sqrt(1 - DELTA) ~ 0.94
    f = rng.standard_normal((NX, NY))
    return (f - f.mean())/f.std()

def trial_methods(truth, gp_params):
    return {
        'regular+spline': spline_interp(truth, mask_reg),
        'regular+CS-Fourier': cs_reconstruct(truth*mask_reg, mask_reg),
        'regular+kriging': gp_reconstruct(truth*mask_reg, mask_reg, *gp_params)[0],
        'ergodic+spline': spline_interp(truth, mask_erg),
        'ergodic+CS-Fourier': cs_reconstruct(truth*mask_erg, mask_erg),
        'ergodic+CS-symlet': cs_reconstruct_sym(truth*mask_erg, mask_erg),
        'ergodic+kriging': gp_reconstruct(truth*mask_erg, mask_erg, *gp_params)[0],
    }

CLASSES = [('white noise (control)', white_noise_field),
           ('Fourier-sparse', lambda: fourier_sparse_field(N_MODES)),
           ('potential-field', lambda: powerlaw_field(BETA))]
print(f"\n--- Monte-Carlo rms error, mean +/- std over {N_TRIALS} trials per truth class ---")
results = {}
for cname, gen in CLASSES:
    t0 = gen()
    gp_params = gp_fit_params(t0*mask_erg, mask_erg)   # truth-free model selection, once per class
    rr = {}
    for k in range(N_TRIALS):
        truth = t0 if k == 0 else gen()
        for m, x in trial_methods(truth, gp_params).items():
            rr.setdefault(m, []).append(rms(x, truth))
    results[cname] = {m: (np.mean(v), np.std(v)) for m, v in rr.items()}
print('method'.ljust(20) + ''.join(c.rjust(24) for c, _ in CLASSES))
for m in results[CLASSES[0][0]]:
    print(m.ljust(20) + ''.join(f"{results[c][m][0]:.3f} +/- {results[c][m][1]:.3f}".rjust(24)
                                for c, _ in CLASSES))

# --- figure 4: the Monte-Carlo table as a chart with error bars ---
DISPLAY = {'regular+spline': 'Regular + spline interpolation',
           'regular+CS-Fourier': 'Regular + compressive sensing (Fourier)',
           'regular+kriging': 'Regular + kriging (Gaussian process)',
           'ergodic+spline': 'Ergodic + spline interpolation',
           'ergodic+CS-Fourier': 'Ergodic + compressive sensing (Fourier)',
           'ergodic+CS-symlet': 'Ergodic + compressive sensing (symlet wavelet)',
           'ergodic+kriging': 'Ergodic + kriging (Gaussian process)'}
methods = list(DISPLAY)
ypos = np.arange(len(methods))[::-1]
xmax = max(results[c][m][0] + results[c][m][1] for c, _ in CLASSES for m in methods)
fig, axes = plt.subplots(1, 3, figsize=(15.5, 5.8), sharey=True, constrained_layout=True)
for a, (cname, _) in zip(axes, CLASSES):
    mus = np.array([results[cname][m][0] for m in methods])
    sds = np.array([results[cname][m][1] for m in methods])
    # bar color encodes the sampling pattern (blue = regular, orange = ergodic); the method
    # labels already carry this, so no legend is needed
    a.barh(ypos, mus, xerr=sds, height=0.6,
           color=['#9ecae1' if m.startswith('regular') else '#fdbe85' for m in methods],
           error_kw=dict(ecolor='k', lw=1, capsize=3))
    for yp, mu, sd in zip(ypos, mus, sds):
        a.text(mu + sd + 0.02*xmax, yp, f'{mu:.3f} ± {sd:.3f}', va='center', fontsize=8)
    t = cname[0].upper() + cname[1:] + ' truth'*(not cname.endswith(')'))
    if cname == CLASSES[0][0]:
        t += f'\n(dashed line: best possible = predict zero everywhere, {np.sqrt(1 - DELTA):.2f})'
    a.set_title(t, fontsize=10)
    a.set_xlabel('root-mean-square error\n(units of the truth standard deviation)')
    a.set_xlim(0, xmax*1.32)
    a.grid(axis='x', lw=0.4, alpha=0.4)
    a.set_axisbelow(True)
axes[0].set_yticks(ypos, [DISPLAY[m] for m in methods])
axes[0].axvline(np.sqrt(1 - DELTA), ls='--', lw=1, color='k')
fig.suptitle('Monte-Carlo root-mean-square reconstruction error: '
             f'mean ± standard deviation over {N_TRIALS} trials per truth class '
             f'({DELTA:.1%} of samples, error bars = ± one standard deviation)')
plt.savefig('./plt4.png', dpi=120)