#!/usr/bin/env python3
"""Phase 5d: WavLM-SV numerical parity check.

Compares the rtx-csm WavLM-SV port against the HF reference. Run once
the user has a Python env with `transformers` + `torch` + `safetensors`
available; emits a JSON report with the per-stage comparisons we can
diff against the Rust outputs.

Workflow:
    pip install transformers torch torchaudio soundfile

    # Generate the reference embeddings:
    python3 scripts/wavlm_sv_parity.py \\
        --wav-a /tmp/csm_24k.wav \\
        --wav-b /tmp/test_24k.wav \\
        --out /tmp/wavlm_sv_reference.json

    # Then compare against the Rust outputs (TODO: extend wavlm_sv_demo
    # to dump embeddings to JSON for diffing).

Why this lives in scripts/ rather than tests/: it requires PyTorch +
HF transformers as a heavy build-time dependency, and the comparison
itself is a one-shot numerical-validation step, not a regression gate.
"""

import argparse
import json
import sys
from pathlib import Path


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--wav-a", required=True, type=Path)
    parser.add_argument("--wav-b", required=True, type=Path)
    parser.add_argument("--out", type=Path, default=Path("/tmp/wavlm_sv_reference.json"))
    parser.add_argument("--model", default="microsoft/wavlm-base-plus-sv")
    args = parser.parse_args()

    try:
        import torch
        import soundfile as sf
        from transformers import WavLMForXVector, AutoFeatureExtractor
    except ImportError as e:
        print(f"missing dependency: {e}", file=sys.stderr)
        print("  pip install transformers torch soundfile", file=sys.stderr)
        return 1

    print(f"loading {args.model} (cached on first run)...")
    extractor = AutoFeatureExtractor.from_pretrained(args.model)
    model = WavLMForXVector.from_pretrained(args.model).eval()

    def embed_path(path: Path) -> torch.Tensor:
        # soundfile returns (n,) for mono WAV; resample if needed.
        samples, sr = sf.read(str(path), dtype="float32")
        if samples.ndim > 1:
            samples = samples.mean(axis=1)
        if sr != 16000:
            import torchaudio
            samples = torch.from_numpy(samples).unsqueeze(0)
            samples = torchaudio.functional.resample(samples, sr, 16000)[0].numpy()
            sr = 16000
        inputs = extractor(samples, sampling_rate=sr, return_tensors="pt")
        with torch.no_grad():
            output = model(**inputs)
        return output.embeddings[0]  # (512,)

    emb_a = embed_path(args.wav_a)
    emb_b = embed_path(args.wav_b)
    cos = torch.nn.functional.cosine_similarity(
        emb_a.unsqueeze(0), emb_b.unsqueeze(0)
    ).item()

    report = {
        "model": args.model,
        "wav_a": str(args.wav_a),
        "wav_b": str(args.wav_b),
        "embedding_dim": emb_a.numel(),
        "cosine_similarity_hf": cos,
        "embedding_a_norm": emb_a.norm().item(),
        "embedding_b_norm": emb_b.norm().item(),
        # First and last 8 elements as a coarse fingerprint we can
        # diff against the Rust output to catch axis/permute bugs.
        "embedding_a_head": emb_a[:8].tolist(),
        "embedding_a_tail": emb_a[-8:].tolist(),
        "embedding_b_head": emb_b[:8].tolist(),
        "embedding_b_tail": emb_b[-8:].tolist(),
    }
    args.out.write_text(json.dumps(report, indent=2))
    print(f"wrote {args.out}")
    print(f"HF cosine_similarity = {cos:.4f}")
    print(f"|emb_a| = {report['embedding_a_norm']:.4f}, |emb_b| = {report['embedding_b_norm']:.4f}")
    return 0


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