"""Dynamic ASR: transcribe Nigerian English, Pidgin, Igbo, Yoruba, Hausa and code-switched speech.

  pip install "omnilingual-asr>=0.2.0" huggingface_hub soundfile scipy
  python transcribe.py audio.wav [more.wav ...] [--model real|codeswitch] [--no-latin-mask]

Model: Meta omniASR_CTC_1B_v2 (Apache-2.0) fine-tuned by Modular. Greedy CTC decoding, no language hint needed.
Output: lowercase text without punctuation (Igbo without tone marks; Yoruba with tone marks).
"""
import argparse, unicodedata
import numpy as np, torch, torch.nn.functional as F
import soundfile as sf
from scipy.signal import resample_poly
from huggingface_hub import hf_hub_download
from omnilingual_asr.models.inference.pipeline import ASRInferencePipeline  # noqa: F401 (registers omniASR cards)
from fairseq2.models.hub import load_model
from fairseq2.data.tokenizers.hub import load_tokenizer
from fairseq2.nn.batch_layout import BatchLayout

REPO = "Modularcomputing/dynamic-asr"
FILES = {"real": "real/model.pt", "codeswitch": "codeswitch-experimental/model.pt"}
BASE = "omniASR_CTC_1B_v2"


def latin_ok(s):
    for c in s:
        cat = unicodedata.category(c)
        if cat[0] == "M":
            if not ("\u0300" <= c <= "\u036f"):
                return False
        elif c.isalpha() and "LATIN" not in unicodedata.name(c, ""):
            return False
    return True


class DynamicASR:
    def __init__(self, model="real", device="cuda" if torch.cuda.is_available() else "cpu", latin_mask=True, ckpt=None):
        self.dev = torch.device(device)
        dtype = torch.bfloat16 if self.dev.type == "cuda" else torch.float32
        self.model = load_model(BASE, device=self.dev, dtype=dtype).eval()   # base weights, then overwritten
        ck = torch.load(ckpt or hf_hub_download(REPO, FILES[model]), map_location="cpu", weights_only=False)
        self.model.load_state_dict({k: v.to(dtype) for k, v in ck["model"].items()})
        self.dtype = dtype
        self.dec = load_tokenizer(BASE).create_decoder(skip_special_tokens=True)
        V = self.model.final_proj.weight.shape[0]
        self.bad = torch.tensor([i for i in range(V) if not latin_ok(str(self.dec(torch.tensor([i]))))], device=self.dev) if latin_mask else None

    @torch.no_grad()
    def __call__(self, wav, sr=16000):
        y = np.asarray(wav, dtype=np.float32)
        if y.ndim > 1: y = y.mean(1)
        if sr != 16000: y = resample_poly(y, 16000, sr).astype(np.float32)
        x = torch.from_numpy(y); x = F.layer_norm(x, x.shape)[None].to(self.dev, self.dtype)
        logits, lay = self.model(x, BatchLayout(x.shape, seq_lens=[x.shape[1]], device=self.dev))
        if self.bad is not None: logits[..., self.bad] = -float("inf")
        s = logits[0, : lay.seq_lens[0]].argmax(-1); m = torch.ones_like(s, dtype=torch.bool); m[1:] = s[1:] != s[:-1]
        return str(self.dec(s[m]))


if __name__ == "__main__":
    ap = argparse.ArgumentParser(); ap.add_argument("audio", nargs="+"); ap.add_argument("--model", default="real", choices=list(FILES))
    ap.add_argument("--no-latin-mask", action="store_true"); ap.add_argument("--ckpt", default=None, help="local model.pt instead of the Hub")
    a = ap.parse_args()
    asr = DynamicASR(a.model, latin_mask=not a.no_latin_mask, ckpt=a.ckpt)
    for f in a.audio:
        w, sr = sf.read(f, dtype="float32"); print(f"{f}\t{asr(w, sr)}")