# -*- coding: utf-8 -*-
"""Transcribes the final WAVs and compares them ignoring how numbers are written."""
import json, os, re, subprocess, sys, unicodedata
from difflib import SequenceMatcher

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.insert(0, os.path.join(ROOT, "scripts"))
from narration import SCRIPTS
from numspell import expand

NUM = set("""cero uno una dos tres cuatro cinco seis siete ocho nueve diez once doce trece
catorce quince dieciseis diecisiete dieciocho diecinueve veinte veintiuno veintidos veintitres
treinta cuarenta cincuenta sesenta setenta ochenta noventa cien ciento doscientos trescientos
cuatrocientos quinientos mil millon millones coma periodico por ciento y tercio tercios
primero segundo""".split())


def norm(s, drop_numbers=False):
    s = expand(s)
    s = unicodedata.normalize("NFD", s.lower())
    s = "".join(c for c in s if unicodedata.category(c) != "Mn")
    s = re.sub(r"[^a-z0-9ñ ]+", " ", s)
    toks = s.split()
    if drop_numbers:
        toks = [t for t in toks if not t.isdigit() and t not in NUM]
    return " ".join(toks)


def main(names):
    for name in names:
        d = json.load(open(os.path.join(ROOT, "audio", name, "segments.json")))
        wavs = [s["wav"] for s in d["segments"]]
        work = os.path.join(ROOT, "tmp", name + "_v")
        os.makedirs(work, exist_ok=True)
        subprocess.run(["conda", "run", "-n", "stt", "--no-capture-output",
                        "whisper-ctranslate2", "--language", "es", "--model", "small",
                        "--output_dir", work, "--output_format", "txt", *wavs],
                       stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, check=True)
        print(f"\n=== {name} ===")
        bad = []
        for s in d["segments"]:
            t = os.path.join(work, os.path.splitext(os.path.basename(s["wav"]))[0] + ".txt")
            got = open(t, encoding="utf-8").read().strip() if os.path.exists(t) else ""
            r = SequenceMatcher(None, norm(s["tts"]), norm(got)).ratio()
            flag = "OK " if r >= 0.88 else "REV"
            if r < 0.88:
                bad.append(s["i"])
            print(f" {s['i']:02d} {flag} {r:.2f} | {got}")
        print(f" -> revisar: {bad}")


if __name__ == "__main__":
    main(sys.argv[1:] or list(SCRIPTS))