"""Envelope cross-correlation between a dry reference and its a2a render.

Reproduces the timing measurement published in the AKUSPACE model card, so new
examples can be added to that table on the same basis rather than asserted.

Method: amplitude envelope of each signal in 2 ms RMS windows, mean-removed and
normalised, cross-correlated over +/-250 ms. Reports the offset of peak
correlation and the peak value.

Reading the result: peak r measures how much the envelope CHANGED, and reverb
changes it by design — most on transient-dense material, where the tail fills
the gaps between hits. A low r on percussion is the effect working, not drift.
The offset is the timing claim; r is not a quality score.
"""

import argparse
import subprocess
import sys
import wave
from pathlib import Path

import numpy as np

WIN_MS = 2.0
MAX_LAG_MS = 250.0


def load_mono(path: Path, sr: int = 48000) -> np.ndarray:
    """Decode anything ffmpeg reads into mono float at sr."""
    out = subprocess.run(
        ["ffmpeg", "-v", "error", "-i", str(path), "-ac", "1", "-ar", str(sr),
         "-f", "wav", "-c:a", "pcm_s16le", "-"],
        capture_output=True, check=True).stdout
    import io
    with wave.open(io.BytesIO(out), "rb") as w:
        raw = w.readframes(w.getnframes())
    return np.frombuffer(raw, dtype="<i2").astype(np.float64) / 32768.0


def envelope(x: np.ndarray, sr: int) -> np.ndarray:
    n = max(1, int(sr * WIN_MS / 1000.0))
    trimmed = x[: len(x) - len(x) % n]
    return np.sqrt((trimmed.reshape(-1, n) ** 2).mean(axis=1) + 1e-12)


def xcorr(a: np.ndarray, b: np.ndarray, max_lag: int) -> tuple[int, float]:
    n = min(len(a), len(b))
    a, b = a[:n], b[:n]
    a = (a - a.mean()) / (a.std() or 1e-12)
    b = (b - b.mean()) / (b.std() or 1e-12)
    best_lag, best_r = 0, -2.0
    for lag in range(-max_lag, max_lag + 1):
        if lag < 0:
            x, y = a[-lag:], b[: n + lag]
        elif lag > 0:
            x, y = a[: n - lag], b[lag:]
        else:
            x, y = a, b
        if len(x) < 10:
            continue
        r = float((x * y).mean())
        if r > best_r:
            best_r, best_lag = r, lag
    return best_lag, best_r


def main() -> int:
    p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    p.add_argument("--dry", required=True)
    p.add_argument("--wet", nargs="+", required=True)
    p.add_argument("--sr", type=int, default=48000)
    a = p.parse_args()

    dry = envelope(load_mono(Path(a.dry), a.sr), a.sr)
    max_lag = int(MAX_LAG_MS / WIN_MS)
    print(f"{'example':38s} {'offset':>9s} {'peak r':>8s}")
    for w in a.wet:
        wet = envelope(load_mono(Path(w), a.sr), a.sr)
        lag, r = xcorr(dry, wet, max_lag)
        # positive lag = wet later than dry; report with the card's sign convention
        print(f"{Path(w).stem:38s} {lag * WIN_MS:+8.0f}ms {r:8.2f}")
    return 0


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