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):
    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'], interactive=True, label="Evaluation Strategies")
        question_input = gr.Textbox(label="Question", lines=1, placeholder="Enter your question here...")

    instruction = gr.Textbox(label="Instructions", visible=False, interactive=False)
    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()