"""Character and Word Error Rate evaluation for OCR output.

Provides CER, WER, digit-specific CER, and normalized edit distance.

Usage:
    from benchmarks.evaluate_cer import compute_cer, compute_digit_cer, evaluate_file
"""

from __future__ import annotations

import re
import unicodedata
from pathlib import Path

import numpy as np


IMAGE_EXTENSIONS = {".jpg", ".jpeg", ".png", ".tiff", ".tif", ".webp", ".bmp"}


def extract_cjk_characters(text: str) -> list[str]:
    """Return Chinese, Japanese, and Korean script characters from text."""
    ranges = (
        (0x3400, 0x4DBF),    # CJK Extension A
        (0x4E00, 0x9FFF),    # CJK Unified Ideographs
        (0xF900, 0xFAFF),    # CJK Compatibility Ideographs
        (0x20000, 0x2FA1F),  # CJK supplementary extensions
        (0x30000, 0x323AF),  # CJK Extensions G and H
        (0x3040, 0x309F),    # Hiragana
        (0x30A0, 0x30FF),    # Katakana
        (0x31F0, 0x31FF),    # Katakana phonetic extensions
        (0xFF66, 0xFF9D),    # Half-width Katakana
        (0x1B000, 0x1B16F),  # Kana supplements and small kana
        (0x1100, 0x11FF),    # Hangul Jamo
        (0x3130, 0x318F),    # Hangul compatibility Jamo
        (0xA960, 0xA97F),    # Hangul Jamo Extended-A
        (0xAC00, 0xD7AF),    # Hangul syllables
        (0xD7B0, 0xD7FF),    # Hangul Jamo Extended-B
    )
    return [
        char
        for char in text
        if any(start <= ord(char) <= end for start, end in ranges)
    ]


def normalize_text(text: str) -> str:
    """Normalize text for fair comparison.

    - Unicode NFC normalization
    - Collapse whitespace
    - Strip leading/trailing whitespace
    """
    text = unicodedata.normalize("NFC", text)
    text = re.sub(r"\s+", " ", text)
    return text.strip()


def edit_distance(ref: list, hyp: list) -> int:
    """Compute Levenshtein edit distance between two sequences."""
    n, m = len(ref), len(hyp)
    if n == 0:
        return m
    if m == 0:
        return n

    # Use two rows for memory efficiency
    prev = list(range(m + 1))
    curr = [0] * (m + 1)

    for i in range(1, n + 1):
        curr[0] = i
        for j in range(1, m + 1):
            if ref[i - 1] == hyp[j - 1]:
                curr[j] = prev[j - 1]
            else:
                curr[j] = 1 + min(prev[j], curr[j - 1], prev[j - 1])
        prev, curr = curr, prev

    return prev[m]


def compute_cer(reference: str, hypothesis: str, normalize: bool = True) -> float:
    """Compute Character Error Rate.

    CER = edit_distance(ref_chars, hyp_chars) / len(ref_chars)
    """
    if normalize:
        reference = normalize_text(reference)
        hypothesis = normalize_text(hypothesis)

    # Remove spaces for character-level comparison
    ref_chars = list(reference.replace(" ", ""))
    hyp_chars = list(hypothesis.replace(" ", ""))

    if len(ref_chars) == 0:
        return 0.0 if len(hyp_chars) == 0 else 1.0

    dist = edit_distance(ref_chars, hyp_chars)
    return dist / len(ref_chars)


def compute_wer(reference: str, hypothesis: str, normalize: bool = True) -> float:
    """Compute Word Error Rate.

    WER = edit_distance(ref_words, hyp_words) / len(ref_words)
    """
    if normalize:
        reference = normalize_text(reference)
        hypothesis = normalize_text(hypothesis)

    ref_words = reference.split()
    hyp_words = hypothesis.split()

    if len(ref_words) == 0:
        return 0.0 if len(hyp_words) == 0 else 1.0

    dist = edit_distance(ref_words, hyp_words)
    return dist / len(ref_words)


def compute_digit_cer(reference: str, hypothesis: str) -> float:
    """Compute CER only on digit characters (0-9).

    Critical for financial documents, invoices, and forms.
    """
    ref_digits = re.sub(r"[^0-9]", "", reference)
    hyp_digits = re.sub(r"[^0-9]", "", hypothesis)

    if len(ref_digits) == 0:
        return 0.0 if len(hyp_digits) == 0 else 1.0

    dist = edit_distance(list(ref_digits), list(hyp_digits))
    return dist / len(ref_digits)


def compute_cjk_cer(reference: str, hypothesis: str) -> float:
    """Compute CER specifically for CJK characters."""
    ref_cjk = extract_cjk_characters(reference)
    hyp_cjk = extract_cjk_characters(hypothesis)

    if len(ref_cjk) == 0:
        return 0.0 if len(hyp_cjk) == 0 else 1.0

    dist = edit_distance(ref_cjk, hyp_cjk)
    return dist / len(ref_cjk)


def detect_repetition(text: str, ngram_size: int = 35, threshold: int = 3) -> float:
    """Detect repetition rate in generated text.

    Returns the fraction of n-grams that appear more than `threshold` times.
    """
    if not isinstance(ngram_size, int) or isinstance(ngram_size, bool) or ngram_size < 1:
        raise ValueError("ngram_size must be a positive integer")
    if not isinstance(threshold, int) or isinstance(threshold, bool) or threshold < 0:
        raise ValueError("threshold must be a non-negative integer")

    words = text.split()
    if len(words) < ngram_size:
        return 0.0

    ngrams = []
    for i in range(len(words) - ngram_size + 1):
        ngrams.append(tuple(words[i:i + ngram_size]))

    if not ngrams:
        return 0.0

    from collections import Counter
    counts = Counter(ngrams)
    repeated = sum(1 for c in counts.values() if c > threshold)
    return repeated / len(counts)


def evaluate_text(reference: str, hypothesis: str) -> dict:
    """Evaluate a reference/hypothesis text pair."""
    return {
        "cer": compute_cer(reference, hypothesis),
        "wer": compute_wer(reference, hypothesis),
        "digit_cer": compute_digit_cer(reference, hypothesis),
        "cjk_cer": compute_cjk_cer(reference, hypothesis),
        "repetition_rate": detect_repetition(hypothesis),
        "ref_length": len(reference),
        "hyp_length": len(hypothesis),
        "ref_digit_count": sum(char.isascii() and char.isdigit() for char in reference),
        "ref_cjk_count": len(extract_cjk_characters(reference)),
    }


def evaluate_file(reference_path: str | Path, hypothesis_path: str | Path) -> dict:
    """Evaluate a single reference/hypothesis pair."""
    reference = Path(reference_path).read_text(encoding="utf-8")
    hypothesis = Path(hypothesis_path).read_text(encoding="utf-8")
    return evaluate_text(reference, hypothesis)


def evaluate_directory(
    images_dir: str | Path,
    ground_truth_dir: str | Path,
    hypothesis_dir: str | Path,
) -> dict:
    """Evaluate all files in a directory.

    Returns aggregate metrics.
    """
    image_dir = Path(images_dir)
    gt_dir = Path(ground_truth_dir)
    hyp_dir = Path(hypothesis_dir)

    for label, directory in (
        ("images", image_dir),
        ("ground truth", gt_dir),
        ("hypothesis", hyp_dir),
    ):
        if not directory.is_dir():
            raise NotADirectoryError(f"{label.title()} directory not found: {directory}")

    image_files = sorted(
        path
        for path in image_dir.iterdir()
        if path.is_file() and path.suffix.lower() in IMAGE_EXTENSIONS
    )
    image_stems = [path.stem for path in image_files]
    duplicate_stems = sorted({stem for stem in image_stems if image_stems.count(stem) > 1})
    if duplicate_stems:
        raise ValueError(
            "Multiple input images share the same stem: " + ", ".join(duplicate_stems)
        )

    results = []
    missing_ground_truth = []
    missing_hypotheses = []
    for image_file in image_files:
        gt_file = gt_dir / f"{image_file.stem}.txt"
        if not gt_file.is_file():
            missing_ground_truth.append(image_file.name)
            continue
        hyp_file = hyp_dir / gt_file.name
        reference = gt_file.read_text(encoding="utf-8")
        if hyp_file.is_file():
            hypothesis = hyp_file.read_text(encoding="utf-8")
        else:
            # A missing prediction is an empty prediction, not a sample that
            # can be silently removed from the accuracy denominator.
            hypothesis = ""
            missing_hypotheses.append(image_file.name)
        result = evaluate_text(reference, hypothesis)
        result["file"] = gt_file.stem
        result["image"] = image_file.name
        result["missing_hypothesis"] = not hyp_file.is_file()
        results.append(result)

    if not results:
        return {
            "error": "No images with matching ground-truth files found",
            "num_images": len(image_files),
            "num_samples": 0,
            "missing_ground_truth": missing_ground_truth,
            "missing_hypotheses": missing_hypotheses,
        }

    digit_results = [r["digit_cer"] for r in results if r["ref_digit_count"] > 0]
    cjk_results = [r["cjk_cer"] for r in results if r["ref_cjk_count"] > 0]

    return {
        "num_images": len(image_files),
        "num_samples": len(results),
        "num_digit_samples": len(digit_results),
        "num_cjk_samples": len(cjk_results),
        "missing_ground_truth": missing_ground_truth,
        "missing_hypotheses": missing_hypotheses,
        "mean_cer": float(np.mean([r["cer"] for r in results])),
        "mean_wer": float(np.mean([r["wer"] for r in results])),
        "mean_digit_cer": float(np.mean(digit_results)) if digit_results else None,
        "mean_cjk_cer": float(np.mean(cjk_results)) if cjk_results else None,
        "mean_repetition_rate": float(np.mean([r["repetition_rate"] for r in results])),
        "per_file": results,
    }