#!/usr/bin/env python3
"""Dynamic ASR Pocket (300M, CPU): offline transcription with onnxruntime and greedy CTC decoding.

  pip install onnxruntime numpy soundfile scipy
  python transcribe_pocket.py clip.wav [more.wav ...]            # int8 model (default)
  python transcribe_pocket.py clip.wav --model model_fp32.onnx   # fp32 model

Input: any wav; it is mixed to mono and resampled to 16 kHz. Clips longer than ~20 s are cut at quiet points.
Model files are looked up next to this script: model_int8.onnx (or model_fp32.onnx) and vocab.json.
Code licence: Apache-2.0. Model weights: CC-BY-SA-4.0.
"""
import argparse, json, os, time
import numpy as np

HERE = os.path.dirname(os.path.abspath(__file__))
SR = 16000


class PocketASR:
    def __init__(self, model=None, vocab=None, threads=4):
        import onnxruntime as ort
        so = ort.SessionOptions()
        so.intra_op_num_threads = threads
        so.inter_op_num_threads = 1
        so.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
        self.sess = ort.InferenceSession(model or os.path.join(HERE, "model_int8.onnx"), so,
                                         providers=["CPUExecutionProvider"])
        v = json.load(open(vocab or os.path.join(HERE, "vocab.json"), encoding="utf-8"))
        self.pieces = v["pieces"]          # index 0 = CTC blank; 1-3 = special tokens (never emitted)

    def logprobs(self, wav):
        # input "wav": float32 [1, T] at 16 kHz -> output "logprobs": [1, frames, vocab]
        return self.sess.run(None, {"wav": np.asarray(wav, dtype=np.float32)[None]})[0][0]

    def greedy(self, lp):
        out, prev = [], -1
        for p in lp.argmax(-1):
            if p != prev and p > 3:
                out.append(self.pieces[p])
            prev = p
        return " ".join("".join(out).split())

    def transcribe(self, wav):
        if len(wav) < SR // 10:
            return ""
        return self.greedy(self.logprobs(wav))


def load_wav(path):
    import soundfile as sf
    y, sr = sf.read(path, dtype="float32", always_2d=True)
    y = y.mean(1)
    if sr != SR:
        from math import gcd
        from scipy.signal import resample_poly
        g = gcd(sr, SR)
        y = resample_poly(y, SR // g, sr // g).astype(np.float32)
    return y


def transcribe_long(asr, y, win=20.0):
    """Audio over ~20 s: cut at the quietest 100 ms in the last 4 s of each window, transcribe pieces, join."""
    if len(y) <= win * SR:
        return asr.transcribe(y)
    out, s = [], 0
    while s < len(y):
        e = min(len(y), s + int(win * SR))
        if e < len(y):
            lo = s + int((win - 4) * SR)
            fr = y[lo:e][: (e - lo) // 1600 * 1600].reshape(-1, 1600)
            e = lo + int(np.argmin((fr ** 2).mean(1))) * 1600 + 800
        out.append(asr.transcribe(y[s:e]))
        s = e
    return " ".join(t for t in out if t)


if __name__ == "__main__":
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("wav", nargs="+")
    ap.add_argument("--model", default=None, help="path to model_int8.onnx (default) or model_fp32.onnx")
    ap.add_argument("--vocab", default=None)
    ap.add_argument("--threads", type=int, default=4)
    a = ap.parse_args()
    asr = PocketASR(a.model, a.vocab, a.threads)
    for p in a.wav:
        y = load_wav(p)
        t = time.perf_counter()
        txt = transcribe_long(asr, y)
        dt = time.perf_counter() - t
        print(f"{p}\t{txt}\t[{len(y) / SR:.1f} s audio, {dt * 1000:.0f} ms, RTF {dt / max(len(y) / SR, 1e-3):.3f}]")