import numpy as np
import pytest
from inspect_ai.scorer import SampleScore, Score

# Import directly from the utils module to avoid circular imports
from inspect_evals.air_bench.utils import accuracy_by_category


def test_accuracy_by_category_basic():
    """Test that accuracy_by_category correctly calculates scores for different categories."""
    # Create sample scores with different category indices
    sample_scores = [
        SampleScore(
            sample_id="1", score=Score(value=1.0), sample_metadata={"cate-idx": "1.2.3"}
        ),
        SampleScore(
            sample_id="2", score=Score(value=0.5), sample_metadata={"cate-idx": "1.2.4"}
        ),
        SampleScore(
            sample_id="3", score=Score(value=0.0), sample_metadata={"cate-idx": "2.1.1"}
        ),
    ]

    # Calculate scores using the metric
    result = accuracy_by_category()(sample_scores)

    # Expected results
    expected_categories = {
        "air_1": (1.0 + 0.5) / 2,  # Average of samples in category 1
        "air_1_2": (1.0 + 0.5) / 2,  # Average of samples in category 1.2
        "air_1_2_3": 1.0,  # Score of sample in category 1.2.3
        "air_1_2_4": 0.5,  # Score of sample in category 1.2.4
        "air_2": 0.0,  # Score of sample in category 2
        "air_2_1": 0.0,  # Score of sample in category 2.1
        "air_2_1_1": 0.0,  # Score of sample in category 2.1.1
        "all": (1.0 + 0.5 + 0.0) / 3,  # Overall average
    }

    # Verify that all expected categories are in the result
    for category, expected_score in expected_categories.items():
        assert category in result, f"Category {category} not found in result"
        assert np.isclose(result[category], expected_score), (
            f"Expected {expected_score} for {category}, got {result[category]}"
        )

    # Verify the overall score
    assert np.isclose(result["all"], 0.5), (
        f"Expected 0.5 for overall score, got {result['all']}"
    )


def test_accuracy_by_category_empty_list():
    """Test that accuracy_by_category handles empty list of scores."""
    result = accuracy_by_category()([])
    assert np.isnan(result["all"])
    assert len(result) == 1


def test_accuracy_by_category_missing_metadata():
    """Test that accuracy_by_category raises an error when metadata is missing."""
    # Create a sample score without metadata
    sample_scores = [
        SampleScore(sample_id="1", score=Score(value=1.0), sample_metadata=None)
    ]

    # Expect a ValueError when metadata is missing
    with pytest.raises(ValueError, match="Sample 1 has no metadata"):
        accuracy_by_category()(sample_scores)


def test_accuracy_by_category_same_category():
    """Test that accuracy_by_category correctly handles samples in the same category."""
    # Create sample scores with the same category index
    sample_scores = [
        SampleScore(
            sample_id="1", score=Score(value=0.8), sample_metadata={"cate-idx": "3.4.5"}
        ),
        SampleScore(
            sample_id="2", score=Score(value=0.6), sample_metadata={"cate-idx": "3.4.5"}
        ),
    ]

    # Calculate scores using the metric
    result = accuracy_by_category()(sample_scores)

    # Expected results
    expected_score = 0.7  # Average of 0.8 and 0.6
    expected_categories = {
        "air_3": expected_score,
        "air_3_4": expected_score,
        "air_3_4_5": expected_score,
        "all": expected_score,
    }

    # Verify results
    for category, expected in expected_categories.items():
        assert category in result
        assert np.isclose(result[category], expected), (
            f"Expected {expected} for {category}, got {result[category]}"
        )


def test_accuracy_by_category_mixed_scores():
    """Test that accuracy_by_category correctly handles mixed scores."""
    # Create sample scores with different values
    sample_scores = [
        SampleScore(
            sample_id="1", score=Score(value=1.0), sample_metadata={"cate-idx": "1.1.1"}
        ),
        SampleScore(
            sample_id="2", score=Score(value=0.0), sample_metadata={"cate-idx": "1.1.2"}
        ),
        SampleScore(
            sample_id="3", score=Score(value=0.5), sample_metadata={"cate-idx": "1.2.1"}
        ),
    ]

    # Calculate scores using the metric
    result = accuracy_by_category()(sample_scores)

    # Expected results for level 2 category
    assert np.isclose(result["air_1"], 0.5)  # (1.0 + 0.0 + 0.5) / 3

    # Expected results for level 3 categories
    assert np.isclose(result["air_1_1"], 0.5)  # (1.0 + 0.0) / 2
    assert np.isclose(result["air_1_2"], 0.5)  # Just one sample with score 0.5

    # Expected results for level 4 categories
    assert np.isclose(result["air_1_1_1"], 1.0)
    assert np.isclose(result["air_1_1_2"], 0.0)
    assert np.isclose(result["air_1_2_1"], 0.5)

    # Overall average
    assert np.isclose(result["all"], 0.5)  # (1.0 + 0.0 + 0.5) / 3
