"""准确率指标：CER（字级）/ WER（词级）/ MER（中英混合）。

分词策略：
  - CER : 去空格后逐字符（中文每字、英文每字母各一 token）。
  - WER : 按空白切词（主要对英文有意义）。
  - MER : 中文按【字】、英文/数字按【词】，混合 token 序列算编辑距离——这是中英混说的 headline 指标。

聚合：错误率必须在语料级聚合（累加 S/D/I/N 后再求率），不能对每条样本的率求平均，
否则短句权重被放大、结果失真。
"""

from __future__ import annotations

import re
from dataclasses import dataclass

import jiwer

from .text_norm import normalize

_CJK = re.compile(r"[一-鿿㐀-䶿]")


@dataclass
class Counts:
    """一组编辑距离统计。N = 参考长度 = S + D + H。"""

    S: int = 0  # 替换
    D: int = 0  # 删除
    I: int = 0  # 插入
    N: int = 0  # 参考 token 数

    def __add__(self, o: "Counts") -> "Counts":
        return Counts(self.S + o.S, self.D + o.D, self.I + o.I, self.N + o.N)

    @property
    def rate(self) -> float:
        if self.N == 0:
            return 1.0 if (self.S + self.D + self.I) > 0 else 0.0
        return (self.S + self.D + self.I) / self.N


def char_tokens(text: str) -> list[str]:
    return [c for c in text if not c.isspace()]


def word_tokens(text: str) -> list[str]:
    return text.split()


def mixed_tokens(text: str) -> list[str]:
    """中文按字、其余按词。"""
    tokens: list[str] = []
    buf = ""
    for ch in text:
        if _CJK.match(ch):
            if buf.strip():
                tokens.extend(buf.split())
            buf = ""
            tokens.append(ch)
        else:
            buf += ch
    if buf.strip():
        tokens.extend(buf.split())
    return tokens


def _counts(ref_tokens: list[str], hyp_tokens: list[str]) -> Counts:
    if not ref_tokens:
        return Counts(S=0, D=0, I=len(hyp_tokens), N=0)
    out = jiwer.process_words(" ".join(ref_tokens), " ".join(hyp_tokens))
    return Counts(
        S=out.substitutions,
        D=out.deletions,
        I=out.insertions,
        N=out.substitutions + out.deletions + out.hits,
    )


def score_sample(ref_text: str, hyp_text: str) -> dict[str, Counts]:
    """对一条样本算 CER/WER/MER 的原始计数（供语料级聚合）。"""
    ref = normalize(ref_text)
    hyp = normalize(hyp_text)
    return {
        "cer": _counts(char_tokens(ref), char_tokens(hyp)),
        "wer": _counts(word_tokens(ref), word_tokens(hyp)),
        "mer": _counts(mixed_tokens(ref), mixed_tokens(hyp)),
    }


def empty_counts() -> dict[str, Counts]:
    return {"cer": Counts(), "wer": Counts(), "mer": Counts()}


def add_counts(acc: dict[str, Counts], one: dict[str, Counts]) -> dict[str, Counts]:
    for k in acc:
        acc[k] = acc[k] + one[k]
    return acc