from __future__ import annotations

import unittest

from app.simulation.components.amesim.junctions.nodes import (
    AmesimP4Node2,
    AmesimPn3Node2,
)
from app.simulation.core.medium import IdealGasMedium


class AmesimPneumaticNodeComponentTests(unittest.TestCase):
    def test_pn3node2_declares_bidirectional_pneumatic_ports(self) -> None:
        node = AmesimPn3Node2.create(
            name="pn3_1",
            medium=IdealGasMedium(),
            parameters={},
        )

        self.assertEqual(set(node.ports), {"port_1", "port_2", "port_3"})
        self.assertTrue(
            all(
                definition.domain == "pneumatic"
                and definition.nominal_role == "bidirectional"
                for definition in node.port_definitions
            )
        )
        self.assertEqual(node.parameter_values, {})

    def test_pn3node2_residuals_use_port_2_as_pressure_reference(self) -> None:
        node = AmesimPn3Node2("pn3_1")
        node.port_1.p = 101000.0
        node.port_2.p = 100000.0
        node.port_3.p = 99000.0
        node.port_1.m_flow = 0.2
        node.port_2.m_flow = -0.3
        node.port_3.m_flow = 0.1

        residuals = {residual.id.rsplit(":", 1)[1]: residual for residual in node.pressure_flow_equation_residuals()}

        self.assertEqual(set(residuals), {"port_1_pressure_reference", "port_3_pressure_reference", "mass_flow_balance"})
        self.assertEqual(residuals["port_1_pressure_reference"].value, 1000.0)
        self.assertEqual(residuals["port_3_pressure_reference"].value, -1000.0)
        self.assertAlmostEqual(residuals["mass_flow_balance"].value, 0.0)

    def test_p4node2_residuals_cover_four_ports(self) -> None:
        node = AmesimP4Node2("p4_1")
        node.port_1.p = 100000.0
        node.port_2.p = 100000.0
        node.port_3.p = 100000.0
        node.port_4.p = 100000.0
        node.port_1.m_flow = 0.1
        node.port_2.m_flow = 0.2
        node.port_3.m_flow = -0.4
        node.port_4.m_flow = 0.1

        residuals = node.pressure_flow_equation_residuals()

        self.assertEqual(len(residuals), 4)
        self.assertTrue(all(residual.value == 0.0 for residual in residuals))
        self.assertEqual(
            residuals[-1].variables,
            (
                "p4_1.port_1.m_flow",
                "p4_1.port_2.m_flow",
                "p4_1.port_3.m_flow",
                "p4_1.port_4.m_flow",
            ),
        )

    def test_node_uses_port_2_temperature_for_non_reference_outlets(self) -> None:
        node = AmesimP4Node2("p4_1")
        node.port_1.m_flow = 0.25
        node.port_2.m_flow = 0.75
        node.port_3.m_flow = -0.5
        node.port_4.m_flow = -0.5

        node.update_stream_outflows(
            {
                "port_1": 100.0,
                "port_2": 300.0,
                "port_3": 900.0,
                "port_4": 1200.0,
            }
        )

        self.assertEqual(node.temperature_reference_h, 300.0)
        self.assertEqual(node.port_1.h_outflow, 300.0)
        self.assertEqual(node.port_2.h_outflow, 250.0)
        self.assertEqual(node.port_3.h_outflow, 300.0)
        self.assertEqual(node.port_4.h_outflow, 300.0)

    def test_port_2_outlet_closes_node_energy_balance(self) -> None:
        node = AmesimP4Node2("p4_1")
        node.port_1.m_flow = -0.01823
        node.port_2.m_flow = -2.56077
        node.port_3.m_flow = 2.579
        node.port_4.m_flow = 0.0

        node.update_stream_outflows(
            {
                "port_1": 2_630_000.0,
                "port_2": 140_900.0,
                "port_3": 63_960.0,
                "port_4": 0.0,
            }
        )

        self.assertEqual(node.port_1.h_outflow, 140_900.0)
        energy_flow = sum(
            port.m_flow
            * (
                {
                    "port_1": 2_630_000.0,
                    "port_2": 140_900.0,
                    "port_3": 63_960.0,
                    "port_4": 0.0,
                }[name]
                if port.m_flow > 0.0
                else port.h_outflow
            )
            for name, port in node.ports.items()
        )
        self.assertAlmostEqual(energy_flow, 0.0, places=10)


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