import torch
import gradio as gr

import os 
import spaces
from PIL import Image

from transformers import TextStreamer
from utils import title_markdown
from utils import block_css
from utils import tos_markdown
from utils import learn_more_markdown



textbox = gr.Textbox(
    show_label = False, placeholder = "Enter text and press ENTER", container = False
)

with gr.Blocks(title = '' ) as demo:
    gr.Markdown(title_markdown)


@spaces.GPU
with gr.Blocks(title='MoE-LLaVA🚀')) as demo:
    gr.Markdown(title_markdown)
    state = gr.State()
    state_ = gr.State()
    first_run = gr.State()
    images_tensor = gr.State()

    with gr.Row():
        with gr.Column(scale=3):
            image1 = gr.Image(label="Input Document", type="filepath")

            cur_dir = os.path.dirname(os.path.abspath(__file__))
            print(cur_dir)
            gr.Examples(
                examples=[
                    [
                        f"demo.jfif",
                        "What is unusual about this image?",
                    ],
                    [
                        f"demo.jfif",
                        "What are the things I should be cautious about when I visit here?",
                    ],
                    [
                        f"demo.jfif",
                        "If there are factual errors in the questions, point it out; if not, proceed answering the question. What’s happening in the desert?",
                    ],
                    [
                        f"demo.jfif",
                        "What is the title of this book?",
                    ],
                    [
                        f"demo.jfif",
                        "What type of food is the girl holding?",
                    ],
                    [
                        f"demo.jfif",
                        "What color is the train?",
                    ],
                    [
                        f"demo.jfif",
                        "What is the girl looking at?",
                    ],
                    [
                        f"demo.jfif",
                        "What might be the reason for the dog's aggressive behavior?",
                    ],
                ],
                inputs=[image1, textbox],
            )

        # with gr.Column(scale=7):
        #     #chatbot = gr.Chatbot(label="MoE-LLaVA", bubble_full_width=True).style(height=750)
        #     with gr.Row():
        #         with gr.Column(scale=8):
        #             textbox.render()
        #         with gr.Column(scale=1, min_width=50):
        #             submit_btn = gr.Button(
        #                 value="Send", variant="primary", interactive=True
        #             )
        #     with gr.Row(elem_id="buttons") as button_row:
        #         upvote_btn = gr.Button(value="👍  Upvote", interactive=True)
        #         downvote_btn = gr.Button(value="👎  Downvote", interactive=True)
        #         flag_btn = gr.Button(value="⚠️  Flag", interactive=True)
        #         # stop_btn = gr.Button(value="⏹️  Stop Generation", interactive=False)
        #         regenerate_btn = gr.Button(value="🔄  Regenerate", interactive=True)
        #         clear_btn = gr.Button(value="🗑️  Clear history", interactive=True)

#    gr.Markdown(tos_markdown)
#    gr.Markdown(learn_more_markdown)

    # submit_btn.click(generate, [image1, textbox, first_run, state, state_, images_tensor],
    #                  [state, state_, chatbot, first_run, textbox, images_tensor, image1])

    # regenerate_btn.click(regenerate, [state, state_], [state, state_, chatbot, first_run]).then(
    #     generate, [image1, textbox, first_run, state, state_, images_tensor],
    #     [state, state_, chatbot, first_run, textbox, images_tensor, image1])

    # clear_btn.click(clear_history, [state, state_],
    #                 [image1, textbox, first_run, state, state_, chatbot, images_tensor])

# app = gr.mount_gradio_app(app, demo, path="/")
demo.launch()