"""RS-fast item-id files (T0.3) and, when the pinned parquet is on disk, the lm-eval loaders."""
import json
import os
import sys
from pathlib import Path

import pytest

REPO = Path(__file__).resolve().parent.parent
SUBSETS = REPO / "configs" / "eval" / "subsets"
DATA = Path(os.environ.get("MT_DATA", "/path/to/data")) / "evalsets"
RS_FAST = ["zh", "ar", "ko", "hu", "bn"]


def test_subset_files():
    redux = json.loads((SUBSETS / "mmlu_redux_1000.json").read_text())
    assert redux["n"] == sum(map(len, redux["ids"].values())) == 1000 and len(redux["ids"]) == 57
    lite = json.loads((SUBSETS / "global_mmlu_lite_rs5.json").read_text())
    assert lite["languages"] == RS_FAST and all(len(lite["ids"][lang]) == 400 for lang in RS_FAST)
    bele = json.loads((SUBSETS / "belebele_200.json").read_text())
    assert bele["n"] == len({(k["link"], k["question_number"]) for k in bele["ids"]}) == 200
    gsm = json.loads((SUBSETS / "gsm8k_250.json").read_text())
    assert gsm["n"] == len(set(gsm["ids"])) == 250


@pytest.mark.skipif(not (DATA / "mmlu_redux").exists(), reason="pinned eval parquet not downloaded")
def test_loaders_match_subsets():
    sys.path.insert(0, str(REPO / "configs" / "eval" / "lm_eval_tasks"))
    import rs_utils as u
    assert len(u.mmlu_redux()["test"]) == 1000
    for lang in RS_FAST:
        g, b = u.global_mmlu_lite(lang=lang)["test"], u.belebele(lang=lang)["test"]
        assert len(g) == 400 and len(b) == 200
        assert set(g["answer"]) <= {0, 1, 2, 3} and set(b["answer"]) <= {0, 1, 2, 3}
    # BELEBELE is parallel: the same 200 questions in every language
    assert len({tuple(u.belebele(lang=lang)["test"]["item_id"]) for lang in RS_FAST}) == 1


def _u():
    sys.path.insert(0, str(REPO / "configs" / "eval" / "lm_eval_tasks"))
    import rs_utils
    return rs_utils


def test_rs_full_suite_lists_every_task_file():
    import yaml
    suite = yaml.safe_load((REPO / "configs" / "eval" / "rs_full.yaml").read_text())
    listed = {t for run in suite["runs"] for t in run["tasks"]}
    files = {p.stem for p in (REPO / "configs" / "eval" / "lm_eval_tasks").glob("rsf_*.yaml")}
    assert listed == files and len(listed) == 2 + 24 + 29 + 4
    gm = json.loads((SUBSETS / "global_mmlu_1000.json").read_text())
    assert len(gm["ids"]) == len(set(gm["ids"])) == 1000 and len(gm["languages"]) == 24


def test_mmlu_pro_extract():
    u = _u()
    assert u.mmlu_pro_extract("... so the answer is (C).") == "C"
    assert u.mmlu_pro_extract("The answer is J") == "J"
    assert u.mmlu_pro_extract("I think (A) at first. Answer: (B)") == "B"
    assert u.mmlu_pro_extract("The answer is (A). Wait, the answer is (D).") == "D"
    assert u.mmlu_pro_extract("Option E is right.") == "E"
    assert u.mmlu_pro_extract("no letter here") is None
    assert u.mmlu_pro_process_results({"gold": "C"}, ["the answer is (C)"]) == {"exact_match": 1.0}


def test_math500_scorer():
    u = _u()
    ok = lambda gold, resp: u.math500_process_results({"gold": gold}, [resp])["exact_match"]  # noqa: E731
    assert ok("\\frac{14}{3}", "so we get \\boxed{\\dfrac{14}{3}}") == 1.0
    assert ok("\\left( 3, \\frac{\\pi}{2} \\right)", "The point is \\boxed{(3, \\frac{\\pi}{2})}") == 1.0
    assert ok("10", "First 7, then \\boxed{10}") == 1.0
    assert ok("10", "\\boxed{11}") == 0.0
    assert ok("\\text{(C)}", "\\boxed{\\text{(C)}}") == 1.0


@pytest.mark.skipif(not (DATA.parent / "nltk_data").exists(), reason="pinned NLTK data not downloaded")
def test_ifeval_scorer_offline():
    u = _u()
    doc = {"key": 1, "prompt": "Answer without commas.", "instruction_id_list": ["punctuation:no_comma"],
           "kwargs": [{"num_highlights": None, "relation": None}]}
    assert u.ifeval_process_results(doc, ["No commas here at all."])["prompt_level_strict_acc"] is True
    assert u.ifeval_process_results(doc, ["Well, a comma."])["prompt_level_strict_acc"] is False


@pytest.mark.skipif(not (DATA / "global_mmlu").exists(), reason="RS-full parquet not downloaded")
def test_rs_full_loaders():
    u = _u()
    assert len(u.mmlu_5shot()["test"]) == 14042 and len(u.mmlu_redux(subset="mmlu_redux_ok.json")["test"]) == 5330
    shot = u.mmlu_5shot()["test"][0]["prompt"]
    assert shot.count("\nAnswer: ") == 5 and shot.endswith("\nAnswer:")
    assert len(u.mmlu_pro()["test"]) == 2000 and len(u.math500()["test"]) == 500 and len(u.ifeval()["test"]) == 541
    assert len(u.gsm8k_250(subset="gsm8k_full.json")["test"]) == 1319
    # parallel sets: the same items in every language
    assert len({tuple(u.global_mmlu(lang=lg)["test"]["item_id"]) for lg in ("en", "ja", "ne")}) == 1
    assert len({tuple(u.belebele(lang=lg, subset="belebele_full.json")["test"]["item_id"]) for lg in ("en", "az", "kk")}) == 1