"""TopEdu project 5: The missing notes of a closed pipe

Sound in a pipe closed at one end and open at the other obeys the wave equation. The closed end reflects pressure
as it is; the open end (ideally) forces pressure to zero. So the pipe only fits odd quarter-waves, and its notes are
f_n = (2n - 1) c / 4L: 1, 3, 5, ... times the lowest note, with no even harmonics.

Stage 1: run this file. It starts a pressure pulse in a 0.6 m pipe, listens at the closed end, finds the notes with
a Fourier transform, and checks the first two against the formula. The CHECK line should say PASS.
Stage 2: record real PVC pipes, find their notes with the same code, and measure what the ideal model leaves out.

Run:  python 05_pipe_resonance.py  [--show]      or     python 05_pipe_resonance.py my_recording.wav
Needs: numpy and matplotlib (reading .wav uses Python's own wave module)
"""
import os
import sys
import wave

import numpy as np
import matplotlib

if "--show" not in sys.argv:
    matplotlib.use("Agg")
import matplotlib.pyplot as plt

OUT = os.path.join(os.path.dirname(os.path.abspath(__file__)), "output")
os.makedirs(OUT, exist_ok=True)

LENGTH = 0.6    # metres
C = 343.0       # speed of sound, m/s (air at about 20 C)
N = 129         # grid points along the pipe


def simulate(length=LENGTH, c=C, n=N, round_trips=60):
    """Leapfrog steps of p_tt = c^2 p_xx. x = 0 is the closed end (mirror point), x = L the open end (p = 0).
    Returns the time step and the pressure heard at the closed end."""
    x = np.linspace(0, length, n)
    h = x[1] - x[0]
    dt = 0.9 * h / c                      # the CFL rule: a wave must not cross more than one cell per step
    steps = int(round_trips * 4 * length / c / dt)
    p_old = np.exp(-(5 * x / length) ** 2)     # a pressure bump at the closed end, starting at rest
    p_old[-1] = 0.0
    r2 = (c * dt / h) ** 2
    lap = lambda p: np.concatenate(([2 * (p[1] - p[0])], p[2:] - 2 * p[1:-1] + p[:-2], [0.0]))
    p = p_old + 0.5 * r2 * lap(p_old)          # first step, from zero starting speed
    heard = [p_old[0], p[0]]
    for _ in range(steps):
        p, p_old = 2 * p - p_old + r2 * lap(p), p
        p[-1] = 0.0
        heard.append(p[0])
    return dt, np.array(heard)


def peaks(signal, dt, count=4, lowest=20.0):
    """The strongest frequencies in a signal, refined between FFT bins by fitting a parabola to each peak."""
    signal = (signal - signal.mean()) * np.hanning(len(signal))
    size = 8 * len(signal)                                  # zero padding gives a smoother spectrum
    spectrum = np.abs(np.fft.rfft(signal, size))
    freqs = np.fft.rfftfreq(size, dt)
    found = []
    for i in range(1, len(spectrum) - 1):
        if freqs[i] > lowest and spectrum[i] > spectrum[i - 1] and spectrum[i] >= spectrum[i + 1] and spectrum[i] > 0.05 * spectrum.max():
            a, b, c = np.log(spectrum[i - 1:i + 2])
            shift = 0.5 * (a - c) / (a - 2 * b + c)
            found.append((spectrum[i], freqs[i] + shift * (freqs[1] - freqs[0])))
    found.sort(reverse=True)
    return sorted(f for _, f in found[:count]), freqs, spectrum


def read_wav(path):
    """A mono or stereo 16-bit .wav file (most phone recorders can export one). Returns dt and the samples."""
    with wave.open(path) as handle:
        rate, channels = handle.getframerate(), handle.getnchannels()
        data = np.frombuffer(handle.readframes(handle.getnframes()), dtype=np.int16).astype(float)
    return 1.0 / rate, data.reshape(-1, channels).mean(axis=1)


def main():
    dt, heard = simulate()
    notes, freqs, spectrum = peaks(heard, dt)
    formula = [(2 * k - 1) * C / (4 * LENGTH) for k in (1, 2, 3, 4)]
    worst = 0.0
    for k, f in enumerate(notes[:2]):
        error = abs(f - formula[k]) / formula[k]
        worst = max(worst, error)
        print(f"note {k + 1}: simulated {f:7.2f} Hz, formula {formula[k]:7.2f} Hz, error {100 * error:.3f}%")
    print(f"ratio of note 2 to note 1: {notes[1] / notes[0]:.3f} (the ideal pipe says 3: no even harmonic)")
    print(f"CHECK first two notes against (2n - 1) c / 4L -> {'PASS' if worst < 0.01 else 'FAIL'} (limit 1%)")

    fig, ax = plt.subplots(figsize=(6, 3.6))
    ax.semilogy(freqs, spectrum + 1e-12)
    for f in formula:
        ax.axvline(f, color="gray", lw=0.6, ls="--")
    ax.set_xlim(0, 1400); ax.set_xlabel("frequency (Hz)"); ax.set_ylabel("strength")
    ax.set_title("Simulated pipe: peaks sit on the dashed odd harmonics")
    fig.tight_layout(); fig.savefig(os.path.join(OUT, "05_spectrum.png"), dpi=150)

    recordings = [arg for arg in sys.argv[1:] if arg.lower().endswith(".wav")]
    for path in recordings:
        rdt, samples = read_wav(path)
        found, _, _ = peaks(samples, rdt, count=6)
        print(f"{os.path.basename(path)}: strongest notes " + ", ".join(f"{f:.1f} Hz" for f in found))
    print(f"plots saved in {OUT}")

    # ---- YOUR TURN ----
    # 1. Predict, then record: a 0.6 m PVC pipe closed at one end (a palm or a cap). Is its lowest note above or
    #    below the formula? Run:  python 05_pipe_resonance.py my_pipe.wav
    # 2. End correction: a real open end acts as if the pipe were longer, by about delta. From each recording,
    #    delta = c / (4 f_1) - L. Measure it for pipes of three or more widths. How does delta depend on the radius?
    #    (Textbooks quote about 0.6 times the radius. Do your pipes agree?)
    # 3. Look for even harmonics in your recordings. If any appear, what could make them?
    # 4. The speed of sound changes with temperature: c is about 331 + 0.6 T(C) m/s. Note the room temperature.
    if "--show" in sys.argv:
        plt.show()


if __name__ == "__main__":
    main()