from __future__ import annotations

import json
import os
from pathlib import Path
import unittest

import numpy as np

from ap1_m1_log_amplitude_transport_reference import (
    EXPECTED_AUTHORITIES,
    FREEZE_POLICY,
    ReferenceConfig,
    file_sha256,
    load_inputs,
    seed_comoving_momenta,
    seed_mode_state,
    short_diagnostic_trajectory,
    transport_log_riccati,
)


ROOT = Path(os.environ["APEIRON_AP1_ROOT"])
PILOT = ROOT / "AP1/APEIRON_AP1_M1_LOG_AMPLITUDE_TRANSPORT_PILOT_20260902T190935Z.json"
FREEZE = ROOT / "AP1/APEIRON_AP1_M1_LOG_AMPLITUDE_TRANSPORT_TOLERANCES_20260902T191009Z.json"
HELDOUT = ROOT / "AP1/APEIRON_AP1_M1_LOG_AMPLITUDE_TRANSPORT_HELDOUT_20260902T191026Z.json"
CODE = ROOT / "AP1/CODE/ap1_m1_log_amplitude_transport_reference.py"


class AP1M1LogAmplitudeTransportReferenceTests(unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        cls.seed_report, cls.arrays, cls.authorities = load_inputs(ROOT)
        cls.pilot = json.loads(PILOT.read_text(encoding="utf-8"))
        cls.freeze = json.loads(FREEZE.read_text(encoding="utf-8"))
        cls.heldout = json.loads(HELDOUT.read_text(encoding="utf-8"))

    def test_frozen_authority_chain_is_exact(self):
        self.assertEqual(self.authorities, EXPECTED_AUTHORITIES)
        self.assertEqual(
            self.heldout["authority_sha256"][CODE.relative_to(ROOT).as_posix()],
            file_sha256(CODE),
        )

    def test_config_is_short_nested_and_shell_sets_are_disjoint(self):
        config = ReferenceConfig()
        config.validate()
        self.assertFalse(
            set(config.pilot_mode_indices) & set(config.heldout_mode_indices)
        )
        self.assertLessEqual(max(config.heldout_endpoints_relative_N), 0.125)

    def test_seed_physical_to_comoving_map_roundtrips(self):
        indices = ReferenceConfig().heldout_mode_indices
        q_seed, k_comoving = seed_comoving_momenta(self.arrays, indices)
        N_seed = float(self.arrays["m1_state_physical_N"][5])
        self.assertTrue(
            np.allclose(q_seed, k_comoving * np.exp(-N_seed), rtol=4e-15)
        )

    def test_seed_state_is_transported_not_reinitialized(self):
        state = seed_mode_state(
            self.arrays, ReferenceConfig().heldout_mode_indices
        )
        invariant = (
            np.log(2.0)
            + 2.0 * state["log_amplitude"]
            + state["log_abs_riccati_imag"]
            - np.log(state["wronskian"])
        )
        self.assertLess(float(np.max(np.abs(invariant))), 2.0e-15)

    def test_fresh_short_log_transport_preserves_log_wronskian(self):
        indices = (1,)
        _q_seed, k_comoving = seed_comoving_momenta(self.arrays, indices)
        initial = seed_mode_state(self.arrays, indices)
        trajectory = short_diagnostic_trajectory(
            self.arrays, 0.002, 2.5e-5, 2.0e-11
        )
        result = transport_log_riccati(
            trajectory, k_comoving, initial, 5.0e-10
        )
        self.assertLess(result["log_wronskian_relative_error"], 1.0e-11)

    def test_pilot_is_complete_and_strictly_pre_tachyonic(self):
        self.assertTrue(self.pilot["all_pilot_gates_pass"])
        self.assertEqual(self.pilot["pilot_aggregate"]["tachyonic_windows"], 0)
        self.assertFalse(self.pilot["checkpoint_eligible"])

    def test_tolerances_were_frozen_before_heldout(self):
        self.assertTrue(self.freeze["tolerances_frozen_before_heldout"])
        self.assertFalse(self.freeze["heldout_evaluated"])
        self.assertEqual(
            self.freeze["input_sha256"][PILOT.name], file_sha256(PILOT)
        )
        self.assertEqual(
            self.freeze["predeclared_freeze_policy"], FREEZE_POLICY
        )

    def test_all_three_heldout_windows_are_tachyonic(self):
        self.assertEqual(self.heldout["heldout_aggregate"]["windows"], 3)
        self.assertEqual(
            self.heldout["heldout_aggregate"]["tachyonic_windows"], 3
        )

    def test_all_frozen_heldout_gates_pass(self):
        self.assertTrue(self.heldout["all_heldout_gates_pass"])
        self.assertTrue(all(self.heldout["gates"].values()))
        self.assertTrue(self.heldout["checkpoint_eligible"])

    def test_log_amplitude_and_convergence_are_below_frozen_caps(self):
        aggregate = self.heldout["heldout_aggregate"]
        thresholds = self.heldout["frozen_thresholds"]
        self.assertLessEqual(
            aggregate["max_abs_log_amplitude_vs_decimal80"],
            thresholds["max_abs_log_amplitude_vs_decimal80"],
        )
        self.assertLessEqual(
            aggregate["max_decimal_time_grid_relative_delta"],
            thresholds["max_decimal_time_grid_relative_delta"],
        )
        self.assertLessEqual(
            aggregate["max_riccati_refinement_log_amplitude_delta"],
            thresholds["max_riccati_refinement_log_amplitude_delta"],
        )

    def test_log_and_decimal_wronskians_pass_existing_cap(self):
        aggregate = self.heldout["heldout_aggregate"]
        cap = self.heldout["frozen_thresholds"][
            "max_wronskian_relative_error"
        ]
        self.assertLessEqual(
            aggregate["max_log_wronskian_relative_error"], cap
        )
        self.assertLessEqual(
            aggregate["max_decimal_wronskian_relative_error"], cap
        )
        self.assertIn("excluded from the gate", self.heldout["wronskian_audit"])

    def test_claim_boundary_keeps_long_physics_locked(self):
        self.assertEqual(self.heldout["new_AP1_M1_background_runs"], 0)
        self.assertFalse(self.heldout["background_candidate_assessed"])
        self.assertFalse(self.heldout["background_tolerances_frozen"])
        self.assertFalse(self.heldout["physical_response_kernel_started"])
        self.assertIn("UV-tail-to-resolved", self.heldout["unvalidated_scope"])


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