#!/usr/bin/env python3
"""Test: full-page vs block-mode OCR en páginas bengalí-densas de CC (mad1).
Mide tiempo y muestra salida, para decidir si el modo por bloques evita el
bucle de repetición del modelo con el bengalí."""
import os, re, subprocess, tempfile, time, json
from pathlib import Path

os.environ.setdefault("RECOGNITION_BATCH_SIZE", "64")
os.environ.setdefault("DETECTOR_BATCH_SIZE", "8")

HERE = Path(__file__).resolve().parent
CFG = json.loads((HERE / "books.json").read_text())
ROOT = Path(CFG["_scan_root"])
PDF = ROOT / "originals" / "mad1.pdf"

from PIL import Image
import torch
print("device mps:", torch.backends.mps.is_available(), flush=True)
from surya.recognition import RecognitionPredictor
from surya.layout import LayoutPredictor

REC = RecognitionPredictor()
LAY = LayoutPredictor()

def render(pg, dpi=300):
    with tempfile.TemporaryDirectory() as td:
        pre = str(Path(td) / "r")
        subprocess.run(["pdftoppm", "-f", str(pg), "-l", str(pg), "-r", str(dpi),
                        "-png", str(PDF), pre], capture_output=True, timeout=120)
        imgs = sorted(Path(td).glob("r*.png"))
        return Image.open(imgs[0]).convert("RGB") if imgs else None

def page_text(page):
    out = []
    for blk in page.blocks:
        if getattr(blk, 'skipped', False) or getattr(blk, 'error', False):
            continue
        txt = re.sub(r'<[^>]+>', ' ', blk.html or '').strip()
        if txt:
            out.append(txt)
    return "\n".join(out)

for pg in (226, 230, 250):
    print(f"\n{'='*70}\nPÁGINA {pg}", flush=True)
    img = render(pg)
    if img is None:
        print("  no renderizó"); continue

    # 1) full-page (lo que hace producción)
    t0 = time.time()
    try:
        res = REC([img], full_page=True)
        txt = page_text(res[0])
        print(f"  [full_page] {time.time()-t0:6.1f}s  {len(txt)} chars", flush=True)
        print("   >>", repr(txt[:200]))
    except Exception as e:
        print(f"  [full_page] {time.time()-t0:6.1f}s  ERROR: {e}", flush=True)

    # 2) block-mode (layout → OCR por bloque)
    t0 = time.time()
    try:
        layout = LAY([img])
        t_lay = time.time() - t0
        res = REC([img], layout_results=layout, full_page=False)
        txt = page_text(res[0])
        print(f"  [block]     {time.time()-t0:6.1f}s (layout {t_lay:.1f}s)  {len(txt)} chars", flush=True)
        print("   >>", repr(txt[:200]))
        # etiquetas detectadas
        labels = [b.label for b in layout[0].bboxes]
        print("   labels:", labels)
    except Exception as e:
        print(f"  [block]     {time.time()-t0:6.1f}s  ERROR: {e}", flush=True)