#!/usr/bin/env python3
"""Transcribes FLEURS Vietnamese test clips with the exported ONNX files and reports WER.

Runs the same loop a mobile app does: encoder once, then the merged decoder step by step with its
KV cache (use_cache_branch false on the first step, true after), greedy, prompt
<|startoftranscript|><|vi|><|transcribe|><|notimestamps|>. Writes out/<size>/wer.json (read by publish.py)
and fails when WER exceeds --max-wer.

Usage:
  .venv/bin/python -I validate.py tiny [--count 30] [--max-wer 0.4]
"""

import argparse
import csv
import json
import re
import sys
import time
import unicodedata
from pathlib import Path

import jiwer
import numpy as np
import onnxruntime as ort
import soundfile as sf
from transformers import WhisperFeatureExtractor, WhisperTokenizer

HERE = Path(__file__).resolve().parent
FLEURS = HERE / "data"
MAX_TOKENS = 224


def clips(count: int) -> list[tuple[Path, str]]:
    rows = []
    with open(FLEURS / "fleurs/data/vi_vn/test.tsv", encoding="utf-8") as f:
        for row in csv.reader(f, delimiter="\t", quoting=csv.QUOTE_NONE):
            wav = next((FLEURS / "vi").glob(f"*/{row[1]}"), None)
            if wav is not None:
                rows.append((wav, row[2]))
            if len(rows) == count:
                break
    return rows


def normalize(text: str) -> str:
    text = unicodedata.normalize("NFC", text.lower())
    return re.sub(r"\s+", " ", re.sub(r"[^\w\s]", " ", text)).strip()


class Transcriber:
    def __init__(self, repo: Path, quantized: bool):
        suffix = "_quantized" if quantized else ""
        opts = ort.SessionOptions()
        opts.intra_op_num_threads = 8
        self.encoder = ort.InferenceSession(str(repo / f"onnx/encoder_model{suffix}.onnx"), opts)
        self.decoder = ort.InferenceSession(str(repo / f"onnx/decoder_model_merged{suffix}.onnx"), opts)
        self.features = WhisperFeatureExtractor.from_pretrained(repo)
        self.tokenizer = WhisperTokenizer.from_pretrained(repo)
        names = ["<|startoftranscript|>", "<|vi|>", "<|transcribe|>", "<|notimestamps|>"]
        self.prompt = self.tokenizer.convert_tokens_to_ids(names)
        self.eot = self.tokenizer.eos_token_id
        self.past = [i for i in self.decoder.get_inputs() if i.name.startswith("past_key_values.")]

    def __call__(self, audio: np.ndarray) -> str:
        feats = self.features(audio, sampling_rate=16000, return_tensors="np").input_features
        hidden = self.encoder.run(None, {"input_features": feats.astype(np.float32)})[0]
        # Shape of an empty cache: [1, heads, 0, head_dim]; dims come from the graph's symbolic shapes.
        cache = {i.name: np.zeros([1, *[d if isinstance(d, int) else 0 for d in i.shape[1:]]], np.float32)
                 for i in self.past}
        for i in self.past:
            if not isinstance(i.shape[1], int) or not isinstance(i.shape[3], int):
                sys.exit(f"error: cannot size the empty cache for {i.name} {i.shape}")
        tokens, ids = list(self.prompt), np.array([self.prompt], np.int64)
        outputs = [o.name for o in self.decoder.get_outputs()]
        for step in range(MAX_TOKENS):
            feed = {"input_ids": ids, "encoder_hidden_states": hidden,
                    "use_cache_branch": np.array([step > 0]), **cache}
            result = dict(zip(outputs, self.decoder.run(None, feed)))
            nxt = int(result["logits"][0, -1].argmax())
            if nxt == self.eot:
                break
            tokens.append(nxt)
            for name in cache:
                present = result[name.replace("past_key_values.", "present.")]
                # The encoder (cross-attention) cache only comes out of the first step.
                if step == 0 or ".decoder." in name:
                    cache[name] = present
            ids = np.array([[nxt]], np.int64)
        return self.tokenizer.decode(tokens, skip_special_tokens=True).strip()


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("size")
    parser.add_argument("--count", type=int, default=30)
    parser.add_argument("--max-wer", type=float, default=0.4)
    args = parser.parse_args()
    repo = HERE / "out" / args.size / "repo"
    samples = clips(args.count)
    failed, results = False, {"clips": len(samples)}
    for quantized in (False, True):
        run = Transcriber(repo, quantized)
        hyps, refs, start = [], [], time.time()
        for wav, ref in samples:
            audio, rate = sf.read(wav, dtype="float32")
            assert rate == 16000, wav
            hyps.append(normalize(run(audio)))
            refs.append(normalize(ref))
        wer = jiwer.wer(refs, hyps)
        label = "int8" if quantized else "fp32"
        print(f"{args.size} {label}: WER {wer:.3f} on {len(samples)} clips in {time.time() - start:.0f} s")
        print(f"  ref: {refs[0]}\n  hyp: {hyps[0]}")
        results[label] = round(wer, 4)
        failed |= wer > args.max_wer
    (repo.parent / "wer.json").write_text(json.dumps(results) + "\n")
    if failed:
        sys.exit(f"error: WER above {args.max_wer}")


if __name__ == "__main__":
    main()