import gradio as gr
import torch
import spaces
import json
import base64
from io import BytesIO
from transformers import SamHQModel, SamHQProcessor, SamModel, SamProcessor
import os
import pandas as pd
from utils import *
from PIL import Image

# Carga de modelos
sam_hq_model = SamHQModel.from_pretrained("syscv-community/sam-hq-vit-huge")
sam_hq_processor = SamHQProcessor.from_pretrained("syscv-community/sam-hq-vit-huge")

sam_model = SamModel.from_pretrained("facebook/sam-vit-huge")
sam_processor = SamProcessor.from_pretrained("facebook/sam-vit-huge")

@spaces.GPU
def predict_masks_and_scores(model, processor, raw_image, input_points=None, input_boxes=None):
    if input_boxes is not None:
        input_boxes = [input_boxes]
    inputs = processor(raw_image, input_boxes=input_boxes, input_points=input_points, return_tensors="pt")
    with torch.no_grad():
        outputs = model(**inputs)

    masks = processor.image_processor.post_process_masks(
        outputs.pred_masks.cpu(), inputs["original_sizes"].cpu(), inputs["reshaped_input_sizes"].cpu()
    )
    scores = outputs.iou_scores
    return masks, scores

def encode_pil_to_base64(pil_image):
    buffer = BytesIO()
    pil_image.save(buffer, format="PNG")
    return base64.b64encode(buffer.getvalue()).decode("utf-8")

def compare_images_points_and_masks(user_image, input_boxes, input_points):
    for example_path, example_data in example_data_map.items():
        if example_data["size"] == list(user_image.size):
            user_image = Image.open(example_data['original_image_path'])
    input_boxes = input_boxes.values.tolist()
    input_points = input_points.values.tolist()

    input_boxes = [[[int(coord) for coord in box] for box in input_boxes if any(box)]]
    input_points = [[[int(coord) for coord in point] for point in input_points if any(point)]]

    input_boxes = input_boxes if input_boxes[0] else None
    input_points = input_points if input_points[0] else None

    sam_masks, sam_scores = predict_masks_and_scores(sam_model, sam_processor, user_image, input_boxes=input_boxes, input_points=input_points)
    sam_hq_masks, sam_hq_scores = predict_masks_and_scores(sam_hq_model, sam_hq_processor, user_image, input_boxes=input_boxes, input_points=input_points)

    if input_boxes and input_points:
        img1_b64 = show_all_annotations_on_image_base64(user_image, sam_masks[0][0], sam_scores[:, 0, :], input_boxes[0], input_points[0], model_name='SAM')
        img2_b64 = show_all_annotations_on_image_base64(user_image, sam_hq_masks[0][0], sam_hq_scores[:, 0, :], input_boxes[0], input_points[0], model_name='SAM_HQ')
    elif input_boxes:
        img1_b64 = show_all_annotations_on_image_base64(user_image, sam_masks[0][0], sam_scores[:, 0, :], input_boxes[0], None, model_name='SAM')
        img2_b64 = show_all_annotations_on_image_base64(user_image, sam_hq_masks[0][0], sam_hq_scores[:, 0, :], input_boxes[0], None, model_name='SAM_HQ')
    elif input_points:
        img1_b64 = show_all_annotations_on_image_base64(user_image, sam_masks[0][0], sam_scores[:, 0, :], None, input_points[0], model_name='SAM')
        img2_b64 = show_all_annotations_on_image_base64(user_image, sam_hq_masks[0][0], sam_hq_scores[:, 0, :], None, input_points[0], model_name='SAM_HQ')

    print('user_image', user_image)
    print("img1_b64", img1_b64)
    print("img2_b64", img2_b64)

    html_code = f"""
    <div style="position: relative; width: 100%; max-width: 600px; margin: 0 auto;" id="imageCompareContainer">
        <div style="position: relative; width: 100%;">
            <img src="data:image/png;base64,{img1_b64}" style="width:100%; display:block;">
            <div id="topWrapper" style="position:absolute; top:0; left:0; width:100%; overflow:hidden;">
                <img id="topImage" src="data:image/png;base64,{img2_b64}" style="width:100%;">
            </div>
            <div id="sliderLine" style="position:absolute; top:0; left:0; width:2px; height:100%; background-color:red; pointer-events:none;"></div>
        </div>
        <input type="range" min="0" max="100" value="0" 
            style="width:100%; margin-top: 10px;" 
            oninput="
            const val = this.value;
            const container = document.getElementById('imageCompareContainer');
            const width = container.offsetWidth;
            const clipValue = 100 - val;
            document.getElementById('topImage').style.clipPath = 'inset(0 ' + clipValue + '% 0 0)';
            document.getElementById('sliderLine').style.left = (width * val / 100) + 'px';
            ">
    </div>
    """
    return html_code

def load_examples(json_file="examples.json"):
    with open(json_file, "r") as f:
        examples = json.load(f)
    return examples

examples = load_examples()
example_paths = [example["image_path"] for example in examples]
example_data_map = {
    example["image_path"]: {
        "original_image_path": example["original_image_path"],
        "points": example["points"],
        "boxes": example["boxes"],
        "size": example["size"]
    }
    for example in examples
}

theme = gr.themes.Soft(primary_hue="indigo", secondary_hue="emerald")
with gr.Blocks(theme=theme, title="🔍 Compare SAM vs SAM-HQ") as demo:
    image_path_box = gr.Textbox(visible=False)
    gr.Markdown("## 🔍 Compare SAM vs SAM-HQ")
    gr.Markdown("Compare the performance of SAM and SAM-HQ on various images. Click on an example to load it")
    gr.Markdown("[SAM-HQ](https://huggingface.co/syscv-community/sam-hq-vit-huge) - [SAM](https://huggingface.co/facebook/sam-vit-huge)")

    with gr.Row():
        image_input = gr.Image(
            type="pil",
            label="Example image (click below to load)",
            interactive=False,
            height=500,
            show_label=True
        )
    
    gr.Examples(
        examples=example_paths,
        inputs=[image_input],
        label="Click an example to try 👇",
    )

    result_html = gr.HTML(elem_id="result-html")

    with gr.Row():
        points_input = gr.Dataframe(
            headers=["x", "y"],
            label="Points",
            datatype=["number", "number"],
            col_count=(2, "fixed")
        )
        boxes_input = gr.Dataframe(
            headers=["x0", "y0", "x1", "y1"],
            label="Boxes",
            datatype=["number", "number", "number", "number"],
            col_count=(4, "fixed")
        )

    def on_image_change(image):
        for example_path, example_data in example_data_map.items():
            print(image.size)
            if example_data["size"] == list(image.size):
                return example_data["points"], example_data["boxes"]
        return [], []

    image_input.change(
        fn=on_image_change,
        inputs=[image_input],
        outputs=[points_input, boxes_input]
    )

    compare_button = gr.Button("Compare points and masks")
    compare_button.click(fn=compare_images_points_and_masks, inputs=[image_input, boxes_input, points_input], outputs=result_html)

    gr.HTML("""
        <style>
        #result-html {
            min-height: 500px;
            border: 1px solid #ccc;
            padding: 10px;
            box-sizing: border-box;
            background-color: #fff;
            border-radius: 8px;
            box-shadow: 0 2px 6px rgba(0, 0, 0, 0.1);
        }
        </style>
    """)

demo.launch()