#!/usr/bin/env python3
"""detector_cascade.py — validated simulator of the soothe2 detector cascade (529c60).

Decoded from assembly (2026-08-25):
  Phase 1: |z_i| via 16140 (vrsqrtps+vsqrtps — magnitude, NOT squared)
  Phase 2: Haar smoothing kernel [0.25, 0.5, 0.25], ctx[0x1b0] iterations
  Phase 3: peak→sin-mod→max-clamp→ratio→pow→log→FMA-blend→memcpy

Validated on chain_samples.pkl (2-frame ptrace capture):
  - op A output matches |z| (max diff 5.4e-6)
  - 2 Haar iterations + scalar blend: rms=0.30, corr=0.998 vs COUT
  - ctx[0x1b0]=2 (Haar iterations) — derived from best-fit

Unknowns (require live capture):
  - ctx[0x54087c] — sin modulation parameter (controls sin_peak clamp)
  - ctx[0x24], ctx[0x1a0], ctx[0x1ac] — ratio parameters for w computation
  - w is currently fitted empirically (≈0.015 for this test signal)
"""
import numpy as np

N = 2049  # FFT bins (NFRAME/2 + 1)


def haar_one_pass(b):
    """One Haar smoothing pass (kernel [0.25, 0.5, 0.25]).

    Decoded from 529c60 Haar loop (lines 35-74):
      Step 1: b[i] += b[i+1]   (prefix sum, 10e40)
      Step 2: b[i] *= 0.5      (scalar mul, ffe0)
      Step 3: scratch[i] = b[i+1] + b[i]  (3-op add, 11580)
      Step 4: b[i+1] = 0.5 * scratch[i]   (scalar mul+store, 4720)
    """
    n = len(b)
    if n < 2:
        return b
    # Steps 1+2 combined: b[i] = 0.5 * (b[i] + b[i+1]) for i < n-1
    # Note: b[n-1] is unchanged by steps 1+2
    b[:-1] = 0.5 * (b[:-1] + b[1:])
    # Steps 3+4: b[i+1] = 0.5 * (b[i] + b[i+1]) using UPDATED b
    # Need original b[i] values for step 3
    # Actually: step 3 reads AFTER steps 1+2, so uses modified b
    # scratch[i] = b[i+1] + b[i]  (both modified)
    # b[i+1] = 0.5 * scratch[i]
    # This means: b_new[i+1] = 0.5 * (b_modified[i+1] + b_modified[i])
    b6f8 = b[1:] + b[:-1]
    b[1:] = 0.5 * b6f8
    return b


def haar_smooth(magnitudes, n_iters):
    """Haar smoothing: iterate Haar passes.

    Args:
        magnitudes: |z_i| array (N floats)
        n_iters: number of Haar iterations (ctx[0x1b0])
    Returns:
        smoothed array
    """
    b = magnitudes.copy()
    for _ in range(n_iters):
        haar_one_pass(b)
    return b


def cascade_detect(complex_state, n_iters=2, w=0.015, sin_peak_floor=0.0):
    """Full detector cascade (529c60) simulation.

    Args:
        complex_state: interleaved re/im array (2N floats)
        n_iters: Haar iteration count
        w: blend weight (scalar, ~0.015 for typical settings)
        sin_peak_floor: minimum from sin modulation (0 = disabled)
    Returns:
        bands_output: smoothed detector curve (N floats)
    """
    n = len(complex_state) // 2
    re = complex_state[0::2]
    im = complex_state[1::2]

    # Phase 1: magnitudes via 16140
    magnitudes = np.sqrt(re**2 + im**2)

    # Phase 2: Haar smoothing
    curve = haar_smooth(magnitudes, n_iters)

    # Phase 3 (partial — unknown ctx params):
    # peak = max(curve)  [4d56b0]
    # sin_peak = sin(ctx[0x54087c]*30 - 90) * 0.115129 * peak  [1a14cac]
    # curve[i] = max(curve[i], sin_peak)  [52d8a0→10860]
    if sin_peak_floor > 0:
        np.maximum(curve, sin_peak_floor, out=curve)

    # Blend: output = curve * (1-w) + accumulator * w
    # 5407a8 (accumulator) = 0 in steady state → output = curve * (1-w)
    # The blend chain:
    #   52d920: 5407a8[i] *= w       (array scalar mul)
    #   52dae0: 5407a8[i] += curve[i] * (1-w)  (FMA)
    #   52dbc0: memcpy 5407a8 → 540678
    bands_output = curve * (1.0 - w)

    return bands_output


def validate():
    """Validate against ptrace capture (chain_samples.pkl)."""
    import pickle
    path = '/tmp/opencode/winetrace_casc/chain_samples.pkl'
    with open(path, 'rb') as f:
        data = pickle.load(f)

    s = data['samples']
    cin = s[0]
    cout = s[3]

    trk = np.array(cin['trk'], dtype=np.float64)
    b0_cout = np.array(cout['bands0'], dtype=np.float64)

    # Fit w and n_iters
    best_rms = 1e10
    best_params = None

    for n_iters in range(1, 11):
        magnitudes = np.zeros(len(trk) // 2)
        re = trk[0::2]; im = trk[1::2]
        magnitudes = np.sqrt(re**2 + im**2)

        curve = haar_smooth(magnitudes, n_iters)

        sig = (curve > 0.5) & (b0_cout > 0.5)
        if sig.sum() < 10:
            continue

        w_vals = 1.0 - b0_cout[sig] / curve[sig]
        w = float(np.median(w_vals))

        predicted = curve * (1.0 - w)
        rms = float(np.sqrt(np.mean((predicted - b0_cout) ** 2)))
        corr = float(np.corrcoef(curve[sig], b0_cout[sig])[0, 1])

        if rms < best_rms:
            best_rms = rms
            best_params = (n_iters, w, corr)

        print(f'  iters={n_iters:2d}: w={w:.6f}, rms={rms:.4f}, corr={corr:.6f}')

    n_iters, w, corr = best_params
    print(f'\nBest: iters={n_iters}, w={w:.6f}, rms={best_rms:.4f}, corr={corr:.6f}')
    return n_iters, w


if __name__ == '__main__':
    import sys
    if '--validate' in sys.argv:
        validate()
    else:
        print('Usage: detector_cascade.py --validate')