"""Transcribe audio files with a pruned omniASR CTC model (greedy CTC, no language model).

    python inference.py audio1.wav [audio2.flac ...] [--model-dir .] [--device cuda]
"""

from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path

import soundfile as sf
import torch
import torchaudio.functional as AF
from safetensors.torch import load_file

sys.path.insert(0, str(Path(__file__).resolve().parent))  # the bundled omniprune package
from omniprune.model import ArchConfig, PrunableCTC, normalize_audio_
from omniprune.text import CharVocab


def load(model_dir: str | Path = ".", device: str = "cpu") -> tuple[PrunableCTC, CharVocab]:
    d = Path(model_dir)
    model = PrunableCTC(ArchConfig.from_dict(json.loads((d / "config.json").read_text())))
    model.load_state_dict(load_file(str(d / "model.safetensors")))
    return model.eval().to(device), CharVocab.load(d / "vocab.json")


def read_audio(path: str) -> torch.Tensor:
    wav, sr = sf.read(path, dtype="float32", always_2d=True)
    wav = torch.from_numpy(wav).mean(1)
    return AF.resample(wav, sr, 16000) if sr != 16000 else wav


@torch.no_grad()
def transcribe(model: PrunableCTC, vocab: CharVocab, paths: list[str]) -> list[str]:
    device = next(model.parameters()).device
    out = []
    for p in paths:
        wav = read_audio(p)[None].to(device)
        lens = torch.tensor([wav.shape[1]], device=device)
        normalize_audio_(wav, lens)
        with torch.autocast(device.type, dtype=torch.bfloat16, enabled=device.type == "cuda"):
            res = model(wav, lens)
        ids = res["logits"][0, : int(res["frame_lens"][0])].argmax(-1)
        out.append(vocab.greedy(ids.tolist()))
    return out


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("audio", nargs="+")
    ap.add_argument("--model-dir", default=str(Path(__file__).parent))
    ap.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu")
    args = ap.parse_args()
    model, vocab = load(args.model_dir, args.device)
    for p, text in zip(args.audio, transcribe(model, vocab, args.audio)):
        print(f"{p}\t{text}")


if __name__ == "__main__":
    main()