from __future__ import annotations

from app.simulation.core.errors import RecoverableTrialStateError

from dataclasses import dataclass
from math import acos, cos, isfinite, log, pi, sqrt

UNIVERSAL_GAS_CONSTANT = 8.31446261815324
# Simcenter Amesim 2404 ``sag_reinit_eos_`` keeps more digits than the
# commonly printed Peng-Robinson constants 0.45724 and 0.07780.
PENG_ROBINSON_A_COEFFICIENT = 0.457235583
PENG_ROBINSON_B_COEFFICIENT = 0.07779607


@dataclass(frozen=True)
class PengRobinsonFluid:
    """Pure-fluid Peng-Robinson equation-of-state helper.

    The class covers the equation-of-state layer plus the enthalpy departure
    needed to compare AMESim pneumatic ``pn2hpti`` reference enthalpy flows.
    """

    name: str
    molar_mass: float
    critical_temperature: float
    critical_pressure: float
    acentric_factor: float

    @property
    def specific_gas_constant(self) -> float:
        return UNIVERSAL_GAS_CONSTANT / self.molar_mass

    @property
    def a_parameter(self) -> float:
        return (
            PENG_ROBINSON_A_COEFFICIENT
            * UNIVERSAL_GAS_CONSTANT
            * UNIVERSAL_GAS_CONSTANT
            * self.critical_temperature
            * self.critical_temperature
            / self.critical_pressure
        )

    @property
    def b_parameter(self) -> float:
        return (
            PENG_ROBINSON_B_COEFFICIENT
            * UNIVERSAL_GAS_CONSTANT
            * self.critical_temperature
            / self.critical_pressure
        )

    @property
    def kappa(self) -> float:
        omega = self.acentric_factor
        return 0.37464 + 1.54226 * omega - 0.26992 * omega * omega

    def alpha(self, temperature: float) -> float:
        self._validate_temperature(temperature)
        reduced_temperature = temperature / self.critical_temperature
        return (1.0 + self.kappa * (1.0 - sqrt(reduced_temperature))) ** 2.0

    def alpha_temperature_derivative(self, temperature: float) -> float:
        self._validate_temperature(temperature)
        reduced_temperature = temperature / self.critical_temperature
        sqrt_reduced_temperature = sqrt(reduced_temperature)
        alpha_base = 1.0 + self.kappa * (1.0 - sqrt_reduced_temperature)
        return -(
            alpha_base
            * self.kappa
            / (self.critical_temperature * sqrt_reduced_temperature)
        )

    def alpha_temperature_second_derivative(self, temperature: float) -> float:
        self._validate_temperature(temperature)
        reduced_temperature = temperature / self.critical_temperature
        sqrt_reduced_temperature = sqrt(reduced_temperature)
        alpha_base = 1.0 + self.kappa * (1.0 - sqrt_reduced_temperature)
        return (
            self.kappa
            / (2.0 * self.critical_temperature * self.critical_temperature)
            * (
                self.kappa / reduced_temperature
                + alpha_base / (reduced_temperature * sqrt_reduced_temperature)
            )
        )

    def attractive_parameter(self, temperature: float) -> float:
        return self.a_parameter * self.alpha(temperature)

    def attractive_parameter_temperature_derivative(self, temperature: float) -> float:
        return self.a_parameter * self.alpha_temperature_derivative(temperature)

    def attractive_parameter_temperature_second_derivative(
        self,
        temperature: float,
    ) -> float:
        return self.a_parameter * self.alpha_temperature_second_derivative(temperature)

    def pressure_from_molar_volume(self, temperature: float, molar_volume: float) -> float:
        self._validate_temperature(temperature)
        if molar_volume <= self.b_parameter:
            raise RecoverableTrialStateError("Molar volume must be larger than Peng-Robinson b parameter.")
        a_alpha = self.attractive_parameter(temperature)
        b = self.b_parameter
        repulsive = UNIVERSAL_GAS_CONSTANT * temperature / (molar_volume - b)
        attractive = a_alpha / (molar_volume * (molar_volume + b) + b * (molar_volume - b))
        return repulsive - attractive

    def pressure_from_density(self, temperature: float, density: float) -> float:
        if density <= 0.0:
            raise ValueError("Density must be positive.")
        return self.pressure_from_molar_volume(temperature, self.molar_mass / density)

    def pressure_temperature_derivative_at_density(
        self,
        temperature: float,
        density: float,
    ) -> float:
        self._validate_temperature(temperature)
        if density <= 0.0:
            raise ValueError("Density must be positive.")
        molar_volume = self.molar_mass / density
        if molar_volume <= self.b_parameter:
            raise RecoverableTrialStateError(
                "Molar volume must be larger than Peng-Robinson b parameter."
            )
        b = self.b_parameter
        denominator = molar_volume * (molar_volume + b) + b * (molar_volume - b)
        return (
            UNIVERSAL_GAS_CONSTANT / (molar_volume - b)
            - self.attractive_parameter_temperature_derivative(temperature) / denominator
        )

    def pressure_density_derivative_at_temperature(
        self,
        temperature: float,
        density: float,
    ) -> float:
        self._validate_temperature(temperature)
        if density <= 0.0:
            raise ValueError("Density must be positive.")
        molar_volume = self.molar_mass / density
        if molar_volume <= self.b_parameter:
            raise RecoverableTrialStateError(
                "Molar volume must be larger than Peng-Robinson b parameter."
            )
        b = self.b_parameter
        denominator = molar_volume * (molar_volume + b) + b * (molar_volume - b)
        pressure_molar_volume_derivative = (
            -UNIVERSAL_GAS_CONSTANT * temperature / (molar_volume - b) ** 2
            + self.attractive_parameter(temperature)
            * 2.0
            * (molar_volume + b)
            / denominator**2
        )
        molar_volume_density_derivative = -self.molar_mass / (density * density)
        return pressure_molar_volume_derivative * molar_volume_density_derivative

    def reduced_parameters(self, pressure: float, temperature: float) -> tuple[float, float]:
        self._validate_pressure_temperature(pressure, temperature)
        a_alpha = self.attractive_parameter(temperature)
        b = self.b_parameter
        A = a_alpha * pressure / (UNIVERSAL_GAS_CONSTANT * UNIVERSAL_GAS_CONSTANT * temperature * temperature)
        B = b * pressure / (UNIVERSAL_GAS_CONSTANT * temperature)
        return A, B

    def compressibility_roots(self, pressure: float, temperature: float) -> tuple[float, ...]:
        A, B = self.reduced_parameters(pressure, temperature)
        coefficients = (
            -(1.0 - B),
            A - 3.0 * B * B - 2.0 * B,
            -(A * B - B * B - B * B * B),
        )
        roots = _real_cubic_roots(*coefficients)
        physical_roots = tuple(sorted(root for root in roots if root > B and isfinite(root)))
        if not physical_roots:
            raise ValueError("Peng-Robinson cubic produced no physical compressibility root.")
        return physical_roots

    def compressibility_factor(
        self,
        pressure: float,
        temperature: float,
        phase: str = "vapor",
    ) -> float:
        roots = self.compressibility_roots(pressure, temperature)
        if phase == "vapor":
            return roots[-1]
        if phase == "liquid":
            return roots[0]
        if phase == "stable-single-root":
            return roots[-1]
        raise ValueError(f"Unsupported phase selector: {phase!r}")

    def molar_volume(
        self,
        pressure: float,
        temperature: float,
        phase: str = "vapor",
    ) -> float:
        z = self.compressibility_factor(pressure, temperature, phase=phase)
        return z * UNIVERSAL_GAS_CONSTANT * temperature / pressure

    def density(
        self,
        pressure: float,
        temperature: float,
        phase: str = "vapor",
    ) -> float:
        return self.molar_mass / self.molar_volume(pressure, temperature, phase=phase)

    def residual_specific_enthalpy(
        self,
        pressure: float,
        temperature: float,
        phase: str = "vapor",
    ) -> float:
        """Return Peng-Robinson enthalpy departure from ideal gas, J/kg."""
        self._validate_pressure_temperature(pressure, temperature)
        z = self.compressibility_factor(pressure, temperature, phase=phase)
        _, B = self.reduced_parameters(pressure, temperature)
        b = self.b_parameter
        attractive = self.attractive_parameter(temperature)
        d_attractive_d_temperature = (
            self.attractive_parameter_temperature_derivative(temperature)
        )
        log_argument = (z + (1.0 + sqrt(2.0)) * B) / (
            z + (1.0 - sqrt(2.0)) * B
        )
        residual_molar_enthalpy = (
            UNIVERSAL_GAS_CONSTANT * temperature * (z - 1.0)
            + (
                temperature * d_attractive_d_temperature
                - attractive
            )
            * log(log_argument)
            / (2.0 * sqrt(2.0) * b)
        )
        return residual_molar_enthalpy / self.molar_mass

    def residual_specific_internal_energy_at_density(
        self,
        temperature: float,
        density: float,
    ) -> float:
        """Return Peng-Robinson internal-energy departure, J/kg."""
        self._validate_temperature(temperature)
        if density <= 0.0:
            raise ValueError("Density must be positive.")
        molar_volume = self.molar_mass / density
        b = self.b_parameter
        if molar_volume <= b:
            raise RecoverableTrialStateError(
                "Molar volume must be larger than Peng-Robinson b parameter."
            )
        attractive = self.attractive_parameter(temperature)
        d_attractive_d_temperature = (
            self.attractive_parameter_temperature_derivative(temperature)
        )
        log_argument = (
            molar_volume + (1.0 + sqrt(2.0)) * b
        ) / (
            molar_volume + (1.0 - sqrt(2.0)) * b
        )
        residual_molar_internal_energy = (
            temperature * d_attractive_d_temperature - attractive
        ) * log(log_argument) / (2.0 * sqrt(2.0) * b)
        return residual_molar_internal_energy / self.molar_mass

    def residual_isochoric_heat_capacity_at_density(
        self,
        temperature: float,
        density: float,
    ) -> float:
        """Return the constant-volume heat-capacity departure, J/kg/K."""
        self._validate_temperature(temperature)
        if density <= 0.0:
            raise ValueError("Density must be positive.")
        molar_volume = self.molar_mass / density
        b = self.b_parameter
        if molar_volume <= b:
            raise RecoverableTrialStateError(
                "Molar volume must be larger than Peng-Robinson b parameter."
            )
        log_argument = (
            molar_volume + (1.0 + sqrt(2.0)) * b
        ) / (
            molar_volume + (1.0 - sqrt(2.0)) * b
        )
        residual_molar_cv = (
            temperature
            * self.attractive_parameter_temperature_second_derivative(temperature)
            * log(log_argument)
            / (2.0 * sqrt(2.0) * b)
        )
        return residual_molar_cv / self.molar_mass

    @staticmethod
    def _validate_temperature(temperature: float) -> None:
        if temperature <= 0.0:
            raise ValueError("Temperature must be positive.")

    @classmethod
    def _validate_pressure_temperature(cls, pressure: float, temperature: float) -> None:
        if pressure <= 0.0:
            raise ValueError("Pressure must be positive.")
        cls._validate_temperature(temperature)

HELIUM_PR = PengRobinsonFluid(
    name="helium",
    molar_mass=0.004002602,
    critical_temperature=5.1953,
    critical_pressure=227_460.0,
    # Simcenter Amesim 2404 helium_eos.data.
    acentric_factor=-0.382,
)

NITROGEN_PR = PengRobinsonFluid(
    name="nitrogen",
    molar_mass=0.0280134,
    critical_temperature=126.192,
    critical_pressure=3.3958e6,
    acentric_factor=0.0372,
)

AIR_PR = PengRobinsonFluid(
    name="air",
    molar_mass=0.02896513,
    critical_temperature=132.5306,
    critical_pressure=3.786e6,
    acentric_factor=0.0335,
)


def _real_cubic_roots(a: float, b: float, c: float) -> tuple[float, ...]:
    """Return real roots for x**3 + a*x**2 + b*x + c = 0."""

    depressed_p = b - a * a / 3.0
    depressed_q = 2.0 * a * a * a / 27.0 - a * b / 3.0 + c
    discriminant = (depressed_q / 2.0) ** 2.0 + (depressed_p / 3.0) ** 3.0
    offset = -a / 3.0
    tolerance = 1e-14

    if discriminant > tolerance:
        sqrt_discriminant = sqrt(discriminant)
        u = _real_cube_root(-depressed_q / 2.0 + sqrt_discriminant)
        v = _real_cube_root(-depressed_q / 2.0 - sqrt_discriminant)
        return (u + v + offset,)

    if abs(discriminant) <= tolerance:
        u = _real_cube_root(-depressed_q / 2.0)
        return tuple(sorted({2.0 * u + offset, -u + offset}))

    if depressed_p >= 0.0:
        raise ValueError("Unexpected cubic state with three real roots and non-negative p.")
    radius = 2.0 * sqrt(-depressed_p / 3.0)
    argument = (3.0 * depressed_q / (2.0 * depressed_p)) * sqrt(-3.0 / depressed_p)
    argument = max(-1.0, min(1.0, argument))
    theta = acos(argument) / 3.0
    roots = [
        radius * cos(theta - 2.0 * pi * index / 3.0) + offset
        for index in range(3)
    ]
    return tuple(sorted(roots))


def _real_cube_root(value: float) -> float:
    if value == 0.0:
        return 0.0
    return (1.0 if value > 0.0 else -1.0) * abs(value) ** (1.0 / 3.0)