#!/usr/bin/env python3
"""Render-parity harness (Phase 5 step 5, dB stage).

Measures the steady-state per-tone reduction directly on the reference wavs in
/home/m/soothe-bt/ and compares with the B.12 bridge model (g*LUT + w*warp^a).

Model: C(f) = g*LUT(log10(L0/res_band(f,fc,Q))) + w*warp(f)^a  ->  red = -20*log10(1-C).
LUT is the frozen PCHIP (B.11 knots). g=1.221, w=0.358, a=3.143 (model_fir.py canonical).

Reference renders: 24-bit WAV, input sources 16-bit mono. Tone amplitude estimated by
Goertzel at the exact tone frequency over a late steady window (3.0-3.75 s region).
"""
import wave
import numpy as np
from scipy.interpolate import PchipInterpolator

BT = '/home/m/soothe-bt/'
FS = 44100.0
GAIN = 4.132  # 10^(sens_dB/40) with sens_dB=24.65
QM = [0.1, 0.2, 0.3, 0.5, 0.7, 1.0, 1.5, 2.0, 3.0, 5.0, 10.0]
FCS = [800., 900., 950., 1000., 1050., 1100., 1200.]
L0D = 10 ** (-7.142 / 20)   # dual input tone level
L0Q = 10 ** (-18.063 / 20)  # t1kq input tone level
L0T = 1.0                   # t1k input tone level (0 dBFS)

# frozen LUT knots (B.11)
LX = np.array([-0.75, -0.5, -0.25, 0, 0.25, 0.5, 0.574, 0.61, 0.75, 1.0])
LY = np.array([0.4402, 0.4552, 0.4813, 0.5072, 0.5329, 0.5332, 0.5645, 0.6471, 0.6562, 0.6670])
LUT = PchipInterpolator(LX, LY)


def load(p):
    w = wave.open(p, 'rb')
    n, ch, sr, sw = w.getnframes(), w.getnchannels(), w.getframerate(), w.getsampwidth()
    b = w.readframes(n)
    N = n * ch
    if sw == 2:
        x = np.frombuffer(b, dtype='<i2').astype(np.float64) / 32768.0
    else:
        raw = np.frombuffer(b, dtype=np.uint8).reshape(N, 3)
        x = (raw[:, 0].astype(np.int64) | (raw[:, 1].astype(np.int64) << 8)
             | (raw[:, 2].astype(np.int64) << 16))
        x = np.where(x >= 0x800000, x - 0x1000000, x).astype(np.float64) / 8388607.0
    return x.reshape(-1, ch)


def tone_amp(p, f):
    a = load(p)
    end = min(len(a), int(3.5 * FS))
    seg = a[:end][-int(0.75 * FS):]
    x = np.mean(seg, axis=1)
    n = len(x)
    w = 2 * np.pi * f / FS
    cw = 2 * np.cos(w)
    s0 = s1 = s2 = 0.0
    for v in x:
        s2 = s1
        s1 = s0
        s0 = v + cw * s1 - s2
    return np.sqrt(abs(s0 * s0 + s1 * s1 - 2 * cw * s0 * s1)) / n


def red(src, out, f):
    return 20 * np.log10(tone_amp(src, f) / tone_amp(out, f))


def res_band(f, fc, Q):
    w0 = fc * 2 * np.pi / FS
    c, s = np.cos(w0), np.sin(w0)
    p = (s * 0.5) / Q
    a, a2 = p * GAIN, p / GAIN
    A = [a + 1, -2 * c, 1 - a]
    B = [a2 + 1, -2 * c, 1 - a2]
    ww = 2 * np.pi * f / FS
    z = np.exp(-1j * ww)
    return abs(2 * (B[0] + B[1] * z + B[2] * z * z) / (A[0] + A[1] * z + A[2] * z * z))


def warp(f):
    x = f / 2000.
    return 0.87 * 7.942 * x / (7.942 + x)


def pred(f, fc, Q, L0):
    xv = np.log10(L0 / res_band(f, fc, Q))
    C = 1.221 * float(LUT(float(np.clip(xv, -1, 1.5)))) + 0.358 * warp(f) ** 3.143
    return -20 * np.log10(1 - min(C, 0.999))


def main():
    meas, pr = [], []
    for q in QM:
        for f in (500, 2000):
            meas.append(red(BT + 'dual.wav', BT + f'dual_b1q_{q}.wav', f))
            pr.append(pred(f, 500, q, L0D))
    for fc in FCS:
        meas.append(red(BT + 'tone1kq.wav', BT + f't1kq_b1f_{int(fc)}.wav', 1000))
        pr.append(pred(1000, fc, 0.9999978, L0Q))
    for fc in FCS:
        meas.append(red(BT + 'tone1k.wav', BT + f't1k_b1f_{int(fc)}.wav', 1000))
        pr.append(pred(1000, fc, 0.9999978, L0T))
    meas = np.array(meas)
    pr = np.array(pr)
    print('RENDER-PARITY vs /home/m/soothe-bt (36 pts, real wavs), B.12 (g=1.221,w=0.358,a=3.143):')
    print('  TOTAL rmse = %.4f dB' % np.sqrt(np.mean((pr - meas) ** 2)))
    for name, sl in [('dual500', slice(0, 22, 2)), ('dual2000', slice(1, 22, 2)),
                     ('t1kq', slice(22, 29)), ('t1k', slice(29, 36))]:
        print('  %-9s rmse=%.4f' % (name, np.sqrt(np.mean((pr[sl] - meas[sl]) ** 2))))
    print('  dual2000 resid:', ' '.join('%+.2f' % x for x in (pr[1:22:2] - meas[1:22:2])))
    print('  t1kq resid:    ', ' '.join('%+.2f' % x for x in (pr[22:29] - meas[22:29])))
    print('  t1k resid:     ', ' '.join('%+.2f' % x for x in (pr[29:] - meas[29:])))


if __name__ == '__main__':
    main()