#!/usr/bin/env python3
"""Are a music bank's two waves the SAME instruments filtered, or DIFFERENT parts?

`structures/bgm-two-stems.md` leaves two readings alive for wave 1 -- the rear
pair of a 4-channel mix, or a second intensity layer -- and notes that runtime
simultaneity cannot separate them, since both predict it.

This tries a static discriminator.  Magnitude-squared coherence is ~1 wherever
one signal is a LINEAR FILTER of the other, and ~0 for independent signals, so
a rear pair modelled as "front pair, filtered" should read high and a different
arrangement layer should read low.

⚠️ READ THE CONTROLS BEFORE THE MEASUREMENT.  The L-vs-R control below is the
one that matters and it is the one that limits this tool: see the docs page.

    bgm_stem_coherence.py <w0.wav> <w1.wav> <control.wav>

Waves come from slb_extract_wave.py + ffmpeg; see the docs page for the exact
offsets and packet counts.
"""
import sys, wave, numpy as np

NFFT, HOP, FS = 8192, 4096, 48000
BANDS = [(0,200),(200,1000),(1000,4000),(4000,12000),(12000,16000),(16000,24000)]

def load(p, nmax, stereo=False):
    w = wave.open(p); n = min(nmax, w.getnframes())
    a = np.frombuffer(w.readframes(n), dtype="<i2").astype(np.float64)
    a = a.reshape(-1, w.getnchannels())
    return (a[:,0], a[:,1]) if stereo else a.mean(axis=1)

def _acc(x, y):
    win = np.hanning(NFFT); Sxx = Syy = Sxy = 0.0; k = 0
    for i in range(0, min(len(x), len(y)) - NFFT, HOP):
        X = np.fft.rfft(x[i:i+NFFT]*win); Y = np.fft.rfft(y[i:i+NFFT]*win)
        Sxx = Sxx + abs(X)**2; Syy = Syy + abs(Y)**2; Sxy = Sxy + X*np.conj(Y); k += 1
    return np.abs(Sxy)**2/(Sxx*Syy+1e-30), Sxx, Syy, k

def coh(x, y): return _acc(x, y)[0]
def spec(x):   return _acc(x, x)[1]

def row(name, C, f):
    print(f"  {name:34} " + "  ".join(
        f"{lo/1000:g}-{hi/1000:g}k {C[(f>=lo)&(f<hi)].mean():.3f}" for lo,hi in BANDS))

if __name__ == "__main__":
    p0, p1, pc = sys.argv[1], sys.argv[2], sys.argv[3]
    D = 60*FS; f = np.fft.rfftfreq(NFFT, 1/FS)
    w0 = load(p0, D); w1 = load(p1, D); ctl = load(pc, D)
    L0,R0 = load(p0, D, True); L1,R1 = load(p1, D, True)

    # POSITIVE control: an actual linear filter of w0 (one-pole LP + 12 ms delay)
    a = 0.06; lp = np.empty_like(w0); acc = 0.0
    for i, v in enumerate(w0):
        acc += a*(v-acc); lp[i] = acc
    d = int(0.012*FS); pos = np.concatenate([np.zeros(d), lp[:-d]])
    sh = np.concatenate([np.zeros(FS), w0[:-FS]])

    print("CONTROLS")
    row("POS w0 vs linear-filter(w0)", coh(w0,pos), f)
    row("NEG w0 vs a different bank",  coh(w0,ctl), f)
    row("NEG w0 vs w0 shifted 1 s",    coh(w0,sh),  f)
    row("REF w0 L vs R (one perf.)",   coh(L0,R0),  f)
    row("REF w1 L vs R (one perf.)",   coh(L1,R1),  f)
    print("\nMEASUREMENT")
    row("w0 vs w1",                    coh(w0,w1),  f)
    print("\nENERGY SHARE")
    for nm,(L,R) in (("w0",(L0,R0)),("w1",(L1,R1))):
        S = spec(L)+spec(R); t = S.sum()
        print(f"  {nm:34} " + "  ".join(
            f"{lo/1000:g}-{hi/1000:g}k {100*S[(f>=lo)&(f<hi)].sum()/t:5.1f}%" for lo,hi in BANDS))