import pickle
import gradio as gr
from datasets import load_dataset
from transformers import AutoModel, AutoFeatureExtractor
import wikipedia


# Only runs once when the script is first run.
with open("./insectarium_768.pickle", "rb") as handle:
    index = pickle.load(handle)

# Load model for computing embeddings.
feature_extractor = AutoFeatureExtractor.from_pretrained("sasha/vit-base-butterflies")
model = AutoModel.from_pretrained("sasha/vit-base-butterflies")

# Candidate images.
dataset = load_dataset("sasha/insectarium-butterflies")
ds = dataset["train"]


def query(image, top_k=4):
    inputs = feature_extractor(image, return_tensors="pt")
    model_output = model(**inputs)
    embedding = model_output.pooler_output.detach()
    results = index.query(embedding, k=top_k)
    inx = results[0][0].tolist()
    logits = results[1][0].tolist()
    images = ds.select(inx)["image"]
    captions = ds.select(inx)["label"]
    images_with_captions = [(i, ds.features["label"].int2str(c)) for i, c in zip(images, captions)]
    labels_with_probs = dict(zip(captions, logits))
    labels_with_probs = {k: 1 - v for k, v in labels_with_probs.items()}
    try:
        description = wikipedia.summary(captions[0], sentences=1)
        description = "### " + description
        url = wikipedia.page(captions[0]).url
        url = " You can learn more about your butterfly [here](" + str(url) + ")!"
        description = description + url
    except:
        description = "### Butterflies are insects in the order Lepidoptera, which also includes moths. Adult butterflies have large, often brightly coloured wings."
        url = "https://en.wikipedia.org/wiki/Butterfly"
        url = " You can learn more about butterflies [here](" + str(url) + ")!"
        description = description + url
    return images_with_captions, labels_with_probs


with gr.Blocks() as demo:
    gr.Markdown("# Find my Butterfly 🦋")
    gr.Markdown(
        "## Use this Space to find your butterfly, based on the [iNaturalist butterfly dataset](https://huggingface.co/datasets/huggan/inat_butterflies_top10k)!"
    )
    with gr.Row():
        with gr.Column(scale=1):
            inputs = gr.Image()
            btn = gr.Button("Find my butterfly!")
            description = gr.Markdown()

        with gr.Column(scale=2):
            outputs = gr.Gallery(rows=2)
            labels = gr.Label()

    gr.Markdown("### Image Examples")
    gr.Examples(
        examples=["elton.jpg", "ken.jpg", "gaga.jpg", "taylor.jpg"],
        inputs=inputs,
        outputs=[outputs, labels],
        fn=query,
        cache_examples=True,
    )
    btn.click(query, inputs, [outputs, labels, description])

demo.launch()