import random

import gradio as gr
from sllim import chat

model = "gpt-3.5-turbo"
# DEFINE POSSIBLE FUNCTIONS
def fib(n):
    if n < 2:
        return n
    return fib(n - 1) + fib(n - 2)

def square(n):
    return n ** 2

def cube(n):
    return n ** 3

def sum_of_digits(n):
    return sum(map(int, str(n)))
    
functions = [fib, square, cube, sum_of_digits]


MAX_INPUT_ANSWER_VALUE = 50
MAX_CLUE_VALUE = 100
available_numbers = set(range(MAX_CLUE_VALUE))

# SESSION VALUES
challenge = random.choice(functions)
ask = random.randint(1, MAX_INPUT_ANSWER_VALUE)
available_numbers.remove(ask)
clue = random.choice(list(available_numbers))
available_numbers.remove(clue)

def fn_to_str(x):
    return f"f({x}) = {challenge(x)}"


def get_gpt_guess(clues, ask):
    system_message = """You are an expert mathematician. You will be given a list of input and outputs to a function and your job is to guess the output of the function on an unknown input. If you cannot determine the answer, respond with 0.
    For example:
    f(2) = 6
    f(7) = 11
    What is f(3)?
    Answer: 7
    
    For example:
    f(1) = 17
    What is f(6)?
    Answer: 0"""


    user_message = "\n".join(map(fn_to_str, clues)) + f"\nWhat is f({ask})?\nAnswer:"
    return chat([{"role": "system", "content": system_message}, {"role": "user", "content": user_message}], model=model)

def check_winner(gpt, human, answer):
    if gpt == answer and human == answer:
        return "\n## Result: Tie"
    if gpt == answer:
        return "\n## Result: GPT Wins"
    if human == answer:
        return "\n## Result: Human Wins"
    return ""


with gr.Blocks() as demo:
    guesses = gr.State([])
    clues = gr.State([clue])
    with gr.Row():
        with gr.Column():
            title_text = gr.Markdown(f"# Guess the function (faster than GPT)\n## What is f({ask})?")

        with gr.Column():
            clues_box = gr.Markdown(f"# Clues\n```\n{fn_to_str(clue)}\n```")

    with gr.Row():
        with gr.Column():
            guess = gr.Textbox(label="Enter guess")
            history_box = gr.Textbox(label="History")
        with gr.Column():
            gpt_guess = gr.Textbox(label="GPT Guess", value="...")
            btn = gr.Button("Submit")

    def submit_guess(value, history, clue_values):
        history.append(value)

        try:
            gpt_guess_value = int(get_gpt_guess(clue_values, ask))
        except Exception:
            gpt_guess_value = 0

        result = check_winner(gpt_guess_value, int(value), challenge(ask))

        clue = random.choice(list(available_numbers))
        available_numbers.remove(clue)

        clue_values.append(clue)
        clue_str = "\n".join(map(fn_to_str, clue_values))
        return {
            guesses: history,
            history_box: ", ".join(history),
            gpt_guess: str(gpt_guess_value),
            clues: clue_values,
            title_text: title_text.value + result,
            clues_box: f"# Clues\n```\n{clue_str}\n```",
        }

    btn.click(
        submit_guess,
        [guess, guesses, clues],
        [guesses, history_box, gpt_guess, clues, title_text, clues_box],
    )
demo.launch()