# pytest: skip-file
# mypy: ignore-errors

"""Validation functions used, for reference"""

from datasets import load_dataset
from inspect_ai.log import read_eval_log
from inspect_ai.scorer import SampleScore, Score

from inspect_evals.livebench import grouped_mean


def compare_logs(
    log_file_path,
    allowed_err=0.1,
    grouped_mean_args: dict = {"group": "category", "all_name": "livebench"},
    model_name: str | None = None,
):
    """Compare original scores for question id and model with scores from an Inspect log file."""
    eval_log = read_eval_log(log_file_path)
    if model_name:
        model = model_name
    else:
        _, model = eval_log.eval.model.split("/")
    print("================================================")
    print(f"Model: {model}")
    print(f"Sample size: {len(eval_log.samples)}")
    print(count_by_group(eval_log.samples, grouped_mean_args["group"]))
    print("================================================")
    log_dict = {eval_sample.id: eval_sample for eval_sample in eval_log.samples}

    q_ids_set = log_dict.keys()
    metric = grouped_mean(**grouped_mean_args)

    log_sample_scores = [
        SampleScore(
            score=eval_sample.scores["decide_scorer"],
            sample_id=eval_sample.id,
            sample_metadata=eval_sample.metadata,
        )
        for eval_sample in eval_log.samples
    ]

    ds = load_dataset("livebench/model_judgment", split="leaderboard")
    livebench_answers_ds = ds.filter(
        lambda row: row["model"] == model and row["question_id"] in q_ids_set
    )
    assert len(livebench_answers_ds) == len(q_ids_set)

    bench_sample_scores = [
        SampleScore(
            score=Score(value=row["score"]),
            sample_id=row["question_id"],
            sample_metadata=log_dict[row["question_id"]].metadata,
        )
        for row in livebench_answers_ds
    ]

    log_results = metric(log_sample_scores)
    bench_results = metric(bench_sample_scores)
    for group, score in log_results.items():
        score_diff = abs(score - bench_results[group])
        # if score_diff > allowed_err:
        print(
            f"Score diff {score_diff}, group: {group}, log score: {score}, bench_results: {bench_results[group]}"
        )


def count_by_group(eval_samples, group: str):
    """Count samples by a specified metadata group."""
    counts = {}
    for sample in eval_samples:
        if (
            not hasattr(sample, "metadata")
            or not sample.metadata
            or group not in sample.metadata
        ):
            continue

        group_value = sample.metadata[group]
        if group_value not in counts:
            counts[group_value] = 0
        counts[group_value] += 1

    return counts


# compare_logs(ALL_SAMPLES_LOG_GPT4omini)

# compare_logs(MATH_SAMPLES_GPT35, grouped_mean_args={"group": "task"})
