import json, numpy as np, soundfile as sf, torch, jiwer
from transformers import Wav2Vec2BertForCTC, Wav2Vec2BertProcessor
SR=16000
proc=Wav2Vec2BertProcessor.from_pretrained("/root/models/joint_best")
model=Wav2Vec2BertForCTC.from_pretrained("/root/models/joint_best",torch_dtype=torch.bfloat16).cuda().eval()
mono={"lin":0.283,"lug":0.101,"sna":0.177}
@torch.inference_mode()
def transcribe(rows):
    rows=sorted(rows,key=lambda r:-r["duration"]); out={}
    b=[];bud=0;bs=[]
    for r in rows:
        if b and bud+r["duration"]>140: bs.append(b);b=[];bud=0
        b.append(r);bud+=r["duration"]
    if b:bs.append(b)
    for bb in bs:
        au=[sf.read(r["audio"],dtype="float32")[0] for r in bb]
        f=proc.feature_extractor(au,sampling_rate=SR,return_tensors="pt",padding=True)
        f={k:v.to("cuda",dtype=torch.bfloat16 if v.dtype==torch.float32 else v.dtype) for k,v in f.items()}
        ids=model(**f).logits.float().argmax(-1).cpu().numpy()
        for r,s in zip(bb,proc.tokenizer.batch_decode(ids)): out[r["id"]]=" ".join(s.replace("|"," ").split())
    return out
print("=== JOINT par langue vs meilleur monolingue ===")
for lang in ("lin","lug","sna"):
    rows=[json.loads(l) for l in open(f"/scratch/prep/manifests/waxal_{lang}_validation.jsonl") if json.loads(l)["text"].strip()]
    for r in rows: r["text"]=" ".join(r["text"].replace("|"," ").split())
    h=transcribe(rows)
    refs=[r["text"] for r in rows]; hyps=[h.get(r["id"],"") for r in rows]
    w,c=jiwer.wer(refs,hyps),jiwer.cer(refs,hyps); comb=0.5*w+0.5*c
    verdict="JOINT GAGNE" if comb<mono[lang] else "mono gagne"
    print(f"  {lang}: joint {comb:.4f} vs mono {mono[lang]:.4f} -> {verdict} ({comb-mono[lang]:+.4f})")
print("EVALJOINT_DONE")