#!/usr/bin/env python3.11
# -*- coding: utf-8 -*-
"""
normalize_training_text.py

Нормализует training_text.txt под символы, которые реально поддерживает
базовая модель rus (читает точный список из lstm-unicharset), заменяет
типографские варианты пробелов/многоточия на обычные, и отбрасывает
строки, где всё равно остались неподдерживаемые символы.
"""

import re
import os

TRAINING_TEXT = "/home/dayhanbiz/public_html/3-damja.science/ТЕСТ_образцы_языков/OCR_дообучение_дореформенный/Корпус_wikisource/training_text.txt"
# Базовый unicharset САМОЙ МОДЕЛИ rus (до объединения с нашими буквами) - это то,
# что модель реально умеет кодировать/распознавать на старте
BASE_UNICHARSET_FILE = "/home/dayhanbiz/public_html/3-damja.science/ТЕСТ_образцы_языков/OCR_дообучение_дореформенный/tesstrain/data/rus/rus_old.lstm-unicharset"
OUTPUT_FILE = TRAINING_TEXT

# Дореформенные буквы, которые мы осознанно добавляем сверх базового набора -
# независимо от того, встретились ли они в какой-то конкретной уже собранной партии
OLD_RU_EXTRA_CHARS = set("ѣѢѳѲіІѵѴ")

# Типографские варианты -> обычные символы (безопасные, всегда есть в базовой кириллице)
REPLACEMENTS = {
    "\u00a0": " ",   # неразрывный пробел
    "\u202f": " ",   # узкий неразрывный пробел
    "\u2009": " ",   # тонкий пробел
    "\u2026": "...", # многоточие одним символом -> три точки
    "\u2013": "-",   # короткое тире (en dash)
}


def load_allowed_chars(unicharset_path):
    allowed = set()
    with open(unicharset_path, encoding="utf-8") as f:
        lines = f.readlines()
    # первая строка - количество символов, дальше построчно: символ и свойства через пробел
    for line in lines[1:]:
        parts = line.split()
        if not parts:
            continue
        symbol = parts[0]
        allowed.update(symbol)
    # пробел и обычные "безопасные" символы добавляем на всякий случай
    allowed.update(" \n\t")
    return allowed


def main():
    if not os.path.exists(BASE_UNICHARSET_FILE):
        print(f"Не найден файл {BASE_UNICHARSET_FILE}")
        print("Он появляется после этапа 'unicharset' в make training - если тренировка")
        print("успела дойти хотя бы до вывода 'Wrote unicharset file...', он должен быть.")
        raise SystemExit(1)

    allowed = load_allowed_chars(BASE_UNICHARSET_FILE)
    allowed |= OLD_RU_EXTRA_CHARS
    print(f"Разрешённых символов (база rus + дореформенные буквы): {len(allowed)}")
    print(f"Дореформенные буквы среди них: {sorted(OLD_RU_EXTRA_CHARS & allowed)}")

    with open(TRAINING_TEXT, encoding="utf-8") as f:
        lines = f.read().split("\n")

    kept = []
    dropped = 0
    for line in lines:
        if not line.strip():
            kept.append(line)
            continue

        fixed = line
        for bad, good in REPLACEMENTS.items():
            fixed = fixed.replace(bad, good)

        unsupported = set(fixed) - allowed
        if unsupported:
            dropped += 1
            continue

        kept.append(fixed)

    with open(OUTPUT_FILE, "w", encoding="utf-8") as f:
        f.write("\n".join(kept))

    print(f"Строк оставлено: {sum(1 for l in kept if l.strip())}")
    print(f"Строк отброшено (остались неподдерживаемые символы): {dropped}")
    print(f"Файл перезаписан: {OUTPUT_FILE}")


if __name__ == "__main__":
    main()