import importlib.util
import unittest
from pathlib import Path


MODULE_PATH = Path(__file__).with_name("audit_ground_truth.py")
SPEC = importlib.util.spec_from_file_location("audit_ground_truth", MODULE_PATH)
audit = importlib.util.module_from_spec(SPEC)
assert SPEC.loader is not None
SPEC.loader.exec_module(audit)


class AuditGroundTruthTest(unittest.TestCase):
    def test_case_summary_exposes_tied_top_set(self):
        rows = [
            {
                "cell_id": "a",
                "score_req_s_per_gpu": 2.0,
                "capacity_lower_bound_req_s_per_gpu": 2.0,
                "capacity_upper_bound_req_s_per_gpu": 2.1,
                "fully_valid": True,
            },
            {
                "cell_id": "b",
                "score_req_s_per_gpu": 2.0,
                "capacity_lower_bound_req_s_per_gpu": 2.0,
                "capacity_upper_bound_req_s_per_gpu": 2.2,
                "fully_valid": True,
            },
            {
                "cell_id": "c",
                "score_req_s_per_gpu": 1.0,
                "capacity_lower_bound_req_s_per_gpu": 1.0,
                "capacity_upper_bound_req_s_per_gpu": 1.5,
                "fully_valid": False,
            },
        ]

        summary = audit.summarize_case("test", rows)

        self.assertEqual(summary["top_set"], ["a", "b"])
        self.assertEqual(summary["distinct_score_count"], 2)
        self.assertEqual(summary["tied_pair_count"], 1)
        self.assertEqual(summary["informative_pair_count"], 2)
        self.assertAlmostEqual(summary["random_top_set_hit_rate"], 2 / 3)
        self.assertEqual(summary["invalid_cells"], ["c"])
        self.assertEqual(
            summary["possibly_optimal_set_from_search_brackets"], ["a", "b"]
        )

    def test_config_gpu_count_includes_data_parallelism(self):
        result = {
            "config_patch": {
                "flag_patch": {
                    "tensor-parallel-size": 2,
                    "data-parallel-size": 4,
                    "expert-parallel-size": 8,
                    "max-num-seqs": 128,
                    "max-num-batched-tokens": 384,
                }
            }
        }

        config = audit.config_from_result(result)

        self.assertEqual(config["gpu_count"], 8)
        self.assertEqual(audit.cell_id(config), "tp2_dp4_ep8_mns128_mbt384")


if __name__ == "__main__":
    unittest.main()