import streamlit as st
import torch
# import numpy as np
import cv2
import wget
import os

from PIL import Image
from streamlit_drawable_canvas import st_canvas

from isegm.inference import clicker as ck
from isegm.inference import utils
from isegm.inference.predictors import get_predictor

# Model Path
prefix = "https://huggingface.co/curt-park/interactive-segmentation/resolve/main"
models = {
    "RITM": "ritm_coco_lvis_h18_itermask.pth",
}

# Items in the sidebar.
model = st.sidebar.selectbox("Select a Model:", tuple(models.keys()))
threshold = st.sidebar.slider("Threshold: ", 0.0, 1.0, 0.5)
marking_type = st.sidebar.radio("Marking Type:", ("positive", "negative"))
image_path = st.sidebar.file_uploader("Background Image:", type=["png", "jpg", "jpeg"])

# Objects for prediction.
clicker = ck.Clicker()
device = torch.device("cpu")
predictor = None
with st.spinner("Wait for downloading a model..."):
    if not os.path.exists(models[model]):
        _ = wget.download(f"{prefix}/{models[model]}")

with st.spinner("Wait for loading a model..."):
    model = utils.load_is_model(models[model], device, cpu_dist_maps=True)
    predictor_params = {"brs_mode": "NoBRS"}
    predictor = get_predictor(model, device=device, **predictor_params)

# Create a canvas component.
#image = None
#if image_path:
#    image = Image.open(image_path)
#canvas_height, canvas_width = 600, 600
#pos_color, neg_color = "#3498DB", "#C70039"
#st.title("Canvas:")
#canvas_result = st_canvas(
#        fill_color="rgba(255, 165, 0, 0.3)",  # Fixed fill color with some opacity
#        stroke_width=3,
#        stroke_color=pos_color if marking_type == "positive" else neg_color,
#        background_color="#eee",
#        background_image=image,
#        update_streamlit=True,
#        drawing_mode="point",
#        point_display_radius=3,
#        key="canvas",
#        width=canvas_width,
#        height=canvas_height,
#)

# Check the user inputs ans execute predictions.
#st.title("Prediction:")
#if canvas_result.json_data and canvas_result.json_data["objects"] and image:
#    objects = canvas_result.json_data["objects"]
#    image_width, image_height = image.size
#    ratio_h, ratio_w = image_height / canvas_height, image_width / canvas_width
#
#    err_x, err_y = 5.5, 1.0
#    pos_clicks, neg_clicks = [], []
#    for click in objects:
#        x, y = (click["left"] + err_x) * ratio_w, (click["top"] + err_y) * ratio_h
#        x, y = min(image_width, max(0, x)), min(image_height, max(0, y))
#
#        is_positive = click["stroke"] == pos_color
#        click = ck.Click(is_positive=is_positive, coords=(y, x))
#        clicker.add_click(click)
#
#    # prediction.
#    pred = None
#    predictor.set_input_image(np.array(image))
#    with st.spinner("Wait for prediction..."):
#        pred = predictor.get_prediction(clicker, prev_mask=None)
#    pred = cv2.resize(pred, dsize=(canvas_height, canvas_width), interpolation=cv2.INTER_CUBIC)
#    pred = np.where(pred > threshold, 1.0, 0)
#    st.image(pred, caption="")