"""Measure word error rate: render sentences, transcribe with Whisper small.

    python benchmark/wer.py --model . --sentences benchmark/eval_sentences.txt

Requires `transformers` and (for the first run) a 460 MB model download.
WER = word-level Levenshtein distance / reference length, lowercase and
punctuation-stripped.
"""
import argparse
import re
import sys
import tempfile
from pathlib import Path

import numpy as np

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
from microtts import MicroTTS  # noqa: E402


def norm(s: str) -> str:
    s = s.lower()
    s = re.sub(r"[^a-z0-9' ]+", " ", s)
    return " ".join(s.split())


def wer(ref: str, hyp: str) -> float:
    r, h = norm(ref).split(), norm(hyp).split()
    if not r:
        return 0.0
    d = np.zeros((len(r) + 1, len(h) + 1), dtype=np.int32)
    d[:, 0] = np.arange(len(r) + 1)
    d[0, :] = np.arange(len(h) + 1)
    for i in range(1, len(r) + 1):
        for j in range(1, len(h) + 1):
            d[i, j] = min(d[i - 1, j] + 1, d[i, j - 1] + 1,
                          d[i - 1, j - 1] + (r[i - 1] != h[j - 1]))
    return float(d[len(r), len(h)] / len(r))


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--model", default=".")
    ap.add_argument("--sentences", default="benchmark/eval_sentences.txt")
    ap.add_argument("--device", default="cpu")
    a = ap.parse_args()

    import soundfile as sf
    from transformers import pipeline

    sents = [l.strip() for l in open(a.sentences, encoding="utf-8") if l.strip()]
    tts = MicroTTS.load(a.model, device=a.device)
    asr = pipeline("automatic-speech-recognition", model="openai/whisper-small",
                   device=0 if a.device == "cuda" else -1)

    tmp = Path(tempfile.mkdtemp())
    total, perfect = 0.0, 0
    for i, s in enumerate(sents):
        fn = tmp / f"s{i:04d}.wav"
        sf.write(str(fn), MicroTTS.normalize(tts.synthesize(s)), 24000)
        hyp = asr(str(fn))["text"].strip()
        e = wer(s, hyp)
        total += e
        perfect += 1 if e == 0 else 0
        print(f"[{i:3d}] WER {e:.3f} | {s[:64]}")
    print(f"\nmean WER {total/len(sents):.4f} | exact {perfect}/{len(sents)}")


if __name__ == "__main__":
    main()