from __future__ import annotations

from dataclasses import asdict, replace
import json
from pathlib import Path
from types import SimpleNamespace
import tempfile
import unittest

import numpy as np

import ap1_m1_coupled_moving_split_background as g29
import ap1_m1_staged_log_background_extension as staged


ROOT = Path(__file__).resolve().parents[2]


class StagedLogBackgroundExtensionTests(unittest.TestCase):
    @classmethod
    def setUpClass(cls) -> None:
        cls.checkpoint, cls.arrays = staged.load_authorities(ROOT)
        cls.code_path = Path(staged.__file__).resolve()

    def test_authority_chain_and_G29_checkpoint_are_exact(self) -> None:
        self.assertTrue(self.checkpoint["all_checkpoint_gates_pass"])
        self.assertEqual(
            self.checkpoint["classification"],
            "M1_BOUNDED_COUPLED_MOVING_SPLIT_BACKGROUND_REFERENCE_PASS_STAGED_LONG_EXTENSION_OPEN_PHYSICAL_KERNEL_BLOCKED",
        )

    def test_stage_plan_is_ordered_disjoint_and_doubles_G29(self) -> None:
        config = staged.StageConfig()
        config.validate()
        self.assertGreater(config.pilot_evaluation_start_N, config.baseline_span_N)
        self.assertGreater(config.heldout_evaluation_start_N, config.pilot_span_N)
        self.assertEqual(config.heldout_span_N / config.baseline_span_N, 2.0)

    def test_fine_and_coarse_stage_spacings_match_G29(self) -> None:
        config = staged.StageConfig()
        g29_config = g29.CoupledConfig()
        expected_fine = g29_config.heldout_span_N / (g29_config.heldout_fine_nodes - 1)
        for span, nodes in (
            (config.pilot_span_N, config.pilot_fine_nodes),
            (config.heldout_span_N, config.heldout_fine_nodes),
        ):
            self.assertAlmostEqual(span / (nodes - 1), expected_fine, places=20)
            coarse_nodes = (nodes + 1) // 2
            self.assertAlmostEqual(
                span / (coarse_nodes - 1), 2.0 * expected_fine, places=20
            )

    def test_config_rejects_physics_or_G29_numerics_change(self) -> None:
        with self.assertRaises(ValueError):
            replace(staged.StageConfig(), qcut_over_Lambda=0.61).validate()
        with self.assertRaises(ValueError):
            replace(staged.StageConfig(), relaxation=0.8).validate()

    def test_heldout_mode_budget_is_predeclared(self) -> None:
        config = staged.StageConfig()
        expected = 128 + config.fine_entry_nodes_per_panel * (
            config.heldout_fine_nodes - 1
        )
        self.assertEqual(expected, 4224)
        self.assertLessEqual(
            expected, staged.STAGE_FREEZE_POLICY["max_heldout_resolved_modes"]
        )

    def test_overlap_point_must_lie_on_nested_grid(self) -> None:
        run = {
            "trajectory": SimpleNamespace(
                N=np.array([0.0, 0.01, 0.02]),
                H=np.array([3.0, 2.0, 1.0]),
            )
        }
        self.assertEqual(staged._grid_value(run, 0.01, "H"), 2.0)
        with self.assertRaises(staged.StagedExtensionError):
            staged._grid_value(run, 0.015, "H")

    def _synthetic_pilot(self) -> dict:
        return {
            "all_pilot_gates_pass": True,
            "heldout_evaluated": False,
            "authority_sha256": {
                "AP1/CODE/ap1_m1_staged_log_background_extension.py": staged.file_sha256(
                    self.code_path
                )
            },
            "config": asdict(staged.StageConfig()),
            "pilot_ensemble": {
                "metrics_for_freeze_or_gate": {
                    name: 1.0e-12
                    for name in staged.STAGE_FREEZE_POLICY["metrics"]
                },
                "stage_endpoint_aggregate_not_a_seed": {
                    "N_relative": 0.03,
                    "H_Mpl": 4.0e-7,
                },
            },
            "new_AP1_M1_background_runs": 2,
        }

    def test_tolerance_freeze_is_chronology_bound(self) -> None:
        with tempfile.TemporaryDirectory() as directory:
            pilot_path = Path(directory) / "pilot.json"
            pilot_path.write_text(json.dumps(self._synthetic_pilot()), encoding="utf-8")
            freeze = staged.build_freeze(pilot_path, self.code_path)
            self.assertTrue(freeze["all_tolerance_freeze_gates_pass"])
            self.assertFalse(freeze["heldout_evaluated"])
            self.assertFalse(freeze["production_background_tolerances_frozen"])
            self.assertEqual(
                freeze["heldout_plan"]["overlap_reference"]["N_relative"], 0.03
            )

    def test_freeze_rejects_changed_code_hash(self) -> None:
        pilot = self._synthetic_pilot()
        pilot["authority_sha256"][
            "AP1/CODE/ap1_m1_staged_log_background_extension.py"
        ] = "0" * 64
        with tempfile.TemporaryDirectory() as directory:
            pilot_path = Path(directory) / "pilot.json"
            pilot_path.write_text(json.dumps(pilot), encoding="utf-8")
            with self.assertRaises(staged.StagedExtensionError):
                staged.build_freeze(pilot_path, self.code_path)

    def test_nonpass_report_is_never_written(self) -> None:
        with tempfile.TemporaryDirectory() as directory:
            target = Path(directory) / "forbidden.json"
            with self.assertRaises(staged.StagedExtensionError):
                staged._write_pass_json({"all_pass": False}, target, "all_pass")
            self.assertFalse(target.exists())

    def test_extended_smoke_remains_finite_without_vacuum_reset(self) -> None:
        config = replace(g29.CoupledConfig(), fixed_point_steps=2)
        run = g29.run_coupled(
            self.arrays,
            2.0e-4,
            17,
            config,
            background_rtol=2.0e-9,
            max_mode_step=2.5e-5,
            entry_nodes_per_panel=1,
            tail_nodes_per_octave=2,
            K_over_Lambda=8.0,
            adiabatic_support_nodes=17,
            adiabatic_momentum_support_nodes=65,
        )
        self.assertTrue(np.all(run["trajectory"].H > 0.0))
        self.assertTrue(np.all(np.isfinite(run["background_rows"])))
        self.assertEqual(run["source_diagnostics"]["resolved"]["vacuum_resets"], 0)

    def test_checkpoint_contract_keeps_production_locked(self) -> None:
        source = self.code_path.read_text(encoding="utf-8")
        self.assertIn('"production_background_tolerances_frozen": False', source)
        self.assertIn('"trajectory_rows_persisted": 0', source)
        self.assertIn('"physical_response_kernel_started": False', source)
        self.assertIn('"seed_released": False', source)


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