from __future__ import annotations

from collections.abc import Mapping
from math import sqrt

from PythonModels.core.base import AlgebraicComponent
from PythonModels.core.equations import EquationResidual
from PythonModels.core.metadata import ParameterDefinition
from PythonModels.core.ports import PortDefinition


class Orifice(AlgebraicComponent):
    """Python port of ModelicaModels.Myorifice."""

    MODEL_TYPE = "orifice"
    PORTS = (
        PortDefinition.pneumatic("port_a", nominal_role="inlet"),
        PortDefinition.pneumatic("port_b", nominal_role="outlet"),
    )
    PARAMETERS = (
        ParameterDefinition(
            "K",
            1e-5,
            label="流量系数",
            quantity="flow_coefficient",
            unit="kg/(s*Pa^0.5)",
            minimum=0.0,
        ),
        ParameterDefinition(
            "opening",
            1.0,
            label="开度",
            minimum=0.0,
            maximum=1.0,
        ),
    )
    RESULT_VARIABLES = ()

    def __init__(self, name: str, opening: float = 1.0, K: float = 1e-5) -> None:
        super().__init__(name=name)
        self.set_parameter_values({"K": K, "opening": opening})
        self.opening = opening
        self.K = K
        self.port_a = self.register_declared_port("port_a")
        self.port_b = self.register_declared_port("port_b")

    @property
    def K_eff(self) -> float:
        return self.K * max(self.opening, 0.001)

    def mass_flow(self, p_a: float, p_b: float) -> float:
        dp = p_a - p_b
        if dp == 0.0:
            return 0.0
        return self.K_eff * sqrt(abs(dp)) * (1.0 if dp > 0.0 else -1.0)

    def pressure_flow_equation_residuals(self) -> tuple[EquationResidual, ...]:
        return (
            EquationResidual(
                id=f"{self.name}:mass_flow_balance",
                owner="component",
                owner_id=self.name,
                relation="sumToZero",
                variables=(
                    f"{self.name}.port_a.m_flow",
                    f"{self.name}.port_b.m_flow",
                ),
                role="flow",
                value=self.port_a.m_flow + self.port_b.m_flow,
            ),
            EquationResidual(
                id=f"{self.name}:pressure_flow_relation",
                owner="component",
                owner_id=self.name,
                relation="constitutive",
                variables=(
                    f"{self.name}.port_a.p",
                    f"{self.name}.port_b.p",
                    f"{self.name}.port_a.m_flow",
                ),
                role="flow",
                value=self.port_a.m_flow
                - self.mass_flow(self.port_a.p, self.port_b.p),
            ),
        )

    def update_stream_outflows(self, connected_h: Mapping[str, float]) -> None:
        self.port_a.h_outflow = connected_h["port_b"]
        self.port_b.h_outflow = connected_h["port_a"]