#!/usr/bin/env python3
"""RESCORING N-BEST : on prend les N meilleures hypotheses du beam KenLM sur joint_cont,
puis on les REORDONNE en ajoutant la vraisemblance CTC d'un 2e modele (joint_cont2).
Contrairement au moyennage de logits (qui avait echoue en brouillant les alignements),
ici on ne fusionne rien : on choisit parmi des hypotheses completes et coherentes.

score(h) = acoustique_cont(h) + lm_kenlm(h) + LAMBDA * acoustique_cont2(h)
Sweep de LAMBDA sur devhard-lin.
"""
import json
import pickle

import jiwer
import numpy as np
import torch
from multiprocessing import Pool
from pyctcdecode import build_ctcdecoder
from transformers import AutoProcessor

M1 = "/root/models/joint_cont_best"
M2 = "/root/models/joint_cont2_best"
ARPA = "/scratch/lm/lin_5g.arpa"
NBEST = 10


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 ctc_logprob(logp, texts, tok):
    """log P(text | logits) par l'algorithme forward CTC, pour une liste de textes."""
    T, V = logp.shape
    lp = torch.from_numpy(logp).unsqueeze(1)                  # (T,1,V)
    outs = []
    for t in texts:
        ids = tok(t.replace(" ", "|")).input_ids if t else []
        ids = [i for i in ids if i != tok.pad_token_id]
        if not ids or len(ids) > T:
            outs.append(-1e9)
            continue
        tgt = torch.tensor(ids, dtype=torch.long).unsqueeze(0)
        loss = torch.nn.functional.ctc_loss(
            lp, tgt, torch.tensor([T]), torch.tensor([len(ids)]),
            blank=tok.pad_token_id, reduction="sum", zero_infinity=True)
        outs.append(-float(loss))
    return outs


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"]
    refs = [r["text"] for r in sub]
    L1 = pickle.load(open("/scratch/lm/logits_lin.pkl", "rb"))
    L2 = pickle.load(open("/scratch/lm/logits_lin_cont2.pkl", "rb"))
    tok = AutoProcessor.from_pretrained(M1).tokenizer
    tok2 = AutoProcessor.from_pretrained(M2).tokenizer
    same_vocab = tok.get_vocab() == tok2.get_vocab()
    print("vocabulaires identiques: %s (|V1|=%d |V2|=%d)"
          % (same_vocab, len(tok.get_vocab()), len(tok2.get_vocab())), flush=True)
    if not same_vocab:
        print("ATTENTION: vocabulaires differents -> rescoring impossible tel quel", flush=True)
        return

    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]

    dec = build_ctcdecoder(lab, kenlm_model_path=ARPA, alpha=0.5, beta=0.5,
                           lm_score_boundary=False)
    with Pool(8) as p:
        allbeams = dec.decode_beams_batch(p, L1, beam_width=64)
    base = [" ".join(bs[0][0].split()) for bs in allbeams]
    base = [(g[:1] + h[1:] if h and g else h) for h, g in zip(base, greedy)]
    w0, c0, m0 = comb(refs, base)
    print("BASE (1-best via decode_beams): combine=%.4f (WER %.4f CER %.4f)" % (m0, w0, c0), flush=True)
    with Pool(8) as p:
        alt = [" ".join(x.split()) for x in dec.decode_batch(p, L1, beam_width=64)]
    alt = [(g[:1] + h[1:] if h and g else h) for h, g in zip(alt, greedy)]
    _, _, ma = comb(refs, alt)
    ident = sum(1 for a, b in zip(base, alt) if a == b)
    print("BASE (1-best via decode_batch) : combine=%.4f  | identiques %d/%d"
          % (ma, ident, len(base)), flush=True)

    # oracle : borne superieure atteignable en choisissant la meilleure des N hypotheses
    oracle = []
    for bs, g, ref in zip(allbeams, greedy, refs):
        cands = [" ".join(b[0].split()) for b in bs[:NBEST]]
        cands = [(g[:1] + c[1:] if c and g else c) for c in cands]
        best = min(cands, key=lambda c: 0.5 * jiwer.wer([ref], [c]) + 0.5 * jiwer.cer([ref], [c]))
        oracle.append(best)
    _, _, mo = comb(refs, oracle)
    print("ORACLE %d-best: %.4f  (marge theorique %+.4f)" % (NBEST, mo, mo - m0), flush=True)

    # scores : acoustique cont (logit_score), kenlm (lm_score - logit_score), acoustique cont2
    prepared = []
    for i, bs in enumerate(allbeams):
        cands = [" ".join(b[0].split()) for b in bs[:NBEST]]
        ac1 = [(b[3] if len(b) > 3 else 0.0) for b in bs[:NBEST]]
        lm = [((b[4] - b[3]) if len(b) > 4 else 0.0) for b in bs[:NBEST]]
        ac2 = ctc_logprob(L2[i], cands, tok)
        prepared.append((cands, np.array(ac1), np.array(lm), np.array(ac2)))
        if (i + 1) % 150 == 0:
            print("  rescoring prepare %d/%d" % (i + 1, len(allbeams)), flush=True)

    print("\n--- sweep LAMBDA (poids du 2e modele) ---", flush=True)
    best = (m0, 0.0)
    for lam in (0.0, 0.1, 0.2, 0.3, 0.5, 0.8, 1.0, 1.5):
        hyps = []
        for (cands, ac1, lm, ac2), g in zip(prepared, greedy):
            tot = ac1 + lm + lam * ac2
            h = cands[int(np.argmax(tot))]
            if h and g:
                h = g[:1] + h[1:]
            hyps.append(h)
        _, _, m = comb(refs, hyps)
        print("  lambda=%.1f : %.4f  (%+.4f)%s" % (lam, m, m - m0, "  <<<" if m < best[0] else ""), flush=True)
        if m < best[0]:
            best = (m, lam)
    print("\nBEST_RESCORE %.4f  lambda=%.1f  (base %.4f, oracle %.4f)"
          % (best[0], best[1], m0, mo), flush=True)
    json.dump({"combine": best[0], "lambda": best[1], "base": m0, "oracle": mo},
              open("/root/nbest_best.json", "w"))
    print("NBEST_DONE", flush=True)


if __name__ == "__main__":
    main()