#!/usr/bin/env python3
"""resalpha.py — find detector res-exponent alpha from existing references.

Model: lvl_case(alpha) = lvl_case(alpha=1) * res_case^(1-alpha)
       real cut = L(log2 lvl) with ONE monotone law L shared by all cases.
For an alpha grid -> global LSQ of cut on (1, log2 lvl) -> mean|resid|(alpha).
The alpha minimising residuals is the detector's true res-weighting.

lvl(1), res and am at the tone bin come from RT_DUMP_BIN tracts of the current
build (bare-chain env; detector path identical). comb skipped (multiband).
"""
import os
import subprocess
import sys

import 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 FS
import corpus as _corpus_mod
load = _corpus_mod.load_mono

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_DUMP_BIN': '/tmp/opencode/resalpha_tract.txt'}
OUT = '/tmp/opencode/resalpha_pairs.csv'


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


def ta(x, f, sr=44100):
    x = x[-int(0.75 * sr):].astype(np.float64)
    t = np.arange(len(x)) / sr
    w = 2 * np.pi * f
    return np.hypot(2 * np.sum(x * np.cos(w * t)) / len(x),
                    2 * np.sum(x * np.sin(w * t)) / len(x))


def db(a):
    return 20 * np.log10(max(a, 1e-12))


def main():
    all_cases = []
    for name, inp, args, ref, f in corpus.build_cases():
        joined = [','.join(args)] if len(args) == 3 else args
        all_cases.append((name, inp, joined, ref, f))

    uniq = {}
    case_meta = {c[0]: c for c in all_cases}
    for name, inp, args, ref, f in all_cases:
        uniq.setdefault((inp, tuple(args)), []).append((name, ref, f))

    rows = []
    out_wav = '/tmp/opencode/resalpha_out.wav'
    for (inp, args), items in sorted(uniq.items()):
        tract = '/tmp/opencode/resalpha_tract.txt'
        if os.path.exists(tract):
            os.remove(tract)
        env = {**os.environ, **ENVBASE,
               'RT_DUMP_FRAME': str(120 * max(1, len(args)) - 1)}
        subprocess.run([RB, inp, out_wav] + list(args),
                       capture_output=True, text=True, env=env,
                       cwd='/home/m/re-tools')
        try:
            d = np.loadtxt(tract)
        except Exception as ex:
            print('load fail', items[0][0], ex, flush=True)
            continue
        if d.ndim == 1:
            continue
        k = d[:, 0].astype(int)
        am, res, lvl = d[:, 1], d[:, 2], d[:, 3]
        for name, ref, f in items:
            b = min(bin_of(f), len(lvl) - 1)
            rows.append((name.split('_')[0], name, f, float(lvl[b]),
                         float(res[b]), float(am[b]), inp, ref))
        print(f'{items[0][0]:>18}: dumped', flush=True)

    final = []
    for g, name, f, lvl1, res, am, inp, ref in rows:
        if g == 'comb':
            continue
        cut = db(ta(load(inp), f)) - db(ta(load(ref), f))
        final.append((g, name, f, lvl1, res, cut))

    with open(OUT, 'w') as fh:
        fh.write('group,case,freq,lvl_a1,res,cut_db\n')
        for g, name, f, lvl1, res, cut in final:
            fh.write(f'{g},{name},{f},{lvl1:.9g},{res:.9g},{cut:.6f}\n')
    print(f'\nwrote {len(final)} rows -> {OUT}')

    # ---- alpha scan ----
    lx = np.array([np.log2(max(r[3], 1e-9)) for r in final])
    lr = np.array([np.log2(max(r[4], 1e-12)) for r in final])
    y = np.array([r[5] for r in final])
    gs = np.array([r[0] for r in final])

    print(f'\n{"alpha":>6} {"mean|resid|":>11}  per-group resid mean')
    best = None
    for alpha in np.arange(-0.5, 1.51, 0.05):
        x = lx + (1 - alpha) * lr
        A = np.stack([np.ones_like(x), x], axis=1)
        coef, *_ = np.linalg.lstsq(A, y, rcond=None)
        resid = y - A @ coef
        m = float(np.mean(np.abs(resid)))
        pergrp = ' '.join(
            f'{g_}:{np.mean(resid[gs == g_]):+.2f}' for g_ in ['t1kq', 't1k', 'al', 'res', 'dual'])
        print(f'{alpha:>6.2f} {m:>11.3f}  {pergrp}')
        if best is None or m < best[1]:
            best = (float(alpha), m)
    print(f'\nBEST alpha={best[0]:.2f} mean|resid|={best[1]:.3f}')


if __name__ == '__main__':
    main()