# Worked examples for "Underwater acoustics for machine-learning engineers: pressure, levels and
# spectra, with Python" (Sunil Tyagi, 27 October 2026): https://suniltyagi.in/notes/underwater-acoustics-for-ml-engineers/
# Run the cells in order in one Python session (Python 3 with NumPy and SciPy).

# %% Setup
import numpy as np
from scipy.signal import welch

fs = 16_000                 # sampling frequency, Hz (as in the synthetic dataset)
P_REF = 1e-6                # reference pressure in water: 1 µPa, in Pa
rng = np.random.default_rng(seed=1)

# %% Example 1: blade-rate frequency
def blade_rate(blades, shaft_rpm):
    """Blade-rate frequency in Hz: f_BR = B * n_shaft / 60."""
    return blades * shaft_rpm / 60

for vessel, B, rpm in [("tanker", 3, 75), ("cargo ship", 4, 120),
                       ("fishing vessel", 5, 360), ("small craft", 3, 1200)]:
    f_br = blade_rate(B, rpm)
    print(f"{vessel:14s} B = {B}, {rpm:4d} rpm -> f_BR = {f_br:5.2f} Hz "
          f"(harmonics {2 * f_br:.2f}, {3 * f_br:.2f} Hz)")

# %% Example 2: SPL from pressure samples
T = 5.0                     # clip length, s
N = int(fs * T)             # 80,000 samples
t = np.arange(N) / fs

# A 120 Hz machinery tonal (amplitude 0.2 Pa) in white noise (RMS 0.05 Pa)
p = 0.2 * np.sin(2 * np.pi * 120 * t) + 0.05 * rng.standard_normal(N)

p_rms = np.sqrt(np.mean(p**2))
spl = 20 * np.log10(p_rms / P_REF)
print(f"p_rms = {p_rms:.4f} Pa")
print(f"SPL   = {spl:.1f} dB re 1 µPa")

# %% Example 3: Welch PSD with the right scaling
f, S = welch(p, fs=fs, window="hann", nperseg=fs, noverlap=fs // 2,
             scaling="density")            # one-sided PSD in Pa²/Hz, 1 Hz bins
df = f[1] - f[0]
print(f"Frequency resolution {df:.1f} Hz, {f.size} bins up to {f[-1]:.0f} Hz")
print(f"Mean-square pressure from the samples: {np.mean(p**2):.5f} Pa²")
print(f"Integral of the PSD (sum × Δf):        {np.sum(S) * df:.5f} Pa²")

level = 10 * np.log10(S / P_REF**2)        # dB re 1 µPa²/Hz
k = np.argmax(S)
print(f"Peak at {f[k]:.0f} Hz: {level[k]:.1f} dB re 1 µPa²/Hz")
print(f"Median level (the broadband floor): {np.median(level):.1f} dB re 1 µPa²/Hz")

# %% Example 4: band level from the PSD
def band_level(f, S, f1, f2):
    """Band level in dB re 1 µPa from a one-sided PSD in Pa²/Hz, over f1 <= f < f2."""
    m = (f >= f1) & (f < f2)
    ms = np.sum(S[m]) * (f[1] - f[0])      # mean-square pressure in the band, Pa²
    return 10 * np.log10(ms / P_REF**2)

L_broad = band_level(f, S, 1000, 2000)     # 1 kHz of broadband noise
floor = np.median(level[(f >= 1000) & (f < 2000)])
print(f"Band 1000-2000 Hz:                {L_broad:.1f} dB re 1 µPa")
print(f"PSD level + 10·log10(1000 Hz):    {floor + 10 * np.log10(1000):.1f} dB re 1 µPa")
L_tone = band_level(f, S, 110, 130)        # 20 Hz band around the tonal
print(f"Band 110-130 Hz (holds the tonal): {L_tone:.1f} dB re 1 µPa")
print(f"Tonal alone (0.2/√2 Pa RMS):       {20 * np.log10(0.2 / np.sqrt(2) / P_REF):.1f} dB re 1 µPa")

# %% Example 5: noise level (dB re 1 µPa²/Hz) <-> PSD (Pa²/Hz)
def nl_to_psd(nl_db):
    """Spectral density level in dB re 1 µPa²/Hz -> PSD in Pa²/Hz."""
    return P_REF**2 * 10 ** (np.asarray(nl_db) / 10)

def psd_to_nl(S_pa2_per_hz):
    """PSD in Pa²/Hz -> spectral density level in dB re 1 µPa²/Hz."""
    return 10 * np.log10(np.asarray(S_pa2_per_hz) / P_REF**2)

print(f"NL = 60 dB re 1 µPa²/Hz  ->  S = {nl_to_psd(60):.1e} Pa²/Hz")
print(f"S = 3.1e-07 Pa²/Hz       ->  NL = {psd_to_nl(3.1e-7):.1f} dB re 1 µPa²/Hz")
f0 = 1500
print(f"At {f0} Hz: NL = {level[f == f0][0]:.1f} dB re 1 µPa²/Hz; "
      f"level of the 1 Hz band = {band_level(f, S, f0 - 0.5, f0 + 0.5):.1f} dB re 1 µPa")

# %% Example 6: spherical-spreading transmission loss
SL = 150.0                                  # illustrative source level, dB re 1 µPa at 1 m
for r in [1, 10, 100, 1_000, 10_000]:
    TL = 20 * np.log10(r / 1.0)             # r0 = 1 m
    RL = SL - TL
    print(f"r = {r:6d} m: TL = {TL:5.1f} dB, RL = {RL:5.1f} dB re 1 µPa, "
          f"p_rms = {P_REF * 10 ** (RL / 20):.2e} Pa")

# %% Example 7: what squaring does to a tone
x = np.sin(2 * np.pi * 120 * t)             # a pure 120 Hz tone
freqs = np.fft.rfftfreq(N, 1 / fs)
for name, sig in [("p(t)", x), ("p²(t)", x**2)]:
    spec = np.abs(np.fft.rfft(sig)) / N
    lines = freqs[spec > 0.01]
    print(f"{name:6s} spectral lines at: {', '.join(f'{v:.0f} Hz' for v in lines)}")

# %% Example 8: synthesise p(t) from a target PSD
def synthesise(S_target, fs, N, rng):
    """Pressure waveform p[n] in Pa whose one-sided PSD is S_target(f) in Pa²/Hz."""
    freqs = np.fft.rfftfreq(N, 1 / fs)
    df = fs / N
    amp = np.sqrt(S_target(freqs) * df)            # |P[k]| = sqrt(S_pp(f_k) * Δf)
    amp[0] = 0.0                                   # no DC
    amp[-1] = 0.0                                  # leave the Nyquist bin empty
    phase = rng.uniform(0, 2 * np.pi, freqs.size)  # random phases
    # irfft imposes Hermitian symmetry (a real output) and divides by N; N/√2 restores
    # pascals, with the one-sided PSD shared between positive and negative frequencies
    return np.fft.irfft(amp * np.exp(1j * phase), n=N) * N / np.sqrt(2)

def target(freq):
    """70 dB re 1 µPa²/Hz at 1 kHz, falling 5 dB per octave, from 100 Hz to 7 kHz."""
    freq = np.asarray(freq, dtype=float)
    nl = 70 - 5 * np.log2(np.maximum(freq, 1.0) / 1000)
    return np.where((freq >= 100) & (freq <= 7000), nl_to_psd(nl), 0.0)

p_syn = synthesise(target, fs, N, rng)
f2, S2 = welch(p_syn, fs=fs, window="hann", nperseg=4096, scaling="density")
for fc in (250, 1000, 4000):
    band = (f2 >= 0.95 * fc) & (f2 <= 1.05 * fc)
    print(f"{fc:5d} Hz: target {psd_to_nl(target(fc)):.1f}, "
          f"synthesised {psd_to_nl(np.mean(S2[band])):.1f} dB re 1 µPa²/Hz")
expected = np.sqrt(np.sum(target(freqs)) * fs / N)
print(f"RMS of the synthesised p(t): {np.sqrt(np.mean(p_syn**2)):.3f} Pa (expected {expected:.3f} Pa)")