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.medium import IdealGasMedium
from PythonModels.core.ports import PortState


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

    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.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 = PortState.pneumatic("port_a", nominal_role="inlet")
        self.port_a.p = p0
        self.port_a.h_outflow = initial_h
        self.register_port(self.port_a)

        self.port_b = PortState.pneumatic("port_b", nominal_role="outlet")
        self.port_b.p = p0
        self.port_b.h_outflow = initial_h
        self.register_port(self.port_b)

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