import pytest
from inspect_ai.scorer import SampleScore, Score

from inspect_evals.scicode.metrics import (
    percentage_main_problems_solved,
    percentage_subproblems_solved,
    total_main_problems_solved,
    total_subproblems_solved,
)


def parametrize_scores():
    return pytest.mark.parametrize(
        "score_list, total_main, total_sub, percentage_main, percentage_sub",
        [
            (
                [SampleScore(score=Score(value={"1.1": 0}))],
                0,
                0,
                0.0,
                0.0,
            ),
            (
                [
                    SampleScore(score=Score(value={"1.1": 1, "1.2": 0})),
                    SampleScore(score=Score(value={"2.1": 1, "2.2": 1})),
                ],
                1,
                3,
                50.0,
                75.0,
            ),
            (
                [
                    SampleScore(score=Score(value={"1.1": 1})),
                    SampleScore(score=Score(value={"2.1": 1, "2.2": 1})),
                ],
                2,
                3,
                100.0,
                100.0,
            ),
        ],
    )


@parametrize_scores()
def test_total_main_problems_solved(
    score_list, total_main, total_sub, percentage_main, percentage_sub
):
    assert total_main_problems_solved()(score_list) == total_main


@parametrize_scores()
def test_total_subproblems_solved(
    score_list, total_main, total_sub, percentage_main, percentage_sub
):
    assert total_subproblems_solved()(score_list) == total_sub


@parametrize_scores()
def test_percentage_main_problems_solved(
    score_list, total_main, total_sub, percentage_main, percentage_sub
):
    assert percentage_main_problems_solved()(score_list) == percentage_main


@parametrize_scores()
def test_percentage_subproblems_solved(
    score_list, total_main, total_sub, percentage_main, percentage_sub
):
    assert percentage_subproblems_solved()(score_list) == percentage_sub
