from __future__ import annotations

import unittest

from app.simulation.core.peng_robinson import AIR_PR, HELIUM_PR, NITROGEN_PR, PengRobinsonFluid


class PengRobinsonTest(unittest.TestCase):
    def test_helium_near_ideal_at_atmospheric_condition(self) -> None:
        z = HELIUM_PR.compressibility_factor(101_325.0, 300.0)
        density = HELIUM_PR.density(101_325.0, 300.0)

        self.assertAlmostEqual(z, 1.0, delta=1.0e-3)
        self.assertAlmostEqual(density, 0.1625, delta=0.002)

    def test_helium_density_pressure_round_trip(self) -> None:
        pressure = 15.3e6
        temperature = 293.15
        density = HELIUM_PR.density(pressure, temperature)

        self.assertGreater(HELIUM_PR.compressibility_factor(pressure, temperature), 1.0)
        self.assertAlmostEqual(
            HELIUM_PR.pressure_from_density(temperature, density),
            pressure,
            delta=pressure * 1.0e-12,
        )

    def test_helium_residual_enthalpy_is_small_at_atmosphere(self) -> None:
        self.assertAlmostEqual(
            HELIUM_PR.residual_specific_enthalpy(101_300.0, 298.15),
            17.834,
            delta=0.01,
        )

    def test_helium_residual_enthalpy_captures_high_pressure_departure(self) -> None:
        self.assertAlmostEqual(
            HELIUM_PR.residual_specific_enthalpy(13_839_965.0, 287.7322),
            11953.9,
            delta=0.1,
        )

    def test_air_reference_remains_available_for_other_models(self) -> None:
        z = AIR_PR.compressibility_factor(101_325.0, 300.0)
        density = AIR_PR.density(101_325.0, 300.0)

        self.assertAlmostEqual(z, 1.0, delta=1.0e-3)
        self.assertAlmostEqual(density, 1.177, delta=0.01)

    def test_nitrogen_can_return_multiple_roots(self) -> None:
        roots = NITROGEN_PR.compressibility_roots(1.0e6, 100.0)

        self.assertEqual(len(roots), 3)
        self.assertLess(roots[0], roots[-1])
        self.assertEqual(NITROGEN_PR.compressibility_factor(1.0e6, 100.0, phase="liquid"), roots[0])
        self.assertEqual(NITROGEN_PR.compressibility_factor(1.0e6, 100.0, phase="vapor"), roots[-1])

    def test_custom_fluid_can_be_defined_explicitly(self) -> None:
        fluid = PengRobinsonFluid(
            name="custom_helium",
            molar_mass=0.004002602,
            critical_temperature=5.1953,
            critical_pressure=227_460.0,
            acentric_factor=-0.385,
        )

        self.assertAlmostEqual(fluid.specific_gas_constant, 2077.3, delta=0.5)
        self.assertAlmostEqual(fluid.compressibility_factor(1.0e6, 298.15), HELIUM_PR.compressibility_factor(1.0e6, 298.15))

    def test_rejects_invalid_states(self) -> None:
        with self.assertRaises(ValueError):
            HELIUM_PR.compressibility_factor(0.0, 300.0)
        with self.assertRaises(ValueError):
            HELIUM_PR.density(101_325.0, 0.0)
        with self.assertRaises(ValueError):
            HELIUM_PR.pressure_from_density(300.0, -1.0)


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