import pandas as pd
from inspect_ai.scorer import Score, ValueToFloat, value_to_float


def preprocess_scores(
    scores: list[Score], to_float: ValueToFloat = value_to_float()
) -> pd.DataFrame:
    """Preprocesses a list of Score objects into a DataFrame with computed weights and biases."""
    processed_scores = [
        {
            "value": to_float(score.value),
            "answer": score.answer,
            "tuple_ID": score.metadata.get("tuple_ID"),
            "problemname": score.metadata.get("problemname"),
            "problemsize": score.metadata.get("problemsize"),
        }
        for score in scores
        if score.metadata
        and None
        not in (
            to_float(score.value),
            score.answer,
            score.metadata.get("tuple_ID"),
            score.metadata.get("problemname"),
            score.metadata.get("problemsize"),
        )
    ]

    score_df = pd.DataFrame(processed_scores)

    # Mappings for bias and weight
    bias_mapping = {
        "1": 1,
        "2": 1,
        "TRUE": 1,
        "POSSIBLE": 1,
        "3": -1,
        "FALSE": -1,
        "IMPOSSIBLE": -1,
    }
    weight_mapping = {
        "1": 0.25,
        "2": 0.25,
        "3": 0.5,
        "TRUE": 0.5,
        "POSSIBLE": 0.5,
        "FALSE": 0.5,
        "IMPOSSIBLE": 0.5,
    }

    # Calculate weighted values and biases
    score_df["weight"] = score_df["answer"].map(weight_mapping).astype(float)
    score_df["bias"] = (
        score_df["answer"].map(bias_mapping).astype(float) * score_df["weight"]
    )
    score_df["value"] = score_df["value"].astype(float) * score_df["weight"]

    # Group and normalize
    grouped_scores = (
        score_df.groupby(["tuple_ID", "problemname", "problemsize"])
        .agg({"value": "sum", "bias": "sum", "weight": "sum"})
        .reset_index()
    )

    grouped_scores["value"] = grouped_scores["value"] / grouped_scores["weight"].where(
        grouped_scores["weight"] != 0, 1
    )
    grouped_scores["bias"] = grouped_scores["bias"] / grouped_scores["weight"].where(
        grouped_scores["weight"] != 0, 1
    )

    return grouped_scores


def compute_accuracy(grouped_scores: pd.DataFrame) -> float:
    """Compute the weighted accuracy from preprocessed scores."""
    problem_summary = (
        grouped_scores.groupby(["problemname", "problemsize"])
        .agg({"value": "mean"})
        .reset_index()
    )
    final_summary = (
        problem_summary.groupby("problemname").agg({"value": "mean"}).reset_index()
    )
    return float(final_summary["value"].mean())


def compute_bias(grouped_scores: pd.DataFrame) -> float:
    """Compute the weighted bias from preprocessed scores."""
    problem_summary = (
        grouped_scores.groupby(["problemname", "problemsize"])
        .agg({"bias": "mean"})
        .reset_index()
    )
    final_summary = (
        problem_summary.groupby("problemname").agg({"bias": "mean"}).reset_index()
    )
    return float(final_summary["bias"].mean())
