"""
Minimal streaming inference example for Silero VAD v6 (MLX, 16 kHz).

Reads a 16 kHz mono PCM-16 WAV file, runs streaming VAD, and prints per-chunk
speech probabilities for the first 25 chunks plus the speech ratio at threshold
0.5 across the whole clip.

Usage:
    uv pip install mlx safetensors numpy huggingface_hub
    uv run python example.py path/to/audio_16k_mono.wav
"""
from __future__ import annotations

import sys
import wave
from pathlib import Path

import mlx.core as mx
import numpy as np
from huggingface_hub import hf_hub_download
from safetensors.numpy import load_file


CHUNK_SAMPLES = 512
CONTEXT_SAMPLES = 64
STFT_HOP = 128
STFT_PAD_RIGHT = 64
LSTM_HIDDEN = 128
SAMPLE_RATE = 16_000

KEYS = [
    "vad_16k.stft_conv.weight",
    "vad_16k.conv1.weight", "vad_16k.conv1.bias",
    "vad_16k.conv2.weight", "vad_16k.conv2.bias",
    "vad_16k.conv3.weight", "vad_16k.conv3.bias",
    "vad_16k.conv4.weight", "vad_16k.conv4.bias",
    "vad_16k.lstm.Wx", "vad_16k.lstm.Wh", "vad_16k.lstm.bias",
    "vad_16k.final_conv.weight", "vad_16k.final_conv.bias",
]


def load_weights(path: str) -> dict[str, mx.array]:
    raw = load_file(path)
    return {k: mx.array(raw[k]) for k in KEYS}


def reflect_pad_right_576(samples_576: np.ndarray) -> np.ndarray:
    """PyTorch ReflectionPad1d(right=64): out[L+i] = in[L-2-i] for i in 0..63."""
    L = samples_576.shape[-1]
    pad = np.empty(STFT_PAD_RIGHT, dtype=samples_576.dtype)
    for i in range(STFT_PAD_RIGHT):
        pad[i] = samples_576[L - 2 - i]
    return np.concatenate([samples_576, pad], axis=-1)


def predict_chunk(audio_640: mx.array, h: mx.array, c: mx.array, w: dict
                   ) -> tuple[mx.array, mx.array, mx.array]:
    z = mx.conv1d(audio_640, w["vad_16k.stft_conv.weight"], stride=STFT_HOP, padding=0)
    real, imag = z[:, :, :129], z[:, :, 129:]
    x = mx.sqrt(real * real + imag * imag + 1e-12)
    cfgs = [(1, 1), (2, 1), (2, 1), (1, 1)]
    for i, (stride, padding) in enumerate(cfgs, start=1):
        x = mx.conv1d(x, w[f"vad_16k.conv{i}.weight"], stride=stride, padding=padding)
        x = x + w[f"vad_16k.conv{i}.bias"]
        x = mx.maximum(x, 0)
    feat = x[:, 0, :]
    gates = feat @ w["vad_16k.lstm.Wx"].T + h @ w["vad_16k.lstm.Wh"].T + w["vad_16k.lstm.bias"]
    i_g, f_g, g_g, o_g = mx.split(gates, 4, axis=-1)
    c_new = mx.sigmoid(f_g) * c + mx.sigmoid(i_g) * mx.tanh(g_g)
    h_new = mx.sigmoid(o_g) * mx.tanh(c_new)
    dec = mx.conv1d(
        mx.maximum(h_new, 0)[:, None, :],
        w["vad_16k.final_conv.weight"], stride=1, padding=0,
    )
    dec = dec + w["vad_16k.final_conv.bias"]
    return mx.sigmoid(dec), h_new, c_new


def load_wav_16k_mono(path: str) -> np.ndarray:
    with wave.open(path, "rb") as w:
        if w.getframerate() != SAMPLE_RATE:
            raise SystemExit(f"expected 16 kHz, got {w.getframerate()}")
        if w.getnchannels() != 1:
            raise SystemExit(f"expected mono, got {w.getnchannels()} channels")
        if w.getsampwidth() != 2:
            raise SystemExit(f"expected 16-bit PCM, got {w.getsampwidth()*8}-bit")
        n = w.getnframes()
        raw = w.readframes(n)
    return (np.frombuffer(raw, dtype=np.int16).astype(np.float32) / 32768.0).copy()


def main() -> None:
    if len(sys.argv) != 2:
        print(__doc__, file=sys.stderr)
        sys.exit(1)

    weights_path = hf_hub_download(repo_id="mlx-community/silero-vad-v6",
                                    filename="model.safetensors")
    w = load_weights(weights_path)

    samples = load_wav_16k_mono(sys.argv[1])
    n_chunks = samples.shape[0] // CHUNK_SAMPLES
    h = mx.zeros((1, LSTM_HIDDEN))
    c = mx.zeros((1, LSTM_HIDDEN))
    context = np.zeros(CONTEXT_SAMPLES, dtype=np.float32)

    n_speech = 0
    threshold = 0.5
    print(f"chunks={n_chunks}, duration={samples.shape[0]/SAMPLE_RATE:.2f}s")
    print(f"{'idx':>5} {'time(s)':>8} {'p(speech)':>10}")
    for ci in range(n_chunks):
        chunk = samples[ci * CHUNK_SAMPLES : (ci + 1) * CHUNK_SAMPLES]
        merged = np.concatenate([context, chunk])
        padded = reflect_pad_right_576(merged)
        audio = mx.array(padded[None, :, None])
        prob, h, c = predict_chunk(audio, h, c, w)
        mx.eval(prob, h, c)
        p = float(prob[0, 0, 0].item())
        if p >= threshold:
            n_speech += 1
        if ci < 25:
            print(f"{ci:>5d} {ci*CHUNK_SAMPLES/SAMPLE_RATE:>8.3f} {p:>10.4f}")
        context = chunk[-CONTEXT_SAMPLES:].copy()

    print(f"\nspeech chunks (threshold {threshold}): {n_speech}/{n_chunks} "
          f"({n_speech/n_chunks*100:.1f}%)")


if __name__ == "__main__":
    main()