"""
Minimal 256ms unified inference example for Silero VAD v6 (MLX, 16 kHz).

Same weights as 32ms streaming mode (model.safetensors), but processes 8×
512-sample chunks per outer call sharing LSTM state, then aggregates the 8
internal probabilities via noisy-OR: 1 - prod(1 - p_i).

Designed for offline / batch use cases like ASR preprocessing where 256ms
decision granularity is acceptable. Roughly 1.5-2× faster wall time than the
32ms streaming inference on the same audio (fewer mx.eval barriers).

Usage:
    uv pip install mlx safetensors numpy huggingface_hub
    uv run python example_256ms.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
BLOCK_TOTAL = 4096
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 predict_block_256ms(block_audio_4096: np.ndarray,
                         initial_ctx: np.ndarray,
                         h: mx.array, c: mx.array,
                         w: dict
                         ) -> tuple[float, np.ndarray, mx.array, mx.array]:
    """Process one 256ms block (8× 32ms internal chunks) with LSTM state shared
    across the chunks, return aggregated probability via noisy-OR."""
    rolling_ctx = initial_ctx
    chunk_probs: list[mx.array] = []
    for ci in range(8):
        src = block_audio_4096[ci * CHUNK_SAMPLES : (ci + 1) * CHUNK_SAMPLES]
        merged = np.concatenate([rolling_ctx, src])
        padded = reflect_pad_right_576(merged)
        audio = mx.array(padded[None, :, None])
        prob, h, c = predict_chunk(audio, h, c, w)
        chunk_probs.append(prob)
        rolling_ctx = src[-CONTEXT_SAMPLES:].copy()

    product = (1.0 - chunk_probs[0])
    for ci in range(1, 8):
        product = product * (1.0 - chunk_probs[ci])
    agg = 1.0 - product
    mx.eval(agg, h, c)
    new_ctx = block_audio_4096[-CONTEXT_SAMPLES:].copy()
    return float(agg[0, 0, 0].item()), new_ctx, h, c


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])
    leftover = samples.shape[0] % BLOCK_TOTAL
    if leftover:
        samples = np.concatenate([samples, np.zeros(BLOCK_TOTAL - leftover, dtype=np.float32)])
    n_blocks = samples.shape[0] // BLOCK_TOTAL

    h = mx.zeros((1, LSTM_HIDDEN))
    c = mx.zeros((1, LSTM_HIDDEN))
    initial_ctx = np.zeros(CONTEXT_SAMPLES, dtype=np.float32)
    threshold = 0.5
    n_speech_blocks = 0
    print(f"blocks={n_blocks}, duration={samples.shape[0]/SAMPLE_RATE:.2f}s, block_dur=256.0ms")
    print(f"{'idx':>5} {'time(s)':>8} {'p(speech)':>10}")

    for bi in range(n_blocks):
        block_audio = samples[bi * BLOCK_TOTAL : (bi + 1) * BLOCK_TOTAL]
        prob, initial_ctx, h, c = predict_block_256ms(block_audio, initial_ctx, h, c, w)
        if prob >= threshold:
            n_speech_blocks += 1
        if bi < 25:
            print(f"{bi:>5d} {bi*BLOCK_TOTAL/SAMPLE_RATE:>8.3f} {prob:>10.4f}")

    print(f"\nspeech blocks (threshold {threshold}): {n_speech_blocks}/{n_blocks} "
          f"({n_speech_blocks/n_blocks*100:.1f}%)")


if __name__ == "__main__":
    main()