"""Transcribe an audio file with Anis Transcribe 1.0 (ONNX int8, CPU).

    pip install onnxruntime "transformers>=5.19" soundfile numpy
    python transcribe.py recording.wav            # any sample rate / channels; ffmpeg-decodable formats need ffmpeg

Greedy CTC decoding, then spoken numbers -> digits (numwords.itn). Audio is read in 25 s windows, as in training.
"""
import json
import subprocess
import sys
from pathlib import Path

import numpy as np
import onnxruntime as ort
import soundfile as sf
from transformers import AutoFeatureExtractor

import numwords

HERE = Path(__file__).parent
SR, WINDOW = 16000, 25 * 16000


def read16k(path: str) -> np.ndarray:
    try:
        wav, sr = sf.read(path, dtype="float32")
        if wav.ndim > 1:
            wav = wav.mean(1)
        if sr == SR:
            return wav
    except Exception:
        pass
    raw = subprocess.run(["ffmpeg", "-v", "error", "-i", path, "-ac", "1", "-ar", str(SR), "-f", "f32le", "-"],
                         check=True, capture_output=True).stdout
    return np.frombuffer(raw, dtype=np.float32)


class AnisTranscribe:
    def __init__(self, threads: int = 4):
        so = ort.SessionOptions()
        so.intra_op_num_threads = threads
        self.sess = ort.InferenceSession(str(HERE / "anis-transcribe-1.0-int8.onnx"), so, providers=["CPUExecutionProvider"])
        self.fe = AutoFeatureExtractor.from_pretrained(str(HERE))
        self.vocab = json.loads((HERE / "vocab.json").read_text(encoding="utf-8"))

    def _greedy(self, lp: np.ndarray) -> str:
        prev, out = 0, []
        for t in lp.argmax(-1):
            if t != prev and t != 0:
                out.append(self.vocab[t])
            prev = t
        return " ".join("".join(out).split())

    def __call__(self, wav: np.ndarray) -> str:
        parts = []
        for s in range(0, max(1, len(wav)), WINDOW):
            seg = wav[s:s + WINDOW]
            if len(seg) < SR // 10:
                continue
            f = self.fe([seg], sampling_rate=SR, return_tensors="np", padding="longest")
            lp = self.sess.run(None, {"features": f["input_features"].astype(np.float32)})[0][0]
            parts.append(self._greedy(lp))
        return numwords.itn(" ".join(p for p in parts if p))


if __name__ == "__main__":
    model = AnisTranscribe()
    for path in sys.argv[1:]:
        print(model(read16k(path)))