#!/usr/bin/env python3
"""
Quick VLAW parameter grid search - test fewer combos per case.
"""

import numpy as np
import os
import sys
import subprocess
import json

sys.path.insert(0, '/home/m/re-tools/scripts')
import corpus

corpus.RB = '/home/m/re-tools/dsp/build/render48k'

with open('scripts/baseline_bridge.json') as f:
    REF_ERRORS = json.load(f)

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

def run_vlaw(inp, args, alpha, beta, c, delta):
    out = f'/tmp/vlaw_{alpha}_{beta}_{c}_{delta}_{os.path.basename(inp)}.wav'
    env = {
        **os.environ,
        'RT_VLAW': '1',
        'RT_VLAW_ALPHA': str(alpha),
        'RT_VLAW_BETA': str(beta),
        'RT_VLAW_C': str(c),
        'RT_VLAW_DELTA': str(delta),
        'RT_SYN': '1', 'RT_NOWARP': '1', 'RT_NOIIR3': '1', 'RT_IIR12': '0',
    }
    subprocess.run([corpus.RB, inp, out] + args, capture_output=True, env=env, cwd='/home/m/re-tools')
    return out

def eval_error(out, ref, f):
    if not os.path.exists(out) or os.path.getsize(out) == 0:
        return None
    ref_sig = corpus.load_mono(ref)
    out_sig = corpus.load_mono(out)
    min_len = min(len(ref_sig), len(out_sig))
    ref_sig = ref_sig[-min_len:]
    out_sig = out_sig[-min_len:]
    ref_ta = corpus.ta(ref_sig, f)
    out_ta = corpus.ta(out_sig, f)
    return corpus.db(out_ta / ref_ta)

def group_key(name):
    return name.split('_')[0]

all_cases = structural_cases()
groups = {}
for name, inp, args, ref, f in all_cases:
    g = group_key(name)
    groups.setdefault(g, []).append((name, inp, args, ref, f))

# Pick one case per group
rep_cases = {}
for g in ['t1kq', 't1k', 'al', 'res', 'dual']:
    if g in groups:
        # Pick middle-ish case
        cases = groups[g]
        rep_cases[g] = cases[len(cases)//2]

print("Representative cases:")
for g, (name, inp, args, ref, f) in rep_cases.items():
    print(f"  {g}: {name}")

# Test a small grid around dual params
dual_params = (3.2193, 0.4927, 0.5423, 6.9177)

print("\n=== Grid search per group ===")
results = {}

for g, (name, inp, args, ref, f) in rep_cases.items():
    print(f"\n--- {g} ({name}) ---")
    best = None
    best_err = float('inf')
    
    # Coarse grid
    alphas = np.linspace(1.0, 5.0, 5)
    betas = np.linspace(0.2, 0.8, 5)
    cs = np.linspace(-0.5, 2.0, 5)
    deltas = np.linspace(0.0, 12.0, 5)
    
    for alpha in alphas:
        for beta in betas:
            for c in cs:
                for delta in deltas:
                    out = run_vlaw(inp, args, alpha, beta, c, delta)
                    err = eval_error(out, ref, f)
                    if err is not None and abs(err) < best_err:
                        best_err = abs(err)
                        best = (alpha, beta, c, delta, err)
                        print(f'  {name}: α={alpha:.3f}, β={beta:.3f}, c={c:.3f}, Δ={delta:.3f} => {err:.3f} dB')
    
    if best:
        results[g] = best[:4]
        print(f'  BEST {g}: α={best[0]:.4f}, β={best[1]:.4f}, c={best[2]:.4f}, Δ={best[3]:.4f} => {best[4]:.3f} dB')

print("\n=== SUMMARY ===")
for g, (a, b, c, d) in results.items():
    print(f'{g}: α={a:.4f}, β={b:.4f}, c={c:.4f}, Δ={d:.4f}')

with open('/tmp/opencode/vlaw_params.json', 'w') as f:
    json.dump({g: {'alpha': a, 'beta': b, 'c': c, 'delta': d} 
               for g, (a, b, c, d) in results.items()}, f, indent=2)
print('\nSaved to /tmp/opencode/vlaw_params.json')