import warnings
warnings.filterwarnings("ignore")

import os, json
import gradio as gr
import pandas as pd

from backend import Backend



QUESTIONS = [
    "Animal Type",
    "Exposure Age",
    "Behavior Test",
    "Intervention 1",
    "Intervention 2",
    "Genetic Chain",
]

with gr.Blocks(theme="dark") as demo:
    backend = Backend()
    with gr.Row():
        with gr.Row():
            # Update
            with gr.Group():
                gr.Markdown(f'<center><h1>Input</h1></center>')
                gr.Markdown(f'<center><p>Please First Upload the File</p></center>')

                openai_key = gr.Textbox(
                        label='Enter your OpenAI API key here', 
                        type='password')
                
                file = gr.File(label='Upload your .txt or .pdf file here', file_types=['.txt', '.pdf'], file_count = 'multiple')
                
                questions = gr.CheckboxGroup(choices = QUESTIONS, value = QUESTIONS, label="Questions", info="Please select the question you want to ask")

                btn_submit_txt = gr.Button(value='Submit')
                btn_submit_txt.style(full_width=True)

            # Output
            with gr.Group():
                gr.Markdown(f'<center><h1>Output</h1></center>')
                gr.Markdown(f'<center><p>The answer to your question is :</p></center>')
                filename_box = gr.Textbox(label = "File")
                question_box = gr.Textbox(label='Question')
                answer_box = gr.Textbox(label='Answer')
                # reference_box = gr.Textbox(label='Reference')

                highlighted_text = gr.outputs.HTML(label="Highlighted Text")
                with gr.Group():
                    gr.Markdown("<center><h4>Please select different questions</h4></center>")
                    with gr.Row():
                        btn_last_question = gr.Button(value='Last Question')
                        btn_next_question = gr.Button(value='Next Question')

                with gr.Group():
                    gr.Markdown("<center><h4>Please select different passages</h4></center>")
                    with gr.Row():
                        btn_last_passage = gr.Button(value='Last Passage')
                        btn_next_passage = gr.Button(value='Next Passage')
            
            # Correctness
            with gr.Group():
                gr.Markdown(f'<center><h1>Correct the Result</h1></center>')
                gr.Markdown(f'<center><p>Please Correct the Results</p></center>')

                with gr.Row():
                    save_results = gr.Textbox(placeholder = "Still need to click the button above to save the results", label = 'Save Results')
                with gr.Group():
                    gr.Markdown(f'<center><p>Please Choose: </p></center>')
                    answer_correct = gr.Radio(choices = ["Correct", "Incorrect"], label='Is the Generated Answer Correct?', info="Pease select whether the generated text is correct")
                    correct_answer = gr.Textbox(placeholder = "Please judge on the generated answer", label = 'Correct Answer', interactive = True)
        
                    reference_correct = gr.Radio(choices = ["Correct", "Incorrect"], label="Is the Reference Correct?", info="Pease select whether the reference is correct")
                    correct_reference = gr.Textbox(placeholder = "Please judge on the generated answer", label = 'Correct Reference', interactive = True)
                
                    btn_submit_correctness = gr.Button(value='Submit Correctness')
                    btn_submit_correctness.style(full_width=True)
            
            # Download
            with gr.Group():
                gr.Markdown(f'<center><h1>Download</h1></center>')
                gr.Markdown(f'<center><p>Download the original LLM answers and corrected LLM answers</p></center>')
                answer_file = gr.File(label='Download original LLM answers', file_types=['.xlsx'])
                btn_download_answer = gr.Button(value='Download original LLM answers')
                btn_download_answer.style(full_width=True)
                corrected_file = gr.File(label='Download corrected data', file_types=['.xlsx'])
                btn_download_corrected = gr.Button(value='Download corrected LLM answers')
                btn_download_corrected.style(full_width=True)
            
        
    with gr.Row():
        reset = gr.Button(value='Reset')
        reset.style(full_width=True)

    # Answer change
    answer_correct.input(
        backend.change_correct_answer,
        inputs = [answer_correct],
        outputs = [correct_answer],
    )

    reference_correct.input(
        backend.change_correct_reference,
        inputs = [reference_correct],
        outputs = [correct_reference],
    )
    

    # Submit button
    btn_submit_txt.click(
            backend.process_file,
            inputs=[file, questions, openai_key],
            outputs=[filename_box, question_box, answer_box, highlighted_text, correct_answer, correct_reference],
        )
    
    btn_submit_correctness.click(   # TODO
            backend.process_results,
            inputs=[answer_correct, correct_answer, reference_correct, correct_reference],
            outputs=[save_results],
        )
    
    # Switch question button
    btn_last_question.click(
        backend.process_last,
        outputs=[filename_box, question_box, answer_box, highlighted_text, correct_answer, correct_reference, save_results, answer_correct, reference_correct],
    )

    btn_next_question.click(
            backend.process_next,
            outputs=[filename_box, question_box, answer_box, highlighted_text, correct_answer, correct_reference, save_results, answer_correct, reference_correct],
        )
    
    # Switch passwage button
    btn_last_passage.click(
        backend.switch_last_passage,
        outputs=[filename_box, question_box, answer_box, highlighted_text, correct_answer, correct_reference, save_results, answer_correct, reference_correct],
    )
    btn_next_passage.click(
        backend.switch_next_passage,
        outputs=[filename_box, question_box, answer_box, highlighted_text, correct_answer, correct_reference, save_results, answer_correct, reference_correct],
    )

    # Download button
    btn_download_answer.click(
            backend.download_answer,
            outputs=[answer_file],
        )
    
    btn_download_corrected.click(
            backend.download_corrected,
            outputs=[corrected_file],
        )

demo.queue()
demo.launch(show_error=True, show_tips=True)