"""AFFINAGE du decodage KenLM (modele FIXE = joint_cont sur lin). Logits en cache.
Etape 1: grille fine alpha x beta (o5, beam64). Etape 2: beam large au meilleur point.
Etape 3: ordre 4 au meilleur point. Etape 4: sna a alpha tres faible."""
import json,pickle,jiwer,itertools
from transformers import AutoProcessor
from pyctcdecode import build_ctcdecoder
from multiprocessing import Pool
ROWS=[json.loads(l) for l in open("/root/devhard/devhard_linsna.jsonl")]
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
def setup(mdl,pkl,lang):
    SUB=[r for r in ROWS if r["lang"]==lang]
    LOG=pickle.load(open(pkl,"rb"))
    tok=AutoProcessor.from_pretrained(mdl).tokenizer
    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]=""
    g=[" ".join(tok.decode(l.argmax(-1)).replace("|"," ").split()) for l in LOG]
    return [r["text"] for r in SUB],LOG,lab,tok,g
def run(lab,LOG,tok,greedy,refs,arpa,a,b,bw,casecopy=True):
    dec=build_ctcdecoder(lab,kenlm_model_path=arpa,alpha=a,beta=b)
    with Pool(8) as p: h=dec.decode_batch(p,LOG,beam_width=bw)
    h=[" ".join(x.split()) for x in h]
    if casecopy: h=[(g[:1]+x[1:] if x and g else x) for x,g in zip(h,greedy)]
    return comb(refs,h)

print("########## LIN ##########",flush=True)
refs,LOG,lab,tok,greedy=setup("/root/models/joint_cont_best","/scratch/lm/logits_lin.pkl","lin")
w,c,g0=comb(refs,greedy); print(f"greedy: {g0:.4f}",flush=True)
best=(9,None)
print("--- etape 1: grille alpha x beta (o5, beam64, casse-greedy) ---",flush=True)
for a in [0.3,0.4,0.5,0.6,0.7]:
    row=[]
    for b in [0.0,0.5,1.0,1.5,2.0]:
        _,_,m=run(lab,LOG,tok,greedy,refs,"/scratch/lm/lin_5g.arpa",a,b,64)
        row.append(f"b{b}={m:.4f}")
        if m<best[0]: best=(m,(5,a,b,64))
    print(f"  a={a}: "+"  ".join(row),flush=True)
print(f"  >>> meilleur etape1: {best[0]:.4f} {best[1]}",flush=True)
o,a,b,_=best[1]
print("--- etape 2: beam large ---",flush=True)
for bw in [128,256]:
    _,_,m=run(lab,LOG,tok,greedy,refs,"/scratch/lm/lin_5g.arpa",a,b,bw)
    print(f"  beam={bw}: {m:.4f}",flush=True)
    if m<best[0]: best=(m,(5,a,b,bw))
print("--- etape 3: ordre 4 et 3 au meilleur point ---",flush=True)
for oo in [4,3]:
    _,_,m=run(lab,LOG,tok,greedy,refs,f"/scratch/lm/lin_{oo}g.arpa",a,b,best[1][3])
    print(f"  o={oo}: {m:.4f}",flush=True)
    if m<best[0]: best=(m,(oo,a,b,best[1][3]))
print(f"BEST_LIN {best[0]:.4f} cfg={best[1]} (greedy {g0:.4f}, gain {g0-best[0]:+.4f})",flush=True)
json.dump({"combine":best[0],"cfg":best[1],"greedy":g0},open("/root/refine_lin.json","w"))

print("\n########## SNA (alpha tres faible) ##########",flush=True)
refs2,LOG2,lab2,tok2,greedy2=setup("/root/models/sna_ps_best","/scratch/lm/logits_sna_sna_ps_best.pkl","sna")
w,c,gs=comb(refs2,greedy2); print(f"greedy sna: {gs:.4f}",flush=True)
bs=(gs,"greedy")
for a2 in [0.05,0.1,0.2]:
    for b2 in [0.0,1.0]:
        _,_,m=run(lab2,LOG2,tok2,greedy2,refs2,"/scratch/lm/sna_5g.arpa",a2,b2,64)
        print(f"  a={a2} b={b2}: {m:.4f}"+("  <<<" if m<bs[0] else ""),flush=True)
        if m<bs[0]: bs=(m,(5,a2,b2))
print(f"BEST_SNA {bs[0]:.4f} cfg={bs[1]}",flush=True)
json.dump({"combine":bs[0],"cfg":str(bs[1]),"greedy":gs},open("/root/refine_sna.json","w"))
print("REFINE_DONE",flush=True)