"""Beam + KenLM (LM = texte WAXAL train seulement). Logits mis en cache -> sweep rapide."""
import json, os, pickle, numpy as np, soundfile as sf, torch, jiwer
from transformers import AutoModelForCTC, AutoProcessor
from pyctcdecode import build_ctcdecoder
from multiprocessing import Pool

LANG=os.environ.get("LG","lin"); MDL=os.environ.get("MDL","/root/models/joint_cont_best")
ROWS=[json.loads(l) for l in open("/root/devhard/devhard_linsna.jsonl")]
SUB=[r for r in ROWS if r["lang"]==LANG]; REFS=[r["text"] for r in SUB]
LMDIR="/scratch/lm"; CACHE=f"/scratch/lm/logits_{LANG}_{os.path.basename(MDL)}.pkl"

def comb(a,b):
    pr=[(x,y) for x,y in zip(a,b) if x.strip()]
    A=[x for x,_ in pr]; B=[y for _,y in pr]
    w=jiwer.wer(A,B); c=jiwer.cer(A,B); return w,c,0.5*w+0.5*c

proc=AutoProcessor.from_pretrained(MDL); tok=proc.tokenizer
if os.path.exists(CACHE):
    LOG=pickle.load(open(CACHE,"rb")); print("logits (cache):",len(LOG),flush=True)
else:
    model=AutoModelForCTC.from_pretrained(MDL,dtype=torch.float32).cuda().eval()
    LOG=[]
    with torch.inference_mode():
        for i in range(0,len(SUB),4):
            b=SUB[i:i+4]
            au=[sf.read(r["audio"],dtype="float32")[0] for r in b]
            x=proc(au,sampling_rate=16000,return_tensors="pt",padding=True)
            x={k:v.cuda() for k,v in x.items()}
            lg=model(**x).logits.log_softmax(-1).cpu().numpy().astype(np.float32)
            n=[len(a) for a in au]
            for j in range(len(b)): LOG.append(lg[j])
    del model; torch.cuda.empty_cache()
    pickle.dump(LOG,open(CACHE,"wb")); print("logits calcules:",len(LOG),flush=True)

greedy=[" ".join(tok.decode(l.argmax(-1)).replace("|"," ").split()) for l in LOG]
w,c,gm=comb(REFS,greedy); print(f"GREEDY {LANG}: combine={gm:.4f} (WER {w:.4f} CER {c:.4f})",flush=True)

# ---- alphabet pyctcdecode : UN SEUL "" (blank) ----
v=tok.get_vocab(); lab=[None]*len(v)
for t,i in v.items(): lab[i]=t
lab[tok.word_delimiter_token_id]=" "
lab[tok.unk_token_id]="⁇"
lab[tok.pad_token_id]=""
assert len(lab)==len(set(lab)), "doublons restants"
print("alphabet OK, taille",len(lab),flush=True)

best=(gm,"greedy",None)
for o in [3,4,5]:
    arpa=f"{LMDIR}/{LANG}_{o}g.arpa"
    if not os.path.exists(arpa): continue
    for alpha in [0.3,0.5,0.8]:
        for beta in [0.0,1.0]:
            dec=build_ctcdecoder(lab,kenlm_model_path=arpa,alpha=alpha,beta=beta)
            with Pool(8) as p:
                hyps=dec.decode_batch(p,LOG,beam_width=64)
            hyps=[" ".join(h.split()) for h in hyps]
            w,c,mm=comb(REFS,hyps)
            flag=" <<<" if mm<best[0] else ""
            print(f"  o={o} a={alpha} b={beta}: combine={mm:.4f} (WER {w:.4f} CER {c:.4f}){flag}",flush=True)
            if mm<best[0]: best=(mm,f"o{o}_a{alpha}_b{beta}",(o,alpha,beta))
print(f"BEST {LANG}: {best[0]:.4f} <- {best[1]}   (greedy {gm:.4f}, gain {gm-best[0]:+.4f})",flush=True)
json.dump({"best":best[0],"cfg":best[1],"params":best[2],"greedy":gm},open(f"/root/kenlm_best_{LANG}.json","w"))
print("KENLM_SWEEP2_DONE",flush=True)