import gradio as gr
from google.colab import userdata
from openai import OpenAI
import gradio as gr
import requests
import os
from PIL import Image
import numpy as np
import ipadic
import MeCab
import difflib
import io
import os

api_key = os.getenvs('OPENAI_API_KEY')
client = OpenAI(api_key=api_key)

def generate_image(text):
    image_path = f"/content/images/{text}.png"
    if not os.path.exists(image_path):
        response = client.images.generate(
          model="dall-e-3",
          prompt=text,
          size="1024x1024",
          quality="standard",
          n=1,
        )
        image_url = response.data[0].url
        image_data = requests.get(image_url).content
        img = Image.open(io.BytesIO((image_data)))
        img = img.resize((512, 512))
        img.save(image_path)
    return image_path

def calulate_similarity_score(ori_text, text):
    if ori_text != text:
        model_name = "text-embedding-3-small"
        response = client.embeddings.create(input = [ori_text, text], model=model_name)
        score = cos_sim(response.data[0].embedding, response.data[1].embedding)
        score = int(round(score, 2) * 100)
        if score == 100:
            score = 99
    else:
      score = 100
    return score

def cos_sim(v1, v2):
    return np.dot(v1, v2) / (np.linalg.norm(v1) * np.linalg.norm(v2))

def tokenize_text(text):
    mecab = MeCab.Tagger(f"-Ochasen {ipadic.MECAB_ARGS}")
    return [
        t.split()[0]
        for t in mecab.parse(text).splitlines()[:-1]
    ]

def create_match_words(ori_text, text):
    ori_words = tokenize_text(ori_text)
    words = tokenize_text(text)
    match_words = [w for w in words if w in ori_words]
    return match_words

def create_hint_text(ori_text, text):
    response = list(difflib.ndiff(list(text), list(ori_text)))
    output = ""
    for r in response:
      if r[:2] == "- ":
          continue
      elif r[:2] == "+ ":
          output += "^"
      else:
          output += r.strip()
    return output

def update_question(selected_option):
    if selected_option == "Q1":
        return "/content/images/白い猫が木の上で休んでいる.png"
    elif selected_option == "Q2":
        return "/content/images/サメが海の中で暴れている.png"
    elif selected_option == "Q3":
        return "/content/images/東京スカイツリーの近くで花火大会が行われている.png"
    elif selected_option == "Q4":
        return "/content/images/イカとタイがいた都会.png"
    elif selected_option == "Q5":
        return "/content/images/赤いきつねと緑のたぬき.png"
    if selected_option == "Q6":
        return "/content/images/宇宙に向かってたい焼きが空を飛んでいる.png"
    elif selected_option == "Q7":
        return "/content/images/イケメンが海岸でクリームパンを眺めている.png"
    elif selected_option == "Q8":
        return "/content/images/生麦生米生卵生麦生米生卵生麦生米生卵.png"
    elif selected_option == "Q9":
        return "/content/images/サイバーエージェントで働く人たち.png"
    elif selected_option == "Q10":
        return "/content/images/柿くへば鐘が鳴るなり法隆寺.png"
    elif selected_option == "Q11":
        return "/content/images/鳴くよウグイス平安京.png"
    else:
        return "/content/images/abc.png"

def main(image, text, option):
    ori_text = update_question(option).split("/")[-1].split(".png")[0]
    image_path = generate_image(text)
    score = calulate_similarity_score(ori_text, text)

    if score < 80:
        match_words = create_match_words(ori_text, text)
        hint_text = "一致している単語リスト: "+" ".join(match_words)
    elif 80 <= score < 100:
        hint_text = "一致していない箇所: "+create_hint_text(ori_text, text)
    else:
        hint_text = ""
    return image_path, f"{score}点", hint_text

with gr.Blocks() as demo:
    with gr.Row():
        with gr.Column():
            gr.Markdown(
                "# プロンプトを当てるゲーム \n これは表示されている画像のプロンプトを当てるゲームです。プロンプトを入力するとそれに対応した画像とスコアとヒントが表示されます。スコア100点を目指して頑張ってください！ \n\nヒントは80点未満の場合は当たっている単語、80点以上の場合は足りない文字を「^」で示した文字列を表示しています。",
            )
            selected_option = gr.components.Radio(["Q1", "Q2", "Q3", "Q4", "Q5", "Q6", "Q7", "Q8", "Q9", "Q10", "Q11"], label="問題を選んでください！")
            output_title_image = gr.components.Image(type="filepath", label="お題")
            selected_option.change(update_question, inputs=[selected_option], outputs=[output_title_image])

            input_text = gr.components.Textbox(lines=1, label="画像にマッチするテキストを入力して！")
            submit_button = gr.Button("Submit")
        with gr.Column():
            output_image = gr.components.Image(type="filepath", label="生成画像")
            output_score = gr.components.Textbox(lines=1, label="スコア")
            output_hint_text = gr.components.Textbox(lines=1, label="ヒント")

    submit_button.click(main, inputs=[output_title_image, input_text, selected_option], outputs=[output_image, output_score, output_hint_text])
demo.launch(server_port=8892)
demo.launch()