import os
os.environ["GRADIO_TEMP_DIR"] = "./tmp"

import sys
import torch
import gradio as gr
import numpy as np
from PIL import Image, ImageDraw, ImageFont
from transformers import (
    DFineForObjectDetection,
    RTDetrImageProcessor,
)

# == select device ==
device = 'cuda' if torch.cuda.is_available() else 'cpu'

# Available models
MODELS = {
    "Egret XLarge": "ds4sd/docling-layout-egret-xlarge",
    "Egret Large": "ds4sd/docling-layout-egret-large", 
    "Egret Medium": "ds4sd/docling-layout-egret-medium",
    "Heron 101": "ds4sd/docling-layout-heron-101",
    "Heron": "ds4sd/docling-layout-heron"
}

# Classes mapping for the docling model
classes_map = {
    0: "Caption",
    1: "Footnote", 
    2: "Formula",
    3: "List-item",
    4: "Page-footer",
    5: "Page-header",
    6: "Picture",
    7: "Section-header",
    8: "Table",
    9: "Text",
    10: "Title",
    11: "Document Index",
    12: "Code",
    13: "Checkbox-Selected",
    14: "Checkbox-Unselected",
    15: "Form",
    16: "Key-Value Region",
}

# Color mapping for visualization
colors = [
    "#FF6B6B", "#4ECDC4", "#45B7D1", "#96CEB4", "#FECA57",
    "#FF9FF3", "#54A0FF", "#5F27CD", "#00D2D3", "#FF9F43",
    "#10AC84", "#EE5A24", "#0ABDE3", "#006BA6", "#F79F1F",
    "#A3CB38", "#FDA7DF"
]

# Global variables for model
current_model = None
current_processor = None
current_model_name = None

def iomin(box1, box2):
    """
    Intersection over Minimum (IoMin)
    box1: Tensor[1, 4]
    box2: Tensor[N, 4]
    Returns: Tensor[N]
    """
    # Intersection
    x1 = torch.max(box1[:, 0], box2[:, 0])
    y1 = torch.max(box1[:, 1], box2[:, 1])
    x2 = torch.min(box1[:, 2], box2[:, 2])
    y2 = torch.min(box1[:, 3], box2[:, 3])
    inter_area = torch.clamp(x2 - x1, min=0) * torch.clamp(y2 - y1, min=0)

    # Areas
    box1_area = (box1[:, 2] - box1[:, 0]) * (box1[:, 3] - box1[:, 1])
    box2_area = (box2[:, 2] - box2[:, 0]) * (box2[:, 3] - box2[:, 1])
    min_area = torch.min(box1_area, box2_area)

    return inter_area / min_area

def nms(boxes, scores, iou_threshold=0.5):
    """
    Custom NMS implementation using IoMin
    """
    keep = []
    _, order = scores.sort(descending=True)

    while order.numel() > 0:
        i = order[0]
        keep.append(i.item())

        if order.numel() == 1:
            break

        box_i = boxes[i].unsqueeze(0)  # [1, 4]
        rest = order[1:]
        ious = iomin(box_i, boxes[rest])

        mask = (ious <= iou_threshold)
        order = order[1:][mask]

    return torch.tensor(keep, dtype=torch.long)

def load_model(model_name):
    """
    Load the selected model
    """
    global current_model, current_processor, current_model_name
    
    if current_model_name == model_name:
        return f"✅ Model {model_name} is already loaded!"
    
    try:
        print(f"Loading model: {model_name}")
        model_path = MODELS[model_name]
        
        processor = RTDetrImageProcessor.from_pretrained(model_path)
        model = DFineForObjectDetection.from_pretrained(model_path)
        model = model.to(device)
        model.eval()
        
        current_processor = processor
        current_model = model
        current_model_name = model_name
        
        return f"✅ Successfully loaded {model_name}!"
        
    except Exception as e:
        return f"❌ Error loading {model_name}: {str(e)}"

def visualize_bbox(image, boxes, labels, scores, classes_map, colors):
    """
    Visualize bounding boxes on image
    """
    if isinstance(image, np.ndarray):
        image = Image.fromarray(image)
    elif not isinstance(image, Image.Image):
        raise ValueError("Input image must be PIL Image or numpy array")
    
    # Create a copy to draw on
    draw_image = image.copy()
    draw = ImageDraw.Draw(draw_image)
    
    # Try to use a font, fallback to default if not available
    try:
        font = ImageFont.truetype("arial.ttf", 20)
    except:
        try:
            font = ImageFont.load_default()
        except:
            font = None
    
    for box, label_id, score in zip(boxes, labels, scores):
        # Convert tensor to int if needed
        if torch.is_tensor(label_id):
            label_id = label_id.item()
        if torch.is_tensor(score):
            score = score.item()
            
        label = classes_map.get(int(label_id), f"Class_{label_id}")
        color = colors[int(label_id) % len(colors)]
        
        # Convert box coordinates to integers
        x1, y1, x2, y2 = [int(coord) for coord in box]
        
        # Draw rectangle
        draw.rectangle([x1, y1, x2, y2], outline=color, width=3)
        
        # Draw label background
        text = f"{label}: {score:.2f}"
        if font:
            bbox = draw.textbbox((x1, y1), text, font=font)
            text_width = bbox[2] - bbox[0]
            text_height = bbox[3] - bbox[1]
        else:
            # Estimate text size if no font available
            text_width = len(text) * 10
            text_height = 20
            
        draw.rectangle([x1, y1-text_height-4, x1+text_width+4, y1], fill=color)
        draw.text((x1+2, y1-text_height-2), text, fill="white", font=font)
    
    return np.array(draw_image)

def recognize_image(input_img, conf_threshold, iou_threshold, nms_method):
    """
    Process image with docling layout model
    """
    if input_img is None:
        return None, "Please upload an image first."
        
    if current_model is None or current_processor is None:
        return None, "Please load a model first."
        
    try:
        # Ensure image is PIL Image
        if isinstance(input_img, np.ndarray):
            input_img = Image.fromarray(input_img)
        
        # Convert to RGB if needed
        if input_img.mode != 'RGB':
            input_img = input_img.convert('RGB')
        
        # Process image
        inputs = current_processor(images=[input_img], return_tensors="pt")
        inputs = {k: v.to(device) for k, v in inputs.items()}
        
        # Run inference
        with torch.no_grad():
            outputs = current_model(**inputs)
            
        # Post-process results
        results = current_processor.post_process_object_detection(
            outputs,
            target_sizes=torch.tensor([input_img.size[::-1]]),
            threshold=conf_threshold,
        )
        
        if not results or len(results) == 0:
            return np.array(input_img), "No detections found."
            
        result = results[0]
        
        # Get results
        boxes = result["boxes"]
        scores = result["scores"] 
        labels = result["labels"]
        
        if len(boxes) == 0:
            return np.array(input_img), "No detections above confidence threshold."
        
        # Apply NMS if requested
        if iou_threshold < 1.0:
            if nms_method == "Custom IoMin":
                # Use custom NMS with IoMin
                keep_indices = nms(
                    boxes=boxes,
                    scores=scores, 
                    iou_threshold=iou_threshold
                )
            else:
                # Use standard torchvision NMS
                keep_indices = torch.ops.torchvision.nms(
                    boxes=boxes,
                    scores=scores,
                    iou_threshold=iou_threshold
                )
            
            boxes = boxes[keep_indices]
            scores = scores[keep_indices]
            labels = labels[keep_indices]
        
        # Handle single detection case
        if len(boxes.shape) == 1:
            boxes = boxes.unsqueeze(0)
            scores = scores.unsqueeze(0)
            labels = labels.unsqueeze(0)
            
        # Visualize results
        output = visualize_bbox(
            input_img, 
            boxes, 
            labels, 
            scores, 
            classes_map,
            colors
        )
        
        detection_info = f"Found {len(boxes)} detections after NMS ({nms_method})"
        return output, detection_info
            
    except Exception as e:
        print(f"[ERROR] recognize_image failed: {e}")
        error_msg = f"Error during processing: {str(e)}"
        # Return original image on error
        if input_img is not None:
            return np.array(input_img), error_msg
        return np.zeros((512, 512, 3), dtype=np.uint8), error_msg

def gradio_reset():
    return gr.update(value=None), gr.update(value=None), gr.update(value="")

if __name__ == "__main__":
    print(f"Using device: {device}")
    
    # Create header HTML
    header_html = """
    <div style="text-align: center; margin-bottom: 20px;">
        <h1>🔍 Document Layout Analysis</h1>
        <p>Using Docling Layout Models for document structure detection</p>
        <p>Select a model, upload an image and adjust the parameters to detect document elements</p>
    </div>
    """
    
    with gr.Blocks(title="Document Layout Analysis", theme=gr.themes.Soft()) as demo:
        gr.HTML(header_html)
        
        with gr.Row():
            with gr.Column():
                # Model selection
                model_dropdown = gr.Dropdown(
                    choices=list(MODELS.keys()),
                    value="Egret XLarge",
                    label="🤖 Select Model",
                    info="Choose which Docling model to use"
                )
                
                load_btn = gr.Button("📥 Load Model", variant="secondary")
                model_status = gr.Textbox(
                    label="Model Status", 
                    interactive=False,
                    value="No model loaded"
                )
                
                input_img = gr.Image(
                    label="📄 Upload Document Image", 
                    interactive=True,
                    type="pil"
                )
                
                with gr.Row():
                    clear = gr.Button("🗑️ Clear")
                    predict = gr.Button("🔍 Detect Layout", interactive=True, variant="primary")
                    
                with gr.Row():
                    conf_threshold = gr.Slider(
                        label="Confidence Threshold",
                        minimum=0.0,
                        maximum=1.0,
                        step=0.05,
                        value=0.6,
                        info="Minimum confidence score for detections"
                    )
                    
                with gr.Row():
                    iou_threshold = gr.Slider(
                        label="NMS IoU Threshold", 
                        minimum=0.0,
                        maximum=1.0,
                        step=0.05,
                        value=0.5,
                        info="IoU threshold for Non-Maximum Suppression"
                    )
                
                nms_method = gr.Radio(
                    choices=["Custom IoMin", "Standard IoU"],
                    value="Custom IoMin",
                    label="NMS Method",
                    info="Choose NMS algorithm"
                )
                
                # Legend
                with gr.Accordion("📋 Detected Classes", open=False):
                    legend_html = "<div style='display: grid; grid-template-columns: repeat(2, 1fr); gap: 10px;'>"
                    for class_id, class_name in classes_map.items():
                        color = colors[class_id % len(colors)]
                        legend_html += f"""
                        <div style='display: flex; align-items: center; padding: 5px;'>
                            <div style='width: 20px; height: 20px; background-color: {color}; margin-right: 10px; border: 1px solid #ccc;'></div>
                            <span>{class_name}</span>
                        </div>
                        """
                    legend_html += "</div>"
                    gr.HTML(legend_html)
                    
            with gr.Column():
                gr.HTML("<h3>🎯 Detection Results</h3>")
                output_img = gr.Image(
                    label="Detected Layout", 
                    interactive=False,
                    type="numpy"
                )
                
                detection_info = gr.Textbox(
                    label="Detection Info",
                    interactive=False,
                    value=""
                )
        
        # Event handlers
        load_btn.click(
            load_model,
            inputs=[model_dropdown],
            outputs=[model_status]
        )
        
        clear.click(
            gradio_reset, 
            inputs=None, 
            outputs=[input_img, output_img, detection_info]
        )
        
        predict.click(
            recognize_image, 
            inputs=[input_img, conf_threshold, iou_threshold, nms_method], 
            outputs=[output_img, detection_info]
        )
    
    # Launch the demo
    demo.launch(
        server_name="0.0.0.0", 
        server_port=7860, 
        debug=True,
        share=False
    )