"""Insertion CHIRURGICALE de virgules : on GARDE tout de lASR (casse + points),
on ajoute UNE virgule seulement si le modele est tres confiant ET quil ny a pas deja
de ponctuation a cet endroit. Balayage de seuils, mesure offline sur devhard."""
import json, re, torch, jiwer
from transformers import AutoTokenizer, AutoModelForTokenClassification
M="/scratch/ftruns/punct_mbert"
tok=AutoTokenizer.from_pretrained(M)
model=AutoModelForTokenClassification.from_pretrained(M).cuda().eval()
I2L=model.config.id2label
COMMA_IDS=[int(i) for i,l in I2L.items() if l.split("|")[1]=="COMMA"]
D=json.load(open("/root/devhard_joint_hyps.json",encoding="utf-8"))

def comma_probs(words):
    enc=tok([words],is_split_into_words=True,truncation=True,max_length=128,return_tensors="pt").to("cuda")
    with torch.inference_mode(): lg=model(**enc).logits[0].softmax(-1)
    wids=enc.word_ids(batch_index=0); seen=set(); out=[]
    for j,wid in enumerate(wids):
        if wid is not None and wid not in seen:
            out.append(float(sum(lg[j][c] for c in COMMA_IDS))); seen.add(wid)
    return out+[0.0]*(len(words)-len(out))

def insert(text,thr):
    toks=text.split()
    if len(toks)<3: return text
    bare=[re.sub(r"[\".,!?;:()]+$","",t) for t in toks]
    ps=comma_probs(bare)
    out=[]
    for i,t in enumerate(toks):
        # ne touche pas au dernier mot ni si ponctuation deja presente
        if i<len(toks)-1 and not re.search(r"[.,!?;:]$",t) and i<len(ps) and ps[i]>=thr:
            out.append(t+",")
        else: out.append(t)
    return " ".join(out)

def comb(R,H):
    pr=[(r,h) for r,h in zip(R,H) if r.strip()]
    r=[x for x,_ in pr]; h=[x for _,x in pr]
    return 0.5*jiwer.wer(r,h)+0.5*jiwer.cer(r,h)

R=[d["ref"] for d in D]; H=[d["hyp"] for d in D]
print(f"BASELINE combine={comb(R,H):.4f}")
cache={}
for thr in [0.5,0.6,0.7,0.8,0.9,0.95]:
    Hn=[insert(h,thr) for h in H]
    c=comb(R,Hn)
    added=sum(hn.count(",")-h.count(",") for h,hn in zip(H,Hn))
    print(f"  thr={thr:.2f} combine={c:.4f}  (virgules ajoutees={added})",flush=True)
print("COMMA_SWEEP_DONE")