"""Evaluate a fine-tuned Whisper checkpoint on the test split (WER / CER).

    uv run python evaluate_model.py --model ./whisper-small-ar-dialectal
    uv run python evaluate_model.py --model openai/whisper-small   # baseline
"""
from __future__ import annotations

import argparse

import torch
import jiwer
from datasets import load_dataset
from transformers import WhisperProcessor, WhisperForConditionalGeneration

from normalize import clean_text
from train import keep_row

DATASET = "oddadmix/dialectal-arabic-lahgtna-v2-smaller-augmented"


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--model", required=True)
    ap.add_argument("--dataset", default=DATASET)
    ap.add_argument("--split", default="test")
    ap.add_argument("--batch_size", type=int, default=16)
    ap.add_argument("--limit", type=int, default=0, help="0 = full split")
    ap.add_argument("--normalize_letters", action="store_true")
    args = ap.parse_args()

    device = "cuda" if torch.cuda.is_available() else "cpu"
    proc = WhisperProcessor.from_pretrained(args.model, language="ar", task="transcribe")
    model = WhisperForConditionalGeneration.from_pretrained(
        args.model, torch_dtype=torch.bfloat16
    ).to(device).eval()

    nl = args.normalize_letters
    ds = load_dataset(args.dataset, split=args.split)
    ds = ds.filter(lambda text, duration: keep_row(text, duration, nl),
                   input_columns=["text", "duration"])
    if args.limit:
        ds = ds.select(range(min(args.limit, len(ds))))
    print(f"evaluating on {len(ds)} clips")

    fe = proc.feature_extractor
    preds: list[str] = []
    refs: list[str] = []
    for i in range(0, len(ds), args.batch_size):
        rows = ds[i : i + args.batch_size]
        arrays = [a["array"] for a in rows["audio"]]
        feats = fe(arrays, sampling_rate=16000, return_tensors="pt").input_features
        feats = feats.to(device, dtype=torch.bfloat16)
        with torch.no_grad():
            gen = model.generate(input_features=feats, max_new_tokens=225)
        preds += proc.batch_decode(gen, skip_special_tokens=True)
        refs += rows["text"]
        print(f"  {min(i + args.batch_size, len(ds))}/{len(ds)}", end="\r")

    preds = [clean_text(p, nl) for p in preds]
    refs = [clean_text(r, nl) for r in refs]
    pairs = [(p, r) for p, r in zip(preds, refs) if r.strip()]
    preds, refs = map(list, zip(*pairs))

    print("\n" + "=" * 50)
    print(f"model : {args.model}")
    print(f"WER   : {jiwer.wer(refs, preds):.4f}")
    print(f"CER   : {jiwer.cer(refs, preds):.4f}")
    print("--- examples ---")
    for r, p in list(zip(refs, preds))[:5]:
        print(f"  REF: {r[:90]}")
        print(f"  HYP: {p[:90]}\n")


if __name__ == "__main__":
    main()