import numpy as np

from inspect_evals.core_bench.utils import (
    calculate_prediction_intervals,
    categorize_keys,
    check_numeric_answer,
    clean_agent_results,
    count_questions,
    strip_keys,
)


def test_categorize_keys(sample_gt_result):
    """Test categorize_keys function correctly identifies key types."""
    numeric_keys, list_keys, string_keys = categorize_keys(sample_gt_result)

    assert numeric_keys == [
        "numeric_question1",
        "numeric_question2",
        "fig_numeric_question",
    ]
    assert list_keys == ["list_question"]
    assert string_keys == ["string_question"]


def test_count_questions():
    """Test count_questions correctly counts written and vision questions."""
    numeric_keys = ["numeric_q1", "numeric_q2", "fig_numeric_q"]
    list_keys = ["list_q", "fig_list_q"]
    string_keys = ["string_q", "fig_string_q"]

    total_written, total_vision = count_questions(numeric_keys, list_keys, string_keys)

    assert total_written == 4  # numeric_q1, numeric_q2, list_q, string_q
    assert total_vision == 3  # fig_numeric_q, fig_list_q, fig_string_q


def test_clean_agent_results():
    """Test clean_agent_results properly handles different value formats."""
    # Fixed sample data with only string or numeric values, no lists
    agent_result = {
        "numeric_question1": "5.2",
        "numeric_question2": "9.8%",
        "fig_numeric_question": "15.2",
        "string_question": "Hello",
    }

    cleaned = clean_agent_results(agent_result)

    assert cleaned["numeric_question1"] == 5.2
    assert cleaned["numeric_question2"] == 9.8  # % sign removed
    assert cleaned["fig_numeric_question"] == 15.2
    assert (
        cleaned["string_question"] == "Hello"
    )  # Failed conversion to float, stays as is


def test_check_numeric_answer():
    """Test check_numeric_answer correctly determines if a value is within range."""
    # Within range
    assert check_numeric_answer(5.2, (5.0, 5.5)) is True

    # At lower bound
    assert check_numeric_answer(5.0, (5.0, 5.5)) is True

    # At upper bound
    assert check_numeric_answer(5.5, (5.0, 5.5)) is True

    # Below range
    assert check_numeric_answer(4.9, (5.0, 5.5)) is False

    # Above range
    assert check_numeric_answer(5.6, (5.0, 5.5)) is False

    # String number representations
    assert check_numeric_answer("5.2", (5.0, 5.5)) is True
    assert check_numeric_answer("5", (5.0, 5.5)) is True
    assert check_numeric_answer("4.9", (5.0, 5.5)) is False

    # Invalid string values
    assert check_numeric_answer("five point two", (5.0, 5.5)) is False
    assert check_numeric_answer("", (5.0, 5.5)) is False

    # Non-numeric values
    assert check_numeric_answer(None, (5.0, 5.5)) is False
    assert check_numeric_answer([5.2], (5.0, 5.5)) is False
    assert check_numeric_answer({"value": 5.2}, (5.0, 5.5)) is False


def test_calculate_prediction_intervals(sample_gt_result):
    """Test calculate_prediction_intervals correctly computes intervals."""
    numeric_keys = ["numeric_question1", "numeric_question2", "fig_numeric_question"]
    intervals = calculate_prediction_intervals(sample_gt_result, numeric_keys)

    # Check that all expected keys are present
    assert set(intervals.keys()) == set(numeric_keys)

    # Check interval properties for each key
    for key in numeric_keys:
        lower, upper = intervals[key]
        values = [record[key] for record in sample_gt_result]
        mean = np.mean(values)

        # Mean should be within the interval
        assert lower <= mean <= upper

        # Lower bound should be less than upper bound
        assert lower < upper


def test_strip_keys():
    """Test strip_keys removes trailing punctuation from dictionary keys."""
    sample_dict = {
        "question?": 1,
        "another_question.": "test",
        "no_punctuation": [1, 2, 3],
    }
    expected = {
        "question": 1,
        "another_question": "test",
        "no_punctuation": [1, 2, 3],
    }
    # strip_keys should remove any trailing punctuation characters.
    assert strip_keys(sample_dict) == expected
