import gradio as gr
import cv2
import numpy as np
from pathlib import Path

# Import the existing homography and warping function from your existing script
from billboard import apply_homography_and_warp


# Function to select points on the image 
def get_point_interface(original_frame, points, evt: gr.SelectData):
    x, y = evt.index
    if points is None:
        points = []
    points = points.copy()  # Make a copy to avoid modifying in-place
    points.append((x, y))

    # Draw the points and lines on the image
    roi_image = original_frame.copy()
    for pt in points:
        cv2.circle(roi_image, pt, 5, (255, 0, 0), -1)
    if len(points) > 1:
        for i in range(len(points) - 1):
            cv2.line(roi_image, points[i], points[i + 1], (255, 0, 0), 2)

    return roi_image, points


# Function to process and apply homography (this will call your existing function)
def process_images(source_image, dest_image, roi_points):
    # Ensure that exactly four points are selected
    if len(roi_points) != 4:
        return "Error: Please select exactly four points."

    # Convert images to correct format (if not already in uint8)
    img_src = np.array(source_image)
    img_dst = np.array(dest_image)

    # Apply homography and get the warped image
    result_img = apply_homography_and_warp(img_src, img_dst, np.array(roi_points))

    # Return the processed image
    return result_img


# Gradio interface setup
with gr.Blocks(title="Homography Warping App", theme=gr.themes.Soft()) as demo:
    
    gr.HTML(
        """
    <h1 style='text-align: center'>
    Homography Warping App
    </h1>
    """
    )
    gr.HTML(
        """
        <h3 style='text-align: center'>
        <a href='https://opencv.org/university/' target='_blank'>OpenCV Courses</a> | <a href='https://github.com/OpenCV-University' target='_blank'>Github</a>
        </h3>
        """
    )
    gr.Markdown(
        "Upload two images (source and destination), select the four points for ROI on the destination image, and click 'Process Images' to warp the source image onto the destination."
    )
    
    with gr.Row():
        with gr.Column(scale=1):
            with gr.Column():
                source_image_input = gr.Image(label="Upload Source Image")
                dest_image_input = gr.Image(label="Upload Destination Image")
        with gr.Column(scale=1):
            dest_image_copy = gr.Image(label="Select ROI points on destination image")
            original_frame_state = gr.State(None)
            selected_points = gr.State([])
            clear_points_button = gr.Button("Clear Points")
            process_button = gr.Button("Process Images")
        with gr.Column(scale=1):
            output_image = gr.Image(label="Wrapped Image")

    examples = [
        ["./images/Apollo-8-Launch.png", "./images/times_square.jpg"],
    ]

    with gr.Row():
        gr.Examples(
            examples=examples,
            inputs=[source_image_input, dest_image_input],
            label="Load Example Images",
        )

    dest_image_input.upload(
        lambda img:img, inputs=dest_image_input, outputs=dest_image_copy  # Simply return the uploaded image
    )

    # Adding the ROI point selection functionality
    dest_image_copy.select(get_point_interface, inputs=[dest_image_copy, selected_points], outputs=[dest_image_copy, selected_points])

    # Clear the selected points
    clear_points_button.click(
        fn=lambda original_frame: (original_frame, []),
        inputs=dest_image_input,
        outputs=[dest_image_input, selected_points],
    )

    # Callback for processing the images
    process_button.click(
        process_images, inputs=[source_image_input, dest_image_input, selected_points], outputs=output_image
    )


# Launch the Gradio app
demo.launch()