from __future__ import annotations

import json
import os
from pathlib import Path
import unittest

import numpy as np

import ap1_m1_coupled_moving_split_background as g29
from ap1_m1_balanced_chart_transport import (
    BalancedChartError,
    BalancedState,
    EXPECTED_AUTHORITIES,
    REFERENCE_POLICY,
    _apply_transfer,
    file_sha256,
    from_u_riccati,
    log_wronskian,
    to_u_quadratics,
    verify_authorities,
)
from ap1_m1_log_amplitude_transport_reference import seed_mode_state


ROOT = Path(os.environ["APEIRON_AP1_ROOT"])
REPORT = ROOT / "AP1/APEIRON_AP1_M1_BALANCED_CHART_TRANSPORT_REFERENCE_CHECKPOINT_20260902T214232Z.json"
CODE = ROOT / "AP1/CODE/ap1_m1_balanced_chart_transport.py"


class AP1M1BalancedChartTransportTests(unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        cls.arrays, cls.authorities = verify_authorities(ROOT)
        cls.report = json.loads(REPORT.read_text(encoding="utf-8"))

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

    def test_reference_policy_keeps_existing_wronskian_cap(self):
        self.assertEqual(REFERENCE_POLICY["wronskian_relative_cap"], g29.WRONSKIAN_LIMIT)
        self.assertEqual(REFERENCE_POLICY["entry_nodes_per_panel"], 1)
        self.assertLess(
            REFERENCE_POLICY["fine_mode_step_N"],
            REFERENCE_POLICY["coarse_mode_step_N"],
        )

    def test_seed_conversion_preserves_all_quadratics(self):
        seed = seed_mode_state(self.arrays, tuple(range(128)))
        state = from_u_riccati(
            seed["log_amplitude"],
            seed["riccati_real"],
            seed["log_abs_riccati_imag"],
            0.0023,
        )
        u2, cross, v2 = to_u_quadratics(state, 0.0023)
        expected_u2 = np.exp(2.0 * np.asarray(seed["log_amplitude"], dtype=np.longdouble))
        beta = np.exp(np.asarray(seed["log_abs_riccati_imag"], dtype=np.longdouble))
        x = np.asarray(seed["riccati_real"], dtype=np.longdouble)
        self.assertLess(float(np.max(np.abs(u2 / expected_u2 - 1.0))), 3.0e-18)
        self.assertLess(
            float(np.max(np.abs(cross - x * expected_u2)))
            / max(float(np.max(np.abs(x * expected_u2))), 1.0e-300),
            3.0e-18,
        )
        expected_v2 = (x * x + beta * beta) * expected_u2
        self.assertLess(float(np.max(np.abs(v2 / expected_v2 - 1.0))), 3.0e-18)

    def test_seed_conversion_preserves_log_wronskian(self):
        seed = seed_mode_state(self.arrays, tuple(range(128)))
        state = from_u_riccati(
            seed["log_amplitude"],
            seed["riccati_real"],
            seed["log_abs_riccati_imag"],
            0.0023,
        )
        error = log_wronskian(state, 0.0023) - np.log(
            np.asarray(seed["wronskian"], dtype=np.longdouble)
        )
        self.assertLess(float(np.max(np.abs(error))), 2.0e-15)

    def test_exact_symplectic_rotation_pivots_charts(self):
        shape = (4, 1)
        state = BalancedState(
            np.zeros(shape, dtype=bool),
            np.zeros(shape, dtype=np.longdouble),
            np.full(shape, -0.25j, dtype=np.clongdouble),
        ).validated()
        zero = np.zeros(shape, dtype=np.longdouble)
        one = np.ones(shape, dtype=np.longdouble)
        updated, switches = _apply_transfer(state, (zero, one, -one, zero))
        self.assertEqual(switches, 4)
        self.assertTrue(np.all(updated.q_chart))
        self.assertTrue(np.allclose(updated.ratio, 0.25j, rtol=0.0, atol=1e-19))
        u2, cross, v2 = to_u_quadratics(updated, 0.0023)
        self.assertTrue(np.allclose(u2, 0.0625, rtol=0.0, atol=1e-19))
        self.assertTrue(np.allclose(cross, 0.0, rtol=0.0, atol=1e-19))
        self.assertTrue(np.allclose(v2, 0.0023**2, rtol=1e-16, atol=0.0))

    def test_invalid_canonical_sign_fails_closed(self):
        shape = (4, 1)
        with self.assertRaises(BalancedChartError):
            BalancedState(
                np.zeros(shape, dtype=bool),
                np.zeros(shape, dtype=np.longdouble),
                np.full(shape, 0.25j, dtype=np.clongdouble),
            ).validated()

    def test_formal_reference_all_gates_pass(self):
        self.assertTrue(self.report["all_reference_gates_pass"])
        self.assertTrue(all(self.report["gates"].values()))
        self.assertTrue(self.report["checkpoint_eligible"])

    def test_pole_crossing_is_refined_and_reset_free(self):
        metrics = self.report["metrics"]
        fine = metrics["fine_balanced_diagnostics"]
        self.assertTrue(metrics["old_single_chart_failed_closed"])
        self.assertGreater(fine["dynamic_chart_switches"], 0)
        self.assertEqual(fine["vacuum_resets"], 0)
        self.assertLessEqual(
            metrics["pole_traversal_mode_step_source_gross_relative_delta"],
            REFERENCE_POLICY["mode_step_source_gross_relative_cap"],
        )
        self.assertLessEqual(
            fine["max_wronskian_relative_error"],
            REFERENCE_POLICY["wronskian_relative_cap"],
        )

    def test_claim_boundary_keeps_physical_work_locked(self):
        self.assertEqual(self.report["new_AP1_M1_background_candidates"], 0)
        self.assertEqual(self.report["trajectory_rows_persisted"], 0)
        self.assertFalse(self.report["physical_response_kernel_started"])
        self.assertFalse(self.report["seed_released"])
        self.assertFalse(self.report["existing_gate_thresholds_changed"])
        self.assertIn("not a coupled long background", self.report["claim_boundary"])


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