#!/usr/bin/env python3
"""Burst analyzer for soothe2 renders.
Reads a 24-bit stereo wav (fx) and dry, computes bin magnitude at fc with a
10ms sliding window, then tracks NOTCH DEPTH over time:
    reduction(t) = -20*log10(g_fx(t)/g_dry(t))   (>=0 when suppressed)
and fits:
  - settled depth (mean over plateau window)
  - attack tau    (fit reduction: A*(1-exp(-(t-t0)/tau)) on burst onset)
  - release tau   (fit reduction: A*exp(-(t-b1)/tau) after burst end)
Works best when a weak probe tone at fc is present during the whole file so
the notch remains observable after the burst (see synth_burstp.py).
Usage: measure.py <fx.wav> <dry.wav> --fc 500 [--burst 0.5 1.5] [--name X]
"""
import argparse, os
import numpy as np

def read_wav(path):
    import wave
    w = wave.open(path, 'rb')
    sw, nc, n = w.getsampwidth(), w.getnchannels(), w.getnframes()
    d = np.frombuffer(w.readframes(n), dtype=np.uint8).reshape(n, nc, sw)
    w.close()
    ch = d[:, 0, :]
    v = ch[:, 0].astype(np.int64) | (ch[:, 1].astype(np.int64) << 8) | (ch[:, 2].astype(np.int64) << 16)
    v = (v ^ (1 << 23)) - (1 << 23)
    return v.astype(np.float64) / (1 << 23)

def mag_at(x, fc, sr, win=0.010):
    w = int(sr * win)
    n = x.size // w
    if n == 0: return np.array([])
    xw = x[:n * w].reshape(n, w)
    X = np.fft.rfft(xw, axis=1)
    k = int(round(fc * w / sr))
    return 2.0 * np.abs(X[:, k]) / w

def fit_tau(t, red, t0, A, lo, hi, t1=None, invert=False, floor=None):
    """fit red(t) = A*(1-exp(-(t-t0)/tau)) [invert=False] or A*exp(-(t-t0)/tau) [invert].
    Window [t0, t1]; grid search tau on [lo,hi] by r2. Returns (tau, r2)."""
    if t1 is None:
        t1 = t0 + 1.5
    m = (t >= t0) & (t <= t1)
    tt, rr = t[m], red[m]
    if floor is not None:
        rr = rr - floor
    if rr.size < 4: return (float('nan'), 0.0)
    best, bestr2 = None, -1e9
    for tau in np.geomspace(lo, hi, 200):
        if invert:
            pred = A * np.exp(-(tt - t0) / tau)
        else:
            pred = A * (1 - np.exp(-(tt - t0) / tau))
        ss = 1 - np.sum((rr - pred) ** 2) / np.sum((rr - rr.mean()) ** 2)
        if ss > bestr2:
            bestr2, best = ss, tau
    return (best, bestr2)

def profile(fx_path, dry_path, freqs, win0=1.6, win1=3.0, sr=44100, win=0.010):
    """Notch transfer profile via probe tones: red(f) in window. Returns dict f->dB."""
    fx = read_wav(fx_path); dry = read_wav(dry_path)
    n = min(fx.size, dry.size)
    fx, dry = fx[:n], dry[:n]
    out = {}
    for f in freqs:
        mf = mag_at(fx, f, sr, win); md = mag_at(dry, f, sr, win)
        tw = (np.arange(min(mf.size, md.size)) + 0.5) * win
        m = (tw >= win0) & (tw <= win1)
        if m.sum() == 0:
            out[f] = float('nan'); continue
        out[f] = -20 * np.log10(
            np.maximum(mf[m].mean(), 1e-9) / np.maximum(md[m].mean(), 1e-9))
    return out

def measure(fx_path, dry_path, fc=500.0, burst=(0.5, 1.5), sr=44100, win=0.010):
    fx = read_wav(fx_path); dry = read_wav(dry_path)
    n = min(fx.size, dry.size)
    fx, dry = fx[:n], dry[:n]
    mf = mag_at(fx, fc, sr, win); md = mag_at(dry, fc, sr, win)
    tw = (np.arange(mf.size) + 0.5) * win
    nmin = min(mf.size, md.size)
    tw, mf, md = tw[:nmin], mf[:nmin], md[:nmin]
    guard = 1e-9
    red = -20 * np.log10(np.maximum(mf, guard) / np.maximum(md, guard))

    b0, b1 = burst
    pl0, pl1 = b1 - 0.30, b1 - 0.01
    pm = (tw >= pl0) & (tw <= pl1)
    settled = red[pm].mean() if pm.sum() else float('nan')

    # attack: relative rise of reduction from burst start; A=settled
    t0a = b0
    ta, r2a = fit_tau(tw, red, t0a, max(settled, 0.001), 0.001, 1.0, t1=b1 - 0.03)
    # release: decay (invert) after burst end; A=settled, toward a floor
    # floor = mean of red in the last stable probe window (well after burst)
    tail = (tw >= b1 + 0.35) & (tw <= b1 + 0.85)
    floor = red[tail].mean() if tail.sum() else 0.0
    tr, r2r = fit_tau(tw, red, b1, max(settled - floor, 0.001), 0.001, 5.0,
                      t1=b1 + 0.35, invert=True, floor=floor)

    return dict(depth_db=settled, attack_tau=ta, attack_r2=r2a,
                release_tau=tr, release_r2=r2r)

if __name__ == '__main__':
    ap = argparse.ArgumentParser()
    ap.add_argument('fx'); ap.add_argument('dry')
    ap.add_argument('--fc', type=float, default=500.0)
    ap.add_argument('--burst', nargs=2, type=float, default=[0.5, 1.5])
    ap.add_argument('--name')
    a = ap.parse_args()
    r = measure(a.fx, a.dry, a.fc, tuple(a.burst))
    nm = a.name or os.path.basename(a.fx)
    print(f"{nm}: depth={r['depth_db']:7.2f}dB  att_tau={r['attack_tau']*1000:7.2f}ms "
          f"(r2={r['attack_r2']:.3f})  rel_tau={r['release_tau']*1000:7.2f}ms "
          f"(r2={r['release_r2']:.3f})")