import numpy as np
import gradio as gr
import segment_anything
import base64
import torch
import typing
import os
import subprocess
import requests
import PIL.Image
import urllib.parse


def download_image(url) -> PIL.Image.Image:
    """Download an image from a URL and return it as a PIL image."""
    return PIL.Image.open(requests.get(url, stream=True).raw)


def image_to_sam_image_embedding(
    image_url: str,
    # model_size: typing.Literal["base", "large", "huge"] = "base",
    model_size: str = "base",
) -> str:
    """Generate an image embedding."""

    image_url = urllib.parse.unquote(image_url)

    try:
        image = download_image(image_url)
    except:
        raise gr.Error(f"Could not find image with URL: {image_url}.")
    image = image.convert("RGB")
    image = np.asarray(image)

    # Select model size
    if model_size == "base":
        predictor = base_predictor
    elif model_size == "large":
        predictor = large_predictor
    elif model_size == "huge":
        predictor = huge_predictor

    # Run model
    predictor.set_image(image)
    # Output shape is (1, 256, 64, 64)
    image_embedding = predictor.get_image_embedding().cpu().numpy()

    # Flatten the array to a 1D array
    flat_arr = image_embedding.flatten()
    # Convert the 1D array to bytes
    bytes_arr = flat_arr.astype(np.float32).tobytes()
    # Encode the bytes to base64
    base64_str = base64.b64encode(bytes_arr).decode("utf-8")

    return base64_str


if __name__ == "__main__":

    # Load the model into memory to make running multiple predictions efficient
    device = "cuda" if torch.cuda.is_available() else "cpu"

    base_sam_checkpoint = "sam_vit_b_01ec64.pth"  # 375 MB
    large_sam_checkpoint = "sam_vit_l_0b3195.pth"  # 1.25 GB
    huge_sam_checkpoint = "sam_vit_h_4b8939.pth"  # 2.56 GB

    # Download the model checkpoints
    for model in [base_sam_checkpoint, large_sam_checkpoint, huge_sam_checkpoint]:
        if not os.path.exists(f"./{model}"):
            result = subprocess.run(
                ["wget", f"https://dl.fbaipublicfiles.com/segment_anything/{model}"],
                check=True,
            )
            print(f"wget {model} result = {result}")

    base_sam = segment_anything.sam_model_registry["vit_b"](
        checkpoint=base_sam_checkpoint
    )
    large_sam = segment_anything.sam_model_registry["vit_l"](
        checkpoint=large_sam_checkpoint
    )
    huge_sam = segment_anything.sam_model_registry["vit_h"](
        checkpoint=huge_sam_checkpoint
    )

    base_sam.to(device=device)
    large_sam.to(device=device)
    huge_sam.to(device=device)

    base_predictor = segment_anything.SamPredictor(base_sam)
    large_predictor = segment_anything.SamPredictor(large_sam)
    huge_predictor = segment_anything.SamPredictor(huge_sam)

    # Gradio app
    app = gr.Interface(
        fn=image_to_sam_image_embedding,
        inputs=[
            gr.components.Textbox(label="Image URL"),
            gr.components.Radio(
                choices=["base", "large", "huge"], label="Model Size", value="base"
            ),
        ],
        outputs="text",
    )
    app.launch()