"""Torch-free real-time streaming with the ONNX tick graphs — pure onnxruntime + numpy.

Two graphs, both **strictly causal** (transformer runs frame-by-frame internally — no future audio) and
**decision-equivalent to offline** (match the batch model to ~0.02-0.03 in fp32; bit-identical to each other):
  • `stream_tick.onnx`     — 1 frame  per call ([2, 1920] @ 24 kHz = 80 ms),  ~47 ms/call on CPU
  • `stream_tick_240.onnx` — 3 frames per call ([2, 5760] @ 24 kHz = 240 ms), ~71 ms/call on CPU

Running the 240 ms graph is ~2× cheaper per second of audio than calling the 80 ms graph three times
(one shared encoder pass instead of three), at the cost of only acting every 240 ms. Pick by how often you
need to react vs how much CPU you want to spend.

    from onnx_streaming import DualTurnONNXStreamer, DualTurnONNXStreamer240

    s = DualTurnONNXStreamer()                    # 80 ms; auto-downloads stream_tick.onnx
    for chunk in chunks_80ms:                     # each [2, 1920] float32 @ 24 kHz (CH0=user, CH1=agent)
        out = s.push(chunk)                       # -> {"eot": float, "vad": [2], "fvad": [2,4]}
        # fire your end-of-turn policy on out["eot"] (threshold + short sustain)
    s.reset()

    s = DualTurnONNXStreamer240()                 # 240 ms; auto-downloads stream_tick_240.onnx
    for chunk in chunks_240ms:                     # each [2, 5760] float32 @ 24 kHz
        frames = s.push(chunk)                    # -> list of 3 {"eot","vad","fvad"} (the 3 × 80 ms frames)
        for fr in frames: ...                     # fr["eot"], fr["vad"], fr["fvad"]
    s.reset()

The graphs carry all state (conv buffer, transformer KV window, downsample history, LSTM state) internally;
you just pass the per-tick state dict back in. No PyTorch / transformers needed.
"""
import numpy as np

try:
    import onnxruntime as ort
except ImportError as e:
    raise ImportError("DualTurn ONNX streaming needs onnxruntime: pip install onnxruntime") from e


class _BaseStreamer:
    RF = 7680; NL = 8; H = 8; HD = 64; TFK = 260          # state shapes (fixed by the exported graph)
    ONNX = None; N_FRAMES = 1; N_SAMPLES = 1920           # overridden per subclass

    def __init__(self, onnx_path=None, repo="anyreach-ai/dualturn-endpointing", threads=4):
        if onnx_path is None:
            from huggingface_hub import hf_hub_download
            onnx_path = hf_hub_download(repo, self.ONNX)
        so = ort.SessionOptions(); so.intra_op_num_threads = threads
        so.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
        self.sess = ort.InferenceSession(onnx_path, sess_options=so, providers=["CPUExecutionProvider"])
        self.reset()

    def reset(self):
        z = np.zeros
        self.state = {"audio_ctx": z((2, 0), np.float32),                       # EMPTY: grow from true start (no
                      "pk": z((self.NL, 2, self.H, 0, self.HD), np.float32),   #  zero-prefix; graph caps at RF). KV likewise grows/caps in-graph
                      "pv": z((self.NL, 2, self.H, 0, self.HD), np.float32),
                      "tf_hist": z((2, self.TFK, 512), np.float32),
                      "h": z((2, 1, 256), np.float32), "c": z((2, 1, 256), np.float32),
                      "pos": np.array([0], np.int64)}

    def _run(self, chunk):
        chunk = np.asarray(chunk, dtype=np.float32)
        if chunk.ndim == 1:
            chunk = np.stack([chunk, np.zeros_like(chunk)])
        if chunk.shape != (2, self.N_SAMPLES):
            raise ValueError(f"chunk must be [2,{self.N_SAMPLES}] @24kHz ({self.N_FRAMES*80} ms); got {chunk.shape}")
        eot, vad, fvad, ac, pk, pv, th, h, c, pos = self.sess.run(
            None, {"chunk": chunk, "audio_ctx": self.state["audio_ctx"], "pk": self.state["pk"],
                   "pv": self.state["pv"], "tf_hist": self.state["tf_hist"], "h": self.state["h"],
                   "c": self.state["c"], "pos": self.state["pos"]})
        self.state = {"audio_ctx": ac, "pk": pk, "pv": pv, "tf_hist": th, "h": h, "c": c, "pos": pos}
        return eot, vad, fvad


class DualTurnONNXStreamer(_BaseStreamer):
    """80 ms tick. push([2,1920]) -> {"eot": float, "vad": [2], "fvad": [2,4]}."""
    ONNX = "stream_tick.onnx"; N_FRAMES = 1; N_SAMPLES = 1920

    def push(self, chunk):
        eot, vad, fvad = self._run(chunk)
        return {"eot": float(eot[0, 0, 0]), "vad": vad[0, 0], "fvad": fvad[0, 0]}


class DualTurnONNXStreamer240(_BaseStreamer):
    """240 ms tick (~2× cheaper/sec of audio). push([2,5760]) -> list of 3 {"eot","vad","fvad"} frames."""
    ONNX = "stream_tick_240.onnx"; N_FRAMES = 3; N_SAMPLES = 5760

    def push(self, chunk):
        eot, vad, fvad = self._run(chunk)
        return [{"eot": float(eot[0, j, 0]), "vad": vad[0, j], "fvad": fvad[0, j]} for j in range(self.N_FRAMES)]