import gradio as gr
import clip_chat


def logit2sentence(logit, slider_value):
    sentence = ""
    if logit < slider_value / 2.5:
        sentence = "Nope. Not at all."
    elif slider_value / 2.5 < logit < slider_value / 1.56:
        sentence = "Not really..."
    elif slider_value / 1.56 < logit < slider_value / 1.36:
        sentence = "Close but not there."
    elif slider_value / 1.36 < logit < slider_value / 1.14:
        sentence = "That's quite close."
    elif slider_value / 1.14 < logit < slider_value:
        sentence = "Almost guessed."
    elif logit >= slider_value:
        sentence = "YES!!"
    return sentence


def give_up():
    image = clip_chat.image_org
    return image, None, "You lost... (Press \"Reset\" to play again)"


def update_difficulty(x):
    if not has_started:
        clip_chat.goal = x
        return clip_chat.goal
    return clip_chat.goal


has_started = False
best_guess = None


def respond(message, chat_history, label_value, image_value):
    global has_started, best_guess

    logits, is_better = clip_chat.answer(message)
    bot_message = logit2sentence(logits, clip_chat.goal)

    if is_better == 3:
        best_guess = {f"Best Guess: \"{message}\"": float(logits) / clip_chat.goal}
    if float(logits) >= clip_chat.goal:
        bot_message = "YES!"
        best_guess = "YOU WIN! (Press \"Reset\" to play again)"
        image_value = clip_chat.image_org
    else:
        if has_started:
            if is_better == -1:
                bot_message += ""
            elif is_better == 0:
                bot_message += "You did worse than the last one."
            elif is_better == 1 or is_better == 3:
                bot_message += "You did better than the last one."

    if not has_started:
        has_started = True

    label_value = best_guess

    chat_history.append((message, bot_message))
    return "", chat_history, label_value, image_value


def reset_everything():
    global has_started, best_guess
    clip_chat.reset_everything()
    has_started = False
    best_guess = None
    return clip_chat.goal, None, "This is a \"Guess the Image\" game. I'm thinking of a picture and you have to guess using the chat above.", None


with gr.Blocks() as demo:
    with gr.Row():
        with gr.Column():
            chatbot = gr.Chatbot()
            msg = gr.Textbox()
            slider = gr.inputs.Slider(minimum=18, maximum=30, default=23, label="Difficulty (18 - Easy, 30 - Expert)")
            description = gr.Markdown("Start writing in the textbox above. The percentage in the label in the right shows how close you are to the answer.\n\n**TIPS:**\n- Write in English\n- Start with simple general stuff like \"A person\", \"A vehicle\", \"An object\", ...\n- If you feel you are not getting any closer to the answer, try a different approach completly! (maybe a 60% of animal is a 80% a vehicle).\n- Adding \"A photo of ... outdoors/indoors\" might get better results.")
        with gr.Column():
            label = gr.Label("This is a \"Guess the Image\" game. I'm thinking of a picture and you have to guess using the chat above.")
            image_output = gr.outputs.Image(type="pil")
            show_image_button = gr.Button("Give Up...")
            reset_button = gr.Button("Reset")

    msg.submit(respond, [msg, chatbot], [msg, chatbot, label, image_output])
    slider.release(update_difficulty, inputs=[slider], outputs=[slider])
    show_image_button.click(give_up, outputs=[image_output, chatbot, label], queue=False)
    reset_button.click(reset_everything, outputs=[slider, image_output, label, chatbot], queue=False)

if __name__ == "__main__":
    demo.title = "CLIP Guess the Image"
    demo.launch(share=False)