from __future__ import annotations

from dataclasses import replace
import json
from pathlib import Path
import tempfile
import unittest
from unittest.mock import patch

import numpy as np

import ap1_m1_balanced_coupled_stage_preflight as preflight
import ap1_m1_coupled_moving_split_background as g29


ROOT = Path(__file__).resolve().parents[2]
CODE = ROOT / "AP1/CODE/ap1_m1_balanced_coupled_stage_preflight.py"
TEST = ROOT / "AP1/CODE/test_ap1_m1_balanced_coupled_stage_preflight.py"


class BalancedCoupledStagePreflightTests(unittest.TestCase):
    def test_authority_chain_is_hash_exact(self) -> None:
        g30_checkpoint, g31_checkpoint = preflight.verify_authorities(ROOT)
        self.assertTrue(g30_checkpoint["all_checkpoint_gates_pass"])
        self.assertTrue(g31_checkpoint["all_reference_gates_pass"])

    def test_plan_has_three_exactly_nested_levels(self) -> None:
        plan = preflight.BalancedCoupledStagePlan()
        plan.validate()
        for phase in ("pilot", "heldout"):
            coarse, middle, fine = plan._level_nodes(phase)
            self.assertEqual(fine - 1, 2 * (middle - 1))
            self.assertEqual(middle - 1, 2 * (coarse - 1))

    def test_future_heldout_is_beyond_all_prior_diagnostic_support(self) -> None:
        plan = preflight.BalancedCoupledStagePlan()
        self.assertEqual(plan.pilot_span_N, plan.prior_diagnostic_endpoint_N)
        self.assertGreater(
            plan.heldout_evaluation_start_N, plan.prior_diagnostic_endpoint_N
        )

    def test_physics_setting_change_is_rejected(self) -> None:
        with self.assertRaises(ValueError):
            replace(
                preflight.BalancedCoupledStagePlan(), qcut_over_Lambda=0.61
            ).validate()
        with self.assertRaises(ValueError):
            replace(
                preflight.BalancedCoupledStagePlan(), K_over_Lambda=32.0
            ).validate()

    def test_refinement_changes_only_picard_count(self) -> None:
        plan = preflight.BalancedCoupledStagePlan()
        frozen = g29.CoupledConfig()
        self.assertGreater(plan.fixed_point_steps, frozen.fixed_point_steps)
        self.assertEqual(plan.relaxation, frozen.relaxation)
        self.assertEqual(plan.max_mode_step_N, frozen.fine_mode_step_N)
        self.assertEqual(plan.background_rtol, frozen.fine_background_rtol)

    def test_existing_metric_policy_is_copied_not_relaxed(self) -> None:
        self.assertEqual(
            preflight.PREDECLARED_POLICY["inherited_metric_freeze_policy"],
            preflight.g30.STAGE_FREEZE_POLICY["metrics"],
        )

    def test_mode_budgets_are_predeclared_and_bounded(self) -> None:
        plan = preflight.BalancedCoupledStagePlan()
        pilot = 128 + plan.entry_nodes_per_panel * (plan.pilot_nodes_fine - 1)
        heldout = 128 + plan.entry_nodes_per_panel * (plan.heldout_nodes_fine - 1)
        self.assertEqual(pilot, 8320)
        self.assertEqual(heldout, 10368)
        self.assertLessEqual(pilot, plan.max_pilot_resolved_modes)
        self.assertLessEqual(heldout, plan.max_heldout_resolved_modes)

    def test_balanced_source_composition_enforces_stored_junction(self) -> None:
        nodes = 3
        trajectory = preflight.ChiTrajectory(
            N=np.linspace(0.0, 2.0e-4, nodes),
            H=np.full(nodes, 4.0e-7),
            Hdot=np.full(nodes, -1.0e-14),
            sigma=np.full(nodes, 0.015),
            theta=np.full(nodes, -0.69),
        )
        zeros = {key: np.zeros((nodes, 4), dtype=np.longdouble) for key in g29.SOURCE_KEYS}
        resolved = {
            "sector_total": {key: value.copy() for key, value in zeros.items()},
            "diagnostics": {
                "max_wronskian_relative_error": 0.0,
                "dynamic_chart_switches": 1,
                "vacuum_resets": 0,
            },
        }
        tail = {
            "sector_finite": {key: value.copy() for key, value in zeros.items()},
            "asymptotic_signed": {
                key: np.zeros(nodes, dtype=np.longdouble) for key in g29.SOURCE_KEYS
            },
            "diagnostics": {},
        }
        arrays = {"quantum_rho_pressure_chi2": np.array([1.0, 2.0, 3.0])}
        with patch.object(
            preflight.g31, "transport_resolved_modes_balanced", return_value=resolved
        ), patch.object(preflight.g29, "moving_adiabatic_tail", return_value=tail):
            source, diagnostics, internal = preflight.assemble_balanced_quantum_sources(
                trajectory, arrays, 2, 4, 64.0, 1.25e-5, nodes, 257
            )
        self.assertTrue(diagnostics["balanced_chart_used"])
        self.assertTrue(diagnostics["junction_values_enforced_exactly"])
        self.assertEqual(
            [source[key][0] for key in g29.SOURCE_KEYS], [1.0, 2.0, 3.0]
        )
        self.assertIs(internal["resolved"], resolved)
        self.assertIs(internal["tail"], tail)

    def test_preflight_passes_without_running_background(self) -> None:
        report = preflight.build_preflight(ROOT, CODE, TEST)
        self.assertTrue(report["all_preflight_gates_pass"])
        self.assertTrue(all(report["gates"].values()))
        self.assertEqual(report["background_runs"], 0)
        self.assertEqual(report["physical_response_kernel_runs"], 0)

    def test_nonpass_preflight_is_never_written(self) -> None:
        with tempfile.TemporaryDirectory() as directory:
            output = Path(directory) / "forbidden.json"
            with self.assertRaises(preflight.BalancedCoupledPreflightError):
                preflight.write_pass_preflight(
                    {"all_preflight_gates_pass": False}, output
                )
            self.assertFalse(output.exists())

    def test_claim_boundary_keeps_every_physical_output_locked(self) -> None:
        report = preflight.build_preflight(ROOT, CODE, TEST)
        self.assertFalse(report["seed_released"])
        self.assertFalse(report["production_background_tolerances_frozen"])
        self.assertFalse(report["nonpass_stored_or_used_as_seed"])
        self.assertEqual(report["trajectory_rows_persisted"], 0)
        self.assertEqual(report["AP1_status"], "ORANGE")


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