"""Smart Turn v3 — semantic end-of-turn detection.

Analyses a raw 16 kHz mono waveform and predicts whether the speaker has
finished their turn. Backbone: Whisper-tiny encoder + linear classifier (ONNX).

The log-mel feature extraction and the front-padding convention (audio sits at
the END of the 8 s window, zeros padded at the BEGINNING) are ported verbatim
from pipecat's reference implementation so results match the trained model:
  pipecat/audio/turn/smart_turn/_whisper_features.py
  pipecat/audio/turn/smart_turn/local_smart_turn_v3.py

Model: https://huggingface.co/pipecat-ai/smart-turn-v3
"""

import numpy as np
import onnxruntime as ort
from numpy.lib.stride_tricks import sliding_window_view

SAMPLE_RATE = 16000
MAX_SECONDS = 8
MAX_SAMPLES = SAMPLE_RATE * MAX_SECONDS  # 128000

_N_FFT = 400
_HOP_LENGTH = 160
_N_MELS = 80
_MEL_FLOOR = 1e-10
_NORM_VARIANCE_EPS = 1e-7


# --- mel filterbank (Slaney), vendored from pipecat / transformers ---
def _hertz_to_mel_slaney(freq):
    min_log_hertz, min_log_mel = 1000.0, 15.0
    logstep = 27.0 / np.log(6.4)
    freq = np.atleast_1d(np.asarray(freq, dtype=np.float64))
    mels = 3.0 * freq / 200.0
    lr = freq >= min_log_hertz
    mels[lr] = min_log_mel + np.log(freq[lr] / min_log_hertz) * logstep
    return mels


def _mel_to_hertz_slaney(mels):
    min_log_hertz, min_log_mel = 1000.0, 15.0
    logstep = np.log(6.4) / 27.0
    mels = np.atleast_1d(np.asarray(mels, dtype=np.float64))
    freq = 200.0 * mels / 3.0
    lr = mels >= min_log_mel
    freq[lr] = min_log_hertz * np.exp(logstep * (mels[lr] - min_log_mel))
    return freq


def _build_mel_filterbank(num_frequency_bins, num_mel_filters, min_freq, max_freq, sr):
    mel_min = float(_hertz_to_mel_slaney(np.array([min_freq]))[0])
    mel_max = float(_hertz_to_mel_slaney(np.array([max_freq]))[0])
    mel_freqs = np.linspace(mel_min, mel_max, num_mel_filters + 2)
    filter_freqs = _mel_to_hertz_slaney(mel_freqs)
    fft_freqs = np.linspace(0, sr // 2, num_frequency_bins)
    diff = np.diff(filter_freqs)
    slopes = np.expand_dims(filter_freqs, 0) - np.expand_dims(fft_freqs, 1)
    down = -slopes[:, :-2] / diff[:-1]
    up = slopes[:, 2:] / diff[1:]
    mel = np.maximum(np.zeros(1), np.minimum(down, up))
    enorm = 2.0 / (filter_freqs[2 : num_mel_filters + 2] - filter_freqs[:num_mel_filters])
    mel *= np.expand_dims(enorm, 0)
    return mel


_HANN_WINDOW = np.hanning(_N_FFT + 1)[:-1]
_MEL_FILTERS = _build_mel_filterbank(_N_FFT // 2 + 1, _N_MELS, 0.0, SAMPLE_RATE / 2.0, SAMPLE_RATE)


def _power_spectrogram(waveform):
    pad = _N_FFT // 2
    padded = np.pad(waveform.astype(np.float64), (pad, pad), mode="reflect")
    win = _HANN_WINDOW.astype(np.float64)
    windows = sliding_window_view(padded, _N_FFT)[::_HOP_LENGTH]
    spec = np.fft.rfft(windows * win, axis=-1)
    return (np.abs(spec) ** 2).T


def compute_whisper_log_mel_features(audio, do_normalize=True):
    """Whisper-style log-mel features -> (80, 800), matching Smart Turn v3."""
    x = np.asarray(audio, dtype=np.float32)
    if x.size < MAX_SAMPLES:
        x = np.pad(x, (0, MAX_SAMPLES - x.size), mode="constant")
    elif x.size > MAX_SAMPLES:
        x = x[:MAX_SAMPLES]
    if do_normalize:
        x = (x - x.mean()) / np.sqrt(x.var() + _NORM_VARIANCE_EPS)
    mags = _power_spectrogram(x)
    mel = np.maximum(_MEL_FLOOR, _MEL_FILTERS.T @ mags)
    log_spec = np.log10(mel)[:, :-1]
    log_spec = np.maximum(log_spec, log_spec.max() - 8.0)
    log_spec = (log_spec + 4.0) / 4.0
    return log_spec.astype(np.float32)


def _fit_to_window(audio):
    """Keep the LAST 8 s; if shorter, pad zeros at the BEGINNING (audio at end)."""
    if len(audio) > MAX_SAMPLES:
        return audio[-MAX_SAMPLES:]
    if len(audio) < MAX_SAMPLES:
        return np.pad(audio, (MAX_SAMPLES - len(audio), 0), mode="constant")
    return audio


class SmartTurn:
    def __init__(self, model_path="models/smart-turn-v3.2-cpu.onnx", threshold=0.5):
        self.session = ort.InferenceSession(model_path, providers=["CPUExecutionProvider"])
        self.threshold = threshold

    def _preprocess(self, audio, sample_rate=SAMPLE_RATE):
        audio = np.asarray(audio, dtype=np.float32)
        if sample_rate != SAMPLE_RATE:
            import librosa
            audio = librosa.resample(audio, orig_sr=sample_rate, target_sr=SAMPLE_RATE)
        audio = _fit_to_window(audio)                       # audio at END of 8 s window
        feats = compute_whisper_log_mel_features(audio)     # (80, 800)
        return feats[np.newaxis, :, :].astype(np.float32)   # (1, 80, 800)

    def predict(self, audio, sample_rate=SAMPLE_RATE):
        """Return {'probability', 'is_complete'} for a mono waveform.

        Pass sample_rate if not 16 kHz (resampled). Audio > 8 s is truncated to
        the last 8 s; shorter audio is front-padded so it sits at the window end.
        """
        input_features = self._preprocess(audio, sample_rate)
        # The ONNX graph already applies the sigmoid: output is a probability in
        # [0,1], not a logit. Do NOT apply sigmoid again (matches the reference
        # inference.py: probability = outputs[0][0]).
        prob = float(self.session.run(None, {"input_features": input_features})[0][0][0])
        return {"probability": prob, "is_complete": prob > self.threshold}


if __name__ == "__main__":
    model = SmartTurn()
    for name, audio in [
        ("2s silence", np.zeros(SAMPLE_RATE * 2, dtype=np.float32)),
        ("3s noise", (np.random.randn(SAMPLE_RATE * 3) * 0.1).astype(np.float32)),
    ]:
        r = model.predict(audio)
        print(f"{name:12s} -> prob={r['probability']:.4f} complete={r['is_complete']}")