"""How much of the envelope match is a static equalizer?

Usage: python3 eq.py REFERENCE.wav SIMULATED.f32 OUT_PREFIX

Recomputes the envelope-mode gain table of match.py, splits it into a static part
(the per-band gain at 0.1 s) and a time-varying remainder, reports the size of the
remainder, and writes
  OUT_eq.wav    the simulation through the static per-band equalizer (one gain per band)
  OUT_eq4.wav   the same with the equalizer smoothed to a cubic in log frequency (4 numbers)
"""
import sys, numpy as np
from scipy.signal import stft, istft
import match as M

ref_path, sim_path, out = sys.argv[1:4]
FS, NPERSEG, HOP = M.FS, M.NPERSEG, M.HOP
ref = M.load(ref_path)
sim = M.load(sim_path)
o_ref, o_sim = M.onset(ref), M.onset(sim)
noise = ref[: max(int(0.2 * FS), o_ref - int(0.05 * FS))]
ref, sim = ref[o_ref:], sim[o_sim:]
ref_full = ref
n = min(len(ref), len(sim))
ref, sim = ref[:n], sim[:n]
f, t, Zr = stft(ref, FS, nperseg=NPERSEG, noverlap=NPERSEG - HOP)
_, _, Zs = stft(sim, FS, nperseg=NPERSEG, noverlap=NPERSEG - HOP)
Er, Es = M.band_envelopes(np.abs(Zr) ** 2, f), M.band_envelopes(np.abs(Zs) ** 2, f)
_, _, Zn = stft(noise, FS, nperseg=NPERSEG, noverlap=NPERSEG - HOP)
En = M.band_envelopes(np.abs(Zn) ** 2, f)
_, _, Zt = stft(ref_full[-int(0.5 * FS):], FS, nperseg=NPERSEG, noverlap=NPERSEG - HOP)
Et = M.band_envelopes(np.abs(Zt) ** 2, f)
floor = np.maximum(np.median(En, axis=1) + 3, np.median(Et, axis=1) - 3)
i_ref = np.searchsorted(t, 0.1)
C = M.CENTERS

G0 = (Er - Es)[:, i_ref]  # static part: balance at 0.1 s
print("static equalizer (dB gain applied to the simulation per third-octave band):")
for c, g in zip(C, G0):
    print(f"  {c:7.0f} Hz  {g:+6.1f}")
# Time-varying remainder over the frames where the recording is above its floor.
print("\nremainder after the static part (recording above floor + 10 dB):")
print("  band     trusted until   remainder RMS   remainder at end of trusted range")
for b in range(len(C)):
    below = np.where(Er[b, i_ref:] < floor[b] + 10)[0]
    last = i_ref + (below[0] if len(below) else len(t) - i_ref) - 1
    if last <= i_ref + 5:
        continue
    r = (Er - Es)[b, i_ref:last] - G0[b]
    print(f"  {C[b]:6.0f} Hz   {t[last]:5.2f} s        {np.sqrt(np.mean(r ** 2)):5.1f} dB        {r[-1]:+6.1f} dB")

# Static equalizer applied to the simulation, full and smoothed.
lf = np.log(np.maximum(f, 1.0))
lc = np.log(C)
usable = (C > 100) & (C < 6000)  # where the recording informs the balance
coef = np.polyfit(lc[usable], G0[usable], 3)
G4 = np.polyval(coef, np.clip(lc, lc[usable].min(), lc[usable].max()))
print("\ncubic in log frequency (4 coefficients) versus the 28 band gains, RMS difference over 100 Hz to 6 kHz:",
      f"{np.sqrt(np.mean((G4 - G0)[usable] ** 2)):.1f} dB")
i_a, i_b = int(0.1 * FS), int(1.0 * FS)
body = lambda x: np.sqrt(np.mean(x[i_a:i_b] ** 2))
target = body(ref)
for name, G in (("eq", G0), ("eq4", G4)):
    Gf = np.interp(lf, lc, G)
    _, y = istft(Zs * 10 ** (Gf / 20)[:, None], FS, nperseg=NPERSEG, noverlap=NPERSEG - HOP)
    y = y[:n]
    M.save(f"{out}_{name}.wav", np.tanh(y * (target / body(y))))
print(f"wrote {out}_eq.wav and {out}_eq4.wav")