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

Собирает финальный тренировочный текстовый файл для tesstrain из скачанного
корпуса Викитеки - с ограничением по объёму (для дообучения не нужны
все 400 страниц, хватит representative выборки) и базовой чисткой
(убирает пустые строки, слишком короткие "строки-мусор").
"""

import os
import re
import csv

CORPUS_DIR = "/home/dayhanbiz/public_html/3-damja.science/ТЕСТ_образцы_языков/OCR_дообучение_дореформенный/Корпус_wikisource"
MANIFEST_FILE = os.path.join(CORPUS_DIR, "manifest.csv")
OUTPUT_FILE = os.path.join(CORPUS_DIR, "training_text.txt")

TARGET_WORDS = 250000  # больше данных для второго захода дообучения
MIN_LINE_CHARS = 8    # пропускаем совсем короткие "строки" (заголовки-огрызки и т.п.)

WORD_RE = re.compile(r"\w+", re.UNICODE)


def clean_lines(text):
    lines = text.split("\n")
    result = []
    for line in lines:
        line = line.strip()
        if len(line) < MIN_LINE_CHARS:
            continue
        # убираем цифровые сноски-огрызки типа "12" в начале строки, оставленные разметкой
        line = re.sub(r"^\d+\s+", "", line)
        result.append(line)
    return result


def main():
    with open(MANIFEST_FILE, encoding="utf-8", newline="") as f:
        reader = csv.DictReader(f, delimiter=";")
        rows = list(reader)

    # тематические файлы сначала
    rows.sort(key=lambda r: r["topical"] != "True")

    total_words = 0
    out_lines = []
    used_files = 0

    for row in rows:
        if total_words >= TARGET_WORDS:
            break
        path = os.path.join(CORPUS_DIR, row["filename"])
        if not os.path.exists(path):
            continue
        with open(path, encoding="utf-8") as f:
            text = f.read()

        lines = clean_lines(text)
        file_word_count = sum(len(WORD_RE.findall(l)) for l in lines)
        if file_word_count == 0:
            continue

        out_lines.extend(lines)
        out_lines.append("")  # разделитель между произведениями
        total_words += file_word_count
        used_files += 1

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

    print(f"Файлов использовано: {used_files}")
    print(f"Слов итого: {total_words}")
    print(f"Строк итого: {len(out_lines)}")
    print(f"Сохранено: {OUTPUT_FILE}")


if __name__ == "__main__":
    main()