#!/usr/bin/env python3
"""Score Thai ASR output the way the OpenThai2.0-ASR-Streaming model card does.

  python score.py refs.jsonl hyps_A.jsonl [hyps_B.jsonl] [--boot 1000]

refs.jsonl   one line per clip: {"id": ..., "text": <reference>, "cluster": <video/recording/speaker/meeting id>}
             (build it from the public dataset; the ids we used are the ones in hypotheses/*.jsonl)
hyps_*.jsonl one line per clip: {"id": ..., "pred_text": <model output>}; every reference id must be present

Prints the character error rate (CER, %) of A, and with B the paired difference A - B with a 95% bootstrap interval
that resamples whole clusters. Normalisation, applied identically to both sides:
  chinda_norm.normalize(text, "th")  Thai digits to Arabic, curly quotes to ASCII, SARA AM decomposed, spaces between
                                     two Thai characters removed (chinda_norm.py)
  num_norm.normalize_numbers         spoken Thai numbers to digits, on whole pythainlp newmm tokens (num_norm.py)
  then every space is removed, and CER = Levenshtein edits / reference characters over the whole set.
Requires: pythainlp (5.3.5 used for the card) and editdistance.
"""
import argparse, json, os, random, sys
from collections import defaultdict

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import editdistance  # noqa: E402
from chinda_norm import normalize  # noqa: E402
from num_norm import normalize_numbers  # noqa: E402


def prep(t):
    return normalize_numbers(normalize(t or "", "th")).replace(" ", "")


def load(path, key):
    return {str(r["id"]): r[key] for r in map(json.loads, open(path, encoding="utf-8")) if r}


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("refs"); ap.add_argument("a"); ap.add_argument("b", nargs="?")
    ap.add_argument("--boot", type=int, default=1000)
    a = ap.parse_args()
    refs = {str(r["id"]): r for r in map(json.loads, open(a.refs, encoding="utf-8")) if r}
    # Take the clip order from hypothesis file A (ours are in the order we decoded them): the bootstrap draws
    # clusters in order of first appearance, so this reproduces the intervals on the model card exactly.
    order = [str(r["id"]) for r in map(json.loads, open(a.a, encoding="utf-8")) if r]
    seen = set(order)
    ids = [i for i in order if i in refs] + [i for i in refs if i not in seen]
    refs = [refs[i] for i in ids]
    R = [prep(r["text"]) for r in refs]
    cl = [str(r.get("cluster", r["id"])) for r in refs]

    def errs(path):
        h = load(path, "pred_text")
        missing = [i for i in ids if i not in h]
        if missing:
            sys.exit(f"{path}: {len(missing)} reference ids missing (first: {missing[0]})")
        return [(editdistance.eval(r, prep(h[i])), len(r)) for i, r in zip(ids, R)]

    A = errs(a.a)
    cer = lambda E: 100 * sum(e for e, _ in E) / sum(n for _, n in E)
    print(f"A: CER {cer(A):.2f}%  ({len(ids)} clips, {len(set(cl))} clusters)")
    if not a.b:
        return
    B = errs(a.b)
    print(f"B: CER {cer(B):.2f}%")
    groups = defaultdict(list)
    for i, c in enumerate(cl):
        groups[c].append(i)
    G = list(groups.values()); rnd = random.Random(1); ds = []
    for _ in range(a.boot):
        idx = [i for _ in G for i in G[rnd.randrange(len(G))]]
        ds.append(cer([A[i] for i in idx]) - cer([B[i] for i in idx]))
    ds.sort(); lo, hi = ds[int(.025 * a.boot)], ds[int(.975 * a.boot)]
    verdict = "A lower" if hi < 0 else "B lower" if lo > 0 else "tie"
    print(f"A - B: {cer(A) - cer(B):+.2f}  95% interval [{lo:+.2f}, {hi:+.2f}]  -> {verdict}")


if __name__ == "__main__":
    main()