from pathlib import Path

import pytest

from src.models import EvalResult, FullEvalResult

cur_fp = Path(__file__)


# Ref: https://github.com/AIR-Bench/AIR-Bench/blob/4b27b8a8f2047a963805fcf6fb9d74be51ec440c/docs/available_tasks.md
# 24.05
# | Task | dev | test |
# | ---- | --- | ---- |
# | Long-Doc | 4 | 11 |
# | QA | 54 | 53 |
#
# 24.04
# | Task | test |
# | ---- | ---- |
# | Long-Doc | 15 |
# | QA | 13 |
NUM_QA_BENCHMARKS_24_05 = 53
NUM_DOC_BENCHMARKS_24_05 = 11
NUM_QA_BENCHMARKS_24_04 = 13
NUM_DOC_BENCHMARKS_24_04 = 15


def test_eval_result():
    EvalResult(
        eval_name="eval_name",
        retrieval_model="bge-m3",
        reranking_model="NoReranking",
        results=[{"domain": "law", "lang": "en", "dataset": "lex_files_500K-600K", "value": 0.45723}],
        task="qa",
        metric="ndcg_at_3",
        timestamp="2024-05-14T03:09:08Z",
        revision="1e243f14bd295ccdea7a118fe847399d",
        is_anonymous=True,
    )


@pytest.mark.parametrize(
    "file_path",
    [
        "AIR-Bench_24.04/bge-m3/jina-reranker-v2-base-multilingual/results.json",
        "AIR-Bench_24.05/bge-m3/NoReranker/results.json",
    ],
)
def test_full_eval_result_init_from_json_file(file_path):
    json_fp = cur_fp.parents[1] / "toydata/eval_results/" / file_path
    full_eval_result = FullEvalResult.init_from_json_file(json_fp)
    assert json_fp.parents[0].stem == full_eval_result.reranking_model
    assert json_fp.parents[1].stem == full_eval_result.retrieval_model
    assert len(full_eval_result.results) == 70


@pytest.mark.parametrize(
    "file_path, task, expected_num_results",
    [
        ("AIR-Bench_24.04/bge-m3/jina-reranker-v2-base-multilingual/results.json", "qa", NUM_QA_BENCHMARKS_24_04),
        (
            "AIR-Bench_24.04/bge-m3/jina-reranker-v2-base-multilingual/results.json",
            "long-doc",
            NUM_DOC_BENCHMARKS_24_04,
        ),
        ("AIR-Bench_24.05/bge-m3/NoReranker/results.json", "qa", NUM_QA_BENCHMARKS_24_05),
        ("AIR-Bench_24.05/bge-m3/NoReranker/results.json", "long-doc", NUM_DOC_BENCHMARKS_24_05),
    ],
)
def test_full_eval_result_to_dict(file_path, task, expected_num_results):
    json_fp = cur_fp.parents[1] / "toydata/eval_results/" / file_path
    full_eval_result = FullEvalResult.init_from_json_file(json_fp)
    result_dict_list = full_eval_result.to_dict(task)
    assert len(result_dict_list) == 1
    result = result_dict_list[0]
    attr_list = frozenset(
        [
            "eval_name",
            "Retrieval Method",
            "Reranking Model",
            "Retrieval Model LINK",
            "Reranking Model LINK",
            "Revision",
            "Submission Date",
            "Anonymous Submission",
        ]
    )
    result_cols = list(result.keys())
    assert len(result_cols) == (expected_num_results + len(attr_list))