from __future__ import annotations

import os
from pathlib import Path
import tempfile
import unittest

import numpy as np

from ap1_m1_background_preflight import (
    EXPECTED_STATE_SHA256,
    IntegrityError,
    MissingClosureError,
    SharedLateParameters,
    build_preflight_report,
    decode_hard_pass_state,
    h0_to_mpl,
    require_new_trajectory_closures,
    standard_fluid_background,
)


STATE = Path(os.environ["APEIRON_AP1_STATE_NPZ"])


class AP1M1BackgroundPreflightTests(unittest.TestCase):
    def test_exact_hard_pass_state_decodes(self):
        y = decode_hard_pass_state(STATE)
        self.assertEqual(y.shape, (15361, 6))
        self.assertAlmostEqual(y[-1, 4], 8.312498343027667e-8, places=20)
        self.assertAlmostEqual(y[-1, 5], 0.09959140184968283, places=15)

    def test_tampered_state_is_rejected(self):
        with tempfile.NamedTemporaryFile(suffix=".npz") as tmp:
            np.savez_compressed(tmp.name, x=np.zeros(15361 * 6))
            with self.assertRaises(IntegrityError):
                decode_hard_pass_state(Path(tmp.name))

    def test_present_scale_is_far_below_frozen_interval(self):
        ratio = 8.312498343027667e-8 / h0_to_mpl(67.4)
        self.assertGreater(ratio, 1.0e53)
        self.assertLess(ratio, 2.0e53)

    def test_standard_fluid_scalings(self):
        N = np.array([0.0, 1.0])
        synthetic = SharedLateParameters(
            H0_km_s_Mpc=67.4,
            Omega_b0=0.05,
            Omega_c0=0.25,
            Omega_gamma0=5.0e-5,
            Omega_nu_rel0=3.0e-5,
        )
        f = standard_fluid_background(N, synthetic)
        self.assertAlmostEqual(f["rho_b"][1] / f["rho_b"][0], np.exp(-3.0), places=15)
        self.assertAlmostEqual(f["rho_gamma"][1] / f["rho_gamma"][0], np.exp(-4.0), places=15)

    def test_missing_new_closures_fail_closed(self):
        with self.assertRaises(MissingClosureError):
            require_new_trajectory_closures(None, None)

    def test_report_does_not_claim_a_run(self):
        report = build_preflight_report(STATE)
        self.assertEqual(report["frozen_v7_13"]["state_sha256"], EXPECTED_STATE_SHA256)
        self.assertFalse(report["new_run"]["started"])
        self.assertFalse(report["new_run"]["old_solver_or_physical_map_called"])
        self.assertFalse(report["new_run"]["runnable"])
        self.assertTrue(report["M1_standard_fluids"]["parameter_vector_complete"])
        self.assertEqual(
            report["M1_standard_fluids"]["massive_neutrino_phase_space_closure"],
            "PASS",
        )
        self.assertEqual(
            report["new_run"]["blocking_closures"],
            ["renormalized chi background closure on the new AP1-M1 trajectory"],
        )


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