import pandas as pd

from openai import OpenAI

FILE_NAME = "data/gaia_validation_20.jsonl"

FINAL_ANSWER_MODEL = "gpt-4.5-preview"

def get_questions():
    df = pd.read_json(FILE_NAME, lines=True)
    result=[]
    
    for index, row in df.iterrows():
        result.append([row["Level"], row["Question"], row["file_name"], row["Final answer"]])

    return result

def get_final_answer(question, initial_answer):
    prompt_template = """
        You are an expert question answering assistant. Given a question and an initial answer, your task is to provide the final answer.
        Your final answer must be a number and/or string OR as few words as possible OR a comma-separated list of numbers and/or strings.
        If you are asked for a number, don't use comma to write your number neither use units such as $ or % unless specified otherwise.
        If you are asked for a string, don't use articles, neither abbreviations, and write the digits in plain text unless specified otherwise.
        If you are asked for a comma-separated list, apply the above rules depending of whether the element to be put in the list is a number or a string.
        If the final answer is a single word, start uppercase.

        **Question:** """ + question + """
        
        **Initial answer:** """ + initial_answer + """
        
        **Example 1:** How many 'r's are in strawberry? 3
        **Example 2:** What is the opposite of black? White
        **Example 3:** What is the biggest city in California? Los Angeles
        **Example 4:** What is the superlative of good? Best
        **Example 5:** How many states are in the USA? 50
        
        **Final answer:** 

        """

    client = OpenAI()
    completion = client.chat.completions.create(
        messages=[{"role": "user", "content": [{"type": "text", "text": prompt_template}]}],
        model=FINAL_ANSWER_MODEL
    )

    final_answer = completion.choices[0].message.content
    
    print("###")
    print(question)
    print(initial_answer)
    print(final_answer)
    print("###")
    
    return final_answer