"""WER benchmark: open-source Whisper models on FLEURS test sets.

Streams FLEURS samples (no full dataset download), transcribes with
mlx-whisper, scores WER/CER with Whisper-paper-style normalization.
Appends per-utterance results to results.jsonl and prints progress.

Usage:
  python bench.py                          # full run (all models x langs)
  python bench.py --models small --langs en_us --n 3   # pilot
"""

import argparse
import io
import json
import time
import unicodedata
from pathlib import Path

import jiwer
import numpy as np
import soundfile as sf
from datasets import Audio, load_dataset
from transformers.models.whisper.english_normalizer import (
    BasicTextNormalizer,
    EnglishTextNormalizer,
)

import mlx_whisper

# FLEURS config -> (whisper language code, primary metric)
LANGS = {
    'en_us': ('en', 'wer'),
    'es_419': ('es', 'wer'),
    'fr_fr': ('fr', 'wer'),
    'de_de': ('de', 'wer'),
    'it_it': ('it', 'wer'),
    'pt_br': ('pt', 'wer'),
    'ja_jp': ('ja', 'cer'),
    'ko_kr': ('ko', 'wer'),
    'cmn_hans_cn': ('zh', 'cer'),
    'yue_hant_hk': ('yue', 'cer'),
    'hi_in': ('hi', 'wer'),
    'ar_eg': ('ar', 'wer'),
}

MODELS = {
    'large-v3': 'mlx-community/whisper-large-v3-mlx',
    'large-v3-turbo': 'mlx-community/whisper-large-v3-turbo',
    'small': 'mlx-community/whisper-small-mlx',
}

SEED = 42
SHUFFLE_BUFFER = 500

english_norm = EnglishTextNormalizer({})
basic_norm = BasicTextNormalizer()


def normalize(text: str, lang: str) -> str:
    text = unicodedata.normalize('NFKC', text)
    if lang == 'en':
        return english_norm(text)
    return basic_norm(text)


def score(ref: str, hyp: str, lang: str, metric: str) -> dict:
    ref_n = normalize(ref, lang)
    hyp_n = normalize(hyp, lang)
    if not ref_n:
        return {'wer': None, 'cer': None}
    out = {'wer': jiwer.wer(ref_n, hyp_n) if ref_n.split() else None}
    if metric == 'cer':
        ref_c = ref_n.replace(' ', '')
        hyp_c = hyp_n.replace(' ', '')
        out['cer'] = jiwer.cer(ref_c, hyp_c) if ref_c else None
    else:
        out['cer'] = None
    return out


def load_samples(fleurs_config: str, n: int) -> list[dict]:
    ds = load_dataset('google/fleurs', fleurs_config, split='test', streaming=True)
    ds = ds.cast_column('audio', Audio(decode=False))
    ds = ds.shuffle(seed=SEED, buffer_size=SHUFFLE_BUFFER)
    samples = []
    for ex in ds:
        raw = ex['audio']['bytes']
        if raw is None:
            with open(ex['audio']['path'], 'rb') as f:
                raw = f.read()
        audio, sr = sf.read(io.BytesIO(raw), dtype='float32')
        if audio.ndim > 1:
            audio = audio.mean(axis=1)
        samples.append(
            {
                'id': ex['id'],
                'audio': audio,
                'sr': sr,
                'ref': ex['transcription'],
                'duration': len(audio) / sr,
            }
        )
        if len(samples) >= n:
            break
    return samples


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--models', nargs='*', default=list(MODELS))
    ap.add_argument('--langs', nargs='*', default=list(LANGS))
    ap.add_argument('--n', type=int, default=50)
    ap.add_argument('--out', default='results.jsonl')
    args = ap.parse_args()

    out_path = Path(args.out)
    done = set()
    if out_path.exists():
        for line in out_path.read_text().splitlines():
            try:
                r = json.loads(line)
                done.add((r['model'], r['lang_config'], r['id']))
            except json.JSONDecodeError:
                pass

    for fleurs_config in args.langs:
        whisper_lang, metric = LANGS[fleurs_config]
        print(f'[{fleurs_config}] loading {args.n} samples...', flush=True)
        # FLEURS is streamed from HF; the connection occasionally drops mid-download.
        # Retry with backoff so a transient network error doesn't abort the run.
        for attempt in range(6):
            try:
                samples = load_samples(fleurs_config, args.n)
                break
            except Exception as e:
                if attempt == 5:
                    raise
                wait = 10 * (attempt + 1)
                print(f'[{fleurs_config}] load failed ({type(e).__name__}); retry {attempt + 1}/5 in {wait}s', flush=True)
                time.sleep(wait)
        total_audio = sum(s['duration'] for s in samples)
        print(f'[{fleurs_config}] {len(samples)} samples, {total_audio / 60:.1f} min audio', flush=True)

        for model_name in args.models:
            repo = MODELS[model_name]
            pending = [s for s in samples if (model_name, fleurs_config, s['id']) not in done]
            if not pending:
                print(f'[{fleurs_config}] {model_name}: already done, skip', flush=True)
                continue
            t0 = time.time()
            with out_path.open('a') as f:
                for i, s in enumerate(pending):
                    try:
                        result = mlx_whisper.transcribe(
                            s['audio'],
                            path_or_hf_repo=repo,
                            language=whisper_lang,
                            task='transcribe',
                            fp16=True,
                            temperature=0.0,
                            condition_on_previous_text=False,
                        )
                    except ValueError as e:
                        # Some languages (e.g. Cantonese 'yue') exist only in the
                        # large-v3 tokenizer; smaller models raise here. Skip the
                        # whole (language, model) cell and continue the run.
                        print(f'[{fleurs_config}] {model_name}: unsupported ({e}); skipping cell', flush=True)
                        break
                    hyp = result['text']
                    metrics = score(s['ref'], hyp, whisper_lang, metric)
                    f.write(
                        json.dumps(
                            {
                                'model': model_name,
                                'lang_config': fleurs_config,
                                'lang': whisper_lang,
                                'id': s['id'],
                                'duration': round(s['duration'], 2),
                                'ref': s['ref'],
                                'hyp': hyp,
                                **metrics,
                            },
                            ensure_ascii=False,
                        )
                        + '\n'
                    )
                    f.flush()
                    if (i + 1) % 10 == 0:
                        print(f'[{fleurs_config}] {model_name}: {i + 1}/{len(pending)}', flush=True)
            dt = time.time() - t0
            rtf = sum(s['duration'] for s in pending) / dt if dt > 0 else 0
            print(f'[{fleurs_config}] {model_name}: done in {dt / 60:.1f} min (RTF {rtf:.1f}x)', flush=True)

    print('ALL DONE', flush=True)


if __name__ == '__main__':
    main()