"""Azerbaijani line & page OCR — CRNN + CTC recogniser (LocalDoc/azerbaijani-ocr-crnn).

Two modes:
  * Line  — the recogniser alone on a cropped line (the model's native input).
  * Page  — PP-OCRv5 detection finds the lines, the recogniser reads them.

Vendored from the model repo (CC-BY-4.0): `modeling.py` and the detection /
reading-order logic of `pipeline.py`.
"""

import os

# Must precede the torch import: allocator fragmentation under the transient
# pressure of many wide crops (see CUDACachingAllocator NVML assert).
os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
os.environ.setdefault("FLAGS_use_mkldnn", "0")   # oneDNN clashes with the PIR executor
os.environ.setdefault("GLOG_minloglevel", "2")

import spaces  # noqa: F401  — must precede any CUDA-touching import
import tempfile
import time
import threading

import numpy as np
import torch
from PIL import Image
from huggingface_hub import hf_hub_download

import gradio as gr

from modeling import load, preprocess, ctc_greedy, HEIGHT, DOWNSAMPLE

MODEL_ID = "LocalDoc/azerbaijani-ocr-crnn"
DET_MODEL = "PP-OCRv5_server_det"
UNCLIP = 1.5            # box inflation; raise if ç ş ğ descenders get clipped
PAD_FRAC = 0.12         # vertical headroom for descenders
MAX_WIDTH = 1200
BATCH = 12

# --------------------------------------------------------------------------- model
_path = hf_hub_download(MODEL_ID, "model.pt")
_model, _itos = load(_path, "cuda")     # eager .to("cuda"); ZeroGPU packs it at startup
_model.eval()
print(f"recogniser ready (training={_model.training})", flush=True)

# ------------------------------------------------------------------- detector (lazy)
_det = None
_det_lock = threading.Lock()


def _detector():
    """PP-OCRv5 detector, built once and reused by the warm worker.

    Constructed lazily: `import spaces` has already run by the time this is
    first called, so importing paddle inside a `@spaces.GPU` call is safe and
    paddle never becomes a hard module-scope dependency.
    """
    global _det
    if _det is None:
        with _det_lock:
            if _det is None:
                from paddleocr import TextDetection
                _det = TextDetection(model_name=DET_MODEL, enable_mkldnn=False,
                                     unclip_ratio=UNCLIP)
    return _det


# Warm the detector's weights into the image cache at build/startup time, so a
# cold GPU worker doesn't spend its first request downloading them. Paddle is a
# CPU build, so this does not initialise CUDA. Non-fatal if it fails: page mode
# then pays the download once on its first call.
try:
    _detector()
    print("PP-OCRv5 detector ready", flush=True)
except Exception as _e:  # noqa: BLE001
    print(f"detector pre-warm skipped ({_e!r}); will build lazily", flush=True)


# ------------------------------------------------------- reading order (from pipeline.py)
def _reading_order(boxes):
    """Cluster boxes into rows by y, then left to right within each row."""
    if not boxes:
        return []
    arr = [np.array(b) for b in boxes]
    tol = max(6.0, float(np.median([np.ptp(b[:, 1]) for b in arr])) * 0.6)
    idx = sorted(range(len(arr)), key=lambda i: arr[i][:, 1].mean())
    rows, cur = [], [idx[0]]
    for i in idx[1:]:
        if abs(arr[i][:, 1].mean() - arr[cur[-1]][:, 1].mean()) <= tol:
            cur.append(i)
        else:
            rows.append(cur)
            cur = [i]
    rows.append(cur)
    out = []
    for row in rows:
        out += sorted(row, key=lambda i: arr[i][:, 0].min())
    return [boxes[i] for i in out]


def _pad(boxes, height):
    """Pad above and below, never more than half the gap to the neighbouring line."""
    order = sorted(range(len(boxes)), key=lambda i: boxes[i][1])
    out = list(boxes)
    for k, i in enumerate(order):
        x0, y0, x1, y1 = boxes[i]
        want = (y1 - y0) * PAD_FRAC
        up = y0 - boxes[order[k - 1]][3] if k > 0 else want * 2
        dn = boxes[order[k + 1]][1] - y1 if k + 1 < len(order) else want * 2
        out[i] = (x0, max(0.0, y0 - min(want, max(0.0, up) * 0.45)),
                  x1, min(float(height), y1 + min(want * 1.4, max(0.0, dn) * 0.45)))
    return out


def _forward(imgs):
    """One forward pass on a list of HxW grayscale crops -> list of T x C arrays.

    Each crop is decoded from its own T columns only, so sub-batches may pad to
    different widths and still be merged as a plain list.
    """
    W = max(im.width for im in imgs)
    W = int(np.ceil(W / DOWNSAMPLE) * DOWNSAMPLE)
    x = torch.ones(len(imgs), 1, HEIGHT, W)
    for k, im in enumerate(imgs):
        a = (np.array(im, dtype=np.float32) / 255.0 - 0.5) / 0.5
        x[k, 0, :, :im.width] = torch.from_numpy(a)
    with torch.no_grad():
        ids = _model(x.to("cuda")).cpu().argmax(-1).numpy()
    return [ids[k][:max(1, im.width // DOWNSAMPLE)] for k, im in enumerate(imgs)]


def _batched(part):
    """Recognise a batch, halving it if the GPU refuses the allocation.

    A very wide batch can exhaust the worker's memory pool (the CUDA caching
    allocator raises an NVML internal assert rather than a clean OOM); retrying
    with smaller sub-batches keeps a big page from failing outright.
    """
    try:
        return _forward(part)
    except RuntimeError:
        if len(part) == 1:
            raise
        mid = len(part) // 2
        return _batched(part[:mid]) + _batched(part[mid:])


def _recognise(crops):
    """Batched recognition (width-sorted so each batch pads to one width)."""
    texts = []
    order = sorted(range(len(crops)), key=lambda i: crops[i].width)
    for s in range(0, len(order), BATCH):
        part = order[s:s + BATCH]
        rows = _batched([crops[i] for i in part])
        for k, i in enumerate(part):
            prev, s_ = -1, []
            for c in rows[k]:
                if c != prev and c != 0:
                    s_.append(_itos[c])
                prev = c
            texts.append((i, "".join(s_)))
    return [t for _, t in sorted(texts)]


def _detect_polys(image):
    """Run the detector via a unique temp file (no fixed path: handlers run concurrently)."""
    with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as fh:
        image.convert("RGB").save(fh.name)
        tmp = fh.name
    try:
        polys = []
        for res in _detector().predict(tmp):
            d = res if isinstance(res, dict) else (getattr(res, "json", {}) or {}).get("res", {})
            polys += [np.array(p).reshape(-1, 2) for p in d.get("dt_polys", [])]
        return polys
    finally:
        try:
            os.remove(tmp)
        except OSError:
            pass


# ------------------------------------------------------------------------ inference
@spaces.GPU(duration=25)
def read_line(image: Image.Image, pad: int = 8) -> str:
    """Read a single cropped line of Azerbaijani text (Latin or Cyrillic script).

    Args:
        image: a crop containing exactly one line of text, dark text on a light background.
        pad: extra white margin, in pixels, wrapped around the crop before recognition.
    """
    if image is None:
        raise gr.Error("Upload an image of one text line first.")
    img = image.convert("L")
    pad = max(0, int(pad))
    if pad:
        canvas = Image.new("L", (img.width + 2 * pad, img.height + 2 * pad), 255)
        canvas.paste(img, (pad, pad))
        img = canvas
    with torch.no_grad():
        logits = _model(preprocess(img).to("cuda"))
    return ctc_greedy(logits.cpu(), _itos)


@spaces.GPU(duration=60)
def read_page(image: Image.Image, col: str = "single",
              join_hyphens: bool = False) -> str:
    """Read a whole page scan: detect the lines, then recognise each one.

    Reading order assumes a single text column. For a two-column page, crop the
    left half in "Left column" mode and the right half in "Right column" mode
    (the columns are then read in the right order).

    Args:
        image: a page scan, roughly upright, at least ~150 dpi.
        col: "single" for a one-column page, "left" or "right" to isolate one column of a two-column page.
        join_hyphens: reattach words broken across a line break.
    """
    if image is None:
        raise gr.Error("Upload a page image first.")
    img = image.convert("RGB")
    if col == "left":
        img = img.crop((0, 0, img.width // 2, img.height))
    elif col == "right":
        img = img.crop((img.width // 2, 0, img.width, img.height))

    polys = _detect_polys(img)
    if not polys:
        return "(no text detected)"
    polys = _reading_order(polys)
    boxes = [(float(p[:, 0].min()), float(p[:, 1].min()),
              float(p[:, 0].max()), float(p[:, 1].max())) for p in polys]
    boxes = _pad(boxes, img.height)

    crops = []
    for x0, y0, x1, y1 in boxes:
        if (y1 - y0) >= (x1 - x0):          # vertical text, not a line
            continue
        c = img.convert("L").crop((int(x0), int(y0), int(x1), int(y1)))
        if c.width < 8 or c.height < 8:
            continue
        r = HEIGHT / c.height
        c = c.resize((max(8, int(c.width * r)), HEIGHT), Image.LANCZOS)
        if c.width > MAX_WIDTH:
            c = c.crop((0, 0, MAX_WIDTH, HEIGHT))
        crops.append(c)

    lines = _recognise(crops)
    if not join_hyphens:
        return "\n".join(lines)
    out = []
    for ln in lines:
        if out and out[-1].endswith("-"):
            out[-1] = out[-1][:-1] + ln.lstrip()
        else:
            out.append(ln)
    return "\n".join(out)


# ------------------------------------------------------------------------------- UI
CSS = """
#col-container { max-width: 1100px; margin: 0 auto; }
.dark .gradio-container { color: var(--body-text-color); }
"""

with gr.Blocks(title="Azerbaijani OCR") as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown(
            "# Azerbaijani OCR — CRNN + CTC\n"
            "An 8.5M-parameter recogniser for Azerbaijani print, **Latin and Cyrillic** "
            "script. It reads one line at a time; page mode pairs it with a text "
            "detector to do whole scans.\n\n"
            "[model card](https://huggingface.co/LocalDoc/azerbaijani-ocr-crnn) · "
            "[benchmark](https://huggingface.co/datasets/LocalDoc/azerbaijani-ocr-benchmark) "
            "(CER 0.0242 / WER 0.0565)"
        )

        with gr.Tabs():
            with gr.Tab("Line"):
                with gr.Row():
                    with gr.Column(scale=1):
                        line_in = gr.Image(type="pil", label="One text line",
                                           sources=["upload", "clipboard"])
                        line_btn = gr.Button("Read line", variant="primary")
                        with gr.Accordion("Options", open=False):
                            line_pad = gr.Slider(0, 24, value=8, step=1,
                                                 label="White margin (px)")
                    with gr.Column(scale=1):
                        line_out = gr.Textbox(label="Recognised text", lines=4,
                                              buttons=["copy"])

            with gr.Tab("Page"):
                with gr.Row():
                    with gr.Column(scale=1):
                        page_in = gr.Image(type="pil", label="Page scan",
                                           sources=["upload", "clipboard"])
                        page_btn = gr.Button("Read page", variant="primary")
                        with gr.Accordion("Options", open=False):
                            page_col = gr.Radio(["single", "left", "right"],
                                                value="single",
                                                label="Column layout",
                                                info="Two-column pages: read the left "
                                                     "half, then the right half.")
                            page_hyph = gr.Checkbox(value=False,
                                                    label="Rejoin hyphenated words")
                    with gr.Column(scale=1):
                        page_out = gr.Textbox(label="Page text", lines=16,
                                              buttons=["copy"])

        line_btn.click(read_line, inputs=[line_in, line_pad], outputs=line_out,
                       api_name="read_line")
        page_btn.click(read_page, inputs=[page_in, page_col, page_hyph],
                       outputs=page_out, api_name="read_page")

        gr.Markdown("### Examples")
        gr.Markdown(
            "Real line crops and pages from "
            "[LocalDoc/azerbaijani-ocr-benchmark](https://huggingface.co/datasets/LocalDoc/azerbaijani-ocr-benchmark) "
            "(CC-BY-4.0), the authors' own evaluation set."
        )
        gr.Examples(
            examples=[["assets/line_latin_0.png"], ["assets/line_cyrillic_0.png"]],
            inputs=[line_in], outputs=line_out, fn=read_line,
            cache_examples=True, cache_mode="lazy",
            label="Line mode",
        )
        gr.Examples(
            examples=[["assets/page_latin.png", "single", False],
                      ["assets/page_cyrillic.png", "single", False]],
            inputs=[page_in, page_col, page_hyph], outputs=page_out, fn=read_page,
            cache_examples=True, cache_mode="lazy",
            label="Page mode",
        )

if __name__ == "__main__":
    demo.launch(theme=gr.themes.Citrus(), css=CSS, mcp_server=True)