#!/usr/bin/env python3
"""Juge final ground-truth-free : décode N clips avec les 3 modèles w2v-BERT + affiche
les 2 prédictions LID. La sortie COHÉRENTE (vrais mots) révèle la vraie langue."""
import glob, json
import soundfile as sf, torch
from transformers import (AutoModelForCTC, AutoProcessor,
                          Wav2Vec2ForSequenceClassification, AutoFeatureExtractor,
                          Wav2Vec2BertForSequenceClassification, SeamlessM4TFeatureExtractor)
SR=16000; N=10
files=sorted(glob.glob("/root/phase2_audio/audio/*.wav"))[:N]
def norm(s): return " ".join(str(s).replace("|"," ").split())

# LID interne
fe2=SeamlessM4TFeatureExtractor.from_pretrained("/root/models/lid_best")
lid=Wav2Vec2BertForSequenceClassification.from_pretrained("/root/models/lid_best",torch_dtype=torch.bfloat16).cuda().eval()
i2l=json.load(open("/root/models/lid_best/config.json"))["id2label"]
our={}
with torch.inference_mode():
    for f in files:
        w=sf.read(f,dtype="float32")[0][:20*SR]
        x=fe2([w],sampling_rate=SR,return_tensors="pt",padding=True)
        x={k:v.to("cuda",dtype=torch.bfloat16 if v.dtype==torch.float32 else v.dtype) for k,v in x.items()}
        our[f]=i2l[str(int(lid(**x).logits.float().argmax(-1)))]
del lid; torch.cuda.empty_cache()
# MMS-LID
fe3=AutoFeatureExtractor.from_pretrained("facebook/mms-lid-256")
mms=Wav2Vec2ForSequenceClassification.from_pretrained("facebook/mms-lid-256",torch_dtype=torch.bfloat16).cuda().eval()
ml2i={str(v):int(k) for k,v in mms.config.id2label.items()}
idx=torch.tensor([ml2i["lin"],ml2i["lug"],ml2i["sna"]]).cuda(); order=["lin","lug","sna"]
mmsp={}
with torch.inference_mode():
    for f in files:
        w=sf.read(f,dtype="float32")[0][:20*SR]
        x=fe3([w],sampling_rate=SR,return_tensors="pt")
        x={k:v.to("cuda",dtype=torch.bfloat16 if v.dtype==torch.float32 else v.dtype) for k,v in x.items()}
        mmsp[f]=order[int(mms(**x).logits.float()[0][idx].argmax())]
del mms; torch.cuda.empty_cache()
# 3 transcriptions
CTC={"lin":"/root/models/lin_s4_best","sna":"/root/models/sna_ps_best","lug":"/root/models/lug_ps_best"}
hyp={f:{} for f in files}
for lang,path in CTC.items():
    proc=AutoProcessor.from_pretrained(path); m=AutoModelForCTC.from_pretrained(path,torch_dtype=torch.bfloat16).cuda().eval()
    with torch.inference_mode():
        for f in files:
            w=sf.read(f,dtype="float32")[0]
            x=proc(w,sampling_rate=SR,return_tensors="pt",padding=True)
            x={k:v.to("cuda",dtype=torch.bfloat16 if v.dtype==torch.float32 else v.dtype) for k,v in x.items()}
            ids=m(**x).logits.float().argmax(-1).cpu().numpy()
            hyp[f][lang]=norm(proc.batch_decode(ids)[0])
    del m; torch.cuda.empty_cache()
for f in files:
    print(f"\n### {f.split('/')[-1]}  | LID_interne={our[f]}  MMS-LID={mmsp[f]}")
    for lang in ["lin","sna","lug"]:
        print(f"  [{lang}] {hyp[f][lang][:90]}")
print("\nDECODE_DONE")