"""Word Error Rate evaluation for speech recognition models.

Evaluates transcription quality on LibriSpeech test-clean using
Qwen3-ASR (via mlx-audio) for inference and edit distance for WER computation.
"""

from dataclasses import dataclass

import numpy as np


@dataclass
class WERResult:
    n_samples: int
    wer: float  # Word Error Rate
    n_correct_words: int
    n_total_words: int
    n_substitutions: int
    n_insertions: int
    n_deletions: int


def _edit_distance(ref: list[str], hyp: list[str]) -> tuple[int, int, int]:
    """Compute edit distance between reference and hypothesis word lists.

    Returns (substitutions, insertions, deletions).
    """
    n = len(ref)
    m = len(hyp)

    # dp[i][j] = (cost, subs, ins, dels) for ref[:i] vs hyp[:j]
    dp = [[(0, 0, 0, 0)] * (m + 1) for _ in range(n + 1)]
    for i in range(1, n + 1):
        dp[i][0] = (i, 0, 0, i)
    for j in range(1, m + 1):
        dp[0][j] = (j, 0, j, 0)

    for i in range(1, n + 1):
        for j in range(1, m + 1):
            if ref[i - 1] == hyp[j - 1]:
                dp[i][j] = dp[i - 1][j - 1]
            else:
                sub = dp[i - 1][j - 1]
                ins = dp[i][j - 1]
                dele = dp[i - 1][j]
                candidates = [
                    (sub[0] + 1, sub[1] + 1, sub[2], sub[3]),
                    (ins[0] + 1, ins[1], ins[2] + 1, ins[3]),
                    (dele[0] + 1, dele[1], dele[2], dele[3] + 1),
                ]
                dp[i][j] = min(candidates, key=lambda x: x[0])

    _, subs, ins, dels = dp[n][m]
    return subs, ins, dels


def evaluate_wer_mlx(
    model,
    n_samples: int = 100,
    seed: int = 42,
) -> WERResult:
    """Evaluate WER on LibriSpeech test-clean using a Qwen3-ASR model.

    Args:
        model: MLX Qwen3-ASR model with .generate() method
        n_samples: Number of audio samples to evaluate
        seed: Random seed for sample selection

    Returns:
        WERResult with WER metrics
    """
    from datasets import load_dataset
    from tqdm import tqdm

    rng = np.random.RandomState(seed)

    print(f"  Loading LibriSpeech test-clean...")
    ds = load_dataset("librispeech_asr", "clean", split="test")

    indices = rng.choice(len(ds), size=min(n_samples, len(ds)), replace=False)
    indices.sort()

    total_subs = 0
    total_ins = 0
    total_dels = 0
    total_words = 0

    for idx in tqdm(indices, desc="  WER eval"):
        item = ds[int(idx)]
        audio = item["audio"]["array"].astype(np.float32)
        ref_text = item["text"].lower().strip()
        ref_words = ref_text.split()

        if not ref_words:
            continue

        total_words += len(ref_words)

        try:
            result = model.generate(
                audio,
                language="English",
                temperature=0.0,
            )
            hyp_text = result.text.lower().strip()
            hyp_words = hyp_text.split()

            subs, ins, dels = _edit_distance(ref_words, hyp_words)
            total_subs += subs
            total_ins += ins
            total_dels += dels
        except Exception as e:
            # Count all words as deletions on failure
            total_dels += len(ref_words)

    wer = (total_subs + total_ins + total_dels) / max(total_words, 1)
    n_correct = total_words - total_subs - total_dels

    print(f"  WER: {wer:.4f} ({total_subs}S + {total_ins}I + {total_dels}D / {total_words} words)")

    return WERResult(
        n_samples=len(indices),
        wer=wer,
        n_correct_words=max(n_correct, 0),
        n_total_words=total_words,
        n_substitutions=total_subs,
        n_insertions=total_ins,
        n_deletions=total_dels,
    )


def print_wer_report(result: WERResult):
    """Print WER evaluation results."""
    print(f"\n  WER Report")
    print(f"  {'=' * 50}")
    print(f"  Samples: {result.n_samples}")
    print(f"  WER: {result.wer:.4f} ({result.wer * 100:.2f}%)")
    print(f"  Total words: {result.n_total_words}")
    print(f"  Errors: {result.n_substitutions}S + {result.n_insertions}I + {result.n_deletions}D")
    print(f"  {'=' * 50}")