"""SOUMISSION HYBRIDE : LID(lid_best) -> lin=lin_s4_best, sna=sna_ps_best. Clips 16k caches."""
import csv, glob, os, json, soundfile as sf, torch
from transformers import (AutoModelForAudioClassification, AutoFeatureExtractor,
                          AutoModelForCTC, AutoProcessor)
SR=16000; CACHE="/scratch/p2_16k"
files=sorted(glob.glob(os.path.join(CACHE,"*.wav")))
ids=[os.path.splitext(os.path.basename(f))[0] for f in files]
def norm(s): return " ".join(str(s).replace("|"," ").split())

# ---- 1) LID ----
fe=AutoFeatureExtractor.from_pretrained("/root/models/lid_best")
lm=AutoModelForAudioClassification.from_pretrained("/root/models/lid_best",dtype=torch.bfloat16).cuda().eval()
i2l=lm.config.id2label
lang={}
with torch.inference_mode():
    for i in range(0,len(files),8):
        b=files[i:i+8]
        au=[sf.read(f,dtype="float32")[0][:SR*20] for f in b]
        x=fe(au,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()}
        pr=lm(**x).logits.float().softmax(-1)
        for f,p in zip(b,pr):
            k=int(p.argmax()); lab=i2l[k] if k in i2l else i2l[str(k)]
            if lab=="lug": lab="lin"     # lug absent du test
            lang[os.path.splitext(os.path.basename(f))[0]]=lab
del lm; torch.cuda.empty_cache()
import collections; print("LID:",dict(collections.Counter(lang.values())),flush=True)

# ---- 2) decode par specialiste ----
out={}
for lg,mdl in [("lin","/root/models/lin_s4_best"),("sna","/root/models/sna_ps_best")]:
    sel=[f for f in files if lang[os.path.splitext(os.path.basename(f))[0]]==lg]
    if not sel: continue
    proc=AutoProcessor.from_pretrained(mdl)
    m=AutoModelForCTC.from_pretrained(mdl,dtype=torch.bfloat16).cuda().eval()
    durs={f:sf.info(f).duration for f in sel}
    sel.sort(key=lambda f:-durs[f])
    bs,cur,bud=[],[],0.0
    for f in sel:
        if cur and bud+durs[f]>140: bs.append(cur); cur,bud=[],0.0
        cur.append(f); bud+=durs[f]
    if cur: bs.append(cur)
    with torch.inference_mode():
        for b in bs:
            au=[sf.read(f,dtype="float32")[0] for f in b]
            x=proc(au,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()}
            pid=m(**x).logits.float().argmax(-1).cpu().numpy()
            for f,s in zip(b,proc.batch_decode(pid)):
                out[os.path.splitext(os.path.basename(f))[0]]=norm(s)
    del m; torch.cuda.empty_cache()
    print(f"{lg}: {len(sel)} clips decodes avec {os.path.basename(mdl)}",flush=True)

OUT="/root/sub_p2_hybride.csv"
with open(OUT,"w",newline="",encoding="utf-8") as fo:
    w=csv.writer(fo); w.writerow(["ID","Target"])
    for i in ids: w.writerow([i,out.get(i) or "a"])
emp=sum(1 for i in ids if not out.get(i,"").strip())
print(f"HYBRIDE_DONE {OUT} | {len(ids)} IDs | vides={emp}",flush=True)
for i in ids[:3]: print("  ",i,":",out.get(i,"")[:70])