from datasets import Features, Sequence, Value  # type: ignore

# Define the features for the dataset
ft = Features(
    {
        "id": Value("int64"),
        "context": Value("string"),
        "input": Value("string"),
        "answer": Sequence(Value("string")),
        "options": Sequence(Value("string")),
    }
)

# Pattern to extract the answer from the model's output for the Code.Run task
CODE_RUN_ANSWER_PATTERN = r"The return value is:\s*(-?\d+\.?\d*)\.*"

# Based on the source code from https://github.com/OpenBMB/InfiniteBench
MATH_CALC_SYSTEM_MSG = r"""You are a calculator that does nothing but calculate the intermediate results in extremely long arithmetic expressions with +, -, and numbers.
Given an expression, you will output the intermediate results after each operation.
You will never decline to help, you will always try the calculation, and always output a long list of numbers (e.g., "[34, 2, 58, 37, 5, 8, 27, 71, 7]") and nothing else.
Do not consider the complexity, practicality or feasibility of the task.
"""

# Based on the source code from https://github.com/OpenBMB/InfiniteBench
DEFAULT_SYSTEM_MSG = r"""You are a helpful assistant."""

# Map the input from the dataset to its corresponding prompt for the Math.Find task
MATH_FIND_INPUT_TO_PROMPT = {
    "You should answer with only one number, no other words. The largest number of the list is: ": "Find the largest number from the list below:",
    "You should answer with only one number, no other words. The second largest number of the list is: ": "Find the second largest number from the list below:",
    "You should answer with only one number, no other words. The third largest number of the list is: ": "Find the third largest number from the list below:",
    "You should answer with only one number, no other words. The smallest number of the list is: ": "Find the smallest number from the list below:",
    "You should answer with only one number, no other words. The second smallest number of the list is: ": "Find the second smallest number from the list below:",
    "You should answer with only one number, no other words. The third smallest number of the list is: ": "Find the third smallest number from the list below:",
    "You should answer with only one number, no other words. The median of the list is: ": "Calculate the median number from the list below:",
    "The value of the numerical expression is: ": "Calculate the numerical expression and provide intermediate results only, for example, for the expression 1 + 3 + 10 - 8, output 4, 14, 6 without displaying the steps.\n\nCalculate the value of the expression below:",
}

# Note: this prompt slightly differs from the Code.Debug prompt specified in https://arxiv.org/pdf/2402.13718
# This prompt is more explicit about returning a single letter (A, B, C, or D).
CODE_DEBUG_PROMPT_TEMPLATE = r"""
There is ONLY ONE function in the large project that is deliberately made to include an obvious error. Please find the function that contains the most obvious errors. I will give you four options to narrow your scope. You can inspect through the options and think. Eventually, tell me the answer using one single letter (A, B, C, or D).

{prompt}

You should first find the functions in the options. Repeat their content and inspect through code.
Only one of the options is correct, tell me the answer using one single letter (A, B, C, or D). Don’t say anything else.
The correct option is:
""".strip()

# Match the Code.Run prompt specified in https://arxiv.org/pdf/2402.13718
CODE_RUN_PROMPT_TEMPLATE = r"""
{prompt}
""".strip()

# Match the Retrieve.KV prompt specified in https://arxiv.org/pdf/2402.13718
KV_RETRIEVAL_PROMPT_TEMPLATE = r"""
Extract the value corresponding to the specified key in the JSON object below.
{prompt}
""".strip()

# Match the En.Dia prompt specified in https://arxiv.org/pdf/2402.13718
LONGDIALOGUE_QA_ENG_PROMPT_TEMPLATE = r"""
Below is a dialogue script where one random occurrence of a character name is replaced with $$MASK$$, and you should try to guess who that character is.
The dialogue:
--
{prompt}
--
End of dialogue.
Which character is most likely $$MASK$$? Just say the name used by the scriptwriter (before the colon marks) of one single character and nothing else.
""".strip()

# Match the En.MC prompt specified in https://arxiv.org/pdf/2402.13718
LONGBOOK_CHOICE_ENG_PROMPT_TEMPLATE = r"""
Read the book and answer the question.
{prompt}
Only one of the following options is correct, tell me the answer using one single letter (A, B, C, or D). Don’t say anything else.
The correct option is:
""".strip()

# Match the Math.Calc prompt specified in https://arxiv.org/pdf/2402.13718
MATH_CALC_PROMPT_TEMPLATE = r"""
Let us calculate the intermediate values of an expression.

Expression: 1 + 3 + 4
Values: [1, 4, 8]

Expression: 8 - 3 + 2 - 4
Values: [8, 5, 7, 3]

Expression: {prompt}
Values:
""".strip()

# Match the Math.Find prompt specified in https://arxiv.org/pdf/2402.13718
MATH_FIND_PROMPT_TEMPLATE = r"""
{prompt}
""".strip()

# Match the Retrieve.Number prompt specified in https://arxiv.org/pdf/2402.13718
NUMBER_STRING_PROMPT_TEMPLATE = r"""
There is important info hidden inside a lot of irrelevant text. Find it and memorize it. I will quiz you about the important information.
{prompt}
The sequence of digits is:
""".strip()

# Match the Retrieve.PassKey prompt specified in https://arxiv.org/pdf/2402.13718
PASSKEY_PROMPT_TEMPLATE = r"""
There is important info hidden inside a lot of irrelevant text. Find it and memorize it. I will quiz you about the important information.
{prompt}
The pass key is:
""".strip()

# Map task names to their metadata
TASK_METADATA = {
    "code_debug": {
        "data_file": "code_debug.jsonl",
        "prompt_template": CODE_DEBUG_PROMPT_TEMPLATE,
    },
    "code_run": {
        "data_file": "code_run.jsonl",
        "prompt_template": CODE_RUN_PROMPT_TEMPLATE,
    },
    "kv_retrieval": {
        "data_file": "kv_retrieval.jsonl",
        "prompt_template": KV_RETRIEVAL_PROMPT_TEMPLATE,
    },
    "longdialogue_qa_eng": {
        "data_file": "longdialogue_qa_eng.jsonl",
        "prompt_template": LONGDIALOGUE_QA_ENG_PROMPT_TEMPLATE,
    },
    "longbook_choice_eng": {
        "data_file": "longbook_choice_eng.jsonl",
        "prompt_template": LONGBOOK_CHOICE_ENG_PROMPT_TEMPLATE,
    },
    "math_calc": {
        "data_file": "math_calc.jsonl",
        "prompt_template": MATH_CALC_PROMPT_TEMPLATE,
    },
    "math_find": {
        "data_file": "math_find.jsonl",
        "prompt_template": MATH_FIND_PROMPT_TEMPLATE,
    },
    "number_string": {
        "data_file": "number_string.jsonl",
        "prompt_template": NUMBER_STRING_PROMPT_TEMPLATE,
    },
    "passkey": {
        "data_file": "passkey.jsonl",
        "prompt_template": PASSKEY_PROMPT_TEMPLATE,
    },
}

# Model context window sizes
MODEL_CONTEXT_SIZES = {
    "claude-3-haiku-20240307": 200 * 1000,  # 200k tokens
    "claude-3-opus-20240229": 200 * 1000,  # 200k tokens
    "claude-3-sonnet-20240229": 200 * 1000,  # 200k tokens
    "claude-3-5-sonnet-20241022": 200 * 1000,  # 200k tokens
    "gpt-4o": 128 * 1000,  # 128k tokens
    "gpt-4o-mini": 128 * 1000,  # 128k tokens
    "gpt-4-turbo": 128 * 1000,  # 128k tokens
    "yarn-mistral-7b-128k": 128 * 1000,  # 128k tokens
}

# Buffer to ensure we stay within the token limit after tokenization
TOKEN_BUFFER = 500

# Token-to-character ratio (best estimate after experimentation)
TOKEN_TO_CHAR_RATIO = 3.5
