"""Standalone CTC-submodel transcription on the dosing audio.

Loads Canary, accesses model.timestamps_asr_model directly, calls its own
transcribe() — bypassing the AED decoder entirely. Goal: see whether the CTC
submodel renders the spoken digit sequence "eins null eins" as separate
tokens or also collapses it.
"""
import json
import subprocess
import tempfile

import torch

from nemo.collections.asr.models import EncDecMultiTaskModel


WAV_SRC = "/sweep/case_c414cf52_15-01-39.m4a"
OUT_PATH = "/sweep/ctc_standalone_dosing.json"


def transcode(src: str) -> str:
    dst = tempfile.NamedTemporaryFile(suffix=".wav", delete=False)
    dst.close()
    subprocess.run(["ffmpeg", "-y", "-loglevel", "error", "-i", src,
                    "-ac", "1", "-ar", "16000", "-f", "wav", dst.name],
                   check=True)
    return dst.name


def to_jsonable(o):
    if hasattr(o, "tolist"):
        return o.tolist()
    if hasattr(o, "item"):
        return o.item()
    return str(o)


def main():
    print("loading Canary 1B v2 (just for the CTC submodel)", flush=True)
    model = EncDecMultiTaskModel.from_pretrained("nvidia/canary-1b-v2")
    model = model.to("cuda").to(torch.bfloat16)
    model.eval()

    tdt = model.timestamps_asr_model
    print(f"submodel: {type(tdt).__name__}", flush=True)

    wav = transcode(WAV_SRC)
    print(f"running tdt.transcribe(timestamps=True) on {WAV_SRC}", flush=True)
    with torch.amp.autocast("cuda", enabled=True, dtype=torch.bfloat16):
        with torch.no_grad():
            hyps = tdt.transcribe([wav], timestamps=True)
    h = hyps[0]

    out = {"text": getattr(h, "text", None),
           "timestamp": getattr(h, "timestamp", None)}
    with open(OUT_PATH, "w") as f:
        json.dump(out, f, default=to_jsonable, ensure_ascii=False, indent=2)

    print(f"\nstandalone CTC text:\n{h.text}\n", flush=True)
    ts = h.timestamp if isinstance(h.timestamp, dict) else {}
    for k, v in ts.items():
        n = len(v) if hasattr(v, "__len__") else "?"
        print(f"timestamp[{k}]: {n} entries", flush=True)
    print(f"\nfull dump -> {OUT_PATH}", flush=True)


if __name__ == "__main__":
    main()