import os 
import gc
import cv2
import torch
import traceback
import numpy as np
import gradio as gr
from itertools import chain
from huggingface_hub import hf_hub_download
from segment_anything import SamPredictor, sam_model_registry


#hf_hub_download(repo_id="vmoras/sam_api", filename="sam_vit_h.pth", token=os.environ.get('model_token'), local_dir="./") 

sam_checkpoint = "sam_vit_h_0.pth"
model_type = "vit_h"
device = 'cuda' if torch.cuda.is_available() else 'cpu'


def set_predictor(image):
    """
    Creates a Sam predictor object based on a given image and model.
    """
    sam = sam_model_registry[model_type](checkpoint=sam_checkpoint)
    sam.to(device=device)

    predictor = SamPredictor(sam)
    predictor.set_image(image)

    if device == 'cuda':
        gc.collect()
        torch.cuda.empty_cache()

    return [image, predictor, 'Done']


def get_polygon(points, image, predictor):
    """
    Returns the points of the polygon given a bounding box and a prediction
    made by Sam.
    """
    points = list(chain.from_iterable(points))

    input_box = np.array(points)

    masks, _, _ = predictor.predict(
        box=input_box[None, :],
        multimask_output=False,
    )

    img = masks[0].astype(np.uint8)
    contours, hierarchy = cv2.findContours(img, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)

    if len(contours) == 0:
        return [], img

    points = contours[0]

    polygon = []
    for point in points:
        for x, y in point:
            polygon.append([int(x), int(y)])

    mask = np.zeros(image.shape, dtype='uint8')
    poly = np.array(polygon)
    cv2.fillPoly(mask, [poly], (0, 255, 0))

    return polygon, mask


def add_bbox(bbox, evt: gr.SelectData):
    if bbox[0] == [0, 0]:
        bbox[0] = [evt.index[0], evt.index[1]]
        return bbox, bbox

    bbox[1] = [evt.index[0], evt.index[1]]
    return bbox, bbox


def clear_bbox(bbox):
    updated_bbox = [[0, 0], [0, 0]]
    return updated_bbox, updated_bbox


with gr.Blocks() as demo:
    gr.Markdown(
    """
    # Instructions
    1. Upload the image and press 'Send Image'.
    2. Wait until the word 'Done' appears on the 'Status' box.
    3. Click on the image where the upper left corner of the bbox should be.
    4. Click on the image where the lower right corner of the bbox should be.
    5. Check the coordinates using the 'bbox' box.
    6. Click on 'Send bounding box'.
    7. On the right side you will see the binary mask '\*'.
    8. On the lower side you will see the points that made up the polygon '\*'.
    9. Click on 'Clear bbox' to send another bounding box and repeat the steps from the thrid step.
    10. Repeat steps 3 to 9 until all the segments for this image are done.
    11. Click on the right corner of the image to remove it and repeat all the steps with the next 
    image.

    '\*' If the binary mask is all black and the polygon is an empty list, it means the program did 
    not find any segment in the bbox. Make the bbox a little big bigger if that happens.
    """)


    image = gr.State()
    embedding = gr.State()
    bbox = gr.State([[0, 0], [0, 0]])

    with gr.Row():
        input_image = gr.Image(label='Image')
        mask = gr.Image(label='Mask')

    with gr.Row():
        with gr.Column():
            output_status = gr.Textbox(label='Status')
            
        with gr.Column():            
            predictor_button = gr.Button('Send Image')
 
    with gr.Row():
        with gr.Column():
            bbox_box = gr.Textbox(label="bbox")

        with gr.Column():
            bbox_button = gr.Button('Clear bbox')

    with gr.Row():
        with gr.Column():
            polygon = gr.Textbox(label='Polygon')

        with gr.Column():
            points_button = gr.Button('Send bounding box')


    predictor_button.click(
        set_predictor, 
        input_image,
        [image, embedding, output_status],
    )

    points_button.click(
        get_polygon, 
        [bbox, image, embedding],
        [polygon, mask],
    )

    bbox_button.click(
        clear_bbox, 
        bbox,
        [bbox, bbox_box],
    )    

    input_image.select(
        add_bbox,
        bbox,
        [bbox, bbox_box]
    )


demo.launch(debug=True)