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

Генерирует обучающие пары (картинка строки .tif + текст .gt.txt) для tesstrain
из нашего подготовленного текста, используя text2image - по каждой строке,
для каждого из проверенных шрифтов. Результат кладётся в
data/MODEL_NAME-ground-truth/, как того ожидает tesstrain (см. его README).
"""

import os
import re
import subprocess
import tempfile
from concurrent.futures import ProcessPoolExecutor, as_completed

TRAINING_TEXT = "/home/dayhanbiz/public_html/3-damja.science/ТЕСТ_образцы_языков/OCR_дообучение_дореформенный/Корпус_wikisource/training_text.txt"
GROUND_TRUTH_DIR = "/home/dayhanbiz/public_html/3-damja.science/ТЕСТ_образцы_языков/OCR_дообучение_дореформенный/tesstrain/data/rus_old-ground-truth"
FONTS_DIR = "/usr/share/fonts"

FONTS = ["Liberation Serif", "DejaVu Serif", "Nimbus Roman", "P052", "URW Bookman Light", "C059", "Carlito"]

MAX_WORKERS = 6
MIN_LINE_CHARS = 8
MAX_LINE_CHARS = 150  # слишком длинные строки режем - для одной строки-картинки не годятся


def load_lines():
    with open(TRAINING_TEXT, encoding="utf-8") as f:
        raw = f.read().split("\n")
    lines = []
    for line in raw:
        line = line.strip()
        if MIN_LINE_CHARS <= len(line) <= MAX_LINE_CHARS:
            lines.append(line)
    return lines


def gen_one(args):
    idx, line, font = args
    font_safe = font.replace(" ", "_")
    base_name = f"line{idx:05d}.{font_safe}.exp0"
    base_path = os.path.join(GROUND_TRUTH_DIR, base_name)

    gt_path = base_path + ".gt.txt"
    tif_path = base_path + ".tif"

    if os.path.exists(tif_path) and os.path.exists(gt_path):
        return None  # уже сделано - пропускаем (можно продолжать прерванную генерацию)

    with tempfile.NamedTemporaryFile("w", suffix=".txt", delete=False, encoding="utf-8") as tf:
        tf.write(line + "\n")
        tmp_text_path = tf.name

    try:
        env = os.environ.copy()
        env["OMP_THREAD_LIMIT"] = "1"
        cmd = [
            "text2image",
            "--strip_unrenderable_words",
            "--xsize=3000", "--ysize=200",
            "--ptsize=18",
            "--leading=32", "--margin=12",
            "--char_spacing=0.0", "--exposure=0",
            "--max_pages=1",
            f"--fonts_dir={FONTS_DIR}",
            f"--font={font}",
            f"--text={tmp_text_path}",
            f"--outputbase={base_path}",
        ]
        result = subprocess.run(cmd, capture_output=True, text=True, timeout=60, env=env)
        if result.returncode != 0:
            return f"ошибка text2image (line{idx}, {font}): {result.stderr.strip()[:200]}"

        with open(gt_path, "w", encoding="utf-8") as f:
            f.write(line + "\n")
        return None
    finally:
        os.unlink(tmp_text_path)


def main():
    os.makedirs(GROUND_TRUTH_DIR, exist_ok=True)

    lines = load_lines()
    print(f"Строк для генерации: {len(lines)}")
    print(f"Шрифтов: {len(FONTS)}")
    total_pairs = len(lines) * len(FONTS)
    print(f"Всего пар изображение+текст к созданию: {total_pairs}")

    tasks = [(idx, line, font) for idx, line in enumerate(lines) for font in FONTS]

    done = 0
    errors = 0
    with ProcessPoolExecutor(max_workers=MAX_WORKERS) as executor:
        futures = {executor.submit(gen_one, t): t for t in tasks}
        for fut in as_completed(futures):
            err = fut.result()
            done += 1
            if err:
                errors += 1
                print(f"  {err}")
            if done % 200 == 0:
                print(f"[{done}/{total_pairs}] готово (ошибок: {errors})", flush=True)

    print(f"\n🎉 Готово! Создано пар: {done - errors}, ошибок: {errors}")
    print(f"📁 Ground truth: {GROUND_TRUTH_DIR}")


if __name__ == "__main__":
    main()