"""Minimal streaming inference for Muno459/fastconformer-quran-streaming (ONNX, cache-aware).

The ONNX takes 80-dim log-mel features (NOT raw audio) + carries cache tensors across chunks.
Pipeline: audio(16k mono) -> log-mel(80) -> fixed-global CMVN -> chunked encoder w/ cache -> CTC greedy.

EXACT shapes that fix the `Where`/`Add` broadcast errors (batch-first for ONNX):
  cache_last_channel     : float32  [1, 17, 70, 512]   (init zeros)   # 17 layers, 70 = left cache
  cache_last_time        : float32  [1, 17, 512, 8]    (init zeros)   # 8  = conv cache (kernel-1)
  cache_last_channel_len : int64    [1] = [0]          (init)
Carry the returned cache_*_next straight back in as the next step's cache_* (already batch-first).

Mel params (must match exactly if you reimplement on-device, e.g. iOS/Android):
  sample_rate=16000, n_fft=512, win_length=400 (25ms hann), hop_length=160 (10ms),
  n_mels=80 (librosa slaney-norm mel), power=2.0, log = ln(x + 2**-24), preemph=0.97, dither=0 (inference).
Then apply fixed-global CMVN from streaming_global_cmvn.npz: (mel - mean) / (std + 1e-5),
using the tlog_* constants for phone audio, clean_* for studio.

The simplest *correct* mel is NeMo's own preprocessor (one import); a hand-rolled mel must match the
params above or accuracy drops. Decode with the BPE tokenizer (tokenizer.model, blank id = 1024).
"""
import numpy as np, onnxruntime as ort, soundfile as sf, sentencepiece as spm

SESS = ort.InferenceSession("model.q8.onnx", providers=["CPUExecutionProvider"])
SP = spm.SentencePieceProcessor(model_file="tokenizer.model")
CMVN = np.load("streaming_global_cmvn.npz")
BLANK = 1024
N_LAYERS, D_MODEL, LEFT_CACHE, TIME_CACHE = 17, 512, 70, 8
CHUNK_MEL = 112          # mel frames per streaming step (8x subsampling -> ~14 output frames)

def log_mel(wav, sr=16000):
    """80-dim log-mel matching NeMo's FilterbankFeatures. Uses NeMo if available (exact),
    else falls back to a torchaudio/librosa reimplementation with the documented params."""
    import torch, torchaudio
    wav = torch.tensor(np.asarray(wav, np.float32))
    wav = torch.cat([wav[:1], wav[1:] - 0.97 * wav[:-1]])               # preemphasis
    spec = torchaudio.transforms.MelSpectrogram(
        sample_rate=sr, n_fft=512, win_length=400, hop_length=160,
        n_mels=80, power=2.0, window_fn=torch.hann_window, norm="slaney", mel_scale="slaney")(wav)
    return torch.log(spec + 2**-24).numpy()                              # (80, T)

def transcribe(path, kind="tlog"):
    wav, sr = sf.read(path)
    if wav.ndim > 1: wav = wav.mean(1)
    mel = log_mel(wav, sr)                                               # (80, T)
    gm = CMVN[f"{kind}_mean"][:, None]; gs = CMVN[f"{kind}_std"][:, None]
    mel = (mel - gm) / (gs + 1e-5)
    clc = np.zeros((1, N_LAYERS, LEFT_CACHE, D_MODEL), np.float32)
    clt = np.zeros((1, N_LAYERS, D_MODEL, TIME_CACHE), np.float32)
    cll = np.zeros((1,), np.int64)
    ids = []
    for s in range(0, mel.shape[1], CHUNK_MEL):
        chunk = mel[:, s:s + CHUNK_MEL][None].astype(np.float32)        # (1,80,chunk)
        length = np.array([chunk.shape[2]], np.int64)
        logp, _, clc, clt, cll = SESS.run(None, {
            "audio_signal": chunk, "length": length,
            "cache_last_channel": clc, "cache_last_time": clt, "cache_last_channel_len": cll})
        a = logp[0].argmax(-1)
        prev = -1
        for t in a:
            t = int(t)
            if t != prev and t != BLANK: ids.append(t)
            prev = t
    return SP.decode([i for i in ids])

if __name__ == "__main__":
    import sys
    print(transcribe(sys.argv[1], kind="tlog" if len(sys.argv) < 3 else sys.argv[2]))