#!/usr/bin/env python3
"""Phase B step 1: spatial max-pooling hypothesis scan (offline, no rebuilds).

lvl'(t,k) = pool(lvl)(t,k) with width w, then law -> chain -> median/rms gain
at metric bin. Variants: pool_am (pool raw am then /res), pool_lvl, pool_db.
"""
import re
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[:, 1], d[:, 2], d[:, 6]  # am,res,W
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)

def iir_fwd_m(x, A, B):
    """x [C,T,N] vectorized over C,T; sequential over N."""
    C, T, N = x.shape
    acc = np.zeros((C, T)); y = np.empty_like(x)
    for i in range(N):
        acc = A[i]*acc + B[i]*x[:, :, i]
        y[:, :, i] = acc
    return y

def gains_batch(c, W, bm):
    y = iir_fwd_m(c, A1, B1)
    y = iir_fwd_m(y, A2, B2)
    y = 0.8*np.exp2(-y) * W[None, None, :]
    for _ in range(2):
        st = np.zeros(y.shape[:2])
        for i in range(y.shape[2]):
            st = y[:, :, i]*B3[i] + st*A3[i]; y[:, :, i] = st
        st = y[:, :, -1].copy()
        for i in range(y.shape[2]-2, 0, -1):
            st = y[:, :, i]*B3[i] + st*A3[i]; y[:, :, i] = st
    return y[:, :, bm]

def slide_max(x, w):
    """sliding max over last axis, width w (odd), 'same' edges."""
    if w <= 1: return x.copy()
    h = w//2
    xp = np.pad(x, ((0,0),(0,0),(h,h)), mode='edge')
    win = np.lib.stride_tricks.sliding_window_view(xp, w, axis=2)
    return win.max(axis=-1)

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}
SCALE = 15.0 * 440.95 / 2048.0

CASES = {}
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)]:
    am, res_k, W = load_dump(DUMPDIR[dump]+dump)
    T = load_traj('/tmp/opencode/'+traj)[WIN[name][0]:WIN[name][1]]
    CASES[name] = (T, res_k, W, bm)

LAW = dict(new=lambda dB: np.maximum(1.8+0.11*dB, 0))
OLD = lambda dB: np.clip((dB+13.78)/82.07, 0, 1)**0.344*4.2

def agg(g, name):
    return np.sqrt(np.mean(g**2)) if name == 'res_500' else np.median(g)

print(f'{"variant":>18} {"w":>3} ' + ' '.join(f'{n:>9}' for n in CASES) + '   (pred err, dB)')
# baselines on CORRECT lvl (traj stores lvl_raw already)
GBASE = {}
for name,(T,res_k,W,bm) in CASES.items():
    dB = np.log10(np.maximum(T, 1e-12))
    GBASE[name] = agg(gains_batch(OLD(dB)[None], W, bm)[0], name)
print('sanity new@w=1 (vs validated):')
row=[]
for name,(T,res_k,W,bm) in CASES.items():
    dB = np.log10(np.maximum(T, 1e-12))
    gn = agg(gains_batch(LAW['new'](dB)[None], W, bm)[0], name)
    row.append(MEAS_OLD[name] + 20*np.log10(gn/GBASE[name]))
print(f'{"new":>18} {1:>3} ' + ' '.join(f'{v:+9.2f}' for v in row))

for variant in ['pool_lvl', 'pool_am', 'pool_db']:
    for w in [3, 5, 9, 17]:
        row = []
        for name,(T,res_k,W,bm) in CASES.items():
            am = T  # stored lvl_raw = am/res*scale -> recover am = lvl/res*scale... careful
            # stored lvl_raw = am/res_k * SCALE => am = lvl_raw * res_k / SCALE
            am_abs = T * res_k[None,:] / SCALE
            if variant == 'pool_am':
                lv = slide_max(am_abs[None], w)[0] / res_k[None,:] * SCALE
            elif variant == 'pool_lvl':
                lv = slide_max(T[None], w)[0]
            else:
                db_ = np.log10(np.maximum(T, 1e-12))*20
                lv = 10**(slide_max(db_[None], w)[0]/20)
            dB = np.log10(np.maximum(lv, 1e-12))*20
            g = gains_batch(LAW['new'](dB)[None], W, bm)
            row.append(MEAS_OLD[name] + 20*np.log10(agg(g[0],name)/GBASE[name]))
        print(f'{variant:>18} {w:>3} ' + ' '.join(f'{v:+9.2f}' for v in row))