#!/usr/bin/env python3
"""Decodage Phase 2 ROUTE PAR LANGUE (100% modeles pre-entraines, 0 dataset externe).
LID (mms-lid-256, contraint a lin/lug/sna) -> lug=Sunbird Whisper-salt ; lin/sna=notre MMS CTC.
Garde-fous anti-hallucination Whisper -> fallback MMS. Sortie CSV Zindi.
"""
import argparse, csv, glob, os
from collections import Counter
import soundfile as sf, torch
from transformers import (AutoModelForCTC, AutoProcessor,
                          WhisperForConditionalGeneration, WhisperProcessor,
                          AutoModelForAudioClassification, AutoFeatureExtractor)
SR = 16000


def norm(s):
    return " ".join(str(s).replace("|", " ").split())


def is_halluc(hyp, dur):
    w = hyp.split()
    if not w:
        return True
    if len(w) >= 6 and Counter(w).most_common(1)[0][1] / len(w) > 0.5:
        return True
    if dur > 0 and len(w) / dur > 6.0:
        return True
    return False


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--whisper_model", default="Sunbird/asr-whisper-large-v3-salt")
    ap.add_argument("--mms_model", required=True)
    ap.add_argument("--lid_model", default="facebook/mms-lid-256")
    ap.add_argument("--audio_dir", default="/root/phase2_audio/audio")
    ap.add_argument("--test_csv", default="/root/Test_phase2.csv")
    ap.add_argument("--out", required=True)
    ap.add_argument("--lang_token", type=int, default=50355)  # Luganda dans salt
    ap.add_argument("--targets", default="lin,lug,sna")
    ap.add_argument("--chunk_s", type=float, default=28.0)
    ap.add_argument("--beams", type=int, default=1)
    ap.add_argument("--no_repeat", type=int, default=0)
    a = ap.parse_args()
    dev = "cuda"

    lid = AutoModelForAudioClassification.from_pretrained(a.lid_model).to(dev).eval()
    lidfe = AutoFeatureExtractor.from_pretrained(a.lid_model)
    id2label = lid.config.id2label
    label2id = {v: k for k, v in id2label.items()}
    tgt = [t for t in a.targets.split(",")]
    missing = [t for t in tgt if t not in label2id]
    assert not missing, f"LID labels manquants {missing} ; ex labels={list(id2label.values())[:8]}"
    tgt_ids = torch.tensor([label2id[t] for t in tgt])

    mms = AutoModelForCTC.from_pretrained(a.mms_model, torch_dtype=torch.float16).to(dev).eval()
    mmsp = AutoProcessor.from_pretrained(a.mms_model)
    wp = WhisperProcessor.from_pretrained(a.whisper_model)
    wm = WhisperForConditionalGeneration.from_pretrained(a.whisper_model, torch_dtype=torch.float16).to(dev).eval()
    tr = wp.tokenizer.convert_tokens_to_ids("<|transcribe|>")
    nt = wp.tokenizer.convert_tokens_to_ids("<|notimestamps|>")
    forced = [[1, a.lang_token], [2, tr], [3, nt]]

    test_ids = [r["ID"] for r in csv.DictReader(open(a.test_csv, encoding="utf-8"))]

    def mms_decode(au):
        x = mmsp(au, sampling_rate=SR, return_tensors="pt")
        x = {k: (v.to(dev, torch.float16) if v.dtype == torch.float32 else v.to(dev)) for k, v in x.items()}
        ids = mms(**x).logits.argmax(-1).cpu().numpy()
        return norm(mmsp.batch_decode(ids)[0])

    def whisper_decode(au):
        win = int(a.chunk_s * SR)
        pieces = [au[j:j + win] for j in range(0, len(au), win)] or [au]
        out = []
        for p in pieces:
            f = wp(p, sampling_rate=SR, return_tensors="pt").input_features.to(dev, torch.float16)
            ids = wm.generate(f, forced_decoder_ids=forced, max_new_tokens=220, num_beams=a.beams, no_repeat_ngram_size=a.no_repeat)
            out.append(wp.batch_decode(ids, skip_special_tokens=True)[0])
        return norm(" ".join(out))

    out = {}; nlug = 0; nfb = 0; langcount = Counter()
    with torch.inference_mode():
        for i, tid in enumerate(test_ids):
            f = os.path.join(a.audio_dir, tid + ".wav")
            if not os.path.exists(f):
                cand = glob.glob(os.path.join(a.audio_dir, tid + ".*"))
                f = cand[0] if cand else None
            if not f:
                out[tid] = "a"; continue
            au = sf.read(f, dtype="float32")[0]; dur = len(au) / SR
            inp = lidfe(au[:SR * 20], sampling_rate=SR, return_tensors="pt")
            inp = {k: v.to(dev) for k, v in inp.items()}
            logits = lid(**inp).logits[0]
            lang = tgt[int(torch.argmax(logits[tgt_ids]).item())]
            langcount[lang] += 1
            if lang == "lug":
                nlug += 1
                h = whisper_decode(au)
                if is_halluc(h, dur):
                    h = mms_decode(au); nfb += 1
            else:
                h = mms_decode(au)
            out[tid] = h or "a"
            if (i + 1) % 200 == 0:
                print(f"{i+1}/{len(test_ids)} langs={dict(langcount)} fb={nfb}", flush=True)
    empt = sum(1 for t in test_ids if not out.get(t, "").strip())
    with open(a.out, "w", newline="", encoding="utf-8") as fo:
        w = csv.writer(fo); w.writerow(["ID", "Target"])
        for t in test_ids:
            w.writerow([t, out.get(t) or "a"])
    print(f"ROUTED_DONE {a.out} | {len(test_ids)} IDs | langs={dict(langcount)} "
          f"halluc_fallback={nfb} vides={empt}", flush=True)


if __name__ == "__main__":
    main()