#!/usr/bin/env python3
"""Phase A step 3 (fast): combo-vectorized law grid-search."""
import re, sys
import numpy as np

SRC = '/home/m/re-tools/dsp/rt_mask_tables.cpp'
src = open(SRC).read()
def tab(name):
    m = re.search(r'const double %s\[\] = \{(.*?)\};' % name, src, re.S)
    return np.array([float(x) for x in re.findall(r'[-+0-9.eE]+', m.group(1))])
A1, B1 = tab('kIIR_A1'), tab('kIIR_B1')
A2, B2 = tab('kIIR_A2'), tab('kIIR_B2')
A3, B3 = tab('kIIR_A3'), tab('kIIR_B3')
DUMPDIR = {'dump_res_new.bin': '/tmp/', 'dump_t1k.bin': '/tmp/'}

def load_dump(p):
    d = np.loadtxt(p, skiprows=1); return d[:, 2], d[:, 6]
def load_traj(p):
    b = open(p, 'rb').read(); off = 0; fr = []
    while off < len(b):
        _, nb = np.frombuffer(b, dtype=np.int32, count=2, offset=off); off += 8
        fr.append(np.frombuffer(b, dtype='<f4', count=int(nb), offset=off).astype(np.float64)); off += 4*int(nb)
    return np.array(fr)

WIN = {'res_500': (55, 90), 'al_12': (243, 278), 'al_24': (243, 278), 't1k_1000': (243, 278)}
MEAS_OLD = {'res_500': 0.219, 'al_12': 0.450, 'al_24': -1.822, 't1k_1000': 1.816}

DATA = {}
for name, traj, dump, bm in [
        ('res_500','traj_res500.bin','dump_res_new.bin',85),
        ('al_12','traj_al12.bin','dump_t1k.bin',85),
        ('al_24','traj_al24.bin','dump_t1k.bin',85),
        ('t1k_1000','traj_t1k.bin','dump_t1k.bin',85)]:
    res_k, W = load_dump(DUMPDIR[dump]+dump)
    T = load_traj('/tmp/opencode/'+traj)[WIN[name][0]:WIN[name][1]]
    dB = np.log10(np.maximum(T, 1e-12))*20.0
    if name != 'res_500':
        lv = dB[:, bm]; keep_dB = dB[lv >= lv.max()-6]; keep_n = (keep_dB.shape[0],)
    else:
        keep_dB = dB; keep_n = None
    DATA[name] = (np.ascontiguousarray(keep_dB, dtype=np.float64), W.astype(np.float64), bm, keep_n)

def gains_batch(dB, W, bm, X0s, SLs, CMs, C_pre=None, FLs=None):
    if FLs is None: FLs = np.zeros_like(X0s)
    """dB [T,nbin]; returns G [C,T] gain at bm for each combo."""
    import sys
    print('gains_batch shapes:', dB.shape, W.shape, bm, X0s.shape, file=sys.stderr)
    C, T, N = len(X0s), dB.shape[0], dB.shape[1]
    if C_pre is not None:
        c = np.broadcast_to(C_pre, (C, T, N))
    else:
        X0 = X0s[:, None, None]; SL = SLs[:, None, None]; CM = CMs[:, None, None]
        FL = FLs[:, None, None]
        c = np.clip(X0 + SL*dB, FL, CM)                       # [C,T,N]
    acc = np.zeros((C, T))
    y = np.empty_like(c)
    for i in range(N):
        acc = A1[i]*acc + B1[i]*c[:, :, i]
        y[:, :, i] = acc
    acc = np.zeros((C, T))
    for i in range(N):
        acc = A2[i]*acc + B2[i]*y[:, :, i]
        y[:, :, i] = 0.8*np.exp2(-acc)*W[i]
    # IIR3 bidi x2 on y
    for _ in range(2):
        st = np.zeros((C, T))
        for i in range(N):
            st = y[:, :, i]*B3[i] + st*A3[i]
            y[:, :, i] = st
        st = y[:, :, -1].copy()
        for i in range(N-2, 0, -1):
            st = y[:, :, i]*B3[i] + st*A3[i]
            y[:, :, i] = st
    return y[:, :, bm]

# old-law reference gains
GO = {}
for name,(dB,W,bm,_) in DATA.items():
    c = np.clip((dB+13.78)/82.07, 0, 1)**0.344*4.2
    c_old = np.clip((dB+13.78)/82.07, 0, 1)**0.344*4.2
    g = gains_batch(dB, W, bm, np.array([0.]), np.array([0.]), np.array([99.]), C_pre=c_old)
    GO[name] = float(g[0].mean()) if name=='res_500' else float(np.median(g[0]))

def evaluate(X0s, SLs, CMs, FLs=None):
    out = {}
    for name,(dB,W,bm,_) in DATA.items():
        g = gains_batch(dB, W, bm, X0s, SLs, CMs, FLs=FLs) # [C,T]
        agg = np.sqrt(np.mean(g**2, axis=1)) if name=='res_500' else np.median(g, axis=1)
        out[name] = MEAS_OLD[name] + 20*np.log10(agg/GO[name])
    return out

X0g = np.arange(1.85, 2.35, 0.05); SLg = np.arange(0.065, 0.102, 0.0025); CMg = np.array([99.])
FLg = np.array([0., 0.15, 0.3, 0.45, 0.6])
X0f, SLf, CMf, FLf = [j.ravel() for j in np.meshgrid(X0g, SLg, CMg, FLg, indexing='ij')]
names = list(DATA)

recs = []
CH = 120
for s in range(0, len(X0f), CH):
    sl = slice(s, s+CH)
    ev = evaluate(X0f[sl], SLf[sl], CMf[sl], FLf[sl])
    for j in range(len(X0f[sl])):
        e = {n: ev[n][j] for n in names}
        recs.append((sum(v*v for v in e.values())/len(names),
                     X0f[sl][j], SLf[sl][j], CMf[sl][j], e, FLf[sl][j]))
recs.sort(key=lambda r: r[0])
print('refined top-12:')
for tot, X0, SL, CM, e, FL in recs[:12]:
    print(f'  X0={X0:.2f} S={SL:.4f} FL={FL:.2f} rms={np.sqrt(tot):.3f}  ' +
          ' '.join(f'{n[:5]}:{v:+.2f}' for n,v in e.items()))
MEAS_NEW={'res_500':2.019,'al_12':0.244,'al_24':-0.069,'t1k_1000':-0.321}
ev18=evaluate(np.array([1.8]),np.array([0.11]),np.array([99.]))
print('new(1.8,.11) model-pred vs measured:')
for n in names:
    print(f'  {n:10} pred{ev18[n][0]:+.3f}  meas{MEAS_NEW[n]:+.3f}')