import spaces  # ZeroGPU: must be imported before torch

import time
import gradio as gr
import torch
from paddleocr import PaddleOCR
from transformers import (
    AutoModelForImageTextToText,
    AutoProcessor,
    Florence2ForConditionalGeneration,
)

GOT_ID = "stepfun-ai/GOT-OCR-2.0-hf"      # 0.6B, end-to-end OCR-2.0 VLM
FLORENCE_ID = "microsoft/Florence-2-large"  # 0.77B, unified vision model, <OCR> task

# ---------------------------------------------------------------------------
# Model loading (module level, on cuda — ZeroGPU optimizes startup placement)
# ---------------------------------------------------------------------------

print("Loading GOT-OCR-2.0-hf ...", flush=True)
got_model = AutoModelForImageTextToText.from_pretrained(
    GOT_ID, dtype=torch.bfloat16, low_cpu_mem_usage=True, device_map="cuda"
)
got_processor = AutoProcessor.from_pretrained(GOT_ID, use_fast=True)

print("Loading Florence-2-large ...", flush=True)
try:
    florence_model = Florence2ForConditionalGeneration.from_pretrained(
        FLORENCE_ID, dtype=torch.float16
    )
except Exception as e:  # native class unavailable / load issue -> community pattern
    print(f"Native Florence2 load failed ({e}); falling back to trust_remote_code", flush=True)
    from transformers import AutoModelForCausalLM

    florence_model = AutoModelForCausalLM.from_pretrained(
        FLORENCE_ID, dtype=torch.float16, trust_remote_code=True
    )
florence_model = florence_model.to("cuda")
florence_processor = AutoProcessor.from_pretrained(
    FLORENCE_ID, trust_remote_code=True
)

print("Initializing PaddleOCR PP-OCRv5 (runs on CPU) ...", flush=True)
paddle_ocr = PaddleOCR(
    text_detection_model_name="PP-OCRv5_server_det",
    text_recognition_model_name="PP-OCRv5_server_rec",
    use_doc_orientation_classify=False,
    use_doc_unwarping=False,
    use_textline_orientation=False,
    device="cpu",  # paddle leg stays on CPU; ZeroGPU quota is for the VLMs
)
print("All models ready.", flush=True)


# ---------------------------------------------------------------------------
# Inference
# ---------------------------------------------------------------------------

@spaces.GPU(duration=120)
def run_vlm_ocr(image: Image.Image):
    """Run GOT-OCR-2.0 and Florence-2 on the same image; returns (got_text, florence_text)."""
    # --- GOT-OCR-2.0: plain page OCR (processor injects the "OCR: " prompt) ---
    got_inputs = got_processor(image, return_tensors="pt").to("cuda")
    with torch.inference_mode():
        got_ids = got_model.generate(
            **got_inputs,
            do_sample=False,
            tokenizer=got_processor.tokenizer,  # required for stop_strings
            stop_strings="<|im_end|>",  # required, else generation runs to the cap
            max_new_tokens=4096,
        )
    got_text = got_processor.decode(
        got_ids[0, got_inputs["input_ids"].shape[1]:], skip_special_tokens=True
    )

    # --- Florence-2: <OCR> task ---
    task_prompt = "<OCR>"
    flo_inputs = florence_processor(
        text=task_prompt, images=image, return_tensors="pt"
    ).to("cuda", torch.float16)
    with torch.inference_mode():
        flo_ids = florence_model.generate(
            **flo_inputs, max_new_tokens=1024, num_beams=3, do_sample=False
        )
    flo_decoded = florence_processor.batch_decode(flo_ids, skip_special_tokens=False)[0]
    flo_parsed = florence_processor.post_process_generation(
        flo_decoded, task=task_prompt, image_size=(image.width, image.height)
    )
    florence_text = flo_parsed[task_prompt] or "(no text detected)"

    return got_text.strip() or "(no text detected)", florence_text


def run_paddle_ocr(image: Image.Image):
    """Run PaddleOCR PP-OCRv5 det+rec on CPU."""
    result = paddle_ocr.predict(image)
    lines = []
    for res in result:
        texts = res.get("rec_texts", []) if hasattr(res, "get") else res["rec_texts"]
        lines.extend(texts)
    return "\n".join(lines) if lines else "(no text detected)"


def compare(image: Image.Image):
    if image is None:
        raise gr.Error("Please upload or paste an image first.")

    t0 = time.perf_counter()
    paddle_text = run_paddle_ocr(image)
    paddle_time = time.perf_counter() - t0

    t1 = time.perf_counter()
    got_text, florence_text = run_vlm_ocr(image)
    vlm_time = time.perf_counter() - t1

    timings = (
        f"PaddleOCR PP-OCRv5: {paddle_time:.2f}s (CPU)  |  "
        f"GOT-OCR-2.0 + Florence-2: {vlm_time:.2f}s (GPU)"
    )
    return paddle_text, got_text, florence_text, timings


# ---------------------------------------------------------------------------
# UI
# ---------------------------------------------------------------------------

DESCRIPTION = """
# OCR Model Comparison — 3 models under 3B

Upload an image and run it through all three models at once:

| Model | Params | Type |
| --- | --- | --- |
| [PaddleOCR PP-OCRv5](https://huggingface.co/PaddlePaddle/PP-OCRv5_server_det) | < 0.1B | classic det+rec pipeline (CPU) |
| [GOT-OCR-2.0-hf](https://huggingface.co/stepfun-ai/GOT-OCR-2.0-hf) | 0.6B | end-to-end OCR VLM |
| [Florence-2-large](https://huggingface.co/microsoft/Florence-2-large) | 0.77B | unified vision model, `<OCR>` task |

Text is shown in reading order per model; lines may differ between models.
"""

with gr.Blocks(title="OCR Model Comparison") as demo:
    gr.Markdown(DESCRIPTION)

    with gr.Row():
        image_in = gr.Image(type="pil", label="Input image", height=420)

    run_btn = gr.Button("Run all three models", variant="primary")

    with gr.Row():
        paddle_out = gr.Textbox(label="PaddleOCR PP-OCRv5 (< 0.1B)", lines=14, show_copy_button=True)
        got_out = gr.Textbox(label="GOT-OCR-2.0-hf (0.6B)", lines=14, show_copy_button=True)
        florence_out = gr.Textbox(label="Florence-2-large (0.77B)", lines=14, show_copy_button=True)

    timings_out = gr.Textbox(label="Runtime", interactive=False)

    run_btn.click(
        fn=compare,
        inputs=image_in,
        outputs=[paddle_out, got_out, florence_out, timings_out],
    )

if __name__ == "__main__":
    demo.launch()