from inspect_evals.livebench import record_to_sample

# Coding
EXAMPLE_CODING_RECORD_LCB = {
    "question_id": "03741bcada4e970eede",
    "category": "coding",
    "task": "LCB_generation",
    "turns": [
        "### Instructions: You are an expert Python programmer. You will be given a question (problem specification) and will generate a correct Python program that matches the specification and passes all tests..."
    ],
    "question_title": "atcoder_Christmas Present",
    "public_test_cases": [
        {"input": "300 100\n", "output": "Bat\n", "testtype": "stdin"},
        {"input": "334 343\n", "output": "Glove\n", "testtype": "stdin"},
    ],
    "private_test_cases": "eJxrYJnay8QABhFtQEZ0tVJmXkFpiZKVgpKhgYGBgmFMnpKOglJ+aQlU1CmxBCJUklpcUlJZkAoSLC5JycxTqtVRQNJubmqgYG5BtnYTE1MFEyMzsvWbmlkoGJsYkK3fzNxEwcTMnGz9hgqg8EPX7p6TX5ZKnAHGQBOMjCkwwNTISMHU3JQCE0xMzBXMjCwpMQHoC1MLEoMhdooeAPgnxbY=",
    "original_json": {
        "question_title": "Christmas Present",
        "question_content": "Takahashi, a young baseball enthusiast, has been a very good boy this year, so Santa has decided to give him a bat or a glove, whichever is more expensive.\nIf a bat costs B yen and a glove costs G yen (B\\neq G), which one will Santa give to Takahashi?\n\nInput\n\nThe input is given from Standard Input in the following format:\nB G\n\nOutput\n\nIf Santa gives Takahashi a bat, print Bat; if Santa gives him a glove, print Glove.\n\nConstraints\n\n\n- B and G are different integers between 1 and 1000, inclusive.\n\nSample Input 1\n\n300 100\n\nSample Output 1\n\nBat\n\nThe bat is more expensive than the glove, so Santa will give Takahashi the bat.\n\nSample Input 2\n\n334 343\n\nSample Output 2\n\nGlove\n\nThe glove is more expensive than the bat, so Santa will give Takahashi the glove.",
        "platform": "atcoder",
        "question_id": "abc334_a",
        "contest_id": "abc334",
        "contest_date": "2023-12-23T00:00:00",
        "starter_code": "",
        "difficulty": "easy",
        "metadata": "{}",
    },
    "release_date": "2023-12-23T00:00:00",
    "livebench_release_date": "2024-07-26T00:00:00",
    "livebench_removal_date": "",
}


def test_record_to_sample_coding():
    """Test that the record_to_sample function works for the coding category.."""
    sample = record_to_sample(EXAMPLE_CODING_RECORD_LCB)

    assert sample.id == EXAMPLE_CODING_RECORD_LCB["question_id"]
    assert sample.input == EXAMPLE_CODING_RECORD_LCB["turns"][0]
    assert not sample.target

    assert sample.metadata["category"] == "coding"
    assert sample.metadata["task"] == EXAMPLE_CODING_RECORD_LCB["task"]
    assert (
        sample.metadata["livebench_release_date"]
        == EXAMPLE_CODING_RECORD_LCB["livebench_release_date"]
    )
    assert (
        sample.metadata["livebench_removal_date"]
        == EXAMPLE_CODING_RECORD_LCB["livebench_removal_date"]
    )
    assert (
        sample.metadata["coding"]["question_title"]
        == EXAMPLE_CODING_RECORD_LCB["question_title"]
    )
    assert (
        sample.metadata["coding"]["public_test_cases"]
        == EXAMPLE_CODING_RECORD_LCB["public_test_cases"]
    )
    assert (
        sample.metadata["coding"]["private_test_cases"]
        == EXAMPLE_CODING_RECORD_LCB["private_test_cases"]
    )
    assert (
        sample.metadata["coding"]["original_json"]
        == EXAMPLE_CODING_RECORD_LCB["original_json"]
    )
    assert not sample.metadata["coding"]["partial_solution"]

    assert not sample.metadata["instruction_following"]["instruction_id_list"]
    assert not sample.metadata["instruction_following"]["kwargs"]


# Instruction following
EXAMPLE_IF_RECORD = {
    "question_id": "cdd5cee29c50448a74db8085ff4d9",
    "category": "instruction_following",
    "task": "summarize",
    "instruction_id_list": [
        "detectable_format:title",
        "startend:end_checker",
        "length_constraints:number_paragraphs",
    ],
    "turns": [
        "The following are the beginning sentences of a news article from the Guardian.\n-------\n Half-submerged by a jungle and populated by cows, chickens and goats, it is one of the more unusual tourist attractions on the South Pacific island of Rarotonga: the “cursed” Sheraton hotel complex, a long-abandoned luxury resort that has lain derelict for a quarter of a century. The crumbling, graffiti-covered buildings are overgrown with creepers. The only human visitors to the 70-acre site in Vaimaanga pay a small fee to explore what has become known as the biggest white elephant in the South Pacific. The Italian-backed hotel project, which began in the late 1980s, was designed to revolutionise tourism in the remote Cook Islands by providing the tiny nation with its first five-star resort. But the house of cards collapsed at the 11th hour, with the project 80% complete, amid allegations of mafia involvement and money laundering back in Italy. The Cook Islands government was left with a debt of $120m and was almost bankrupted. In the years since, repeated attempts to finish the project have foundered, but now a new effort to salvage the buildings is under way, and could have an even bigger impact on the tiny Pacific state than the original plans were expected to. Tim Tepaki, a Cook Islander who has been involved in property development in New Zealand, claims to have hundreds of millions of dollars in Chinese backing. He hopes to rebuild the complex and turn the Pacific backwater into a favourite destination for tourists from China. Until now tourism in the Cooks – a nation of 15 islands scattered across an area the size of India – has relied mainly on visitors from Australia and New Zealand, who tend to focus on the capital Rarotonga and the atoll of Aitutaki, a 50-minute flight north.\n-------\nPlease summarize based on the sentences provided. Your answer must contain a title, wrapped in double angular brackets, such as <<poem of joy>>. Finish your response with this exact phrase Any other questions?. No other words should follow this phrase. There should be 4 paragraphs. Paragraphs are separated with the markdown divider: ***"
    ],
    "instructions_id_list": [
        "detectable_format:title",
        "startend:end_checker",
        "length_constraints:number_paragraphs",
    ],
    "kwargs": [
        {
            "num_sentences": None,
            "relation": None,
            "section_spliter": None,
            "num_sections": None,
            "keywords": None,
            "num_words": None,
            "num_bullets": None,
            "forbidden_words": None,
            "end_phrase": None,
            "num_paragraphs": None,
            "nth_paragraph": None,
            "first_word": None,
            "postscript_marker": None,
            "prompt_to_repeat": None,
        },
        {
            "num_sentences": None,
            "relation": None,
            "section_spliter": None,
            "num_sections": None,
            "keywords": None,
            "num_words": None,
            "num_bullets": None,
            "forbidden_words": None,
            "end_phrase": "Any other questions?",
            "num_paragraphs": None,
            "nth_paragraph": None,
            "first_word": None,
            "postscript_marker": None,
            "prompt_to_repeat": None,
        },
        {
            "num_sentences": None,
            "relation": None,
            "section_spliter": None,
            "num_sections": None,
            "keywords": None,
            "num_words": None,
            "num_bullets": None,
            "forbidden_words": None,
            "end_phrase": None,
            "num_paragraphs": 4,
            "nth_paragraph": None,
            "first_word": None,
            "postscript_marker": None,
            "prompt_to_repeat": None,
        },
    ],
    "task_prompt": "Please summarize based on the sentences provided.",
    "release_date": "2023-12-23T00:00:00",
    "livebench_release_date": "2024-06-24T00:00:00",
    "livebench_removal_date": "",
}


def test_record_to_sample_instruction_following():
    """Test that the record_to_sample function works for the instruction following category.."""
    sample = record_to_sample(EXAMPLE_IF_RECORD)

    assert sample.id == EXAMPLE_IF_RECORD["question_id"]
    assert sample.input == EXAMPLE_IF_RECORD["turns"][0]
    assert not sample.target

    assert sample.metadata["category"] == "instruction_following"
    assert sample.metadata["task"] == EXAMPLE_IF_RECORD["task"]

    assert (
        sample.metadata["instruction_following"]["instruction_id_list"]
        == EXAMPLE_IF_RECORD["instructions_id_list"]
    )
    assert sample.metadata["instruction_following"]["kwargs"] == [
        {},
        {"end_phrase": "Any other questions?"},
        {"num_paragraphs": 4},
    ]

    assert not sample.metadata["coding"]["question_title"]
    assert not sample.metadata["coding"]["partial_solution"]
    assert not sample.metadata["coding"]["public_test_cases"]
    assert not sample.metadata["coding"]["private_test_cases"]
    assert not sample.metadata["coding"]["original_json"]


EXAMPLE_REASONING_RECORD = {
    "question_id": "0daa7ca38beec4441",
    "ground_truth": "no, yes, yes",
    "category": "reasoning",
    "task": "web_of_lies_v2",
    "turns": [
        "In this question, assume each person either always tells the truth or always lies. Tala is at the movie theater. The person at the restaurant says the person at the aquarium lies. Ayaan is at the aquarium. Ryan is at the botanical garden. The person at the park says the person at the art gallery lies. The person at the museum tells the truth. Zara is at the museum. Jake is at the art gallery. The person at the art gallery says the person at the theater lies. Beatriz is at the park. The person at the movie theater says the person at the train station lies. Nadia is at the campground. The person at the campground says the person at the art gallery tells the truth. The person at the theater lies. The person at the amusement park says the person at the aquarium tells the truth. Grace is at the restaurant. The person at the aquarium thinks their friend is lying. Nia is at the theater. Kehinde is at the train station. The person at the theater thinks their friend is lying. The person at the botanical garden says the person at the train station tells the truth. The person at the aquarium says the person at the campground tells the truth. The person at the aquarium saw a firetruck. The person at the train station says the person at the amusement park lies. Mateo is at the amusement park. Does the person at the train station tell the truth? Does the person at the amusement park tell the truth? Does the person at the aquarium tell the truth? Think step by step, and then put your answer in **bold** as a list of three words, yes or no (for example, **yes, no, yes**). If you don't know, guess."
    ],
    "livebench_release_date": "2024-06-24T00:00:00",
    "livebench_removal_date": "",
}


def test_record_to_sample_others():
    """Test that the record_to_sample function works for the standard format"""
    sample = record_to_sample(EXAMPLE_REASONING_RECORD)

    assert sample.id == EXAMPLE_REASONING_RECORD["question_id"]
    assert sample.input == EXAMPLE_REASONING_RECORD["turns"][0]
    assert sample.target == EXAMPLE_REASONING_RECORD["ground_truth"]

    assert sample.metadata["category"] == "reasoning"
    assert sample.metadata["task"] == EXAMPLE_REASONING_RECORD["task"]

    assert not sample.metadata["coding"]["question_title"]
    assert not sample.metadata["coding"]["partial_solution"]
    assert not sample.metadata["coding"]["public_test_cases"]
    assert not sample.metadata["coding"]["private_test_cases"]
    assert not sample.metadata["coding"]["original_json"]

    assert not sample.metadata["instruction_following"]["instruction_id_list"]
    assert not sample.metadata["instruction_following"]["kwargs"]
