"""omniASR candidates on the eval set: transcripts + single-clip latency + CTC streaming (prefix re-decode) sim.
~/venv_omni/bin/python asr_omni.py <card> -> /scratch/duplex/asr_eval/hyp_<card>.jsonl"""
import json, sys, time, numpy as np, soundfile as sf, torch
from omnilingual_asr.models.inference.pipeline import ASRInferencePipeline
LC = {"ibo": "ibo_Latn", "yor": "yor_Latn", "hau": "hau_Latn", "pcm": "pcm_Latn", "eng": "eng_Latn"}
card = sys.argv[1]; D = "/scratch/duplex/asr_eval"
R = [json.loads(l) for l in open(f"{D}/list.jsonl")]
p = ASRInferencePipeline(model_card=card)
ctc = "CTC" in card
def tr(w, lang):
    kw = {} if ctc or lang is None else dict(lang=[LC[lang]])
    return p.transcribe([{"waveform": w, "sample_rate": 16000}], batch_size=1, **kw)[0]
w0 = sf.read(R[0]["wav"], dtype="float32")[0]
for _ in range(3): tr(w0, "eng")
out = open(f"{D}/hyp_{card}.jsonl", "w")
for r in R:
    w = sf.read(r["wav"], dtype="float32")[0]
    lang = r["lang"] if r["lang"] != "cs" else None     # LLM cards: CS clips get no language hint
    torch.cuda.synchronize(); t = time.perf_counter(); h = tr(w, lang); torch.cuda.synchronize(); dt = time.perf_counter() - t
    rec = dict(id=r["id"], hyp=h, lat=round(dt, 4), dur=r["dur"])
    if ctc and r["dur"] > 4 and int(r["id"][-3:]) % 4 == 0:     # streaming sim: re-decode growing prefix every 0.64 s
        parts = []
        for e in np.arange(0.64, r["dur"], 0.64):
            t = time.perf_counter(); hp = tr(w[: int(e * 16000)], None); parts.append((round(float(e), 2), hp, round(time.perf_counter() - t, 4)))
        rec["partials"] = parts
    out.write(json.dumps(rec, ensure_ascii=False) + "\n"); out.flush()
print("OMNI_DONE", card)