from __future__ import annotations

from math import sqrt

from PythonModels.core.base import AlgebraicComponent
from PythonModels.core.ports import PortState


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

    def __init__(self, name: str, opening: float = 1.0, K: float = 1e-7) -> None:
        super().__init__(name=name)
        self.opening = opening
        self.K = K
        self.port_a = PortState()
        self.port_b = PortState()

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