#!/usr/bin/env python3
"""Gate du pari Whisper : décode devhard-lin (439 clips, locuteurs held-out) avec le modèle
fine-tuné et compare aux références de la campagne.
Cibles à battre : 0.3534 (joint_cont greedy) et surtout **0.3457** = la config DU RECORD
(joint_cont + KenLM beam a=0.5 b=1.0 lsb=True + casse copiée du greedy).
Mesure aussi le taux de virgules produit (le déficit mesuré coûte +0.0144).
"""
import json, os
import jiwer, soundfile as sf, torch
from transformers import WhisperForConditionalGeneration, WhisperProcessor

MODEL = os.environ.get("MODEL", "/scratch/runs/whisper_lin/final")
LANG = os.environ.get("LANG_ASR", "lin")
WLANG = {"lin": "ln", "sna": "sn"}[LANG]
BEAMS = int(os.environ.get("BEAMS", "5"))
NOREP = int(os.environ.get("NOREP", "0"))
REPPEN = float(os.environ.get("REPPEN", "1.0"))
AUD = "/root/devhard_audio"
SR = 16000


def comb(refs, hyps):
    pr = [(r, h) for r, h in zip(refs, hyps) if r.strip()]
    a = [x for x, _ in pr]; b = [y for _, y in pr]
    w = jiwer.wer(a, b); c = jiwer.cer(a, b)
    return w, c, 0.5 * w + 0.5 * c


def main():
    rows = [json.loads(l) for l in open("/root/devhard/devhard_linsna.jsonl", encoding="utf-8")]
    sub = [r for r in rows if r["lang"] == LANG]
    for r in sub:
        r["audio"] = os.path.join(AUD, os.path.basename(r["audio"]))
    sub = [r for r in sub if os.path.exists(r["audio"])]
    refs = [r["text"] for r in sub]
    print("devhard-%s : %d clips | modele %s | beams %d" % (LANG, len(sub), MODEL, BEAMS), flush=True)

    proc = WhisperProcessor.from_pretrained(MODEL, language=WLANG, task="transcribe")
    m = WhisperForConditionalGeneration.from_pretrained(MODEL, dtype=torch.bfloat16).cuda().eval()
    m.generation_config.language = WLANG
    m.generation_config.task = "transcribe"
    m.generation_config.forced_decoder_ids = None

    hyps = []
    with torch.inference_mode():
        for i in range(0, len(sub), 8):
            b = sub[i:i + 8]
            au = [sf.read(r["audio"], dtype="float32")[0] for r in b]
            au = [a.mean(1) if a.ndim > 1 else a for a in au]
            x = proc(au, sampling_rate=SR, return_tensors="pt").input_features.cuda().to(torch.bfloat16)
            gk = dict(num_beams=BEAMS, max_new_tokens=200, language=WLANG, task="transcribe")
            if NOREP:
                gk["no_repeat_ngram_size"] = NOREP
            if REPPEN != 1.0:
                gk["repetition_penalty"] = REPPEN
            g = m.generate(x, **gk)
            hyps += [" ".join(t.strip().split()) for t in proc.batch_decode(g, skip_special_tokens=True)]
            if (i + 8) % 80 == 0:
                print("  %d/%d" % (min(i + 8, len(sub)), len(sub)), flush=True)

    w, c, k = comb(refs, hyps)
    print("\n=== WHISPER-FT %s ===" % LANG, flush=True)
    print("WER %.4f | CER %.4f | COMBINE %.4f" % (w, c, k), flush=True)
    for name, ref in (("joint_cont greedy", 0.3534), ("CONFIG DU RECORD (KenLM beam)", 0.3457)):
        print("  vs %-32s %.4f : %+.4f %s" % (name, ref, k - ref, "✅ MIEUX" if k < ref else "❌"), flush=True)
    nr = sum(t.count(",") for t in refs); nh = sum(t.count(",") for t in hyps)
    print("virgules : refs %d | whisper %d  (notre CTC en produisait 20)" % (nr, nh), flush=True)
    print("\nexemples :", flush=True)
    for i in range(3):
        print("  REF : %s" % refs[i][:95], flush=True)
        print("  HYP : %s" % hyps[i][:95], flush=True)
    json.dump({"wer": w, "cer": c, "combine": k, "commas": nh},
              open("/root/whisper_eval_%s.json" % LANG, "w"))
    print("WHISPER_EVAL_DONE", flush=True)


if __name__ == "__main__":
    main()