#!/usr/bin/env python3
"""pairfit22q.py — collect (model lvl_raw, real cut_dB) pairs across corpus."""
import sys, os, subprocess, numpy as np
sys.path.insert(0, "/home/m/re-tools/scripts")
sys.path.insert(0, "/home/m/re-tools/handoff")
import corpus
from render_parity import load, FS

RB = "/home/m/re-tools/dsp/build/render48k"
ENVBASE = {"RT_LUT_OFF": "1", "RT_IIR12": "0", "RT_NOWARP": "1",
           "RT_NOBLEND": "1", "RT_NOIIR3": "1", "RT_POOL": "0",
           "RT_SCALE_M": "1.0", "RT_FLOOR": "0",
           "RT_DUMP_BIN": "/tmp/opencode/tract_pair.txt",
           "RT_DUMP_FRAME": "120"}

def amp(p, f):
    a = load(p)
    seg = np.mean(a[:min(len(a), int(3.5*FS))][-int(0.75*FS):], axis=1)
    n = len(seg); tt = np.arange(n)/FS
    return 2*np.abs(np.dot(seg, np.exp(-2j*np.pi*f*tt)))/n

def bin_of(f):
    return int(round(f/48000*4096))

def main():
    e = {**os.environ, **ENVBASE}
    pairs = []
    seen_render = {}
    cases = corpus.build_cases()
    # unique renders first (input,args): dump tract once, remember lvl per freq
    uniq = {}
    for name, inp, args, ref, f in cases:
        key = (inp, ",".join(args) if isinstance(args, list) else args)
        uniq.setdefault(key, []).append((name, ref, f))
    n_done = 0
    for (inp, argstr), items in sorted(uniq.items()):
        out_wav = "/tmp/opencode/pf_model.wav"
        env = {**e}
        bands = argstr.split(",") if "," in argstr else [argstr]
        nbands = max(1, len(bands)//3)
        env["RT_DUMP_FRAME"] = str(120 * nbands - 1)  # land mid-frame of last band
        r = subprocess.run([RB, inp, out_wav] + ([argstr] if "," in argstr else [" ".join([])]),
                           capture_output=True, text=True, env=env, cwd="/home/m/re-tools")
        if not os.path.exists("/tmp/opencode/tract_pair.txt"):
            print("no tract for", argstr); continue
        try:
            d = np.loadtxt("/tmp/opencode/tract_pair.txt")
        except Exception as ex:
            print("load fail", ex); continue
        if d.ndim == 1: continue
        k, am, res, lvl = d.T[0], d.T[1], d.T[2], d.T[3]
        for name, ref, f in items:
            b = bin_of(f)
            if b >= len(lvl): continue
            lv = float(lvl[b])
            i_db = 20*np.log10(max(amp(inp, f), 1e-12))
            r_db = 20*np.log10(max(amp(ref, f), 1e-12))
            cut = i_db - r_db
            pairs.append((name, f, lv, cut))
        n_done += 1
        os.remove("/tmp/opencode/tract_pair.txt")
    print("collected %d pairs from %d renders" % (len(pairs), n_done))
    with open("/tmp/opencode/pairs.csv", "w") as fh:
        fh.write("case,freq,lvl,cut_db\n")
        for nm, f, lv, c in pairs:
            fh.write("%s,%.1f,%.4f,%.3f\n" % (nm, f, lv, c))
    arr = np.array([(lv, c) for _, _, lv, c in pairs])
    o = np.argsort(arr[:, 0]); s = arr[o]
    print("lvl -> cut samples:")
    step = max(1, len(s)//18)
    for i in range(0, len(s), step):
        print("  %8.3f -> %+8.2f dB" % tuple(s[i]))

if __name__ == "__main__":
    main()