"""Greedy-CTC inference for anshulsc/w2v-bert-ethio-v5.

Pipeline matches the evaluation harness exactly:
  1. 16 kHz mono audio -> SeamlessM4TFeatureExtractor (from facebook/w2v-bert-2.0)
  2. Wav2Vec2BertForCTC forward pass -> argmax over logits
  3. Mask the six [LID:*] token ids (311-316) to the CTC blank id
  4. Wav2Vec2CTCTokenizer.decode (collapses CTC repeats, maps | to spaces)

Needs: torch transformers soundfile librosa
"""

import sys

import soundfile as sf
import torch
from transformers import SeamlessM4TFeatureExtractor, Wav2Vec2BertForCTC, Wav2Vec2CTCTokenizer

MODEL_ID = "anshulsc/w2v-bert-ethio-v5"
LID_TOKEN_IDS = [311, 312, 313, 314, 315, 316]  # [LID:*] tokens must not survive decoding


def transcribe(audio_path: str) -> str:
    speech, sr = sf.read(audio_path, dtype="float32")
    if speech.ndim > 1:  # downmix stereo to mono
        speech = speech.mean(axis=1)

    feature_extractor = SeamlessM4TFeatureExtractor.from_pretrained("facebook/w2v-bert-2.0")
    if sr != feature_extractor.sampling_rate:
        import librosa

        speech = librosa.resample(speech, orig_sr=sr, target_sr=feature_extractor.sampling_rate)

    inputs = feature_extractor(
        speech,
        sampling_rate=feature_extractor.sampling_rate,
        return_tensors="pt",
        return_attention_mask=True,
    )

    device = "cuda" if torch.cuda.is_available() else "cpu"
    model = Wav2Vec2BertForCTC.from_pretrained(MODEL_ID, torch_dtype=torch.float16).to(device).eval()
    inputs = {k: v.to(device) for k, v in inputs.items()}

    tokenizer = Wav2Vec2CTCTokenizer.from_pretrained(MODEL_ID)
    with torch.no_grad():
        logits = model(**inputs).logits
    pred_ids = torch.argmax(logits, dim=-1)

    lid = torch.tensor(LID_TOKEN_IDS, device=pred_ids.device)
    pred_ids = pred_ids.masked_fill(torch.isin(pred_ids, lid), tokenizer.pad_token_id)

    return tokenizer.decode(pred_ids[0]).strip()


if __name__ == "__main__":
    print(transcribe(sys.argv[1]))