"""Score ASR candidates: WER/CER per language (lowercased, punctuation stripped), diacritic-folded CER, latency, streaming
partial stability, and how well the mix estimate survives ASR (TVD between mix(ref text) and mix(hyp)). ~/venv/bin/python asr_score.py"""
import json, glob, re, unicodedata, statistics as st, os, sys
import jiwer
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import mix as M
D = "/scratch/duplex/asr_eval"
R = {json.loads(l)["id"]: json.loads(l) for l in open(f"{D}/list.jsonl")}
def norm(s):
    s = unicodedata.normalize("NFC", s.lower()); s = re.sub(r"[^\w\s']", " ", s); return re.sub(r"\s+", " ", s).strip()
def fold(s):
    return "".join(c for c in unicodedata.normalize("NFD", s) if unicodedata.category(c) != "Mn")
def q(x, p):
    x = sorted(x); return x[min(len(x) - 1, int(p * len(x)))] if x else None
out = {}
for f in sorted(glob.glob(f"{D}/hyp_*.jsonl")):
    if len(sys.argv) > 1 and sys.argv[1] not in f: continue
    tag = f.split("hyp_")[1][:-6]; H = [json.loads(l) for l in open(f)]
    if len(H) < len(R) * 0.9: print("incomplete", tag, len(H)); continue
    res = {}
    for lg in ["eng", "pcm", "ibo", "yor", "hau", "cs"]:
        S = [h for h in H if R[h["id"]]["lang"] == lg]
        refs = [norm(R[h["id"]]["text"]) for h in S]; hyps = [norm(h["hyp"]) or "x" for h in S]
        res[lg] = dict(n=len(S), wer=round(jiwer.wer(refs, hyps), 3), cer=round(jiwer.cer(refs, hyps), 3),
                       cer_nodiac=round(jiwer.cer([fold(r) for r in refs], [fold(h) for h in hyps]), 3))
    lat = [h["lat"] for h in H]; rtf = [h["lat"] / h["dur"] for h in H]
    res["latency"] = dict(p50_s=round(q(lat, .5), 3), p90_s=round(q(lat, .9), 3), rtf_p50=round(q(rtf, .5), 4))
    P = [h for h in H if h.get("partials")]
    if P:   # stability: share of final words already present (in order prefix) in the partial 1.28 s before the end
        stab, plat = [], []
        for h in P:
            fin = norm(h["hyp"]).split(); parts = h["partials"]
            if len(parts) >= 2 and fin:
                pw = norm(parts[-2][1]).split(); k = 0
                while k < min(len(pw), len(fin)) and pw[k] == fin[k]: k += 1
                stab.append(k / max(1, len(pw)))
            plat += [p[2] for p in parts]
        res["streaming"] = dict(n=len(P), partial_prefix_stability=round(st.mean(stab), 3) if stab else None, partial_call_p50_s=round(q(plat, .5), 4), partial_call_p90_s=round(q(plat, .9), 4))
    # mix estimate after ASR (CS clips): TVD(mix(ref), mix(hyp)), rules-only and with GlotLID
    tv = {}
    for lid in ("none", "glotlid"):
        v = []
        for h in H:
            r = R[h["id"]]
            if r["lang"] != "cs": continue
            ref_mix = {}
            for l in r["wlang"]: ref_mix[l] = ref_mix.get(l, 0) + 1 / len(r["wlang"])
            v.append(M.tvd(ref_mix, M.estimate(h["hyp"], lid=lid, asr=True)["mix"]))
        tv[lid] = round(st.mean(v), 3) if v else None
        mono = [M.estimate(h["hyp"], lid=lid, asr=True)["mix"].get(R[h["id"]]["lang"], 0) for h in H if R[h["id"]]["lang"] in ("ibo", "yor", "hau", "pcm")]
        tv[lid + "_mono_true_share"] = round(st.mean(mono), 3)
    res["cs_mix_tvd_after_asr"] = tv
    out[tag] = res
# LID on reference text (upper bound for the mix estimate)
tv = {}
for lid in ("none", "glotlid"):
    v = []; acc = []
    for r in R.values():
        if r["lang"] != "cs": continue
        ref_mix = {}
        for l in r["wlang"]: ref_mix[l] = ref_mix.get(l, 0) + 1 / len(r["wlang"])
        e = M.estimate(r["text"], lid=lid); v.append(M.tvd(ref_mix, e["mix"]))
        if len(e["labels"]) == len(r["wlang"]): acc += [a == b for a, b in zip(e["labels"], r["wlang"])]
    tv[lid] = dict(mix_tvd=round(st.mean(v), 3), word_acc=round(sum(acc) / len(acc), 3) if acc else None, n_words=len(acc))
out["_lid_on_reference_text"] = tv
json.dump(out, open(f"{D}/asr_scores.json", "w"), indent=1); print(json.dumps(out, indent=1))