#!/usr/bin/env python3
import sys, numpy as np
sys.path.insert(0, '/home/m/re-tools')
from render_parity import load
from scipy.optimize import minimize
import wave

BT='/home/m/soothe-bt/'; FS=44100.0; GAIN=4.132

def bandres(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]
    w=2*np.pi*np.asarray(f)/FS; z=np.exp(-1j*w)
    return np.abs(2.0*(B[0]+B[1]*z+B[2]*z*z)/(A[0]+A[1]*z+A[2]*z*z))

def warp(f): x=np.asarray(f)/2000.0; return 0.87*7.942*x/(7.942+x)

def lut_powerlaw(xv, C, A=0.0, B=1.0):
    """FUN_180563440 power-law mode: centered=2*x-1, result=sign*10^(log10|/C)"""
    x = np.clip(xv, A, B)
    centered = 2.0 * (x - A) / (B - A) - 1.0  # map [A,B] → [-1,1]
    abs_c = np.abs(centered)
    result = np.where(abs_c > 1e-9,
                      np.sign(centered) * np.power(10.0, np.log10(np.maximum(abs_c, 1e-9)) / C),
                      0.0)
    return (result + 1.0) / 2.0  # remap [-1,1] → [0,1]

def lut_linear(xv, A=0.483, B=0.717):
    """Linear mode: (B-A)*x+A"""
    return np.clip(A + (B - A) * np.clip(xv, 0, 1), 0, 1)

def frames(x, fc, Q, G_, W_, A_, C_lut=10.0, mode='powerlaw', NS=2048, hop=512, tatt=0.011, trel=0.08):
    win=np.sqrt(np.hanning(NS)); wsum=win.sum(); n=len(x)
    nfr=max(1,int(np.ceil((n-NS)/hop))+1)
    X=np.empty((nfr,NS//2+1),dtype=complex); freqs=np.fft.rfftfreq(NS,1/FS)
    res=bandres(freqs,fc,Q)
    att=np.exp(-hop/(tatt*FS)); rel=np.exp(-hop/(trel*FS))
    am=np.zeros(freqs.size); G=np.empty(X.shape)
    for m in range(nfr):
        s=m*hop; seg=np.zeros(NS); kk=min(NS,n-s); seg[:kk]=x[s:s+kk]
        F=np.fft.rfft(win*seg); X[m]=F; ac=2*np.abs(F)/wsum
        am=np.where(ac>am, att*am+(1-att)*ac, rel*am+(1-rel)*ac)
        xv=np.log10(np.maximum(am/np.maximum(res,1e-12),1e-9))
        if mode=='powerlaw':
            L = lut_powerlaw(xv, C_lut)
        else:
            L = lut_linear(xv)
        C = G_*L + W_*warp(freqs)**A_
        G[m] = np.maximum(1-np.minimum(C,0.95),1e-9)
    return X,G,win,hop,n

def synthe(X,G,win,hop,n):
    out=np.zeros(n); acc=np.zeros(n); NS=len(win)
    for m in range(X.shape[0]):
        seg=np.fft.irfft(X[m]*G[m])*win; s=m*hop; lay=min(NS,n-s)
        out[s:s+lay]+=seg[:lay]; acc[s:s+lay]+=(win*win)[:lay]
    return out/np.maximum(acc,1e-12)

def tone_cmp(x,f,seglen=0.75*FS):
    x=np.asarray(x)[-int(seglen):]; n=len(x); t=np.arange(n)/FS; w=2*np.pi*f
    return np.hypot(2*np.sum(x*np.cos(w*t))/n,2*np.sum(x*np.sin(w*t))/n)

def dB(v): return 20*np.log10(np.clip(v,1e-9,None))

x=np.mean(load(BT+'dual.wav'),axis=1)
refs=[(q,f'dual_b1q_{q}.wav',f) for q in [0.1,1.0,10.0] for f in (500,2000)]
tone_ref={(r,f):dB(tone_cmp(np.mean(load(BT+r),axis=1),f)) for _,r,f in refs}

def score(GWA, C_lut, mode):
    G_,W_,A_=GWA; tot=[]
    for q,ref,f in refs:
        X,G,win,hop,n=frames(x,500.0,q,G_,W_,A_,C_lut,mode)
        y=synthe(X,G,win,hop,n); tot.append(dB(tone_cmp(y,f))-tone_ref[(ref,f)])
    return np.array(tot)

# Test power-law mode with different C values
print("=== POWER-LAW MODE ===")
for C in [2.0, 3.0, 5.0, 8.0, 10.0, 15.0, 20.0]:
    def obj(p):
        return np.mean(np.abs(score(p, C, 'powerlaw')))
    r = minimize(obj, [1.0, 0.3, 1.0], method='Nelder-Mead', options=dict(maxiter=500))
    err = score(r.x, C, 'powerlaw')
    print(f'  C={C:5.1f} G={r.x[0]:.3f} W={r.x[1]:.3f} A={r.x[2]:.3f} mean={np.mean(np.abs(err)):.3f} '
          f'errs=[{",".join(f"{e:+.2f}" for e in err)}]')

# Also test linear mode
print("=== LINEAR MODE ===")
def obj_lin(p):
    return np.mean(np.abs(score(p, 0, 'linear')))
r = minimize(obj_lin, [1.0, 0.3, 1.0], method='Nelder-Mead', options=dict(maxiter=500))
err = score(r.x, 0, 'linear')
print(f'  G={r.x[0]:.3f} W={r.x[1]:.3f} A={r.x[2]:.3f} mean={np.mean(np.abs(err)):.3f} '
      f'errs=[{",".join(f"{e:+.2f}" for e in err)}]')