"""
Example: Extract speaker embeddings using the MLX model.

Requirements:
    pip install mlx numpy soundfile

Usage:
    python example_usage.py audio.wav
"""

import sys
import numpy as np
import mlx.core as mx
import soundfile as sf
from pathlib import Path


def compute_fbank(wav: np.ndarray, sr: int = 16000, n_mels: int = 80) -> np.ndarray:
    """Compute log Mel filterbank features (matching Kaldi defaults)."""
    import numpy as np
    
    frame_len = int(0.025 * sr)  # 25ms
    frame_shift = int(0.01 * sr)  # 10ms
    nfft = 512
    
    # Pre-emphasis
    wav = np.append(wav[0], wav[1:] - 0.97 * wav[:-1])
    
    # Framing
    n_frames = 1 + (len(wav) - frame_len) // frame_shift
    frames = np.stack([wav[i * frame_shift:i * frame_shift + frame_len] for i in range(n_frames)])
    
    # Windowing
    window = np.hamming(frame_len).astype(np.float32)
    frames *= window
    
    # FFT
    spec = np.abs(np.fft.rfft(frames, n=nfft)) ** 2
    
    # Mel filterbank
    low_freq, high_freq = 20, sr / 2
    low_mel = 2595 * np.log10(1 + low_freq / 700)
    high_mel = 2595 * np.log10(1 + high_freq / 700)
    mel_points = np.linspace(low_mel, high_mel, n_mels + 2)
    hz_points = 700 * (10 ** (mel_points / 2595) - 1)
    bins = np.floor((nfft + 1) * hz_points / sr).astype(int)
    
    fbank = np.zeros((n_mels, nfft // 2 + 1))
    for i in range(n_mels):
        for j in range(bins[i], bins[i + 1]):
            fbank[i, j] = (j - bins[i]) / (bins[i + 1] - bins[i])
        for j in range(bins[i + 1], bins[i + 2]):
            fbank[i, j] = (bins[i + 2] - j) / (bins[i + 2] - bins[i + 1])
    
    features = np.dot(spec, fbank.T)
    features = np.where(features > 0, np.log(features), np.log(1e-10))
    
    # CMVN (cepstral mean and variance normalization)
    features = (features - features.mean(axis=0)) / (features.std(axis=0) + 1e-7)
    
    return features.astype(np.float32)


def load_model(model_dir: str):
    """Load the MLX ResNet34 embedding model."""
    import importlib.util
    
    model_dir = Path(model_dir)
    
    # Import model class
    spec = importlib.util.spec_from_file_location("resnet", model_dir / "resnet_embedding.py")
    mod = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(mod)
    
    model = mod.ResNet34Embedding()
    
    # Load weights
    weights = np.load(str(model_dir / "weights.npz"))
    for key in weights.files:
        path = key.split(".")
        module = model
        for attr in path[:-1]:
            if attr.isdigit():
                module = module[int(attr)]
            elif attr == "layers":
                module = module.layers
            else:
                module = getattr(module, attr)
        setattr(module, path[-1], mx.array(weights[key]))
    model.eval()
    
    return model


def extract_embedding(model, audio_path: str) -> np.ndarray:
    """Extract speaker embedding from audio file."""
    wav, sr = sf.read(audio_path, dtype="float32")
    if sr != 16000:
        import subprocess, io
        result = subprocess.run(
            ["ffmpeg", "-i", audio_path, "-ar", "16000", "-ac", "1", "-f", "wav", "-"],
            capture_output=True
        )
        wav, sr = sf.read(io.BytesIO(result.stdout), dtype="float32")
    
    if wav.ndim > 1:
        wav = wav.mean(axis=1)
    
    features = compute_fbank(wav, sr)
    embedding = model(mx.array(features[np.newaxis, :, :]))
    mx.eval(embedding)
    
    # L2 normalize
    emb = np.array(embedding).flatten()
    emb = emb / (np.linalg.norm(emb) + 1e-8)
    return emb


if __name__ == "__main__":
    if len(sys.argv) < 2:
        print("Usage: python example_usage.py <audio.wav> [audio2.wav]")
        sys.exit(1)
    
    model = load_model(Path(__file__).parent)
    
    embeddings = []
    for path in sys.argv[1:]:
        emb = extract_embedding(model, path)
        print(f"{path}: embedding shape = {emb.shape}")
        embeddings.append(emb)
    
    if len(embeddings) == 2:
        sim = np.dot(embeddings[0], embeddings[1])
        print(f"\nCosine similarity: {sim:.6f}")
        print(f"Same speaker: {'Yes' if sim > 0.65 else 'No'}")