#!/usr/bin/env python3
"""fit_maskaxis.py — brute-force подбор: какая формула полосы даёт широкую кривую t1kq.

Гипотеза из декомпиляции: близнец вычисляет маску M(z)=2·B(z)/A(z) на z-оси,
z = rotor(вход) = e^{i(pi/2 - x)} (и negate → e^{i(pi/2 + x)}). Вход x — частотная
ось бинов (в рад). Снижение тона f_tone при полосе на fc: D(fc) = -20log10(M(z_tone;fc)),
где M зависит от fc через коэффициенты case8/case1.

Ищем формулу + Q_eff, воспроизводящие измеренную широкую кривую.
"""
import numpy as np

FS = 44100.0

# измеренные кривые (fc, D_dB) от bandshape.py
T1KQ = np.array([
    (800, 7.868), (900, 8.536), (950, 8.726), (980, 8.779),
    (1000, 8.788), (1020, 8.778), (1050, 8.729), (1100, 8.575), (1200, 8.115)])
T1K = np.array([
    (500, 11.674), (800, 14.548), (900, 15.332), (950, 15.553),
    (1000, 15.626), (1050, 15.557), (1100, 15.378), (1200, 14.840),
    (1500, 13.209), (2000, 11.668)])


def ab_case8(fc, Q, gain):
    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 A, B


def ab_case8_sin(fc, Q, gain):
    w0 = fc * 2 * np.pi / FS
    c, s = np.cos(w0), np.sin(w0)
    p = (s * 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 A, B


def ab_case1(fc, Q, gain=None):
    w = 1.0 / np.sin(fc * np.pi / FS)
    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 A, B


def ab_case1_noQ(fc, Q, gain=None):
    w = 1.0 / np.sin(fc * np.pi / FS)
    w2 = w * w
    a = 1.0 / (w + 1.0 + w2)   # k=(1/Q)*w с Q=1
    A = [a, 2 * a, a]
    B = [1.0, (1.0 - w2) * 2 * a, (1.0 - w + w2) * a]
    return A, B


def z_at_tone(f_tone, variant, zsign=1.0):
    wt = 2 * np.pi * f_tone / FS
    if variant == 'unit':
        return np.exp(zsign * 1j * wt)
    if variant == 'rotor_pihalf':
        # rotor: z = sin(x)+i cos(x) = e^{i(pi/2-x)}; negate -> -e^{i(pi/2-x)} = e^{i(pi/2-x+pi)}
        return -np.exp(1j * (np.pi / 2 - wt))
    if variant == 'rotor_pihalf_pos':
        return np.exp(1j * (np.pi / 2 - wt))
    raise ValueError(variant)


def h_at(A, B, z):
    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 Dcurve(rows, ab_fn, variant, Q, gain, zsign=1.0, f_tone=1000.0):
    z = z_at_tone(f_tone, variant, zsign)
    ds = []
    for fc, _ in rows:
        A, B = ab_fn(fc, Q, gain)
        ds.append(-20 * np.log10(max(h_at(A, B, z), 1e-12)))
    return np.array(ds)


def fit_shape(rows, ab_fn, variant, gain, zsign=1.0, f_tone=1000.0):
    meas = np.array([r[1] for r in rows])
    best = None
    for Q in np.logspace(-2, 1.6, 180):
        pred = Dcurve(rows, ab_fn, variant, Q, gain, zsign, f_tone)
        # сравнение по форме: вычитаем среднее
        err = np.mean((pred - pred.mean() - (meas - meas.mean()))**2)
        rmse = np.sqrt(err)
        if best is None or rmse < best[0]:
            best = (rmse, Q, pred)
    return best


def run():
    print('=== t1kq_only1 (tone 1k, q=0.9999978) ===')
    rows = T1KQ
    meas = np.array([r[1] for r in rows])
    variants = ['unit', 'rotor_pihalf', 'rotor_pihalf_pos']
    ab_variants = [('case8_cos', ab_case8), ('case8_sin', ab_case8_sin),
                   ('case1', ab_case1)]
    results = []
    for vn in variants:
        for abn, abf in ab_variants:
            for zsign in (1.0, -1.0):
                for gain in (1.0, 10.0**(12.0 / 20.0)):
                    rmse, Q, pred = fit_shape(rows, abf, vn, gain, zsign)
                    results.append((rmse, vn, abn, zsign, gain, Q, pred))
    results.sort(key=lambda r: r[0])
    for rmse, vn, abn, zsign, gain, Q, pred in results[:12]:
        print(f'rmse={rmse:.4f} z={vn} zs={zsign:+} {abn} gain={gain:.2f} Q_eff={Q:.5f}')
        print('   pred: ' + ' '.join(f'{p:6.3f}' for p in pred))
        print('   meas: ' + ' '.join(f'{m:6.3f}' for m in meas))
        print()


if __name__ == '__main__':
    run()