#!/usr/bin/env python3
"""WHISPER COMME RESCOREUR DES N-BEST DU CTC (angle jamais testé).
Principe §4c : un rescoreur n'a pas besoin du même vocabulaire, seulement d'évaluer
log P(texte | audio). Les rescoreurs testés jusqu'ici étaient soit du TEXTE SEUL
(charLM -> échec), soit des CTC de la MÊME famille (cont2 -> n'aide que le shona).
Whisper est le seul à la fois ANCRÉ DANS L'AUDIO et doté d'un vrai modèle de langue
(décodeur autorégressif, contexte phrase entière) => signal réellement décorrélé.
Cible : la marge d'oracle lin (0.3158 vs 0.3457 = -0.030) qu'aucune méthode n'a entamée.

score(h) = ac_ctc(h) + lm_kenlm(h) + w2*ac_cont2(h) + mu*logP_whisper(h|audio)
Gate : devhard-lin, config du RECORD (alpha=0.5, beta=1.0, lsb=True).
"""
import json, os, pickle
import jiwer, numpy as np, soundfile as sf, torch
from multiprocessing import Pool
from pyctcdecode import build_ctcdecoder
from transformers import (AutoModelForCTC, AutoProcessor,
                          WhisperForConditionalGeneration, WhisperProcessor)

M1 = "/root/models/joint_cont_best"
ARPA = "/scratch/lm/lin_5g.arpa"
WM = os.environ.get("WMODEL", "/scratch/runs/whisper_lin_v2/final")
R = "/scratch/restore"
NBEST = int(os.environ.get("NBEST", "10"))
AUD = "/root/devhard_audio"
SR = 16000


def comb(refs, hyps):
    pr = [(r, h) for r, h in zip(refs, hyps) if r.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 encode_for(tok, text):
    v = tok.get_vocab()
    d = getattr(tok, "word_delimiter_token", "|")
    s = text.replace(" ", d)
    keep = "".join(c for c in s if c in v)
    if not keep:
        keep = "".join(c for c in text.lower().replace(" ", d) if c in v)
    return [v[c] for c in keep if v[c] != tok.pad_token_id]


def ctc_score(logp, ids, blank):
    T = logp.shape[0]
    if not ids or len(ids) > T:
        return -1e9
    lp = torch.from_numpy(logp).unsqueeze(1)
    return -float(torch.nn.functional.ctc_loss(
        lp, torch.tensor(ids).unsqueeze(0), torch.tensor([T]), torch.tensor([len(ids)]),
        blank=blank, reduction="sum", zero_infinity=True))


def main():
    rows = [json.loads(l) for l in open("/root/devhard/devhard_linsna.jsonl", encoding="utf-8")]
    sub = [r for r in rows if r["lang"] == "lin"]
    for r in sub:
        r["audio"] = os.path.join(AUD, os.path.basename(r["audio"]))
    sub = [r for r in sub if os.path.exists(r["audio"])]
    refs = [r["text"] for r in sub]
    L1 = pickle.load(open("/scratch/lm/logits_lin.pkl", "rb"))
    tok = AutoProcessor.from_pretrained(M1).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] = ""
    greedy = [" ".join(tok.decode(l.argmax(-1)).replace("|", " ").split()) for l in L1]
    cc = lambda h, g: (g[:1] + h[1:]) if (h and g) else h

    dec = build_ctcdecoder(lab, kenlm_model_path=ARPA, alpha=0.5, beta=1.0, lm_score_boundary=True)
    with Pool(8) as p:
        allb = dec.decode_beams_batch(p, L1, beam_width=64)
    with Pool(8) as p:
        db = [" ".join(x.split()) for x in dec.decode_batch(p, L1, beam_width=64)]
    REF = comb(refs, [cc(h, g) for h, g in zip(db, greedy)])[2]
    print("REFERENCE (config record) : %.4f" % REF, flush=True)

    cands, AC1, LMS = [], [], []
    for i, bs in enumerate(allb):
        c = [" ".join(b[0].split()) for b in bs[:NBEST]]
        a = [(b[3] if len(b) > 3 else 0.0) for b in bs[:NBEST]]
        l = [((b[4] - b[3]) if len(b) > 4 else 0.0) for b in bs[:NBEST]]
        if db[i] not in c:
            c.append(db[i])
            a.append(ctc_score(L1[i], encode_for(tok, db[i]), tok.pad_token_id))
            l.append(float(np.mean(l)) if l else 0.0)
        cands.append(c); AC1.append(np.array(a)); LMS.append(np.array(l))
    orc = [min(cands[i], key=lambda h: comb([refs[i]], [h])[2] if refs[i].strip() else 0)
           for i in range(len(cands))]
    print("ORACLE %d-best : %.4f (marge %+.4f)" % (NBEST, comb(refs, orc)[2],
                                                   comb(refs, orc)[2] - REF), flush=True)

    # --- rescoreur CTC cont2 (référence connue) ---
    SC = {}
    proc2 = AutoProcessor.from_pretrained(R + "/joint_cont2_best")
    m2 = AutoModelForCTC.from_pretrained(R + "/joint_cont2_best", dtype=torch.float32).cuda().eval()
    LG = []
    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 = proc2(au, sampling_rate=SR, return_tensors="pt", padding=True)
            x = {k: vv.cuda() for k, vv in x.items()}
            lgt = m2(**x).logits.log_softmax(-1).float().cpu().numpy()
            for j in range(len(b)):
                LG.append(lgt[j])
    t2 = proc2.tokenizer
    SC["cont2"] = [np.array([ctc_score(LG[i], encode_for(t2, x), t2.pad_token_id)
                             for x in cands[i]]) for i in range(len(cands))]
    del m2; torch.cuda.empty_cache()
    print("cont2 OK", flush=True)

    # --- WHISPER comme rescoreur : log P(texte | audio) par teacher forcing ---
    wp = WhisperProcessor.from_pretrained(WM, language="ln", task="transcribe")
    wm = WhisperForConditionalGeneration.from_pretrained(WM, dtype=torch.float32).cuda().eval()
    WS = []
    with torch.inference_mode():
        for i, r in enumerate(sub):
            au = sf.read(r["audio"], dtype="float32")[0]
            if au.ndim > 1:
                au = au.mean(1)
            feat = wp.feature_extractor(au, sampling_rate=SR, return_tensors="pt").input_features.cuda()
            enc = wm.model.encoder(feat)
            sc_i = []
            for h in cands[i]:
                ids = wp.tokenizer(h, return_tensors="pt").input_ids.cuda()
                out = wm(encoder_outputs=enc, decoder_input_ids=ids[:, :-1])
                lp = out.logits.log_softmax(-1)
                tgt = ids[:, 1:]
                sc_i.append(float(lp.gather(-1, tgt.unsqueeze(-1)).squeeze(-1).sum()))
            WS.append(np.array(sc_i))
            if (i + 1) % 80 == 0:
                print("  whisper %d/%d" % (i + 1, len(sub)), flush=True)
    print("whisper OK", flush=True)

    def ev(w2=0.0, mu=0.0):
        hyps = []
        for i in range(len(cands)):
            tot = AC1[i] + LMS[i] + w2 * SC["cont2"][i] + mu * WS[i]
            hyps.append(cc(cands[i][int(np.argmax(tot))], greedy[i]))
        return comb(refs, hyps)[2]

    print("\n--- Whisper SEUL comme rescoreur (mu) ---", flush=True)
    best_mu = (9, 0)
    for mu in (0.05, 0.1, 0.2, 0.4, 0.8, 1.5, 3.0):
        m = ev(0.0, mu)
        if m < best_mu[0]:
            best_mu = (m, mu)
        print("  mu=%.2f : %.4f (%+.4f)%s" % (mu, m, m - REF, "  <-- GAIN" if m < REF else ""), flush=True)
    print("\n--- cont2 seul (rappel) ---", flush=True)
    best_w = (9, 0)
    for w2 in (1.0, 2.5):
        m = ev(w2, 0.0)
        if m < best_w[0]:
            best_w = (m, w2)
        print("  w2=%.1f : %.4f (%+.4f)" % (w2, m, m - REF), flush=True)
    print("\n--- cont2 + Whisper ---", flush=True)
    best = (9, None, None)
    for w2 in (0.0, 1.0, 2.5):
        for mu in (0.0, 0.05, 0.1, 0.2, 0.4, 0.8):
            m = ev(w2, mu)
            if m < best[0]:
                best = (m, w2, mu)
    print("  BEST %.4f (cont2=%s mu=%s)  %+.4f vs record" % (best[0], best[1], best[2], best[0] - REF), flush=True)
    json.dump({"ref": REF, "best": best[0], "w2": best[1], "mu": best[2],
               "whisper_solo": list(best_mu)}, open("/root/whisper_rescore.json", "w"))
    print("WHISPER_RESCORE_DONE", flush=True)


if __name__ == "__main__":
    main()