from __future__ import annotations

import importlib.util
from pathlib import Path
import unittest


ROOT = Path(__file__).resolve().parents[1]
SPEC = importlib.util.spec_from_file_location("analyze_study4", ROOT / "scripts" / "analyze_study4.py")
assert SPEC is not None and SPEC.loader is not None
ANALYSIS = importlib.util.module_from_spec(SPEC)
SPEC.loader.exec_module(ANALYSIS)


class Study4AnalysisTests(unittest.TestCase):
    def test_cluster_sign_flip_known_one_cluster(self) -> None:
        self.assertEqual(ANALYSIS.cluster_sign_flip([1.0]), 1.0)

    def test_cluster_sign_flip_strong_same_direction(self) -> None:
        self.assertAlmostEqual(ANALYSIS.cluster_sign_flip([1.0] * 5), 2 / 32)

    def test_exact_mcnemar_and_holm(self) -> None:
        self.assertEqual(ANALYSIS.exact_mcnemar([1] * 8 + [0] * 2, [0] * 8 + [1] * 2), (8, 2, 0.109375))
        self.assertEqual(ANALYSIS.holm_adjust([0.04, 0.01, 0.20, 0.03]), [0.09, 0.04, 0.20, 0.09])

    def test_primary_clusters_models_within_task(self) -> None:
        rows = []
        for task_index in range(2):
            task = f"T{task_index}"
            for model in ANALYSIS.MODELS:
                for harness in ANALYSIS.HARNESSES:
                    rows.append(
                        {
                            "task_id": task,
                            "model_id": model,
                            "harness_id": harness,
                            "resolved_at_1": int(harness == "H007" and task_index == 0),
                        }
                    )
        result = ANALYSIS.primary(rows)
        self.assertEqual(result["independent_clusters"], 2)
        self.assertEqual(result["model_task_pairs"], 6)
        self.assertEqual(result["h007_count"], 3)
        self.assertEqual(result["h000_count"], 0)
        self.assertAlmostEqual(result["paired_risk_difference"], 0.5)


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