from __future__ import annotations

from collections.abc import Callable, Sequence
from dataclasses import dataclass
from math import exp, isfinite, log
from typing import ClassVar

from app.simulation.core.errors import RecoverableTrialStateError
from app.simulation.core.medium import (
    GasMedium,
    IdealGasMedium,
    ThermodynamicProperties,
    ThermodynamicPropertiesLinearization,
    ThermodynamicPropertyTangents,
)
from app.simulation.core.peng_robinson import HELIUM_PR, PengRobinsonFluid
from app.simulation.performance import profile_property, record_property_iterations
from app.simulation.property_cache import cache_property_calculation


@dataclass(frozen=True)
class AmesimIdealAirMedium(IdealGasMedium):
    """AMESim air properties evaluated with the ideal-gas method.

    Substance identity and property method are part of the concrete Python
    type.  A future air correlation or helium Peng-Robinson implementation can
    therefore coexist as a sibling type without turning ``gi`` into a fluid
    enumeration.
    """

    SUBSTANCE_ID: ClassVar[str] = "air"
    PROPERTY_METHOD_ID: ClassVar[str] = "ideal_gas"

    name: str = "AMESimAirIdealGas"
    R_gas: float = 287.0
    cp_ref: float = 1005.0
    T_ref: float = 300.0
    cp_slope: float = 0.0
    viscosity_ref: float = 1.82e-5
    viscosity_T_ref: float = 293.15
    sutherland_constant: float = 110.4


@dataclass(frozen=True)
class AmesimHeliumPengRobinsonMedium(IdealGasMedium):
    """AMESim helium with a Peng-Robinson mechanical equation of state.

    The pressure-density-temperature relation is evaluated by the shared
    ``HELIUM_PR`` fluid.  The caloric reference follows the constant NASA
    polynomial from Simcenter Amesim 2404 ``helium_cp_h_s.data``.
    """

    SUBSTANCE_ID: ClassVar[str] = "helium"
    PROPERTY_METHOD_ID: ClassVar[str] = "peng_robinson"
    fluid: ClassVar[PengRobinsonFluid] = HELIUM_PR
    nasa_cp_over_R: ClassVar[float] = 2.5
    nasa_enthalpy_constant_K: ClassVar[float] = -745.375
    nasa_viscosity_coefficients: ClassVar[tuple[float, float, float, float]] = (
        0.7501594,
        35.76324,
        -2212.129,
        0.9212635,
    )

    name: str = "AMESimHeliumPengRobinson"
    R_gas: float = HELIUM_PR.specific_gas_constant
    cp_ref: float = nasa_cp_over_R * HELIUM_PR.specific_gas_constant
    T_ref: float = 293.15
    cp_slope: float = 0.0
    viscosity_ref: float = 1.96e-5
    viscosity_T_ref: float = 293.15
    sutherland_constant: float = 79.4

    @property
    def cv(self) -> float:
        return (self.nasa_cp_over_R - 1.0) * self.R_gas

    def cv_at_temperature(self, T: float) -> float:
        del T
        return self.cv

    def diagnostic_dynamic_viscosity(self, T: float) -> float:
        """Return the AMESim NASA-table viscosity used by pipe diagnostics.

        pn2pipefr reports Reynolds number with sagum viscosity. Keep this
        separate from dynamic_viscosity so matching that diagnostic cannot
        alter the already-validated pipe flow or friction dynamics.
        """

        if T <= 0.0:
            raise ValueError("Temperature must be positive.")
        a, b, c, d = self.nasa_viscosity_coefficients
        return 1.0e-7 * exp(a * log(T) + b / T + c / (T * T) + d)

    @profile_property("density")
    @cache_property_calculation("density")
    def density(self, p: float, T: float) -> float:
        return self.fluid.density(p, T)

    def _real_heat_capacities(
        self,
        p: float,
        T: float,
    ) -> tuple[float, float, float, float, float]:
        density = self.density(p, T)
        pressure_density_derivative = (
            self.fluid.pressure_density_derivative_at_temperature(
                T,
                density,
            )
        )
        pressure_temperature_derivative = (
            self.fluid.pressure_temperature_derivative_at_density(
                T,
                density,
            )
        )
        cv = (
            self.cv_at_temperature(T)
            + self.fluid.residual_isochoric_heat_capacity_at_density(T, density)
        )
        cp = (
            cv
            + T
            * pressure_temperature_derivative
            * pressure_temperature_derivative
            / (density * density * pressure_density_derivative)
        )
        if cp <= 0.0 or cv <= 0.0:
            raise ValueError("Real-gas heat capacities must be positive.")
        return (
            cp,
            cv,
            density,
            pressure_density_derivative,
            pressure_temperature_derivative,
        )

    def _local_isentropic_density_pressure_factor(
        self,
        p: float,
        T: float,
    ) -> tuple[float, float]:
        cp, cv, density, pressure_density_derivative, pressure_temperature_derivative = (
            self._real_heat_capacities(p, T)
        )
        heat_capacity_ratio = cp / cv
        factor = p / (
            density * pressure_density_derivative * heat_capacity_ratio
        )
        exponent = (
            p
            * (heat_capacity_ratio - 1.0)
            / (
                heat_capacity_ratio
                * T
                * pressure_temperature_derivative
            )
        )
        return factor, exponent

    @profile_property("isentropic_density_pressure_factor")
    @cache_property_calculation("isentropic_density_pressure_factor")
    def isentropic_density_pressure_factor(
        self,
        p: float,
        T: float,
        downstream_pressure: float | None = None,
    ) -> float:
        upstream_factor, isentropic_temperature_exponent = (
            self._local_isentropic_density_pressure_factor(p, T)
        )
        if downstream_pressure is None or downstream_pressure >= p:
            return upstream_factor

        pressure_ratio = max(downstream_pressure / p, 1.0e-12)
        isentropic_temperature = max(
            T * pressure_ratio**isentropic_temperature_exponent,
            2.2,
        )
        downstream_factor, _unused_exponent = (
            self._local_isentropic_density_pressure_factor(
                max(downstream_pressure, 1.0),
                isentropic_temperature,
            )
        )
        # AMESim 2404 saggs_ evaluates the local factor at the upstream
        # state and at an approximate isentropic downstream state.
        return 0.5 * (upstream_factor + downstream_factor)

    def pressure(self, m: float, T: float, V: float) -> float:
        if V <= 0.0:
            raise ValueError("Volume must stay positive.")
        return self.fluid.pressure_from_density(T, m / V)

    @profile_property("specific_internal_energy")
    def specific_internal_energy(self, T: float) -> float:
        return self.R_gas * (
            (self.nasa_cp_over_R - 1.0) * T
            + self.nasa_enthalpy_constant_K
        )

    @profile_property("specific_internal_energy_at_pressure")
    def specific_internal_energy_at_pressure(self, p: float, T: float) -> float:
        density = self.density(p, T)
        return (
            self.specific_internal_energy(T)
            + self.fluid.residual_specific_internal_energy_at_density(T, density)
        )

    @profile_property("specific_enthalpy")
    def specific_enthalpy(self, T: float) -> float:
        return self.R_gas * (
            self.nasa_cp_over_R * T
            + self.nasa_enthalpy_constant_K
        )

    @profile_property("specific_enthalpy_at_pressure")
    def specific_enthalpy_at_pressure(self, p: float, T: float) -> float:
        return self.specific_enthalpy(T) + self.fluid.residual_specific_enthalpy(p, T)

    def temperature_from_internal_energy(self, u: float) -> float:
        return (
            u / self.R_gas - self.nasa_enthalpy_constant_K
        ) / (self.nasa_cp_over_R - 1.0)

    def temperature_from_enthalpy(self, h: float) -> float:
        return (
            h / self.R_gas - self.nasa_enthalpy_constant_K
        ) / self.nasa_cp_over_R

    @profile_property("temperature_from_pressure_enthalpy")
    @cache_property_calculation("temperature_from_pressure_enthalpy")
    def temperature_from_pressure_enthalpy(self, p: float, h: float) -> float:
        temperature = max(self.temperature_from_enthalpy(h), 2.2)
        for _iteration in range(16):
            residual_enthalpy = self.fluid.residual_specific_enthalpy(p, temperature)
            next_temperature = max(
                self.temperature_from_enthalpy(h - residual_enthalpy),
                2.2,
            )
            if abs(next_temperature - temperature) <= 1.0e-10 * max(
                temperature,
                1.0,
            ):
                record_property_iterations(
                    "temperature_from_pressure_enthalpy",
                    _iteration + 1,
                    True,
                )
                return next_temperature
            temperature = next_temperature
        record_property_iterations(
            "temperature_from_pressure_enthalpy",
            16,
            False,
        )
        return temperature

    def temperature_from_mass_internal_energy(self, m: float, U: float) -> float:
        if m <= 0.0:
            raise RecoverableTrialStateError(
                "Mass must stay positive when recovering temperature."
            )
        return self.temperature_from_internal_energy(U / m)

    @profile_property("properties_from_mU")
    @cache_property_calculation("properties_from_mU")
    def properties_from_mU(
        self,
        m: float,
        U: float,
        V: float,
    ) -> ThermodynamicProperties:
        """Recover a real-gas state, reusing exact repeated evaluations.

        Implicit integration asks several component interfaces for the same
        ``(m, U, V)`` state while closing one RHS evaluation and while building
        finite-difference Jacobians. The calculation is pure and its result is
        immutable, so an exact-key bounded cache avoids repeating the
        Peng-Robinson temperature iteration without changing model semantics.
        """
        if m <= 0.0:
            raise RecoverableTrialStateError(
                "Mass must stay positive when recovering temperature."
            )
        if V <= 0.0:
            raise ValueError("Volume must stay positive.")
        density = m / V
        target_internal_energy = U / m
        temperature = max(
            self.temperature_from_internal_energy(target_internal_energy),
            2.2,
        )
        converged = False
        for _iteration in range(16):
            residual_internal_energy = (
                self.fluid.residual_specific_internal_energy_at_density(
                    temperature,
                    density,
                )
            )
            next_temperature = max(
                self.temperature_from_internal_energy(
                    target_internal_energy - residual_internal_energy
                ),
                2.2,
            )
            if abs(next_temperature - temperature) <= 1.0e-10 * max(
                temperature,
                1.0,
            ):
                temperature = next_temperature
                converged = True
                break
            temperature = next_temperature
        record_property_iterations(
            "properties_from_mU",
            _iteration + 1,
            converged,
        )
        pressure = self.fluid.pressure_from_density(temperature, density)
        return ThermodynamicProperties(
            p=pressure,
            T=temperature,
            rho=density,
            u=target_internal_energy,
            h=self.specific_enthalpy_at_pressure(
                pressure,
                temperature,
            ),
        )

    def linearize_properties_from_mU(
        self,
        m: float,
        U: float,
        V: float,
        dm: Sequence[float],
        dU: Sequence[float],
        dV: Sequence[float],
        *,
        properties: ThermodynamicProperties | None = None,
    ) -> ThermodynamicPropertiesLinearization:
        """Implicitly differentiate the Peng-Robinson m/U/V recovery."""

        dm_values = tuple(float(value) for value in dm)
        dU_values = tuple(float(value) for value in dU)
        dV_values = tuple(float(value) for value in dV)
        if not (len(dm_values) == len(dU_values) == len(dV_values)):
            raise ValueError("Thermodynamic tangent vectors must have equal lengths.")
        props = properties or self.properties_from_mU(m, U, V)
        width = len(dm_values)

        def invalid(reason: str) -> ThermodynamicPropertiesLinearization:
            return ThermodynamicPropertiesLinearization(
                properties=props,
                tangents=ThermodynamicPropertyTangents.zeros(width),
                valid=False,
                reason=reason,
            )

        expected_density = m / V
        expected_internal_energy = U / m
        if (
            abs(props.rho - expected_density)
            > 1.0e-12 * max(abs(expected_density), 1.0)
            or abs(props.u - expected_internal_energy)
            > 1.0e-12 * max(abs(expected_internal_energy), 1.0)
        ):
            return invalid("properties_primal_mismatch")
        if not all(
            isfinite(value)
            for values in (dm_values, dU_values, dV_values)
            for value in values
        ):
            return invalid("non_finite_tangent_input")
        if props.T <= 2.2 * (1.0 + 1.0e-10):
            return invalid("temperature_floor_boundary")

        pressure_temperature_derivative = (
            self.fluid.pressure_temperature_derivative_at_density(
                props.T,
                props.rho,
            )
        )
        pressure_density_derivative = (
            self.fluid.pressure_density_derivative_at_temperature(
                props.T,
                props.rho,
            )
        )
        cv = (
            self.cv_at_temperature(props.T)
            + self.fluid.residual_isochoric_heat_capacity_at_density(
                props.T,
                props.rho,
            )
        )
        recovered_internal_energy = (
            self.specific_internal_energy(props.T)
            + self.fluid.residual_specific_internal_energy_at_density(
                props.T,
                props.rho,
            )
        )
        recovery_scale = max(
            abs(props.u),
            abs(cv * props.T) if isfinite(cv) else 0.0,
            1.0,
        )
        if (
            not all(
                isfinite(value)
                for value in (
                    pressure_temperature_derivative,
                    pressure_density_derivative,
                    cv,
                    recovered_internal_energy,
                )
            )
            or cv <= 0.0
        ):
            return invalid("invalid_peng_robinson_derivative")
        if abs(recovered_internal_energy - props.u) > 1.0e-8 * recovery_scale:
            return invalid("properties_recovery_not_converged")

        internal_energy_density_derivative = (
            props.p - props.T * pressure_temperature_derivative
        ) / (props.rho * props.rho)
        drho: list[float] = []
        du: list[float] = []
        dT: list[float] = []
        dp: list[float] = []
        dh: list[float] = []
        for mass_tangent, energy_tangent, volume_tangent in zip(
            dm_values,
            dU_values,
            dV_values,
            strict=True,
        ):
            density_tangent = (
                mass_tangent / V - m * volume_tangent / (V * V)
            )
            internal_energy_tangent = (
                energy_tangent / m - U * mass_tangent / (m * m)
            )
            temperature_tangent = (
                internal_energy_tangent
                - internal_energy_density_derivative * density_tangent
            ) / cv
            pressure_tangent = (
                pressure_temperature_derivative * temperature_tangent
                + pressure_density_derivative * density_tangent
            )
            enthalpy_tangent = (
                internal_energy_tangent
                + pressure_tangent / props.rho
                - props.p * density_tangent / (props.rho * props.rho)
            )
            drho.append(density_tangent)
            du.append(internal_energy_tangent)
            dT.append(temperature_tangent)
            dp.append(pressure_tangent)
            dh.append(enthalpy_tangent)

        tangent_values = (*drho, *du, *dT, *dp, *dh)
        if not all(isfinite(value) for value in tangent_values):
            return invalid("non_finite_property_tangent")
        return ThermodynamicPropertiesLinearization(
            properties=props,
            tangents=ThermodynamicPropertyTangents(
                p=tuple(dp),
                T=tuple(dT),
                rho=tuple(drho),
                u=tuple(du),
                h=tuple(dh),
            ),
        )


@dataclass(frozen=True)
class AmesimGasPropertyModelSpec:
    """A selectable calculation method for one AMESim gas substance."""

    value: int
    label: str
    method_id: str
    factory: Callable[[], GasMedium]
    eos_type: int

    def build_medium(self) -> GasMedium:
        return self.factory()


AMESIM_AIR_IDEAL_GAS_PROPERTY_MODEL = 0
AMESIM_AIR_PROPERTY_MODELS = (
    AmesimGasPropertyModelSpec(
        value=AMESIM_AIR_IDEAL_GAS_PROPERTY_MODEL,
        label="理想气体",
        method_id=AmesimIdealAirMedium.PROPERTY_METHOD_ID,
        factory=AmesimIdealAirMedium,
        eos_type=1,
    ),
)


AMESIM_HELIUM_PENG_ROBINSON_PROPERTY_MODEL = 0
AMESIM_HELIUM_PROPERTY_MODELS = (
    AmesimGasPropertyModelSpec(
        value=AMESIM_HELIUM_PENG_ROBINSON_PROPERTY_MODEL,
        label="Peng–Robinson",
        method_id=AmesimHeliumPengRobinsonMedium.PROPERTY_METHOD_ID,
        factory=AmesimHeliumPengRobinsonMedium,
        eos_type=6,
    ),
)