import os

import gradio as gr
import numpy as np
import torch
from lavis.models import load_model_and_preprocess
from PIL import Image

device = torch.device("cuda") if torch.cuda.is_available() else "cpu"

model, vis_processors, _ = load_model_and_preprocess(
    name="blip2_opt", model_type="pretrain_opt2.7b", is_eval=True, device=device
)


def generate_caption(image, caption_type):
    image = vis_processors["eval"](image).unsqueeze(0).to(device)

    if caption_type == "Beam Search":
        caption = model.generate({"image": image})
    else:
        caption = model.generate(
            {"image": image}, use_nucleus_sampling=True, num_captions=3
        )

    caption = "\n".join(caption)

    if torch.cuda.is_available():
        torch.cuda.empty_cache()

    return caption


def chat(input_image, question, history):
    history = history or []
    question = question.lower()

    image = vis_processors["eval"](input_image).unsqueeze(0).to(device)

    clean = lambda x: x.replace("<p>", "").replace("</p>", "").replace("\n", "")
    clean_h = lambda x: (clean(x[0]), clean(x[1]))
    context = list(map(clean_h, history))
    template = "Question: {} Answer: {}."
    prompt = (
        " ".join(
            [template.format(context[i][0], context[i][1]) for i in range(len(context))]
        )
        + " Question: "
        + question
        + " Answer:"
    )

    response = model.generate({"image": image, "prompt": prompt})
    history.append((question, response[0]))

    return history, history


def clear_chat(history):
    return [], []


with gr.Blocks() as demo:
    gr.Markdown(
        "### BLIP-2: Bootstrapping Language-Image Pre-training with Frozen Image Encoders and Large Language Models"
    )
    gr.Markdown(
        "This demo uses the `pretrain_opt2.7b` weights. For more information please visit [Github](https://github.com/salesforce/LAVIS/tree/main/projects/blip2) or [Paper](https://arxiv.org/abs/2301.12597)."
    )

    with gr.Row():
        with gr.Column():
            input_image = gr.Image(label="Image", type="pil")
            caption_type = gr.Radio(
                ["Beam Search", "Nucleus Sampling"],
                label="Caption Decoding Strategy",
                value="Beam Search",
            )
            btn_caption = gr.Button("Generate Caption")
            output_text = gr.Textbox(label="Answer", lines=5)

        with gr.Column():
            chatbot = gr.Chatbot().style(color_map=("green", "pink"))
            chat_state = gr.State()

            question_txt = gr.Textbox(label="Question", lines=1)
            btn_answer = gr.Button("Generate Answer")
            btn_clear = gr.Button("Clear Chat")

    btn_caption.click(
        generate_caption, inputs=[input_image, caption_type], outputs=[output_text]
    )

    btn_answer.click(
        chat,
        inputs=[input_image, question_txt, chat_state],
        outputs=[chatbot, chat_state],
    )

    btn_clear.click(clear_chat, inputs=[chat_state], outputs=[chatbot, chat_state])

    gr.Examples(
        [
            ["./merlion.png", "Beam Search", "which city is this?"],
            [
                "./Blue_Jay_0044_62759.jpg",
                "Beam Search",
                "what is the name of this bird?",
            ],
            ["./5kstbz-0001.png", "Beam Search", "where is the man standing?"],
            [
                "ILSVRC2012_val_00000008.JPEG",
                "Beam Search",
                "Name the colors of macarons you see in the image.",
            ],
        ],
        inputs=[input_image, caption_type, question_txt],
    )

    gr.Markdown(
        "Sample images are taken from [ImageNet](https://paperswithcode.com/sota/image-classification-on-imagenet), [CUB](https://paperswithcode.com/dataset/cub-200-2011) and [GamePhysics](https://asgaardlab.github.io/CLIPxGamePhysics/) datasets."
    )

demo.launch()