#!/usr/bin/env python3
"""Full-corpus honest metric harness for soothe2-re.

Protects against the Phase B regression: when re-transcribing the structural
FUN_180529fe0 chain (BITEXACT_PLAN step 1), each build must NOT regress below the
committed bridge baseline. This harness:

  * runs every dataset case through dsp/build/framed_test,
  * measures trimmed 24-bit tone-amplitude error (dB) vs reference wavs,
  * writes the results to a JSON file (-> commit as the bridge baseline),
  * with --compare <base.json> prints per-group deltas and exits non-zero if the
    max per-group degradation exceeds --tol (default 0.25 dB).

Usage:
  python3 scripts/corpus.py [--out results.json] [TAG]
  python3 scripts/corpus.py --compare baseline.json [--tol 0.25]

Metric (must match AGENTS.md / render_parity.py):
  - 24-bit refs decoded with x>=0x800000 => x-0x1000000 (NOT OR-0xFF000000).
  - Goertzel steady-state tone amplitude, last 0.75 s window, trimmed length.
"""
import json, os, subprocess, sys
import numpy as np
import wave

RB = '/home/m/re-tools/dsp/build/framed_test'
SB = '/home/m/soothe-bt'


def load_mono(p):
    w = wave.open(p, 'rb'); n = w.getnframes(); d = w.readframes(n)
    ch = w.getnchannels(); b = w.getsampwidth()
    if b == 2:
        return np.frombuffer(d, dtype=np.int16).astype(np.float64).reshape(-1, ch).mean(1) / 32768
    # 24-bit: handle misaligned data (extra bytes from bext/junk chunks)
    n_samples = len(d) // (ch * 3)
    d = d[:n_samples * ch * 3]
    raw = np.frombuffer(d, dtype=np.uint8).reshape(-1, ch, 3)
    s = raw[:, :, 0].astype(np.int64) | (raw[:, :, 1].astype(np.int64) << 8) | (raw[:, :, 2].astype(np.int64) << 16)
    return np.where(s >= 0x800000, s - 0x1000000, s).mean(1).astype(np.float64) / 8388608.0


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 run(inp, out, args):
    subprocess.run([RB, inp, out] + args, capture_output=True)


def tone_err(out, ref, f):
    return db(ta(load_mono(out), f) / ta(load_mono(ref), f))


# (name, input, [band args], ref, tone)
def build_cases():
    cases = []
    for fc in ['800', '900', '950', '1000', '1050', '1100', '1200']:
        cases.append((f't1kq_{fc}', f'{SB}/tone1kq.wav', [f'{fc},0.99999785,12'],
                      f'{SB}/t1kq_only1_{fc}.wav', 1000.0))
    for fc in ['500', '678.7611083984375', '800', '900', '950', '1000', '1050',
               '1100', '1200', '1500', '2000']:
        cases.append((f't1k_{fc}', f'{SB}/tone1k.wav', [f'{fc},1.0,12'],
                      f'{SB}/t1k_b1f_{fc}.wav', 1000.0))
    for lv in ['3', '6', '9', '12', '18', '24']:
        cases.append((f'al_{lv}', f'{SB}/lvl_tone_lv{lv}.wav', ['1000,1.0,12'],
                      f'{SB}/al_{lv}.wav', 1000.0))
    for fc in ['300', '400', '450', '475', '490', '500', '510', '525', '550', '600', '700']:
        cases.append((f'res_{fc}', f'{SB}/resonant.wav', [f'{fc},1.0,12'],
                      f'{SB}/res_only1_{fc}.wav', 1000.0))
    for q in ['0.1', '0.2', '0.3', '0.5', '0.7', '1.0', '1.5', '2.0', '3.0', '5.0', '10.0']:
        cases.append((f'dual_{q}_500', f'{SB}/dual.wav', ['500', q, '12'],
                      f'{SB}/dual_b1q_{q}.wav', 500.0))
        cases.append((f'dual_{q}_2000', f'{SB}/dual.wav', ['500', q, '12'],
                      f'{SB}/dual_b1q_{q}.wav', 2000.0))
    comb = ['1000,1.0,12', '1778.7,4.5,-12', '195.1,0.49,2.24', '13408,1,12']
    for fc, tone in [('500', 500.0), ('1000', 1000.0), ('1500', 1500.0),
                     ('2000', 2000.0), ('3000', 3000.0)]:
        cases.append((f'comb_{fc}', f'{SB}/comb.wav', comb, f'{SB}/comb_ref.wav', tone))
    return cases


def group_stats(res):
    groups = {}
    for k, v in res.items():
        g = k.split('_')[0]
        groups.setdefault(g, []).append(v)
    out = {}
    for g, vs in groups.items():
        out[g] = {'n': len(vs), 'mean_abs': float(np.mean(np.abs(vs))),
                  'max_abs': float(np.max(np.abs(vs)))}
    allv = list(res.values())
    out['TOTAL'] = {'n': len(allv), 'mean_abs': float(np.mean(np.abs(allv))),
                    'max_abs': float(np.max(np.abs(allv)))}
    return out


def main():
    args = sys.argv[1:]
    compare = None; out_path = None; tol = 0.25
    for i in range(len(args)):
        if args[i] == '--compare': compare = args[i + 1]
        elif args[i] == '--out': out_path = args[i + 1]
        elif args[i] == '--tol': tol = float(args[i + 1])

    if compare:
        base = json.load(open(compare))
        bstats = group_stats(base)
        # re-run current build through the harness
        res, _ = run_all()
        print(f'{"group":>12} {"n":>3} {"base_mean":>9} {"cur_mean":>9} {"d":>9} {"max_d":>8}')
        worst = 0.0
        for g in sorted(set(bstats) | set(group_stats(res))):
            bs = bstats.get(g, {'n': 0, 'mean_abs': 0.0})
            cs = group_stats(res).get(g, {'n': 0, 'mean_abs': 0.0})
            d = cs['mean_abs'] - bs['mean_abs']
            worst = max(worst, d)
            print(f'{g:>12} {cs["n"]:>3} {bs["mean_abs"]:>9.3f} {cs["mean_abs"]:>9.3f} '
                  f'{d:>+9.3f} {cs["max_abs"]:>8.3f}')
        print(f'\nmax per-group mean|err| degradation vs baseline: {worst:+.3f} dB (tol {tol})')
        return 1 if worst > tol else 0

    res, _ = run_all()
    stats = group_stats(res)
    print(f'{"group":>12} {"n":>3} {"mean|e|":>8} {"max":>7}')
    for g in ['t1kq', 't1k', 'al', 'res', 'dual', 'comb', 'TOTAL']:
        if g in stats:
            s = stats[g]
            print(f'{g:>12} {s["n"]:>3} {s["mean_abs"]:>8.3f} {s["max_abs"]:>7.3f}')
    if out_path:
        json.dump(res, open(out_path, 'w'), indent=2, sort_keys=True)
        print(f'\nwrote {out_path}')
    return 0


def run_all():
    res = {}
    cases = build_cases()
    for name, inp, args, ref, f in cases:
        out = f'/tmp/corpus_{name}.wav'
        run(inp, out, args)
        res[name] = tone_err(out, ref, f)
    return res, res


if __name__ == '__main__':
    sys.exit(main())