from __future__ import annotations

from collections.abc import Mapping

from PythonModels.core.base import DynamicComponent
from PythonModels.core.equations import EquationResidual
from PythonModels.core.medium import IdealGasMedium, ThermodynamicProperties
from PythonModels.core.ports import PortState
from PythonModels.core.state import VolumeState


class Tank(DynamicComponent):
    """Python port of ModelicaModels.Mytank."""

    def __init__(
        self,
        name: str,
        medium: IdealGasMedium,
        V: float = 0.1,
        p0: float = 1e5,
        T0: float = 300.0,
    ) -> None:
        super().__init__(name=name)
        self.medium = medium
        self.V = V
        m0 = p0 * V / (medium.R_gas * T0)
        U0 = m0 * medium.specific_internal_energy(T0)
        self.state = VolumeState(m=m0, U=U0)
        self.port_a = self.register_port(
            PortState.pneumatic("port_a", nominal_role="inlet")
        )

    def get_state_vector(self) -> list[float]:
        return self.state.as_vector()

    def set_state_vector(self, values: list[float]) -> None:
        self.state = VolumeState.from_vector(values)

    def properties(self) -> ThermodynamicProperties:
        props = self.medium.properties_from_mU(self.state.m, self.state.U, self.V)
        self.port_a.p = props.p
        self.port_a.h_outflow = props.h
        return props

    def refresh_thermodynamic_ports(self) -> ThermodynamicProperties:
        return self.properties()

    def state_derivative_from_ports(
        self,
        connected_h: Mapping[str, float],
    ) -> list[float]:
        properties = self.properties()
        derivative = self.derivatives_from_connection(
            connected_h=connected_h["port_a"],
            port_m_flow=self.port_a.m_flow,
            internal_h=properties.h,
        )
        return derivative.as_vector()

    def pressure_flow_equation_residuals(self) -> tuple[EquationResidual, ...]:
        pressure = self.medium.properties_from_mU(
            self.state.m,
            self.state.U,
            self.V,
        ).p
        return (
            EquationResidual(
                id=f"{self.name}:port_a_pressure_state",
                owner="component",
                owner_id=self.name,
                relation="state",
                variables=(f"{self.name}.port_a.p", f"{self.name}.state"),
                role="effort",
                value=self.port_a.p - pressure,
            ),
        )

    def derivatives_from_connection(
        self,
        *,
        connected_h: float,
        port_m_flow: float,
        internal_h: float,
    ) -> VolumeState:
        inlet_h = self.connection_inlet_enthalpy(
            port_m_flow=port_m_flow,
            connected_h=connected_h,
            internal_h=internal_h,
        )
        return self.derivatives(inlet_h, port_m_flow)

    def derivatives(self, inlet_h: float, m_flow: float) -> VolumeState:
        return VolumeState(m=m_flow, U=m_flow * inlet_h)