import gradio as gr
import pandas as pd
import text_gen as gen


NUM_QUESTIONS = 10
answers = {}
separator = '<[._.]>'

df = pd.read_csv('lines_2.txt', sep='*')
df = pd.concat([df[df.Part == 'Start'].sample(1),
                df[df.Part == 'Middle_1'].sample(1),
                df[df.Part == 'Middle_2'].sample(1),
                df[df.Part == 'Middle_3'].sample(1),
                df[df.Part == 'Middle_4'].sample(NUM_QUESTIONS-5),
                df[df.Part == 'End'].sample(1)])

questions = dict(zip(df.Question, df.Answer))
context = dict(zip(df.Question, df.Context))

def get_answer(question, answer, options):
    global answers
    global separator
    answer = options.split(separator)[answer]
    answers.update({question:int(answer == questions[question])})


def set_score():
    global answers
    global NUM_QUESTIONS
    score = sum(answers.values())
    start = '## <p style="text-align: center;">'
    end = '</p>'
    if score == 0: return f'{start}Not a single right answer. Are you doing this on purpose?{end}'
    elif score <= NUM_QUESTIONS / 2 + 1: return f'{start}Only {score} right answers out of {NUM_QUESTIONS}. You ought to pay more attention.{end}'
    elif score == NUM_QUESTIONS: return f'{start}Perfect score!{end}'
    else: return f'{start}{score} right answers out of {NUM_QUESTIONS}. It\'s probably alright.{end}'


wrong_answers = {q:[gen.generate_text(
                    q,
                    context[q],
                    gen.model_names[i],
                    gen.model[i],
                    gen.tokenizers[i],
                    minimum=len(questions[q].split())+8)
                    for i in range(3)]
                for q in questions.keys()}


with gr.Blocks(theme='glass') as demo:
    with gr.Row():

        with gr.Column(scale=1):
            pass

        with gr.Column(scale=2):

            gr.Markdown(f'## <p style="text-align: center;">IMITATION GAME</p>\n' +
                        f'### <p style="text-align: center;">Choose answers NOT given by an AI model.')

            for ind, question in enumerate(list(questions.keys()), start=1):
                letters = list('ABCD')
                options = list(set(wrong_answers[question] + [questions[question]]))

                gr.Markdown(f'### <p>{ind}. {question}</p>')

                for letter, option in zip(letters, options):
                    gr.Markdown(f'{letter}. {option}')

                options = gr.State(separator.join(options))

                radio = gr.Radio(letters, type='index', show_label=False)
                question = gr.State(question)
                radio.change(fn=get_answer, inputs=[question, radio, options])

            button = gr.Button(value='Get score')
            score = gr.Markdown()
            button.click(fn=set_score, outputs=score)

        with gr.Column(scale=1):
            pass