import gradio as gr
from pal import solve_pal
from mathprompter import solve_mp
from TA import solve_ta
from utils import run_code


def run(token, question, method):
    if len(token) == 0:
        raise gr.Error("Please provide a valid Hugging Face access token")
    if method == "PAL":
        code_op, generated_code = solve_pal(question, token)
        if code_op is not None:
            return code_op, gr.Code.update(value=generated_code, interactive=True), gr.Button.update(visible=True)
        else:
            return (
                "Code execution failed, please review it from below and re-run it or try-asking again with more details",
                gr.Code.update(value=generated_code, interactive=True), gr.Button.update(visible=True))

    elif method == "TA":
        code_op, generated_code = solve_ta(question, token)
        if code_op is not None:
            return code_op, gr.Code.update(value=generated_code, interactive=True), gr.Button.update(visible=True)
        else:
            return (
                "Code execution failed, please review it from below and re-run it or try-asking again with more details",
                gr.Code.update(value=generated_code, interactive=True), gr.Button.update(visible=True))

    elif method == "MathPrompter":
        exp_op, code_op, generated_code = solve_mp(question, token)

        if code_op is not None:
            ans = f"Answer from Expression execution: {exp_op} \nAnswer from Code execution: {code_op}  "
            return ans, gr.Code.update(value=generated_code, interactive=True), gr.Button.update(visible=True)
        else:
            return (
                "Code execution failed, please review it from below and re-run it or try-asking again with more details",
                gr.Code.update(value=generated_code, interactive=True), gr.Button.update(visible=True))

    else:
        raise gr.Error("Please select the evaluating strategy from dropdown")


def run_edits(code):
    if "input(" in code:
        return "Code execution failed, Please remove any input statement or bugs", code
    try:
        code_op = run_code(code)
        return code_op, code
    except:
        return "Code execution failed, please review it from below and re-run it or try-asking again with more details", code


# User Interface Part

TITLE = "Reasoning with StarCoder 💫"
demo = gr.Blocks(title=TITLE, theme=gr.themes.Default())


def render_instruction(mtd):
    if mtd == "PAL":
        return gr.Textbox.update(
            value='''
            💫 Query can be an instruction or a direct question starting with What, Find, etc.
            💫 Use numbers wherever possible, i.e use 2 instead of two
            💫 Example: What is the value of sin(30)?
            ''',
            visible=True
        )
    if mtd == "TA":
        return gr.Textbox.update(
            value='''
            💫 Query should be a direct instruction starting with *write the code* and the main question.
            💫 Use numbers wherever possible, i.e use 2 instead of two
            💫 Example: Write the  code to find 7th Fibonacci number.
            ''',
            visible=True
        )
    if mtd == "MathPrompter":
        return gr.Textbox.update(
            value='''
            💫 Query should be a direct question, can start by giving a context and then asking the intended segment.
            💫 Use numbers wherever possible, i.e use 2 instead of two
            💫 Example: The dimensions of the input image are 80x80, if the filter of 3x3 size is convoluted on the image, what are the dimensions of output image?
            ''',
            visible=True
        )


with demo:
    gr.HTML(f"<center><h1>{TITLE}<h1></center>")

    access_token = gr.Textbox(type="password", label="Hugging Face Access Token")
    with gr.Row():
        methods = gr.Dropdown(choices=['PAL', 'TA', 'MathPrompter'], value="PAL",interactive=True, label="Evaluation Strategies")
        question_input = gr.Textbox(label="Question", lines=1, placeholder="Enter your question here...")

    instruction = gr.Textbox(label="Instructions", visible=True, interactive=False, value=render_instruction("PAL")['value'])
    methods.change(fn=render_instruction, inputs=methods, outputs=instruction)

    question_output = gr.Textbox(label="Answer", interactive=False)
    code = gr.Code(language="python", interactive=False, label="Generated Code (Editable)")
    submit_btn = gr.Button("Submit")
    edit_btn = gr.Button("Run the edited code", visible=False)

    submit_btn.click(run, inputs=[access_token, question_input, methods], outputs=[question_output, code, edit_btn])
    edit_btn.click(run_edits, inputs=code, outputs=[question_output, code])
    gr.Examples(
        examples=[
            [
                "The dimensions of the input image are 80x80, if the filter of 3x3 size is convoluted on the image, what are the dimensions of output image?",
                "PAL"],
            [
                "The dimensions of the input image are 80x80, if the filter of 3x3 size is convoluted on the image, what are the dimensions of output image?",
                "MathPrompter"],
            ["What is the value of sin(30)?", "PAL"],
            ["Write the  code to find 7th Fibonacci number.", "TA"],
            ["Write a program to filter all the odd numbers from a python list", "PAL"],
        ],
        inputs=[question_input, methods],
        outputs=[question_output, code, edit_btn],
        fn=run,
        cache_examples=False,
        label="Sample Questions",
    )

demo.launch()

'''
Carol was playing a trivia game. In the first round she scored 17 points and in the second round she scored 6 points. In the last round she lost 16 points. How many points did she have at the end of the game?
While on vacation, Debby took 24 pictures at the zoo and 12 at the museum. If she later deleted 14 of the pictures, how many pictures from her vacation did she still have?
I had 5 apples, I gave 3 to my brother and 4 to my uncle. How many apples I am left with? 
'''