from __future__ import annotations

import json
import tempfile
import unittest
from pathlib import Path

import analyze_phase3 as analysis


class Phase3AnalysisTests(unittest.TestCase):
    def test_ap36_stability_formula(self):
        with tempfile.TemporaryDirectory() as tmp:
            root = Path(tmp)
            (root / "client").mkdir()
            (root / "opprof").mkdir()
            (root / "client/result.json").write_text(
                json.dumps({"t0_mono_ns": 0, "warmup_seconds": 60})
            )
            (root / "client/requests.jsonl").write_text(
                "".join(
                    json.dumps({"success": True, "completed_s": index + 1}) + "\n"
                    for index in range(16)
                )
            )
            records = []
            for bin_index in range(3):
                for step in range(16):
                    records.append(
                        {
                            "step_index": len(records),
                            "model_executed": True,
                            "submit_mono_ns": int(
                                (45 + 5 * bin_index + (step + 0.5) / 16 * 5)
                                * 1e9
                            ),
                            "prefill_tokens": 100,
                            "decode_tokens": 0,
                        }
                    )
            (root / "opprof/test.jsonl").write_text(
                "".join(json.dumps(item) + "\n" for item in records)
            )
            result = analysis.ap36_warmup_stability(root)
            self.assertTrue(result["passes"])
            self.assertEqual(result["normalized_drift"], 0)

    def test_ap37_partial_verdict_can_confirm_but_not_refute(self):
        self.assertEqual(analysis.partial_verdict(True), "PASS")
        self.assertEqual(analysis.partial_verdict(False), "INCONCLUSIVE")

    def test_accepted_markers_come_only_from_complete_stages(self):
        with tempfile.TemporaryDirectory() as tmp:
            root = Path(tmp)
            complete = root / "stages/primary-01-saturation"
            complete.mkdir(parents=True)
            (complete / "stage-complete.json").write_text(
                json.dumps({"runs": ["P01-C00-saturation"]})
            )
            accepted = root / "primary/P01-C00/saturation"
            accepted.mkdir(parents=True)
            (accepted / "run-complete.json").write_text(
                json.dumps({"run_id": "P01-C00-saturation"})
            )
            unaccepted = root / "primary/P05-C00/saturation"
            unaccepted.mkdir(parents=True)
            (unaccepted / "run-complete.json").write_text(
                json.dumps({"run_id": "P05-C00-saturation"})
            )
            primary, confirmations, stages, excluded = analysis.accepted_marker_paths(
                root
            )
            self.assertEqual(primary, [accepted / "run-complete.json"])
            self.assertEqual(confirmations, [])
            self.assertEqual(stages, [complete])
            self.assertEqual(excluded, ["P05-C00-saturation"])

    def test_r64_is_ratio_of_cohort_sums(self):
        with tempfile.TemporaryDirectory() as tmp:
            path = Path(tmp) / "manifest.jsonl"
            rows = [{"input_tokens": value} for value in (1, 3, 2, 2)]
            path.write_text("".join(json.dumps(row) + "\n" for row in rows))
            value, pieces = analysis.manifest_raggedness(path, 2)
            self.assertEqual(pieces, [(2.0, 6.0), (0.0, 4.0)])
            self.assertAlmostEqual(value, 0.2)

    def test_one_percentage_point_ranking_ties(self):
        shares = dict.fromkeys(analysis.FAMILIES, 0.0)
        shares.update(attention=0.40, moe_gemm=0.395, moe_router=0.20)
        ranked = {item["family"]: item["rank"] for item in analysis.ranked_families(shares)}
        self.assertEqual(ranked["attention"], ranked["moe_gemm"])
        self.assertGreater(ranked["moe_router"], ranked["attention"])

    def test_holm_uses_declared_total_test_family(self):
        values = [{"p": 0.001}, {"p": 0.01}]
        analysis.holm(values, total_tests=10)
        self.assertAlmostEqual(values[0]["p_holm"], 0.01)
        self.assertAlmostEqual(values[1]["p_holm"], 0.09)

    def test_robust_spline_prediction_is_nonnegative(self):
        rows = [(float(x), float(n), float(2 * x + n)) for x in range(1, 20) for n in (1, 4)]
        predict, hull = analysis.fit_nonnegative_robust(rows)
        self.assertGreaterEqual(predict(3, 2), 0)
        self.assertTrue(analysis.inside_convex(hull, (3, 2)))
        self.assertFalse(analysis.inside_convex(hull, (100, 2)))


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