#!/bin/bash
# Phase 5d parity check: runs HF reference and Rust port on the same audio
# pair, then prints side-by-side cosines, embedding norms, and a head/tail
# fingerprint cosine (same-utterance HF vs Rust similarity).
#
# Prerequisites:
#   - .venv created via `uv venv .venv --python 3.13`
#   - .venv has transformers, torch, soundfile, torchaudio, scipy
#   - /tmp/wavlm_sv.safetensors produced by `cargo run --example wavlm_sv_convert`
#
# Usage:
#   ./wavlm_sv_parity_check.sh /path/to/a.wav /path/to/b.wav

set -e
cd "$(dirname "$0")"

A="${1:-/tmp/csm_24k.wav}"
B="${2:-/tmp/pipeline_out.wav}"
WEIGHTS="${WAVLM_SV_SAFETENSORS:-/tmp/wavlm_sv.safetensors}"

if [ ! -f "$A" ] || [ ! -f "$B" ]; then
    echo "missing audio: $A or $B" >&2
    exit 1
fi

if [ ! -d .venv ]; then
    echo "missing .venv — run: uv venv .venv --python 3.13" >&2
    echo "                       uv pip install --python .venv/bin/python transformers torch soundfile torchaudio scipy" >&2
    exit 1
fi

PY=.venv/bin/python

echo "== Phase 5d parity check =="
echo "  audio_a: $A"
echo "  audio_b: $B"
echo "  weights: $WEIGHTS"

echo "[1/3] HF reference embeddings..."
$PY wavlm_sv_parity.py --wav-a "$A" --wav-b "$B" --out /tmp/wavlm_sv_ref.json 2>/dev/null | tail -2

echo "[2/3] Rust port embeddings..."
( cd ../../../.. && cargo run -p rtx-csm --release --features metal --example wavlm_sv_demo -- \
    --weights "$WEIGHTS" --a "$A" --b "$B" --parity-json /tmp/wavlm_sv_rust.json ) 2>/dev/null | tail -2

echo "[3/3] Parity report:"
$PY <<'PYEOF'
import json, numpy as np
hf = json.load(open('/tmp/wavlm_sv_ref.json'))
rs = json.load(open('/tmp/wavlm_sv_rust.json'))
def cos(a, b):
    a = np.array(a); b = np.array(b)
    return float((a @ b) / (np.linalg.norm(a) * np.linalg.norm(b) + 1e-9))
print(f"  pairwise cosine — HF: {hf['cosine_similarity_hf']:+.4f}  Rust: {rs['cosine_similarity_rust']:+.4f}  delta: {rs['cosine_similarity_rust'] - hf['cosine_similarity_hf']:+.4f}")
print(f"  |emb_a| ratio (Rust/HF): {rs['embedding_a_norm']/hf['embedding_a_norm']:.4f}")
print(f"  |emb_b| ratio (Rust/HF): {rs['embedding_b_norm']/hf['embedding_b_norm']:.4f}")
hfa = hf['embedding_a_head'] + hf['embedding_a_tail']
rsa = rs['embedding_a_head'] + rs['embedding_a_tail']
hfb = hf['embedding_b_head'] + hf['embedding_b_tail']
rsb = rs['embedding_b_head'] + rs['embedding_b_tail']
print(f"  same-utterance HF↔Rust cosine — utt_a: {cos(hfa, rsa):.5f}  utt_b: {cos(hfb, rsb):.5f}  (1.0 = perfect parity)")
parity_a = cos(hfa, rsa)
parity_b = cos(hfb, rsb)
if parity_a > 0.99 and parity_b > 0.99:
    print("  VERDICT: parity within tolerance (>0.99 same-utterance HF↔Rust cosine)")
else:
    print("  VERDICT: NUMERICAL DRIFT — investigate intermediate tensors")
PYEOF