# By default, Julia/LLVM does not use fused multiply-add operations (FMAs).
# Since these FMAs can increase the performance of many numerical algorithms,
# we need to opt-in explicitly.
# See https://ranocha.de/blog/Optimizing_EC_Trixi for further details.
@muladd begin
#! format: noindent

@doc raw"""
    CompressibleNavierStokesDiffusion1D(equations; mu, Pr,
                                        gradient_variables=GradientVariablesPrimitive())

Contains the diffusion (i.e. parabolic) terms applied
to mass, momenta, and total energy together with the advective terms from
the [`CompressibleEulerEquations1D`](@ref).

- `equations`: instance of the [`CompressibleEulerEquations1D`](@ref)
- `mu`: dynamic viscosity,
- `Pr`: Prandtl number,
- `gradient_variables`: which variables the gradients are taken with respect to.
                        Defaults to [`GradientVariablesPrimitive()`](@ref).
                        For an entropy stable formulation, use [`GradientVariablesEntropy()`](@ref).

Fluid properties such as the dynamic viscosity ``\mu`` can be provided in any consistent unit system, e.g.,
[``\mu``] = kg m⁻¹ s⁻¹.
The viscosity ``\mu`` may be a constant or a function of the current state, e.g.,
depending on temperature (Sutherland's law): ``\mu = \mu(T)``.
In the latter case, the function `mu` needs to have the signature `mu(u, equations)`.

The particular form of the compressible Navier-Stokes implemented is
```math
\frac{\partial}{\partial t}
\begin{pmatrix}
\rho \\ \rho v_1 \\ \rho e_{\text{total}}
\end{pmatrix}
+
\frac{\partial}{\partial x}
\begin{pmatrix}
 \rho v_1 \\ \rho v_1^2 + p \\ (\rho e_{\text{total}} + p) v_1
\end{pmatrix}
=
\frac{\partial}{\partial x}
\begin{pmatrix}
0 \\ \tau \\ \tau v_1 - q
\end{pmatrix}
```
where the system is closed with the ideal gas assumption giving
```math
p = (\gamma - 1) \left( \rho e_{\text{total}} - \frac{1}{2} \rho v_1^2 \right)
```
as the pressure. The value of the adiabatic constant `gamma` is taken from the [`CompressibleEulerEquations1D`](@ref).
The terms on the right hand side of the system above
are built from the viscous stress
```math
\tau = \mu \frac{\partial}{\partial x} v_1
```
where the heat flux is
```math
q = -\kappa \frac{\partial}{\partial x} \left(T\right),\quad T = \frac{p}{R\rho}
```
where ``T`` is the temperature and ``\kappa`` is the thermal conductivity for Fourier's law.
Under the assumption that the gas has a constant Prandtl number,
the thermal conductivity is
```math
\kappa = \frac{\gamma \mu R}{(\gamma - 1)\textrm{Pr}}.
```
From this combination of temperature ``T`` and thermal conductivity ``\kappa`` we see
that the gas constant `R` cancels and the heat flux becomes
```math
q = -\kappa \frac{\partial}{\partial x} \left(T\right) = -\frac{\gamma \mu}{(\gamma - 1)\textrm{Pr}} \frac{\partial}{\partial x} \left(\frac{p}{\rho}\right)
```
which is the form implemented below in the [`flux`](@ref) function.

In one spatial dimensions we require gradients for two quantities, e.g.,
primitive quantities
```math
\frac{\partial}{\partial x} v_1,\, \frac{\partial}{\partial x} T
```
or the entropy variables
```math
\frac{\partial}{\partial x} w_2,\, \frac{\partial}{\partial x} w_3
```
where
```math
w_2 = \frac{\rho v_1}{p},\, w_3 = -\frac{\rho}{p}
```
"""
struct CompressibleNavierStokesDiffusion1D{GradientVariables, RealT <: Real, Mu,
                                           E <: AbstractCompressibleEulerEquations{1}} <:
       AbstractCompressibleNavierStokesDiffusion{1, 3, GradientVariables}
    # TODO: parabolic
    # Add NGRADS as a type parameter here and in AbstractEquationsParabolic, add `ngradients(...)` accessor function

    mu::Mu                     # viscosity
    Pr::RealT                  # Prandtl number
    kappa::RealT               # thermal diffusivity for Fourier's law
    max_1_kappa::RealT         # max(1, kappa) used for parabolic cfl => `max_diffusivity`

    equations_hyperbolic::E    # CompressibleEulerEquations1D
    gradient_variables::GradientVariables # GradientVariablesPrimitive or GradientVariablesEntropy
end

# default to primitive gradient variables
function CompressibleNavierStokesDiffusion1D(equations::CompressibleEulerEquations1D;
                                             mu, Prandtl,
                                             gradient_variables = GradientVariablesPrimitive())
    @unpack gamma, inv_gamma_minus_one = equations

    Pr = promote_type(typeof(gamma), typeof(Prandtl))(Prandtl)
    # Under the assumption of constant Prandtl number the thermal conductivity
    # constant is kappa = gamma μ / ((gamma-1) Prandtl).
    # Important note! Factor of μ is accounted for later in `flux`.
    # This avoids recomputation of kappa for non-constant μ.
    kappa = gamma * inv_gamma_minus_one / Pr

    return CompressibleNavierStokesDiffusion1D{typeof(gradient_variables),
                                               typeof(Pr), typeof(mu),
                                               typeof(equations)}(mu, Pr, kappa,
                                                                  max(one(kappa),
                                                                      kappa),
                                                                  equations,
                                                                  gradient_variables)
end

# Together with our specialization of `Adapt.adapt_structure`,
# this allows to move semidiscretizations and their components including
# the equations to GPUs and adapt the floating point type, e.g.,
# to `Float32` to improve performance on GPUs.
function Base.similar(equations::CompressibleNavierStokesDiffusion1D,
                      ::Type{NewRealT}) where {NewRealT}
    mu = equations.mu isa Real ? convert(NewRealT, equations.mu) : equations.mu
    return CompressibleNavierStokesDiffusion1D(similar(equations.equations_hyperbolic,
                                                       NewRealT);
                                               mu = mu,
                                               Prandtl = convert(NewRealT,
                                                                 equations.Pr),
                                               gradient_variables = equations.gradient_variables)
end

# TODO: parabolic
# This is the flexibility a user should have to select the different gradient variable types
# varnames(::typeof(cons2prim)   , ::CompressibleNavierStokesDiffusion1D) = ("v1", "T")
# varnames(::typeof(cons2entropy), ::CompressibleNavierStokesDiffusion1D) = ("w2", "w3")

function varnames(variable_mapping,
                  equations_parabolic::CompressibleNavierStokesDiffusion1D)
    return varnames(variable_mapping, equations_parabolic.equations_hyperbolic)
end

# we specialize this function to compute gradients of primitive variables instead of
# conservative variables.
function gradient_variable_transformation(::CompressibleNavierStokesDiffusion1D{GradientVariablesPrimitive})
    return cons2prim_temperature
end
function gradient_variable_transformation(::CompressibleNavierStokesDiffusion1D{GradientVariablesEntropy})
    return cons2entropy
end

# Explicit formulas for the diffusive Navier-Stokes fluxes are available, e.g., in Section 2
# of the paper by Rueda-Ramírez, Hennemann, Hindenlang, Winters, and Gassner
# "An Entropy Stable Nodal Discontinuous Galerkin Method for the resistive
#  MHD Equations. Part II: Subcell Finite Volume Shock Capturing"
# where one sets the magnetic field components equal to 0.
function flux(u, gradients, orientation::Integer,
              equations::CompressibleNavierStokesDiffusion1D)
    # Here, `u` is assumed to be the "transformed" variables specified by `gradient_variable_transformation`.
    v1, _ = convert_transformed_to_velocity_temperature(u, equations)
    # Here `gradients` is assumed to contain the gradients of the primitive variables (rho, v1, T)
    # either computed directly or reverse engineered from the gradient of the entropy variables
    # by way of the `convert_gradient_variables` function.
    _, dv1dx, dTdx = convert_derivative_to_primitive(u, gradients[1], equations)

    # Viscous stress (tensor)
    tau_11 = dv1dx

    # Fourier's law q = -kappa * grad(T) = -kappa * grad(p / (R rho))
    # with thermal diffusivity constant kappa = gamma μ R / ((gamma-1) Pr)
    # Note, the gas constant cancels under this formulation, so it is not present
    # in the implementation
    q1 = equations.kappa * dTdx

    # In the simplest cases, the user passed in `mu` or `mu()`
    # (which returns just a constant) but
    # more complex functions like Sutherland's law are possible.
    # `dynamic_viscosity` is a helper function that handles both cases
    # by dispatching on the type of `equations.mu`.
    mu = dynamic_viscosity(u, equations)

    # parabolic flux components in the x-direction
    f1 = 0
    f2 = tau_11 * mu
    f3 = (v1 * tau_11 + q1) * mu

    return SVector(f1, f2, f3)
end

@doc raw"""
    max_diffusivity(u, equations_parabolic::CompressibleNavierStokesDiffusion1D)

# Returns
- `dynamic_viscosity(u, equations_parabolic) / u[1] * equations_parabolic.max_1_kappa`
where `max_1_kappa = max(one(kappa), kappa)` is computed in the constructor.

For the diffusive estimate we use the eigenvalues of the diffusivity matrix,
as suggested in Section 3.5 of
- Krais et. al (2021)
  FLEXI: A high order discontinuous Galerkin framework for hyperbolic–parabolic conservation laws
  [DOI: 10.1016/j.camwa.2020.05.004](https://doi.org/10.1016/j.camwa.2020.05.004)

For details on the derivation of eigenvalues of the diffusivity matrix
for the compressible Navier-Stokes equations see for instance
- Richard P. Dwight (2006)
  Efficiency improvements of RANS-based analysis and optimization using implicit and adjoint methods on unstructured grids
  PhD Thesis, University of Manchester
  https://elib.dlr.de/50794/1/rdwight-PhDThesis-ImplicitAndAdjoint.pdf
  See especially equations (2.79), (3.24), and (3.25) from Chapter 3.2.3

The eigenvalues of the diffusivity matrix in 1D are
``-\frac{\mu}{\rho} \{0, 1, \kappa\}``
and thus the largest absolute eigenvalue is
``\frac{\mu}{\rho} \max(1, \kappa)``.
"""
@inline function max_diffusivity(u,
                                 equations_parabolic::CompressibleNavierStokesDiffusion1D)
    # See for instance also the computation in FLUXO:
    # https://github.com/project-fluxo/fluxo/blob/c7e0cc9b7fd4569dcab67bbb6e5a25c0a84859f1/src/equation/navierstokes/calctimestep.f90#L122-L128
    #
    # Accordingly, the spectral radius/largest absolute eigenvalue can be computed as:
    return dynamic_viscosity(u, equations_parabolic) / u[1] *
           equations_parabolic.max_1_kappa
end

"""
    cons2prim_temperature(u, equations::CompressibleNavierStokesDiffusion1D)

Convert conservative variables `u` to primitive variables `(rho, v1, T)`.
In contrast to [`cons2prim`](@ref), this function returns temperature as the last variable instead of pressure.

!!! warning "Experimental code"
    This function is experimental and may change in any future release.
"""
@inline function cons2prim_temperature(u,
                                       equations::CompressibleNavierStokesDiffusion1D)
    rho, rho_v1, _ = u

    v1 = rho_v1 / rho
    T = temperature(u, equations)

    return SVector(rho, v1, T)
end

# Convert conservative variables to entropy
# TODO: parabolic. We can improve efficiency by not computing w_1, which involves logarithms
# This can be done by specializing `cons2entropy` and `entropy2cons` to `CompressibleNavierStokesDiffusion1D`,
# but this may be confusing to new users.
function cons2entropy(u, equations::CompressibleNavierStokesDiffusion1D)
    return cons2entropy(u, equations.equations_hyperbolic)
end
function entropy2cons(w, equations::CompressibleNavierStokesDiffusion1D)
    return entropy2cons(w, equations.equations_hyperbolic)
end

"""
    entropy2velocity_temperature(w, equations::AbstractCompressibleNavierStokesDiffusion{1, 3})

This directly converts entropy variables `w` to velocity and temperature, which are computed 
from the entropy variables via 
``T = -1/w_3`` and ``v_1 = -w_2/w_3``, where ``w_3 = -\\rho/p`` following

- Hughes, Franca, Mallet (1986) 
  A new finite element formulation for CFD
  [DOI: 10.1016/0045-7825(86)90127-1](https://doi.org/10.1016/0045-7825(86)90127-1)
"""
@inline function entropy2velocity_temperature(w,
                                              ::AbstractCompressibleNavierStokesDiffusion{1,
                                                                                          3})
    inv_w3 = inv(w[3])
    T = -inv_w3
    v1 = -w[2] * inv_w3
    return SVector(v1, T)
end

"""
    convert_transformed_to_velocity_temperature(u_transformed, equations::CompressibleNavierStokesDiffusion1D)

Convert transformed gradient variables from [`gradient_variable_transformation`](@ref) to `(v_1, T)`. For
[`CompressibleNavierStokesDiffusion1D`](@ref), gradients are always converted to gradients of primitive variables, 
so the parabolic fluxes only require velocity and temperature to evaluate.
"""
@inline function convert_transformed_to_velocity_temperature(u_transformed,
                                                             equations::CompressibleNavierStokesDiffusion1D{GradientVariablesPrimitive})
    _, v1, T = u_transformed
    return SVector(v1, T)
end

@inline function convert_transformed_to_velocity_temperature(u_transformed,
                                                             equations::CompressibleNavierStokesDiffusion1D{GradientVariablesEntropy})
    return entropy2velocity_temperature(u_transformed, equations)
end

# Takes the solution values `u` and gradient of the entropy variables (w_2, w_3) and
# reverse engineers the gradients to be terms of the primitive variables (v1, T).
# Helpful because then the diffusive fluxes have the same form as on paper.
# Note, the first component of `gradient_entropy_vars` contains gradient(rho) which is unused.
# TODO: parabolic; entropy stable parabolic terms
@inline function convert_derivative_to_primitive(u, gradient,
                                                 ::CompressibleNavierStokesDiffusion1D{GradientVariablesPrimitive})
    return gradient
end

# the first argument is always the "transformed" variables.
@inline function convert_derivative_to_primitive(w, gradient_entropy_vars,
                                                 equations::CompressibleNavierStokesDiffusion1D{GradientVariablesEntropy})
    v1, T = entropy2velocity_temperature(w, equations)

    return SVector(gradient_entropy_vars[1],
                   T * (gradient_entropy_vars[2] + v1 * gradient_entropy_vars[3]), # grad(u) = T*(grad(w_2)+v1*grad(w_3))
                   T * T * gradient_entropy_vars[3])
end

"""
    cons2prim(u, equations::CompressibleNavierStokesDiffusion1D)

Forwards to [`cons2prim(u, equations::CompressibleEulerEquations1D)`](@ref)
to convert conservative variables to primitive variables.
"""
@inline function cons2prim(u, equations::CompressibleNavierStokesDiffusion1D)
    return cons2prim(u, equations.equations_hyperbolic)
end
"""
    prim2cons(u, equations::CompressibleNavierStokesDiffusion1D)

Forwards to [`prim2cons(u, equations::CompressibleEulerEquations1D)`](@ref)
to convert primitive variables to conservative variables.
"""
@inline function prim2cons(u, equations::CompressibleNavierStokesDiffusion1D)
    return prim2cons(u, equations.equations_hyperbolic)
end

"""
    temperature(u, equations::CompressibleNavierStokesDiffusion1D)

Compute the temperature from the conservative variables `u`.
In particular, this assumes a specific gas constant ``R = 1``:
```math
T = \\frac{p}{\\rho}
```
"""
@inline function temperature(u, equations::CompressibleNavierStokesDiffusion1D)
    rho, rho_v1, rho_e_total = u
    @unpack gamma = equations

    p = (gamma - 1) * (rho_e_total - 0.5f0 * rho_v1^2 / rho)
    T = p / rho # Corresponds to a specific gas constant R = 1
    return T
end

@inline function velocity(u, equations::CompressibleNavierStokesDiffusion1D)
    rho = u[1]
    v1 = u[2] / rho
    return v1
end

@inline function (boundary_condition::BoundaryConditionNavierStokesWall{<:NoSlip,
                                                                        <:Adiabatic})(flux_inner,
                                                                                      u_inner,
                                                                                      orientation::Integer,
                                                                                      direction,
                                                                                      x,
                                                                                      t,
                                                                                      operator_type::Gradient,
                                                                                      equations::CompressibleNavierStokesDiffusion1D{GradientVariablesPrimitive})
    v1 = boundary_condition.boundary_condition_velocity.boundary_value_function(x, t,
                                                                                equations)
    return SVector(u_inner[1], v1, u_inner[3])
end

@inline function (boundary_condition::BoundaryConditionNavierStokesWall{<:NoSlip,
                                                                        <:Adiabatic})(flux_inner,
                                                                                      u_inner,
                                                                                      orientation::Integer,
                                                                                      direction,
                                                                                      x,
                                                                                      t,
                                                                                      operator_type::Divergence,
                                                                                      equations::CompressibleNavierStokesDiffusion1D{GradientVariablesPrimitive})
    normal_heat_flux = boundary_condition.boundary_condition_heat_flux.boundary_value_normal_flux_function(x,
                                                                                                           t,
                                                                                                           equations)
    v1 = boundary_condition.boundary_condition_velocity.boundary_value_function(x, t,
                                                                                equations)
    _, tau_1n, _ = flux_inner # extract fluxes for 2nd equation
    normal_energy_flux = v1 * tau_1n + normal_heat_flux
    return SVector(flux_inner[1], flux_inner[2], normal_energy_flux)
end

@inline function (boundary_condition::BoundaryConditionNavierStokesWall{<:NoSlip,
                                                                        <:Isothermal})(flux_inner,
                                                                                       u_inner,
                                                                                       orientation::Integer,
                                                                                       direction,
                                                                                       x,
                                                                                       t,
                                                                                       operator_type::Gradient,
                                                                                       equations::CompressibleNavierStokesDiffusion1D{GradientVariablesPrimitive})
    v1 = boundary_condition.boundary_condition_velocity.boundary_value_function(x, t,
                                                                                equations)
    T = boundary_condition.boundary_condition_heat_flux.boundary_value_function(x, t,
                                                                                equations)
    return SVector(u_inner[1], v1, T)
end

@inline function (boundary_condition::BoundaryConditionNavierStokesWall{<:NoSlip,
                                                                        <:Isothermal})(flux_inner,
                                                                                       u_inner,
                                                                                       orientation::Integer,
                                                                                       direction,
                                                                                       x,
                                                                                       t,
                                                                                       operator_type::Divergence,
                                                                                       equations::CompressibleNavierStokesDiffusion1D{GradientVariablesPrimitive})
    return flux_inner
end

# specialized BC impositions for GradientVariablesEntropy.

# This should return a SVector containing the boundary values of entropy variables.
# Here, `w_inner` are the transformed variables (e.g., entropy variables).
#
# Taken from "Entropy stable modal discontinuous Galerkin schemes and wall boundary conditions
#             for the compressible Navier-Stokes equations" by Chan, Lin, Warburton 2022.
# DOI: 10.1016/j.jcp.2021.110723
@inline function (boundary_condition::BoundaryConditionNavierStokesWall{<:NoSlip,
                                                                        <:Adiabatic})(flux_inner,
                                                                                      w_inner,
                                                                                      orientation::Integer,
                                                                                      direction,
                                                                                      x,
                                                                                      t,
                                                                                      operator_type::Gradient,
                                                                                      equations::CompressibleNavierStokesDiffusion1D{GradientVariablesEntropy})
    v1 = boundary_condition.boundary_condition_velocity.boundary_value_function(x, t,
                                                                                equations)
    negative_rho_inv_p = w_inner[3] # w_3 = -rho / p
    return SVector(w_inner[1], -v1 * negative_rho_inv_p, negative_rho_inv_p)
end

# this is actually identical to the specialization for GradientVariablesPrimitive, but included for completeness.
@inline function (boundary_condition::BoundaryConditionNavierStokesWall{<:NoSlip,
                                                                        <:Adiabatic})(flux_inner,
                                                                                      w_inner,
                                                                                      orientation::Integer,
                                                                                      direction,
                                                                                      x,
                                                                                      t,
                                                                                      operator_type::Divergence,
                                                                                      equations::CompressibleNavierStokesDiffusion1D{GradientVariablesEntropy})
    normal_heat_flux = boundary_condition.boundary_condition_heat_flux.boundary_value_normal_flux_function(x,
                                                                                                           t,
                                                                                                           equations)
    v1 = boundary_condition.boundary_condition_velocity.boundary_value_function(x, t,
                                                                                equations)
    _, tau_1n, _ = flux_inner # extract fluxes for 2nd equation
    normal_energy_flux = v1 * tau_1n + normal_heat_flux
    return SVector(flux_inner[1], flux_inner[2], normal_energy_flux)
end

@inline function (boundary_condition::BoundaryConditionNavierStokesWall{<:NoSlip,
                                                                        <:Isothermal})(flux_inner,
                                                                                       w_inner,
                                                                                       orientation::Integer,
                                                                                       direction,
                                                                                       x,
                                                                                       t,
                                                                                       operator_type::Gradient,
                                                                                       equations::CompressibleNavierStokesDiffusion1D{GradientVariablesEntropy})
    v1 = boundary_condition.boundary_condition_velocity.boundary_value_function(x, t,
                                                                                equations)
    T = boundary_condition.boundary_condition_heat_flux.boundary_value_function(x, t,
                                                                                equations)

    # the entropy variables w2 = rho * v1 / p = v1 / T = -v1 * w3.
    w3 = -1 / T
    return SVector(w_inner[1], -v1 * w3, w3)
end

@inline function (boundary_condition::BoundaryConditionNavierStokesWall{<:NoSlip,
                                                                        <:Isothermal})(flux_inner,
                                                                                       w_inner,
                                                                                       orientation::Integer,
                                                                                       direction,
                                                                                       x,
                                                                                       t,
                                                                                       operator_type::Divergence,
                                                                                       equations::CompressibleNavierStokesDiffusion1D{GradientVariablesEntropy})
    return SVector(flux_inner[1], flux_inner[2], flux_inner[3])
end
end # @muladd
