"""Довідкова реалізація нормалізатора — повний конвеєр, як він задуманий.

Три кроки, і кожен із них потрібен. Якщо викинути перший або другий, якість
падає на конкретних класах входу, і це заміряно (див. README, розділ
«Препроцесинг обов'язковий»).

    from normalize import Normalizer
    n = Normalizer("skypro1111/<repo>")
    n("Рахунок на 1111 386,40 грн треба сплатити до 25.07.2026.")
    # 'Рахунок на один мільйон сто одинадцять тисяч триста вісімдесят шість
    #  гривень сорок копійок треба сплатити до двадцять п'ятого липня
    #  дві тисячі двадцять шостого року.'

Модель навчена на ОДНОМУ реченні за раз. Довгий текст ріжте на речення й
подавайте окремо — інакше впертеся в ліміт нових токенів.
"""
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

from digit_router import route as _digit_route
from bignum_spacing import add_spacing as _add_spacing

MAX_NEW_TOKENS = 192


class Normalizer:
    def __init__(self, model_id: str, device: str = None, dtype=torch.bfloat16):
        self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")
        self.tok = AutoTokenizer.from_pretrained(model_id)
        self.model = AutoModelForCausalLM.from_pretrained(
            model_id, dtype=dtype).to(self.device).eval()

        # ban-digits: модель має ПИСАТИ словами, тож будь-який токен із цифрою
        # заборонено. Без цього вона зрідка копіює цифру з входу, і рядок стає
        # непридатним для синтезу мовлення.
        self.bad_words = [[i] for t, i in self.tok.get_vocab().items()
                          if any(ch.isdigit() for ch in t)]

    @torch.inference_mode()
    def _generate(self, text: str) -> str:
        msgs = [{"role": "user", "content": text}]
        ids = self.tok.apply_chat_template(
            msgs, add_generation_prompt=True, return_tensors="pt").to(self.device)
        out = self.model.generate(
            ids,
            max_new_tokens=MAX_NEW_TOKENS,
            do_sample=False,                 # greedy: жодного семплювання
            bad_words_ids=self.bad_words,
            pad_token_id=self.tok.pad_token_id,
        )
        return self.tok.decode(out[0, ids.shape[-1]:],
                               skip_special_tokens=True).strip()

    def __call__(self, text: str) -> str:
        text = (text or "").strip()
        if not text:
            return ""

        # 1) голий рядок цифр читається по цифрі детерміновано, повз модель:
        #    рахувати однакові цифри вона не вміє («000000» -> чотири «нуль»)
        handled, routed = _digit_route(text)
        if handled:
            return routed

        # 2) довгі числа надійніше читаються групами по три — вставляємо
        #    роздільники розрядів ПЕРЕД подачею
        text = _add_spacing(text)

        # 3) власне модель
        return self._generate(text)


if __name__ == "__main__":
    import sys
    n = Normalizer(sys.argv[1] if len(sys.argv) > 1 else ".")
    for t in ("Зустріч о 08:05 у кабінеті 12.",
              "Збитки оцінили в 7000000000 доларів.",
              "Пиши на support@shop.com.ua або дзвони 067 123 45 67.",
              "Борг сягнув 2625414042726 гривень."):
        print(f"IN : {t}\nOUT: {n(t)}\n")