"""MiniCPM-OCR document parsing for a Hugging Face ZeroGPU Space.

Both models are PyTorch. ZeroGPU only lends a GPU inside the Gradio function
decorated with @spaces.GPU, so layout and recognition happen in that one call.
"""

from __future__ import annotations

import os

# ZeroGPU patches torch at import time. Import it before torch on a Space,
# and keep a no-op decorator for a normal local process.
if os.environ.get("SPACE_ID"):
    import spaces
else:

    class spaces:  # noqa: N801
        @staticmethod
        def GPU(*_args, **_kwargs):
            def decorate(fn):
                return fn

            return decorate


import numpy as np
import torch
from PIL import Image
from transformers import (
    AutoImageProcessor,
    AutoModelForImageClassification,
    AutoModelForObjectDetection,
    AutoProcessor,
    MiniCPMOcrForConditionalGeneration,
)

from layout import detect as detect_layout
from postprocess import format_block

OCR_ID = os.environ.get("OCR_MODEL_ID", "openbmb/MiniCPM-OCR-1.0")
LAYOUT_ID = os.environ.get("LAYOUT_MODEL_ID", "PaddlePaddle/PP-DocLayoutV3_safetensors")
ORI_ID = os.environ.get("ORI_MODEL_ID", "PaddlePaddle/PP-LCNet_x1_0_doc_ori_safetensors")

PROMPTS = {
    "text": ("Text Recognition:", 2048),
    "content": ("Text Recognition:", 2048),
    "abstract": ("Text Recognition:", 2048),
    "reference": ("Text Recognition:", 1024),
    "reference_content": ("Text Recognition:", 1024),
    "vertical_text": ("Text Recognition:", 1024),
    "vision_footnote": ("Text Recognition:", 512),
    "algorithm": ("Text Recognition:", 2048),
    "doc_title": ("Text Recognition:", 256),
    "paragraph_title": ("Text Recognition:", 256),
    "figure_title": ("Text Recognition:", 256),
    "table": ("Table Recognition:", 2048),
    "formula": ("Formula Recognition:", 1024),
    "display_formula": ("Formula Recognition:", 1024),
    "inline_formula": ("Text Recognition:", 512),
    "formula_number": ("Text Recognition:", 128),
    "header": ("Text Recognition:", 256),
    "footer": ("Text Recognition:", 256),
    "footnote": ("Text Recognition:", 512),
    "aside_text": ("Text Recognition:", 512),
    "number": ("Text Recognition:", 128),
    "chart": ("Chart to Table:", 2048),
    "seal": ("Seal Recognition:", 256),
}
SKIP_MARKDOWN = {"number", "footnote", "header", "footer", "aside_text", "formula_number"}
SKIP_VLM = {"image", "header_image", "footer_image"}
COPY = {
    "en": {
        "language_label": "Language",
        "intro": "# Document parsing\nUpload an image.",
        "limit_note": (
            'Because of ZeroGPU limits, this Space can only run Transformers serially on a single image. '
            'For a smoother, more complete experience, use the <a href="https://platform.modelbest.ai/console/playground/document-parse" target="_blank" rel="noopener noreferrer">official demo</a>.'
        ),
        "image": "Image",
        "upload": "Drop an image here\n- or -\nclick to upload",
        "upright": "Auto-rotate",
        "upright_info": (
            "Corrects a small skew and turns the page to 0°, 90°, 180°, or 270° "
            "before layout detection and recognition."
        ),
        "seal": "Recognize seals",
        "seal_info": "Reads the text inside seals. Turn off to skip them.",
        "chart": "Recognize charts",
        "chart_info": "Converts charts into tables. Turn off to skip them.",
        "parse": "Parse",
        "result": "Recognition result",
        "view_label": "Display",
        "rendered": "Rendered",
        "source": "Markdown",
        "parsing": "Parsing…",
        "need_image": "Please upload an image.",
        "empty": "No content was recognized on this page.",
    },
    "zh": {
        "language_label": "语言",
        "intro": "# 文档解析\n上传一张图片。",
        "limit_note": (
            '由于 ZeroGPU 的限制，这里只能用 Transformers 串行、单图推理。需要更流畅、更全面的体验，'
            '请前往<a href="https://platform.modelbest.ai/console/playground/document-parse" target="_blank" rel="noopener noreferrer">官方 Demo</a>。'
        ),
        "image": "图片",
        "upload": "将图片拖到这里\n- 或 -\n点击上传",
        "upright": "自适应旋转",
        "upright_info": "先做小角度摆正和 0°/90°/180°/270° 转向，再做版面检测和识别。",
        "seal": "印章识别",
        "seal_info": "识别印章里的文字。关闭后跳过印章。",
        "chart": "图表识别",
        "chart_info": "把图表转成表格。关闭后跳过图表。",
        "parse": "解析",
        "result": "识别结果",
        "view_label": "展示",
        "rendered": "渲染",
        "source": "原始 Markdown",
        "parsing": "正在解析…",
        "need_image": "请上传一张图片。",
        "empty": "这一页没有识别出内容。",
    },
}
RESULT_CSS = """
.ocr-columns { align-items: flex-start !important; gap: 16px; }
.ocr-columns > .column { flex: 1 1 0 !important; min-width: 0 !important; width: calc(50% - 8px) !important; }
.ocr-col > .block, .ocr-col > .form { width: 100% !important; }
.html:has(.ocr-note),
.html-container:has(.ocr-note) { padding: 0 !important; }
.ocr-note { margin: 0 0 12px; padding: 0; color: #6b7280; font-size: 13px; line-height: 1.5; }
.ocr-note a { color: #2563eb; margin: 0; padding: 0; }
.ocr-result table { border-collapse: collapse; width: 100%; margin: 8px 0; }
.ocr-result td, .ocr-result th { border: 1px solid #d1d5db; padding: 4px 8px; vertical-align: top; }
"""

layout_model = AutoModelForObjectDetection.from_pretrained(LAYOUT_ID, dtype=torch.float32).to("cuda").eval()
ori_processor = AutoImageProcessor.from_pretrained(ORI_ID)
ori_model = AutoModelForImageClassification.from_pretrained(ORI_ID, dtype=torch.float32).to("cuda").eval()
ocr_processor = AutoProcessor.from_pretrained(OCR_ID)
ocr_model = MiniCPMOcrForConditionalGeneration.from_pretrained(OCR_ID, dtype=torch.bfloat16).to("cuda").eval()


def _copy(language: str) -> dict:
    return COPY["zh"] if language == "zh" else COPY["en"]


def _note(language: str) -> str:
    return f'<p class="ocr-note">{_copy(language)["limit_note"]}</p>'


def _duration(image, upright, recognize_seal, recognize_chart, language):
    del image, upright, recognize_seal, recognize_chart, language
    return 90


def _rotate(image: Image.Image, angle: float) -> Image.Image:
    if abs(angle) < 0.3 or abs(abs(angle) - 360) < 0.3:
        return image
    nearest = round(angle / 90.0) * 90.0
    if abs(angle - nearest) < 0.3:
        turns = int(nearest) % 360
        return image.rotate(turns, expand=True, fillcolor=(255, 255, 255))
    return image.rotate(angle, expand=True, fillcolor=(255, 255, 255), resample=Image.Resampling.BICUBIC)


def _projection_score(binary: np.ndarray) -> float:
    profile = binary.sum(axis=1).astype(np.float64)
    mean = float(profile.mean())
    if mean < 1e-6:
        return 0.0
    delta = np.diff(profile)
    return float((delta * delta).mean() / (mean * mean))


def _deskew_angle(image: Image.Image) -> float:
    import cv2

    gray = np.asarray(image.convert("L"))
    long_side = max(gray.shape)
    if long_side > 900:
        scale = 900 / long_side
        gray = cv2.resize(
            gray,
            (max(1, int(gray.shape[1] * scale)), max(1, int(gray.shape[0] * scale))),
            interpolation=cv2.INTER_AREA,
        )
    block = max(15, (min(gray.shape) // 35) | 1)
    binary = cv2.adaptiveThreshold(
        gray, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, cv2.THRESH_BINARY_INV, block, 15
    )

    def score(angle: float) -> float:
        height, width = binary.shape
        matrix = cv2.getRotationMatrix2D((width / 2, height / 2), angle, 1.0)
        rotated = cv2.warpAffine(binary, matrix, (width, height), flags=cv2.INTER_NEAREST, borderValue=0)
        return _projection_score(rotated)

    angles = np.linspace(-15, 15, 31)
    scores = [score(float(angle)) for angle in angles]
    best = float(angles[int(np.argmax(scores))])
    straight = scores[len(scores) // 2]
    if abs(best) <= 0.5 and straight >= max(scores) * 0.98:
        return 0.0
    fine = np.linspace(best - 1.0, best + 1.0, 11)
    fine_scores = [score(float(angle)) for angle in fine]
    return float(fine[int(np.argmax(fine_scores))])


def _orientation_degrees(image: Image.Image) -> int:
    preview = image
    long_side = max(image.size)
    if long_side > 1600:
        scale = 1600 / long_side
        preview = image.resize((max(1, int(image.width * scale)), max(1, int(image.height * scale))))
    inputs = ori_processor(images=preview, return_tensors="pt")
    inputs = {key: value.to("cuda") for key, value in inputs.items()}
    with torch.inference_mode():
        label_id = int(ori_model(**inputs).last_hidden_state.argmax(-1).item())
    label = str(ori_model.config.id2label[label_id])
    digits = "".join(char for char in label if char.isdigit())
    return int(digits or "0") % 360


def _upright(image: Image.Image) -> Image.Image:
    image = image.convert("RGB")
    skew = _deskew_angle(image)
    leveled = _rotate(image, skew)
    quadrant = _orientation_degrees(leveled)
    return _rotate(leveled, float(quadrant))


def _crop(image: Image.Image, box):
    width, height = image.size
    x1 = max(0, int(box[0]))
    y1 = max(0, int(box[1]))
    x2 = min(width, int(box[2]) + 1)
    y2 = min(height, int(box[3]) + 1)
    if x2 - x1 < 2 or y2 - y1 < 2:
        return None
    return image.crop((x1, y1, x2, y2))


def _strip_fence(text: str) -> str:
    lines = text.strip().splitlines()
    if len(lines) >= 2 and lines[0].startswith("```") and lines[-1].startswith("```"):
        return "\n".join(lines[1:-1]).strip()
    return text.strip()


def _recognize(image: Image.Image, prompt: str, max_new_tokens: int) -> str:
    messages = [
        {
            "role": "user",
            "content": [
                {"type": "image", "image": image},
                {"type": "text", "text": prompt},
            ],
        }
    ]
    inputs = ocr_processor.apply_chat_template(
        messages,
        add_generation_prompt=True,
        tokenize=True,
        return_dict=True,
        return_tensors="pt",
        enable_thinking=False,
        processor_kwargs={"max_slice_nums": 9, "downsample_mode": "4x"},
    ).to("cuda", dtype=torch.bfloat16)
    with torch.inference_mode():
        output = ocr_model.generate(
            **inputs,
            max_new_tokens=max_new_tokens,
            do_sample=False,
            repetition_penalty=1.0,
            no_repeat_ngram_size=100,
        )
    prompt_len = int(inputs["input_ids"].shape[1])
    return _strip_fence(ocr_processor.decode(output[0, prompt_len:], skip_special_tokens=True))


def _plain(text: str) -> str:
    return text.replace("#", "").replace("*", "").strip()


def _format_block(label: str, text: str) -> str:
    return format_block(label, text)


def _detect(image: Image.Image):
    return detect_layout(layout_model, image)


def _progress_view(step: int, total: int | None) -> str:
    if not total or step <= 0:
        percent = 0
        label = "0%"
    else:
        percent = min(100, round(100 * step / total))
        label = f"{percent}%"
    return (
        "<div style=\"margin:4px 0 12px\">"
        f"<div style=\"font-size:14px;margin-bottom:6px\">{label}</div>"
        "<div style=\"height:10px;background:#e5e7eb;border-radius:999px;overflow:hidden\">"
        f"<div style=\"width:{percent}%;height:100%;background:#2563eb\"></div>"
        "</div></div>"
    )


def _jobs(image: Image.Image, recognize_seal: bool, recognize_chart: bool):
    skipped = set(SKIP_VLM)
    if not recognize_seal:
        skipped.add("seal")
    if not recognize_chart:
        skipped.add("chart")
    pieces = []
    jobs = []
    for label, box in _detect(image):
        if label in skipped:
            if label == "image":
                pieces.append(_format_block(label, ""))
            continue
        crop = _crop(image, box)
        if crop is None:
            continue
        jobs.append((label, crop))
    if not jobs and not pieces:
        jobs.append(("text", image))
    return pieces, jobs


@spaces.GPU(duration=_duration)
def parse(image, upright, recognize_seal, recognize_chart, language):
    text = _copy(language)
    page = image.convert("RGB")
    if upright:
        page = _upright(page)
    pieces, jobs = _jobs(page, bool(recognize_seal), bool(recognize_chart))
    total = len(jobs) + 1
    yield _progress_view(1, total), text["parsing"]
    for index, (label, crop) in enumerate(jobs, start=1):
        prompt, limit = PROMPTS.get(label, ("Text Recognition:", 1024))
        recognized = _recognize(crop, prompt, limit)
        if recognized and label not in SKIP_MARKDOWN:
            formatted = _format_block(label, recognized)
            if not pieces or _plain(formatted) != _plain(pieces[-1]):
                pieces.append(formatted)
        yield _progress_view(1 + index, total), text["parsing"]
    yield _progress_view(total, total), "\n\n".join(pieces) or text["empty"]


def run(image, upright, recognize_seal, recognize_chart, language):
    text = _copy(language)
    if image is None:
        raise gr.Error(text["need_image"])
    notice = text["parsing"]
    yield _progress_view(0, None), notice, notice
    for progress, body in parse(image, bool(upright), bool(recognize_seal), bool(recognize_chart), language):
        yield progress, body, body


if os.environ.get("OCR_SKIP_LAUNCH") != "1":
    import gradio as gr

    english = COPY["en"]

    def apply_language(language, view):
        text = _copy(language)
        return (
            gr.update(label=text["language_label"]),
            gr.update(value=text["intro"]),
            gr.update(value=_note(language)),
            gr.update(label=text["image"], placeholder=text["upload"]),
            gr.update(label=text["upright"], info=text["upright_info"]),
            gr.update(label=text["seal"], info=text["seal_info"]),
            gr.update(label=text["chart"], info=text["chart_info"]),
            gr.update(value=text["parse"]),
            gr.update(
                label=text["view_label"],
                choices=[(text["rendered"], "rendered"), (text["source"], "source")],
                value=view,
            ),
        )

    def apply_view(view):
        rendered = view != "source"
        return gr.update(visible=rendered), gr.update(visible=not rendered)

    with gr.Blocks(title="MiniCPM-OCR") as demo:
        intro = gr.Markdown(english["intro"])
        note = gr.HTML(_note("en"))
        with gr.Row(elem_classes=["ocr-columns"]):
            with gr.Column(elem_classes=["ocr-col"], scale=1, min_width=0):
                language = gr.Radio(
                    choices=[("English", "en"), ("中文", "zh")],
                    value="en",
                    label=english["language_label"],
                )
                image = gr.Image(type="pil", label=english["image"], placeholder=english["upload"])
                with gr.Group():
                    upright = gr.Checkbox(label=english["upright"], info=english["upright_info"], value=True)
                    recognize_seal = gr.Checkbox(label=english["seal"], info=english["seal_info"], value=True)
                    recognize_chart = gr.Checkbox(label=english["chart"], info=english["chart_info"], value=True)
                button = gr.Button(english["parse"], variant="primary")
            with gr.Column(elem_classes=["ocr-col"], scale=1, min_width=0):
                view = gr.Radio(
                    choices=[(english["rendered"], "rendered"), (english["source"], "source")],
                    value="rendered",
                    label=english["view_label"],
                )
                progress = gr.HTML()
                markdown = gr.Markdown(
                    elem_classes=["ocr-result"],
                    sanitize_html=False,
                    latex_delimiters=[
                        {"left": "$$", "right": "$$", "display": True},
                        {"left": "$", "right": "$", "display": False},
                    ],
                    max_height=640,
                )
                source = gr.Code(language="markdown", visible=False, lines=24, show_label=False)
        language.change(
            apply_language,
            inputs=[language, view],
            outputs=[language, intro, note, image, upright, recognize_seal, recognize_chart, button, view],
        )
        view.change(apply_view, inputs=[view], outputs=[markdown, source])
        button.click(
            run,
            inputs=[image, upright, recognize_seal, recognize_chart, language],
            outputs=[progress, markdown, source],
            show_progress="full",
        )

    demo.launch(css=RESULT_CSS)