import cv2
import numpy as np
from PIL import Image
import gradio as gr
from huggingface_hub import snapshot_download
import onnxruntime as ort
import os
import time

MODEL_REPO = "fredcallagan/uvdoc-grid-onnx"

print("Downloading ONNX model files...")
model_dir = snapshot_download(
    repo_id=MODEL_REPO,
    local_dir="./uvdoc_model",
    local_dir_use_symlinks=False,
)
onnx_path = os.path.join(model_dir, "UVDoc_grid.onnx")
print(f"Model path: {onnx_path}")

print("Creating ONNX Runtime session...")
session = ort.InferenceSession(
    onnx_path, providers=['CPUExecutionProvider']
)
input_name = session.get_inputs()[0].name
print("Session ready.")


def _preprocess(image: Image.Image, target_size=(496, 720)):
    img_rgb = np.array(image)
    h_orig, w_orig = img_rgb.shape[:2]
    resized = cv2.resize(img_rgb, target_size)
    normalized = resized.astype(np.float32) / 255.0
    transposed = np.transpose(normalized, (2, 0, 1))
    batched = np.expand_dims(transposed, axis=0)
    return batched, h_orig, w_orig


def unwarp(image, interpolation="cubic", border_mode="replicate"):
    if image is None:
        return None, None, "No image provided"

    start_time = time.time()

    try:
        interp_map = {
            "nearest": cv2.INTER_NEAREST,
            "linear": cv2.INTER_LINEAR,
            "cubic": cv2.INTER_CUBIC,
            "lanczos": cv2.INTER_LANCZOS4,
        }
        border_map = {
            "replicate": cv2.BORDER_REPLICATE,
            "constant": cv2.BORDER_CONSTANT,
            "reflect": cv2.BORDER_REFLECT,
        }
        interp = interp_map.get(interpolation, cv2.INTER_CUBIC)
        border = border_map.get(border_mode, cv2.BORDER_REPLICATE)

        img_bgr = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
        h_orig, w_orig = img_bgr.shape[:2]

        input_tensor, _, _ = _preprocess(image)
        result = session.run(None, {input_name: input_tensor})[0]
        grid = np.transpose(result[0], (1, 2, 0))  # (45, 31, 2)

        grid_up = cv2.resize(grid, (w_orig, h_orig), interpolation=cv2.INTER_LINEAR)
        map_x = ((grid_up[..., 0] + 1) / 2) * (w_orig - 1)
        map_y = ((grid_up[..., 1] + 1) / 2) * (h_orig - 1)

        unwarped = cv2.remap(
            img_bgr,
            map_x.astype(np.float32),
            map_y.astype(np.float32),
            interpolation=interp,
            borderMode=border,
        )
        unwarped_rgb = cv2.cvtColor(unwarped, cv2.COLOR_BGR2RGB)

        # Draw corners on original image using the 45x31 grid
        corners_grid = [
            (0, 0),      # top-left
            (0, 30),     # top-right
            (44, 30),    # bottom-right
            (44, 0),     # bottom-left
        ]

        img_corners = img_bgr.copy()
        corner_colors = [(0, 255, 0), (0, 0, 255), (255, 0, 0), (255, 255, 0)]  # BGR

        pixel_corners = []
        for idx, (gy, gx) in enumerate(corners_grid):
            norm_x, norm_y = grid[gy, gx]
            px = int(((norm_x + 1) / 2) * (w_orig - 1))
            py = int(((norm_y + 1) / 2) * (h_orig - 1))
            pixel_corners.append((px, py))
            radius = max(8, min(w_orig, h_orig) // 60)
            cv2.circle(img_corners, (px, py), radius=radius, color=corner_colors[idx], thickness=-1)
            cv2.putText(img_corners, str(idx + 1), (px - 5, py + 5),
                        cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 2)

        # Connect corners with lines
        for i in range(len(pixel_corners)):
            pt1 = pixel_corners[i]
            pt2 = pixel_corners[(i + 1) % len(pixel_corners)]
            cv2.line(img_corners, pt1, pt2, (0, 255, 255), 3)

        corners_rgb = cv2.cvtColor(img_corners, cv2.COLOR_BGR2RGB)

        elapsed = time.time() - start_time
        timer_text = f"Processing completed in {elapsed:.2f} seconds"

        return Image.fromarray(unwarped_rgb), Image.fromarray(corners_rgb), timer_text

    except Exception as e:
        import traceback
        print(traceback.format_exc())
        raise gr.Error(f"Inference failed: {e}")


with gr.Blocks() as demo:
    gr.Markdown("# UVDoc Document Unwarping (High-Resolution)")
    gr.Markdown(
        "Upload a warped or curved document photo and UVDoc will geometrically "
        "rectify it at your **original image resolution**."
    )
    with gr.Row():
        with gr.Column():
            input_image = gr.Image(type="pil", label="Warped Document")
            interpolation = gr.Dropdown(
                choices=["nearest", "linear", "cubic", "lanczos"],
                value="cubic",
                label="Remap Interpolation",
            )
            border_mode = gr.Dropdown(
                choices=["replicate", "constant", "reflect"],
                value="replicate",
                label="Border Mode",
            )
            btn = gr.Button("Unwarp")
        with gr.Column():
            with gr.Row():
                output_image = gr.Image(type="pil", label="Rectified Document")
            with gr.Row():
                corners_image = gr.Image(type="pil", label="Detected Corners")
            with gr.Row():
                timer_text = gr.Textbox(label="Timer", interactive=False)
    btn.click(
        fn=unwarp,
        inputs=[input_image, interpolation, border_mode],
        outputs=[output_image, corners_image, timer_text],
    )

demo.launch()