#!/usr/bin/env python3
"""Vérifie s'il y a du LUGANDA dans la nouvelle data mal routé par le LID.
Décode N clips avec les 3 champions, confiance CTC (mean max-softmax) → langue argmax,
compare au LID. Si des clips 'gagnent' en lug alors que LID dit lin/sna → LID rate le lug."""
import glob, json
import soundfile as sf, torch, numpy as np, subprocess
from transformers import (AutoModelForCTC, AutoProcessor,
                          Wav2Vec2BertForSequenceClassification, SeamlessM4TFeatureExtractor)
SR=16000; N=120
CH={"lin":"/root/models/mms1b_lin_best","sna":"/root/models/sna_ps_best","lug":"/root/models/lug_ps_best"}
files=sorted(glob.glob("/scratch/p2new/newaudios/*.wav"))[:N]

def load16(f):
    w,sr=sf.read(f,dtype="float32")
    if sr==SR: return w
    p=subprocess.run(["ffmpeg","-v","error","-i",f,"-f","f32le","-ac","1","-ar",str(SR),"pipe:1"],capture_output=True)
    return np.frombuffer(p.stdout,dtype=np.float32)

# LID
fe=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"]
wavs={f:load16(f) for f in files}
lidp={}; lidprob={}
with torch.inference_mode():
    for f in files:
        x=fe([wavs[f][:20*SR]],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()}
        lo=lid(**x).logits.float()[0]; p=torch.softmax(lo,-1)
        lidp[f]=i2l[str(int(lo.argmax()))]; lidprob[f]={i2l[str(j)]:round(float(p[j]),3) for j in range(len(p))}
del lid; torch.cuda.empty_cache()
# confiance CTC
conf={f:{} for f in files}
for lang,path in CH.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:
            x=proc(wavs[f],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()}
            conf[f][lang]=float(m(**x).logits.float().softmax(-1)[0].max(-1).values.mean())
    del m; torch.cuda.empty_cache()
ctc={f:max(conf[f],key=conf[f].get) for f in files}
from collections import Counter
print("N=",len(files))
print("LID  distribution:",dict(Counter(lidp.values())))
print("CTC-conf distribution:",dict(Counter(ctc.values())))
print("clips ou CTC dit LUG:",sum(1 for f in files if ctc[f]=="lug"))
print("proba lug médiane (LID):",round(float(np.median([lidprob[f]['lug'] for f in files])),3),
      "| max:",round(float(np.max([lidprob[f]['lug'] for f in files])),3))
# clips ou lug a la plus haute confiance CTC -> candidats luganda ratés
lugc=[f for f in files if ctc[f]=="lug"]
print("\nclips CTC=lug (candidats luganda), LID dit:",dict(Counter(lidp[f] for f in lugc)))
for f in lugc[:5]: print("  ",f.split("/")[-1],"conf:",{k:round(conf[f][k],2) for k in CH},"LIDprob:",lidprob[f])
print("CHECK_DONE")