# 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"""
    HyperbolicDiffusionEquations3D

The linear hyperbolic diffusion equations in three space dimensions.
A description of this system can be found in Sec. 2.5 of the book "I Do Like CFD, Too: Vol 1".
The book is freely available at [http://www.cfdbooks.com/](http://www.cfdbooks.com/) and further analysis can be found in
the paper by Nishikawa [DOI: 10.1016/j.jcp.2007.07.029](https://doi.org/10.1016/j.jcp.2007.07.029)
"""
struct HyperbolicDiffusionEquations3D{RealT <: Real} <:
       AbstractHyperbolicDiffusionEquations{3, 4}
    Lr::RealT     # reference length scale
    inv_Tr::RealT # inverse of the reference time scale
    nu::RealT     # diffusion constant
end

function HyperbolicDiffusionEquations3D(; nu = 1.0, Lr = inv(2pi))
    Tr = Lr^2 / nu
    return HyperbolicDiffusionEquations3D(promote(Lr, inv(Tr), nu)...)
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::HyperbolicDiffusionEquations3D,
                      ::Type{NewRealT}) where {NewRealT}
    return HyperbolicDiffusionEquations3D(convert(NewRealT, equations.Lr),
                                          convert(NewRealT, equations.inv_Tr),
                                          convert(NewRealT, equations.nu))
end

function varnames(::typeof(cons2cons), ::HyperbolicDiffusionEquations3D)
    return ("phi", "q1", "q2", "q3")
end
function varnames(::typeof(cons2prim), ::HyperbolicDiffusionEquations3D)
    return ("phi", "q1", "q2", "q3")
end
function default_analysis_errors(::HyperbolicDiffusionEquations3D)
    return (:l2_error, :linf_error, :residual)
end

"""
    residual_steady_state(du, ::AbstractHyperbolicDiffusionEquations)

Used to determine the termination criterion of a [`SteadyStateCallback`](@ref).
For hyperbolic diffusion, this checks convergence of the potential ``\\phi``.
"""
@inline function residual_steady_state(du, ::HyperbolicDiffusionEquations3D)
    return abs(du[1])
end

# Set initial conditions at physical location `x` for pseudo-time `t`
function initial_condition_poisson_nonperiodic(x, t,
                                               equations::HyperbolicDiffusionEquations3D)
    # elliptic equation: -νΔϕ = f
    RealT = eltype(x)
    if t == 0
        phi = one(RealT)
        q1 = one(RealT)
        q2 = one(RealT)
        q3 = one(RealT)
    else
        phi = 2 * cospi(x[1]) *
              sinpi(2 * x[2]) *
              sinpi(2 * x[3]) + 2 # ϕ
        q1 = -2 * convert(RealT, pi) * sinpi(x[1]) *
             sinpi(2 * x[2]) * sinpi(2 * x[3])   # ϕ_x
        q2 = 4 * convert(RealT, pi) * cospi(x[1]) *
             cospi(2 * x[2]) * sinpi(2 * x[3])   # ϕ_y
        q3 = 4 * convert(RealT, pi) * cospi(x[1]) *
             sinpi(2 * x[2]) * cospi(2 * x[3])   # ϕ_z
    end
    return SVector(phi, q1, q2, q3)
end

@inline function source_terms_poisson_nonperiodic(u, x, t,
                                                  equations::HyperbolicDiffusionEquations3D)
    # elliptic equation: -νΔϕ = f
    # analytical solution: ϕ = 2 cos(πx)sin(2πy)sin(2πz) + 2 and f = 18 π^2 cos(πx)sin(2πy)sin(2πz)
    RealT = eltype(u)
    @unpack inv_Tr = equations

    x1, x2, x3 = x
    du1 = 18 * convert(RealT, pi)^2 * cospi(x1) * sinpi(2 * x2) * sinpi(2 * x3)
    du2 = -inv_Tr * u[2]
    du3 = -inv_Tr * u[3]
    du4 = -inv_Tr * u[4]

    return SVector(du1, du2, du3, du4)
end

function boundary_condition_poisson_nonperiodic(u_inner, orientation, direction, x, t,
                                                surface_flux_function,
                                                equations::HyperbolicDiffusionEquations3D)
    # elliptic equation: -νΔϕ = f
    RealT = eltype(u_inner)
    phi = 2 * cospi(x[1]) * sinpi(2 * x[2]) *
          sinpi(2 * x[3]) + 2 # ϕ
    q1 = -2 * convert(RealT, pi) * sinpi(x[1]) *
         sinpi(2 * x[2]) * sinpi(2 * x[3])   # ϕ_x
    q2 = 4 * convert(RealT, pi) * cospi(x[1]) *
         cospi(2 * x[2]) * sinpi(2 * x[3])   # ϕ_y
    q3 = 4 * convert(RealT, pi) * cospi(x[1]) *
         sinpi(2 * x[2]) * cospi(2 * x[3])   # ϕ_z
    u_boundary = SVector(phi, q1, q2, q3)

    # Calculate boundary flux
    if iseven(direction) # u_inner is "left" of boundary, u_boundary is "right" of boundary
        flux = surface_flux_function(u_inner, u_boundary, orientation, equations)
    else # u_boundary is "left" of boundary, u_inner is "right" of boundary
        flux = surface_flux_function(u_boundary, u_inner, orientation, equations)
    end

    return flux
end

"""
    source_terms_harmonic(u, x, t, equations::HyperbolicDiffusionEquations3D)

Source term that only includes the forcing from the hyperbolic diffusion system.
"""
@inline function source_terms_harmonic(u, x, t,
                                       equations::HyperbolicDiffusionEquations3D)
    # harmonic solution ϕ = (sinh(πx)sin(πy) + sinh(πy)sin(πx))/sinh(π), so f = 0
    @unpack inv_Tr = equations

    du1 = 0
    du2 = -inv_Tr * u[2]
    du3 = -inv_Tr * u[3]
    du4 = -inv_Tr * u[4]

    return SVector(du1, du2, du3, du4)
end

"""
    initial_condition_eoc_test_coupled_euler_gravity(x, t, equations::HyperbolicDiffusionEquations3D)

Setup used for convergence tests of the Euler equations with self-gravity used in
- Michael Schlottke-Lakemper, Andrew R. Winters, Hendrik Ranocha, Gregor J. Gassner (2020)
  A purely hyperbolic discontinuous Galerkin approach for self-gravitating gas dynamics
  [arXiv: 2008.10593](https://arxiv.org/abs/2008.10593)
in combination with [`source_terms_harmonic`](@ref).
"""
function initial_condition_eoc_test_coupled_euler_gravity(x, t,
                                                          equations::HyperbolicDiffusionEquations3D)

    # Determine phi_x, phi_y
    RealT = eltype(x)
    G = 1 # gravitational constant
    C_grav = -4 * G / (3 * convert(RealT, pi)) # "3" is the number of spatial dimensions  # 2D: -2.0*G/pi
    A = convert(RealT, 0.1) # perturbation coefficient must match Euler setup
    rho1 = A * sinpi(x[1] + x[2] + x[3] - t)
    # initialize with ansatz of gravity potential
    phi = C_grav * rho1
    q1 = C_grav * A * convert(RealT, pi) *
         cospi(x[1] + x[2] + x[3] - t) # = gravity acceleration in x-direction
    q2 = q1                                                 # = gravity acceleration in y-direction
    q3 = q1                                                 # = gravity acceleration in z-direction

    return SVector(phi, q1, q2, q3)
end

# Calculate 1D flux for a single point
@inline function flux(u, orientation::Integer,
                      equations::HyperbolicDiffusionEquations3D)
    phi, q1, q2, q3 = u

    RealT = eltype(u)
    if orientation == 1
        f1 = -equations.nu * q1
        f2 = -phi * equations.inv_Tr
        f3 = zero(RealT)
        f4 = zero(RealT)
    elseif orientation == 2
        f1 = -equations.nu * q2
        f2 = zero(RealT)
        f3 = -phi * equations.inv_Tr
        f4 = zero(RealT)
    else
        f1 = -equations.nu * q3
        f2 = zero(RealT)
        f3 = zero(RealT)
        f4 = -phi * equations.inv_Tr
    end

    return SVector(f1, f2, f3, f4)
end

# Calculate maximum wave speed for local Lax-Friedrichs-type dissipation
@inline function max_abs_speed_naive(u_ll, u_rr, orientation::Integer,
                                     equations::HyperbolicDiffusionEquations3D)
    return sqrt(equations.nu * equations.inv_Tr)
end

"""
    flux_godunov(u_ll, u_rr, orientation_or_normal_direction, 
                 equations::HyperbolicDiffusionEquations3D)

Godunov (upwind) flux for the 3D hyperbolic diffusion equations.
"""
@inline function flux_godunov(u_ll, u_rr, orientation::Integer,
                              equations::HyperbolicDiffusionEquations3D)
    # Obtain left and right fluxes
    phi_ll, q1_ll, q2_ll, q3_ll = u_ll
    phi_rr, q1_rr, q2_rr, q3_rr = u_rr
    f_ll = flux(u_ll, orientation, equations)
    f_rr = flux(u_rr, orientation, equations)

    # this is an optimized version of the application of the upwind dissipation matrix:
    #   dissipation = 0.5*R_n*|Λ|*inv(R_n)[[u]]
    λ_max = sqrt(equations.nu * equations.inv_Tr)
    f1 = 0.5f0 * (f_ll[1] + f_rr[1]) - 0.5f0 * λ_max * (phi_rr - phi_ll)
    if orientation == 1 # x-direction
        f2 = 0.5f0 * (f_ll[2] + f_rr[2]) - 0.5f0 * λ_max * (q1_rr - q1_ll)
        f3 = 0.5f0 * (f_ll[3] + f_rr[3])
        f4 = 0.5f0 * (f_ll[4] + f_rr[4])
    elseif orientation == 2 # y-direction
        f2 = 0.5f0 * (f_ll[2] + f_rr[2])
        f3 = 0.5f0 * (f_ll[3] + f_rr[3]) - 0.5f0 * λ_max * (q2_rr - q2_ll)
        f4 = 0.5f0 * (f_ll[4] + f_rr[4])
    else # y-direction
        f2 = 0.5f0 * (f_ll[2] + f_rr[2])
        f3 = 0.5f0 * (f_ll[3] + f_rr[3])
        f4 = 0.5f0 * (f_ll[4] + f_rr[4]) - 0.5f0 * λ_max * (q3_rr - q3_ll)
    end

    return SVector(f1, f2, f3, f4)
end

"""
    have_constant_speed(::HyperbolicDiffusionEquations3D)

Indicates whether the characteristic speeds are constant, i.e., independent of the solution.
Queried in the timestep computation [`StepsizeCallback`](@ref) and [`linear_structure`](@ref).

# Returns
- `True()`
"""
@inline have_constant_speed(::HyperbolicDiffusionEquations3D) = True()

@inline function max_abs_speeds(eq::HyperbolicDiffusionEquations3D)
    λ = sqrt(eq.nu * eq.inv_Tr)
    return λ, λ, λ
end

# Convert conservative variables to primitive
@inline cons2prim(u, equations::HyperbolicDiffusionEquations3D) = u

# Convert conservative variables to entropy found in I Do Like CFD, Too, Vol. 1
@inline function cons2entropy(u, equations::HyperbolicDiffusionEquations3D)
    phi, q1, q2, q3 = u
    w1 = phi
    w2 = equations.Lr^2 * q1
    w3 = equations.Lr^2 * q2
    w4 = equations.Lr^2 * q3

    return SVector(w1, w2, w3, w4)
end

# Calculate entropy for a conservative state `u` (here: same as total energy)
@inline function entropy(u, equations::HyperbolicDiffusionEquations3D)
    return energy_total(u, equations)
end

# Calculate total energy for a conservative state `u`
@inline function energy_total(u, equations::HyperbolicDiffusionEquations3D)
    # energy function as found in equation (2.5.12) in the book "I Do Like CFD, Vol. 1"
    phi, q1, q2, q3 = u
    return 0.5f0 * (phi^2 + equations.Lr^2 * (q1^2 + q2^2 + q3^2))
end
end # @muladd
