elastic_kernels.py

"""
================================================================================
  ELASTIC SENSITIVITY KERNELS FOR Vp AND Vs
  Firedrake-Adjoint Framework (Pure-Python FEM/FD Implementation)
================================================================================

  This code implements the Firedrake-adjoint workflow for computing
  elastic sensitivity (Fréchet) kernels.  Every class and function is
  named and documented to match the real Firedrake/Spyro API so the code
  can be ported to a Firedrake environment with minimal changes.

  PHYSICS (Firedrake UFL notation)
  ─────────────────────────────────
  Forward elastic wave equation (weak form):

      ∫ ρ ü·v dx  +  ∫ σ(u):ε(v) dx  =  ∫ f·v dx
      σ = λ(∇·u)I + 2με(u)
      ε(u) = ½(∇u + ∇uᵀ)

  Lamé parameters from velocities:
      λ = ρ(Vp² − 2Vs²)
      μ = ρ Vs²

  ADJOINT STATE METHOD
  ─────────────────────
  FWI objective:
      J(m) = ½ Σᵣ ∫₀ᵀ |u_syn(xᵣ,t) − d_obs(xᵣ,t)|² dt

  Adjoint equation (time-reversed):
      ρ λ̈ = ∇·σ(λ) + f_adj,    f_adj = Σᵣ [u−d_obs] δ(x−xᵣ)

  SENSITIVITY KERNELS (Firedrake-adjoint output):
  ─────────────────────────────────────────────────
      K_λ(x) =  ∫₀ᵀ [∇·u_fwd(x,t)][∇·λ_adj(x,t)] dt   (λ-kernel)
      K_μ(x) =  ∫₀ᵀ 2 εᵢⱼ(u_fwd)·εᵢⱼ(λ_adj) dt          (μ-kernel)

  Chain rule to Vp and Vs:
      K_Vp(x) = ∂J/∂Vp = 2ρVp · K_λ            (P-wave kernel)
      K_Vs(x) = ∂J/∂Vs = −4ρVs · K_λ + 4ρVs · K_μ   (S-wave kernel)
      K_ρ(x)  = ∂J/∂ρ  = (Vp²−2Vs²) K_λ + Vs² K_μ    (density kernel)

  KERNEL CONDITIONING
  ────────────────────
  1. Source-point muting:  taper = 1 − exp(−|x−xₛ|²/R²)
     Zeros the kernel within radius R of the source to remove singularity.
  2. Receiver-point muting: same Gaussian taper at each receiver.
  3. Illumination preconditioning:  K_precond = K / (P_illum + ε)
     P_illum = ∫₀ᵀ |u_fwd|² dt   (source-energy proxy)
  4. Depth preconditioning:  K_depth = K · (1 + z/z₀)²
     Compensates geometrical amplitude decay.
  5. Gaussian smoothing (spatial regularisation).

  FRESNEL ZONE ANALYSIS
  ──────────────────────
  Theoretical Fresnel zone half-width at depth z:
      W_F(z) = √(λ·z/2) = √(Vp/(2·f₀) · z)
  Measured from kernel: half-width of the central positive lobe at each depth.
  Saved as CSV for quantitative comparison.

  OUTPUT FILES
  ─────────────
  Matplotlib:
      01_model_geometry.png        — domain, source, receivers, anomaly
      02_ricker_wavelet.png        — source signature + spectrum
      03_forward_snapshots.png     — 9-panel vx wavefield
      04_forward_wavefield.gif     — animated propagation
      05_seismograms.png           — recorded vx, vz
      06_adjoint_snapshots.png     — 9-panel adjoint wavefield
      07_raw_kernels.png           — K_Vp, K_Vs, K_λ, K_μ (raw)
      08_conditioned_kernels.png   — after muting + preconditioning
      09_kernel_comparison.png     — raw vs conditioned side-by-side
      10_fresnel_analysis.png      — measured vs theoretical Fresnel widths
      11_depth_profiles.png        — 1D vertical kernel profiles
      12_horizontal_profiles.png   — 1D horizontal slices at target depth
      13_multi_frequency.png       — kernels at 3 Hz, 5 Hz, 8 Hz
      14_source_receiver_muting.png— muting function visualisation
      15_illumination.png          — source illumination map
      16_kernel_summary.png        — publication-quality 6-panel summary

  CSV:
      fresnel_zone_table.csv       — depth vs measured/theoretical width
      kernel_statistics.csv        — per-kernel amplitude statistics
================================================================================
"""

import numpy as np
import matplotlib; matplotlib.use("Agg")
import matplotlib.pyplot as plt
import matplotlib.gridspec as gridspec
import matplotlib.colors as mcolors
from matplotlib.colors import TwoSlopeNorm
import matplotlib.patches as mpatches
from scipy.ndimage import gaussian_filter
from scipy.signal import butter, filtfilt
from PIL import Image
import io, os, gc, time, csv, warnings
warnings.filterwarnings("ignore")

os.makedirs("kernel_results", exist_ok=True)
od = "kernel_results"

# ══════════════════════════════════════════════════════════════════════════════
# FIREDRAKE-ADJOINT DICTIONARY  (mirrors spyro.tools.Dictionary)
# ══════════════════════════════════════════════════════════════════════════════
class Dictionary:
    """
    Central configuration — mirrors Firedrake/Spyro Dictionary.

    In real Firedrake/Spyro:
        dictionary = {
            "options":   {"method": "mass_lumped_triangle", "degree": 1},
            "mesh":      {"Lx": 3000.0, "Lz": 3000.0, "nx": 120, "nz": 120},
            "acquisition": {"source_type": "Ricker", "frequency": 5.0,
                            "delay": 0.3, "amplitude": 1e10,
                            "receiver_locations": [...]},
            "time_axis": {"final_time": 1.5, "dt": 0.001},
        }
    """
    # ── Domain ────────────────────────────────────────────────────────────
    Lx, Lz   = 3000., 3000.   # domain [m]
    Nx, Nz   = 100,   100     # grid resolution

    # ── Elastic background ────────────────────────────────────────────────
    Vp_bg    = 3000.;  Vs_bg  = 1732.;  rho_bg  = 2500.

    # ── Circular anomaly (Camembert coin) ─────────────────────────────────
    cx,  cz  = 1500.,  1500.  # anomaly centre [m]
    radius   = 300.           # anomaly radius [m]
    Vp_anom  = 4000.          # anomaly Vp [m/s]  (+33%)
    Vs_anom  = 2309.          # anomaly Vs [m/s]  (+33%)
    rho_anom = 2800.          # anomaly density

    # ── Source ────────────────────────────────────────────────────────────
    f0, t0   = 5.0,   0.3     # Ricker Hz, delay [s]
    src_amp  = 1e10
    T        = 1.5            # total simulation time [s]
    src_x    = 1500.; src_z = 50.   # single surface source

    # ── Receivers (surface line array) ───────────────────────────────────
    n_recv   = 40;   recv_z = 50.    # 40 receivers near surface

    # ── ABC sponge ────────────────────────────────────────────────────────
    abc_w    = 20             # sponge width [grid points]

    # ── Kernel conditioning ───────────────────────────────────────────────
    mute_R   = 200.           # source/receiver mute radius [m]
    depth_z0 = 500.           # depth preconditioning length scale [m]
    illum_eps = 1e-6          # illumination normalisation floor

    # ── Multi-frequency analysis ──────────────────────────────────────────
    freqs    = [3.0, 5.0, 8.0]  # Hz for Fresnel analysis


cfg = Dictionary()

# Derived grid
dx  = cfg.Lx / (cfg.Nx - 1)
dz  = cfg.Lz / (cfg.Nz - 1)
dt  = 0.9 * min(dx, dz) / (2**0.5 * max(cfg.Vp_bg, cfg.Vp_anom))
nt  = int(cfg.T / dt) + 1
t_  = np.linspace(0, cfg.T, nt)
x_  = np.linspace(0, cfg.Lx, cfg.Nx)
z_  = np.linspace(0, cfg.Lz, cfg.Nz)
X, Z = np.meshgrid(x_, z_, indexing="ij")

print(f"  Dictionary: Nx={cfg.Nx} Nz={cfg.Nz} dx={dx:.1f}m "
      f"dt={dt*1e4:.2f}×10⁻⁴s nt={nt}")

# ══════════════════════════════════════════════════════════════════════════════
# MESH + FUNCTION SPACE  (mirrors Firedrake RectangleMesh + FunctionSpace)
# ══════════════════════════════════════════════════════════════════════════════
# In Firedrake:
#   mesh = RectangleMesh(Nx, Nz, Lx, Lz)
#   Q = FunctionSpace(mesh, "CG", 1)      # scalar P1
#   V = VectorFunctionSpace(mesh, "CG", 1) # vector P1

# Here we work directly with [Nx, Nz] numpy arrays (equivalent on structured grid)


# ══════════════════════════════════════════════════════════════════════════════
# ELASTIC MODEL  (mirrors Firedrake Function.interpolate + conditional)
# ══════════════════════════════════════════════════════════════════════════════
def build_model(anomaly=True, Vp_arr=None):
    """
    Build elastic parameter arrays.

    Firedrake equivalent:
        x, z = SpatialCoordinate(mesh)
        r = sqrt((x-cx)**2 + (z-cz)**2)
        Vp = Function(Q).interpolate(
            conditional(r < radius, Constant(Vp_anom), Constant(Vp_bg)))
        lam = Function(Q).interpolate(rho*(Vp**2 - 2*Vs**2))
        mu  = Function(Q).interpolate(rho*Vs**2)
    """
    Vp  = np.full((cfg.Nx, cfg.Nz), cfg.Vp_bg,  dtype=np.float64)
    Vs  = np.full((cfg.Nx, cfg.Nz), cfg.Vs_bg,  dtype=np.float64)
    rho = np.full((cfg.Nx, cfg.Nz), cfg.rho_bg, dtype=np.float64)
    if anomaly:
        R2 = (X - cfg.cx)**2 + (Z - cfg.cz)**2
        mask = R2 < cfg.radius**2
        Vp[mask]  = cfg.Vp_anom
        Vs[mask]  = cfg.Vs_anom
        rho[mask] = cfg.rho_anom
    if Vp_arr is not None:
        Vp = np.clip(Vp_arr, cfg.Vp_bg*0.5, cfg.Vp_anom*1.5)
        Vs = np.minimum(Vs, Vp/np.sqrt(3.)*0.98)
    lam = rho*(Vp**2 - 2.*Vs**2)
    mu  = rho*Vs**2
    return {"Vp":Vp,"Vs":Vs,"rho":rho,"lam":lam,"mu":mu}

# ══════════════════════════════════════════════════════════════════════════════
# SOURCE + RECEIVERS  (mirrors spyro.Sources + spyro.Receivers)
# ══════════════════════════════════════════════════════════════════════════════
class Acquisition:
    """
    Single source + surface receiver array.
    Mirrors spyro.Sources / spyro.Receivers geometry objects.
    """
    def __init__(self, f0=None):
        f0 = f0 or cfg.f0
        self.f0  = f0
        self.SX  = int(round(cfg.src_x / dx))
        self.SZ  = int(round(cfg.src_z / dz))
        self.RX  = np.linspace(int(cfg.Nx*0.04), int(cfg.Nx*0.96),
                               cfg.n_recv, dtype=int)
        self.RZ  = np.full(cfg.n_recv, int(round(cfg.recv_z/dz)), dtype=int)
        tau       = np.pi * f0 * (t_ - cfg.t0)
        self.wav  = cfg.src_amp * (1. - 2*tau**2) * np.exp(-tau**2)

    def __repr__(self):
        return (f"Acquisition(src=({self.SX*dx:.0f},{self.SZ*dz:.0f})m, "
                f"n_recv={cfg.n_recv}, f0={self.f0}Hz)")


# ══════════════════════════════════════════════════════════════════════════════
# SPONGE ABC  (mirrors Firedrake boundary integral term)
# ══════════════════════════════════════════════════════════════════════════════
def make_sponge(Vp):
    W = cfg.abc_w; d = np.ones((cfg.Nx, cfg.Nz))
    v = float(np.mean(Vp[:W, :])); a = -np.log(1e-6)/(W*dx)
    for i in range(W):
        f = ((W-i)/W)**2; q = np.exp(-a*f*v*dt/dx)
        d[i,:]*=q; d[cfg.Nx-1-i,:]*=q
    for j in range(W):
        f = ((W-j)/W)**2; q = np.exp(-a*f*v*dt/dz)
        d[:,cfg.Nz-1-j]*=q
    return d

# ══════════════════════════════════════════════════════════════════════════════
# FORWARD SOLVER  (mirrors spyro.solvers.forward)
# ══════════════════════════════════════════════════════════════════════════════
def forward(model, acq, store=False, snap_n=16, tag="fwd"):
    """
    Staggered-grid elastic forward solver (Virieux 1986).

    Firedrake/Spyro equivalent:
        u_syn, seismograms = spyro.solvers.forward(
            model_dict, mesh, comm, c, excitations, receivers)

    Physics:
        σxx += Δt[(λ+2μ)∂vx/∂x + λ∂vz/∂z]
        σzz += Δt[λ∂vx/∂x + (λ+2μ)∂vz/∂z]
        σxz += Δt·μ(∂vx/∂z + ∂vz/∂x)
        vx  += (Δt/ρ)(∂σxx/∂x + ∂σxz/∂z)
        vz  += (Δt/ρ)(∂σxz/∂x + ∂σzz/∂z)

    Free surface at z=0: σzz=0, σxz=0

    Parameters
    ----------
    model   : dict from build_model()
    acq     : Acquisition object
    store   : if True, return full wavefield [Nx,Nz,nt] (for adjoint)
    snap_n  : number of snapshots for visualisation
    tag     : progress label

    Returns
    -------
    svx, svz   : [n_recv, nt] seismograms
    snaps      : [(t, vx_2d), ...]
    VXA, VZA   : [Nx,Nz,nt] float32 | None
    illum      : [Nx,Nz] source illumination ∫|vx|² dt
    """
    lam,mu,rho,Vp = model["lam"],model["mu"],model["rho"],model["Vp"]
    sponge = make_sponge(Vp)

    vx=np.zeros((cfg.Nx,cfg.Nz)); vz=np.zeros_like(vx)
    sxx=np.zeros_like(vx); szz=np.zeros_like(vx); sxz=np.zeros_like(vx)

    l2m=lam+2.*mu; dtr=dt/rho; dtl=dt*lam; dtm=dt*mu; dtl2=dt*l2m
    svx=np.zeros((cfg.n_recv,nt)); svz=np.zeros_like(svx)
    illum=np.zeros((cfg.Nx,cfg.Nz))
    si=max(1,nt//snap_n); snaps=[]

    if store:
        VXA=np.zeros((cfg.Nx,cfg.Nz,nt),dtype=np.float32)
        VZA=np.zeros_like(VXA)
    else:
        VXA=VZA=None

    print(f"  [{tag}] {nt} steps ...",end="",flush=True)
    for it in range(nt):
        # Stress update
        dvx_dx=(vx[1:,:-1]-vx[:-1,:-1])/dx; dvz_dz=(vz[:-1,1:]-vz[:-1,:-1])/dz
        sxx[:-1,:-1]+=dtl2[:-1,:-1]*dvx_dx+dtl[:-1,:-1]*dvz_dz
        szz[:-1,:-1]+=dtl[:-1,:-1]*dvx_dx+dtl2[:-1,:-1]*dvz_dz
        dvx_dz=(vx[:-1,1:]-vx[:-1,:-1])/dz; dvz_dx=(vz[1:,:-1]-vz[:-1,:-1])/dx
        sxz[:-1,:-1]+=dtm[:-1,:-1]*(dvx_dz+dvz_dx)
        # Velocity update
        d1=(sxx[1:,1:]-sxx[:-1,1:])/dx+(sxz[1:,1:]-sxz[1:,:-1])/dz
        d2=(sxz[1:,1:]-sxz[:-1,1:])/dx+(szz[1:,1:]-szz[1:,:-1])/dz
        vx[1:,1:]+=dtr[1:,1:]*d1; vz[1:,1:]+=dtr[1:,1:]*d2
        # Source injection (Ricker)
        vx[acq.SX,acq.SZ]+=dtr[acq.SX,acq.SZ]*acq.wav[it]
        # Free surface
        szz[:,0]=0.; sxz[:,0]=0.
        # Sponge ABC
        vx*=sponge; vz*=sponge; sxx*=sponge; szz*=sponge; sxz*=sponge
        # Record seismograms
        for ir in range(cfg.n_recv):
            svx[ir,it]=vx[int(acq.RX[ir]),int(acq.RZ[ir])]
            svz[ir,it]=vz[int(acq.RX[ir]),int(acq.RZ[ir])]
        # Store wavefield
        if store:
            VXA[:,:,it]=vx.astype(np.float32)
            VZA[:,:,it]=vz.astype(np.float32)
        # Illumination: ∫ |vx|² dt
        illum += vx**2 * dt
        if it%si==0: snaps.append((t_[it],vx.copy()))

    print(f"  max|vx|={np.max(np.abs(vx)):.2e}")
    return svx,svz,snaps,VXA,VZA,illum


# ══════════════════════════════════════════════════════════════════════════════
# ADJOINT SOLVER + KERNEL COMPUTATION  (mirrors spyro.solvers.gradient)
# ══════════════════════════════════════════════════════════════════════════════
def adjoint_kernels(model, acq, res_vx, res_vz, VXA, VZA, snap_n=16, tag="adj"):
    """
    Adjoint-state solver for elastic sensitivity kernels.

    Firedrake-adjoint equivalent:
        J = assemble(misfit_functional)
        dJdm = compute_gradient(J, Control(m))

    The adjoint equation (time-reversed, identical operator):
        ρ λ̈ = ∇·σ(λ) + f_adj(t)
        f_adj(xᵣ,t) = [u_syn(xᵣ,T-t) - d_obs(xᵣ,T-t)]

    KERNEL FORMULAS (cross-correlation imaging condition):
    ────────────────────────────────────────────────────────
    λ-kernel  (sensitivity to λ = ρ(Vp²-2Vs²)):
        K_λ(x) = ∫₀ᵀ [∇·u_fwd(t)] · [∇·λ_adj(t)] dt

    μ-kernel  (sensitivity to μ = ρVs²):
        K_μ(x) = ∫₀ᵀ 2 [ε_ij(u_fwd)] · [ε_ij(λ_adj)] dt
               = ∫₀ᵀ 2[(∂ux/∂x)(∂λx/∂x) + (∂uz/∂z)(∂λz/∂z)
                        + ½(∂ux/∂z+∂uz/∂x)(∂λx/∂z+∂λz/∂x)] dt

    Chain rule to physical parameters:
        K_Vp = ∂J/∂Vp = 2ρVp · K_λ
        K_Vs = ∂J/∂Vs = -4ρVs·K_λ  +  4ρVs·K_μ
        K_ρ  = ∂J/∂ρ  = (Vp²-2Vs²)·K_λ  +  Vs²·K_μ
    """
    lam,mu,rho = model["lam"],model["mu"],model["rho"]
    Vp,Vs      = model["Vp"],model["Vs"]
    sponge     = make_sponge(Vp)

    vxA=np.zeros((cfg.Nx,cfg.Nz)); vzA=np.zeros_like(vxA)
    sxA=np.zeros_like(vxA); szA=np.zeros_like(vxA); sxzA=np.zeros_like(vxA)
    l2m=lam+2.*mu; dtr=dt/rho; dtl=dt*lam; dtm=dt*mu; dtl2=dt*l2m

    # Gradient accumulators
    K_lam = np.zeros((cfg.Nx,cfg.Nz))
    K_mu  = np.zeros((cfg.Nx,cfg.Nz))

    si=max(1,nt//snap_n); snaps=[]

    print(f"  [{tag}] backward {nt} steps ...",end="",flush=True)
    for rev in range(nt):
        it = nt-1-rev   # reversed time index

        # Adjoint stress update
        dvx_dx=(vxA[1:,:-1]-vxA[:-1,:-1])/dx; dvz_dz=(vzA[:-1,1:]-vzA[:-1,:-1])/dz
        sxA[:-1,:-1]+=dtl2[:-1,:-1]*dvx_dx+dtl[:-1,:-1]*dvz_dz
        szA[:-1,:-1]+=dtl[:-1,:-1]*dvx_dx+dtl2[:-1,:-1]*dvz_dz
        dvx_dz=(vxA[:-1,1:]-vxA[:-1,:-1])/dz; dvz_dx=(vzA[1:,:-1]-vzA[:-1,:-1])/dx
        sxzA[:-1,:-1]+=dtm[:-1,:-1]*(dvx_dz+dvz_dx)
        # Adjoint velocity update
        d1=(sxA[1:,1:]-sxA[:-1,1:])/dx+(sxzA[1:,1:]-sxzA[1:,:-1])/dz
        d2=(sxzA[1:,1:]-sxzA[:-1,1:])/dx+(szA[1:,1:]-szA[1:,:-1])/dz
        vxA[1:,1:]+=dtr[1:,1:]*d1; vzA[1:,1:]+=dtr[1:,1:]*d2
        # Adjoint source: time-reversed residuals at receivers
        for ir in range(cfg.n_recv):
            ix2,iz2 = int(acq.RX[ir]),int(acq.RZ[ir])
            vxA[ix2,iz2]+=dtr[ix2,iz2]*res_vx[ir,it]
            vzA[ix2,iz2]+=dtr[ix2,iz2]*res_vz[ir,it]
        # Free surface
        szA[:,0]=0.; sxzA[:,0]=0.
        # Sponge
        vxA*=sponge; vzA*=sponge; sxA*=sponge; szA*=sponge; sxzA*=sponge

        # ── KERNEL CROSS-CORRELATION IMAGING CONDITION ──────────────────
        if VXA is not None:
            vxf = VXA[:,:,it].astype(np.float64)
            vzf = VZA[:,:,it].astype(np.float64)

            # Divergence of forward field: ∇·u_fwd
            div_fwd = np.zeros((cfg.Nx,cfg.Nz))
            div_fwd[1:,1:] = (vxf[1:,1:]-vxf[:-1,1:])/dx \
                            + (vzf[1:,1:]-vzf[1:,:-1])/dz

            # Divergence of adjoint field: ∇·λ_adj
            div_adj = np.zeros((cfg.Nx,cfg.Nz))
            div_adj[1:,1:] = (vxA[1:,1:]-vxA[:-1,1:])/dx \
                            + (vzA[1:,1:]-vzA[1:,:-1])/dz

            # λ-kernel: K_λ += (∇·u_fwd)(∇·λ_adj) Δt
            K_lam += div_fwd * div_adj * dt

            # Strain tensors for μ-kernel
            # ε_xx = ∂ux/∂x,  ε_zz = ∂uz/∂z,  ε_xz = ½(∂ux/∂z + ∂uz/∂x)
            eps_xx_f = np.zeros((cfg.Nx,cfg.Nz))
            eps_zz_f = np.zeros((cfg.Nx,cfg.Nz))
            eps_xz_f = np.zeros((cfg.Nx,cfg.Nz))
            eps_xx_f[1:,1:]  = (vxf[1:,1:]-vxf[:-1,1:])/dx
            eps_zz_f[1:,1:]  = (vzf[1:,1:]-vzf[1:,:-1])/dz
            eps_xz_f[:-1,:-1]= 0.5*((vxf[:-1,1:]-vxf[:-1,:-1])/dz
                                    +(vzf[1:,:-1]-vzf[:-1,:-1])/dx)

            eps_xx_a = np.zeros((cfg.Nx,cfg.Nz))
            eps_zz_a = np.zeros((cfg.Nx,cfg.Nz))
            eps_xz_a = np.zeros((cfg.Nx,cfg.Nz))
            eps_xx_a[1:,1:]  = (vxA[1:,1:]-vxA[:-1,1:])/dx
            eps_zz_a[1:,1:]  = (vzA[1:,1:]-vzA[1:,:-1])/dz
            eps_xz_a[:-1,:-1]= 0.5*((vxA[:-1,1:]-vxA[:-1,:-1])/dz
                                    +(vzA[1:,:-1]-vzA[:-1,:-1])/dx)

            # μ-kernel: K_μ += 2[ε_xx·εA_xx + ε_zz·εA_zz + 2·ε_xz·εA_xz] Δt
            K_mu += 2.*(eps_xx_f*eps_xx_a + eps_zz_f*eps_zz_a
                       + 2.*eps_xz_f*eps_xz_a) * dt

        if rev%si==0: snaps.append((t_[it], vxA.copy()))

    print(f"  max|vxA|={np.max(np.abs(vxA)):.2e}")

    # ── CHAIN RULE: λ,μ → Vp,Vs,ρ ─────────────────────────────────────
    # ∂J/∂Vp = ∂J/∂λ · ∂λ/∂Vp  where  ∂λ/∂Vp = 2ρVp
    # ∂J/∂Vs = ∂J/∂λ · ∂λ/∂Vs + ∂J/∂μ · ∂μ/∂Vs
    #        where ∂λ/∂Vs = -4ρVs,  ∂μ/∂Vs = 2ρVs
    # ∂J/∂ρ  = ∂J/∂λ · ∂λ/∂ρ + ∂J/∂μ · ∂μ/∂ρ
    #        where ∂λ/∂ρ = Vp²-2Vs², ∂μ/∂ρ = Vs²
    K_Vp = 2.*rho*Vp*K_lam
    K_Vs = (-4.*rho*Vs*K_lam) + (4.*rho*Vs*K_mu)
    K_rho= ((Vp**2 - 2.*Vs**2)*K_lam) + (Vs**2*K_mu)

    return K_Vp, K_Vs, K_rho, K_lam, K_mu, snaps


# ══════════════════════════════════════════════════════════════════════════════
# KERNEL CONDITIONING  (Muting + Preconditioning)
# ══════════════════════════════════════════════════════════════════════════════
class KernelConditioner:
    """
    Apply muting and preconditioning to raw sensitivity kernels.

    Three conditioning steps mirror standard Firedrake/Spyro FWI practice:

    1. SOURCE-POINT MUTING
       taper_src = 1 - exp(-|x-xs|²/R²)
       Removes the singular source-point artefact (the 'rabbit ear' at xs).
       Physical justification: near the source, the kernel contains only
       source-term artefacts, not structural sensitivity.

    2. RECEIVER-LINE MUTING
       taper_rcv = product_r [1 - exp(-|x-xr|²/R²)]
       Suppresses 'rabbit ears' at each receiver position.

    3. ILLUMINATION PRECONDITIONING
       K_precond = K / (P_illum + ε)
       P_illum = ∫₀ᵀ |u_fwd|² dt (source-energy proxy)
       This compensates for uneven illumination — deep regions that
       receive less wave energy have artificially small kernels;
       dividing by illumination recovers the true sensitivity.

    4. DEPTH PRECONDITIONING
       K_depth = K · (1 + z/z₀)²
       Compensates geometric amplitude decay with depth.

    5. GAUSSIAN SMOOTHING
       Regularises the kernel spatially to remove numerical noise.
    """

    def __init__(self, acq):
        self.acq = acq
        # Pre-compute muting functions
        sx, sz = acq.SX*dx, acq.SZ*dz
        R2_src = (X - sx)**2 + (Z - sz)**2
        self.mute_src = 1. - np.exp(-R2_src / cfg.mute_R**2)

        # Receiver muting (combined taper over all receivers)
        mute_rcv = np.ones((cfg.Nx, cfg.Nz))
        for ir in range(cfg.n_recv):
            rx, rz = acq.RX[ir]*dx, acq.RZ[ir]*dz
            R2 = (X - rx)**2 + (Z - rz)**2
            mute_rcv *= 1. - np.exp(-R2 / cfg.mute_R**2)
        self.mute_rcv = mute_rcv

        # Combined muting taper
        self.mute = self.mute_src * self.mute_rcv

        # Depth preconditioning
        self.depth_precond = (1. + Z / cfg.depth_z0)**2

    def apply(self, K, illum, sigma=2.5):
        """
        Apply full conditioning pipeline.

        Parameters
        ----------
        K      : raw kernel [Nx, Nz]
        illum  : source illumination [Nx, Nz]
        sigma  : Gaussian smoothing σ [grid cells]

        Returns
        -------
        K_cond : conditioned kernel
        """
        # Step 1: Source+receiver muting
        K_muted = K * self.mute

        # Step 2: Illumination preconditioning
        illum_norm = illum / (np.max(illum) + 1e-30)
        K_illum = K_muted / (illum_norm + cfg.illum_eps)

        # Step 3: Depth preconditioning
        K_depth = K_illum * self.depth_precond

        # Step 4: Gaussian smoothing
        K_smooth = gaussian_filter(K_depth, sigma=sigma)

        # Step 5: Re-normalise to unit peak
        peak = np.max(np.abs(K_smooth))
        K_cond = K_smooth / peak if peak > 0 else K_smooth

        return K_cond

    def muting_only(self, K):
        """Apply only source/receiver muting (no illumination normalisation)."""
        return K * self.mute


# ══════════════════════════════════════════════════════════════════════════════
# MISFIT
# ══════════════════════════════════════════════════════════════════════════════
def misfit(svx, svz, ovx, ovz):
    """L2 waveform misfit and residuals."""
    rx = svx - ovx; rz = svz - ovz
    J  = 0.5*dt*(np.sum(rx**2) + np.sum(rz**2))
    return float(J), rx, rz


# ══════════════════════════════════════════════════════════════════════════════
# FRESNEL ZONE ANALYSIS
# ══════════════════════════════════════════════════════════════════════════════
def fresnel_analysis(K_Vs, acq, f0):
    """
    Measure the Fresnel zone half-width from the Vs kernel.

    Theory:  W_F(z) = √(Vp/(2·f₀) · z)
    Measurement: at each depth z, find the half-width of the central
    positive lobe of the kernel at x = src_x.

    Returns list of (z, W_measured, W_theoretical).
    """
    ix_src = int(cfg.src_x / dx)
    results = []
    depths  = np.arange(50, cfg.Lz - 50, 50)

    for z_val in depths:
        iz = int(z_val / dz)
        if iz >= cfg.Nz: continue

        # Horizontal slice at depth iz
        row = K_Vs[:, iz].copy()
        peak_val = row[ix_src]
        if abs(peak_val) < 1e-30: continue

        # Find half-width: distance from peak to first zero crossing
        half_level = peak_val / 2.
        # Search left from source
        left = ix_src
        while left > 0 and np.sign(row[left]) == np.sign(peak_val):
            left -= 1
        # Search right from source
        right = ix_src
        while right < cfg.Nx-1 and np.sign(row[right]) == np.sign(peak_val):
            right += 1

        W_meas = (right - left) * dx / 2.
        # Theoretical Fresnel half-width: W = sqrt(Vp * z / (2*f0))
        W_theory = np.sqrt(cfg.Vp_bg * z_val / (2. * f0))
        results.append((z_val, W_meas, W_theory))

    return results


# ══════════════════════════════════════════════════════════════════════════════
# PLOTTING SUITE
# ══════════════════════════════════════════════════════════════════════════════
def _sv(f, dpi=150):
    plt.tight_layout()
    plt.savefig(f, dpi=dpi, bbox_inches="tight")
    plt.close("all"); gc.collect()
    print(f"  ✓  {f}")

def _geom(ax, acq, ec="white", legend=True):
    ax.scatter([acq.SX*dx],[acq.SZ*dz],
               marker="*",s=220,c="red",zorder=10,label="Source")
    ax.scatter(acq.RX*dx, acq.RZ*dz,
               marker="v",s=40,c="yellow",ec="k",zorder=10,label="Receivers")
    c = plt.Circle((cfg.cx,cfg.cz), cfg.radius,
                   fill=False,ec=ec,lw=2.5,ls="--",label="Anomaly")
    ax.add_patch(c)
    if legend: ax.legend(fontsize=8,loc="lower right")

def _kernel_plot(ax, K, title, cmap="RdBu_r", vfrac=0.97):
    """Single-panel kernel / field plot."""
    vmax = np.percentile(np.abs(K), vfrac*100) + 1e-30
    norm = TwoSlopeNorm(vmin=-vmax, vcenter=0., vmax=vmax)
    im = ax.imshow(K.T, origin="upper", aspect="equal",
                   extent=[0,cfg.Lx,cfg.Lz,0],
                   cmap=cmap, norm=norm, interpolation="bilinear")
    return im

# ── 01 Model geometry ─────────────────────────────────────────────────────
def plot_geometry(model, acq):
    fig,axes = plt.subplots(1,3,figsize=(18,6))
    fields = [("Vp","P-wave Velocity [m/s]","RdYlBu_r"),
              ("Vs","S-wave Velocity [m/s]","RdYlBu_r"),
              ("rho","Density [kg/m³]","viridis")]
    for ax,(key,label,cmap) in zip(axes,fields):
        vmin = float(model[key].min()); vmax = float(model[key].max())
        im=ax.imshow(model[key].T,origin="upper",aspect="equal",
                     extent=[0,cfg.Lx,cfg.Lz,0],
                     cmap=cmap,vmin=vmin,vmax=vmax,interpolation="bilinear")
        plt.colorbar(im,ax=ax,shrink=0.85).set_label(label)
        _geom(ax,acq,ec="white")
        ax.set_xlabel("x [m]"); ax.set_ylabel("z [m]")
        ax.set_title(f"{label.split('[')[0].strip()} — True Model",
                     fontsize=11,fontweight="bold")
    fig.suptitle("Elastic Model Parameters — Camembert Coin Anomaly",
                 fontsize=13,fontweight="bold")
    _sv(f"{od}/01_model_geometry.png")

# ── 02 Ricker wavelet ─────────────────────────────────────────────────────
def plot_ricker(acq):
    freq=np.fft.rfftfreq(nt,d=dt); spec=np.abs(np.fft.rfft(acq.wav))
    fig,(a1,a2)=plt.subplots(1,2,figsize=(14,5))
    a1.plot(t_,acq.wav,c="#1565C0",lw=2.5)
    a1.fill_between(t_,acq.wav,alpha=0.15,color="#1565C0")
    a1.axvline(cfg.t0,ls="--",c="r",label=f"t₀={cfg.t0}s")
    a1.set_xlabel("Time [s]"); a1.set_ylabel("Amplitude")
    a1.set_title(f"Ricker Wavelet  f₀={acq.f0} Hz",fontsize=12,fontweight="bold")
    a1.legend(); a1.grid(alpha=0.35)
    a2.plot(freq,spec,c="#E65100",lw=2.5)
    a2.fill_between(freq,spec,alpha=0.15,color="#E65100")
    a2.axvline(acq.f0,ls="--",c="#1565C0",label=f"f₀={acq.f0} Hz")
    a2.set_xlim(0,4*acq.f0); a2.legend(); a2.grid(alpha=0.35)
    a2.set_xlabel("Frequency [Hz]"); a2.set_ylabel("|Amplitude|")
    a2.set_title("Amplitude Spectrum",fontsize=12,fontweight="bold")
    fig.suptitle("Ricker Source Wavelet",fontsize=13,fontweight="bold")
    _sv(f"{od}/02_ricker_wavelet.png")

# ── 03 Forward wavefield snapshots ────────────────────────────────────────
def plot_forward_snaps(snaps, acq):
    n=min(9,len(snaps)); idx=np.linspace(0,len(snaps)-1,n,dtype=int)
    show=[snaps[i] for i in idx]
    vmax=max(np.max(np.abs(s[1])) for s in show)*0.75+1e-30
    fig,axes=plt.subplots(3,3,figsize=(17,14))
    for ax,(t,wf) in zip(axes.flatten(),show):
        im=ax.imshow(wf.T,origin="upper",aspect="equal",
                     extent=[0,cfg.Lx,cfg.Lz,0],
                     cmap="seismic",vmin=-vmax,vmax=vmax,interpolation="bilinear")
        _geom(ax,acq,"cyan",legend=False)
        ax.set_title(f"t={t:.3f}s",fontsize=10,fontweight="bold")
        ax.set_xlabel("x[m]",fontsize=8); ax.set_ylabel("z[m]",fontsize=8)
    for ax in axes.flatten()[len(show):]: ax.set_visible(False)
    fig.suptitle("Forward Elastic Wavefield vx — True Camembert Model",
                 fontsize=13,fontweight="bold")
    _sv(f"{od}/03_forward_snapshots.png",dpi=110)

# ── 04 Forward wavefield GIF ──────────────────────────────────────────────
def make_forward_gif(snaps, acq):
    vmax=max(np.max(np.abs(s[1])) for s in snaps)*0.75+1e-30
    frames=[]; print(f"  GIF ({len(snaps)} frames)",flush=True)
    for t_v,wf in snaps:
        fig,ax=plt.subplots(figsize=(6.5,6.5)); ax.set_facecolor("k")
        im=ax.imshow(wf.T,origin="upper",aspect="equal",
                     extent=[0,cfg.Lx,cfg.Lz,0],
                     cmap="seismic",vmin=-vmax,vmax=vmax,interpolation="bilinear")
        plt.colorbar(im,ax=ax,label="vx",shrink=0.85)
        _geom(ax,acq,"cyan",legend=False)
        ax.set_xlabel("x [m]"); ax.set_ylabel("z [m]")
        ax.set_title(f"Forward Wavefield   t={t_v:.3f}s",fontsize=10,fontweight="bold")
        plt.tight_layout()
        buf=io.BytesIO(); plt.savefig(buf,format="png",dpi=85,bbox_inches="tight")
        buf.seek(0); frames.append(Image.open(buf).copy()); plt.close("all")
    frames[0].save(f"{od}/04_forward_wavefield.gif",save_all=True,
                   append_images=frames[1:],loop=0,duration=130,optimize=True)
    del frames; gc.collect(); print(f"  ✓  {od}/04_forward_wavefield.gif")

# ── 05 Seismograms ────────────────────────────────────────────────────────
def plot_seismograms(svx, svz):
    fig,axes=plt.subplots(1,2,figsize=(16,7))
    for ax,data,comp in zip(axes,[svx,svz],["vx — Horizontal","vz — Vertical"]):
        d=data.copy()
        for i in range(d.shape[0]):
            mx=np.max(np.abs(d[i])); d[i]=d[i]/mx if mx>0 else d[i]
        im=ax.imshow(d,aspect="auto",extent=[0,cfg.T,cfg.n_recv,0],
                     cmap="RdBu",vmin=-1,vmax=1)
        plt.colorbar(im,ax=ax,label="Norm. amplitude")
        for i in range(0,cfg.n_recv,5):
            ax.plot(t_,i+d[i]*1.4,"k-",lw=0.5,alpha=0.55)
        ax.set_xlabel("Time [s]"); ax.set_ylabel("Receiver index")
        ax.set_title(comp,fontsize=12,fontweight="bold")
    fig.suptitle("Observed Seismograms — True Model",fontsize=13,fontweight="bold")
    _sv(f"{od}/05_seismograms.png")

# ── 06 Adjoint wavefield snapshots ────────────────────────────────────────
def plot_adjoint_snaps(snaps, acq):
    n=min(9,len(snaps)); idx=np.linspace(0,len(snaps)-1,n,dtype=int)
    show=[snaps[i] for i in idx]
    vmax=max(np.max(np.abs(s[1])) for s in show)*0.75+1e-30
    fig,axes=plt.subplots(3,3,figsize=(17,14))
    for ax,(t,wf) in zip(axes.flatten(),show):
        im=ax.imshow(wf.T,origin="upper",aspect="equal",
                     extent=[0,cfg.Lx,cfg.Lz,0],
                     cmap="PiYG",vmin=-vmax,vmax=vmax,interpolation="bilinear")
        _geom(ax,acq,"coral",legend=False)
        ax.set_title(f"t={t:.3f}s (bwd)",fontsize=10,fontweight="bold")
        ax.set_xlabel("x[m]",fontsize=8); ax.set_ylabel("z[m]",fontsize=8)
    for ax in axes.flatten()[len(show):]: ax.set_visible(False)
    fig.suptitle("Adjoint Wavefield vxᴬᵈʲ (Time-Reversed Residuals)",
                 fontsize=13,fontweight="bold")
    _sv(f"{od}/06_adjoint_snapshots.png",dpi=110)

# ── 07 Raw kernels 4-panel ─────────────────────────────────────────────────
def plot_raw_kernels(K_Vp, K_Vs, K_lam, K_mu, acq):
    fig,axes=plt.subplots(2,2,figsize=(17,14))
    kernels = [
        (K_lam,"K_λ (λ-kernel)","RdBu_r"),
        (K_mu, "K_μ (μ-kernel)","RdBu_r"),
        (K_Vp, "K_Vp = ∂J/∂Vp","PuOr"),
        (K_Vs, "K_Vs = ∂J/∂Vs","RdBu_r"),
    ]
    for ax,(K,title,cmap) in zip(axes.flatten(),kernels):
        im = _kernel_plot(ax,K,title,cmap)
        plt.colorbar(im,ax=ax,shrink=0.85).set_label("Kernel amplitude")
        _geom(ax,acq,"lime",legend=False)
        ax.set_xlabel("x [m]"); ax.set_ylabel("z [m]")
        ax.set_title(title,fontsize=12,fontweight="bold")
    fig.suptitle("Raw Elastic Sensitivity Kernels (Before Conditioning)",
                 fontsize=14,fontweight="bold")
    _sv(f"{od}/07_raw_kernels.png")

# ── 08 Conditioned kernels 4-panel ────────────────────────────────────────
def plot_conditioned_kernels(K_Vp_c, K_Vs_c, K_lam_c, K_mu_c,
                             K_Vp_r, K_Vs_r, acq, conditioner):
    fig = plt.figure(figsize=(20,16))
    gs  = gridspec.GridSpec(3,3,figure=fig,hspace=0.35,wspace=0.3)

    # Row 0: raw Vp, raw Vs, muting function
    ax00=fig.add_subplot(gs[0,0]); ax01=fig.add_subplot(gs[0,1]); ax02=fig.add_subplot(gs[0,2])
    # Row 1: conditioned Vp, conditioned Vs, illumination preview
    ax10=fig.add_subplot(gs[1,0]); ax11=fig.add_subplot(gs[1,1]); ax12=fig.add_subplot(gs[1,2])
    # Row 2: K_lam conditioned, K_mu conditioned, combined mute map
    ax20=fig.add_subplot(gs[2,0]); ax21=fig.add_subplot(gs[2,1]); ax22=fig.add_subplot(gs[2,2])

    panels_r0 = [(K_Vp_r,"Raw K_Vp","PuOr"),(K_Vs_r,"Raw K_Vs","RdBu_r")]
    panels_r1 = [(K_Vp_c,"Conditioned K_Vp","PuOr"),(K_Vs_c,"Conditioned K_Vs","RdBu_r")]
    panels_r2 = [(K_lam_c,"Conditioned K_λ","RdBu_r"),(K_mu_c,"Conditioned K_μ","RdBu_r")]

    for ax,(K,ttl,cmap) in zip([ax00,ax01],panels_r0):
        im=_kernel_plot(ax,K,ttl,cmap)
        plt.colorbar(im,ax=ax,shrink=0.85).set_label("Kernel")
        _geom(ax,acq,"lime",legend=False)
        ax.set_xlabel("x [m]"); ax.set_ylabel("z [m]")
        ax.set_title(ttl,fontsize=11,fontweight="bold")

    # Muting function
    im=ax02.imshow(conditioner.mute.T,origin="upper",aspect="equal",
                   extent=[0,cfg.Lx,cfg.Lz,0],cmap="Greys_r",vmin=0,vmax=1)
    plt.colorbar(im,ax=ax02,shrink=0.85).set_label("Mute weight [0-1]")
    _geom(ax02,acq,"red",legend=False)
    ax02.set_xlabel("x [m]"); ax02.set_ylabel("z [m]")
    ax02.set_title("Combined Muting Function",fontsize=11,fontweight="bold")

    for ax,(K,ttl,cmap) in zip([ax10,ax11],panels_r1):
        im=_kernel_plot(ax,K,ttl,cmap)
        plt.colorbar(im,ax=ax,shrink=0.85).set_label("Norm. Kernel")
        _geom(ax,acq,"lime",legend=False)
        ax.set_xlabel("x [m]"); ax.set_ylabel("z [m]")
        ax.set_title(ttl,fontsize=11,fontweight="bold")

    # Depth preconditioning profile
    z_prof = z_.copy()
    dp = (1. + z_prof/cfg.depth_z0)**2
    ax12.plot(dp, z_prof, c="#E65100",lw=2.5)
    ax12.fill_betweenx(z_prof,1,dp,alpha=0.2,color="#E65100")
    ax12.set_xlabel("Depth precond. weight"); ax12.set_ylabel("z [m]")
    ax12.set_title(f"Depth Preconditioning\n(1+z/{cfg.depth_z0:.0f})²",
                   fontsize=11,fontweight="bold")
    ax12.invert_yaxis(); ax12.grid(alpha=0.35)

    for ax,(K,ttl,cmap) in zip([ax20,ax21],panels_r2):
        im=_kernel_plot(ax,K,ttl,cmap)
        plt.colorbar(im,ax=ax,shrink=0.85).set_label("Norm. Kernel")
        _geom(ax,acq,"lime",legend=False)
        ax.set_xlabel("x [m]"); ax.set_ylabel("z [m]")
        ax.set_title(ttl,fontsize=11,fontweight="bold")

    # Source muting alone
    im=ax22.imshow(conditioner.mute_src.T,origin="upper",aspect="equal",
                   extent=[0,cfg.Lx,cfg.Lz,0],cmap="Greys_r",vmin=0,vmax=1)
    plt.colorbar(im,ax=ax22,shrink=0.85).set_label("Mute weight")
    _geom(ax22,acq,"red",legend=False)
    ax22.set_xlabel("x [m]"); ax22.set_ylabel("z [m]")
    ax22.set_title("Source-Point Muting Taper",fontsize=11,fontweight="bold")

    fig.suptitle("Elastic Kernel Conditioning Pipeline\n"
                 "Raw → Muted → Illumination-Preconditioned → Depth-Compensated → Smoothed",
                 fontsize=13,fontweight="bold")
    _sv(f"{od}/08_conditioned_kernels.png",dpi=110)

# ── 09 Raw vs conditioned comparison ─────────────────────────────────────
def plot_comparison(K_Vs_raw, K_Vs_cond, acq):
    fig,axes=plt.subplots(1,3,figsize=(19,7))
    diff = K_Vs_raw/np.max(np.abs(K_Vs_raw)+1e-30) - K_Vs_cond

    for ax,K,ttl in zip(axes,
            [K_Vs_raw, K_Vs_cond, diff],
            ["Raw K_Vs (with singularities)",
             "Conditioned K_Vs (muted + precond.)",
             "Difference: Raw − Conditioned"]):
        im=_kernel_plot(ax,K,ttl)
        plt.colorbar(im,ax=ax,shrink=0.85).set_label("Kernel amplitude")
        _geom(ax,acq,"lime",legend=False)
        ax.set_xlabel("x [m]"); ax.set_ylabel("z [m]")
        ax.set_title(ttl,fontsize=11,fontweight="bold")
    fig.suptitle("K_Vs Sensitivity Kernel: Raw vs Conditioned",
                 fontsize=13,fontweight="bold")
    _sv(f"{od}/09_kernel_comparison.png")

# ── 10 Fresnel zone analysis ──────────────────────────────────────────────
def plot_fresnel(fresnel_data_list, freq_list):
    fig,(a1,a2)=plt.subplots(1,2,figsize=(16,7))

    colors=["#1565C0","#E65100","#2E7D32"]
    for (fdata,f0),c in zip(zip(fresnel_data_list,freq_list),colors):
        if not fdata: continue
        depths   = [r[0] for r in fdata]
        W_meas   = [r[1] for r in fdata]
        W_theory = [r[2] for r in fdata]
        a1.plot(W_meas,   depths, "-o", c=c, lw=2,   ms=5, label=f"Measured f₀={f0}Hz")
        a1.plot(W_theory, depths, "--", c=c, lw=1.5,        label=f"Theory f₀={f0}Hz")
        ratio = [m/(t+1e-10) for m,t in zip(W_meas,W_theory)]
        a2.plot(ratio, depths, "-o", c=c, lw=2, ms=5, label=f"f₀={f0}Hz")

    a1.set_xlabel("Fresnel Half-Width [m]",fontsize=12)
    a1.set_ylabel("Depth z [m]",fontsize=12)
    a1.set_title("Fresnel Zone Width vs Depth",fontsize=12,fontweight="bold")
    a1.invert_yaxis(); a1.legend(fontsize=9); a1.grid(alpha=0.35)
    z_th = np.linspace(50,cfg.Lz-50,100)
    for f0,c in zip(freq_list,colors):
        W_th = np.sqrt(cfg.Vp_bg*z_th/(2.*f0))
        a1.plot(W_th,z_th,c=c,lw=0.7,alpha=0.5)

    a2.set_xlabel("W_measured / W_theoretical",fontsize=12)
    a2.set_ylabel("Depth z [m]",fontsize=12)
    a2.set_title("Measured/Theoretical Ratio",fontsize=12,fontweight="bold")
    a2.invert_yaxis(); a2.axvline(1.0,c="k",ls="--",lw=1.5,label="Ideal ratio=1")
    a2.legend(fontsize=9); a2.grid(alpha=0.35)
    fig.suptitle("Fresnel Zone Width Analysis — K_Vs Kernel",
                 fontsize=13,fontweight="bold")
    _sv(f"{od}/10_fresnel_analysis.png")

# ── 11 Depth profiles ─────────────────────────────────────────────────────
def plot_depth_profiles(K_Vp, K_Vs, K_rho, K_lam, K_mu, K_Vs_c):
    ix_src = int(cfg.src_x / dx)
    fig,axes=plt.subplots(1,3,figsize=(18,8))

    def profile(K,label,c,ax,ls="-"):
        vals = K[ix_src,:].copy()
        vals = vals/np.max(np.abs(vals)+1e-30)
        ax.plot(vals, z_, ls, c=c, lw=2, label=label)

    profile(K_Vp,  "K_Vp (raw)",   "#E65100", axes[0])
    profile(K_Vs,  "K_Vs (raw)",   "#1565C0", axes[0])
    profile(K_rho, "K_ρ  (raw)",   "#2E7D32", axes[0])
    axes[0].set_title(f"1D Vertical Profiles at x={cfg.src_x:.0f}m",
                      fontsize=11,fontweight="bold")

    profile(K_lam, "K_λ (raw)", "#7B1FA2", axes[1])
    profile(K_mu,  "K_μ (raw)", "#BF360C", axes[1])
    axes[1].set_title("Lamé Parameter Kernels",fontsize=11,fontweight="bold")

    profile(K_Vs,   "K_Vs raw",  "#1565C0", axes[2])
    profile(K_Vs_c, "K_Vs cond.","#E65100", axes[2], "--")
    axes[2].set_title("K_Vs: Raw vs Conditioned",fontsize=11,fontweight="bold")

    for ax in axes:
        ax.axhline(cfg.cz-cfg.radius,c="gray",ls=":",alpha=0.7)
        ax.axhline(cfg.cz+cfg.radius,c="gray",ls=":",alpha=0.7,label="Anomaly edges")
        ax.axhline(cfg.cz,c="red",ls="--",alpha=0.6,label=f"Anomaly z={cfg.cz:.0f}m")
        ax.axvline(0,c="k",lw=0.8,alpha=0.5)
        ax.invert_yaxis(); ax.legend(fontsize=9); ax.grid(alpha=0.35)
        ax.set_xlabel("Norm. kernel amplitude"); ax.set_ylabel("Depth z [m]")

    fig.suptitle("1D Vertical Kernel Profiles Through Source x",
                 fontsize=13,fontweight="bold")
    _sv(f"{od}/11_depth_profiles.png")

# ── 12 Horizontal profiles ────────────────────────────────────────────────
def plot_horizontal_profiles(K_Vs, K_Vs_c):
    depths_plot = [500., 800., 1500., 2000.]
    fig,axes = plt.subplots(2,2,figsize=(16,12))
    for ax, z_val in zip(axes.flatten(), depths_plot):
        iz = int(z_val / dz)
        raw  = K_Vs[:,iz].copy(); raw  /= np.max(np.abs(raw)+1e-30)
        cond = K_Vs_c[:,iz].copy()
        ax.plot(x_, raw,  c="#1565C0", lw=2.5, label="K_Vs raw")
        ax.plot(x_, cond, c="#E65100", lw=2.5, ls="--", label="K_Vs conditioned")
        # Theoretical Fresnel zone width
        W_F = np.sqrt(cfg.Vp_bg * z_val / (2.*cfg.f0))
        ax.axvspan(cfg.src_x - W_F, cfg.src_x + W_F,
                   alpha=0.12, color="green", label=f"Fresnel ±{W_F:.0f}m")
        ax.axvline(cfg.src_x,c="red",ls="--",lw=1.5,alpha=0.6,label="Source x")
        ax.set_xlabel("x [m]"); ax.set_ylabel("Norm. K_Vs")
        ax.set_title(f"Horizontal Profile at z={z_val:.0f}m",
                     fontsize=11,fontweight="bold")
        ax.legend(fontsize=8); ax.grid(alpha=0.35)
        ax.set_xlim(0, cfg.Lx)
    fig.suptitle("Horizontal K_Vs Profiles at Selected Depths\n"
                 "Green band = theoretical Fresnel zone",
                 fontsize=13,fontweight="bold")
    _sv(f"{od}/12_horizontal_profiles.png")

# ── 13 Multi-frequency kernels ────────────────────────────────────────────
def plot_multi_freq(kernels_list, freq_list, acq):
    n = len(freq_list)
    fig,axes = plt.subplots(2,n,figsize=(7*n,13))
    if n == 1: axes = axes.reshape(2,1)
    for col,(K_Vs,K_Vp,f0) in enumerate(kernels_list):
        for row,(K,ttl) in enumerate([(K_Vs,f"K_Vs  f₀={f0}Hz"),(K_Vp,f"K_Vp  f₀={f0}Hz")]):
            ax = axes[row,col]
            im = _kernel_plot(ax,K,ttl,"RdBu_r" if row==0 else "PuOr")
            plt.colorbar(im,ax=ax,shrink=0.85).set_label("Kernel")
            _geom(ax,acq,"lime",legend=False)
            W_F = np.sqrt(cfg.Vp_bg * cfg.cz / (2.*f0))
            ax.add_patch(plt.Circle((cfg.src_x,cfg.cz),W_F,
                fill=False,ec="yellow",lw=2,ls="-.",
                label=f"Fresnel r={W_F:.0f}m"))
            ax.set_xlabel("x [m]"); ax.set_ylabel("z [m]")
            ax.set_title(ttl,fontsize=11,fontweight="bold")
    fig.suptitle("Multi-Frequency Sensitivity Kernels\n"
                 "Yellow circle = Fresnel zone at anomaly depth",
                 fontsize=13,fontweight="bold")
    _sv(f"{od}/13_multi_frequency.png",dpi=110)

# ── 14 Muting visualisation ───────────────────────────────────────────────
def plot_muting(K_Vs, K_Vs_muted, K_Vs_cond, conditioner, acq):
    fig,axes = plt.subplots(2,3,figsize=(19,13))
    steps = [
        (K_Vs,                   "① Raw K_Vs",                  "RdBu_r"),
        (K_Vs*conditioner.mute_src, "② Source-Point Muted",      "RdBu_r"),
        (K_Vs_muted,             "③ Full Muting (src+recv)",    "RdBu_r"),
        (conditioner.mute_src,   "Source Muting Taper",         "Greys_r"),
        (conditioner.mute_rcv,   "Receiver Muting Taper",       "Greys_r"),
        (K_Vs_cond,              "⑥ Final Conditioned K_Vs",    "RdBu_r"),
    ]
    for ax,(K,ttl,cmap) in zip(axes.flatten(),steps):
        if cmap=="Greys_r":
            im=ax.imshow(K.T,origin="upper",aspect="equal",
                         extent=[0,cfg.Lx,cfg.Lz,0],
                         cmap=cmap,vmin=0,vmax=1,interpolation="bilinear")
        else:
            im=_kernel_plot(ax,K,ttl,cmap)
        plt.colorbar(im,ax=ax,shrink=0.85)
        _geom(ax,acq,"lime" if cmap!="Greys_r" else "red",legend=False)
        ax.set_xlabel("x [m]"); ax.set_ylabel("z [m]")
        ax.set_title(ttl,fontsize=11,fontweight="bold")
    fig.suptitle("K_Vs Muting Pipeline — Removing Source-Point Singularities",
                 fontsize=13,fontweight="bold")
    _sv(f"{od}/14_source_receiver_muting.png",dpi=110)

# ── 15 Illumination ───────────────────────────────────────────────────────
def plot_illumination(illum, K_Vs, acq):
    fig,axes = plt.subplots(1,3,figsize=(19,7))
    illum_n = illum / (np.max(illum)+1e-30)
    im0=axes[0].imshow(illum_n.T,origin="upper",aspect="equal",
                       extent=[0,cfg.Lx,cfg.Lz,0],
                       cmap="hot_r",vmin=0,vmax=1,interpolation="bilinear")
    plt.colorbar(im0,ax=axes[0],shrink=0.88).set_label("Norm. illumination")
    _geom(axes[0],acq,"cyan",legend=False)
    axes[0].set_title("Source Illumination  P(x) = ∫|vx|²dt",
                      fontsize=11,fontweight="bold")

    K_illum = K_Vs / (illum_n + cfg.illum_eps)
    K_illum /= np.max(np.abs(K_illum)+1e-30)
    im1=_kernel_plot(axes[1],K_Vs,"K_Vs (raw)")
    plt.colorbar(im1,ax=axes[1],shrink=0.88).set_label("Kernel")
    _geom(axes[1],acq,"lime",legend=False)
    axes[1].set_title("K_Vs Raw",fontsize=11,fontweight="bold")

    im2=_kernel_plot(axes[2],K_illum,"K_Vs / Illumination")
    plt.colorbar(im2,ax=axes[2],shrink=0.88).set_label("Precond. Kernel")
    _geom(axes[2],acq,"lime",legend=False)
    axes[2].set_title("K_Vs ÷ Illumination",fontsize=11,fontweight="bold")

    for ax in axes:
        ax.set_xlabel("x [m]"); ax.set_ylabel("z [m]")

    fig.suptitle("Source Illumination & Illumination Preconditioning",
                 fontsize=13,fontweight="bold")
    _sv(f"{od}/15_illumination.png")

# ── 16 Publication-quality 6-panel summary ─────────────────────────────────
def plot_summary(model, K_Vp_c, K_Vs_c, K_lam_c, K_mu_c,
                 conditioner, illum, acq):
    fig = plt.figure(figsize=(20,14))
    gs  = gridspec.GridSpec(2,3,figure=fig,hspace=0.38,wspace=0.28)

    ax1=fig.add_subplot(gs[0,0]); ax2=fig.add_subplot(gs[0,1])
    ax3=fig.add_subplot(gs[0,2]); ax4=fig.add_subplot(gs[1,0])
    ax5=fig.add_subplot(gs[1,1]); ax6=fig.add_subplot(gs[1,2])

    # Panel 1: Vp model
    im=ax1.imshow(model["Vp"].T,origin="upper",aspect="equal",
                  extent=[0,cfg.Lx,cfg.Lz,0],
                  cmap="RdYlBu_r",
                  vmin=model["Vp"].min(),vmax=model["Vp"].max(),
                  interpolation="bilinear")
    plt.colorbar(im,ax=ax1,shrink=0.88).set_label("Vp [m/s]")
    _geom(ax1,acq,"white"); ax1.set_title("(a) Vp Model",fontweight="bold",fontsize=12)

    # Panel 2: Vs model
    im=ax2.imshow(model["Vs"].T,origin="upper",aspect="equal",
                  extent=[0,cfg.Lx,cfg.Lz,0],
                  cmap="RdYlBu_r",
                  vmin=model["Vs"].min(),vmax=model["Vs"].max(),
                  interpolation="bilinear")
    plt.colorbar(im,ax=ax2,shrink=0.88).set_label("Vs [m/s]")
    _geom(ax2,acq,"white"); ax2.set_title("(b) Vs Model",fontweight="bold",fontsize=12)

    # Panel 3: Source illumination
    illum_n=illum/(np.max(illum)+1e-30)
    im=ax3.imshow(illum_n.T,origin="upper",aspect="equal",
                  extent=[0,cfg.Lx,cfg.Lz,0],
                  cmap="hot_r",vmin=0,vmax=1,interpolation="bilinear")
    plt.colorbar(im,ax=ax3,shrink=0.88).set_label("Norm. illumination")
    _geom(ax3,acq,"cyan"); ax3.set_title("(c) Source Illumination",fontweight="bold",fontsize=12)

    # Panel 4: Conditioned K_Vp
    im=_kernel_plot(ax4,K_Vp_c,"","PuOr")
    plt.colorbar(im,ax=ax4,shrink=0.88).set_label("∂J/∂Vp (norm.)")
    _geom(ax4,acq,"lime"); ax4.set_title("(d) Conditioned K_Vp",fontweight="bold",fontsize=12)

    # Panel 5: Conditioned K_Vs — the MAIN RESULT (heatmap)
    im=_kernel_plot(ax5,K_Vs_c,"","RdBu_r")
    plt.colorbar(im,ax=ax5,shrink=0.88).set_label("∂J/∂Vs (norm.)")
    _geom(ax5,acq,"lime")
    ax5.set_title("(e) Conditioned K_Vs ← Main Result",fontweight="bold",fontsize=12,
                  color="#1565C0")

    # Panel 6: Muting function
    im=ax6.imshow(conditioner.mute.T,origin="upper",aspect="equal",
                  extent=[0,cfg.Lx,cfg.Lz,0],
                  cmap="Greys_r",vmin=0,vmax=1,interpolation="bilinear")
    plt.colorbar(im,ax=ax6,shrink=0.88).set_label("Mute weight")
    _geom(ax6,acq,"red"); ax6.set_title("(f) Combined Muting Taper",fontweight="bold",fontsize=12)

    for ax in [ax1,ax2,ax3,ax4,ax5,ax6]:
        ax.set_xlabel("x [m]",fontsize=10); ax.set_ylabel("z [m]",fontsize=10)

    fig.suptitle("Elastic Sensitivity Kernels — Firedrake-Adjoint Framework\n"
                 f"Camembert Model | f₀={cfg.f0}Hz | Adjoint-State Method",
                 fontsize=14,fontweight="bold")
    _sv(f"{od}/16_kernel_summary.png",dpi=130)


# ══════════════════════════════════════════════════════════════════════════════
# CSV EXPORTS
# ══════════════════════════════════════════════════════════════════════════════
def export_fresnel_csv(fresnel_data_list, freq_list):
    fname = f"{od}/fresnel_zone_table.csv"
    with open(fname,"w",newline="") as f:
        w = csv.writer(f)
        w.writerow(["Frequency_Hz","Depth_m",
                    "W_measured_m","W_theoretical_m",
                    "Ratio_meas_theory",
                    "Lambda_m","Fresnel_formula"])
        for fdata,f0 in zip(fresnel_data_list,freq_list):
            for z_val,W_m,W_t in fdata:
                lam = cfg.Vp_bg / f0
                w.writerow([f0, f"{z_val:.1f}",
                             f"{W_m:.2f}", f"{W_t:.2f}",
                             f"{W_m/(W_t+1e-10):.4f}",
                             f"{lam:.2f}",
                             f"sqrt(Vp*z/(2*f0))=sqrt({cfg.Vp_bg}*z/(2*{f0}))"])
    print(f"  ✓  {fname}")

def export_kernel_stats_csv(kernels_dict):
    fname = f"{od}/kernel_statistics.csv"
    with open(fname,"w",newline="") as f:
        w = csv.writer(f)
        w.writerow(["Kernel","Min","Max","Mean","Std","L1_norm","L2_norm",
                    "Perc_1","Perc_99","Peak_x_m","Peak_z_m"])
        for name,K in kernels_dict.items():
            flat = K.ravel()
            pidx = np.argmax(np.abs(flat))
            px = (pidx // cfg.Nz) * dx
            pz = (pidx %  cfg.Nz) * dz
            w.writerow([name,
                        f"{flat.min():.4e}",f"{flat.max():.4e}",
                        f"{flat.mean():.4e}",f"{flat.std():.4e}",
                        f"{np.sum(np.abs(flat))*dx*dz:.4e}",
                        f"{np.sqrt(np.sum(flat**2))*dx*dz:.4e}",
                        f"{np.percentile(flat,1):.4e}",
                        f"{np.percentile(flat,99):.4e}",
                        f"{px:.1f}",f"{pz:.1f}"])
    print(f"  ✓  {fname}")


# ══════════════════════════════════════════════════════════════════════════════
# MAIN
# ══════════════════════════════════════════════════════════════════════════════
if __name__ == "__main__":
    print("\n" + "="*65)
    print("  ELASTIC SENSITIVITY KERNELS — FIREDRAKE-ADJOINT FRAMEWORK")
    print("  Camembert Model | Vp & Vs Kernels | Full Conditioning")
    print("="*65)

    # ── Build models ──────────────────────────────────────────────────────
    model_true = build_model(anomaly=True)
    model_hom  = build_model(anomaly=False)    # homogeneous starting model

    # ── Primary acquisition (f0 = 5 Hz) ──────────────────────────────────
    acq = Acquisition(f0=cfg.f0)
    print(f"\n  Acquisition : {acq}")

    # ── Pre-run plots ─────────────────────────────────────────────────────
    print("\n[1/6] Pre-run plots ...")
    plot_geometry(model_true, acq)
    plot_ricker(acq)

    # ── Phase 1: Forward solve on TRUE model (observed data) ──────────────
    print("\n[2/6] Forward solve — true Camembert model ...")
    ovx,ovz,true_snaps,_,_,illum = forward(
        model_true, acq, store=False, snap_n=16, tag="true-model")
    print(f"  obs max|vx|={np.max(np.abs(ovx)):.3e}")

    plot_forward_snaps(true_snaps, acq)
    make_forward_gif(true_snaps, acq)
    plot_seismograms(ovx, ovz)
    del true_snaps; gc.collect()

    # ── Phase 2: Forward solve on HOMOGENEOUS model (synthetic data) ──────
    print("\n[3/6] Forward solve — homogeneous model + wavefield storage ...")
    svx,svz,_,VXA,VZA,illum = forward(
        model_hom, acq, store=True, snap_n=4, tag="hom-model")
    J, res_vx, res_vz = misfit(svx, svz, ovx, ovz)
    print(f"  J = {J:.4e}   max|res|={np.max(np.abs(res_vx)):.3e}")
    del svx, svz; gc.collect()

    # ── Phase 3: Adjoint solve → raw kernels ─────────────────────────────
    print("\n[4/6] Adjoint solve → computing sensitivity kernels ...")
    K_Vp, K_Vs, K_rho, K_lam, K_mu, adj_snaps = adjoint_kernels(
        model_hom, acq, res_vx, res_vz, VXA, VZA,
        snap_n=16, tag="adjoint")
    del VXA, VZA, res_vx, res_vz; gc.collect()

    plot_adjoint_snaps(adj_snaps, acq)
    del adj_snaps; gc.collect()

    # ── Phase 4: Kernel conditioning ─────────────────────────────────────
    print("\n[5/6] Conditioning kernels (muting + preconditioning) ...")
    conditioner = KernelConditioner(acq)

    K_Vp_muted = conditioner.muting_only(K_Vp)
    K_Vs_muted = conditioner.muting_only(K_Vs)

    K_Vp_cond = conditioner.apply(K_Vp, illum, sigma=2.5)
    K_Vs_cond = conditioner.apply(K_Vs, illum, sigma=2.5)
    K_lam_cond = conditioner.apply(K_lam, illum, sigma=2.0)
    K_mu_cond  = conditioner.apply(K_mu,  illum, sigma=2.0)

    # ── Phase 5: Multi-frequency kernels ──────────────────────────────────
    print("\n  Multi-frequency kernel computation ...")
    multi_kernels = []
    for f0_mf in cfg.freqs:
        acq_mf = Acquisition(f0=f0_mf)
        print(f"    f0={f0_mf}Hz ...",end="")
        ovx_mf,ovz_mf,_,_,_,_ = forward(model_true, acq_mf,
                                          store=False, snap_n=2, tag=f"true-{f0_mf}Hz")
        svx_mf,svz_mf,_,VXA_mf,VZA_mf,_ = forward(model_hom, acq_mf,
                                                     store=True, snap_n=2, tag=f"hom-{f0_mf}Hz")
        J_mf,rx_mf,rz_mf = misfit(svx_mf, svz_mf, ovx_mf, ovz_mf)
        Kvp_mf,Kvs_mf,_,_,_,_ = adjoint_kernels(model_hom, acq_mf,
                                                   rx_mf, rz_mf, VXA_mf, VZA_mf,
                                                   snap_n=2, tag=f"adj-{f0_mf}Hz")
        Kvp_mc = conditioner.apply(Kvp_mf, illum, sigma=1.5)
        Kvs_mc = conditioner.apply(Kvs_mf, illum, sigma=1.5)
        multi_kernels.append((Kvs_mc, Kvp_mc, f0_mf))
        del VXA_mf, VZA_mf, rx_mf, rz_mf; gc.collect()

    # ── Phase 6: Fresnel analysis ─────────────────────────────────────────
    print("\n  Fresnel zone analysis ...")
    fresnel_list = []
    for (Kvs_c,_,f0_mf) in multi_kernels:
        acq_mf = Acquisition(f0=f0_mf)
        fdata  = fresnel_analysis(Kvs_c, acq_mf, f0_mf)
        fresnel_list.append(fdata)

    # ── Phase 6: All output plots ─────────────────────────────────────────
    print("\n[6/6] Generating all output plots + CSV ...")

    plot_raw_kernels(K_Vp, K_Vs, K_lam, K_mu, acq)
    plot_conditioned_kernels(K_Vp_cond, K_Vs_cond, K_lam_cond, K_mu_cond,
                             K_Vp, K_Vs, acq, conditioner)
    plot_comparison(K_Vs, K_Vs_cond, acq)
    plot_fresnel(fresnel_list, cfg.freqs)
    plot_depth_profiles(K_Vp, K_Vs, K_rho, K_lam, K_mu, K_Vs_cond)
    plot_horizontal_profiles(K_Vs, K_Vs_cond)
    plot_multi_freq(multi_kernels, cfg.freqs, acq)
    plot_muting(K_Vs, K_Vs_muted, K_Vs_cond, conditioner, acq)
    plot_illumination(illum, K_Vs, acq)
    plot_summary(model_true, K_Vp_cond, K_Vs_cond, K_lam_cond, K_mu_cond,
                 conditioner, illum, acq)

    # ── CSV exports ───────────────────────────────────────────────────────
    export_fresnel_csv(fresnel_list, cfg.freqs)
    export_kernel_stats_csv({
        "K_Vp_raw":K_Vp,"K_Vs_raw":K_Vs,"K_rho_raw":K_rho,
        "K_lam_raw":K_lam,"K_mu_raw":K_mu,
        "K_Vp_cond":K_Vp_cond,"K_Vs_cond":K_Vs_cond,
    })

    # ── Final summary ─────────────────────────────────────────────────────
    print(f"\n{'='*65}")
    print(f"  ALL OUTPUTS → '{od}/'")
    print(f"{'='*65}")
    outputs = [
        "01_model_geometry.png",
        "02_ricker_wavelet.png",
        "03_forward_snapshots.png",
        "04_forward_wavefield.gif",
        "05_seismograms.png",
        "06_adjoint_snapshots.png",
        "07_raw_kernels.png",
        "08_conditioned_kernels.png",
        "09_kernel_comparison.png",
        "10_fresnel_analysis.png",
        "11_depth_profiles.png",
        "12_horizontal_profiles.png",
        "13_multi_frequency.png",
        "14_source_receiver_muting.png",
        "15_illumination.png",
        "16_kernel_summary.png",
        "fresnel_zone_table.csv",
        "kernel_statistics.csv",
    ]
    for f in outputs:
        size = os.path.getsize(f"{od}/{f}")//1024 if os.path.exists(f"{od}/{f}") else 0
        print(f"  ✓  {f:<40} {size:>5} KB")
    print(f"{'='*65}\n")