"""Единый нормализатор текста для бенчмарка.

Применяется СИММЕТРИЧНО к эталону и ко всем гипотезам — своей и чужим. Асимметрия
здесь дороже любой другой ошибки: она сдвигает WER на единицы процентов и делает
цифры невоспроизводимыми, а бенчмарк без воспроизводимости бесполезен.

Правила публикуются вместе с результатами замера, на странице
https://rechetext.ru/benchmark:

1. нижний регистр;
2. ё → е (в эталонах Common Voice ё стоит, в выходе моделей обычно нет);
3. дефис → пробел («по-русски» и «по русски» — не ошибка распознавания);
4. знаки препинания снимаются полностью;
5. цифры разворачиваются в слова: e2e-модели пишут «5», эталон говорит «пять»;
6. пробелы сжимаются.

Отдельное правило отбора, не нормализации: записи, где ЦИФРЫ ЕСТЬ В ЭТАЛОНЕ,
из основного среза WER исключаются. Так же поступают авторы одного из используемых
наборов данных в собственных замерах, и причина честная: развернуть «1990» можно как «тысяча девятьсот девяносто» и как
«одна тысяча девятьсот девяносто», оба варианта верны, а WER посчитает один из них
ошибкой. Такие записи уходят в отдельный числовой срез (см. score.py).

Границы: порядковые числительные («5-й» → «пятый») не разворачиваются, римские цифры
не разворачиваются, падеж числительного не согласуется с окружением («у пять домов»).
Это осознанно — угадывание падежа внесло бы в нормализатор собственные ошибки.
"""

from __future__ import annotations

import re

# ── числительные ────────────────────────────────────────────────────────────────

_UNITS_M = ('ноль', 'один', 'два', 'три', 'четыре', 'пять', 'шесть', 'семь', 'восемь', 'девять')
_UNITS_F = ('ноль', 'одна', 'две', 'три', 'четыре', 'пять', 'шесть', 'семь', 'восемь', 'девять')
_TEENS = (
    'десять', 'одиннадцать', 'двенадцать', 'тринадцать', 'четырнадцать',
    'пятнадцать', 'шестнадцать', 'семнадцать', 'восемнадцать', 'девятнадцать',
)
_TENS = (
    '', '', 'двадцать', 'тридцать', 'сорок', 'пятьдесят',
    'шестьдесят', 'семьдесят', 'восемьдесят', 'девяносто',
)
_HUNDREDS = (
    '', 'сто', 'двести', 'триста', 'четыреста', 'пятьсот',
    'шестьсот', 'семьсот', 'восемьсот', 'девятьсот',
)
# разряд: (формы слова разряда, женский ли род у единиц этого разряда)
_SCALES: tuple[tuple[tuple[str, str, str], bool], ...] = (
    (('', '', ''), False),
    (('тысяча', 'тысячи', 'тысяч'), True),
    (('миллион', 'миллиона', 'миллионов'), False),
    (('миллиард', 'миллиарда', 'миллиардов'), False),
    (('триллион', 'триллиона', 'триллионов'), False),
)
# сколько цифр после запятой → название доли
_FRACTIONS = {
    1: ('десятая', 'десятых', 'десятых'),
    2: ('сотая', 'сотых', 'сотых'),
    3: ('тысячная', 'тысячных', 'тысячных'),
}
# единицы измерения, которые модель пишет знаком, а человек читает словом
_UNIT_WORDS = {
    '%': ('процент', 'процента', 'процентов'),
    '₽': ('рубль', 'рубля', 'рублей'),
    '$': ('доллар', 'доллара', 'долларов'),
    '€': ('евро', 'евро', 'евро'),
}


def plural(n: int, forms: tuple[str, str, str]) -> str:
    """Русский выбор формы: 1 дом, 2 дома, 5 домов. Считается по последним двум
    цифрам — 11 и 12 идут по третьей форме, а не по первой и второй."""
    n = abs(n) % 100
    if 11 <= n <= 14:
        return forms[2]
    n %= 10
    if n == 1:
        return forms[0]
    if 2 <= n <= 4:
        return forms[1]
    return forms[2]


def _triple(n: int, feminine: bool) -> list[str]:
    """Число 1..999 словами."""
    out: list[str] = []
    if n >= 100:
        out.append(_HUNDREDS[n // 100])
        n %= 100
    if 10 <= n <= 19:
        out.append(_TEENS[n - 10])
        return out
    if n >= 20:
        out.append(_TENS[n // 10])
        n %= 10
    if n:
        out.append((_UNITS_F if feminine else _UNITS_M)[n])
    return out


def int_to_words(n: int) -> str:
    """Целое число словами. Ноль и минус обрабатываются отдельно: в разрядной
    сборке ниже нулевая тройка не даёт ни одного слова."""
    if n == 0:
        return _UNITS_M[0]
    prefix = ['минус'] if n < 0 else []
    n = abs(n)
    triples: list[int] = []
    while n:
        triples.append(n % 1000)
        n //= 1000
    words: list[str] = []
    for idx in range(len(triples) - 1, -1, -1):
        value = triples[idx]
        if not value:
            continue
        forms, feminine = _SCALES[idx] if idx < len(_SCALES) else _SCALES[-1]
        words += _triple(value, feminine)
        if forms[0]:
            words.append(plural(value, forms))
    return ' '.join(prefix + words)


def _feminine_tail(n: int) -> str:
    """Целое, но единицы в женском роде: «одна целая», «две целых»."""
    words = int_to_words(n).split()
    if words and words[-1] in ('один', 'два'):
        words[-1] = 'одна' if words[-1] == 'один' else 'две'
    return ' '.join(words)


def _decimal_to_words(whole: str, frac: str) -> str:
    """«3,5» → «три целых пять десятых». Больше трёх знаков после запятой не
    разворачиваем — читают такие числа кто как, единого способа нет."""
    w, f = int(whole), int(frac)
    names = _FRACTIONS.get(len(frac))
    if names is None:
        return f'{int_to_words(w)} запятая ' + ' '.join(int_to_words(int(d)) for d in frac)
    return f'{_feminine_tail(w)} {plural(w, ("целая", "целых", "целых"))} {_feminine_tail(f)} {plural(f, names)}'


# ── порядковые в датах ──────────────────────────────────────────────────────────
# «21 августа» человек читает «двадцать первое августа», а не «двадцать один
# августа». Без этого правила наш собственный нормализатор добавляет ошибку там, где
# распознано верно — даты в речи встречаются постоянно. Общий случай порядковых
# («5-й дом») не берём: там нужен род и падеж, а угадывать их дороже, чем пропустить.
_MONTHS = (
    'января', 'февраля', 'марта', 'апреля', 'мая', 'июня',
    'июля', 'августа', 'сентября', 'октября', 'ноября', 'декабря',
)
_ORD_UNITS_N = (
    '', 'первое', 'второе', 'третье', 'четвертое', 'пятое',
    'шестое', 'седьмое', 'восьмое', 'девятое',
)
_ORD_TEENS_N = (
    'десятое', 'одиннадцатое', 'двенадцатое', 'тринадцатое', 'четырнадцатое',
    'пятнадцатое', 'шестнадцатое', 'семнадцатое', 'восемнадцатое', 'девятнадцатое',
)
_ORD_TENS_N = {20: 'двадцатое', 30: 'тридцатое'}
_DATE_RE = re.compile(r'\b(\d{1,2})\s+(' + '|'.join(_MONTHS) + r')\b')


def _day_ordinal(n: int) -> str:
    """День месяца порядковым числительным в среднем роде: 1..31."""
    if 10 <= n <= 19:
        return _ORD_TEENS_N[n - 10]
    if n % 10 == 0:
        return _ORD_TENS_N.get(n, '')
    if n > 20:
        return f'{_TENS[n // 10]} {_ORD_UNITS_N[n % 10]}'
    return _ORD_UNITS_N[n]


def expand_dates(text: str) -> str:
    def repl(m: 're.Match[str]') -> str:
        day = int(m.group(1))
        word = _day_ordinal(day) if 1 <= day <= 31 else ''
        return f'{word} {m.group(2)}' if word else m.group(0)

    return _DATE_RE.sub(repl, text)


_NUM_RE = re.compile(r'(\d+)(?:[.,](\d+))?')
# число со знаком единицы: «5%», «5 %», «12 000 ₽». Разделитель разрядов —
# только пробел МЕЖДУ цифрами: `[\d\s]*` съедал бы пробел перед числом и склеивал
# его с предыдущим словом
_SIGN_RE = re.compile(r'(\d+(?:[  ]\d{3})*)\s*([%₽$€])')

# Канонизация после разворота. «1990» читают и «тысяча девятьсот», и «одна тысяча
# девятьсот» — оба варианта верны, поэтому обе формы сводим к короткой. Применяется
# к эталону и к гипотезе одинаково, иначе разворот цифр сам создаёт ошибку там,
# где произнесено правильно.
_CANON = (
    (re.compile(r'\bодна тысяча\b'), 'тысяча'),
    (re.compile(r'\bодин миллион\b'), 'миллион'),
    (re.compile(r'\bодин миллиард\b'), 'миллиард'),
)


def expand_numbers(text: str) -> str:
    """Цифры и знаки единиц — словами. Порядок важен: сначала «5 %» → «пять
    процентов», иначе знак останется висеть без числа и будет снят как пунктуация."""
    def sign(m: 're.Match[str]') -> str:
        digits = re.sub(r'\s+', '', m.group(1))
        if not digits.isdigit():
            return m.group(0)
        forms = _UNIT_WORDS[m.group(2)]
        return f'{int_to_words(int(digits))} {plural(int(digits), forms)}'

    text = _SIGN_RE.sub(sign, text)

    def num(m: 're.Match[str]') -> str:
        whole, frac = m.group(1), m.group(2)
        return _decimal_to_words(whole, frac) if frac else int_to_words(int(whole))

    return _NUM_RE.sub(num, text)


# ── нормализация ────────────────────────────────────────────────────────────────

# всё, что не буква, не цифра и не пробел. Апостроф тоже уходит: в русских
# эталонах он не встречается, а модели иногда ставят его в заимствованиях
_PUNCT_RE = re.compile(r'[^\w\s]|_', re.UNICODE)
_SPACE_RE = re.compile(r'\s+')
_DIGIT_RE = re.compile(r'\d')
# «-5», но не «шоссе-5»: минусом считаем только знак, перед которым нет буквы или цифры
_MINUS_RE = re.compile(r'(?<![\w])-(?=\d)')

# знаки, по которым считается пунктуационный F1 (score.py). Многоточие и
# восклицание сводим к точке: смысловой границы это не меняет, а различать их
# у модели, обученной на разношёрстных данных, бессмысленно
PUNCT_CLASSES: dict[str, str | None] = {
    ',': ',', '.': '.', '!': '.', '?': '?', '…': '.',
    ':': ',', ';': ',', '—': ',', '-': None,
}


def has_digits(text: str) -> bool:
    """Есть ли в тексте цифры. По этому признаку запись уходит из основного среза
    WER в числовой — см. docstring модуля."""
    return bool(_DIGIT_RE.search(text or ''))


def normalize(text: str) -> str:
    """Полный нормализатор. Единственная функция, которой пользуются и эталон,
    и гипотезы; отдельных путей для «своих» и «чужих» текстов нет по определению."""
    t = (text or '').lower().replace('ё', 'е')
    # минус перед числом — слово, а не дефис: снимать его раньше, чем дефисы
    # превратятся в пробелы, иначе «-5» станет «пять» и потеряет знак
    t = _MINUS_RE.sub('минус ', t)
    t = t.replace('-', ' ').replace('–', ' ').replace('—', ' ')
    # даты раньше обычных чисел: иначе «21» уже станет «двадцать один»
    t = expand_dates(t)
    t = expand_numbers(t)
    t = _PUNCT_RE.sub(' ', t)
    t = _SPACE_RE.sub(' ', t).strip()
    for rx, repl in _CANON:
        t = rx.sub(repl, t)
    return t


def words_with_punct(text: str) -> list[tuple[str, str | None]]:
    """Слово плюс знак, стоящий сразу после него. Нужно для пунктуационного F1:
    сама пунктуация из WER выброшена, но проверить её надо — e2e-модель ставит
    знаки сама, и это её преимущество, а не шум.

    Возвращает пары (нормализованное слово, класс знака или None) в порядке текста.
    Слово может развернуться в несколько (цифры) — тогда знак вешается на последнее.
    """
    out: list[tuple[str, str | None]] = []
    t = (text or '').lower().replace('ё', 'е')
    for chunk in _SPACE_RE.split(_MINUS_RE.sub('минус ', t)):
        if not chunk:
            continue
        mark: str | None = None
        # знак берём только с конца слова: внутри слова он либо часть числа,
        # либо мусор, и на разметку границ фразы не влияет
        while chunk and not chunk[-1].isalnum():
            cls = PUNCT_CLASSES.get(chunk[-1])
            if cls and mark is None:
                mark = cls
            chunk = chunk[:-1]
        words = normalize(chunk).split()
        for i, w in enumerate(words):
            out.append((w, mark if i == len(words) - 1 else None))
    return out