import numpy as np
from scipy.io import wavfile
from scipy.signal import resample, fftconvolve
import sys

SR = 8000

def load(path):
    sr, x = wavfile.read(path)
    assert sr == SR, sr
    x = x.astype(np.float64)
    return x

sus = load(sys.argv[1])          # suspect audio (0..3078s)
ref = load(sys.argv[2])          # reference episode (full)
label = sys.argv[3]

# z-score suspect once
S = sus - sus.mean()
# precompute cumulative energy of S for local normalization
Spad = np.concatenate([[0.0], np.cumsum(S*S)])

def ncc_best(template):
    # template: 1D array; returns (best_ncc, lag_samples)
    t = template - template.mean()
    tn = np.linalg.norm(t)
    if tn == 0: return (0.0, 0)
    t = t / tn
    L = len(t)
    # raw cross-correlation via fft: corr[k] = sum_i S[k+i]*t[i]
    corr = fftconvolve(S, t[::-1], mode='valid')   # length len(S)-L+1
    # local L2 norm of S window of length L at each lag
    win_energy = Spad[L:len(S)+1] - Spad[0:len(S)-L+1]
    local_norm = np.sqrt(np.maximum(win_energy, 1e-9))
    ncc = corr / local_norm
    k = int(np.argmax(ncc))
    return (float(ncc[k]), k)

# take template chunks at fractions of the reference (skip intro/outro edges)
fracs = [0.15, 0.30, 0.45, 0.60, 0.75, 0.90]
chunk_s = 45  # seconds
chunk_n0 = int(chunk_s * SR)

speeds = np.round(np.arange(1.00, 1.351, 0.01), 3)

print(f"=== {label} ===")
for fr in fracs:
    start = int(fr * len(ref))
    base = ref[start:start+chunk_n0]
    if len(base) < chunk_n0*0.9:
        continue
    best = (-1, None, None)  # ncc, speed, lag
    for k in speeds:
        # speed up reference chunk by factor k -> shorter by k
        newlen = max(8, int(round(len(base)/k)))
        tmpl = resample(base, newlen)
        ncc, lag = ncc_best(tmpl)
        if ncc > best[0]:
            best = (ncc, k, lag)
    ncc, k, lag = best
    t_ref = start/SR
    t_sus = lag/SR
    print(f" ref@{t_ref:6.1f}s  -> best speed k={k:.3f}  ncc={ncc:.3f}  suspect_lag={t_sus:7.1f}s")