from collections import defaultdict
from typing import Callable, Dict, List, cast

from inspect_ai.model import Model
from inspect_ai.scorer import (
    Metric,
    SampleScore,
    Score,
    Scorer,
    Target,
    Value,
    metric,
    model_graded_qa,
    scorer,
)
from inspect_ai.solver import TaskState


def return_metadata_variable_as_history(
    state: TaskState, metadata_variable: str = "needle_question"
) -> str:
    """
    Returns a specific variable from the metadata as the input through the include_history attribute of the model_graded_qa() function.

    Args:
        state (TaskState): The TaskState for the current Sample being evaluated.
        metadata_variable (str): The name of the metadata variable to use.

    Returns:
        str: The content of the specified metadata variable (converted to a string if a different type).
    """
    if metadata_variable not in state.metadata:
        raise KeyError(
            f"No variable with name `{metadata_variable}` found in metadata."
        )

    metadata_content = state.metadata[metadata_variable]

    return str(metadata_content)


@metric
def subset_accuracy_combinations() -> Metric:
    """
    Custom metric to calculate accuracy for subsets and combinations of experimental parameters (metadata values).

    This metric processes a list of `SampleScore` objects, grouping them by specific experimental parameter values
    (e.g., context length, position), as well as combinations of these parameters, and computes accuracy
    for each group. It also calculates the overall accuracy for the entire dataset.

    Returns:
        Metric: A callable function (`calculate`) that takes a list of `SampleScore` objects and returns a dictionary
                containing accuracy metrics for individual parameters, parameter combinations, and overall accuracy.
    """

    def calculate(scores: list[SampleScore]) -> Value:
        """
        Calculate accuracy metrics for individual and combinations of experimental parameter values, as well as the overall dataset.

        Args:
            scores (List[SampleScore]): A list of `SampleScore` objects, each containing metadata about context length,
                                        position, and the corresponding score value.

        Returns:
            Dict[str, float]: A dictionary containing accuracy metrics, where:
                - Keys represent specific groupings of context length and position (e.g., "context_length_10000_accuracy")
                  or combinations of both (e.g., "context_length_10000_position_5000_accuracy"), as well as "overall_accuracy".
                - Values represent the calculated accuracy for each group or overall.
        """
        # Define the experimental parameters, context length and position
        experimental_params = ["target_context_length", "target_position"]

        # Safely convert value to float
        def to_numeric(val: Value) -> float:
            if isinstance(val, int | float):
                return float(val)
            elif isinstance(val, bool):
                return float(val)
            raise ValueError(f"Cannot calculate accuracy for non-numeric value: {val}")

        # Group scores by individual context length and position values
        individual_scores: Dict[str, List[float]] = defaultdict(list)
        for sample_score in scores:
            metadata = sample_score.score.metadata

            if metadata is None:
                continue

            for key in experimental_params:
                value = metadata[key]
                individual_scores[f"{key}_{value}"].append(
                    to_numeric(sample_score.score.value)
                )

        # Group scores by combinations of context length and position values
        combined_scores: Dict[str, List[float]] = defaultdict(list)
        for sample_score in scores:
            metadata = sample_score.score.metadata

            if metadata is None:
                continue

            param_value_pairs = [
                (param, metadata[param]) for param in experimental_params
            ]
            key = "_".join(f"{param}_{value}" for param, value in param_value_pairs)

            combined_scores[key].append(to_numeric(sample_score.score.value))

        results: Dict[str, float] = {}

        # Calculate accuracy for individual values
        for group, group_scores in individual_scores.items():
            results[f"{group}_accuracy"] = sum(group_scores) / len(group_scores)

        # Calculate accuracy for combinations
        for group, group_scores in combined_scores.items():
            results[f"{group}_accuracy"] = sum(group_scores) / len(group_scores)

        # Calculate overall accuracy
        numeric_scores = [
            to_numeric(sample_score.score.value) for sample_score in scores
        ]
        results["overall_accuracy"] = sum(numeric_scores) / len(numeric_scores)

        return cast(Value, results)

    return cast(Metric, calculate)


@scorer(metrics=[subset_accuracy_combinations()])
def custom_model_graded_qa_with_history_scorer(
    include_history: Callable[[TaskState], str],
    judge_model: list[str | Model] | str | Model | None = None,
) -> Scorer:
    """
    A modified scoring function that extends the base scorer to include task metadata.

    This function calls the 'model_graded_qa' function to compute the initial score and then
    augments the returned score object with additional metadata from the task state.

    Args:
        include_history (Callable[[TaskState], str]): A callable that determines how the chat history is presented.
        judge_model (list[str | Model] | str | Model | None): The model to use for grading. Defaults to None and uses the active model.

    Returns:
        Scorer: A callable scoring function that uses the enhanced scoring logic.
    """
    custom_instructions = """
    Evaluate the answer based on the following criteria:
    Score 1: The answer is completely unrelated to the reference.
    Score 3: The answer has minor relevance but does not align with the reference.
    Score 5: The answer has moderate relevance but contains inaccuracies.
    Score 7: The answer aligns with the reference but has minor omissions.
    Score 10: The answer is completely accurate and aligns perfectly with the reference.
    Only respond with a numerical score.
    """

    base_score_fn = model_graded_qa(
        instructions=custom_instructions,
        grade_pattern=r"(\d+)",
        include_history=include_history,
        partial_credit=True,
        model=judge_model,
    )

    async def score(state: TaskState, target: Target) -> Score:
        base_score: Score = await base_score_fn(state, target)

        return Score(
            value=base_score.value,
            answer=base_score.answer,
            explanation=base_score.explanation,
            metadata=state.metadata if state.metadata is not None else {},
        )

    return score
