#!/usr/bin/env python3
"""lawfit22r.py — offline affine-law fitting infra (NOTES_LEVEL 22r NEXT-1).

The detector path is law-independent (RT_DUMP_ALL in framed_model.cpp), so ONE
trajectory capture per unique render suffices to evaluate ANY affine dB law
    cut(lvl) = K * (A_db + S_db * log2(lvl)),  K = 20*log10(2)/6.0174
offline: mask[frame][band][bin] = exp2(-(A+S*log2 lvl)/6.0174) replayed through
an exact numpy replica of render48k.cpp + spectral.cpp (resample 44.1<->48,
Hann STFT 4096/1024 @48k, pointwise mask, WOLA istft, FULL-BLK chunking with
zero-padded tail = 61 hop-frames per 65536-block) + corpus metric.

NOTE (22s): pointwise per-bin fitting is INVALID for out-of-band evals — the
detector level at far bins is transient during the metric window (slow twin
adaptation, NOTES 22n) and Hann leakage couples neighbouring bins.

Modes:
  collect              capture lvl trajectories -> /tmp/opencode/lawfit_traj.npz
  sanity A S JSON      full-sim group table vs an actual corpus run
  fit A0 S0            coordinate-descent (A,S) on TOTAL + per-group optima

Bare-chain env (NOTES 22k, matches 22r baseline TOTAL 1.931 @ A=7.4,S=1.85):
  RT_LUT_OFF=1 RT_IIR12=0 RT_NOWARP=1 RT_NOBLEND=1 RT_NOIIR3=1   (no FLOOR)
"""
import json
import os
import struct
import subprocess
import sys
import time

import numpy as np
from scipy.signal import resample_poly

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import corpus

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

NFFT = 4096
HOP = 1024
RBIN = NFFT // 2 + 1
BLK = 1 << 16
NF_BLK = (BLK - NFFT) // HOP + 1  # 61 hop-frames per block
TRAJ = '/tmp/opencode/lawfit_traj.npz'
K_DB = 20.0 * np.log10(2.0) / 6.0174  # exp2(-y) dB factor (informational)

BASE_ENV = {
    'RT_LUT_OFF': '1', 'RT_IIR12': '0', 'RT_NOWARP': '1',
    'RT_NOBLEND': '1', 'RT_NOIIR3': '1',
}

_WIN = 0.5 * (1.0 - np.cos(2.0 * np.pi * np.arange(NFFT) / (NFFT - 1)))
_WOLA = float(np.sum(_WIN * _WIN) / HOP)


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 n_bands(argstr_list):
    if len(argstr_list) == 1:
        return max(1, len(argstr_list[0].split(',')) // 3)
    return len(argstr_list)


def read_traj(path):
    recs = []
    with open(path, 'rb') as fh:
        while True:
            hdr = fh.read(8)
            if len(hdr) < 8:
                break
            fr, nb = struct.unpack('<ii', hdr)
            recs.append(np.frombuffer(fh.read(4 * nb), dtype='<f4').astype(np.float32))
    if not recs:
        return None
    return np.stack(recs)


def collect():
    cases = structural_cases()
    uniq = {}
    for name, inp, args, ref, f in cases:
        key = f'{os.path.basename(inp)}|{";".join(args)}'
        uniq.setdefault(key, {'inp': inp, 'args': list(args),
                              'nb': n_bands(list(args))})

    store = {}
    traj_bin = '/tmp/opencode/lawfit_traj.bin'
    resp_bin = '/tmp/opencode/lawfit_res.bin'
    for key, u in sorted(uniq.items()):
        if os.path.exists(traj_bin):
            os.remove(traj_bin)
        if os.path.exists(resp_bin):
            os.remove(resp_bin)
        env = {**os.environ, **BASE_ENV, 'RT_DUMP_ALL': traj_bin,
               'RT_DUMPRESPATH': resp_bin}
        subprocess.run([corpus.RB, u['inp'], '/tmp/opencode/lawfit_out.wav'] + u['args'],
                       capture_output=True, text=True, env=env,
                       cwd='/home/m/re-tools')
        arr = read_traj(traj_bin)
        if arr is None or arr.shape[1] != RBIN or arr.shape[0] % u['nb']:
            print(f'NO/BAD TRAJ: {key}: {None if arr is None else arr.shape}')
            continue
        lvl = arr.reshape(-1, u['nb'], RBIN)  # [frame][band][bin]
        store[key] = lvl
        # static twin-response spectra per band ({i32 band, i32 nbin, f32[nbin]})
        res = None
        if os.path.exists(resp_bin):
            raw = open(resp_bin, 'rb').read()
            off = 0
            bands = []
            while off < len(raw):
                bi, nb = struct.unpack_from('<ii', raw, off)
                off += 8
                r = np.frombuffer(raw, dtype='<f4', count=nb, offset=off)
                off += 4 * nb
                bands.append(r.astype(np.float32))
            if len(bands) == u['nb'] and all(b.size == RBIN for b in bands):
                res = np.stack(bands)
        if res is not None:
            store[key + '|res'] = res
        print(f'{key}: {lvl.shape[0]} frames x {u["nb"]} bands'
              + (' +res' if res is not None else ' NORES'))
    np.savez_compressed(TRAJ, **store)
    print(f'\nwrote {len(store)} entries -> {TRAJ}')


def load_traj():
    z = np.load(TRAJ)
    return {k: z[k] for k in z.files}


# ---------------- exact render48k/spectral replay ----------------

class CaseInput:
    """Law-independent per-render data: resampled+blocked input, frame count."""

    def __init__(self, inp_path):
        self.x44 = corpus.load_mono(inp_path)
        x48 = resample_poly(self.x44, 160, 147)
        self.blocks = [(s, min(BLK, len(x48) - s)) for s in range(0, len(x48), BLK)]
        self.nf = len(self.blocks) * NF_BLK
        need = self.blocks[-1][0] + BLK
        self.xext = np.zeros(need)
        self.xext[:len(x48)] = x48
        self.offs = np.array([b[0] + f * HOP
                              for b in self.blocks for f in range(NF_BLK)],
                             dtype=np.int64)

    def frames_ok(self, lvl):
        return lvl.shape[0] == self.nf


def masks_from_lvl(lvl, nb, A, S, Q=0.0, res=None, rp=0.0):
    """Per-frame lower-half mask exactly as bare chain + LAWAFFINE produces.

    Law families (22t): scalar cut=A+S*log2(lvl); quad adds curvature Q*x^2;
    resrp multiplies per-band mask by res^rp (decomp warp-cascade factor,
    applied BEFORE cross-band min like the C++ warp section).
    """
    nf = lvl.shape[0]
    lv64 = lvl.astype(np.float64)
    x = np.log2(np.maximum(lv64, 1e-12))
    mm = np.exp2(-((A + S * x + Q * x * x) / 6.0174))
    low = lv64 <= 1e-6                       # C++ fallback: exp2(-level)
    if low.any():
        mm[low] = np.exp2(-lv64[low])
    if res is not None and rp:
        mm = mm * res[None].astype(np.float64) ** rp
    if nb > 1:
        lo = np.min(mm, axis=1)                 # min across bands
    else:
        lo = mm[:, 0, :]
    full = np.ones((nf, NFFT))
    full[:, :RBIN] = lo
    full[:, RBIN:] = lo[:, 1:RBIN - 1][:, ::-1]  # C++ 418: mask[k]=mask[nfft-k]
    return full


def replay(ci, lvl, nb, A, S, Q=0.0, res=None, rp=0.0):
    """Return trimmed 44.1k output for law params on prepared CaseInput ci."""
    assert ci.frames_ok(lvl), f'traj {lvl.shape[0]} != frames {ci.nf}'
    full = masks_from_lvl(lvl, nb, A, S, Q, res, rp)
    segs = ci.xext[ci.offs[:, None] + np.arange(NFFT)[None, :]] * _WIN[None, :]
    yspec = np.fft.fft(segs, axis=1) * full
    td = np.real(np.fft.ifft(yspec, axis=1))
    del segs, yspec
    td *= _WIN[None, :]
    L = len(ci.x44)
    y48 = np.zeros(L)
    overlap = np.zeros(NFFT)
    fi = 0
    zeros = np.zeros(HOP)
    for s, n in ci.blocks:
        for f in range(NF_BLK):
            off = s + f * HOP
            acc = overlap + td[fi]
            fi += 1
            e = min(off + HOP, L)
            if e > off:
                y48[off:e] = acc[:e - off] / _WOLA
            overlap[:NFFT - HOP] = acc[HOP:]
            overlap[NFFT - HOP:] = zeros
    y44 = resample_poly(y48, 147, 160)
    return y44[:L]


# ---------------- evaluation ----------------

_CI_CACHE = {}


def case_input(key, inp):
    if key not in _CI_CACHE:
        _CI_CACHE[key] = CaseInput(inp)
    return _CI_CACHE[key]


def build_eval_index(trajs):
    idx = []
    for name, inp, args, ref, f in structural_cases():
        key = f'{os.path.basename(inp)}|{";".join(args)}'
        if key not in trajs:
            continue
        idx.append({'name': name, 'key': key, 'inp': inp, 'ref': ref,
                    'f': f, 'nb': trajs[key].shape[1], 'grp': name.split('_')[0]})
    return idx


def sim_errors(trajs, idx, A, S, Q=0.0, rp=0.0):
    errs = {}
    ycache = {}
    for e in idx:
        ck = (e['key'], e['nb'], round(Q, 6), round(rp, 6))
        if ck not in ycache:
            ci = case_input(e['key'], e['inp'])
            ycache[ck] = replay(ci, trajs[e['key']], e['nb'], A, S,
                                Q, trajs.get(e['key'] + '|res'), rp)
        y44 = ycache[ck]
        errs[e['name']] = corpus.db(corpus.ta(y44, e['f'])) - \
            corpus.db(corpus.ta(corpus.load_mono(e['ref']), e['f']))
    return errs


def group_stats(errs):
    gs = {}
    for k, v in errs.items():
        gs.setdefault(k.split('_')[0], []).append(v)
    out = {g: float(np.mean(np.abs(v))) for g, v in gs.items()}
    out['TOTAL'] = float(np.mean(np.abs(list(errs.values()))))
    return out


def sanity(A, S, json_path):
    trajs = load_traj()
    idx = build_eval_index(trajs)
    refs = json.load(open(json_path))
    t0 = time.time()
    errs = sim_errors(trajs, idx, A, S)
    st = group_stats(errs)
    ast = group_stats(refs)
    print(f'{"group":>8} {"sim":>8} {"actual":>8} {"d":>7}   ({time.time()-t0:.1f}s)')
    for g in ['t1kq', 't1k', 'al', 'res', 'dual', 'comb', 'TOTAL']:
        print(f'{g:>8} {st[g]:>8.3f} {ast[g]:>8.3f} {st[g]-ast[g]:>+7.3f}')
    print('\nper-case worst deltas:')
    deltas = sorted(((abs(errs[k] - refs[k]), k) for k in errs),
                    reverse=True)[:8]
    for d, k in deltas:
        print(f'  {k:>18}: sim {errs[k]:+8.3f} actual {refs[k]:+8.3f} d {errs[k]-refs[k]:+.3f}')


def fit(A0, S0):
    trajs = load_traj()
    idx = build_eval_index(trajs)
    fit2(trajs, idx, A0, S0)


def fit2(trajs, idx, A0, S0):
    """Law-family comparison: scalar / quad / resrp / quad+resrp (NOTES 22t)."""

    def objective(p, sub=None):
        A, S, Q, rp = p
        ii = idx if sub is None else sub
        return group_stats(sim_errors(trajs, ii, A, S, Q, rp))

    def descend(p0, steps, sub=None, label=''):
        best = list(p0)
        bst = objective(tuple(best), sub)
        print(f'{label} start {[round(v, 4) for v in best]}: '
              f'TOTAL={bst["TOTAL"]:.3f}', flush=True)
        while all(s > 1e-4 for s in steps.values()):
            moved = False
            for i, nm in enumerate(['A', 'S', 'Q', 'rp']):
                if nm not in steps:
                    continue
                st = steps[nm]
                for d in (+st, -st):
                    cand = list(best)
                    cand[i] = round(cand[i] + d, 6)
                    if cand[3] < 0 or (nm == 'Q' and abs(cand[2]) > 3):
                        continue
                    s2 = objective(tuple(cand), sub)
                    if s2['TOTAL'] < bst['TOTAL'] - 1e-4:
                        best, bst = cand, s2
                        moved = True
                        print(f'  {label} {nm}{d:+.4g}: TOTAL={s2["TOTAL"]:.3f} '
                              f'{[round(v, 4) for v in best]}', flush=True)
            if not moved:
                for nm in steps:
                    steps[nm] /= 2
        return best, bst

    res_all = {}

    b, t = descend([A0, S0, 0.0, 0.0], {'A': 0.4, 'S': 0.15}, label='[scalar]')
    res_all['scalar'] = (list(b), dict(t))

    bA, bS, _, _ = res_all['scalar'][0]
    b, t = descend([bA, bS, 0.0, 0.0], {'A': 0.3, 'S': 0.15, 'Q': 0.06},
                   label='[quad]')
    res_all['quad'] = (list(b), dict(t))

    b, t = descend([bA, bS, 0.0, 0.03], {'A': 0.3, 'S': 0.15, 'rp': 0.01},
                   label='[resrp]')
    res_all['resrp'] = (list(b), dict(t))

    bA, bS, bQ, _ = res_all['quad'][0]
    b, t = descend([bA, bS, bQ, 0.03], {'A': 0.25, 'S': 0.12, 'Q': 0.05,
                                        'rp': 0.008}, label='[quad+resrp]')
    res_all['quad+resrp'] = (list(b), dict(t))

    print('\n================ LAW FAMILY SUMMARY ================')
    for fam, (p, st) in res_all.items():
        gs = ' '.join(f'{g}={st[g]:.3f}' for g in
                      ['t1kq', 't1k', 'al', 'res', 'dual', 'comb'])
        print(f'{fam:>12}: A={p[0]:7.3f} S={p[1]:6.3f} Q={p[2]:+6.3f} '
              f'rp={p[3]:5.3f}  TOTAL={st["TOTAL"]:.3f}\n{"":>14}{gs}')
    json.dump({f: {'params': p, 'groups': s} for f, (p, s) in res_all.items()},
              open('/tmp/opencode/lawfit_fit2.json', 'w'), indent=1)
    print('\nwrote /tmp/opencode/lawfit_fit2.json')


def freq_of(name):
    for nm, inp, args, ref, f in structural_cases():
        if nm == name:
            return f
    raise KeyError(name)


def percase():
    """Per-render greedy (A,S) optima -> /tmp/opencode/lawfit_percase.json."""
    trajs = load_traj()
    idx = build_eval_index(trajs)
    renders = {}
    for e in idx:
        renders.setdefault(e['key'], {'nb': e['nb'], 'evals': []})['evals'].append(e)

    out = {}
    for key, r in sorted(renders.items()):
        lvl = trajs[key]
        ci = case_input(key, r['evals'][0]['inp'])

        def ev(A, S):
            y44 = replay(ci, lvl, r['nb'], A, S)
            return [corpus.db(corpus.ta(y44, e['f'])) -
                    corpus.db(corpus.ta(corpus.load_mono(e['ref']), e['f']))
                    for e in r['evals']]

        bA, bS, bval = None, None, 1e9
        grid = [(float(A), float(S))
                for A in np.arange(4.0, 11.01, 0.75)
                for S in np.arange(0.25, 3.01, 0.25)]
        for A, S in grid:
            errs = ev(A, S)
            m = float(np.mean(np.abs(errs)))
            if m < bval:
                bA, bS, bval = A, S, m
        # band params from args string(s)
        bands = []
        for a in r['evals'][0]['inp'] and key.split('|')[1].split(';'):
            p = a.split(',')
            bands.append({'fc': float(p[0]), 'q': float(p[1]), 'sens': float(p[2])})
        out[key] = {
            'group': sorted({e['grp'] for e in r['evals']}),
            'bands': bands, 'best_A': bA, 'best_S': bS, 'best_err': bval,
            'evals': [{'name': e['name'], 'freq': e['f']} for e in r['evals']],
        }
        print(f'{key:>44}: A={bA:5.2f} S={bS:5.2f}  mean|e|={bval:.3f}', flush=True)
    json.dump(out, open('/tmp/opencode/lawfit_percase.json', 'w'), indent=1)
    print('\nwrote /tmp/opencode/lawfit_percase.json')


def _bisect_A(ev, f_idx, S, lo=0.0, hi=16.0, iters=11):
    """Find A s.t. signed err at eval f_idx == 0 (monotone decreasing in A)."""
    def e(A):
        return ev(A, S)[f_idx]
    elo, ehi = e(lo), e(hi)
    if elo <= 0:
        return lo
    if ehi >= 0:
        return hi
    for _ in range(iters):
        mid = 0.5 * (lo + hi)
        if e(mid) > 0:
            lo = mid
        else:
            hi = mid
    return 0.5 * (lo + hi)


def lines():
    """Per-render A*(S) at fixed S anchors -> /tmp/opencode/lawfit_lines.json."""
    trajs = load_traj()
    idx = build_eval_index(trajs)
    renders = {}
    for e in idx:
        renders.setdefault(e['key'], {'nb': e['nb'], 'evals': []})['evals'].append(e)

    S_ANCHORS = [0.5, 1.5, 2.5]
    out = {}
    for key, r in sorted(renders.items()):
        lvl = trajs[key]
        ci = case_input(key, r['evals'][0]['inp'])
        rms = float(np.sqrt(np.mean(ci.x44 ** 2)))
        rec = {'group': sorted({e['grp'] for e in r['evals']}),
               'bands': key.split('|')[1], 'rms_db': 20 * np.log10(max(rms, 1e-9)),
               'A_star': {}}
        if len(r['evals']) == 1:
            def ev(A, S):
                y44 = replay(ci, lvl, r['nb'], A, S)
                return [corpus.db(corpus.ta(y44, r['evals'][0]['f'])) -
                        corpus.db(corpus.ta(corpus.load_mono(r['evals'][0]['ref']),
                                            r['evals'][0]['f']))]
            for S in S_ANCHORS:
                rec['A_star'][S] = round(_bisect_A(ev, 0, S), 3)
            rec['err_at_Astar'] = round(abs(ev(rec['A_star'][1.5], 1.5)[0]), 4)
        else:
            # multi-eval render: minimise mean|err| per S anchor (grid+refine)
            refs = [corpus.db(corpus.ta(corpus.load_mono(e['ref']), e['f']))
                    for e in r['evals']]
            for S in S_ANCHORS:
                best = (None, 1e9)
                for A in np.arange(0, 16.01, 0.5):
                    y44 = replay(ci, lvl, r['nb'], A, S)
                    m = float(np.mean([abs(corpus.db(corpus.ta(y44, e['f'])) - rf)
                                       for e, rf in zip(r['evals'], refs)]))
                    if m < best[1]:
                        best = (float(A), m)
                rec['A_star'][S] = round(best[0], 3)
                rec.setdefault('multi_err', {})[S] = round(best[1], 4)
        out[key] = rec
        extra = f" multi={rec.get('multi_err', {}).get(1.5)}" if 'multi_err' in rec else ''
        print(f'{key:>58}: A*=' +
              ','.join(f'{rec["A_star"][S]:6.2f}' for S in S_ANCHORS) +
              f'  rms={rec["rms_db"]:6.1f}{extra}', flush=True)
    json.dump(out, open('/tmp/opencode/lawfit_lines.json', 'w'), indent=1)
    print('\nwrote /tmp/opencode/lawfit_lines.json')


if __name__ == '__main__':
    if not sys.argv[1:]:
        print(__doc__)
        sys.exit(1)
    cmd = sys.argv[1]
    if cmd == 'collect':
        collect()
    elif cmd == 'sanity':
        sanity(float(sys.argv[2]), float(sys.argv[3]), sys.argv[4])
    elif cmd == 'fit':
        fit(float(sys.argv[2]), float(sys.argv[3]))
    elif cmd == 'percase':
        percase()
    elif cmd == 'lines':
        lines()
    else:
        print(f'unknown mode {cmd}')
        sys.exit(1)