from __future__ import annotations

import unittest

from PythonModels.core.solver import SolveIVPConfig
from PythonModels.systems.test_mql import TestMqlSystem


class TestMqlPneumaticChamberSegmentTests(unittest.TestCase):
    def setUp(self) -> None:
        self.system = TestMqlSystem()
        self.discovery = self.system.discover_pneumatic_branch_topology()
        self.spec = self.discovery.chamber_segment_specs[0]

    def test_discovers_four_fixed_chamber_segments_from_real_topology(self) -> None:
        self.assertEqual(len(self.discovery.chamber_segment_specs), 4)
        self.assertEqual(self.spec.volume_alias, "pn_general_chamber")
        self.assertEqual(self.spec.inlet_node_alias, "pn_node3_8")
        self.assertEqual(self.spec.inlet_line_alias, "pneumatic_96")
        self.assertEqual(self.spec.inlet_orifice_alias, "pn_orifice_18")
        self.assertEqual(self.spec.volume_inlet_port, "port_2")
        self.assertEqual(self.spec.outlet_orifice_alias, "pn_orifice_19")
        self.assertEqual(self.spec.outlet_line_alias, "pneumatic_97")
        self.assertEqual(self.spec.outlet_node_alias, "pn_node3_9")
        self.assertEqual(self.spec.source, "amesim-cir-topology")

    def test_segment_closure_writes_real_amesim_port_directions(self) -> None:
        closure = self.system.pneumatic_chamber_segment_closure_from_spec(
            self.spec,
            inlet_pressure_pa=16.0e6,
            outlet_pressure_pa=1.0e6,
        )

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

        self.assertGreater(snapshot.inlet_flow, 0.0)
        self.assertGreater(snapshot.outlet_flow, 0.0)
        self.assertAlmostEqual(
            closure.components.volume.port_b.m_flow,
            snapshot.inlet_flow,
        )
        self.assertAlmostEqual(
            closure.components.volume.port_a.m_flow,
            -snapshot.outlet_flow,
        )
        self.assertAlmostEqual(
            closure.components.inlet_orifice.port_b.m_flow,
            snapshot.inlet_flow,
        )
        self.assertAlmostEqual(
            closure.components.outlet_orifice.port_a.m_flow,
            -snapshot.outlet_flow,
        )
        self.assertEqual(len(rhs), 2)
        self.assertAlmostEqual(rhs[0], snapshot.inlet_flow - snapshot.outlet_flow)

    def test_segment_closure_handles_reverse_boundary_pressures(self) -> None:
        closure = self.system.pneumatic_chamber_segment_closure_from_spec(
            self.spec,
            inlet_pressure_pa=1.0e6,
            outlet_pressure_pa=16.0e6,
        )

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

        self.assertLess(snapshot.inlet_flow, 0.0)
        self.assertLess(snapshot.outlet_flow, 0.0)
        self.assertAlmostEqual(rhs[0], snapshot.inlet_flow - snapshot.outlet_flow)

    def test_simulates_topology_derived_segment_with_prescribed_boundaries(self) -> None:
        solution = self.system.simulate_pneumatic_chamber_segment_from_spec(
            self.spec,
            inlet_pressure_pa=16.0e6,
            outlet_pressure_pa=14.0e6,
            config=SolveIVPConfig(
                t_start=0.0,
                t_stop=1.0e-5,
                max_step=1.0e-6,
            ),
            t_eval=[0.0, 5.0e-6, 1.0e-5],
        )

        self.assertTrue(solution.success)
        self.assertEqual(len(solution.t), 3)
        self.assertEqual(len(solution.y), 2)
        self.assertEqual([len(row) for row in solution.y], [3, 3])


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