from __future__ import annotations

from collections.abc import Mapping
from math import pi

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


class ResistivePipe(AlgebraicComponent):
    """Quasi-steady Darcy resistance used by topology-driven simulation."""

    MODEL_TYPE = "pipe"
    PORTS = (
        PortDefinition.pneumatic("port_a", nominal_role="inlet"),
        PortDefinition.pneumatic("port_b", nominal_role="outlet"),
    )
    PARAMETERS = (
        ParameterDefinition(
            "length",
            5.0,
            label="长度",
            quantity="length",
            unit="m",
            minimum=0.0,
            minimum_exclusive=True,
        ),
        ParameterDefinition(
            "diameter",
            0.02,
            label="直径",
            quantity="length",
            unit="m",
            minimum=0.0,
            minimum_exclusive=True,
        ),
        ParameterDefinition(
            "lambda_darcy",
            0.02,
            label="摩阻系数",
            minimum=0.0,
        ),
        ParameterDefinition(
            "p0",
            1e5,
            label="初始压力",
            quantity="pressure",
            unit="Pa",
            minimum=0.0,
            minimum_exclusive=True,
        ),
        ParameterDefinition(
            "T0",
            300.0,
            label="初始温度",
            quantity="temperature",
            unit="K",
            minimum=0.0,
            minimum_exclusive=True,
        ),
    )
    RESULT_VARIABLES = ()

    def __init__(
        self,
        name: str,
        medium: IdealGasMedium,
        L: float = 5.0,
        D: float = 0.02,
        lambda_darcy: float = 0.02,
        p0: float = 1e5,
        T0: float = 300.0,
    ) -> None:
        super().__init__(name=name)
        self.set_parameter_values(
            {
                "length": L,
                "diameter": D,
                "lambda_darcy": lambda_darcy,
                "p0": p0,
                "T0": T0,
            }
        )
        self.medium = medium
        self.L = L
        self.D = D
        self.lambda_darcy = lambda_darcy
        self.p0 = p0
        self.T0 = T0
        self.area = pi * D * D / 4.0
        initial_h = medium.specific_enthalpy(T0)

        self.port_a = self.register_declared_port("port_a")
        self.port_a.p = p0
        self.port_a.h_outflow = initial_h

        self.port_b = self.register_declared_port("port_b")
        self.port_b.p = p0
        self.port_b.h_outflow = initial_h

    def pressure_drop(self, m_flow_a: float, p_a: float, p_b: float) -> float:
        average_pressure = max(0.5 * (p_a + p_b), 1.0)
        density = max(self.medium.density(average_pressure, self.T0), 1e-12)
        resistance = self.lambda_darcy * (self.L / self.D)
        return (
            resistance
            * m_flow_a
            * abs(m_flow_a)
            / (2.0 * density * self.area * self.area)
        )

    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}:darcy_pressure_loss",
                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="effort",
                value=(
                    self.port_a.p
                    - self.port_b.p
                    - self.pressure_drop(
                        self.port_a.m_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"]