using CUDAnative
using CuArrays
CuArrays.allowscalar(false)
try
  isdefined(NumericalFluxes, :rusanov!)
catch
  include("../../src/DGmethods/NumericalFluxes.jl")
end

include("DGBalanceLawProfileHarness.jl")

const _nstate = 5
const _ρ, _U, _V, _W, _E = 1:_nstate
const stateid = (ρid = _ρ, Uid = _U, Vid = _V, Wid = _W, Eid = _E)
const statenames = ("ρ", "U", "V", "W", "E")

const _nviscstates = 6
const _τ11, _τ22, _τ33, _τ12, _τ13, _τ23 = 1:_nviscstates

const _ngradstates = 3
const _states_for_gradient_transform = (_ρ, _U, _V, _W)

const _nauxstate = 3
const _a_x, _a_y, _a_z = 1:_nauxstate

const γ_exact = 7//5
const μ_exact = 1//100

@inline function preflux(Q, _...)
  γ::eltype(Q) = γ_exact
  @inbounds ρ, U, V, W, E = Q[_ρ], Q[_U], Q[_V], Q[_W], Q[_E]
  ρinv = 1 / ρ
  u, v, w = ρinv * U, ρinv * V, ρinv * W
  ((γ-1)*(E - ρinv * (U^2 + V^2 + W^2) / 2), u, v, w, ρinv)
end

# max eigenvalue
@inline function wavespeed(n, Q, aux, t, P, u, v, w, ρinv)
  γ::eltype(Q) = γ_exact
  @inbounds abs(n[1] * u + n[2] * v + n[3] * w) + sqrt(ρinv * γ * P)
end

# flux function
cns_flux!(F, Q, VF, aux, t) = cns_flux!(F, Q, VF, aux, t, preflux(Q)...)

@inline function cns_flux!(F, Q, VF, aux, t, P, u, v, w, ρinv)
  @inbounds begin
    ρ, U, V, W, E = Q[_ρ], Q[_U], Q[_V], Q[_W], Q[_E]

    τ11, τ22, τ33 = VF[_τ11], VF[_τ22], VF[_τ33]
    τ12 = τ21 = VF[_τ12]
    τ13 = τ31 = VF[_τ13]
    τ23 = τ32 = VF[_τ23]

    # inviscid terms
    F[1, _ρ], F[2, _ρ], F[3, _ρ] = U          , V          , W
    F[1, _U], F[2, _U], F[3, _U] = u * U  + P , v * U      , w * U
    F[1, _V], F[2, _V], F[3, _V] = u * V      , v * V + P  , w * V
    F[1, _W], F[2, _W], F[3, _W] = u * W      , v * W      , w * W + P
    F[1, _E], F[2, _E], F[3, _E] = u * (E + P), v * (E + P), w * (E + P)

    # viscous terms
    F[1, _U] -= τ11; F[2, _U] -= τ12; F[3, _U] -= τ13
    F[1, _V] -= τ21; F[2, _V] -= τ22; F[3, _V] -= τ23
    F[1, _W] -= τ31; F[2, _W] -= τ32; F[3, _W] -= τ33

    F[1, _E] -= u * τ11 + v * τ12 + w * τ13
    F[2, _E] -= u * τ21 + v * τ22 + w * τ23
    F[3, _E] -= u * τ31 + v * τ32 + w * τ33
  end
end

# Compute the velocity from the state
@inline function velocities!(vel, Q, _...)
  @inbounds begin
    # ordering should match states_for_gradient_transform
    ρ, U, V, W = Q[1], Q[2], Q[3], Q[4]
    ρinv = 1 / ρ
    vel[1], vel[2], vel[3] = ρinv * U, ρinv * V, ρinv * W
  end
end

# Visous flux
@inline function compute_stresses!(VF, grad_vel, _...)
  μ::eltype(VF) = μ_exact
  @inbounds begin
    dudx, dudy, dudz = grad_vel[1, 1], grad_vel[2, 1], grad_vel[3, 1]
    dvdx, dvdy, dvdz = grad_vel[1, 2], grad_vel[2, 2], grad_vel[3, 2]
    dwdx, dwdy, dwdz = grad_vel[1, 3], grad_vel[2, 3], grad_vel[3, 3]

    # strains
    ϵ11 = dudx
    ϵ22 = dvdy
    ϵ33 = dwdz
    ϵ12 = (dudy + dvdx) / 2
    ϵ13 = (dudz + dwdx) / 2
    ϵ23 = (dvdz + dwdy) / 2

    # deviatoric stresses
    VF[_τ11] = 2μ * (ϵ11 - (ϵ11 + ϵ22 + ϵ33) / 3)
    VF[_τ22] = 2μ * (ϵ22 - (ϵ11 + ϵ22 + ϵ33) / 3)
    VF[_τ33] = 2μ * (ϵ33 - (ϵ11 + ϵ22 + ϵ33) / 3)
    VF[_τ12] = 2μ * ϵ12
    VF[_τ13] = 2μ * ϵ13
    VF[_τ23] = 2μ * ϵ23
  end
end

@inline function stresses_penalty!(VF, nM, velM, QM, aM, velP, QP, aP, t)
  @inbounds begin
    n_Δvel = similar(VF, Size(3, 3))
    for j = 1:3, i = 1:3
      n_Δvel[i, j] = nM[i] * (velP[j] - velM[j]) / 2
    end
    compute_stresses!(VF, n_Δvel)
  end
end

numerical_flux!(x...) = NumericalFluxes.rusanov!(x..., cns_flux!, wavespeed,
                                                 preflux)
let
  DFloat = Float64
  dim = 2
  nelem = 100
  N = 4
  cpu_dg = DGProfiler(Array, DFloat, dim, nelem, N, _nstate, cns_flux!,
                      numerical_flux!;
                      stateoffset = ((_E, 20), (_ρ, 1)),
                      ngradstate = _ngradstates,
                      states_grad = _states_for_gradient_transform,
                      nviscstate = _nviscstates,
                      gradient_transform! = velocities!,
                      viscous_transform! = compute_stresses!,
                      viscous_penalty! = stresses_penalty!,
                      nauxstate = _nauxstate)

  gpu_dg = DGProfiler(CuArray, DFloat, dim, nelem, N, _nstate, cns_flux!,
                      numerical_flux!;
                      stateoffset = ((_E, 20), (_ρ, 1)),
                      ngradstate = _ngradstates,
                      states_grad = _states_for_gradient_transform,
                      nviscstate = _nviscstates,
                      gradient_transform! = velocities!,
                      viscous_transform! = compute_stresses!,
                      viscous_penalty! = stresses_penalty!,
                      nauxstate = _nauxstate)

  volumerhs!(cpu_dg)
  facerhs!(cpu_dg)
  volumeviscterms!(cpu_dg)
  faceviscterms!(cpu_dg)

  volumerhs!(gpu_dg)
  facerhs!(gpu_dg)
  volumeviscterms!(gpu_dg)
  faceviscterms!(gpu_dg)

  gpu_rhs = Array(gpu_dg.rhs)
  cpu_rhs = Array(cpu_dg.rhs)
  @show gpu_rhs ≈ cpu_rhs
  nothing
end
