#!/usr/bin/env python3
"""fit_bandshape.py — подбор модели формы полосы к измеренной кривой.

Модели (из декомпиляции, /tmp/det/detector_decoded.md):
  case1 (mode=1 bell): FUN_1805343e0(fs, fc, Q=0.707)
      w = 1/FUN_181a14cfa(fc*pi/fs);  k = (1/Q)*w;  w2=w*w;  a=1/(k+1+w2)
      A=[a,2a,a];  B=[1,(1-w2)*2a,(1-k+w2)*a]
  case8 (RBJ bell+gain): FUN_180533ec0(fs, fc, Q=0.9999978542327881, gain=10^(sens/20))
      w0=fc*2pi/fs; c=cos(w0); s=sin(w0); p=(c*0.5)/Q
      alpha=p*gain; alpha2=p/gain
      A=[alpha+1, s*(-0.5), 1-alpha];  B=[alpha2+1, s*(-0.5), 1-alpha2]
Передаточная функция H(f) = |(A0+A1 z^-1+A2 z^-2)/(B0+B1 z^-1+B2 z^-2)|, z=e^{j2pi f/fs}.
Кривая подавления тона f_tone при полосе на fc: D(fc) = -20log10(|H(f_tone)|), норм.
"""
import numpy as np
from bandshape import measure_series

FS = 44100.0

def cfa_candidates(x):
    return {
        'sin': np.sin(x),
        'sqrt': np.sqrt(x),
        'tan': np.tan(x),
        'asin': np.arcsin(np.clip(x, -1, 1)),
        'csc': 1.0/np.sin(x),
        'cot': 1.0/np.tan(x),
        'none': x,
    }

def cfa_fn(name):
    return cfa_candidates(1.0)[name].__class__ if False else {
        'sin': np.sin, 'sqrt': np.sqrt, 'tan': np.tan,
        'asin': lambda x: np.arcsin(np.clip(x, -1, 1)),
        'csc': lambda x: 1.0/np.sin(x),
        'cot': lambda x: 1.0/np.tan(x),
        'none': lambda x: x,
    }[name]

def h_mag(A, B, f, fs=FS, zsign=-1.0):
    z = np.exp(zsign*2j*np.pi*f/fs)
    num = B[0] + B[1]*z + B[2]*z**2
    den = A[0] + A[1]*z + A[2]*z**2
    return np.abs(2.0*num/den)

def model_case1(fc, f_tone, Q, cfa, fs=FS):
    x = fc*np.pi/fs
    w = 1.0/cfa(x)
    k = (1.0/Q)*w
    w2 = w*w
    a = 1.0/(k+1.0+w2)
    A = [a, 2*a, a]
    B = [1.0, (1.0-w2)*2*a, (1.0-k+w2)*a]
    return h_mag(A, B, f_tone, fs)

def model_case8(fc, f_tone, Q, gain, fs=FS):
    w0 = fc*2*np.pi/fs
    c, s = np.cos(w0), np.sin(w0)
    p = (c*0.5)/Q
    alpha = p*gain
    alpha2 = p/gain
    A = [alpha+1.0, s*(-0.5), 1.0-alpha]
    B = [alpha2+1.0, s*(-0.5), 1.0-alpha2]
    return h_mag(A, B, f_tone, fs)

def norm_db(d):
    d = np.asarray(d, float)
    return d - d[np.argmax(d)]

def fit(rows, pred_fn, f_tone=1000.0):
    fc = np.array([r[0] for r in rows])
    meas = np.array([r[1] for r in rows])
    pred = np.array([pred_fn(f, f_tone) for f in fc])
    pred_db = norm_db(-20*np.log10(np.maximum(pred, 1e-12)))
    meas_db = norm_db(meas)
    se = np.mean((pred_db - meas_db)**2)
    rmse = np.sqrt(se)
    return rmse, pred_db, meas_db, fc

def run(rows, tag):
    print(f'=== {tag} ===')
    best = []
    cfa_fns = cfa_fn
    # case1
    for qname, q in [('Q0.707', 0.707), ('Q1.0', 1.0), ('Q0.99999', 0.9999978542327881)]:
        for cfa_name in ['sin', 'sqrt', 'tan', 'asin', 'csc', 'cot', 'none']:
            def make(cfa_name, q):
                def f_(fc, ft):
                    return model_case1(fc, ft, q, cfa_fns(cfa_name))
                return f_
            rmse, p, m, fc = fit(rows, make(cfa_name, q))
            best.append((rmse, f'case1 {qname} cfa={cfa_name}', p))
    # case8
    for qname, q in [('Q0.99999', 0.9999978542327881), ('Q1.0', 1.0), ('Q0.707', 0.707)]:
        for gain in [10.0**(12.0/20.0), 1.0]:
            def make(q, gain):
                def f_(fc, ft):
                    return model_case8(fc, ft, q, gain)
                return f_
            rmse, p, m, fc = fit(rows, make(q, gain))
            best.append((rmse, f'case8 {qname} gain={gain:.2f}', p))
    best.sort()
    print(f'{"model":45s} rmse')
    for rmse, name, _ in best[:8]:
        print(f'{name:45s} {rmse:.4f}')
    print()
    # best curve vs meas
    rmse, name, p = best[0]
    m = np.array([r[1] for r in rows])
    print(f'best={name} rmse={rmse:.4f}')
    for i, r in enumerate(rows):
        print(f'  fc={r[0]:8.1f} meas={m[i]:7.3f} model={p[i]:7.3f} dB')
    return best

if __name__ == '__main__':
    rows = measure_series('t1kq_only1',
                   '/home/m/soothe-bt/tone1kq.wav',
                   [800, 900, 950, 980, 1000, 1020, 1050, 1100, 1200], tag='t1kq_only1')
    run(rows, 't1kq_only1')