#!/usr/bin/env python3
"""Soumission avec RESCORING N-BEST sur le lingala :
  lin : beams KenLM de joint_cont (+ hypothese decode_batch) reordonnes par
        score = ac_cont + lm_kenlm + LAMBDA * ac_cont2   (LAMBDA par env, defaut 1.5)
  sna : sna_ps greedy
Casse du 1er caractere copiee du greedy. Routage par LANGF.
"""
import csv
import glob
import json
import os

import numpy as np
import soundfile as sf
import torch
from multiprocessing import Pool
from pyctcdecode import build_ctcdecoder
from transformers import AutoModelForCTC, AutoProcessor

SR = 16000
CACHE = "/scratch/p2_16k"
M1 = "/root/models/joint_cont_best"
M2 = "/root/models/joint_cont2_best"
ARPA = os.environ.get("ARPA", "/scratch/lm/lin_5g.arpa")
LAMBDA = float(os.environ.get("LAMBDA", "1.5"))
GAMMA = float(os.environ.get("GAMMA", "2.0"))   # bonus par mot (compense le biais des sommes de log-probs)
NBEST = 10
LANGF = os.environ.get("LANGF", "/root/test_lang_gpulid.json")
OUT = os.environ.get("OUT", "/root/sub_rescore.csv")


def norm(s):
    return " ".join(str(s).replace("|", " ").split())


def batches(sel, budget):
    d = {f: sf.info(f).duration for f in sel}
    sel = sorted(sel, key=lambda f: -d[f])
    bs, cur, acc = [], [], 0.0
    for f in sel:
        if cur and acc + d[f] > budget:
            bs.append(cur)
            cur, acc = [], 0.0
        cur.append(f)
        acc += d[f]
    if cur:
        bs.append(cur)
    return bs


def logits_for(model_dir, files, dtype=torch.float32):
    proc = AutoProcessor.from_pretrained(model_dir)
    m = AutoModelForCTC.from_pretrained(model_dir, dtype=dtype).cuda().eval()
    res = {}
    with torch.inference_mode():
        for b in batches(files, 90):
            au = [sf.read(f, dtype="float32")[0] for f in b]
            x = proc(au, sampling_rate=SR, return_tensors="pt", padding=True)
            x = {k: v.cuda() for k, v in x.items()}
            lg = m(**x).logits.log_softmax(-1).float().cpu().numpy()
            for j, f in enumerate(b):
                res[f] = lg[j]
    del m
    torch.cuda.empty_cache()
    return proc, res


def ctc_scores(logp, texts, tok):
    T = logp.shape[0]
    lp = torch.from_numpy(logp).unsqueeze(1)
    out = []
    for t in texts:
        ids = [i for i in tok(t.replace(" ", "|")).input_ids if i != tok.pad_token_id] if t else []
        if not ids or len(ids) > T:
            out.append(-1e9)
            continue
        loss = torch.nn.functional.ctc_loss(
            lp, torch.tensor(ids).unsqueeze(0), torch.tensor([T]), torch.tensor([len(ids)]),
            blank=tok.pad_token_id, reduction="sum", zero_infinity=True)
        out.append(-float(loss))
    return out


def main():
    lang = json.load(open(LANGF))
    files = sorted(glob.glob(os.path.join(CACHE, "*.wav")))
    ids = [os.path.splitext(os.path.basename(f))[0] for f in files]
    lin = [f for f in files if lang[os.path.splitext(os.path.basename(f))[0]] == "lin"]
    sna = [f for f in files if lang[os.path.splitext(os.path.basename(f))[0]] == "sna"]
    print("lin=%d (rescoring lambda=%.1f gamma=%.1f) | sna=%d (sna_ps greedy)"
          % (len(lin), LAMBDA, GAMMA, len(sna)), flush=True)
    out = {}

    proc1, LG1 = logits_for(M1, lin)
    print("logits joint_cont OK", flush=True)
    _, LG2 = logits_for(M2, lin)
    print("logits joint_cont2 OK", flush=True)

    tok = proc1.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] = ""
    dec = build_ctcdecoder(lab, kenlm_model_path=ARPA, alpha=0.5, beta=0.5,
                           lm_score_boundary=False)

    order = lin
    L1 = [LG1[f] for f in order]
    with Pool(8) as p:
        allbeams = 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)]

    for i, f in enumerate(order):
        g = norm(tok.decode(L1[i].argmax(-1)))
        bs = allbeams[i]
        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]]
        if db[i] not in cands:
            cands.append(db[i])
            ac1.append(ctc_scores(L1[i], [db[i]], tok)[0])
            lm.append(float(np.mean(lm)) if lm else 0.0)
        ac2 = ctc_scores(LG2[f], cands, tok)
        nw = np.array([float(len(x.split())) for x in cands])
        tot = np.array(ac1) + np.array(lm) + LAMBDA * np.array(ac2) + GAMMA * nw
        h = norm(cands[int(np.argmax(tot))])
        if h and g:
            h = g[:1] + h[1:]
        out[os.path.splitext(os.path.basename(f))[0]] = h
        if (i + 1) % 150 == 0:
            print("  rescore %d/%d" % (i + 1, len(order)), flush=True)
    print("lin OK", flush=True)

    proc2 = AutoProcessor.from_pretrained("/root/models/sna_ps_best")
    m2 = AutoModelForCTC.from_pretrained("/root/models/sna_ps_best",
                                         dtype=torch.bfloat16).cuda().eval()
    with torch.inference_mode():
        for b in batches(sna, 140):
            au = [sf.read(f, dtype="float32")[0] for f in b]
            x = proc2(au, sampling_rate=SR, return_tensors="pt", padding=True)
            x = {k: v.to("cuda", dtype=torch.bfloat16 if v.dtype == torch.float32 else v.dtype)
                 for k, v in x.items()}
            pid = m2(**x).logits.float().argmax(-1).cpu().numpy()
            for f, s in zip(b, proc2.batch_decode(pid)):
                out[os.path.splitext(os.path.basename(f))[0]] = norm(s)
    del m2
    torch.cuda.empty_cache()
    print("sna OK", flush=True)

    fb = {}
    ref = "/root/sub_p2_KENLM.csv"
    if os.path.exists(ref):
        fb = {r["ID"]: r["Target"] for r in csv.DictReader(open(ref, encoding="utf-8"))}
    filled = 0
    for i in ids:
        if not out.get(i, "").strip() and fb.get(i, "").strip():
            out[i] = fb[i]
            filled += 1
    with open(OUT, "w", newline="", encoding="utf-8") as fo:
        w = csv.writer(fo)
        w.writerow(["ID", "Target"])
        for i in ids:
            w.writerow([i, out.get(i) or "a"])
    print("RESCORE_GEN_DONE %s | %d IDs | vides=%d | combles=%d"
          % (OUT, len(ids), sum(1 for i in ids if not out.get(i, "").strip()), filled), flush=True)


if __name__ == "__main__":
    main()