from __future__ import annotations

import unittest

from PythonModels.components.amesim_pneumatic import (
    AmesimPneumaticOrifice,
    AmesimPneumaticVolume,
)
from PythonModels.systems.test_mql_closure import (
    TestMqlPneumaticBranchComponents,
    TestMqlPneumaticClosure,
)


class TestMqlPneumaticClosureTests(unittest.TestCase):
    def _make_closure(
        self,
        *,
        upstream_pressure: float = 15.3e6,
        downstream_pressure: float = 1.0e5,
    ) -> TestMqlPneumaticClosure:
        upstream = AmesimPneumaticVolume.from_liters(
            name="source_chamber",
            volume_liters=57.0,
            p0=upstream_pressure,
            T0=293.15,
        )
        downstream = AmesimPneumaticVolume.from_liters(
            name="target_chamber",
            volume_liters=15.0,
            p0=downstream_pressure,
            T0=293.15,
        )
        orifice = AmesimPneumaticOrifice.from_mm2(
            name="branch_orifice",
            area_mm2=78.5,
            flow_coefficient=0.9,
        )

        def initial_state_vector() -> list[float]:
            return [*upstream.get_state_vector(), *downstream.get_state_vector()]

        def apply_state_vector(values: list[float]) -> None:
            if len(values) != 4:
                raise ValueError("two-volume closure state vector requires four values")
            upstream.set_state_vector(values[:2])
            downstream.set_state_vector(values[2:])

        return TestMqlPneumaticClosure(
            components=TestMqlPneumaticBranchComponents(
                name="source_to_target",
                upstream_volume=upstream,
                orifice=orifice,
                downstream_volume=downstream,
            ),
            initial_state_vector=initial_state_vector,
            apply_state_vector=apply_state_vector,
        )

    def test_snapshot_writes_modelica_style_port_flow_signs(self) -> None:
        closure = self._make_closure()

        snapshot = closure.snapshot()

        self.assertGreater(snapshot.flow, 0.0)
        self.assertAlmostEqual(
            closure.components.upstream_volume.port_a.m_flow,
            -snapshot.flow,
        )
        self.assertAlmostEqual(
            closure.components.orifice.port_a.m_flow,
            snapshot.flow,
        )
        self.assertAlmostEqual(
            closure.components.orifice.port_b.m_flow,
            -snapshot.flow,
        )
        self.assertAlmostEqual(
            closure.components.downstream_volume.port_a.m_flow,
            snapshot.flow,
        )
        self.assertAlmostEqual(
            closure.components.orifice.port_a.p,
            snapshot.upstream.p,
        )
        self.assertAlmostEqual(
            closure.components.orifice.port_b.p,
            snapshot.downstream.p,
        )

    def test_rhs_uses_canonical_flow_for_volume_mass_balance(self) -> None:
        closure = self._make_closure()
        state = closure.initial_state_vector()

        snapshot = closure.snapshot(state)
        rhs = closure.rhs(state)

        self.assertEqual(len(rhs), 4)
        self.assertAlmostEqual(rhs[0], -snapshot.flow)
        self.assertAlmostEqual(rhs[2], snapshot.flow)
        self.assertAlmostEqual(rhs[0] + rhs[2], 0.0)
        self.assertAlmostEqual(rhs[1], -snapshot.flow * snapshot.upstream.h)
        self.assertAlmostEqual(rhs[3], snapshot.flow * snapshot.upstream.h)

    def test_reverse_pressure_reverses_canonical_flow_and_port_signs(self) -> None:
        closure = self._make_closure(
            upstream_pressure=1.0e5,
            downstream_pressure=15.3e6,
        )

        snapshot = closure.snapshot()
        rhs = closure.rhs(closure.initial_state_vector())

        self.assertLess(snapshot.flow, 0.0)
        self.assertGreater(closure.components.upstream_volume.port_a.m_flow, 0.0)
        self.assertLess(closure.components.downstream_volume.port_a.m_flow, 0.0)
        self.assertAlmostEqual(rhs[0], -snapshot.flow)
        self.assertAlmostEqual(rhs[2], snapshot.flow)
        self.assertAlmostEqual(rhs[0] + rhs[2], 0.0)
        self.assertAlmostEqual(rhs[1], -snapshot.flow * snapshot.downstream.h)
        self.assertAlmostEqual(rhs[3], snapshot.flow * snapshot.downstream.h)

    def test_snapshot_can_apply_explicit_state_vector(self) -> None:
        closure = self._make_closure()
        state = closure.initial_state_vector()
        state[0] *= 0.99

        snapshot = closure.snapshot(state)

        self.assertAlmostEqual(
            closure.components.upstream_volume.state.m,
            state[0],
        )
        self.assertGreater(snapshot.upstream.p, snapshot.downstream.p)


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