from __future__ import annotations

from dataclasses import dataclass
from math import isfinite
from typing import Protocol, Sequence

from app.simulation.core.errors import RecoverableTrialStateError
from app.simulation.performance import profile_property


@dataclass(frozen=True)
class ThermodynamicProperties:
    p: float
    T: float
    rho: float
    u: float
    h: float


@dataclass(frozen=True)
class ThermodynamicPropertyTangents:
    """Directional derivatives of a recovered thermodynamic state."""

    p: tuple[float, ...]
    T: tuple[float, ...]
    rho: tuple[float, ...]
    u: tuple[float, ...]
    h: tuple[float, ...]

    @property
    def width(self) -> int:
        return len(self.p)

    @classmethod
    def zeros(cls, width: int) -> "ThermodynamicPropertyTangents":
        values = (0.0,) * width
        return cls(p=values, T=values, rho=values, u=values, h=values)


@dataclass(frozen=True)
class ThermodynamicPropertiesLinearization:
    """Primal properties and a validity-checked directional linearization."""

    properties: ThermodynamicProperties
    tangents: ThermodynamicPropertyTangents
    valid: bool = True
    reason: str | None = None


class GasMedium(Protocol):
    """Thermodynamic contract required by pneumatic components.

    ``IdealGasMedium`` is the default implementation.  Keeping the component
    boundary structural allows a later helium/Peng-Robinson implementation to
    be registered without changing every AMESim component constructor.
    """

    name: str
    R_gas: float
    cp_ref: float
    T_ref: float

    @property
    def cv(self) -> float: ...

    @property
    def gamma(self) -> float: ...

    def cp_at_temperature(self, T: float) -> float: ...

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

    def density(self, p: float, T: float) -> float: ...

    def isentropic_density_pressure_factor(
        self,
        p: float,
        T: float,
        downstream_pressure: float | None = None,
    ) -> float: ...

    def dynamic_viscosity(self, T: float) -> float: ...

    def specific_internal_energy(self, T: float) -> float: ...

    def specific_internal_energy_at_pressure(self, p: float, T: float) -> float: ...

    def specific_enthalpy(self, T: float) -> float: ...

    def specific_enthalpy_at_pressure(self, p: float, T: float) -> float: ...

    def temperature_from_internal_energy(self, u: float) -> float: ...

    def temperature_from_enthalpy(self, h: float) -> float: ...

    def temperature_from_pressure_enthalpy(self, p: float, h: float) -> float: ...

    def temperature_from_mass_internal_energy(self, m: float, U: float) -> float: ...

    def pressure(self, m: float, T: float, V: float) -> float: ...

    def properties_from_mU(
        self,
        m: float,
        U: float,
        V: float,
    ) -> ThermodynamicProperties: ...

    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: ...


@dataclass(frozen=True)
class IdealGasMedium:
    """Temperature-dependent ideal-gas air approximation.

    This is still not a strict clone of `Modelica.Media.Air.SimpleAir`.
    The small linear `cp(T)` term is kept configurable for calibration, but the
    current default is calibrated against the committed Testmodel baseline and
    therefore falls back to the constant-heat-capacity limit.
    """

    name: str = "SimpleAirApprox"
    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

    @property
    def cv(self) -> float:
        return self.cv_at_temperature(self.T_ref)

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

    def cp_at_temperature(self, T: float) -> float:
        return self.cp_ref + self.cp_slope * (T - self.T_ref)

    def cv_at_temperature(self, T: float) -> float:
        return self.cp_at_temperature(T) - self.R_gas

    @profile_property("density")
    def density(self, p: float, T: float) -> float:
        return p / (self.R_gas * T)

    @profile_property("isentropic_density_pressure_factor")
    def isentropic_density_pressure_factor(
        self,
        p: float,
        T: float,
        downstream_pressure: float | None = None,
    ) -> float:
        del p
        del downstream_pressure
        cp = self.cp_at_temperature(T)
        cv = self.cv_at_temperature(T)
        return cv / cp

    @profile_property("dynamic_viscosity")
    def dynamic_viscosity(self, T: float) -> float:
        """Return dynamic viscosity using the default air Sutherland law."""

        if T <= 0.0:
            raise ValueError("Temperature must be positive.")
        return (
            self.viscosity_ref
            * (T / self.viscosity_T_ref) ** 1.5
            * (self.viscosity_T_ref + self.sutherland_constant)
            / (T + self.sutherland_constant)
        )

    @profile_property("specific_internal_energy")
    def specific_internal_energy(self, T: float) -> float:
        delta_T = T - self.T_ref
        return (
            self.cv * self.T_ref
            + self.cv * delta_T
            + 0.5 * self.cp_slope * delta_T * delta_T
        )

    @profile_property("specific_internal_energy_at_pressure")
    def specific_internal_energy_at_pressure(self, p: float, T: float) -> float:
        del p
        return self.specific_internal_energy(T)

    @profile_property("specific_enthalpy")
    def specific_enthalpy(self, T: float) -> float:
        delta_T = T - self.T_ref
        return (
            self.cp_ref * self.T_ref
            + self.cp_ref * delta_T
            + 0.5 * self.cp_slope * delta_T * delta_T
        )

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

    def temperature_from_internal_energy(self, u: float) -> float:
        reference_internal_energy = self.cv * self.T_ref
        delta_u = u - reference_internal_energy

        if abs(self.cp_slope) <= 1e-15:
            return self.T_ref + delta_u / self.cv

        a = 0.5 * self.cp_slope
        b = self.cv
        c = -delta_u
        discriminant = max(b * b - 4.0 * a * c, 0.0)
        positive_root = (-b + discriminant**0.5) / (2.0 * a)
        negative_root = (-b - discriminant**0.5) / (2.0 * a)
        delta_T = positive_root if abs(positive_root) <= abs(negative_root) else negative_root
        return self.T_ref + delta_T

    def temperature_from_enthalpy(self, h: float) -> float:
        reference_enthalpy = self.cp_ref * self.T_ref
        delta_h = h - reference_enthalpy

        if abs(self.cp_slope) <= 1e-15:
            return self.T_ref + delta_h / self.cp_ref

        a = 0.5 * self.cp_slope
        b = self.cp_ref
        c = -delta_h
        discriminant = max(b * b - 4.0 * a * c, 0.0)
        positive_root = (-b + discriminant**0.5) / (2.0 * a)
        negative_root = (-b - discriminant**0.5) / (2.0 * a)
        delta_T = positive_root if abs(positive_root) <= abs(negative_root) else negative_root
        return self.T_ref + delta_T

    @profile_property("temperature_from_pressure_enthalpy")
    def temperature_from_pressure_enthalpy(self, p: float, h: float) -> float:
        del p
        return self.temperature_from_enthalpy(h)

    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)

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

    @profile_property("properties_from_mU")
    def properties_from_mU(self, m: float, U: float, V: float) -> ThermodynamicProperties:
        T = self.temperature_from_mass_internal_energy(m, U)
        p = self.pressure(m, T, V)
        rho = m / V
        u = U / m
        h = self.specific_enthalpy(T)
        return ThermodynamicProperties(p=p, T=T, rho=rho, u=u, h=h)

    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:
        """Linearize properties_from_mU for several seed directions."""

        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)
        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 ThermodynamicPropertiesLinearization(
                properties=props,
                tangents=ThermodynamicPropertyTangents.zeros(width),
                valid=False,
                reason="properties_primal_mismatch",
            )
        if not all(
            isfinite(value)
            for values in (dm_values, dU_values, dV_values)
            for value in values
        ):
            return ThermodynamicPropertiesLinearization(
                properties=props,
                tangents=ThermodynamicPropertyTangents.zeros(width),
                valid=False,
                reason="non_finite_tangent_input",
            )

        cv = self.cv_at_temperature(props.T)
        cp = self.cp_at_temperature(props.T)
        if not isfinite(cv) or not isfinite(cp) or cv <= 0.0 or cp <= 0.0:
            return ThermodynamicPropertiesLinearization(
                properties=props,
                tangents=ThermodynamicPropertyTangents.zeros(width),
                valid=False,
                reason="non_positive_heat_capacity",
            )

        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 / cv
            pressure_tangent = self.R_gas * (
                props.T * density_tangent + props.rho * temperature_tangent
            )
            enthalpy_tangent = cp * temperature_tangent
            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)
        valid = all(isfinite(value) for value in tangent_values)
        return ThermodynamicPropertiesLinearization(
            properties=props,
            tangents=ThermodynamicPropertyTangents(
                p=tuple(dp),
                T=tuple(dT),
                rho=tuple(drho),
                u=tuple(du),
                h=tuple(dh),
            ),
            valid=valid,
            reason=None if valid else "non_finite_property_tangent",
        )