import cv2
import numpy as np
import gradio as gr

from pydantic import BaseModel


# Points color and marker
color = (0, 255, 0)  # Red color for all points
marker_type = 1  # Cross marker

with gr.Blocks() as demo:
    with gr.Row():
        gr.Markdown('''# Annotate Points!🚀
        Upload an image and click to annotate points on it.
        ''')

    # Annotating points on an image
    with gr.Tab(label='Image'):
        with gr.Row():#.style(equal_height=True):
            with gr.Column():
                # Input image
                original_image = gr.State(value=None)  # store original image without points
                input_image = gr.Image(type="numpy", label="Upload Image")

                # Annotate points
                selected_points = gr.State([])  # store points
                with gr.Row():
                    gr.Markdown('Click on the image to select points.')
                    undo_button = gr.Button('Undo point')

            # Show the image with the annotated points
            with gr.Tab(label='Image+Points'):
                output_image = gr.Image(type='numpy')


        # Store the original image once uploaded
        def store_img(img):
            return img, []  # Reset selected points when a new image is uploaded


        input_image.upload(store_img, [input_image], [original_image, selected_points])


        # Get points when clicked on the image
        def get_point(img, sel_pix, evt: gr.SelectData):
            sel_pix.append(evt.index)  # Append the point's location (coordinates)

            # Draw points on the image
            for point in sel_pix:
                cv2.drawMarker(img, point, color, markerType=marker_type, markerSize=80, thickness=20)
            return img if isinstance(img, np.ndarray) else np.array(img)


        input_image.select(get_point, [input_image, selected_points], [input_image])


        # Undo the last selected point
        def undo_points(orig_img, sel_pix):
            temp = orig_img.copy()
            if len(sel_pix) != 0:
                sel_pix.pop()  # Remove the last point
                for point in sel_pix:
                    cv2.drawMarker(temp, point, color, markerType=marker_type, markerSize=20, thickness=5)
            return temp if isinstance(temp, np.ndarray) else np.array(temp)


        undo_button.click(undo_points, [original_image, selected_points], [input_image])

    # Launch the app
demo.queue().launch(inbrowser=True)