import gradio as gr
import cv2
import numpy as np
from PIL import Image

block = gr.Blocks(
            title="SAM and others",
            # theme="shivi/calm_seafoam@>=0.0.1,<1.0.0",
        )
colors = [(255, 0, 0), (0, 255, 0)]
markers = [1, 5]

def get_point(img, sel_pix, evt: gr.SelectData):
        img = np.array(img, dtype=np.uint8)
        sel_pix.append(evt.index)
        # draw points
        
        print(sel_pix)
        for point in sel_pix:
            cv2.drawMarker(img, point, colors[0], markerType=markers[0], markerSize=6, thickness=2)
        return Image.fromarray(img).convert("RGB")



def undo_button(orig_img, sel_pix):
    temp = orig_img.copy()
    temp = np.array(temp, dtype=np.uint8)
    if len(sel_pix) != 0:
        sel_pix.pop()
        for point in sel_pix:
            cv2.drawMarker(temp, point, colors[0], markerType=markers[0], markerSize=6, thickness=2)
    return Image.fromarray(temp).convert("RGB")

def toggle_button(orig_img, mode):
    print(mode)
    if mode:
        ret = gr.Image(value= orig_img,elem_id="image_upload", type='pil', label="Upload", height=512, tool = "editor")# tool = "sketch", brush_color='#00FFFF', mask_opacity=0.6)
    else:
        ret = gr.Image(value = orig_img, elem_id="image_upload", type='pil', label="Upload", height=512, tool = "sketch", brush_color='#00FFFF', mask_opacity=0.6)
    mode = not mode    
    return ret, mode
    
def store_img(img):
    print("call for store")
    return img, []  # when new image is uploaded, `selected_points` should be empty
    
with block:
    selected_points = gr.State([])
    original_image = gr.State()
    mode = gr.State(True)
    input_image = gr.Image(elem_id="image_upload", type='pil', label="Upload", height=512,)# tool = "sketch", brush_color='#00FFFF', mask_opacity=0.6)
    undo = gr.Button("undo mode", visible=True)
    toggle = gr.Button("toggle mode", visible=True)
    input_image.upload(
        store_img,
        [input_image],
        [original_image, selected_points]
    )

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

    undo.click(fn=undo_button, inputs=[original_image, selected_points], outputs=[input_image])
    toggle.click(fn=toggle_button, inputs=[original_image, mode], outputs=[input_image, mode])
    
block.launch()