from __future__ import annotations

from dataclasses import dataclass

from PythonModels.systems.test_mql 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,
        )


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"
    }