#!/usr/bin/env python3
"""
kzk_sim.py - a simple 2D KZK (Khokhlov-Zabolotskaya-Kuznetsov) simulator
for a pure-tone (sine wave) source, or an amplitude-modulated (AM) carrier
in air: the parametric (directional) speaker.

Geometry: 2D Cartesian. One transverse coordinate x and the propagation
axis z (a strip source that is uniform in y). Time is the retarded time
tau = t - z/c0, so the plotted waveform travels along with the wave.

  d2p/(dz dtau) = (c0/2) d2p/dx2 + (delta/(2 c0^3)) d3p/dtau3 + (beta/(2 rho0 c0^3)) d2(p^2)/dtau2
                  [diffraction]     [absorption]                  [nonlinearity]

Source at z = 0 (one period of a pure tone, Gaussian-shaded across x):

  p(x, 0, tau) = p0 * exp(-x^2 / a^2) * sin(2 pi f0 tau)

AM mode (--am): the carrier f0 is modulated by a tone fm, and the time window
is one period of fm instead of one period of f0:

  p(x, 0, tau) = p0 * exp(-x^2 / a^2) * [1 + m cos(2 pi fm tau)] * sin(2 pi f0 tau)

The nonlinearity of air demodulates the beam, so the on-axis waveform picks up
an audible part at fm, 2 fm, 3 fm, ... ; that part is filtered out (< 20 kHz),
plotted, and its harmonic distortion reported. In air the absorption of every
harmonic comes from ISO 9613-1 (temperature and humidity dependent).

Numerical method: operator splitting in z (Strang), in the spirit of
Lee & Hamilton, J. Acoust. Soc. Am. 97, 906 (1995):
  * diffraction + absorption: solved exactly for every harmonic n*f0 in
    (kx, n) space using FFTs in x and tau
  * nonlinearity: exact implicit solution of dp/dz = (beta/(rho0 c0^3)) p dp/dtau,
    i.e. each point of the waveform slides in tau in proportion to its pressure

Usage:
  python kzk_sim.py 0.3            simulate out to z = 0.3 m
  python kzk_sim.py                asks for the distance
  python kzk_sim.py --validate     compare against analytic solutions
  python kzk_sim.py --am 3         60 kHz carrier modulated by 1 kHz, in air, out to z = 3 m
  python kzk_sim.py --am 3 --fc 40e3 --fm 500 --rh 30     other carrier, tone, humidity
  python kzk_sim.py -h             all options
"""

import argparse
import sys
import time
from dataclasses import dataclass
from pathlib import Path

import numpy as np
from scipy import fft as sfft
from scipy.special import jv


# ----------------------------------------------------------------------------
# Physical parameters
# ----------------------------------------------------------------------------

@dataclass
class Medium:
    c0: float = 1500.0       # small-signal sound speed [m/s]
    rho0: float = 1000.0     # ambient density [kg/m^3]
    beta: float = 3.5        # coefficient of nonlinearity (water ~3.5)
    delta: float = 4.33e-6   # diffusivity of sound [m^2/s] (water, thermoviscous)
    iso9613: bool = False    # air: take the absorption from ISO 9613-1 instead of delta
    temp_c: float = 20.0     # air temperature [deg C] (ISO 9613-1 only)
    rh: float = 50.0         # relative humidity [%] (ISO 9613-1 only)

    @classmethod
    def air(cls, temp_c=20.0, rh=50.0):
        return cls(c0=343.0, rho0=1.21, beta=1.2, delta=0.0, iso9613=True, temp_c=temp_c, rh=rh)


@dataclass
class Source:
    f0: float = 1.0e6        # tone frequency [Hz] (the carrier in AM mode)
    p0: float = 1.0e6        # peak source pressure on axis [Pa] (carrier amplitude in AM mode)
    a: float = 0.015         # Gaussian half-width (1/e of amplitude) [m]
    fm: float = 0.0          # modulation frequency [Hz]; 0 means a pure tone
    m: float = 1.0           # modulation index (0..1), AM mode only


def window_frequency(src):
    """Fundamental of the periodic time window: f0 for a pure tone, fm for AM."""
    if not src.fm:
        return src.f0
    ratio = src.f0 / src.fm
    if abs(ratio - round(ratio)) > 1e-9:
        raise ValueError("the carrier must be a whole multiple of the modulation frequency "
                         "so that the AM signal repeats every 1/fm")
    return src.fm


def am_nt(src, per_cycle=64):
    """Samples in the AM time window: per_cycle samples on every carrier cycle.

    The audible tone is generated by the difference between the slopes of the
    carrier and its sidebands; on a grid that difference picks up a factor
    cos(2 pi / per_cycle), i.e. the audio comes out low by about
    (2 pi / per_cycle)^2 / 2: 1.9 % at 32 samples, 0.5 % at 64.
    """
    return int(round(per_cycle * src.f0 / window_frequency(src)))


def source_waveform(src, tau):
    """Source time signal with unit carrier amplitude."""
    carrier = np.sin(2 * np.pi * src.f0 * tau)
    if src.fm:
        return (1 + src.m * np.cos(2 * np.pi * src.fm * tau)) * carrier
    return carrier


def iso9613_alpha(f, temp_c=20.0, rh=50.0, pa_kpa=101.325):
    """Absorption of sound in air [Np/m] at frequency f [Hz] (ISO 9613-1:1993).

    Classical (viscous + thermal) loss plus the vibrational relaxation of
    oxygen and nitrogen, which dominates at ultrasonic frequencies.
    """
    f = np.asarray(f, dtype=float)
    T, T0, T01 = temp_c + 273.15, 293.15, 273.16
    pa = pa_kpa / 101.325                                          # pa / pr
    psat = 10 ** (-6.8346 * (T01 / T) ** 1.261 + 4.6151)           # psat / pr
    h = rh * psat / pa                                             # water vapour [% molar]
    fr_o = pa * (24 + 4.04e4 * h * (0.02 + h) / (0.391 + h))
    fr_n = pa * (T / T0) ** -0.5 * (9 + 280 * h * np.exp(-4.170 * ((T / T0) ** (-1 / 3) - 1)))
    return f ** 2 * (1.84e-11 / pa * (T / T0) ** 0.5
                     + (T / T0) ** -2.5 * (0.01275 * np.exp(-2239.1 / T) / (fr_o + f ** 2 / fr_o)
                                           + 0.1068 * np.exp(-3352.0 / T) / (fr_n + f ** 2 / fr_n)))


def absorption(med, f):
    """Absorption coefficient [Np/m] at frequency f [Hz]."""
    if med.iso9613:
        return iso9613_alpha(f, med.temp_c, med.rh)
    return med.delta * (2 * np.pi * np.asarray(f, dtype=float)) ** 2 / (2 * med.c0 ** 3)


def derived_scales(med, src):
    """Characteristic lengths that tell you which effect dominates where."""
    w0 = 2 * np.pi * src.f0
    k = w0 / med.c0
    z_rayleigh = k * src.a ** 2 / 2                                  # diffraction length
    x_shock = med.rho0 * med.c0 ** 3 / (med.beta * w0 * src.p0)      # plane-wave shock formation
    alpha1 = float(absorption(med, src.f0))                          # absorption at f0 [Np/m]
    gol_dberg = 1 / (alpha1 * x_shock)                               # nonlinearity vs absorption
    return dict(w0=w0, k=k, z_rayleigh=z_rayleigh, x_shock=x_shock,
                alpha1=alpha1, goldberg=gol_dberg)


# ----------------------------------------------------------------------------
# Solver
# ----------------------------------------------------------------------------

def _nonlinear_step(p, h, b, dtau, n_newton=3, cubic=False):
    """Exact solution of dp/dz = b p dp/dtau over a step h (periodic in tau).

    The solution is p(z+h, tau) = p(z, tau + b h p(z+h, tau)); it is implicit,
    so it is solved with Newton's method on an interpolant of p(z, .).

    Linear interpolation slightly smooths the waveform at every step. That is
    negligible with ~1000 samples per period (pure tone), but not in AM mode,
    where each carrier cycle gets only ~32 samples; there cubic (Catmull-Rom)
    interpolation is used, which loses ~1000x less amplitude per step.
    """
    nt = p.shape[-1]
    c = b * h / dtau                     # shift in grid points per Pa
    j = np.arange(nt)
    q = p.copy()
    for _ in range(n_newton):
        pos = j + c * q
        i0 = np.floor(pos)
        f = pos - i0
        i0 = i0.astype(np.intp) % nt
        i1 = (i0 + 1) % nt
        v0 = np.take_along_axis(p, i0, axis=-1)
        v1 = np.take_along_axis(p, i1, axis=-1)
        if cubic:
            vm = np.take_along_axis(p, (i0 - 1) % nt, axis=-1)
            v2 = np.take_along_axis(p, (i0 + 2) % nt, axis=-1)
            c3 = 0.5 * (v2 - vm) + 1.5 * (v0 - v1)
            c2 = vm - 2.5 * v0 + 2 * v1 - 0.5 * v2
            c1 = 0.5 * (v1 - vm)
            value = ((c3 * f + c2) * f + c1) * f + v0
            slope = (3 * c3 * f + 2 * c2) * f + c1
        else:
            value = v0 + f * (v1 - v0)
            slope = v1 - v0
        q = q - (q - value) / (1 - c * slope)
    return q


def simulate(z_end, med=None, src=None, nt=1024, dx=None, n_snapshots=2,
             n_harmonics=6, diffraction=True, nonlinear=True, h_max=None,
             cfl=0.5, verbose=True):
    """March the KZK equation from z = 0 to z = z_end.

    Returns a dict with the on-axis waveforms at n_snapshots distances
    (equally spaced from 0 to z_end; the default 2 gives source and z_end)
    and the on-axis harmonic amplitudes versus z.

    The time window is one period of window_frequency(src): 1/f0 for a pure
    tone, 1/fm for AM. nt is the number of samples in that window, and
    "harmonic n" below means n times the window frequency (n*fm in AM mode,
    so harmonics 1, 2, 3 are the audible tone and its distortion).
    """
    med = med or Medium()
    src = src or Source()
    s = derived_scales(med, src)
    f_win = window_frequency(src)
    w_win = 2 * np.pi * f_win
    k_win = w_win / med.c0
    T = 1 / f_win
    dtau = T / nt
    tau = np.arange(nt) * dtau
    b = med.beta / (med.rho0 * med.c0 ** 3)

    # transverse grid (periodic for the FFT, with an absorbing sponge at the edges)
    if diffraction:
        w_end = src.a * np.sqrt(1 + (z_end / s["z_rayleigh"]) ** 2)
        # In AM mode the audio spreads far faster than the carrier (its Rayleigh
        # distance k_m a^2/2 is centimetres), so the box must also hold the audio
        # beam, or it reflects off the edges and corrupts the result after a metre
        # or two. 0.75 of the audio's paraxial width was converged in tests (a wider
        # box changed nothing); for a pure tone this term never matters.
        w_end_win = src.a * np.sqrt(1 + (z_end / (k_win * src.a ** 2 / 2)) ** 2)
        x_useful = max(4 * w_end, 6 * src.a, 0.75 * w_end_win)
        x_sponge = 0.35 * x_useful
        x_half = x_useful + x_sponge
        # a/12 gives the same AM results as a/24 (tested) at half the cost
        dx = dx or src.a / (12 if src.fm else 24)
        nx = sfft.next_fast_len(int(np.ceil(2 * x_half / dx)))
        nx += nx % 2
        x = (np.arange(nx) - nx // 2) * dx
        depth = np.clip((np.abs(x) - x_useful) / x_sponge, 0, None)
        # strong enough for the most diffracting component (the audio in AM mode)
        sponge_rate = (40 / (k_win * src.a ** 2 / 2)) * depth ** 2   # [Np/m]
    else:
        x_useful = 0.0
        nx = 1
        x = np.zeros(1)
        sponge_rate = np.zeros(1)
    ix0 = nx // 2                                               # x = 0 (axis)

    # linear operator per harmonic: diffraction phase and absorption
    n = np.arange(nt // 2 + 1)
    absorb = absorption(med, n * f_win)                         # [Np/m] for each harmonic
    kx = 2 * np.pi * sfft.fftfreq(nx, dx if diffraction else 1.0)
    diff = np.zeros((nx, n.size))
    if diffraction:
        diff[:, 1:] = kx[:, None] ** 2 / (2 * k_win * n[None, 1:])  # [rad/m]

    def linear_step(p, h):
        P = sfft.rfft(p, axis=1, workers=-1)
        if diffraction:
            P = sfft.fft(P, axis=0, workers=-1)
        P *= np.exp((1j * diff - absorb[None, :]) * h)
        P[:, 0] = 0.0                                           # no DC in KZK
        if diffraction:
            P = sfft.ifft(P, axis=0, workers=-1)
        p = sfft.irfft(P, n=nt, axis=1, workers=-1)
        if diffraction:
            p *= np.exp(-sponge_rate * h)[:, None]
        return p

    def nonlinear_step(p, h):
        # only rows that carry signal need the (relatively costly) update
        rows = np.flatnonzero(np.abs(p).max(axis=1) > 1e-6 * src.p0)
        if rows.size:
            p[rows] = _nonlinear_step(p[rows], h, b, dtau, cubic=bool(src.fm))
        return p

    # amplitude of harmonic n (n*f_win) is 2|sum p e^{-i n w_win tau}|/nt, i.e. an FFT bin
    # (an FFT rather than a matrix product: numpy 2.0 + Apple Accelerate prints false
    # "divide by zero in matmul" warnings for the latter)
    def harmonics(p_rows):
        return 2 * np.abs(np.fft.rfft(p_rows, axis=-1)[..., 1:n_harmonics + 1]) / nt

    # source
    p = src.p0 * np.exp(-(x / src.a) ** 2)[:, None] * source_waveform(src, tau)[None, :]
    if not diffraction:
        p = source_waveform(src, tau)[None, :] * src.p0

    h_max = h_max or min(s["x_shock"], s["z_rayleigh"] if diffraction else np.inf) / 100
    z_snap = np.linspace(0, z_end, n_snapshots)
    snaps = [p[ix0].copy()]
    next_snap = 1

    z_hist = [0.0]
    harm_hist = [harmonics(p[ix0:ix0 + 1])[0]]

    t_start = time.time()
    z = 0.0
    n_steps = 0
    while z < z_end * (1 - 1e-12):
        h = h_max
        if nonlinear:
            slope = np.abs(np.diff(p, axis=1, append=p[:, :1])).max() / dtau
            if slope > 0:
                h = min(h, cfl / (b * slope))
        h = min(h, z_snap[next_snap] - z)

        if nonlinear:
            p = nonlinear_step(p, h / 2)
        p = linear_step(p, h)
        if nonlinear:
            p = nonlinear_step(p, h / 2)
        z += h
        n_steps += 1

        z_hist.append(z)
        harm_hist.append(harmonics(p[ix0:ix0 + 1])[0])

        if abs(z - z_snap[next_snap]) <= 1e-12 * max(z_end, 1):
            z = z_snap[next_snap]
            snaps.append(p[ix0].copy())
            next_snap += 1
            if verbose:
                print(f"  z = {z:7.4f} m  ({n_steps} steps, {time.time() - t_start:5.1f} s)")
            if next_snap == n_snapshots:
                break

    return dict(
        med=med, src=src, scales=s, nt=nt, tau=tau, T=T, f_win=f_win, x=x,
        z_end=z_end, z_snap=z_snap, snaps=np.array(snaps),
        z_hist=np.array(z_hist), harm_hist=np.array(harm_hist),
        n_steps=n_steps, runtime=time.time() - t_start,
    )


# ----------------------------------------------------------------------------
# Validation against analytic solutions
# ----------------------------------------------------------------------------

def validate():
    """Check each physical effect against a known analytic solution."""
    ok = True
    med = Medium()
    print("Validation")
    print("-" * 66)

    # 1) linear Gaussian beam (diffraction + absorption, no nonlinearity)
    src = Source()
    s = derived_scales(med, src)
    r = simulate(2 * s["z_rayleigh"], med, src, nt=64, n_snapshots=2,
                 nonlinear=False, verbose=False)
    zz, a1 = r["z_hist"], r["harm_hist"][:, 0]
    exact = src.p0 * (1 + (zz / s["z_rayleigh"]) ** 2) ** -0.25 * np.exp(-s["alpha1"] * zz)
    err = np.abs(a1 - exact).max() / src.p0
    ok &= err < 1e-3
    print(f"1) Linear Gaussian beam, 0 <= z <= 2 zR, on-axis amplitude")
    print(f"   max error vs exact (1+(z/zR)^2)^(-1/4) e^(-alpha z): {err:.1e} p0")

    # 2) plane wave, pre-shock: Fubini solution B_n = 2 J_n(n s)/(n s)
    src = Source()
    s = derived_scales(med, src)
    weak = Medium(delta=med.delta / 20)          # Gamma ~ 5000: effectively lossless
    sigma = 0.8
    r = simulate(sigma * s["x_shock"], weak, src, nt=2048, n_snapshots=2,
                 diffraction=False, verbose=False)
    B = r["harm_hist"][-1] / src.p0
    nn = np.arange(1, B.size + 1)
    fub = 2 * jv(nn, nn * sigma) / (nn * sigma)
    err = np.abs(B - fub).max()
    ok &= err < 2e-3
    print(f"2) Plane wave, sigma = z/x_shock = {sigma} (Fubini solution)")
    print("   n      sim     exact")
    for i in range(4):
        print(f"   {i + 1}   {B[i]:.4f}   {fub[i]:.4f}")
    print(f"   max error: {err:.1e} p0")

    # 3) plane wave, old-age sawtooth: Fay solution B_n = (2/G) / sinh(n (1+s)/G)
    lossy = Medium(delta=med.delta * 5)
    src = Source()
    s = derived_scales(lossy, src)
    G = s["goldberg"]
    sigma = 4.0
    r = simulate(sigma * s["x_shock"], lossy, src, nt=1024, n_snapshots=2,
                 diffraction=False, verbose=False)
    B = r["harm_hist"][-1] / src.p0
    fay = (2 / G) / np.sinh(nn * (1 + sigma) / G)
    err = np.abs(B - fay).max()
    ok &= err < 1e-2
    print(f"3) Plane wave, sigma = {sigma}, Goldberg number {G:.0f} (Fay solution)")
    print("   n      sim     exact")
    for i in range(4):
        print(f"   {i + 1}   {B[i]:.4f}   {fay[i]:.4f}")
    print(f"   max error: {err:.1e} p0")

    # 4) AM carrier in air, plane wave, weak nonlinearity: second-order (quasilinear)
    #    theory. Each pair of lines (carrier & sideband, sideband & sideband) beats
    #    and radiates the difference tone, so on the axis the audio at fm and 2fm is
    #    A1 = (b/4) m wm p0^2 [I(a_c + a_+) + I(a_c + a_-)],  A2 = (b/4) m^2 wm p0^2 I(a_+ + a_-)
    #    with b = beta/(rho0 c0^3) and I(g) = (1 - e^{-g z})/g (carrier absorption).
    air = Medium.air()
    src = Source(f0=60e3, fm=1e3, m=0.5, p0=100.0, a=0.1)
    z = 0.1
    r = simulate(z, air, src, nt=am_nt(src), n_snapshots=2, diffraction=False, verbose=False)
    A = r["harm_hist"][-1]
    b = air.beta / (air.rho0 * air.c0 ** 3)
    wm = 2 * np.pi * src.fm
    a_c, a_p, a_m = (float(absorption(air, f)) for f in (src.f0, src.f0 + src.fm, src.f0 - src.fm))
    I = lambda g: (1 - np.exp(-g * z)) / g
    A1 = b / 4 * src.m * wm * src.p0 ** 2 * (I(a_c + a_p) + I(a_c + a_m))
    A2 = b / 4 * src.m ** 2 * wm * src.p0 ** 2 * I(a_p + a_m)
    err = max(abs(A[0] / A1 - 1), abs(A[1] / A2 - 1))
    ok &= err < 1e-2
    print(f"4) AM in air (fc = 60 kHz, fm = 1 kHz, m = {src.m}), plane wave, z = {z} m (quasilinear theory)")
    print("   tone      sim [Pa]   exact [Pa]")
    print(f"   1 kHz     {A[0]:.5f}    {A1:.5f}")
    print(f"   2 kHz     {A[1]:.5f}    {A2:.5f}")
    print(f"   max relative error: {err:.1e}")

    print("-" * 66)
    print("ALL CHECKS PASSED" if ok else "SOME CHECKS FAILED")
    return ok


# ----------------------------------------------------------------------------
# Plotting
# ----------------------------------------------------------------------------

INK = "#0b0b0b"
INK2 = "#52514e"
MUTED = "#898781"
GRID = "#e1e0d9"
AXIS = "#c3c2b7"
SURFACE = "#fcfcfb"
BLUE = "#2a78d6"
RED = "#b0182f"          # the audible part in AM mode


def _style(ax, title):
    ax.set_facecolor(SURFACE)
    ax.set_title(title, loc="left", color=INK, fontsize=11, pad=8)
    for side in ("top", "right"):
        ax.spines[side].set_visible(False)
    for side in ("left", "bottom"):
        ax.spines[side].set_color(AXIS)
    ax.tick_params(colors=INK2, labelsize=9)
    ax.xaxis.label.set_color(INK2)
    ax.yaxis.label.set_color(INK2)
    ax.grid(color=GRID, linewidth=0.6)
    ax.set_axisbelow(True)


def make_figure(r, outdir, show=False, n_bars=10):
    """Time-domain waveform and frequency spectrum on axis at distance z."""
    import matplotlib
    if not show:
        matplotlib.use("Agg")
    import matplotlib.pyplot as plt

    outdir = Path(outdir)
    outdir.mkdir(parents=True, exist_ok=True)

    src, s = r["src"], r["scales"]
    z_end, nt, T = r["z_end"], r["nt"], r["T"]
    MPa = 1e-6
    source_wave, wave = r["snaps"][0], r["snaps"][-1]
    t_us = np.arange(2 * nt + 1) * (T / nt) * 1e6                 # two periods
    tile = lambda w: np.concatenate([w, w, w[:1]])

    fig, (ax_t, ax_f) = plt.subplots(1, 2, figsize=(12.5, 5), facecolor=SURFACE)
    fig.suptitle(
        f"KZK 2D, pure tone f0 = {src.f0 / 1e6:g} MHz, p0 = {src.p0 * MPa:g} MPa, "
        f"a = {src.a * 100:g} cm, water  |  distance z = {z_end:g} m  "
        f"(z / x_shock = {z_end / s['x_shock']:.2f})",
        color=INK, fontsize=12, x=0.01, ha="left")

    # time domain
    _style(ax_t, f"Time: pressure on the beam axis at z = {z_end:g} m")
    if z_end > 0:
        ax_t.plot(t_us, tile(source_wave) * MPa, color=MUTED, lw=1.4, ls="--",
                  label="source sine (z = 0)")
        ax_t.plot(t_us, tile(wave) * MPa, color=BLUE, lw=2, label=f"after z = {z_end:g} m")
    else:
        ax_t.plot(t_us, tile(wave) * MPa, color=BLUE, lw=2, label="source sine (z = 0)")
    amp = max(np.abs(source_wave).max(), np.abs(wave).max()) * MPa
    ax_t.set_xlim(0, t_us[-1])
    ax_t.set_ylim(-1.1 * amp, 1.45 * amp)
    ax_t.set_xlabel("time τ [µs]   (retarded time τ = t − z/c0)")
    ax_t.set_ylabel("pressure p [MPa]")
    ax_t.legend(frameon=False, fontsize=9, labelcolor=INK2, ncol=2, loc="upper center")

    # frequency domain
    _style(ax_f, f"Frequency: amplitude spectrum at z = {z_end:g} m")
    spec = 2 * np.abs(np.fft.rfft(wave))[1:n_bars + 1] / nt
    f_mhz = np.arange(1, n_bars + 1) * src.f0 / 1e6
    ax_f.bar(f_mhz, spec * MPa, width=0.6 * src.f0 / 1e6, color=BLUE, zorder=3)
    ax_f.axhline(src.p0 * MPa, color=MUTED, lw=1.4, ls="--")
    ax_f.text(f_mhz[-1] + 0.4 * src.f0 / 1e6, src.p0 * MPa, "source: p0 at f0 only",
              color=MUTED, fontsize=9, ha="right", va="bottom")
    for i in range(3):
        ax_f.text(f_mhz[i], (spec[i] + 0.012 * src.p0) * MPa, f"{spec[i] * MPa:.3f}", color=INK2,
                  fontsize=9, ha="center", va="bottom", zorder=4,
                  bbox=dict(facecolor=SURFACE, edgecolor="none", pad=1))
    ax_f.set_xticks(f_mhz)
    ax_f.set_xlabel("frequency [MHz]   (harmonics n·f0)")
    ax_f.set_ylabel("amplitude [MPa]")
    ax_f.set_ylim(0, src.p0 * MPa * 1.12)
    ax_f.grid(axis="x", visible=False)

    fig.tight_layout(rect=(0, 0, 1, 0.94))
    path = outdir / f"kzk_z{z_end:g}m.png"
    fig.savefig(path, dpi=130, facecolor=SURFACE)
    if show:
        plt.show()
    plt.close(fig)
    return path


# ----------------------------------------------------------------------------
# AM mode: the audible part of the beam
# ----------------------------------------------------------------------------

AUDIBLE_MAX = 20e3       # upper edge of the audible band [Hz]
P_REF = 20e-6            # reference pressure for dB SPL [Pa]


def spl(amplitude):
    """dB SPL of a sinusoid with the given peak amplitude [Pa]."""
    return 20 * np.log10(np.maximum(amplitude, 1e-12) / np.sqrt(2) / P_REF)


def audible(wave, f_win, f_max=AUDIBLE_MAX):
    """What a listener hears: the waveform low-passed to the audible band.

    Returns the audible waveform and the amplitudes [Pa] of its lines at
    f_win, 2 f_win, ... up to f_max. In AM mode f_win = fm, so these are the
    tone and its harmonic distortion.
    """
    nt = wave.size
    W = np.fft.rfft(wave)
    n_max = int(f_max // f_win)
    W[n_max + 1:] = 0
    W[0] = 0
    return np.fft.irfft(W, n=nt), 2 * np.abs(W[1:n_max + 1]) / nt


def thd(amps):
    """Total harmonic distortion sqrt(A2^2 + A3^2 + ...) / A1 (amps along the last axis)."""
    amps = np.asarray(amps)
    return np.sqrt(np.sum(amps[..., 1:] ** 2, axis=-1)) / np.maximum(amps[..., 0], 1e-300)


def make_am_figure(r, outdir, show=False, n_audio_bars=8):
    """AM mode: the on-axis signal and its spectrum, then the audible part alone.

    Saves two figures: the signal at z (4 panels) and the audible level and
    distortion against distance.
    """
    import matplotlib
    if not show:
        matplotlib.use("Agg")
    import matplotlib.pyplot as plt

    outdir = Path(outdir)
    outdir.mkdir(parents=True, exist_ok=True)

    src, med = r["src"], r["med"]
    z_end, nt, f_win = r["z_end"], r["nt"], r["f_win"]
    wave = r["snaps"][-1]
    audio, amps = audible(wave, f_win)
    silent = amps[0] < 1e-9 * src.p0                          # e.g. --linear: no demodulation
    distortion = thd(amps)
    t_ms = np.arange(nt + 1) / (nt * f_win) * 1e3             # one modulation period
    close = lambda w: np.append(w, w[:1])

    header = (f"KZK 2D, AM in air: {src.f0 / 1e3:g} kHz carrier modulated by {src.fm / 1e3:g} kHz "
              f"(m = {src.m:g}), p0 = {src.p0:g} Pa, a = {src.a * 100:g} cm, "
              f"{med.temp_c:g} °C, {med.rh:g} % RH  |  distance z = {z_end:g} m")

    fig, ((ax_t, ax_f), (ax_a, ax_s)) = plt.subplots(2, 2, figsize=(13, 9.5), facecolor=SURFACE)
    fig.suptitle(header, color=INK, fontsize=12, x=0.01, ha="left")

    # full signal on axis
    _style(ax_t, f"Time: pressure on the beam axis at z = {z_end:g} m")
    envelope = src.p0 * (1 + src.m * np.cos(2 * np.pi * src.fm * t_ms / 1e3))
    ax_t.plot(t_ms, envelope, color=MUTED, lw=1.2, ls="--", label="source envelope (z = 0)")
    ax_t.plot(t_ms, -envelope, color=MUTED, lw=1.2, ls="--")
    ax_t.plot(t_ms, close(wave), color=BLUE, lw=0.5, label=f"after z = {z_end:g} m")
    ax_t.set_xlim(0, t_ms[-1])
    ax_t.set_xlabel("time τ [ms]   (one period of the modulation)")
    ax_t.set_ylabel("pressure p [Pa]")
    ax_t.legend(frameon=False, fontsize=9, labelcolor=INK2, ncol=2, loc="upper center")
    ax_t.set_ylim(-1.15 * envelope.max(), 1.4 * envelope.max())

    # full spectrum in dB SPL: audible lines in red, ultrasonic in blue
    _style(ax_f, f"Frequency: spectrum at z = {z_end:g} m (audible band shaded)")
    W = 2 * np.abs(np.fft.rfft(wave)) / nt
    n_show = min(W.size - 1, int(3.2 * src.f0 / f_win))
    f_khz = np.arange(1, n_show + 1) * f_win / 1e3
    lev = spl(W[1:n_show + 1])
    floor = max(0.0, lev.max() - 110)
    loud = lev > floor
    is_audio = f_khz * 1e3 <= AUDIBLE_MAX
    ax_f.axvspan(0, AUDIBLE_MAX / 1e3, color=RED, alpha=0.07, lw=0)
    ax_f.vlines(f_khz[loud & ~is_audio], floor, lev[loud & ~is_audio], color=BLUE, lw=1.2)
    ax_f.vlines(f_khz[loud & is_audio], floor, lev[loud & is_audio], color=RED, lw=1.8)
    i_c = int(round(src.f0 / f_win)) - 1
    ax_f.annotate(f"carrier {src.f0 / 1e3:g} kHz: {lev[i_c]:.0f} dB", (f_khz[i_c], lev[i_c]),
                  xytext=(8, 0), textcoords="offset points", color=INK2, fontsize=9, va="center")
    if not silent:
        ax_f.annotate(f"audible {src.fm / 1e3:g} kHz: {spl(amps[0]):.0f} dB", (f_khz[0], spl(amps[0])),
                      xytext=(8, 4), textcoords="offset points", color=RED, fontsize=9)
    ax_f.set_xlim(0, f_khz[-1])
    ax_f.set_ylim(floor, lev.max() + 12)
    ax_f.set_xlabel("frequency [kHz]")
    ax_f.set_ylabel("level [dB SPL]")

    if silent:
        for ax in (ax_a, ax_s):
            _style(ax, "Audible sound")
            ax.text(0.5, 0.5, "No audible sound: without nonlinearity\nthe air does not demodulate the beam",
                    transform=ax.transAxes, ha="center", va="center", color=INK2, fontsize=11)
            ax.set_xticks([])
            ax.set_yticks([])
        fig.tight_layout(rect=(0, 0, 1, 0.95))
        path = outdir / f"kzk_am_z{z_end:g}m_linear.png"
        fig.savefig(path, dpi=130, facecolor=SURFACE)
        if show:
            plt.show()
        plt.close(fig)
        return (path,)

    # audible waveform against an undistorted tone
    _style(ax_a, f"Audible sound (below {AUDIBLE_MAX / 1e3:g} kHz): THD = {100 * distortion:.1f} %")
    two = lambda w: np.concatenate([w, w, w[:1]])
    t2 = np.arange(2 * nt + 1) / (nt * f_win) * 1e3
    F = np.fft.rfft(audio)
    F[2:] = 0
    tone = np.fft.irfft(F, n=nt)
    ax_a.plot(t2, two(tone) * 1e3, color=MUTED, lw=1.4, ls="--",
              label=f"its {src.fm / 1e3:g} kHz part alone (undistorted)")
    ax_a.plot(t2, two(audio) * 1e3, color=RED, lw=2, label="what a listener on axis hears")
    peak = np.abs(audio).max() * 1e3
    ax_a.set_xlim(0, t2[-1])
    ax_a.set_ylim(-1.2 * peak, 1.55 * peak)
    ax_a.set_xlabel("time τ [ms]")
    ax_a.set_ylabel("pressure [mPa]")
    ax_a.legend(frameon=False, fontsize=9, labelcolor=INK2, ncol=1, loc="upper right")

    # audible spectrum: the tone and its distortion products
    _style(ax_s, "Audible spectrum: the tone and its distortion")
    nb = min(n_audio_bars, amps.size)
    fa = np.arange(1, nb + 1) * f_win / 1e3
    la = spl(amps[:nb])
    base = min(la.min(), la[0] - 60)
    ax_s.bar(fa, la - base, bottom=base, width=0.6 * f_win / 1e3, color=RED, zorder=3)
    for i in range(min(3, nb)):
        ax_s.text(fa[i], la[i] + 1, f"{la[i]:.1f} dB", color=INK2, fontsize=9, ha="center", va="bottom")
    ax_s.set_xticks(fa)
    ax_s.set_ylim(base, la[0] + 10)
    ax_s.set_xlabel(f"frequency [kHz]   (multiples of {src.fm / 1e3:g} kHz)")
    ax_s.set_ylabel("level [dB SPL]")
    ax_s.grid(axis="x", visible=False)

    fig.tight_layout(rect=(0, 0, 1, 0.95))
    path = outdir / f"kzk_am_z{z_end:g}m.png"
    fig.savefig(path, dpi=130, facecolor=SURFACE)

    # audible level and distortion against distance
    zz, H = r["z_hist"][1:], r["harm_hist"][1:]
    fig2, (ax_l, ax_d) = plt.subplots(1, 2, figsize=(13, 4.6), facecolor=SURFACE)
    fig2.suptitle(header.split("|")[0].strip(), color=INK, fontsize=12, x=0.01, ha="left")
    _style(ax_l, "Audible level on the beam axis against distance")
    for i, (style, color) in enumerate(((("-", 2.2), RED), (("--", 1.5), INK2), ((":", 1.5), MUTED))):
        if i < H.shape[1]:
            ax_l.plot(zz, spl(H[:, i]), color=color, ls=style[0], lw=style[1],
                      label=f"{(i + 1) * f_win / 1e3:g} kHz")
    ax_l.set_xlabel("distance z [m]")
    ax_l.set_ylabel("level [dB SPL]")
    ax_l.set_xlim(0, z_end)
    ax_l.set_ylim(spl(H[:, 0]).max() - 60, spl(H[:, 0]).max() + 8)
    ax_l.legend(frameon=False, fontsize=9, labelcolor=INK2)
    _style(ax_d, f"Distortion of the {src.fm / 1e3:g} kHz tone against distance "
                 f"(harmonics up to {H.shape[1] * f_win / 1e3:g} kHz)")
    ax_d.plot(zz, 100 * thd(H), color=RED, lw=2)
    ax_d.set_xlabel("distance z [m]")
    ax_d.set_ylabel("THD [%]")
    ax_d.set_xlim(0, z_end)
    ax_d.set_ylim(0, None)
    fig2.tight_layout(rect=(0, 0, 1, 0.92))
    path2 = outdir / f"kzk_am_vs_z_{z_end:g}m.png"
    fig2.savefig(path2, dpi=130, facecolor=SURFACE)

    if show:
        plt.show()
    plt.close(fig)
    plt.close(fig2)
    return path, path2


# ----------------------------------------------------------------------------
# Command line
# ----------------------------------------------------------------------------

def ask_distance():
    while True:
        try:
            z = float(input("Enter propagation distance z in metres (z >= 0): "))
        except ValueError:
            print("  please type a number, e.g. 0.3")
            continue
        except EOFError:
            sys.exit("no distance given")
        if z >= 0:
            return z
        print("  the distance cannot be negative")


def main():
    ap = argparse.ArgumentParser(description="Simple 2D KZK simulation of a pure-tone beam in water, "
                                             "or (--am) of an AM parametric speaker in air.")
    ap.add_argument("z", nargs="?", type=float, help="propagation distance in metres (>= 0)")
    ap.add_argument("--f0", type=float, default=1.0e6, help="pure-tone frequency [Hz] (default 1e6)")
    ap.add_argument("--p0", type=float, default=None,
                    help="source pressure [Pa] (default 1e6; AM: carrier amplitude, default 100)")
    ap.add_argument("--a", type=float, default=None,
                    help="source half-width [m] (default 0.015; AM: default 0.1)")
    ap.add_argument("--nt", type=int, default=None,
                    help="time samples per period (default 1024; AM: 64 per carrier cycle)")
    am = ap.add_argument_group("AM mode (parametric speaker in air)")
    am.add_argument("--am", action="store_true", help="AM carrier in air instead of a pure tone in water")
    am.add_argument("--fc", type=float, default=60e3, help="carrier frequency [Hz] (default 60e3)")
    am.add_argument("--fm", type=float, default=1e3, help="modulating (audible) tone [Hz] (default 1e3)")
    am.add_argument("--m", type=float, default=1.0, help="modulation index, 0 < m <= 1 (default 1)")
    am.add_argument("--temp", type=float, default=20.0, help="air temperature [deg C] (default 20)")
    am.add_argument("--rh", type=float, default=50.0, help="relative humidity [%%] (default 50)")
    ap.add_argument("--linear", action="store_true", help="switch nonlinearity off (for comparison)")
    ap.add_argument("--out", default=None,
                    help="folder for the figures (default: 'output' next to this script)")
    ap.add_argument("--no-show", action="store_true", help="save figures without opening a window")
    ap.add_argument("--validate", action="store_true", help="run the analytic checks and exit")
    args = ap.parse_args()

    if args.validate:
        sys.exit(0 if validate() else 1)

    z = args.z if args.z is not None else ask_distance()
    if z < 0:
        sys.exit("the distance cannot be negative")

    # save next to this script, not in whatever folder Python was started from
    outdir = Path(args.out) if args.out else Path(__file__).resolve().parent / "output"

    if args.am:
        run_am(z, args, outdir)
        return

    med = Medium()
    src = Source(f0=args.f0, p0=args.p0 or 1.0e6, a=args.a or 0.015)
    s = derived_scales(med, src)
    print(f"Pure tone f0 = {src.f0 / 1e6:g} MHz, p0 = {src.p0 / 1e6:g} MPa, a = {src.a * 100:g} cm in water")
    print(f"  shock formation distance x_shock = {s['x_shock']:.3f} m")
    print(f"  Rayleigh (diffraction) distance zR = {s['z_rayleigh']:.3f} m")
    print(f"  absorption at f0 = {s['alpha1']:.4f} Np/m, Goldberg number = {s['goldberg']:.0f}")
    print(f"Marching to z = {z:g} m ...")

    r = simulate(z, med, src, nt=args.nt or 1024, nonlinear=not args.linear)
    w = r["snaps"][-1]
    amps = r["harm_hist"][-1]
    print(f"Done: {r['n_steps']} steps in {r['runtime']:.1f} s")
    print(f"On axis at z = {z:g} m:")
    print(f"  peak positive pressure  {w.max() / 1e6:+.3f} MPa")
    print(f"  peak negative pressure  {w.min() / 1e6:+.3f} MPa")
    print("  harmonic amplitudes     " + ", ".join(
        f"{i + 1}f0: {amps[i] / 1e6:.3f}" for i in range(4)) + " MPa")

    path = make_figure(r, outdir, show=not args.no_show)
    print(f"Saved {path}")


def run_am(z, args, outdir):
    """Command-line AM mode: simulate, print the audible result, save the figures."""
    if not 0 < args.m <= 1:
        sys.exit("the modulation index must be in (0, 1]")
    if args.fm >= args.fc:
        sys.exit("the carrier must be higher than the modulating tone")
    med = Medium.air(args.temp, args.rh)
    src = Source(f0=args.fc, p0=args.p0 or 100.0, a=args.a or 0.1, fm=args.fm, m=args.m)
    try:
        nt = args.nt or am_nt(src)
    except ValueError as e:
        sys.exit(str(e))
    s = derived_scales(med, src)
    print(f"AM in air: carrier {src.f0 / 1e3:g} kHz modulated by {src.fm / 1e3:g} kHz (m = {src.m:g}), "
          f"p0 = {src.p0:g} Pa, a = {src.a * 100:g} cm, {med.temp_c:g} deg C, {med.rh:g} % RH")
    print(f"  shock formation distance x_shock = {s['x_shock']:.3f} m (for the carrier amplitude p0)")
    print(f"  Rayleigh (diffraction) distance zR = {s['z_rayleigh']:.3f} m (carrier)")
    print(f"  absorption at the carrier = {s['alpha1']:.4f} Np/m = {8.686 * s['alpha1']:.2f} dB/m (ISO 9613-1)")
    print(f"  time window 1/fm = {1e3 / src.fm:g} ms with nt = {nt} samples")
    print(f"Marching to z = {z:g} m ...")

    r = simulate(z, med, src, nt=nt, nonlinear=not args.linear)
    wave = r["snaps"][-1]
    _, amps = audible(wave, r["f_win"])
    carrier = 2 * np.abs(np.fft.rfft(wave)[int(round(src.f0 / r["f_win"]))]) / nt
    print(f"Done: {r['n_steps']} steps in {r['runtime']:.1f} s")
    print(f"On axis at z = {z:g} m:")
    print(f"  carrier {src.f0 / 1e3:g} kHz        {carrier:8.3f} Pa  ({spl(carrier):.1f} dB SPL)")
    if amps[0] < 1e-9 * src.p0:
        print("  no audible sound: without nonlinearity the air does not demodulate the beam")
    else:
        for i in range(3):
            print(f"  audible {(i + 1) * src.fm / 1e3:g} kHz        {amps[i] * 1e3:8.3f} mPa "
                  f"({spl(amps[i]):.1f} dB SPL)")
        print(f"  THD of the audible tone: {100 * thd(amps):.1f} %")

    paths = make_am_figure(r, outdir, show=not args.no_show)
    for path in paths:
        print(f"Saved {path}")


if __name__ == "__main__":
    main()