from __future__ import annotations

import unittest

from PythonModels.components.amesim_pneumatic_line import (
    AmesimPnl0002Pipe,
    AmesimPnl0003Pipe,
    AmesimPnl00rPipe,
)


class AmesimPnl0002PipeTests(unittest.TestCase):
    def setUp(self) -> None:
        self.pipe = AmesimPnl0002Pipe(
            name="pneumatic_86",
            diameter_mm=20.0,
            length_m=2.0,
            relative_roughness=0.045 / 20.0,
            pctr_0=100_000.0,
            Tctr_0=293.15,
        )

    def test_initial_state_uses_center_compliance(self) -> None:
        properties = self.pipe.properties()

        self.assertAlmostEqual(self.pipe.volume, 6.283185307179586e-4)
        self.assertEqual(len(self.pipe.get_state_vector()), 2)
        self.assertAlmostEqual(properties.p, 100_000.0, delta=1.0e-6)
        self.assertAlmostEqual(properties.T, 293.15)
        self.assertAlmostEqual(self.pipe.port_1.p, properties.p)
        self.assertAlmostEqual(self.pipe.port_2.p, properties.p)

    def test_port_flow_enters_center_from_higher_external_pressure(self) -> None:
        forward = self.pipe.port_mass_flow(
            port_pressure_pa=101_000.0,
            port_temperature_k=293.15,
        )
        reverse = self.pipe.port_mass_flow(
            port_pressure_pa=99_000.0,
            port_temperature_k=293.15,
        )

        self.assertGreater(forward, 0.0)
        self.assertLess(reverse, 0.0)

    def test_connection_derivatives_conserve_two_external_port_flows(self) -> None:
        properties = self.pipe.properties()

        derivative = self.pipe.derivatives_from_connections(
            port_1_m_flow=0.2,
            connected_h_1=properties.h + 1000.0,
            port_2_m_flow=-0.1,
            connected_h_2=properties.h - 1000.0,
        )

        self.assertAlmostEqual(derivative.m, 0.1)


class AmesimPnl0003PipeTests(unittest.TestCase):
    def setUp(self) -> None:
        self.pipe = AmesimPnl0003Pipe(
            name="pneumatic_88",
            diameter_mm=20.0,
            length_m=0.3,
            relative_roughness=0.045 / 20.0,
            p1_0=15.3e6,
            p2_0=15.3e6,
            T1_0=293.15,
            T2_0=293.15,
        )

    def test_initial_state_uses_two_half_volume_compliances(self) -> None:
        port_1 = self.pipe.properties_1()
        port_2 = self.pipe.properties_2()

        self.assertAlmostEqual(self.pipe.volume, 9.424777960769381e-5)
        self.assertAlmostEqual(self.pipe.compliance_volume, self.pipe.volume / 2.0)
        self.assertEqual(len(self.pipe.get_state_vector()), 4)
        self.assertAlmostEqual(port_1.p, 15.3e6, delta=1.0e-5)
        self.assertAlmostEqual(port_2.p, 15.3e6, delta=1.0e-5)
        self.assertAlmostEqual(port_1.T, 293.15)
        self.assertAlmostEqual(port_2.T, 293.15)

    def test_center_resistance_flow_follows_end_pressure_gradient(self) -> None:
        state = self.pipe.get_state_vector()
        state[0] *= 1.01
        state[1] *= 1.01
        self.pipe.set_state_vector(state)

        forward = self.pipe.resistance_mass_flow()
        state[0] /= 1.01 * 1.01
        state[1] /= 1.01 * 1.01
        state[2] *= 1.01
        state[3] *= 1.01
        self.pipe.set_state_vector(state)
        reverse = self.pipe.resistance_mass_flow()

        self.assertGreater(forward, 0.0)
        self.assertLess(reverse, 0.0)

    def test_connection_derivatives_conserve_internal_center_flow_mass(self) -> None:
        state = self.pipe.get_state_vector()
        state[0] *= 1.01
        state[1] *= 1.01
        self.pipe.set_state_vector(state)

        d1, d2 = self.pipe.derivatives_from_connections(
            port_1_m_flow=0.2,
            connected_h_1=self.pipe.properties_1().h + 1000.0,
            port_2_m_flow=-0.1,
            connected_h_2=self.pipe.properties_2().h - 1000.0,
        )

        self.assertAlmostEqual(d1.m + d2.m, 0.1)


class AmesimPnl00rPipeTests(unittest.TestCase):
    def setUp(self) -> None:
        self.pipe = AmesimPnl00rPipe(
            name="pneumatic_100",
            diameter_mm=14.0,
            length_m=1.0,
            relative_roughness=0.045 / 14.0,
        )

    def test_stateless_resistance_flow_follows_pressure_gradient(self) -> None:
        forward = self.pipe.mass_flow(
            port_1_pressure_pa=15.31e6,
            port_1_temperature_k=293.15,
            port_2_pressure_pa=15.29e6,
            port_2_temperature_k=293.15,
        )
        reverse = self.pipe.mass_flow(
            port_1_pressure_pa=15.29e6,
            port_1_temperature_k=293.15,
            port_2_pressure_pa=15.31e6,
            port_2_temperature_k=293.15,
        )

        self.assertGreater(forward, 0.0)
        self.assertLess(reverse, 0.0)
        self.assertAlmostEqual(abs(forward), abs(reverse), delta=abs(forward) * 0.01)

    def test_diagnostics_expose_darcy_terms(self) -> None:
        diagnostics = self.pipe.diagnostics(
            mass_flow_kg_s=1.0e-4,
            pressure_pa=15.3e6,
            temperature_k=293.15,
        )

        self.assertGreater(diagnostics.reynolds_number, 0.0)
        self.assertGreater(diagnostics.gas_velocity_m_s, 0.0)
        self.assertGreater(diagnostics.pressure_drop_pa, 0.0)


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