import gradio as gr
import openai
import pytesseract
import random
import os
import time

Init_system_prompt = "You are an AI Assistant that tries to teach kids various subjects. You are given learning material and you task is to ask questions given the material and then you also grade answers and give feedback how to improve the answers"
system_message = {"role": "system", "content": Init_system_prompt}

system_prompts = {
    "English": "You are an AI Assistant that tries to teach kids various subjects. You are given learning material and you task is to ask questions given the material and then you also grade answers and give feedback.",
    "Finnish": "Olet tekoälyavustaja jonka tehtävänä on auttaa lapsia oppimaan koulussa. Sinulle annetaan oppimateriaalia tekstinä ja sinun tehtäväsi on kysyä kysymyksiä annetusta tekstistä, arvostella vastauksia ja antaa palautetta kuinka vastauksia voidaan parantaa."
}

question_strings = {
    "English": "\n Please ask a question about the previous paragramph: Question:",
    "Finnish": "\n Kysy kysymys edellisen kappaleen perusteella. Kysymys:",
}

lang_mapping = {
    "English": "eng",
    "Finnish": "fin"
}

grading_button_mapping = {
    "English": gr.Button(value="Grade my answer"),
    "Finnish": gr.Button(value="Arvostele vastaukseni")
}

new_question_button_mapping = {
    "English": gr.Button(value="New question"),
    "Finnish": gr.Button(value="Uusi kysymys")
}

clear_button__mapping = {
    "English": gr.Button(value="Clear messages"),
    "Finnish": gr.Button(value="Tyhjennä viestit")
}

grading_prompt_start_mapping = {
    "English": "You are given question and answer. Give grading between one to ten with reasoning and give feedback how answer should be improved to get better grade. Here are rules how to format your answer, it is really important to stick to these rules:\n Your answer should be devided into the following sections separated by newline: Grading, feedback, reasoning for feedback. \n Do not repeat Answer or Question.",
    "Finnish": "Saat kysymyksen ja vastauksen. Arvostele vastaus asteikolla yhdestä kymmeneen perusteluineen ja anna palautetta kuinka vastausta tulisi parantaa paremman arvosanan saamiseksi. \n Tässä ovat ohjeet vastauksesi muotoilulle:\n Vastauksesi tulee olla jaettuna seuraaviin osioihin rivinvaihdolla erotettuna: Arvostelu, palaute, perustelut palautteelle. \n Älä toista seuraavia osioita: Kysymys, Vastaus."
}

new_question_prompt = {
    "English": "Next question please",
    "Finnish": "Saisinko uuden kysymyksen"
}

import os
from PIL import Image

import pytesseract
os.system("rm -f path.txt")
path = os.system("which tesseract >> path.txt")
with open("path.txt", 'r') as file:
    tesseract_path = file.read().replace('\n', '')


########### TAB 1 (UPLOAD) FUNCTIONS  #############################

def print_files(files):
    for file in files:
        print(file.name)


def create_data(files, language_selection):
    question_context = ''
    for file in files:
        if file.name.endswith('png') or file.name.endswith('.jpg') or file.name.endswith('.jpeg'):
            try:
                question_context += (pytesseract.image_to_string(Image.open(file.name), lang=lang_mapping[language_selection])) + '\n\n'
            except Exception as e:
                print(e)
                pass
    system_prompt = system_prompts[language_selection]
    

    return question_context, system_prompt, new_question_button_mapping[language_selection], clear_button__mapping[language_selection]


########### TAB 2 (CHAT) FUNCTIONS  #############################

def user(user_message, history):
        return "", history + [[user_message, None]]

def add_new_question(language_selection, history):
        return history + [[new_question_prompt[language_selection], None]]

def bot(history, messages_history, api_key, system_prompt, teksti_contexti, temperature, max_tokens, chatgpt_model, max_context_size_for_question, language_selection):
    user_message = history[-1][0]
    bot_message, messages_history = ask_gpt(user_message, messages_history, api_key, system_prompt, teksti_contexti, temperature, max_tokens, chatgpt_model, max_context_size_for_question, language_selection)
    messages_history += [{"role": "assistant", "content": bot_message}]
    history[-1][1] = bot_message
    return history, messages_history



def ask_gpt(message, messages_history, api_key, system_prompt, context, temperature, max_tokens, chatgpt_model, max_context_size_for_question, language_selection):
        if message == new_question_prompt[language_selection]:
            max_possible_position = len(context)- max_context_size_for_question
            start = random.randint(0,max_possible_position)
            messages_history += [{"role": "user", "content": context[start:start+max_context_size_for_question] + question_strings[language_selection]}]
            
            openai.api_key = api_key
            response = openai.ChatCompletion.create(
                model=chatgpt_model,
                messages=messages_history,
                temperature=temperature,
                max_tokens=max_tokens
            )
            return response['choices'][0]['message']['content'], messages_history
        else:
            if len(messages_history) <= 1:
                max_possible_position = len(context)- max_context_size_for_question
                start = random.randint(0,max_possible_position)
                messages_history += [{"role": "user", "content": context[start:start+max_context_size_for_question] + question_strings[language_selection]}]
                
                openai.api_key = api_key
                response = openai.ChatCompletion.create(
                    model=chatgpt_model,
                    messages=messages_history,
                    temperature=temperature,
                    max_tokens=max_tokens
                )
                return response['choices'][0]['message']['content'], messages_history
            else:
                question = messages_history[-1]['content']
                if language_selection == 'English':
                    prompt_start = grading_prompt_start_mapping[language_selection]
                    prompt_end = 'Question: ' + question + '\n Answer: ' + message + '\n'
                    full_prompt = prompt_start + prompt_end
                elif language_selection == 'Finnish':
                    prompt_start = grading_prompt_start_mapping[language_selection]
                    prompt_end = 'Kysymys: ' + question + '\n Vastaus: ' + message + '\n'
                    full_prompt = prompt_start + prompt_end

                messages_history += [{"role": "user", "content": full_prompt}]

                openai.api_key = api_key
                response = openai.ChatCompletion.create(
                    model=chatgpt_model,
                    messages=messages_history,
                    temperature=temperature,
                    max_tokens=max_tokens
                )
                return prompt_end.replace('\n', '<br>') + response['choices'][0]['message']['content'].replace('\n', '<br>'), messages_history

def init_history(messages_history, system_prompt):
    messages_history = []
    messages_history += [{"role": "system", "content": system_prompt}]
    msg_log = gr.Textbox(value="Tähän tulee message history")
    system_prompt = gr.Textbox(value=system_prompt, label='Insert system message here')
    return messages_history, system_prompt, msg_log


############# INTERFACE ##########################
with gr.Blocks() as demo:
    gr.Markdown("Exam preparation demo: Upload images of textbooks and start learning!")
    
    
    
    ############# TAB 1 ##########################
    with gr.Tab("Upload documents and create context"):
        with gr.Row():
            api_key = gr.Textbox(value='', type='password', label='Insert OPENAI API-key here')
        with gr.Row():
            language_selection = gr.Dropdown(value='English', choices=["English","Finnish"], label='Select language')
            files = gr.File(file_count='multiple', file_types=['image'], interactivate = True)
            create_context_btn = gr.Button(value='Recognize text and create context')
        with gr.Row():
            gr.Markdown("")
        with gr.Row():
            teksti_contexti = gr.Textbox(value='Tähän tulee konteksti', label='Created context')
    
    ############# TAB 2 ##########################

    with gr.Tab("Chat"):
        gr.Markdown("""<h1><center>ChatGPT 
        ChatBot with Gradio and OpenAI</center></h1>
        """)
        new_question_state_msg = gr.State(value=[["New_question", None]])#, hidden=True)
        with gr.Row():
            system_prompt = gr.Textbox(value=Init_system_prompt, label='Insert system message here')
            chatgpt_model = gr.Dropdown(choices=["gpt-3.5-turbo", "gpt-3.5-turbo-0301", "gpt-3.5-turbo-0613"], value='gpt-3.5-turbo',label='ChatGPT model to use', interactive=True)
            temperature = gr.Slider(minimum=0.0, maximum=1.0, step=0.05, value=0.7, label='Temperature')
            max_tokens = gr.Slider(minimum=10, maximum=600, step=10, value=200, label='Max tokens')
            max_context_size_for_question = gr.Slider(minimum=200, maximum=2000, step=50, value=300, label='Max context for questions')
        with gr.Row():
            chatbot = gr.Chatbot(label='ChatGPT Chat')
            state = gr.State([])
        with gr.Row():
            msg = gr.Textbox()
        with gr.Row():
            new_question = gr.Button(value="New_question")
            clear = gr.Button("Clear")
                    
        
    with gr.Accordion("Klikkaa avataksesi ohjeet"):
        gr.Markdown("Ohjeet tulee tänne")


    # TAB 1 (CREATE CONTEXT) Interactive elements:
    create_context_btn.click(create_data, [files, language_selection], [teksti_contexti, system_prompt, new_question, clear])


    # TAB 2 (CHAT) Interactive elements:
    msg.submit(user, [msg, chatbot], [msg, chatbot], queue=False).then(
        bot, [chatbot, state, api_key, system_prompt, teksti_contexti, temperature, max_tokens, chatgpt_model, max_context_size_for_question, language_selection], [chatbot, state]
    )
    new_question.click(add_new_question, [language_selection, chatbot], [chatbot], queue=False).then(
        bot, [chatbot, state, api_key, system_prompt, teksti_contexti, temperature, max_tokens, chatgpt_model, max_context_size_for_question, language_selection], [chatbot, state]
    )

    clear.click(lambda: None, None, chatbot, queue=False).success(init_history, [state, system_prompt], [state, system_prompt])


demo.launch(debug=True)