from __future__ import annotations

from collections.abc import Mapping
from math import sqrt

from app.simulation.core.base import AlgebraicComponent
from app.simulation.core.catalog import ComponentDisplaySpec, PortDisplaySpec
from app.simulation.core.equations import EquationResidual
from app.simulation.core.metadata import ParameterDefinition
from app.simulation.core.medium import IdealGasMedium
from app.simulation.core.ports import PortDefinition


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

    MODEL_TYPE = "orifice"
    MODEL_VERSION = "1.0.0"
    PRESSURE_FLOW_DEPENDS_ON_STREAM = False
    PRESSURE_FLOW_EXACT_SUM_TO_ZERO_EQUATION_SUFFIXES = frozenset(
        ("mass_flow_balance",)
    )
    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 = ()
    DISPLAY = ComponentDisplaySpec(
        label="孔板/阀门",
        library_id="experimental",
        category_id="flow",
        symbol="orifice",
        ports=(
            PortDisplaySpec("port_a", "left", order=10),
            PortDisplaySpec("port_b", "right", order=20),
        ),
        order=40,
    )

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

    @classmethod
    def create(
        cls,
        *,
        name: str,
        medium: IdealGasMedium,
        parameters: Mapping[str, float],
    ) -> Orifice:
        return cls(
            name=name,
            opening=parameters["opening"],
            K=parameters["K"],
        )

    @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"]