#!/usr/bin/env python3
"""Count the listener's errors on what run.mjs heard.

  python3 tools/asr/eval/score.py RESULTS.json [SUMMARY.json] [--no-end NO_END.json]

For each language: the word error rate (words wrong, left out or added, per
100 words of the references), or for Japanese and Korean the character
error rate, over all its clips together. Both sides are normalised alike
first: Unicode NFKC, lower case, every punctuation mark and symbol made a
space, single spaces; and these spelling variants that do not change a word:
Hindi nukta dropped and chandrabindu written as anusvara (FLEURS's
references and the model disagree on them), Arabic short vowels, tatweel
and the hamza forms of alef written as bare alef. For characters, spaces
are left out. Numbers written as digits on one side and as words on the
other count as errors, as they do for every language alike.

Prints a table, and writes the numbers to SUMMARY.json when given. With
--no-end, the same clips heard with no end-of-speech rule at all, it also
counts the clips the rule cut at a pause: those where what was heard is
less than 70% of what was heard with no end.

Standard library only (Python 3.9 or later).
"""
import json
import sys
import unicodedata

CER = {"ja", "ko", "zh"}
ARABIC_MARKS = {chr(c) for c in range(0x064B, 0x0653)} | {"ٰ", "ـ"}
ARABIC_ALEF = str.maketrans({"آ": "ا", "أ": "ا", "إ": "ا", "ٱ": "ا"})


def norm(s, code):
    s = unicodedata.normalize("NFKC", s).lower()
    if code == "hi":
        s = s.replace("़", "").replace("ँ", "ं")
    if code == "ar":
        s = "".join(ch for ch in s if ch not in ARABIC_MARKS).translate(ARABIC_ALEF)
    s = "".join(" " if unicodedata.category(ch)[0] in "PS" else ch for ch in s)
    return " ".join(s.split())


def edits(a, b):
    prev = list(range(len(b) + 1))
    for i in range(1, len(a) + 1):
        cur = [i] + [0] * len(b)
        for j in range(1, len(b) + 1):
            cur[j] = min(prev[j] + 1, cur[j - 1] + 1, prev[j - 1] + (a[i - 1] != b[j - 1]))
        prev = cur
    return prev[-1]


def score(code, clips):
    errors = total = 0
    for c in clips:
        ref, hyp = norm(c["ref"], code), norm(c["text"], code)
        if code in CER:
            r, h = list(ref.replace(" ", "")), list(hyp.replace(" ", ""))
        else:
            r, h = ref.split(), hyp.split()
        errors += edits(r, h)
        total += len(r)
    return 100 * errors / max(total, 1)


def cut_at_pause(clips, no_end):
    """Clips the end-of-speech rule ended at a pause: much less heard than
    with no end at all."""
    return sum(1 for c, q in zip(clips, no_end) if c["ended_by"] == "silence" and len(c["text"]) < 0.7 * len(q["text"]))


def main(argv):
    no_end = None
    if "--no-end" in argv:
        i = argv.index("--no-end")
        no_end = json.load(open(argv[i + 1], encoding="utf-8"))["languages"]
        argv = argv[:i] + argv[i + 2:]
    if len(argv) < 2:
        print(__doc__, file=sys.stderr)
        return 2
    results = json.load(open(argv[1], encoding="utf-8"))
    summary = {}
    print("| language | clips | measure | errors per 100 | cut at a pause | audio s |")
    print("|---|---|---|---|---|---|")
    for code, lang in sorted(results["languages"].items()):
        clips = lang["clips"]
        rate = round(score(code, clips), 2)
        measure = "cer" if code in CER else "wer"
        cut = cut_at_pause(clips, no_end[code]["clips"]) if no_end and code in no_end else None
        seconds = round(sum(c["seconds"] for c in clips), 1)
        summary[code] = {"measure": measure, "rate": rate, "clips": len(clips), "seconds": seconds,
                         "source": lang["source"], "slot": lang["slot"]}
        if cut is not None:
            summary[code]["cut_at_pause"] = cut
        print(f"| {code} | {len(clips)} | {measure.upper()} | {rate:.2f} | {'-' if cut is None else cut} | {seconds} |")
    if len(argv) > 2:
        with open(argv[2], "w", encoding="utf-8") as f:
            json.dump({"date": results["date"], "runtime": results["runtime"], "languages": summary}, f, indent=1, ensure_ascii=False)
            f.write("\n")
    return 0


if __name__ == "__main__":
    sys.exit(main(sys.argv))