from typing import Any, cast

from inspect_ai.scorer import (
    Metric,
    SampleScore,
    metric,
)


def _count_main_problems_solved(scores: list[SampleScore]) -> int:
    return len(
        [
            sample_score
            for sample_score in scores
            if all(
                value == 1
                for value in cast(dict[str, Any], sample_score.score.value).values()
            )
        ]
    )


@metric
def total_main_problems_solved() -> Metric:
    """Return function to compute total number of _main_ problems solved successfully."""
    return _count_main_problems_solved


@metric
def percentage_main_problems_solved() -> Metric:
    """Return function to compute percentage of _main_ problems solved successfully."""

    def metric(scores: list[SampleScore]) -> float:
        count = _count_main_problems_solved(scores)
        total = len(scores)
        percentage = (count / total) * 100
        return percentage

    return metric


def _count_subproblems_solved(scores: list[SampleScore]) -> int:
    return len(
        [
            v
            for sample_score in scores
            for k, v in cast(dict[str, Any], sample_score.score.value).items()
            if v == 1
        ]
    )


@metric
def total_subproblems_solved() -> Metric:
    """Return function to compute total number of _main_ problems solved successfully."""
    return _count_subproblems_solved


@metric
def percentage_subproblems_solved() -> Metric:
    """Return function to compute percentage of subproblems solved successfully."""

    def metric(scores: list[SampleScore]) -> float:
        count = _count_subproblems_solved(scores)
        total = sum(
            len(cast(dict[str, Any], sample_score.score.value))
            for sample_score in scores
        )
        return (count / total) * 100

    return metric
