import pandas as pd
import pytest
from inspect_ai.scorer import Score

from inspect_evals.worldsense._utils import (
    compute_accuracy,
    compute_bias,
    preprocess_scores,
)
from inspect_evals.worldsense.worldsense import record_to_sample

# Adapted from the first records in the dataset
example_records = [
    {
        "Key": -176741083417243227,
        "problemname": "Compl.trivial",
        "tuple_ID": "shared_tuple_id_for_this_problem",
        "text": "Jeremy is enrolled in 3 courses per week: Jeremy takes philosophy before biology and finance before philosophy. \nChoose one of the following alternatives: (1) Jeremy takes philosophy in between finance and biology, (2) Jeremy takes philosophy outside of the time range between finance and biology, or (3) it is impossible to decide.\nThink carefully, and only respond with one of these possible options (1), (2), or (3).",
        "expectedresp": ["1", "2", "3"],
        "span": 2,
        "descr": "[(1, 2),(0, 1)]",
        "query": "(1, 0, 2)",
        "objects": ["finance", "philosophy", "biology"],
        "problemsize": 3,
        "skin": "courses",
        "qskin": "directQ2",
        "complexity": "Complexity_0",
        "query_len": 3,
        "goldresp_obfusc": "Mark",
    },
    {
        "Key": 1747235547611487721,
        "problemname": "Compl.trivial",
        "tuple_ID": "shared_tuple_id_for_this_problem",
        "text": "Jeremy is enrolled in 3 courses per week: Jeremy takes philosophy before biology and finance before philosophy. \nChoose one of the following alternatives: (1) Jeremy takes finance in between philosophy and biology, (2) Jeremy takes finance outside of the time range between philosophy and biology, or (3) it is impossible to decide.\nThink carefully, and only respond with one of these possible options (1), (2), or (3).",
        "expectedresp": ["1", "2", "3"],
        "span": 2,
        "descr": "[(1, 2),(0, 1)]",
        "query": "(0, 1, 2)",
        "objects": ["finance", "philosophy", "biology"],
        "problemsize": 3,
        "skin": "courses",
        "qskin": "directQ2",
        "complexity": "Complexity_0",
        "query_len": 3,
        "goldresp_obfusc": "Youssef",
    },
    {
        "Key": 1917042815647934077,
        "problemname": "Compl.trivial",
        "tuple_ID": "shared_tuple_id_for_this_problem",
        "text": "Jeremy is enrolled in 3 courses per week: Jeremy takes philosophy before biology and finance before philosophy. \nChoose one of the following alternatives: (1) Jeremy takes geography in between philosophy and biology, (2) Jeremy takes geography outside of the time range between philosophy and biology, or (3) it is impossible to decide.\nThink carefully, and only respond with one of these possible options (1), (2), or (3).",
        "expectedresp": ["1", "2", "3"],
        "span": 2,
        "descr": "[(1, 2),(0, 1)]",
        "query": "(-1, 1, 2)",
        "objects": ["finance", "philosophy", "biology"],
        "problemsize": 3,
        "skin": "courses",
        "qskin": "directQ2",
        "complexity": "Complexity_0",
        "query_len": 3,
        "goldresp_obfusc": "Yoda",
    },
]


@pytest.mark.parametrize(
    "record, expected_target",
    [
        (example_records[0], "1"),
        (example_records[1], "2"),
        (example_records[2], "3"),
    ],
)
def test_record_to_sample(record, expected_target):
    sample = record_to_sample(record)

    assert sample.input == record["text"]
    assert sample.choices == record["expectedresp"]
    assert sample.target == expected_target
    # assert sample.id == record["tuple_ID"]
    assert sample.metadata == {
        "tuple_ID": record["tuple_ID"],
        "problemname": record["problemname"],
        "problemsize": record["problemsize"],
    }


example_scores = [
    Score(
        value=0,
        answer="IMPOSSIBLE",
        explanation="IMPOSSIBLE",
        metadata={
            "tuple_ID": "fast_spec_r20_seed07_results_10-12-2023_06h28m50s_13726",
            "problemname": "Consist.normal",
            "problemsize": 4,
        },
    ),
    Score(
        value=0,
        answer="IMPOSSIBLE",
        explanation="IMPOSSIBLE",
        metadata={
            "tuple_ID": "fast_spec_r20_seed07_results_10-12-2023_06h28m50s_22690",
            "problemname": "Consist.trivial",
            "problemsize": 4,
        },
    ),
    Score(
        value=1.0,
        answer="1",
        explanation="(1)",
        metadata={
            "tuple_ID": "fast_spec_r20_seed07_results_10-12-2023_06h28m50s_25962",
            "problemname": "Compl.normal",
            "problemsize": 5,
        },
    ),
]


def test_preprocess_scores():
    scores = example_scores
    processed_scores_df = preprocess_scores(scores)
    assert isinstance(processed_scores_df, pd.DataFrame)
    assert not processed_scores_df.empty
    assert "value" in processed_scores_df.columns
    assert "bias" in processed_scores_df.columns
    assert "weight" in processed_scores_df.columns

    for col in ["value", "bias", "weight"]:
        assert processed_scores_df[col].dtype == float

    for score, processed_score in zip(scores, processed_scores_df.itertuples()):
        print(processed_score)
        assert score.metadata["tuple_ID"] == processed_score.tuple_ID
        assert score.metadata["problemname"] == processed_score.problemname
        assert score.metadata["problemsize"] == processed_score.problemsize
        assert -1 <= processed_score.bias <= 1
        assert 0 <= processed_score.weight <= 1
        assert 0 <= processed_score.value <= 1


def test_preprocess_scores_grouping():
    scores = [
        Score(
            value="1",
            answer="1",  # weight = 0.25, bias = 1
            metadata={"tuple_ID": 1, "problemname": "Infer.trivial", "problemsize": 3},
        ),
        Score(
            value="0",
            answer="1",  # weight = 0.25, bias = 1
            metadata={"tuple_ID": 1, "problemname": "Infer.trivial", "problemsize": 3},
        ),
    ]
    processed_scores_df = preprocess_scores(scores)

    # Check that the grouped scores match the expected values
    #    tuple_ID    problemname  problemsize  value  bias  weight
    # 0         1  Infer.trivial            3    0.5   1.0     0.5
    expected_data = {
        "tuple_ID": [1],
        "problemname": ["Infer.trivial"],
        "problemsize": [3],
        "value": [0.5],  # (1 * 0.25 + 0 * 0.25) / group weight 0.5
        "bias": [1.0],  # (1 * 0.25 + 1 * 0.25) / group weight 0.5
        "weight": [0.5],  # Sum of wieghts: 0.25 + 0.25
    }
    expected_df = pd.DataFrame(expected_data)
    assert expected_df.equals(processed_scores_df), "DataFrames are not equal"


def test_compute_accuracy():
    # NOTE: Unsure of what constitutes a valid combination of tuple_ID, problemname, and problemsize
    scores = [
        Score(
            value="1",
            answer="1",
            metadata={"tuple_ID": 1, "problemname": "Infer.trivial", "problemsize": 1},
        ),
        Score(
            value="1",
            answer="1",
            metadata={"tuple_ID": 1, "problemname": "Infer.trivial", "problemsize": 3},
        ),
        Score(
            value="0",
            answer="2",
            metadata={"tuple_ID": 1, "problemname": "Infer.trivial", "problemsize": 3},
        ),
        Score(
            value="1",
            answer="3",
            metadata={"tuple_ID": 2, "problemname": "Compl.normal", "problemsize": 2},
        ),
        Score(
            value="0",
            answer="TRUE",
            metadata={"tuple_ID": 2, "problemname": "Compl.normal", "problemsize": 2},
        ),
    ]
    expected_preprocessed_scores = pd.DataFrame(
        {
            "tuple_ID": [1, 1, 2],
            "problemname": ["Infer.trivial", "Infer.trivial", "Compl.normal"],
            "problemsize": [1, 3, 2],
            "value": [1.0, 0.5, 0.5],
            "bias": [1.0, 1.0, 0.0],
            "weight": [0.25, 0.5, 1.0],
        }
    )
    df = preprocess_scores(scores)
    assert df.equals(expected_preprocessed_scores), "DataFrames are not equal"

    accuracy = compute_accuracy(df)
    assert isinstance(accuracy, float)
    assert 0 <= accuracy <= 1
    assert accuracy == 0.625


def test_compute_bias():
    scores = [
        Score(
            value="1",
            answer="1",
            metadata={"tuple_ID": 1, "problemname": "Infer.trivial", "problemsize": 3},
        ),
        Score(
            value="0",
            answer="2",
            metadata={"tuple_ID": 1, "problemname": "Infer.trivial", "problemsize": 3},
        ),
        Score(
            value="1",
            answer="3",
            metadata={"tuple_ID": 2, "problemname": "Compl.normal", "problemsize": 2},
        ),
        Score(
            value="0",
            answer="TRUE",
            metadata={"tuple_ID": 2, "problemname": "Compl.normal", "problemsize": 2},
        ),
    ]
    df = preprocess_scores(scores)
    bias = compute_bias(df)
    assert isinstance(bias, float)
    assert -1 <= bias <= 1
    assert bias == 0.5
