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 = 10
MAX_CLUE_VALUE = 20


def fn_to_str(challenge_idx, x):
    return f"f({x}) = {functions[challenge_idx](x)}"


def get_gpt_guess(fidx, 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(lambda x: fn_to_str(fidx, x), 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:
    # SESSION VALUES
    ask = gr.State(lambda: random.randint(1, MAX_INPUT_ANSWER_VALUE))
    clue = gr.State(random.choice(list(set(range(MAX_CLUE_VALUE)) - set([ask.value]))))

    available_numbers = gr.State(set(range(MAX_CLUE_VALUE)) - set([ask.value, clue.value]))
    challenge_idx = gr.State(random.randint(0, len(functions) - 1))

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

        with gr.Column():
            clues_box = gr.Markdown(f"# Clues\n```\n{fn_to_str(challenge_idx.value, clue.value)}\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")
            rstbtn = gr.Button("Reset")

    def submit_guess(value, history, clue_values, fidx, avail, _ask ):
        history.append(value)

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

        result = check_winner(gpt_guess_value, int(value), functions[fidx](_ask))

        _clue = random.choice(list(avail))
        avail.remove(_clue)

        clue_values.append(_clue)
        clue_str = "\n".join(map(lambda x: fn_to_str(fidx, x), 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```",
            available_numbers: avail,
            challenge_idx: fidx,
            ask: _ask,
        }

    btn.click(
        submit_guess,
        [guess, guesses, clues, challenge_idx, available_numbers, ask],
        [guesses, history_box, gpt_guess, clues, title_text, clues_box, available_numbers, challenge_idx, ask],
    )

    def reset():
        _ask = random.randint(1, MAX_INPUT_ANSWER_VALUE)
        _clue = random.choice(list(set(range(2, MAX_CLUE_VALUE)) - set([_ask])))
        fidx = random.randint(0, len(functions) - 1)
        clue_str = "\n".join(map(lambda x: fn_to_str(fidx, x), [_clue]))
        return {
            guesses: [],
            history_box: "",
            clues: [_clue],
            title_text: f"# Guess the function (faster than GPT)\n## What is f({_ask})?",
            clues_box: f"# Clues\n```\n{clue_str}\n```",
            gpt_guess: "...",
            ask: _ask,
            clue: _clue,
            available_numbers: set(range(MAX_CLUE_VALUE)) - set([_ask, _clue]),
            challenge_idx: fidx,
        }
    
    rstbtn.click(reset, [], [ask, clue, available_numbers, challenge_idx, guesses, history_box, gpt_guess, clues, title_text, clues_box])


demo.launch()