from __future__ import annotations

import unittest

from app.main import compile_reactflow_network
from app.simulation.components.experimental.junctions.tee import Tee
from app.simulation.components.experimental.storage.cylinder import Cylinder
from app.simulation.components.experimental.storage.tank import Tank
from app.simulation.core.base import Component
from app.simulation.core.medium import IdealGasMedium
from app.simulation.core.ports import PortDefinition, PortState, PortVariableDefinition
from app.simulation.systems.network import SimulationNetwork
from tests.test_system_xml_protocol import physical_connection_project


class ComponentInterfaceTests(unittest.TestCase):
    def test_pneumatic_port_exposes_equation_roles_and_signed_flow(self) -> None:
        port = PortState.pneumatic("port_a", nominal_role="bidirectional")
        assert port.definition is not None

        self.assertEqual(
            [
                (variable.name, variable.role, variable.connection_rule)
                for variable in port.definition.variables
            ],
            [
                ("p", "effort", "equal"),
                ("m_flow", "flow", "sumToZero"),
                ("h_outflow", "stream", "streamMix"),
                ("volume", "signal", "directed"),
                ("volume_flow", "signal", "directed"),
            ],
        )

        port.m_flow = 10.0
        self.assertEqual(port.actual_direction(), "in")
        self.assertEqual(port.inflow_rate, 10.0)
        self.assertEqual(port.outflow_rate, 0.0)

        port.m_flow = -4.0
        self.assertEqual(port.actual_direction(), "out")
        self.assertEqual(port.inflow_rate, 0.0)
        self.assertEqual(port.outflow_rate, 4.0)

    def test_physical_connection_is_neutral_and_canonical(self) -> None:
        forward = compile_reactflow_network(physical_connection_project())
        reverse = compile_reactflow_network(
            physical_connection_project(reverse_edge=True)
        )

        self.assertEqual(
            forward.connections[0].undirected_key,
            reverse.connections[0].undirected_key,
        )
        self.assertEqual(
            forward.connections[0].as_interface_dict(),
            reverse.connections[0].as_interface_dict(),
        )
        self.assertNotIn("source", forward.connections[0].as_interface_dict())

    def test_compiler_instantiates_parameters_and_actual_edges(self) -> None:
        project = physical_connection_project()
        network = compile_reactflow_network(project)

        self.assertEqual(set(network.components), {"cylinder_1", "tank_1"})
        self.assertEqual(network.components["cylinder_1"].V, 0.01)
        self.assertEqual(network.components["tank_1"].V, 0.1)
        self.assertEqual(len(network.connections), 1)
        self.assertEqual(network.connections[0].id, "edge-1")
        self.assertEqual(network.connections[0].kind, "physical")
        self.assertEqual(network.connections[0].domain, "pneumatic")
        equation_system = network.as_interface_dict()["pressureFlowSystem"]
        self.assertEqual(equation_system["unknownCount"], 4)
        self.assertEqual(equation_system["equationCount"], 4)
        self.assertTrue(equation_system["isSquare"])

    def test_physical_port_requires_a_tee_for_branching(self) -> None:
        medium = IdealGasMedium()
        network = SimulationNetwork("branch-check")
        network.add_component(Cylinder("cylinder_1", medium))
        network.add_component(Tank("tank_1", medium))
        network.add_component(Tank("tank_2", medium))
        network.connect("cylinder_1", "port_b", "tank_1", "port_a")

        with self.assertRaisesRegex(ValueError, "Use a junction component"):
            network.connect("cylinder_1", "port_b", "tank_2", "port_a")

    def test_connection_requires_the_same_variable_contract(self) -> None:
        medium = IdealGasMedium()
        network = SimulationNetwork("contract-check")
        network.add_component(Cylinder("cylinder_1", medium))
        incompatible = Component("incompatible_1")
        incompatible.register_port(
            PortState(
                definition=PortDefinition(
                    name="port_a",
                    kind="physical",
                    domain="pneumatic",
                    nominal_role="bidirectional",
                    positive_flow_direction="intoComponent",
                    variables=(PortVariableDefinition("p", "effort", "equal"),),
                )
            )
        )
        network.add_component(incompatible)

        with self.assertRaisesRegex(ValueError, "variable contracts do not match"):
            network.connect("cylinder_1", "port_b", "incompatible_1", "port_a")

    def test_connection_equations_declare_endpoint_neutral_constraints(self) -> None:
        medium = IdealGasMedium()
        network = SimulationNetwork("residual-check")
        cylinder = Cylinder("cylinder_1", medium)
        tank = Tank("tank_1", medium)
        network.add_component(cylinder)
        network.add_component(tank)
        network.connect("tank_1", "port_a", "cylinder_1", "port_b")

        cylinder.port_b.p = 2.0e5
        cylinder.port_b.m_flow = -3.0
        tank.port_a.p = 2.0e5
        tank.port_a.m_flow = 3.0

        residuals = {
            residual.variables[0].rsplit(".", 1)[-1]: residual
            for residual in network.connection_equation_definitions()
        }
        self.assertEqual(set(residuals), {"p", "m_flow"})
        self.assertEqual(residuals["p"].relation, "equal")
        self.assertEqual(residuals["m_flow"].relation, "sumToZero")
        self.assertFalse(hasattr(residuals["p"], "value"))

    def test_tee_ports_are_physical_and_bidirectional(self) -> None:
        tee = Tee("tee_1")
        self.assertEqual(
            {definition.name for definition in tee.port_definitions},
            {"port_in", "port_out1", "port_out2"},
        )
        self.assertTrue(
            all(
                definition.kind == "physical"
                and definition.nominal_role == "bidirectional"
                for definition in tee.port_definitions
            )
        )




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