"""Spectrogram comparison of the recording, the simulation and the matched outputs.

Usage: python3 plot.py OUT_PREFIX  (reads the wavs written by match.py)
"""
import sys, wave, numpy as np
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from scipy.signal import stft

FS = 48000


def load(path):
    w = wave.open(path)
    return np.frombuffer(w.readframes(w.getnframes()), dtype="<i2").astype(float) / 32768


out = sys.argv[1]
files = [("recording", f"{out}_reference_aligned.wav"), ("simulation", f"{out}_sim_aligned.wav"),
         ("matched: decay rates", f"{out}_decay.wav"), ("matched: full envelope", f"{out}_envelope.wav")]
fig, axes = plt.subplots(2, 2, figsize=(13, 8), sharex=True, sharey=True)
for ax, (title, path) in zip(axes.flat, files):
    x = load(path)
    f, t, Z = stft(x, FS, nperseg=4096, noverlap=4096 - 512)
    S = 20 * np.log10(np.abs(Z) + 1e-9)
    S -= S.max()
    ax.pcolormesh(t, f, S, vmin=-90, vmax=0, cmap="magma", shading="auto")
    ax.set_yscale("log")
    ax.set_ylim(25, 12000)
    ax.set_title(title)
    ax.set_xlabel("time [s]")
    ax.set_ylabel("frequency [Hz]")
fig.suptitle("Reference hit versus simulated strike (dB relative to each clip's peak, 90 dB range)")
fig.tight_layout()
fig.savefig(f"{out}_spectrograms.png", dpi=110)
print(f"wrote {out}_spectrograms.png")

# Band decay curves for a few bands.
a = np.load(f"{out}_analysis.npz")
t, Er, Es, centers, floor = a["t"], a["Er"], a["Es"], a["centers"], a["floor"]
fig, axes = plt.subplots(2, 4, figsize=(14, 6), sharex=True)
for ax, fc in zip(axes.flat, [35, 89, 178, 356, 713, 1425, 2851, 5702]):
    b = int(np.argmin(np.abs(centers - fc)))
    ax.plot(t, Er[b] - Er[b, np.searchsorted(t, 0.1)], label="recording")
    ax.plot(t, Es[b] - Es[b, np.searchsorted(t, 0.1)], label="simulation")
    ax.axhline(floor[b] - Er[b, np.searchsorted(t, 0.1)], color="gray", ls=":", lw=1, label="recording noise floor")
    ax.set_title(f"{centers[b]:.0f} Hz band")
    ax.set_ylim(-80, 10)
    ax.grid(alpha=0.3)
axes[0, 0].legend(fontsize=8)
for ax in axes[1]:
    ax.set_xlabel("time after onset [s]")
for ax in axes[:, 0]:
    ax.set_ylabel("level rel. 0.1 s [dB]")
fig.tight_layout()
fig.savefig(f"{out}_band_decay.png", dpi=110)
print(f"wrote {out}_band_decay.png")