from __future__ import annotations

from dataclasses import dataclass

from app.simulation.examples.test_mql.system import COMPONENT_SPECS


@dataclass(frozen=True)
class TestMqlPneumaticNode3Balance:
    temperature_k: float
    pressure_pa: float
    port_1_enthalpy_flow_w: float
    port_1_mass_flow_g_s: float
    port_1_volume_derivative_l_min: float
    port_1_volume_cm3: float
    port_2_enthalpy_flow_w: float
    port_2_mass_flow_g_s: float
    port_2_volume_derivative_l_min: float
    port_2_volume_cm3: float
    port_3_enthalpy_flow_w: float
    port_3_mass_flow_g_s: float
    port_3_volume_derivative_l_min: float
    port_3_volume_cm3: float


@dataclass(frozen=True)
class TestMqlPneumaticNode3:
    """Exact algebraic contract of AMESim ``PN3NODE2``.

    Pressure and temperature are fixed by port 2 and duplicated to ports 1 and
    3. Flow and volume signals at port 2 are the sums of ports 1 and 3, matching
    the ``EXPRESS2`` equations stored in ``test_mql_.cir``.
    """

    alias: str

    def balance(
        self,
        *,
        port_2_temperature_k: float,
        port_2_pressure_pa: float,
        port_1_enthalpy_flow_w: float,
        port_1_mass_flow_g_s: float,
        port_3_enthalpy_flow_w: float,
        port_3_mass_flow_g_s: float,
        port_1_volume_derivative_l_min: float = 0.0,
        port_1_volume_cm3: float = 0.0,
        port_3_volume_derivative_l_min: float = 0.0,
        port_3_volume_cm3: float = 0.0,
    ) -> TestMqlPneumaticNode3Balance:
        if port_2_temperature_k <= 0.0:
            raise ValueError("port_2_temperature_k must be positive")
        if port_2_pressure_pa <= 0.0:
            raise ValueError("port_2_pressure_pa must be positive")
        return TestMqlPneumaticNode3Balance(
            temperature_k=port_2_temperature_k,
            pressure_pa=port_2_pressure_pa,
            port_1_enthalpy_flow_w=port_1_enthalpy_flow_w,
            port_1_mass_flow_g_s=port_1_mass_flow_g_s,
            port_1_volume_derivative_l_min=port_1_volume_derivative_l_min,
            port_1_volume_cm3=port_1_volume_cm3,
            port_2_enthalpy_flow_w=(
                port_1_enthalpy_flow_w + port_3_enthalpy_flow_w
            ),
            port_2_mass_flow_g_s=port_1_mass_flow_g_s + port_3_mass_flow_g_s,
            port_2_volume_derivative_l_min=(
                port_1_volume_derivative_l_min + port_3_volume_derivative_l_min
            ),
            port_2_volume_cm3=port_1_volume_cm3 + port_3_volume_cm3,
            port_3_enthalpy_flow_w=port_3_enthalpy_flow_w,
            port_3_mass_flow_g_s=port_3_mass_flow_g_s,
            port_3_volume_derivative_l_min=port_3_volume_derivative_l_min,
            port_3_volume_cm3=port_3_volume_cm3,
        )


@dataclass(frozen=True)
class TestMqlPneumaticNode4Balance:
    temperature_k: float
    pressure_pa: float
    port_1_enthalpy_flow_w: float
    port_1_mass_flow_g_s: float
    port_1_volume_derivative_l_min: float
    port_1_volume_cm3: float
    port_2_enthalpy_flow_w: float
    port_2_mass_flow_g_s: float
    port_2_volume_derivative_l_min: float
    port_2_volume_cm3: float
    port_3_enthalpy_flow_w: float
    port_3_mass_flow_g_s: float
    port_3_volume_derivative_l_min: float
    port_3_volume_cm3: float
    port_4_enthalpy_flow_w: float
    port_4_mass_flow_g_s: float
    port_4_volume_derivative_l_min: float
    port_4_volume_cm3: float


@dataclass(frozen=True)
class TestMqlPneumaticNode4:
    """Exact algebraic contract of AMESim ``P4NODE2``.

    Pressure and temperature are fixed by port 2 and duplicated to ports 1, 3,
    and 4. Flow and volume signals at port 2 are the sums of ports 1, 3, and
    4, matching the saved AMESim variables for ``pnnode4_*`` instances.
    """

    alias: str

    def balance(
        self,
        *,
        port_2_temperature_k: float,
        port_2_pressure_pa: float,
        port_1_enthalpy_flow_w: float,
        port_1_mass_flow_g_s: float,
        port_3_enthalpy_flow_w: float,
        port_3_mass_flow_g_s: float,
        port_4_enthalpy_flow_w: float,
        port_4_mass_flow_g_s: float,
        port_1_volume_derivative_l_min: float = 0.0,
        port_1_volume_cm3: float = 0.0,
        port_3_volume_derivative_l_min: float = 0.0,
        port_3_volume_cm3: float = 0.0,
        port_4_volume_derivative_l_min: float = 0.0,
        port_4_volume_cm3: float = 0.0,
    ) -> TestMqlPneumaticNode4Balance:
        if port_2_temperature_k <= 0.0:
            raise ValueError("port_2_temperature_k must be positive")
        if port_2_pressure_pa <= 0.0:
            raise ValueError("port_2_pressure_pa must be positive")
        return TestMqlPneumaticNode4Balance(
            temperature_k=port_2_temperature_k,
            pressure_pa=port_2_pressure_pa,
            port_1_enthalpy_flow_w=port_1_enthalpy_flow_w,
            port_1_mass_flow_g_s=port_1_mass_flow_g_s,
            port_1_volume_derivative_l_min=port_1_volume_derivative_l_min,
            port_1_volume_cm3=port_1_volume_cm3,
            port_2_enthalpy_flow_w=(
                port_1_enthalpy_flow_w
                + port_3_enthalpy_flow_w
                + port_4_enthalpy_flow_w
            ),
            port_2_mass_flow_g_s=(
                port_1_mass_flow_g_s
                + port_3_mass_flow_g_s
                + port_4_mass_flow_g_s
            ),
            port_2_volume_derivative_l_min=(
                port_1_volume_derivative_l_min
                + port_3_volume_derivative_l_min
                + port_4_volume_derivative_l_min
            ),
            port_2_volume_cm3=(
                port_1_volume_cm3 + port_3_volume_cm3 + port_4_volume_cm3
            ),
            port_3_enthalpy_flow_w=port_3_enthalpy_flow_w,
            port_3_mass_flow_g_s=port_3_mass_flow_g_s,
            port_3_volume_derivative_l_min=port_3_volume_derivative_l_min,
            port_3_volume_cm3=port_3_volume_cm3,
            port_4_enthalpy_flow_w=port_4_enthalpy_flow_w,
            port_4_mass_flow_g_s=port_4_mass_flow_g_s,
            port_4_volume_derivative_l_min=port_4_volume_derivative_l_min,
            port_4_volume_cm3=port_4_volume_cm3,
        )


@dataclass(frozen=True)
class TestMqlP4NodePortConnection:
    line_alias: str
    local_node_alias: str
    local_port: str
    remote_node_alias: str
    remote_port: str


@dataclass(frozen=True)
class TestMqlP4NodePrimaryConnection:
    line_alias: str
    node_alias: str
    node_port: str
    chamber_alias: str
    chamber_port: str


@dataclass(frozen=True)
class TestMqlP4NodeOrificeConnection:
    orifice_alias: str
    node_alias: str
    node_port: str
    direct_line_alias: str


@dataclass(frozen=True)
class TestMqlP4NodeNeighborhood:
    node_alias: str
    primary: TestMqlP4NodePrimaryConnection
    port_1: TestMqlP4NodePortConnection
    port_3: TestMqlP4NodePortConnection
    port_4: TestMqlP4NodeOrificeConnection


def build_test_mql_node3_assembly() -> dict[str, TestMqlPneumaticNode3]:
    return {
        str(spec["alias"]): TestMqlPneumaticNode3(alias=str(spec["alias"]))
        for spec in COMPONENT_SPECS
        if spec["submodel"] == "PN3NODE2"
    }

def build_test_mql_node4_assembly() -> dict[str, TestMqlPneumaticNode4]:
    return {
        str(spec["alias"]): TestMqlPneumaticNode4(alias=str(spec["alias"]))
        for spec in COMPONENT_SPECS
        if spec["submodel"] == "P4NODE2"
    }