from __future__ import annotations

from dataclasses import dataclass
from math import pi, sqrt

from PythonModels.core.base import AlgebraicComponent, DynamicComponent
from PythonModels.core.medium import ThermodynamicProperties
from PythonModels.core.peng_robinson import HELIUM_PR, PengRobinsonFluid
from PythonModels.core.ports import PortState
from PythonModels.core.state import VolumeState


@dataclass(frozen=True)
class AmesimPneumaticGas:
    """Caloric constants plus Peng-Robinson EOS for AMESim pneumatic components."""

    fluid: PengRobinsonFluid = HELIUM_PR
    cp: float = 5193.0
    cv: float = 3116.0

    @property
    def gamma(self) -> float:
        return self.cp / self.cv

    @property
    def R_gas(self) -> float:
        return self.fluid.specific_gas_constant

    def density(self, pressure: float, temperature: float) -> float:
        return self.fluid.density(pressure, temperature)

    def pressure(self, density: float, temperature: float) -> float:
        return self.fluid.pressure_from_density(temperature, density)

    def specific_internal_energy(self, temperature: float) -> float:
        return self.cv * temperature

    def specific_enthalpy(self, temperature: float) -> float:
        return self.cp * temperature

    def temperature_from_internal_energy(self, specific_internal_energy: float) -> float:
        if self.cv <= 0.0:
            raise ValueError("cv must be positive.")
        return specific_internal_energy / self.cv


HELIUM_PNEUMATIC_GAS = AmesimPneumaticGas()


def liters_to_m3(value: float) -> float:
    return value * 1.0e-3


def m3_to_cm3(value: float) -> float:
    return value * 1.0e6


def cm3_to_m3(value: float) -> float:
    return value * 1.0e-6


def kg_to_g(value: float) -> float:
    return value * 1.0e3


def mm2_to_m2(value: float) -> float:
    return value * 1.0e-6


def diameter_mm_to_area_m2(diameter_mm: float) -> float:
    diameter_m = diameter_mm * 1.0e-3
    return pi * diameter_m * diameter_m / 4.0


class AmesimPneumaticVolume(DynamicComponent):
    """First-pass AMESim pneumatic control volume using helium PR pressure closure."""

    def __init__(
        self,
        name: str,
        volume: float,
        gas: AmesimPneumaticGas = HELIUM_PNEUMATIC_GAS,
        p0: float = 101_325.0,
        T0: float = 293.15,
        heat_transfer_coefficient: float = 0.0,
        heat_transfer_area: float = 0.0,
        external_temperature_k: float = 293.15,
    ) -> None:
        if volume <= 0.0:
            raise ValueError("volume must be positive.")
        if heat_transfer_coefficient < 0.0:
            raise ValueError("heat_transfer_coefficient must be non-negative.")
        if heat_transfer_area < 0.0:
            raise ValueError("heat_transfer_area must be non-negative.")
        if external_temperature_k <= 0.0:
            raise ValueError("external_temperature_k must be positive.")
        super().__init__(name=name)
        self.volume = volume
        self.gas = gas
        self.heat_transfer_coefficient = heat_transfer_coefficient
        self.heat_transfer_area = heat_transfer_area
        self.external_temperature = external_temperature_k
        rho0 = gas.density(p0, T0)
        m0 = rho0 * volume
        U0 = m0 * gas.specific_internal_energy(T0)
        self.state = VolumeState(m=m0, U=U0)
        self.port_a = PortState()
        self.port_b = PortState()

    @classmethod
    def from_liters(
        cls,
        name: str,
        volume_liters: float,
        gas: AmesimPneumaticGas = HELIUM_PNEUMATIC_GAS,
        p0: float = 101_325.0,
        T0: float = 293.15,
        heat_transfer_coefficient: float = 0.0,
        heat_transfer_area: float = 0.0,
        external_temperature_k: float = 293.15,
    ) -> "AmesimPneumaticVolume":
        return cls(
            name=name,
            volume=liters_to_m3(volume_liters),
            gas=gas,
            p0=p0,
            T0=T0,
            heat_transfer_coefficient=heat_transfer_coefficient,
            heat_transfer_area=heat_transfer_area,
            external_temperature_k=external_temperature_k,
        )

    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 volume_cm3(self) -> float:
        return m3_to_cm3(self.volume)

    def volume_rate_m3_s(self) -> float:
        return 0.0

    def thermal_energy_flow_w(self, temperature_k: float | None = None) -> float:
        temperature = self.properties().T if temperature_k is None else temperature_k
        return (
            self.heat_transfer_coefficient
            * self.heat_transfer_area
            * (self.external_temperature - temperature)
        )

    def gas_mass_g(self) -> float:
        return kg_to_g(self.state.m)

    def pressure_gauge_pa(self, reference_pressure_pa: float = 101_300.0) -> float:
        return self.properties().p - reference_pressure_pa

    def properties(self) -> ThermodynamicProperties:
        if self.state.m <= 0.0:
            raise ValueError("volume mass must stay positive.")
        T = self.gas.temperature_from_internal_energy(self.state.U / self.state.m)
        rho = self.state.m / self.volume
        p = self.gas.pressure(rho, T)
        u = self.state.U / self.state.m
        h = self.gas.specific_enthalpy(T)
        self.port_a.p = p
        self.port_a.h_outflow = h
        self.port_b.p = p
        self.port_b.h_outflow = h
        return ThermodynamicProperties(p=p, T=T, rho=rho, u=u, h=h)

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

    def derivatives_from_two_connections(
        self,
        *,
        port_a_m_flow: float,
        connected_h_a: float,
        port_b_m_flow: float,
        connected_h_b: float,
        internal_h: float,
        volume_rate_m3_s: float | None = None,
    ) -> VolumeState:
        properties = self.properties()
        inlet_h_a = self.connection_inlet_enthalpy(
            port_m_flow=port_a_m_flow,
            connected_h=connected_h_a,
            internal_h=internal_h,
        )
        inlet_h_b = self.connection_inlet_enthalpy(
            port_m_flow=port_b_m_flow,
            connected_h=connected_h_b,
            internal_h=internal_h,
        )
        return VolumeState(
            m=port_a_m_flow + port_b_m_flow,
            U=(
                port_a_m_flow * inlet_h_a
                + port_b_m_flow * inlet_h_b
                + self.thermal_energy_flow_w(properties.T)
                - properties.p * (
                    self.volume_rate_m3_s()
                    if volume_rate_m3_s is None
                    else volume_rate_m3_s
                )
            ),
        )


class AmesimVariablePneumaticVolume(AmesimPneumaticVolume):
    """PNCH012-style volume with a dead volume plus an external moving volume."""

    def __init__(
        self,
        name: str,
        dead_volume: float,
        gas: AmesimPneumaticGas = HELIUM_PNEUMATIC_GAS,
        p0: float = 101_325.0,
        T0: float = 293.15,
        external_volume: float = 0.0,
        heat_transfer_coefficient: float = 0.0,
        heat_transfer_area: float = 0.0,
        external_temperature_k: float = 293.15,
    ) -> None:
        if dead_volume <= 0.0:
            raise ValueError("dead_volume must be positive.")
        if dead_volume + external_volume <= 0.0:
            raise ValueError("total volume must be positive.")
        self.dead_volume = dead_volume
        self.external_volume = external_volume
        self.external_volume_rate = 0.0
        super().__init__(
            name=name,
            volume=dead_volume + external_volume,
            gas=gas,
            p0=p0,
            T0=T0,
            heat_transfer_coefficient=heat_transfer_coefficient,
            heat_transfer_area=heat_transfer_area,
            external_temperature_k=external_temperature_k,
        )

    @classmethod
    def from_liters(
        cls,
        name: str,
        dead_volume_liters: float,
        gas: AmesimPneumaticGas = HELIUM_PNEUMATIC_GAS,
        p0: float = 101_325.0,
        T0: float = 293.15,
        external_volume_liters: float = 0.0,
        heat_transfer_coefficient: float = 0.0,
        heat_transfer_area: float = 0.0,
        external_temperature_k: float = 293.15,
    ) -> "AmesimVariablePneumaticVolume":
        return cls(
            name=name,
            dead_volume=liters_to_m3(dead_volume_liters),
            gas=gas,
            p0=p0,
            T0=T0,
            external_volume=liters_to_m3(external_volume_liters),
            heat_transfer_coefficient=heat_transfer_coefficient,
            heat_transfer_area=heat_transfer_area,
            external_temperature_k=external_temperature_k,
        )

    def volume_rate_m3_s(self) -> float:
        return self.external_volume_rate

    def set_external_volume_m3(
        self,
        external_volume: float,
        external_volume_rate_m3_s: float = 0.0,
    ) -> None:
        if self.dead_volume + external_volume <= 0.0:
            raise ValueError("total volume must be positive.")
        self.external_volume = external_volume
        self.external_volume_rate = external_volume_rate_m3_s
        self.volume = self.dead_volume + self.external_volume


class AmesimPneumaticOrifice(AlgebraicComponent):
    """First-pass PNOR001/PNVO001-style compressible helium orifice.

    This is a calibrated placeholder boundary for the Python port. It preserves
    AMESim-style area and coefficient inputs, but final parity must be checked
    against AMESim CSV results before treating it as numerically equivalent.
    """

    def __init__(
        self,
        name: str,
        area: float,
        flow_coefficient: float = 1.0,
        gas: AmesimPneumaticGas = HELIUM_PNEUMATIC_GAS,
        opening: float = 1.0,
    ) -> None:
        if area < 0.0:
            raise ValueError("area must be non-negative.")
        if flow_coefficient < 0.0:
            raise ValueError("flow_coefficient must be non-negative.")
        super().__init__(name=name)
        self.area = area
        self.flow_coefficient = flow_coefficient
        self.gas = gas
        self.opening = opening
        self.port_a = PortState()
        self.port_b = PortState()

    @classmethod
    def from_mm2(
        cls,
        name: str,
        area_mm2: float,
        flow_coefficient: float = 1.0,
        gas: AmesimPneumaticGas = HELIUM_PNEUMATIC_GAS,
        opening: float = 1.0,
    ) -> "AmesimPneumaticOrifice":
        return cls(
            name=name,
            area=mm2_to_m2(area_mm2),
            flow_coefficient=flow_coefficient,
            gas=gas,
            opening=opening,
        )

    @property
    def effective_area(self) -> float:
        opening = min(max(self.opening, 0.0), 1.0)
        return self.area * opening

    def mass_flow(self, p_a: float, p_b: float, upstream_temperature: float) -> float:
        if p_a == p_b or self.effective_area == 0.0 or self.flow_coefficient == 0.0:
            return 0.0
        if p_a > p_b:
            return compressible_orifice_mass_flow(
                upstream_pressure=p_a,
                downstream_pressure=p_b,
                upstream_temperature=upstream_temperature,
                area=self.effective_area,
                flow_coefficient=self.flow_coefficient,
                gas=self.gas,
            )
        return -compressible_orifice_mass_flow(
            upstream_pressure=p_b,
            downstream_pressure=p_a,
            upstream_temperature=upstream_temperature,
            area=self.effective_area,
            flow_coefficient=self.flow_coefficient,
            gas=self.gas,
        )


def compressible_orifice_mass_flow(
    *,
    upstream_pressure: float,
    downstream_pressure: float,
    upstream_temperature: float,
    area: float,
    flow_coefficient: float,
    gas: AmesimPneumaticGas = HELIUM_PNEUMATIC_GAS,
) -> float:
    if upstream_pressure <= 0.0 or downstream_pressure < 0.0:
        raise ValueError("pressures must be non-negative and upstream pressure must be positive.")
    if upstream_temperature <= 0.0:
        raise ValueError("upstream_temperature must be positive.")
    if area < 0.0 or flow_coefficient < 0.0:
        raise ValueError("area and flow_coefficient must be non-negative.")
    if downstream_pressure >= upstream_pressure or area == 0.0 or flow_coefficient == 0.0:
        return 0.0

    gamma = gas.gamma
    pressure_ratio = max(downstream_pressure / upstream_pressure, 0.0)
    critical_ratio = (2.0 / (gamma + 1.0)) ** (gamma / (gamma - 1.0))
    coefficient = flow_coefficient * area * upstream_pressure / sqrt(gas.R_gas * upstream_temperature)
    if pressure_ratio <= critical_ratio:
        flow_function = sqrt(gamma) * (2.0 / (gamma + 1.0)) ** ((gamma + 1.0) / (2.0 * (gamma - 1.0)))
    else:
        term = pressure_ratio ** (2.0 / gamma) - pressure_ratio ** ((gamma + 1.0) / gamma)
        flow_function = sqrt((2.0 * gamma / (gamma - 1.0)) * max(term, 0.0))
    return coefficient * flow_function