import gradio as gr
import numpy as np
from transformers import AutoImageProcessor, AutoModel
from transformers.image_utils import to_numpy_array
import torch
import plotly.graph_objects as go
from PIL import Image
import spaces

@spaces.GPU
def process_images(image1, image2):
    """
    Process two images and return a plot of the matching keypoints.
    """
    if image1 is None or image2 is None:
        return None

    images = [image1, image2]
    processor = AutoImageProcessor.from_pretrained("ETH-CVG/lightglue_superpoint")
    model = AutoModel.from_pretrained("ETH-CVG/lightglue_superpoint")

    inputs = processor(images, return_tensors="pt")
    with torch.no_grad():
        outputs = model(**inputs)

    image_sizes = [[(image.height, image.width) for image in images]]
    outputs = processor.post_process_keypoint_matching(
        outputs, image_sizes, threshold=0.2
    )
    output = outputs[0]

    image1 = to_numpy_array(image1)
    image2 = to_numpy_array(image2)

    height0, width0 = image1.shape[:2]
    height1, width1 = image2.shape[:2]

    # Create PIL image from numpy array
    pil_img = Image.fromarray((image1 / 255.0 * 255).astype(np.uint8))
    pil_img2 = Image.fromarray((image2 / 255.0 * 255).astype(np.uint8))

    # Create Plotly figure
    fig = go.Figure()

    # Get keypoints
    keypoints0_x, keypoints0_y = output["keypoints0"].unbind(1)
    keypoints1_x, keypoints1_y = output["keypoints1"].unbind(1)

    # Add a separate trace for each match (line + markers) to enable highlighting
    for keypoint0_x, keypoint0_y, keypoint1_x, keypoint1_y, matching_score in zip(
        keypoints0_x,
        keypoints0_y,
        keypoints1_x,
        keypoints1_y,
        output["matching_scores"],
    ):
        color_val = matching_score.item()
        color = f"rgba({int(255 * (1 - color_val))}, {int(255 * color_val)}, 0, 0.7)"

        hover_text = (
            f"Score: {matching_score.item():.2f}<br>"
            f"Point 1: ({keypoint0_x.item():.1f}, {keypoint0_y.item():.1f})<br>"
            f"Point 2: ({keypoint1_x.item():.1f}, {keypoint1_y.item():.1f})"
        )

        fig.add_trace(
            go.Scatter(
                x=[keypoint0_x.item(), keypoint1_x.item() + width0],
                y=[keypoint0_y.item(), keypoint1_y.item()],
                mode="lines+markers",
                line=dict(color=color, width=2),
                marker=dict(color=color, size=5, opacity=0.8),
                hoverinfo="text",
                hovertext=hover_text,
                showlegend=False,
            )
        )

    # Update layout to use images as background
    fig.update_layout(
        title="LightGlue Keypoint Matching",
        xaxis=dict(
            range=[0, width0 + width1],
            showgrid=False,
            zeroline=False,
            showticklabels=False,
        ),
        yaxis=dict(
            range=[max(height0, height1), 0],
            showgrid=False,
            zeroline=False,
            showticklabels=False,
            scaleanchor="x",
            scaleratio=1,
        ),
        margin=dict(l=0, r=0, t=50, b=0),
        height=max(height0, height1),
        width=width0 + width1,
        images=[
            dict(
                source=pil_img,
                xref="x",
                yref="y",
                x=0,
                y=0,
                sizex=width0,
                sizey=height0,
                sizing="stretch",
                opacity=1,
                layer="below",
            ),
            dict(
                source=pil_img2,
                xref="x",
                yref="y",
                x=width0,
                y=0,
                sizex=width1,
                sizey=height1,
                sizing="stretch",
                opacity=1,
                layer="below",
            ),
        ],
    )

    return fig


# Create the Gradio interface
with gr.Blocks(title="LightGlue Matching Demo") as demo:
    gr.Markdown("# LightGlue Matching Demo")
    gr.Markdown(
        "Upload two images and get a side-by-side matching of your images using LightGlue."
    )
    gr.Markdown("""
        ## How to use:
        1. Upload two images using the file uploaders above
        2. Click the 'Match Images' button
        3. View the matched output image below
        
        The app will create a side-by-side matching of your images using LightGlue.
        You can also select an example image pair from the dataset.
    """)

    with gr.Row():
        # Input images on the same row
        image1 = gr.Image(label="First Image", type="pil")
        image2 = gr.Image(label="Second Image", type="pil")

    # Process button
    process_btn = gr.Button("Match Images", variant="primary")

    # Output plot
    output_plot = gr.Plot(label="Matching Results")

    # Connect the function
    process_btn.click(fn=process_images, inputs=[image1, image2], outputs=output_plot)

    # Add some example usage

    examples = gr.Dataset(
        components=[image1, image2],
        label="Example Image Pairs",
        samples=[
            [
                "https://raw.githubusercontent.com/magicleap/SuperGluePretrainedNetwork/refs/heads/master/assets/phototourism_sample_images/united_states_capitol_98169888_3347710852.jpg",
                "https://raw.githubusercontent.com/magicleap/SuperGluePretrainedNetwork/refs/heads/master/assets/phototourism_sample_images/united_states_capitol_26757027_6717084061.jpg",
            ],
            [
                "https://raw.githubusercontent.com/cvg/LightGlue/refs/heads/main/assets/DSC_0410.JPG",
                "https://raw.githubusercontent.com/cvg/LightGlue/refs/heads/main/assets/DSC_0411.JPG",
            ],
            [
                "https://raw.githubusercontent.com/cvg/LightGlue/refs/heads/main/assets/sacre_coeur1.jpg",
                "https://raw.githubusercontent.com/cvg/LightGlue/refs/heads/main/assets/sacre_coeur2.jpg",
            ],
            [
                "https://raw.githubusercontent.com/magicleap/SuperGluePretrainedNetwork/refs/heads/master/assets/phototourism_sample_images/piazza_san_marco_06795901_3725050516.jpg",
                "https://raw.githubusercontent.com/magicleap/SuperGluePretrainedNetwork/refs/heads/master/assets/phototourism_sample_images/piazza_san_marco_58751010_4849458397.jpg",
            ],
            [
                "https://raw.githubusercontent.com/magicleap/SuperGluePretrainedNetwork/refs/heads/master/assets/phototourism_sample_images/london_bridge_19481797_2295892421.jpg",
                "https://raw.githubusercontent.com/magicleap/SuperGluePretrainedNetwork/refs/heads/master/assets/phototourism_sample_images/london_bridge_78916675_4568141288.jpg",
            ],
        ],
    )

    examples.select(lambda x: (x[0], x[1]), [examples], [image1, image2])

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