from typing import Any

from inspect_ai.dataset import Sample

from inspect_evals.lab_bench.record_to_sample_helpers import (
    UNCERTAIN_ANSWER_CHOICE,
    record_to_sample_base,
    record_to_sample_figqa,
    record_to_sample_protocolqa,
    record_to_sample_suppqa,
    record_to_sample_tableqa,
)

EXAMPLE_LITQA_RECORD: dict[str, Any] = {
    "question": "What effect does expression of the ATPase-deficient E251Q mutant of the Spindle E protein in silkworm cells have on the levels of the mature piwiRNAs piR1712 and piR2986?",
    "ideal": "Decrease",
    "distractors": ["Increase", "No effect"],
}


def test_record_to_sample_litqa():
    """Test that the LitQA record is mapped correctly to Inspect Sample"""
    example_sample = record_to_sample_base(EXAMPLE_LITQA_RECORD)
    assert example_sample == Sample(
        input="What effect does expression of the ATPase-deficient E251Q mutant of the Spindle E protein in silkworm cells have on the levels of the mature piwiRNAs piR1712 and piR2986?",
        choices=[
            "Decrease",
            "Increase",
            "No effect",
            UNCERTAIN_ANSWER_CHOICE,
        ],
        target="A",
    )


EXAMPLE_SUPPQA_RECORD: dict[str, Any] = {
    "paper-title": "Sequencing thousands of single-cell genomes with combinatorial indexing.",
    "source": "https://doi.org/10.1038/nmeth.4154",
    "question": "During the Second Sort and PCR Indexing step, which temperature and duration are required for the incubation to denature transposase for LAND ONLY samples?",
    "ideal": "55°C for 5 minutes",
    "distractors": [
        "72°C for 5 minutes",
        "68°C for 45 minutes",
        "55°C for 45 minutes",
    ],
}


def test_record_to_sample_suppqa():
    """Test that the SuppQA record is mapped correctly to Inspect Sample"""
    example_sample = record_to_sample_suppqa(EXAMPLE_SUPPQA_RECORD)
    assert example_sample == Sample(
        input="Paper title: Sequencing thousands of single-cell genomes with combinatorial indexing.\nDOI: https://doi.org/10.1038/nmeth.4154\nDuring the Second Sort and PCR Indexing step, which temperature and duration are required for the incubation to denature transposase for LAND ONLY samples?",
        choices=[
            "55°C for 5 minutes",
            "72°C for 5 minutes",
            "68°C for 45 minutes",
            "55°C for 45 minutes",
            UNCERTAIN_ANSWER_CHOICE,
        ],
        target="A",
    )


EXAMPLE_FIGQA_RECORD: dict[str, Any] = {
    "question": "According to figure D, which of the following groups has a higher average % total time freezing across three stimulations?",
    "ideal": "Off epoch",
    "distractors": ["On epoch", "Neither option", "Both are equal"],
    "figure": {"bytes": b"\xff\xd8\xff\xe0"},
}


def test_record_to_sample_figqa():
    """Test that the FigQA record is mapped correctly to Inspect Sample"""
    example_sample = record_to_sample_figqa(EXAMPLE_FIGQA_RECORD)
    assert example_sample.input[0].content[0].type == "image"
    assert example_sample.input[0].content[0].image == "data:image/jpeg;base64,/9j/4A=="
    assert (
        example_sample.input[0].content[1].text
        == "According to figure D, which of the following groups has a higher average % total time freezing across three stimulations?"
    )
    assert example_sample.choices == [
        "Off epoch",
        "On epoch",
        "Neither option",
        "Both are equal",
        UNCERTAIN_ANSWER_CHOICE,
    ]
    assert example_sample.target == "A"


EXAMPLE_TABLEQA_RECORD: dict[str, Any] = {
    "question": "How many LNA antisense oligonucleotides are in Phase I trials?",
    "ideal": "1",
    "distractors": ["2", "4", "5", "6"],
    "tables": [{"bytes": b"\x89PNG\r\n\x1a\n"}, {"bytes": b"\x89PNG\r\n\x1a\n"}],
}


def test_record_to_sample_tableqa():
    """Test that the FigQA record is mapped correctly to Inspect Sample"""
    example_sample = record_to_sample_tableqa(EXAMPLE_TABLEQA_RECORD)
    assert example_sample.input[0].content[0].type == "image"
    assert (
        example_sample.input[0].content[0].image == "data:image/png;base64,iVBORw0KGgo="
    )
    assert example_sample.input[0].content[1].type == "image"
    assert (
        example_sample.input[0].content[1].image == "data:image/png;base64,iVBORw0KGgo="
    )
    assert (
        example_sample.input[0].content[2].text
        == "How many LNA antisense oligonucleotides are in Phase I trials?"
    )
    assert example_sample.choices == [
        "1",
        "2",
        "4",
        "5",
        "6",
        UNCERTAIN_ANSWER_CHOICE,
    ]
    assert example_sample.target == "A"


EXAMPLE_DBQA_RECORD: dict[str, Any] = {
    "question": "Which of the following genes is associated with Currarino syndrome (Currarino triad) according to DisGeNet but not according to OMIM?",
    "ideal": "PCSK5",
    "distractors": ["MNX1", "MCTP1", "DMXL1"],
}


def test_record_to_sample_dbqa():
    """Test that the DbQA record is mapped correctly to Inspect Sample"""
    example_sample = record_to_sample_base(EXAMPLE_DBQA_RECORD)
    assert example_sample == Sample(
        input="Which of the following genes is associated with Currarino syndrome (Currarino triad) according to DisGeNet but not according to OMIM?",
        choices=[
            "PCSK5",
            "MNX1",
            "MCTP1",
            "DMXL1",
            UNCERTAIN_ANSWER_CHOICE,
        ],
        target="A",
    )


EXAMPLE_PROTOCOLQA_RECORD: dict[str, Any] = {
    "question": "You complete the listed protocol but notice that there are difficulties dissolving the RNA in step 7. Which of the following may address this issue?",
    "ideal": "The pellet should be dried for 30 minutes in step 6.",
    "distractors": [
        "The pellet should be resuspended in EtOH in step 7.",
        "The pellet should be resuspended in 2 mL of RNase free water in step 7.",
        "The Propan-2-ol should be heated to 37 °C in step 3",
    ],
    "protocol": """Step 1: Make up the volume of your RNA solution to 500μl with RNase free water.

Step 2: Add 50 μl of room temperature 3M Sodium Acetate (pH 8.5)

Step 3: Add 500 μl of room temperature Propan-2-ol.

Step 4: Pellet the precipitated RNA by centrifugation at 12,000 g for 15 mins.

Step 5: Note the pellet will appear clear and glassy after the supernatant has been removed.

Step 6: Allow the pellet to dry with the tube upside down on tissue at room temperature for 2 mins.

Step 7: Resuspend the pellet in RNase free water in an appropriate volume (usually 30-50μl)""",
}


def test_record_to_sample_protocolqa():
    """Test that the ProtocolQA record is mapped correctly to Inspect Sample"""
    example_sample = record_to_sample_protocolqa(EXAMPLE_PROTOCOLQA_RECORD)
    assert example_sample == Sample(
        input="""Protocol: Step 1: Make up the volume of your RNA solution to 500μl with RNase free water.

Step 2: Add 50 μl of room temperature 3M Sodium Acetate (pH 8.5)

Step 3: Add 500 μl of room temperature Propan-2-ol.

Step 4: Pellet the precipitated RNA by centrifugation at 12,000 g for 15 mins.

Step 5: Note the pellet will appear clear and glassy after the supernatant has been removed.

Step 6: Allow the pellet to dry with the tube upside down on tissue at room temperature for 2 mins.

Step 7: Resuspend the pellet in RNase free water in an appropriate volume (usually 30-50μl)

You complete the listed protocol but notice that there are difficulties dissolving the RNA in step 7. Which of the following may address this issue?""",
        choices=[
            "The pellet should be dried for 30 minutes in step 6.",
            "The pellet should be resuspended in EtOH in step 7.",
            "The pellet should be resuspended in 2 mL of RNase free water in step 7.",
            "The Propan-2-ol should be heated to 37 °C in step 3",
            UNCERTAIN_ANSWER_CHOICE,
        ],
        target="A",
    )


EXAMPLE_SEQQA_RECORD: dict[str, Any] = {
    "question": "What is the percent GC of the DNA sequence TAAGTAGTGCTAGAGATGAACCCTACCGGTTTGGCACCCATTCTAGCATGCTAGATC? Round to the nearest integer.",
    "ideal": "47",
    "distractors": ["46", "48", "24"],
}


def test_record_to_sample_seqqa():
    """Test that the SeqQA record is mapped correctly to Inspect Sample"""
    example_sample = record_to_sample_base(EXAMPLE_SEQQA_RECORD)
    assert example_sample == Sample(
        input="What is the percent GC of the DNA sequence TAAGTAGTGCTAGAGATGAACCCTACCGGTTTGGCACCCATTCTAGCATGCTAGATC? Round to the nearest integer.",
        choices=[
            "47",
            "46",
            "48",
            "24",
            UNCERTAIN_ANSWER_CHOICE,
        ],
        target="A",
    )


EXAMPLE_CLONING_SCENARIO_RECORD: dict[str, Any] = {
    "question": "I have a plasmid pLAB050 with sequence TCGGTCTCCAACGATGTGCTTCAGTATTACATTTTTTGCCTTCAACGCCTTGATTGTTCTATTTTTGCTAATAATAAATCTATTTCATCGGACTAAAAGTCCATTAGTTGTAAGCGGATTTAGCTCAGTTGGGAGAGCGCCAGACTGAAGAAAAACTTCGGTCAAGTCATCTGGAGGTCCTGTGTTCGATCCACAGAATTCGCAGATGGCCGGCATGGTCCCAGCCTCCTCGCTGGCGCCGGCTGGGCAACACCTTCGGGTGGCGAATGGGACTTGAGACGGAAAGTGAAACGTGATTTCATGCGTCATTTTGAACATTTTGTAAATCTTATTTAATAATGTGTGCGGCAATTCACATTTAATTTATGAATGTTTTCTTAACATCGCGGCAACTCAAGAAACGGCAGGTTCGGATCTTAGCTACTAGAGAAAGAGGAGAAATACTAGATGCGTAAAGGCGAAGAGCTGTTCACTGGTGTCGTCCCTATTCTGGTGGAACTGGATGGTGATGTCAACGGTCATAAGTTTTCCGTGCGTGGCGAGGGTGAAGGTGACGCAACTAATGGTAAACTGACGCTGAAGTTCATCTGTACTACTGGTAAACTGCCGGTTCCTTGGCCGACTCTGGTAACGACGCTGACTTATGGTGTTCAGTGCTTTGCTCGTTATCCGGACCATATGAAGCAGCATGACTTCTTCAAGTCCGCCATGCCGGAAGGCTATGTGCAGGAACGCACGATTTCCTTTAAGGATGACGGCACGTACAAAACGCGTGCGGAAGTGAAATTTGAAGGCGATACCCTGGTAAACCGCATTGAGCTGAAAGGCATTGACTTTAAAGAGGACGGCAATATCCTGGGCCATAAGCTGGAATACAATTTTAACAGCCACAATGTTTACATCACCGCCGATAAACAAAAAAATGGCATTAAAGCGAATTTTAAAATTCGCCACAACGTGGAGGATGGCAGCGTGCAGCTGGCTGATCACTACCAGCAAAACACTCCAATCGGTGATGGTCCTGTTCTGCTGCCAGACAATCACTATCTGAGCACGCAAAGCGTTCTGTCTAAAGATCCGAACGAGAAACGCGATCATATGGTTCTGCTGGAGTTCGTAACCGCAGCGGGCATCACGCATGGTATGGATGAACTGTACAAATGACCAGGCATCAAATAAAACGAAAGGCTCAGTCGAAAGACTGGGCCTTTCGTTTTATCTGTTGTTTGTCGGTGAACGCTCTCTACTAGAGTCACACTGGCTCACCTTCGGGTGGGCCTTTCTGCGTTTATACGTCTCAGTTTTAGAGCTAGAAATAGCAAGTTAAAATAAGGCTAGTCCGTTATCAACTTGAAAAAGTGGCACCGAGTCGGTGCTTTTTTTATTTTTTGTCACTATTGTTATGTAAAATGCCACCTCTGACAGTATGGAACGCAAACTTCTGTCTAGTGGATAGCTGTGAGACCAGACCAATAAAAAACGCCCGGCGGCAACCGAGCGTTCTGAACAAATCCAGATGGAGTTCTGAGGTCATTACTGGATCTATCAACAGGAGTCCAAGCGAGCTCGATATCAAATTACGCCCCGCCCTGCCACTCATCGCAGTACTGTTGTAATTCATTAAGCATTCTGCCGACATGGAAGCCATCACAAACGGCATGATGAACCTGAATCGCCAGCGGCATCAGCACCTTGTCGCCTTGCGTATAATATTTGCCCATGGTGAAAACGGGGGCGAAGAAGTTGTCCATATTGGCCACGTTTAAATCAAAACTGGTGAAACTCACCCAGGGATTGGCTGAAACGAAAAACATATTCTCAATAAACCCTTTAGGGAAATAGGCCAGGTTTTCACCGTAACACGCCACATCTTGCGAATATATGTGTAGAAACTGCCGGAAATCGTCGTGGTATTCACTCCAGAGCGATGAAAACGTTTCAGTTTGCTCATGGAAAACGGTGTAACAAGGGTGAACACTATCCCATATCACCAGCTCACCGTCTTTCATTGCCATACGAAATTCCGGATGAGCATTCATCAGGCGGGCAAGAATGTGAATAAAGGCCGGATAAAACTTGTGCTTATTTTTCTTTACGGTCTTTAAAAAGGCCGTAATATCCAGCTGAACGGTCTGGTTATAGGTACATTGAGCAACTGACTGAAATGCCTCAAAATGTTCTTTACGATGCCATTGGGATATATCAACGGTGGTATATCCAGTGATTTTTTTCTCCATTTTAGCTTCCTTAGCTCCTGAAAATCTCGATAACTCAAAAAATACGCCCGGTAGTGATCTTATTTCATTATGGTGAAAGTTGGAACCTCTTACGTGCCCGATCAATCATGACCAAAATCCCTTAACGTGAGTTTTCGTTCCACTGAGCGTCAGACCCCGTAGAAAAGATCAAAGGATCTTCTTGAGATCCTTTTTTTCTGCGCGTAATCTGCTGCTTGCAAACAAAAAAACCACCGCTACCAGCGGTGGTTTGTTTGCCGGATCAAGAGCTACCAACTCTTTTTCCGAAGGTAACTGGCTTCAGCAGAGCGCAGATACCAAATACTGTTCTTCTAGTGTAGCCGTAGTTAGGCCACCACTTCAAGAACTCTGTAGCACCGCCTACATACCTCGCTCTGCTAATCCTGTTACCAGTGGCTGCTGCCAGTGGCGATAAGTCGTGTCTTACCGGGTTGGACTCAAGACGATAGTTACCGGATAAGGCGCAGCGGTCGGGCTGAACGGGGGGTTCGTGCACACAGCCCAGCTTGGAGCGAACGACCTACACCGAACTGAGATACCTACAGCGTGAGCTATGAGAAAGCGCCACGCTTCCCGAAGGGAGAAAGGCGGACAGGTATCCGGTAAGCGGCAGGGTCGGAACAGGAGAGCGCACGAGGGAGCTTCCAGGGGGAAACGCCTGGTATCTTTATAGTCCTGTCGGGTTTCGCCACCTCTGACTTGAGCGTCGATTTTTGTGATGCTCGTCAGGGGGGCGGAGCCTATGGAAAAACGCCAGCAACGCGGCCTTTTTACGGTTCCTGGCCTTTTGCTGGCCTTTTGCTCACATGTTCTTTCCTGCGTTATCCCCTGATTCTGTGGATAACCGTAG. I also have two DNA oligos with sequences GACTTTCATCACTATCTTTTCCCCCG and AAACCGGGGGAAAAGATAGTGATGAA. I plan to anneal the two oligos and clone them into the plasmid with Golden Gate cloning. Which restriction enzyme(s) should I use to do so?",
    "ideal": "BsmBI",
    "distractors": ["BsaI", "BbsI", "SapI"],
}


def test_record_to_sample_cloning_scenario():
    """Test that the CloningScenario record is mapped correctly to Inspect Sample"""
    example_sample = record_to_sample_base(EXAMPLE_CLONING_SCENARIO_RECORD)
    assert example_sample == Sample(
        input="I have a plasmid pLAB050 with sequence TCGGTCTCCAACGATGTGCTTCAGTATTACATTTTTTGCCTTCAACGCCTTGATTGTTCTATTTTTGCTAATAATAAATCTATTTCATCGGACTAAAAGTCCATTAGTTGTAAGCGGATTTAGCTCAGTTGGGAGAGCGCCAGACTGAAGAAAAACTTCGGTCAAGTCATCTGGAGGTCCTGTGTTCGATCCACAGAATTCGCAGATGGCCGGCATGGTCCCAGCCTCCTCGCTGGCGCCGGCTGGGCAACACCTTCGGGTGGCGAATGGGACTTGAGACGGAAAGTGAAACGTGATTTCATGCGTCATTTTGAACATTTTGTAAATCTTATTTAATAATGTGTGCGGCAATTCACATTTAATTTATGAATGTTTTCTTAACATCGCGGCAACTCAAGAAACGGCAGGTTCGGATCTTAGCTACTAGAGAAAGAGGAGAAATACTAGATGCGTAAAGGCGAAGAGCTGTTCACTGGTGTCGTCCCTATTCTGGTGGAACTGGATGGTGATGTCAACGGTCATAAGTTTTCCGTGCGTGGCGAGGGTGAAGGTGACGCAACTAATGGTAAACTGACGCTGAAGTTCATCTGTACTACTGGTAAACTGCCGGTTCCTTGGCCGACTCTGGTAACGACGCTGACTTATGGTGTTCAGTGCTTTGCTCGTTATCCGGACCATATGAAGCAGCATGACTTCTTCAAGTCCGCCATGCCGGAAGGCTATGTGCAGGAACGCACGATTTCCTTTAAGGATGACGGCACGTACAAAACGCGTGCGGAAGTGAAATTTGAAGGCGATACCCTGGTAAACCGCATTGAGCTGAAAGGCATTGACTTTAAAGAGGACGGCAATATCCTGGGCCATAAGCTGGAATACAATTTTAACAGCCACAATGTTTACATCACCGCCGATAAACAAAAAAATGGCATTAAAGCGAATTTTAAAATTCGCCACAACGTGGAGGATGGCAGCGTGCAGCTGGCTGATCACTACCAGCAAAACACTCCAATCGGTGATGGTCCTGTTCTGCTGCCAGACAATCACTATCTGAGCACGCAAAGCGTTCTGTCTAAAGATCCGAACGAGAAACGCGATCATATGGTTCTGCTGGAGTTCGTAACCGCAGCGGGCATCACGCATGGTATGGATGAACTGTACAAATGACCAGGCATCAAATAAAACGAAAGGCTCAGTCGAAAGACTGGGCCTTTCGTTTTATCTGTTGTTTGTCGGTGAACGCTCTCTACTAGAGTCACACTGGCTCACCTTCGGGTGGGCCTTTCTGCGTTTATACGTCTCAGTTTTAGAGCTAGAAATAGCAAGTTAAAATAAGGCTAGTCCGTTATCAACTTGAAAAAGTGGCACCGAGTCGGTGCTTTTTTTATTTTTTGTCACTATTGTTATGTAAAATGCCACCTCTGACAGTATGGAACGCAAACTTCTGTCTAGTGGATAGCTGTGAGACCAGACCAATAAAAAACGCCCGGCGGCAACCGAGCGTTCTGAACAAATCCAGATGGAGTTCTGAGGTCATTACTGGATCTATCAACAGGAGTCCAAGCGAGCTCGATATCAAATTACGCCCCGCCCTGCCACTCATCGCAGTACTGTTGTAATTCATTAAGCATTCTGCCGACATGGAAGCCATCACAAACGGCATGATGAACCTGAATCGCCAGCGGCATCAGCACCTTGTCGCCTTGCGTATAATATTTGCCCATGGTGAAAACGGGGGCGAAGAAGTTGTCCATATTGGCCACGTTTAAATCAAAACTGGTGAAACTCACCCAGGGATTGGCTGAAACGAAAAACATATTCTCAATAAACCCTTTAGGGAAATAGGCCAGGTTTTCACCGTAACACGCCACATCTTGCGAATATATGTGTAGAAACTGCCGGAAATCGTCGTGGTATTCACTCCAGAGCGATGAAAACGTTTCAGTTTGCTCATGGAAAACGGTGTAACAAGGGTGAACACTATCCCATATCACCAGCTCACCGTCTTTCATTGCCATACGAAATTCCGGATGAGCATTCATCAGGCGGGCAAGAATGTGAATAAAGGCCGGATAAAACTTGTGCTTATTTTTCTTTACGGTCTTTAAAAAGGCCGTAATATCCAGCTGAACGGTCTGGTTATAGGTACATTGAGCAACTGACTGAAATGCCTCAAAATGTTCTTTACGATGCCATTGGGATATATCAACGGTGGTATATCCAGTGATTTTTTTCTCCATTTTAGCTTCCTTAGCTCCTGAAAATCTCGATAACTCAAAAAATACGCCCGGTAGTGATCTTATTTCATTATGGTGAAAGTTGGAACCTCTTACGTGCCCGATCAATCATGACCAAAATCCCTTAACGTGAGTTTTCGTTCCACTGAGCGTCAGACCCCGTAGAAAAGATCAAAGGATCTTCTTGAGATCCTTTTTTTCTGCGCGTAATCTGCTGCTTGCAAACAAAAAAACCACCGCTACCAGCGGTGGTTTGTTTGCCGGATCAAGAGCTACCAACTCTTTTTCCGAAGGTAACTGGCTTCAGCAGAGCGCAGATACCAAATACTGTTCTTCTAGTGTAGCCGTAGTTAGGCCACCACTTCAAGAACTCTGTAGCACCGCCTACATACCTCGCTCTGCTAATCCTGTTACCAGTGGCTGCTGCCAGTGGCGATAAGTCGTGTCTTACCGGGTTGGACTCAAGACGATAGTTACCGGATAAGGCGCAGCGGTCGGGCTGAACGGGGGGTTCGTGCACACAGCCCAGCTTGGAGCGAACGACCTACACCGAACTGAGATACCTACAGCGTGAGCTATGAGAAAGCGCCACGCTTCCCGAAGGGAGAAAGGCGGACAGGTATCCGGTAAGCGGCAGGGTCGGAACAGGAGAGCGCACGAGGGAGCTTCCAGGGGGAAACGCCTGGTATCTTTATAGTCCTGTCGGGTTTCGCCACCTCTGACTTGAGCGTCGATTTTTGTGATGCTCGTCAGGGGGGCGGAGCCTATGGAAAAACGCCAGCAACGCGGCCTTTTTACGGTTCCTGGCCTTTTGCTGGCCTTTTGCTCACATGTTCTTTCCTGCGTTATCCCCTGATTCTGTGGATAACCGTAG. I also have two DNA oligos with sequences GACTTTCATCACTATCTTTTCCCCCG and AAACCGGGGGAAAAGATAGTGATGAA. I plan to anneal the two oligos and clone them into the plasmid with Golden Gate cloning. Which restriction enzyme(s) should I use to do so?",
        choices=[
            "BsmBI",
            "BsaI",
            "BbsI",
            "SapI",
            UNCERTAIN_ANSWER_CHOICE,
        ],
        target="A",
    )
