import numpy as np, sys
from scipy.io import wavfile
from scipy.signal import resample
from numpy.fft import rfft, irfft
SR=4000
def load(p):
    sr,x=wavfile.read(p); assert sr==SR; return x.astype(np.float64)

def search(haystack, template, speeds):
    S=haystack-haystack.mean(); N=len(S)
    cum=np.concatenate([[0.0],np.cumsum(S*S)])
    NFFT=1<<int(np.ceil(np.log2(N+len(template)+8)))
    FS=rfft(S,NFFT)
    best=(-1,None,None)
    for k in speeds:
        t=resample(template,max(8,int(round(len(template)/k))))
        t=t-t.mean(); n=np.linalg.norm(t)
        if n==0: continue
        t=t/n; L=len(t)
        corr=irfft(FS*rfft(t[::-1],NFFT),NFFT)[L-1:N]
        we=cum[L:N+1]-cum[0:N-L+1]
        ncc=corr/np.sqrt(np.maximum(we,1e-9))
        j=int(np.argmax(ncc))
        if ncc[j]>best[0]: best=(float(ncc[j]),float(k),j/SR)
    return best

sus=load('sus4k.wav'); yo=load('ref_yo4k.wav'); fam=load('ref_family4k.wav')

# Control 1: suspect chunk @ 300s searched within suspect (should be ncc~1, k=1, lag=300)
tmpl=sus[300*SR:345*SR]
print("CTRL suspect-in-suspect:", search(sus,tmpl,[1.0]))

# Control 2: yo chunk searched within yo (ncc~1, k=1)
print("CTRL yo-in-yo:", search(yo,yo[400*SR:445*SR],[1.0]))

# Control 3: make a synthetic 'suspect' = yo sped up 1.25x, then search yo-chunk with sweep -> should recover k=1.25
yo125=resample(yo,int(round(len(yo)/1.25)))
print("CTRL recover-speed (yo@1.25x):", search(yo125, yo[500*SR:545*SR], np.round(np.arange(1.0,1.351,0.01),3)))