text:
# ============================================================
# Gas Kinetic Flux Functions
# ============================================================

"""
    flux_gks(u, μ, dt, a = 0::T, su = 0.0) where {T<:Real}
    flux_gks(uL::T, uR::T, μ, dt, dxL, dxR, a = 0, suL = 0.0, suR = 0.0) where {T<:Real}

Gas kinetic scalar flux

* @args: conservative scalars and their slopes
* @args: viscosity
* @args: time step and cell size
* @return: scalar flux
"""
function flux_gks(u, μ, dt, a = 0, su = 0.0)
    prim = ifelse(a == 0, conserve_prim(u), conserve_prim(u, a))

    Mu = gauss_moments(prim)[1]
    tau = 2.0 * μ

    fa = pdf_slope(u, su)
    Δ = -prim[1] * moments_conserve_slope(fa, Mu, 1)
    faT = pdf_slope(u, Δ)

    # time-integration constants
    Mt = zeros(5)
    Mt[4] = dt
    Mt[5] = -tau * dt * exp(-dt / tau) + tau * Mt[4]

    # flux related to upwind distribution
    Muv = moments_conserve(Mu, 1)
    Mau = moments_conserve_slope(fa, Mu, 2)
    MauT = moments_conserve_slope(faT, Mu, 1)

    fw =
        Mt[4] * prim[1] * Muv - (Mt[5] + tau * Mt[4]) * prim[1] * Mau -
        tau * Mt[4] * prim[1] * MauT

    return fw / dt
end

function flux_gks(
    uL::T,
    uR::T,
    μ,
    dt,
    dxL,
    dxR,
    a = 0,
    suL = 0.0,
    suR = 0.0,
) where {T<:Real}

    primL = ifelse(a == 0, conserve_prim(uL), conserve_prim(uL, a))
    primR = ifelse(a == 0, conserve_prim(uR), conserve_prim(uR, a))

    Mu1, MuL1, MuR1 = gauss_moments(primL)
    Mu2, MuL2, MuR2 = gauss_moments(primR)

    u = primL[1] * moments_conserve(MuL1, 0) + primR[1] * moments_conserve(MuR2, 0)
    prim = ifelse(a == 0, conserve_prim(u), conserve_prim(u, a))
    tau = 2.0 * μ + 1.0 * abs(uL - uR) / (abs(uL) + abs(uR)) * dt

    faL = pdf_slope(uL, suL)
    Δ = -primL[1] * moments_conserve_slope(faL, Mu1, 1)
    faTL = pdf_slope(uL, Δ)

    faR = pdf_slope(uR, suR)
    Δ = -primR[1] * moments_conserve_slope(faR, Mu2, 1)
    faTR = pdf_slope(uR, Δ)

    Mu, MuL, MuR = gauss_moments(prim)
    sw0L = (u - (uL - suL * dxL)) / dxL
    sw0R = ((uR + suR * dxR) - u) / dxR
    gaL = pdf_slope(u, sw0L)
    gaR = pdf_slope(u, sw0R)
    Δ =
        -prim[1] *
        (moments_conserve_slope(gaL, MuL, 1) + moments_conserve_slope(gaR, MuR, 1))
    #sw = (uR - uL) / (dxL + dxR)
    #ga = pdf_slope(u, sw)
    #Δ = -prim[1] .* moments_conserve_slope(ga, Mu, 1)
    gaT = pdf_slope(u, Δ)

    # time-integration constants
    Mt = zeros(5)
    Mt[4] = tau * (1.0 - exp(-dt / tau))
    Mt[5] = -tau * dt * exp(-dt / tau) + tau * Mt[4]
    Mt[1] = dt - Mt[4]
    Mt[2] = -tau * Mt[1] + Mt[5]
    Mt[3] = 0.5 * dt^2 - tau * Mt[1]

    # flux related to central distribution
    Muv = moments_conserve(Mu, 1)
    MauL = moments_conserve_slope(gaL, MuL, 2)
    MauR = moments_conserve_slope(gaR, MuR, 2)
    #Mau = moments_conserve_slope(ga, MuR, 2)
    MauT = moments_conserve_slope(gaT, Mu, 1)

    fw = Mt[1] * prim[1] * Muv + Mt[2] * prim[1] * (MauL + MauR) + Mt[3] * prim[1] * MauT
    #fw = Mt[1] * prim[1] * Muv + Mt[2] * prim[1] * Mau + Mt[3] * prim[1] * MauT

    # flux related to upwind distribution
    MuvL = moments_conserve(MuL1, 1)
    MauL = moments_conserve_slope(faL, MuL1, 2)
    MauLT = moments_conserve_slope(faTL, MuL1, 1)

    MuvR = moments_conserve(MuR2, 1)
    MauR = moments_conserve_slope(faR, MuR2, 2)
    MauRT = moments_conserve_slope(faTR, MuR2, 1)

    fw +=
        Mt[4] * primL[1] * MuvL - (Mt[5] + tau * Mt[4]) * primL[1] * MauL -
        tau * Mt[4] * primL[1] * MauLT + Mt[4] * primR[1] * MuvR -
        (Mt[5] + tau * Mt[4]) * primR[1] * MauR - tau * Mt[4] * primR[1] * MauRT
    # fw += Mt[4] * primL[1] * MuvL + Mt[4] * primR[1] * MuvR

    return fw

end

"""
Gas kinetic Navier-Stokes flux

* @args: conservative variables and their left/right slopes
* @args: molecular and thermodynamic parameters
* @args: time step and cell size

"""
function flux_gks!(
    fw::X,
    w::Y,
    inK::Real,
    γ::Real,
    μᵣ::Real,
    ω::Real,
    sw = zero(w)::Y,
) where {X<:AbstractArray{<:AbstractFloat,1},Y<:AbstractArray{<:AbstractFloat,1}}

    prim = conserve_prim(w, γ)
    mus = gauss_moments(prim, inK)
    tau = vhs_collision_time(prim, μᵣ, ω)

    ax = pdf_slope(prim, sw, inK)
    if length(prim) == 3
        ∂ft = -prim[1] .* moments_conserve_slope(ax, mus[1], mus[2], 1)
        at = pdf_slope(prim, ∂ft, inK)

        Muv = moments_conserve(mus[1], mus[2], 1, 0)
        Mau = moments_conserve_slope(ax, mus[1], mus[2], 2)
        Mtu = moments_conserve_slope(at, mus[1], mus[2], 1)
    elseif length(prim) == 4
        ∂ft = -prim[1] .* moments_conserve_slope(ax, mus[1], mus[2], mus[3], 1, 0)
        at = pdf_slope(prim, ∂ft, inK)

        Muv = moments_conserve(mus[1], mus[2], mus[3], 1, 0, 0)
        Mau = moments_conserve_slope(ax, mus[1], mus[2], mus[3], 2, 0)
        Mtu = moments_conserve_slope(at, mus[1], mus[2], mus[3], 1, 0)
    elseif length(prim) == 5
        ∂ft = -prim[1] .* moments_conserve_slope(ax, mus[1], mus[2], mus[3], 1, 0, 0)
        at = pdf_slope(prim, ∂ft, inK)

        Muv = moments_conserve(mus[1], mus[2], mus[3], 1, 0, 0)
        Mau = moments_conserve_slope(ax, mus[1], mus[2], mus[3], 2, 0, 0)
        Mtu = moments_conserve_slope(at, mus[1], mus[2], mus[3], 1, 0, 0)
    end

    @. fw = prim[1] * (Muv - tau * Mau - tau * Mtu)

    return nothing

end

# ------------------------------------------------------------
# 1D
# ------------------------------------------------------------
function flux_gks!(
    fw::X,
    wL::Y,
    wR::Y,
    inK::Real,
    γ::Real,
    μᵣ::Real,
    ω::Real,
    dt::Real,
    dxL::Real,
    dxR::Real,
    swL = zero(wL)::Y,
    swR = zero(wR)::Y,
) where {X<:AbstractArray{<:AbstractFloat,1},Y<:AbstractArray{<:AbstractFloat,1}}

    primL = conserve_prim(wL, γ)
    primR = conserve_prim(wR, γ)

    Mu1, Mxi1, MuL1, MuR1 = gauss_moments(primL, inK)
    Mu2, Mxi2, MuL2, MuR2 = gauss_moments(primR, inK)

    w =
        primL[1] .* moments_conserve(MuL1, Mxi1, 0, 0) .+
        primR[1] .* moments_conserve(MuR2, Mxi2, 0, 0)
    prim = conserve_prim(w, γ)
    tau =
        vhs_collision_time(prim, μᵣ, ω) +
        2.0 * dt * abs(primL[1] / primL[end] - primR[1] / primR[end]) /
        (primL[1] / primL[end] + primR[1] / primR[end])

    faL = pdf_slope(primL, swL, inK)
    sw = -primL[1] .* moments_conserve_slope(faL, Mu1, Mxi1, 1)
    faTL = pdf_slope(primL, sw, inK)

    faR = pdf_slope(primR, swR, inK)
    sw = -primR[1] .* moments_conserve_slope(faR, Mu2, Mxi2, 1)
    faTR = pdf_slope(primR, sw, inK)

    Mu, Mxi, MuL, MuR = gauss_moments(prim, inK)
    sw0L = (w .- (wL .- swL .* dxL)) ./ dxL
    sw0R = ((wR .+ swR .* dxR) .- w) ./ dxR
    gaL = pdf_slope(prim, sw0L, inK)
    gaR = pdf_slope(prim, sw0R, inK)
    sw =
        -prim[1] .* (
            moments_conserve_slope(gaL, MuL, Mxi, 1) .+
            moments_conserve_slope(gaR, MuR, Mxi, 1)
        )
    # ga = pdf_slope(prim, sw, inK)
    # sw = -prim[1] .* moments_conserve_slope(ga, Mu, Mxi, 1)
    gaT = pdf_slope(prim, sw, inK)

    # time-integration constants
    Mt = zeros(5)
    Mt[4] = tau * (1.0 - exp(-dt / tau))
    Mt[5] = -tau * dt * exp(-dt / tau) + tau * Mt[4]
    Mt[1] = dt - Mt[4]
    Mt[2] = -tau * Mt[1] + Mt[5]
    Mt[3] = 0.5 * dt^2 - tau * Mt[1]

    # flux related to central distribution
    Muv = moments_conserve(Mu, Mxi, 1, 0)
    MauL = moments_conserve_slope(gaL, MuL, Mxi, 2)
    MauR = moments_conserve_slope(gaR, MuR, Mxi, 2)
    # Mau = moments_conserve_slope(ga, MuR, Mxi, 2)
    MauT = moments_conserve_slope(gaT, Mu, Mxi, 1)

    fw .=
        Mt[1] .* prim[1] .* Muv .+ Mt[2] .* prim[1] .* (MauL .+ MauR) .+
        Mt[3] .* prim[1] .* MauT
    # fw .= Mt[1] .* prim[1] .* Muv .+ Mt[2] .* prim[1] .* Mau .+ Mt[3] .* prim[1] .* MauT

    # flux related to upwind distribution
    MuvL = moments_conserve(MuL1, Mxi1, 1, 0)
    MauL = moments_conserve_slope(faL, MuL1, Mxi1, 2)
    MauLT = moments_conserve_slope(faTL, MuL1, Mxi1, 1)

    MuvR = moments_conserve(MuR2, Mxi2, 1, 0)
    MauR = moments_conserve_slope(faR, MuR2, Mxi2, 2)
    MauRT = moments_conserve_slope(faTR, MuR2, Mxi2, 1)

    @. fw +=
        Mt[4] * primL[1] * MuvL - (Mt[5] + tau * Mt[4]) * primL[1] * MauL -
        tau * Mt[4] * primL[1] * MauLT + Mt[4] * primR[1] * MuvR -
        (Mt[5] + tau * Mt[4]) * primR[1] * MauR - tau * Mt[4] * primR[1] * MauRT
    # @. fw += Mt[4] * primL[1] * MuvL + Mt[4] * primR[1] * MuvR

    return nothing

end

#--- mixture ---#
function flux_gks!(
    fw::X,
    wL::Y,
    wR::Y,
    inK,
    γ,
    mi,
    ni,
    me,
    ne,
    Kn,
    dt,
    dxL::Real,
    dxR::Real,
    swL = zero(wL)::Y,
    swR = zero(wR)::Y,
) where {X<:AbstractArray{<:AbstractFloat,2},Y<:AbstractArray{<:Real,2}}

    primL = mixture_conserve_prim(wL, γ)
    primR = mixture_conserve_prim(wR, γ)

    Mu1, Mxi1, MuL1, MuR1 = mixture_gauss_moments(primL, inK)
    Mu2, Mxi2, MuL2, MuR2 = mixture_gauss_moments(primR, inK)
    Muv1 = mixture_moments_conserve(MuL1, Mxi1, 0, 0)
    Muv2 = mixture_moments_conserve(MuR2, Mxi2, 0, 0)

    faL = mixture_pdf_slope(primL, swL, inK)
    mm = mixture_moments_conserve_slope(faL, Mu1, Mxi1, 1)
    sw = similar(swL)
    for j in axes(sw, 2)
        @. sw[:, j] = -primL[1, j] * mm[:, j]
    end
    faTL = mixture_pdf_slope(primL, sw, inK)

    faR = mixture_pdf_slope(primR, swR, inK)
    mm = mixture_moments_conserve_slope(faR, Mu2, Mxi2, 1)
    for j in axes(sw, 2)
        @. sw[:, j] = -primR[1, j] * mm[:, j]
    end
    faTR = mixture_pdf_slope(primR, sw, inK)

    w = similar(wL)
    for j in axes(w, 2)
        @. w[:, j] = primL[1, j] * Muv1[:, j] + primR[1, j] * Muv2[:, j]
    end
    prim = mixture_conserve_prim(w, γ)
    tau = aap_hs_collision_time(prim, mi, ni, me, ne, Kn)
    prim = aap_hs_prim(prim, tau, mi, ni, me, ne, Kn)

    Mu, Mxi, MuL, MuR = mixture_gauss_moments(prim, inK)
    sw0L = (w .- (wL .- swL .* dxL)) ./ dxL
    sw0R = ((wR .+ swR .* dxR) .- w) ./ dxR
    gaL = mixture_pdf_slope(prim, sw0L, inK)
    gaR = mixture_pdf_slope(prim, sw0R, inK)
    mmL = mixture_moments_conserve_slope(gaL, MuL, Mxi, 1)
    mmR = mixture_moments_conserve_slope(gaR, MuR, Mxi, 1)
    for j in axes(sw, 2)
        @. sw[:, j] = -prim[1, j] * (mmL[:, j] + mmR[:, j])
    end
    gaT = mixture_pdf_slope(prim, sw, inK)

    Mt = zeros(5, 2)
    for j in axes(Mt, 2)
        Mt[4, j] = tau[j] * (1.0 - exp(-dt / tau[j]))
        Mt[5, j] = -tau[j] * dt * exp(-dt / tau[j]) + tau[j] * Mt[4]
        Mt[1, j] = dt - Mt[4, j]
        Mt[2, j] = -tau[j] * Mt[1, j] + Mt[5, j]
        Mt[3, j] = 0.5 * dt^2 - tau[j] * Mt[1, j]
    end

    # flux related to central distribution
    Muv = mixture_moments_conserve(Mu, Mxi, 1, 0)
    MauL = mixture_moments_conserve_slope(gaL, MuL, Mxi, 2)
    MauR = mixture_moments_conserve_slope(gaR, MuR, Mxi, 2)
    MauT = mixture_moments_conserve_slope(gaT, Mu, Mxi, 1)
    for j in axes(fw, 2)
        @. fw[:, j] =
            Mt[1, j] * prim[1, j] * Muv[:, j] +
            Mt[2, j] * prim[1, j] * (MauL[:, j] + MauR[:, j]) +
            Mt[3, j] * prim[1, j] * MauT[:, j]
    end

    # flux related to upwind distribution
    MuvL = mixture_moments_conserve(MuL1, Mxi1, 1, 0)
    MauL = mixture_moments_conserve_slope(faL, MuL1, Mxi1, 2)
    MauLT = mixture_moments_conserve_slope(faTL, MuL1, Mxi1, 1)

    MuvR = mixture_moments_conserve(MuR2, Mxi2, 1, 0)
    MauR = mixture_moments_conserve_slope(faR, MuR2, Mxi2, 2)
    MauRT = mixture_moments_conserve_slope(faTR, MuR2, Mxi2, 1)

    for j in axes(fw, 2)
        @. fw[:, j] +=
            Mt[4, j] * primL[1, j] * MuvL[:, j] -
            (Mt[5, j] + tau[j] * Mt[4, j]) * primL[1, j] * MauL[:, j] -
            tau[j] * Mt[4, j] * primL[1, j] * MauLT[:, j] +
            Mt[4, j] * primR[1, j] * MuvR[:, j] -
            (Mt[5, j] + tau[j] * Mt[4, j]) * primR[1, j] * MauR[:, j] -
            tau[j] * Mt[4, j] * primR[1, j] * MauRT[:, j]
    end

    return nothing

end

#--- 1D1F1V ---#
function flux_gks!(
    fw::T1,
    fh::T2,
    wL::T3,
    wR::T3,
    u::T4,
    inK,
    γ,
    visRef,
    visIdx,
    dt,
    dxL,
    dxR,
    swL::T3,
    swR::T3,
) where {
    T1<:AbstractArray{<:AbstractFloat,1},
    T2<:AbstractArray{<:AbstractFloat,1},
    T3<:AbstractArray{<:Real,1},
    T4<:AbstractArray{<:AbstractFloat,1},
}

    primL = conserve_prim(wL, γ)
    primR = conserve_prim(wR, γ)

    Mu1, Mxi1, MuL1, _ = gauss_moments(primL, inK)
    Mu2, Mxi2, _, MuR2 = gauss_moments(primR, inK)

    w =
        primL[1] .* moments_conserve(MuL1, Mxi1, 0, 0) .+
        primR[1] .* moments_conserve(MuR2, Mxi2, 0, 0)
    prim = conserve_prim(w, γ)
    tau =
        vhs_collision_time(prim, visRef, visIdx) +
        2.0 * dt * abs(primL[1] / primL[end] - primR[1] / primR[end]) /
        (primL[1] / primL[end] + primR[1] / primR[end])

    faL = pdf_slope(primL, swL, inK)
    sw = -primL[1] .* moments_conserve_slope(faL, Mu1, Mxi1, 1)
    faTL = pdf_slope(primL, sw, inK)

    faR = pdf_slope(primR, swR, inK)
    sw = -primR[1] .* moments_conserve_slope(faR, Mu2, Mxi2, 1)
    faTR = pdf_slope(primR, sw, inK)

    Mu, Mxi, MuL, MuR = gauss_moments(prim, inK)
    sw0L = (w .- (wL .- swL .* dxL)) ./ dxL
    sw0R = ((wR .+ swR .* dxR) .- w) ./ dxR
    gaL = pdf_slope(prim, sw0L, inK)
    gaR = pdf_slope(prim, sw0R, inK)
    sw =
        -prim[1] .* (
            moments_conserve_slope(gaL, MuL, Mxi, 1) .+
            moments_conserve_slope(gaR, MuR, Mxi, 1)
        )
    # ga = pdf_slope(prim, sw, inK)
    # sw = -prim[1] .* moments_conserve_slope(ga, Mu, Mxi, 1)
    gaT = pdf_slope(prim, sw, inK)

    # time-integration constants
    Mt = zeros(5)
    Mt[4] = tau * (1.0 - exp(-dt / tau))
    Mt[5] = -tau * dt * exp(-dt / tau) + tau * Mt[4]
    Mt[1] = dt - Mt[4]
    Mt[2] = -tau * Mt[1] + Mt[5]
    Mt[3] = 0.5 * dt^2 - tau * Mt[1]

    # flux related to central distribution
    Muv = moments_conserve(Mu, Mxi, 1, 0)
    MauL = moments_conserve_slope(gaL, MuL, Mxi, 2)
    MauR = moments_conserve_slope(gaR, MuR, Mxi, 2)
    # Mau = moments_conserve_slope(ga, MuR, Mxi, 2)
    MauT = moments_conserve_slope(gaT, Mu, Mxi, 1)

    fw .=
        Mt[1] .* prim[1] .* Muv .+ Mt[2] .* prim[1] .* (MauL .+ MauR) .+
        Mt[3] .* prim[1] .* MauT
    # fw .= Mt[1] .* prim[1] .* Muv .+ Mt[2] .* prim[1] .* Mau .+ Mt[3] .* prim[1] .* MauT

    # flux related to upwind distribution
    MuvL = moments_conserve(MuL1, Mxi1, 1, 0)
    MauL = moments_conserve_slope(faL, MuL1, Mxi1, 2)
    MauLT = moments_conserve_slope(faTL, MuL1, Mxi1, 1)

    MuvR = moments_conserve(MuR2, Mxi2, 1, 0)
    MauR = moments_conserve_slope(faR, MuR2, Mxi2, 2)
    MauRT = moments_conserve_slope(faTR, MuR2, Mxi2, 1)

    @. fw +=
        Mt[4] * primL[1] * MuvL - (Mt[5] + tau * Mt[4]) * primL[1] * MauL -
        tau * Mt[4] * primL[1] * MauLT + Mt[4] * primR[1] * MuvR -
        (Mt[5] + tau * Mt[4]) * primR[1] * MauR - tau * Mt[4] * primR[1] * MauRT
    # @. fw += Mt[4] * primL[1] * MuvL + Mt[4] * primR[1] * MuvR

    #--- fluxes of distribution functions ---#
    δ = heaviside.(u)
    HL = maxwellian(u, primL)
    HR = maxwellian(u, primR)
    H = maxwellian(u, prim)

    @. fh =
        Mt[1] * u * H +
        Mt[2] * u^2 * (gaL[1] * H + gaL[2] * u * H + 0.5 * gaL[3] * (u^2 * H)) * δ +
        Mt[2] * u^2 * (gaR[1] * H + gaR[2] * u * H + 0.5 * gaR[3] * (u^2 * H)) * (1.0 - δ) +
        Mt[3] * u * (gaT[1] * H + gaT[2] * u * H + 0.5 * gaT[3] * (u^2 * H)) +
        Mt[4] * u * HL * δ -
        (Mt[5] + tau * Mt[4]) *
        u^2 *
        (faL[1] * HL + faL[2] * u * HL + 0.5 * faL[3] * (u^2 * HL)) *
        δ + Mt[4] * u * HR * (1.0 - δ) -
        (Mt[5] + tau * Mt[4]) *
        u^2 *
        (faR[1] * HR + faR[2] * u * HR + 0.5 * faR[3] * (u^2 * HR)) *
        (1.0 - δ) -
        tau *
        Mt[4] *
        u *
        (faTL[1] * HL + faTL[2] * u * HL + 0.5 * faTL[3] * (u^2 * HL)) *
        δ -
        tau *
        Mt[4] *
        u *
        (faTR[1] * HR + faTR[2] * u * HR + 0.5 * faTR[3] * (u^2 * HR)) *
        (1.0 - δ)

    return nothing

end

#--- 1D2F1V ---#
function flux_gks!(
    fw::T1,
    fh::T2,
    fb::T2,
    wL::T3,
    wR::T3,
    u::T4,
    inK,
    γ,
    visRef,
    visIdx,
    dt,
    dxL,
    dxR,
    swL::T3,
    swR::T3,
) where {
    T1<:AbstractArray{<:AbstractFloat,1},
    T2<:AbstractArray{<:AbstractFloat,1},
    T3<:AbstractArray{<:Real,1},
    T4<:AbstractArray{<:AbstractFloat,1},
}

    primL = conserve_prim(wL, γ)
    primR = conserve_prim(wR, γ)

    Mu1, Mxi1, MuL1, MuR1 = gauss_moments(primL, inK)
    Mu2, Mxi2, MuL2, MuR2 = gauss_moments(primR, inK)

    w =
        primL[1] .* moments_conserve(MuL1, Mxi1, 0, 0) .+
        primR[1] .* moments_conserve(MuR2, Mxi2, 0, 0)
    prim = conserve_prim(w, γ)
    tau =
        vhs_collision_time(prim, visRef, visIdx) +
        2.0 * dt * abs(primL[1] / primL[end] - primR[1] / primR[end]) /
        (primL[1] / primL[end] + primR[1] / primR[end])

    faL = pdf_slope(primL, swL, inK)
    sw = -primL[1] .* moments_conserve_slope(faL, Mu1, Mxi1, 1)
    faTL = pdf_slope(primL, sw, inK)

    faR = pdf_slope(primR, swR, inK)
    sw = -primR[1] .* moments_conserve_slope(faR, Mu2, Mxi2, 1)
    faTR = pdf_slope(primR, sw, inK)

    Mu, Mxi, MuL, MuR = gauss_moments(prim, inK)
    sw0L = (w .- (wL .- swL .* dxL)) ./ dxL
    sw0R = ((wR .+ swR .* dxR) .- w) ./ dxR
    gaL = pdf_slope(prim, sw0L, inK)
    gaR = pdf_slope(prim, sw0R, inK)
    sw =
        -prim[1] .* (
            moments_conserve_slope(gaL, MuL, Mxi, 1) .+
            moments_conserve_slope(gaR, MuR, Mxi, 1)
        )
    # ga = pdf_slope(prim, sw, inK)
    # sw = -prim[1] .* moments_conserve_slope(ga, Mu, Mxi, 1)
    gaT = pdf_slope(prim, sw, inK)

    # time-integration constants
    Mt = zeros(5)
    Mt[4] = tau * (1.0 - exp(-dt / tau))
    Mt[5] = -tau * dt * exp(-dt / tau) + tau * Mt[4]
    Mt[1] = dt - Mt[4]
    Mt[2] = -tau * Mt[1] + Mt[5]
    Mt[3] = 0.5 * dt^2 - tau * Mt[1]

    # flux related to central distribution
    Muv = moments_conserve(Mu, Mxi, 1, 0)
    MauL = moments_conserve_slope(gaL, MuL, Mxi, 2)
    MauR = moments_conserve_slope(gaR, MuR, Mxi, 2)
    # Mau = moments_conserve_slope(ga, MuR, Mxi, 2)
    MauT = moments_conserve_slope(gaT, Mu, Mxi, 1)

    fw .=
        Mt[1] .* prim[1] .* Muv .+ Mt[2] .* prim[1] .* (MauL .+ MauR) .+
        Mt[3] .* prim[1] .* MauT
    # fw .= Mt[1] .* prim[1] .* Muv .+ Mt[2] .* prim[1] .* Mau .+ Mt[3] .* prim[1] .* MauT

    # flux related to upwind distribution
    MuvL = moments_conserve(MuL1, Mxi1, 1, 0)
    MauL = moments_conserve_slope(faL, MuL1, Mxi1, 2)
    MauLT = moments_conserve_slope(faTL, MuL1, Mxi1, 1)

    MuvR = moments_conserve(MuR2, Mxi2, 1, 0)
    MauR = moments_conserve_slope(faR, MuR2, Mxi2, 2)
    MauRT = moments_conserve_slope(faTR, MuR2, Mxi2, 1)

    @. fw +=
        Mt[4] * primL[1] * MuvL - (Mt[5] + tau * Mt[4]) * primL[1] * MauL -
        tau * Mt[4] * primL[1] * MauLT + Mt[4] * primR[1] * MuvR -
        (Mt[5] + tau * Mt[4]) * primR[1] * MauR - tau * Mt[4] * primR[1] * MauRT
    # @. fw += Mt[4] * primL[1] * MuvL + Mt[4] * primR[1] * MuvR

    #--- fluxes of distribution functions ---#
    δ = heaviside.(u)

    HL = maxwellian(u, primL)
    BL = HL .* inK ./ (2.0 * primL[end])
    HR = maxwellian(u, primR)
    BR = HR .* inK ./ (2.0 * primR[end])

    H = maxwellian(u, prim)
    B = H .* inK ./ (2.0 * prim[end])

    @. fh =
        Mt[1] * u * H +
        Mt[2] * u^2 * (gaL[1] * H + gaL[2] * u * H + 0.5 * gaL[3] * (u^2 * H + B)) * δ +
        Mt[2] *
        u^2 *
        (gaR[1] * H + gaR[2] * u * H + 0.5 * gaR[3] * (u^2 * H + B)) *
        (1.0 - δ) +
        Mt[3] * u * (gaT[1] * H + gaT[2] * u * H + 0.5 * gaT[3] * (u^2 * H + B)) +
        Mt[4] * u * HL * δ -
        (Mt[5] + tau * Mt[4]) *
        u^2 *
        (faL[1] * HL + faL[2] * u * HL + 0.5 * faL[3] * (u^2 * HL + BL)) *
        δ + Mt[4] * u * HR * (1.0 - δ) -
        (Mt[5] + tau * Mt[4]) *
        u^2 *
        (faR[1] * HR + faR[2] * u * HR + 0.5 * faR[3] * (u^2 * HR + BR)) *
        (1.0 - δ) -
        tau *
        Mt[4] *
        u *
        (faTL[1] * HL + faTL[2] * u * HL + 0.5 * faTL[3] * (u^2 * HL + BL)) *
        δ -
        tau *
        Mt[4] *
        u *
        (faTR[1] * HR + faTR[2] * u * HR + 0.5 * faTR[3] * (u^2 * HR + BR)) *
        (1.0 - δ)

    @. fb =
        Mt[1] * u * B +
        Mt[2] *
        u^2 *
        (gaL[1] * B + gaL[2] * u * B + 0.5 * gaL[3] * (u^2 * B + Mxi[2] * H)) *
        δ +
        Mt[2] *
        u^2 *
        (gaR[1] * B + gaR[2] * u * B + 0.5 * gaR[3] * (u^2 * B + Mxi[2] * H)) *
        (1.0 - δ) +
        Mt[3] * u * (gaT[1] * H + gaT[2] * u * H + 0.5 * gaT[3] * (u^2 * B + Mxi[2] * H)) +
        Mt[4] * u * BL * δ -
        (Mt[5] + tau * Mt[4]) *
        u^2 *
        (faL[1] * BL + faL[2] * u * BL + 0.5 * faL[3] * (u^2 * BL + Mxi[2] * HL)) *
        δ + Mt[4] * u * BR * (1.0 - δ) -
        (Mt[5] + tau * Mt[4]) *
        u^2 *
        (faR[1] * BR + faR[2] * u * BR + 0.5 * faR[3] * (u^2 * BR + Mxi[2] * HR)) *
        (1.0 - δ) -
        tau *
        Mt[4] *
        u *
        (faTL[1] * BL + faTL[2] * u * BL + 0.5 * faTL[3] * (u^2 * BL + Mxi[2] * HL)) *
        δ -
        tau *
        Mt[4] *
        u *
        (faTR[1] * BR + faTR[2] * u * BR + 0.5 * faTR[3] * (u^2 * BR + Mxi[2] * HR)) *
        (1.0 - δ)

    return nothing

end

# ------------------------------------------------------------
# 2D
# ------------------------------------------------------------

function flux_gks!(
    fw::X,
    wL::Y,
    wR::Y,
    inK::Real,
    γ::Real,
    μᵣ::Real,
    ω::Real,
    dt::Real,
    dxL::Real,
    dxR::Real,
    dy::Real,
    swL = zero(wL)::Y,
    swR = zero(wR)::Y,
) where {X<:AbstractArray{<:AbstractFloat,1},Y<:AbstractArray{<:AbstractFloat,1}}

    primL = conserve_prim(wL, γ)
    primR = conserve_prim(wR, γ)

    Mu1, Mv1, Mxi1, MuL1, MuR1 = gauss_moments(primL, inK)
    Mu2, Mv2, Mxi2, MuL2, MuR2 = gauss_moments(primR, inK)

    w =
        primL[1] .* moments_conserve(MuL1, Mv1, Mxi1, 0, 0, 0) .+
        primR[1] .* moments_conserve(MuR2, Mv2, Mxi2, 0, 0, 0)
    prim = conserve_prim(w, γ)
    tau =
        vhs_collision_time(prim, μᵣ, ω) +
        2.0 * dt * abs(primL[1] / primL[end] - primR[1] / primR[end]) /
        (primL[1] / primL[end] + primR[1] / primR[end])

    faL = pdf_slope(primL, swL, inK)
    sw = -primL[1] .* moments_conserve_slope(faL, Mu1, Mv1, Mxi1, 1, 0)
    faTL = pdf_slope(primL, sw, inK)

    faR = pdf_slope(primR, swR, inK)
    sw = -primR[1] .* moments_conserve_slope(faR, Mu2, Mv1, Mxi2, 1, 0)
    faTR = pdf_slope(primR, sw, inK)

    Mu, Mv, Mxi, MuL, MuR = gauss_moments(prim, inK)
    sw0L = (w .- (wL .- swL .* dxL)) ./ dxL
    sw0R = ((wR .+ swR .* dxR) .- w) ./ dxR
    gaL = pdf_slope(prim, sw0L, inK)
    gaR = pdf_slope(prim, sw0R, inK)
    sw =
        -prim[1] .* (
            moments_conserve_slope(gaL, MuL, Mv, Mxi, 1, 0) .+
            moments_conserve_slope(gaR, MuR, Mv, Mxi, 1, 0)
        )
    # ga = pdf_slope(prim, sw, inK)
    # sw = -prim[1] .* moments_conserve_slope(ga, Mu, Mv, Mxi, 1, 0)
    gaT = pdf_slope(prim, sw, inK)

    # time-integration constants
    Mt = zeros(5)
    Mt[4] = tau * (1.0 - exp(-dt / tau))
    Mt[5] = -tau * dt * exp(-dt / tau) + tau * Mt[4]
    Mt[1] = dt - Mt[4]
    Mt[2] = -tau * Mt[1] + Mt[5]
    Mt[3] = 0.5 * dt^2 - tau * Mt[1]

    # flux related to central distribution
    Muv = moments_conserve(Mu, Mv, Mxi, 1, 0, 0)
    MauL = moments_conserve_slope(gaL, MuL, Mv, Mxi, 2, 0)
    MauR = moments_conserve_slope(gaR, MuR, Mv, Mxi, 2, 0)
    # Mau = moments_conserve_slope(ga, Mu, Mv, Mxi, 2, 0)
    MauT = moments_conserve_slope(gaT, Mu, Mv, Mxi, 1, 0)

    fw .=
        Mt[1] .* prim[1] .* Muv .+ Mt[2] .* prim[1] .* (MauL .+ MauR) .+
        Mt[3] .* prim[1] .* MauT
    # fw .= Mt[1] .* prim[1] .* Muv .+ Mt[2] .* prim[1] .* Mau .+ Mt[3] .* prim[1] .* MauT

    # flux related to upwind distribution
    MuvL = moments_conserve(MuL1, Mv1, Mxi1, 1, 0, 0)
    MauL = moments_conserve_slope(faL, MuL1, Mv1, Mxi1, 2, 0)
    MauLT = moments_conserve_slope(faTL, MuL1, Mv1, Mxi1, 1, 0)

    MuvR = moments_conserve(MuR2, Mv2, Mxi2, 1, 0, 0)
    MauR = moments_conserve_slope(faR, MuR2, Mv2, Mxi2, 2, 0)
    MauRT = moments_conserve_slope(faTR, MuR2, Mv2, Mxi2, 1, 0)

    @. fw +=
        Mt[4] * primL[1] * MuvL - (Mt[5] + tau * Mt[4]) * primL[1] * MauL -
        tau * Mt[4] * primL[1] * MauLT + Mt[4] * primR[1] * MuvR -
        (Mt[5] + tau * Mt[4]) * primR[1] * MauR - tau * Mt[4] * primR[1] * MauRT
    # @. fw += Mt[4] * primL[1] * MuvL + Mt[4] * primR[1] * MuvR

    fw .*= dy

    return nothing

end

#--- mixture ---#
function flux_gks!(
    fw::X,
    wL::Y,
    wR::Y,
    inK,
    γ,
    mi,
    ni,
    me,
    ne,
    Kn,
    dt,
    dxL,
    dxR,
    len,
    swL = zero(wL)::Y,
    swR = zero(wR)::Y,
) where {X<:AbstractArray{<:AbstractFloat,2},Y<:AbstractArray{<:Real,2}}

    primL = mixture_conserve_prim(wL, γ)
    primR = mixture_conserve_prim(wR, γ)

    Mu1, Mv1, Mw1, MuL1, MuR1 = mixture_gauss_moments(primL, inK)
    Mu2, Mv2, Mw2, MuL2, MuR2 = mixture_gauss_moments(primR, inK)

    w =
        primL[1] .* mixture_moments_conserve(MuL1, Mv1, Mw1, 0, 0, 0) .+
        primR[1] .* mixture_moments_conserve(MuR2, Mv2, Mw2, 0, 0, 0)
    prim = mixture_conserve_prim(w, γ)
    tau = aap_hs_collision_time(prim, mi, ni, me, ne, Kn)
    for i in eachindex(tau)
        tau[i] +=
            2.0 * dt * abs(primL[1, i] / primL[end, i] - primR[1, i] / primR[end, i]) /
            (primL[1, i] / primL[end, i] + primR[1, i] / primR[end, i])
    end
    prim = aap_hs_prim(prim, tau, mi, ni, me, ne, Kn) # pseudo primitive variables

    faL = mixture_pdf_slope(primL, swL, inK)
    sw = mixture_moments_conserve_slope(faL, Mu1, Mv1, Mw1, 1, 0)
    for j in axes(sw, 2)
        sw[:, j] .*= -primL[1, j]
    end
    faTL = mixture_pdf_slope(primL, sw, inK)
    faR = mixture_pdf_slope(primR, swR, inK)
    sw = mixture_moments_conserve_slope(faR, Mu2, Mv2, Mw2, 1, 0)
    for j in axes(sw, 2)
        sw[:, j] .*= -primR[1, j]
    end
    faTR = mixture_pdf_slope(primR, sw, inK)

    Mu, Mv, Mw, MuL, MuR = mixture_gauss_moments(prim, inK)
    sw0L = (w .- (wL .- swL .* dxL)) ./ dxL
    sw0R = ((wR .+ swR .* dxR) .- w) ./ dxR
    gaL = mixture_pdf_slope(prim, sw0L, inK)
    gaR = mixture_pdf_slope(prim, sw0R, inK)
    sw =
        mixture_moments_conserve_slope(gaL, MuL, Mv, Mw, 1, 0) .+
        mixture_moments_conserve_slope(gaR, MuR, Mv, Mw, 1, 0)
    for j in axes(sw, 2)
        sw[:, j] .*= -prim[1, j]
    end
    gaT = mixture_pdf_slope(prim, sw, inK)

    # time-integration constants
    Mt = zeros(5, axes(fw, 2))
    for j in axes(Mt, 2)
        Mt[4, j] = tau[j] * (1.0 - exp(-dt / tau[j]))
        Mt[5, j] = -tau[j] * dt * exp(-dt / tau[j]) + tau[j] * Mt[4]
        Mt[1, j] = dt - Mt[4, j]
        Mt[2, j] = -tau[j] * Mt[1, j] + Mt[5, j]
        Mt[3, j] = 0.5 * dt^2 - tau[j] * Mt[1, j]
    end

    # flux related to central distribution
    Muv = mixture_moments_conserve(Mu, Mv, Mw, 1, 0, 0)
    MauL = mixture_moments_conserve_slope(gaL, MuL, Mv, Mw, 2, 0)
    MauR = mixture_moments_conserve_slope(gaR, MuR, Mv, Mw, 2, 0)
    MauT = mixture_moments_conserve_slope(gaT, Mu, Mv, Mw, 1, 0)

    for j in axes(fw, 2)
        @. fw[:, j] =
            Mt[1, j] * prim[1, j] * Muv[:, j] +
            Mt[2, j] * prim[1, j] * (MauL[:, j] + MauR[:, j]) +
            Mt[3, j] * prim[1, j] * MauT[:, j]
    end

    # flux related to upwind distribution
    MuvL = mixture_moments_conserve(MuL1, Mv1, Mw1, 1, 0, 0)
    MauL = mixture_moments_conserve_slope(faL, MuL1, Mv1, Mw1, 2, 0)
    MauLT = mixture_moments_conserve_slope(faTL, MuL1, Mv1, Mw1, 1, 0)

    MuvR = mixture_moments_conserve(MuR2, Mv2, Mw2, 1, 0, 0)
    MauR = mixture_moments_conserve_slope(faR, MuR2, Mv2, Mw2, 2, 0)
    MauRT = mixture_moments_conserve_slope(faTR, MuR2, Mv2, Mw2, 1, 0)

    for j in axes(fw, 2)
        @. fw[:, j] +=
            Mt[4, j] * primL[1, j] * MuvL[:, j] -
            (Mt[5, j] + tau[j] * Mt[4, j]) * primL[1, j] * MauL[:, j] -
            tau[j] * Mt[4, j] * primL[1, j] * MauLT[:, j] +
            Mt[4, j] * primR[1, j] * MuvR[:, j] -
            (Mt[5, j] + tau[j] * Mt[4, j]) * primR[1, j] * MauR[:, j] -
            tau[j] * Mt[4, j] * primR[1, j] * MauRT[:, j]
    end

    @. fw .* len

    return nothing

end

#--- 2D1F2V ---#
function flux_gks!(
    fw::T1,
    ff::T2,
    wL::T3,
    wR::T3,
    u::T4,
    v::T4,
    inK::Real,
    γ::Real,
    μᵣ::Real,
    ω::Real,
    dt::Real,
    dxL::Real,
    dxR::Real,
    dy::Real,
    swL::T3,
    swR::T3,
) where {
    T1<:AbstractVector{<:AbstractFloat},
    T2<:AbstractMatrix{<:AbstractFloat},
    T3<:AbstractVector{<:AbstractFloat},
    T4<:AbstractMatrix{<:AbstractFloat},
}

    primL = conserve_prim(wL, γ)
    primR = conserve_prim(wR, γ)

    Mu1, Mv1, Mxi1, MuL1, MuR1 = gauss_moments(primL, inK)
    Mu2, Mv2, Mxi2, MuL2, MuR2 = gauss_moments(primR, inK)

    w =
        primL[1] .* moments_conserve(MuL1, Mv1, Mxi1, 0, 0, 0) .+
        primR[1] .* moments_conserve(MuR2, Mv2, Mxi2, 0, 0, 0)
    prim = conserve_prim(w, γ)
    tau =
        vhs_collision_time(prim, μᵣ, ω) +
        2.0 * dt * abs(primL[1] / primL[end] - primR[1] / primR[end]) /
        (primL[1] / primL[end] + primR[1] / primR[end])

    faL = pdf_slope(primL, swL, inK)
    sw = -primL[1] .* moments_conserve_slope(faL, Mu1, Mv1, Mxi1, 1, 0)
    faTL = pdf_slope(primL, sw, inK)

    faR = pdf_slope(primR, swR, inK)
    sw = -primR[1] .* moments_conserve_slope(faR, Mu2, Mv1, Mxi2, 1, 0)
    faTR = pdf_slope(primR, sw, inK)

    Mu, Mv, Mxi, MuL, MuR = gauss_moments(prim, inK)
    sw0L = (w .- (wL .- swL .* dxL)) ./ dxL
    sw0R = ((wR .+ swR .* dxR) .- w) ./ dxR
    gaL = pdf_slope(prim, sw0L, inK)
    gaR = pdf_slope(prim, sw0R, inK)
    sw =
        -prim[1] .* (
            moments_conserve_slope(gaL, MuL, Mv, Mxi, 1, 0) .+
            moments_conserve_slope(gaR, MuR, Mv, Mxi, 1, 0)
        )
    # ga = pdf_slope(prim, sw, inK)
    # sw = -prim[1] .* moments_conserve_slope(ga, Mu, Mv, Mxi, 1, 0)
    gaT = pdf_slope(prim, sw, inK)

    # time-integration constants
    Mt = zeros(5)
    Mt[4] = tau * (1.0 - exp(-dt / tau))
    Mt[5] = -tau * dt * exp(-dt / tau) + tau * Mt[4]
    Mt[1] = dt - Mt[4]
    Mt[2] = -tau * Mt[1] + Mt[5]
    Mt[3] = 0.5 * dt^2 - tau * Mt[1]

    # flux related to central distribution
    Muv = moments_conserve(Mu, Mv, Mxi, 1, 0, 0)
    MauL = moments_conserve_slope(gaL, MuL, Mv, Mxi, 2, 0)
    MauR = moments_conserve_slope(gaR, MuR, Mv, Mxi, 2, 0)
    # Mau = moments_conserve_slope(ga, Mu, Mv, Mxi, 2, 0)
    MauT = moments_conserve_slope(gaT, Mu, Mv, Mxi, 1, 0)

    fw .=
        Mt[1] .* prim[1] .* Muv .+ Mt[2] .* prim[1] .* (MauL .+ MauR) .+
        Mt[3] .* prim[1] .* MauT
    # fw .= Mt[1] .* prim[1] .* Muv .+ Mt[2] .* prim[1] .* Mau .+ Mt[3] .* prim[1] .* MauT

    # flux related to upwind distribution
    MuvL = moments_conserve(MuL1, Mv1, Mxi1, 1, 0, 0)
    MauL = moments_conserve_slope(faL, MuL1, Mv1, Mxi1, 2, 0)
    MauLT = moments_conserve_slope(faTL, MuL1, Mv1, Mxi1, 1, 0)

    MuvR = moments_conserve(MuR2, Mv2, Mxi2, 1, 0, 0)
    MauR = moments_conserve_slope(faR, MuR2, Mv2, Mxi2, 2, 0)
    MauRT = moments_conserve_slope(faTR, MuR2, Mv2, Mxi2, 1, 0)

    @. fw +=
        Mt[4] * primL[1] * MuvL - (Mt[5] + tau * Mt[4]) * primL[1] * MauL -
        tau * Mt[4] * primL[1] * MauLT + Mt[4] * primR[1] * MuvR -
        (Mt[5] + tau * Mt[4]) * primR[1] * MauR - tau * Mt[4] * primR[1] * MauRT
    # @. fw += Mt[4] * primL[1] * MuvL + Mt[4] * primR[1] * MuvR

    fw .*= dy

    #--- fluxes of distribution functions ---#
    δ = heaviside.(u)
    HL = maxwellian(u, v, primL)
    HR = maxwellian(u, v, primR)
    H = maxwellian(u, v, prim)

    @. ff =
        Mt[1] * u * H +
        Mt[2] *
        u^2 *
        (gaL[1] * H + gaL[2] * u * H + gaL[3] * v * H + 0.5 * gaL[4] * (u^2 + v^2) * H) *
        δ +
        Mt[2] *
        u^2 *
        (gaR[1] * H + gaR[2] * u * H + gaR[3] * v * H + 0.5 * gaR[4] * (u^2 + v^2) * H) *
        (1.0 - δ) +
        Mt[3] *
        u *
        (gaT[1] * H + gaT[2] * u * H + gaT[3] * v * H + 0.5 * gaT[4] * (u^2 + v^2) * H) +
        Mt[4] * u * HL * δ -
        (Mt[5] + tau * Mt[4]) *
        u^2 *
        (
            faL[1] * HL +
            faL[2] * u * HL +
            faL[3] * v * HL +
            0.5 * faL[4] * (u^2 + v^2) * HL
        ) *
        δ + Mt[4] * u * HR * (1.0 - δ) -
        (Mt[5] + tau * Mt[4]) *
        u^2 *
        (
            faR[1] * HR +
            faR[2] * u * HR +
            faR[3] * v * HL +
            0.5 * faR[4] * (u^2 + v^2) * HR
        ) *
        (1.0 - δ) -
        tau *
        Mt[4] *
        u *
        (
            faTL[1] * HL +
            faTL[2] * u * HL +
            faTL[3] * v * HL +
            0.5 * faTL[4] * (u^2 + v^2) * HL
        ) *
        δ -
        tau *
        Mt[4] *
        u *
        (
            faTR[1] * HR +
            faTR[2] * u * HR +
            faTR[3] * v * HR +
            0.5 * faTR[4] * (u^2 + v^2) * HR
        ) *
        (1.0 - δ)

    ff .*= dy

    return nothing

end

#--- 2D2F2V ---#
function flux_gks!(
    fw::T1,
    fh::T2,
    fb::T2,
    wL::T3,
    wR::T3,
    u::T4,
    v::T4,
    inK::Real,
    γ::Real,
    μᵣ::Real,
    ω::Real,
    dt::Real,
    dxL::Real,
    dxR::Real,
    dy::Real,
    swL::T3,
    swR::T3,
) where {
    T1<:AbstractVector{<:AbstractFloat},
    T2<:AbstractMatrix{<:AbstractFloat},
    T3<:AbstractVector{<:AbstractFloat},
    T4<:AbstractMatrix{<:AbstractFloat},
}

    primL = conserve_prim(wL, γ)
    primR = conserve_prim(wR, γ)

    Mu1, Mv1, Mxi1, MuL1, MuR1 = gauss_moments(primL, inK)
    Mu2, Mv2, Mxi2, MuL2, MuR2 = gauss_moments(primR, inK)

    w =
        primL[1] .* moments_conserve(MuL1, Mv1, Mxi1, 0, 0, 0) .+
        primR[1] .* moments_conserve(MuR2, Mv2, Mxi2, 0, 0, 0)
    prim = conserve_prim(w, γ)
    tau =
        vhs_collision_time(prim, μᵣ, ω) +
        2.0 * dt * abs(primL[1] / primL[end] - primR[1] / primR[end]) /
        (primL[1] / primL[end] + primR[1] / primR[end])

    faL = pdf_slope(primL, swL, inK)
    sw = -primL[1] .* moments_conserve_slope(faL, Mu1, Mv1, Mxi1, 1, 0)
    faTL = pdf_slope(primL, sw, inK)

    faR = pdf_slope(primR, swR, inK)
    sw = -primR[1] .* moments_conserve_slope(faR, Mu2, Mv1, Mxi2, 1, 0)
    faTR = pdf_slope(primR, sw, inK)

    Mu, Mv, Mxi, MuL, MuR = gauss_moments(prim, inK)
    sw0L = (w .- (wL .- swL .* dxL)) ./ dxL
    sw0R = ((wR .+ swR .* dxR) .- w) ./ dxR
    gaL = pdf_slope(prim, sw0L, inK)
    gaR = pdf_slope(prim, sw0R, inK)
    sw =
        -prim[1] .* (
            moments_conserve_slope(gaL, MuL, Mv, Mxi, 1, 0) .+
            moments_conserve_slope(gaR, MuR, Mv, Mxi, 1, 0)
        )
    # ga = pdf_slope(prim, sw, inK)
    # sw = -prim[1] .* moments_conserve_slope(ga, Mu, Mv, Mxi, 1, 0)
    gaT = pdf_slope(prim, sw, inK)

    # time-integration constants
    Mt = zeros(5)
    Mt[4] = tau * (1.0 - exp(-dt / tau))
    Mt[5] = -tau * dt * exp(-dt / tau) + tau * Mt[4]
    Mt[1] = dt - Mt[4]
    Mt[2] = -tau * Mt[1] + Mt[5]
    Mt[3] = 0.5 * dt^2 - tau * Mt[1]

    # flux related to central distribution
    Muv = moments_conserve(Mu, Mv, Mxi, 1, 0, 0)
    MauL = moments_conserve_slope(gaL, MuL, Mv, Mxi, 2, 0)
    MauR = moments_conserve_slope(gaR, MuR, Mv, Mxi, 2, 0)
    # Mau = moments_conserve_slope(ga, Mu, Mv, Mxi, 2, 0)
    MauT = moments_conserve_slope(gaT, Mu, Mv, Mxi, 1, 0)

    fw .=
        Mt[1] .* prim[1] .* Muv .+ Mt[2] .* prim[1] .* (MauL .+ MauR) .+
        Mt[3] .* prim[1] .* MauT
    # fw .= Mt[1] .* prim[1] .* Muv .+ Mt[2] .* prim[1] .* Mau .+ Mt[3] .* prim[1] .* MauT

    # flux related to upwind distribution
    MuvL = moments_conserve(MuL1, Mv1, Mxi1, 1, 0, 0)
    MauL = moments_conserve_slope(faL, MuL1, Mv1, Mxi1, 2, 0)
    MauLT = moments_conserve_slope(faTL, MuL1, Mv1, Mxi1, 1, 0)

    MuvR = moments_conserve(MuR2, Mv2, Mxi2, 1, 0, 0)
    MauR = moments_conserve_slope(faR, MuR2, Mv2, Mxi2, 2, 0)
    MauRT = moments_conserve_slope(faTR, MuR2, Mv2, Mxi2, 1, 0)

    @. fw +=
        Mt[4] * primL[1] * MuvL - (Mt[5] + tau * Mt[4]) * primL[1] * MauL -
        tau * Mt[4] * primL[1] * MauLT + Mt[4] * primR[1] * MuvR -
        (Mt[5] + tau * Mt[4]) * primR[1] * MauR - tau * Mt[4] * primR[1] * MauRT
    # @. fw += Mt[4] * primL[1] * MuvL + Mt[4] * primR[1] * MuvR

    fw .*= dy

    #--- fluxes of distribution functions ---#
    δ = heaviside.(u)
    HL = maxwellian(u, v, primL)
    HR = maxwellian(u, v, primR)
    H = maxwellian(u, v, prim)
    BL = HL .* inK ./ (2.0 * primL[end])
    BR = HR .* inK ./ (2.0 * primR[end])
    B = H .* inK ./ (2.0 * prim[end])

    @. fh =
        Mt[1] * u * H +
        Mt[2] *
        u^2 *
        (
            gaL[1] * H +
            gaL[2] * u * H +
            gaL[3] * v * H +
            0.5 * gaL[4] * ((u^2 + v^2) * H + B)
        ) *
        δ +
        Mt[2] *
        u^2 *
        (
            gaR[1] * H +
            gaR[2] * u * H +
            gaR[3] * v * H +
            0.5 * gaR[4] * ((u^2 + v^2) * H + B)
        ) *
        (1.0 - δ) +
        Mt[3] *
        u *
        (
            gaT[1] * H +
            gaT[2] * u * H +
            gaT[3] * v * H +
            0.5 * gaT[4] * ((u^2 + v^2) * H + B)
        ) +
        Mt[4] * u * HL * δ -
        (Mt[5] + tau * Mt[4]) *
        u^2 *
        (
            faL[1] * HL +
            faL[2] * u * HL +
            faL[3] * v * HL +
            0.5 * faL[4] * ((u^2 + v^2) * HL + BL)
        ) *
        δ + Mt[4] * u * HR * (1.0 - δ) -
        (Mt[5] + tau * Mt[4]) *
        u^2 *
        (
            faR[1] * HR +
            faR[2] * u * HR +
            faR[3] * v * HR +
            0.5 * faR[4] * ((u^2 + v^2) * HR + BR)
        ) *
        (1.0 - δ) -
        tau *
        Mt[4] *
        u *
        (
            faTL[1] * HL +
            faTL[2] * u * HL +
            faTL[3] * v * HL +
            0.5 * faTL[4] * ((u^2 + v^2) * HL + BL)
        ) *
        δ -
        tau *
        Mt[4] *
        u *
        (
            faTR[1] * HR +
            faTR[2] * u * HR +
            faTR[3] * v * HR +
            0.5 * faTR[4] * ((u^2 + v^2) * HR + BR)
        ) *
        (1.0 - δ)
    @. fb =
        Mt[1] * u * B +
        Mt[2] *
        u^2 *
        (
            gaL[1] * B +
            gaL[2] * u * B +
            gaL[3] * v * B +
            0.5 * gaL[4] * ((u^2 + v^2) * B + Mxi[2] * H)
        ) *
        δ +
        Mt[2] *
        u^2 *
        (
            gaR[1] * B +
            gaR[2] * u * B +
            gaR[3] * v * B +
            0.5 * gaR[4] * ((u^2 + v^2) * B + Mxi[2] * H)
        ) *
        (1.0 - δ) +
        Mt[3] *
        u *
        (
            gaT[1] * B +
            gaT[2] * u * B +
            gaT[3] * v * B +
            0.5 * gaT[4] * ((u^2 + v^2) * B + Mxi[2] * H)
        ) +
        Mt[4] * u * BL * δ -
        (Mt[5] + tau * Mt[4]) *
        u^2 *
        (
            faL[1] * BL +
            faL[2] * u * BL +
            faL[3] * v * BL +
            0.5 * faL[4] * ((u^2 + v^2) * BL + Mxi[2] * HL)
        ) *
        δ + Mt[4] * u * BR * (1.0 - δ) -
        (Mt[5] + tau * Mt[4]) *
        u^2 *
        (
            faR[1] * BR +
            faR[2] * u * BR +
            faR[3] * v * BR +
            0.5 * faR[4] * ((u^2 + v^2) * BR + Mxi[2] * HR)
        ) *
        (1.0 - δ) -
        tau *
        Mt[4] *
        u *
        (
            faTL[1] * BL +
            faTL[2] * u * BL +
            faTL[3] * v * BL +
            0.5 * faTL[4] * ((u^2 + v^2) * BL + Mxi[2] * HL)
        ) *
        δ -
        tau *
        Mt[4] *
        u *
        (
            faTR[1] * BR +
            faTR[2] * u * BR +
            faTR[3] * v * BR +
            0.5 * faTR[4] * ((u^2 + v^2) * BR + Mxi[2] * HR)
        ) *
        (1.0 - δ)

    fh .*= dy
    fb .*= dy

    return nothing

end

#--- FR ---#
function flux_gks!(
    fw::X,
    wL::Y,
    wR::Y,
    inK::Real,
    γ::Real,
    μᵣ::Real,
    ω::Real,
    dt::Real,
    swL = zero(wL)::Y,
    swR = zero(wR)::Y,
) where {X<:AbstractArray{<:AbstractFloat,1},Y<:AbstractArray{<:AbstractFloat,1}}
    primL = conserve_prim(wL, γ)
    primR = conserve_prim(wR, γ)

    Mu1, Mv1, Mxi1, MuL1, MuR1 = gauss_moments(primL, inK)
    Mu2, Mv2, Mxi2, MuL2, MuR2 = gauss_moments(primR, inK)

    w =
        primL[1] .* moments_conserve(MuL1, Mv1, Mxi1, 0, 0, 0) .+
        primR[1] .* moments_conserve(MuR2, Mv2, Mxi2, 0, 0, 0)
    prim = conserve_prim(w, γ)
    tau =
        vhs_collision_time(prim, μᵣ, ω) +
        2.0 * dt * abs(primL[1] / primL[end] - primR[1] / primR[end]) /
        (primL[1] / primL[end] + primR[1] / primR[end])

    faL = pdf_slope(primL, swL, inK)
    sw = -primL[1] .* moments_conserve_slope(faL, Mu1, Mv1, Mxi1, 1, 0)
    faTL = pdf_slope(primL, sw, inK)

    faR = pdf_slope(primR, swR, inK)
    sw = -primR[1] .* moments_conserve_slope(faR, Mu2, Mv1, Mxi2, 1, 0)
    faTR = pdf_slope(primR, sw, inK)

    MuvL = moments_conserve(MuL1, Mv1, Mxi1, 1, 0, 0)
    MauL = moments_conserve_slope(faL, MuL1, Mv1, Mxi1, 2, 0)
    MauLT = moments_conserve_slope(faTL, MuL1, Mv1, Mxi1, 1, 0)

    MuvR = moments_conserve(MuR2, Mv2, Mxi2, 1, 0, 0)
    MauR = moments_conserve_slope(faR, MuR2, Mv2, Mxi2, 2, 0)
    MauRT = moments_conserve_slope(faTR, MuR2, Mv2, Mxi2, 1, 0)

    @. fw =
        primL[1] * (MuvL - tau * MauL - tau * MauLT) +
        primR[1] * (MuvR - tau * MauR - tau * MauRT)

    return nothing
end

"""
    flux_ugks!(
        fw::T1,
        fh::T2,
        fb::T2,
        wL::T3,
        hL::T4,
        bL::T4,
        wR::T3,
        hR::T4,
        bR::T4,
        u::T5,
        ω::T5,
        inK,
        γ,
        visRef,
        visIdx,
        pr,
        dt,
        dxL,
        dxR,
        shL = zeros(eltype(hL), axes(hL))::T4,
        sbL = zeros(eltype(bL), axes(bL))::T4,
        shR = zeros(eltype(hR), axes(hR))::T4,
        sbR = zeros(eltype(bR), axes(bR))::T4,
    ) where {
        T1<:AbstractArray{<:AbstractFloat,1},
        T2<:AbstractArray{<:AbstractFloat,1},
        T3<:AbstractArray{<:Real,1},
        T4<:AbstractArray{<:AbstractFloat,1},
        T5<:AbstractArray{<:AbstractFloat,1},
    }

    flux_ugks!(
        fw::T1,
        fh::T2,
        fb::T2,
        wL::T3,
        hL::T4,
        bL::T4,
        wR::T3,
        hR::T4,
        bR::T4,
        u::T5,
        v::T5,
        ω::T5,
        inK,
        γ,
        visRef,
        visIdx,
        pr,
        dt,
        dxL,
        dxR,
        len,
        shL = zeros(eltype(hL), axes(hL))::T4,
        sbL = zeros(eltype(bL), axes(bL))::T4,
        shR = zeros(eltype(hR), axes(hR))::T4,
        sbR = zeros(eltype(bR), axes(bR))::T4,
    ) where {
        T1<:AbstractArray{<:AbstractFloat,1},
        T2<:AbstractArray{<:AbstractFloat,2},
        T3<:AbstractArray{<:Real,1},
        T4<:AbstractArray{<:AbstractFloat,2},
        T5<:AbstractArray{<:AbstractFloat,2},
    }

    flux_ugks!(
        fw::T1,
        fh0::T2,
        fh1::T2,
        fh2::T2,
        wL::T3,
        h0L::T4,
        h1L::T4,
        h2L::T4,
        wR::T3,
        h0R::T4,
        h1R::T4,
        h2R::T4,
        u::T5,
        v::T5,
        ω::T5,
        inK,
        γ,
        mi,
        ni,
        me,
        ne,
        Kn,
        dt,
        dxL,
        dxR,
        len,
        sh0L = zeros(eltype(h0L), axes(h0L))::T4,
        sh1L = zeros(eltype(h1L), axes(h1L))::T4,
        sh2L = zeros(eltype(h2L), axes(h2L))::T4,
        sh0R = zeros(eltype(h0R), axes(h0R))::T4,
        sh1R = zeros(eltype(h1R), axes(h1R))::T4,
        sh2R = zeros(eltype(h2R), axes(h2R))::T4,
    ) where {
        T1<:AbstractArray{<:AbstractFloat,2},
        T2<:AbstractArray{<:AbstractFloat,3},
        T3<:AbstractArray{<:Real,2},
        T4<:AbstractArray{<:AbstractFloat,3},
        T5<:AbstractArray{<:AbstractFloat,3},
    }

Unified gas kinetic scheme (UGKS)

* @args: particle distribution functions and their slopes at left/right sides of interface
* @args: particle velocity quadrature points and weights
* @args: time step

"""
function flux_ugks!(
    fw::T1,
    fh::T2,
    fb::T2,
    wL::T3,
    hL::T4,
    bL::T4,
    wR::T3,
    hR::T4,
    bR::T4,
    u::T5,
    ω::T5,
    inK,
    γ,
    visRef,
    visIdx,
    pr,
    dt,
    dxL,
    dxR,
    shL = zeros(eltype(hL), axes(hL))::T4,
    sbL = zeros(eltype(bL), axes(bL))::T4,
    shR = zeros(eltype(hR), axes(hR))::T4,
    sbR = zeros(eltype(bR), axes(bR))::T4,
) where {
    T1<:AbstractArray{<:AbstractFloat,1},
    T2<:AbstractArray{<:AbstractFloat,1},
    T3<:AbstractArray{<:Real,1},
    T4<:AbstractArray{<:AbstractFloat,1},
    T5<:AbstractArray{<:AbstractFloat,1},
} # 1D2F flux

    #--- reconstruct initial distribution ---#
    δ = heaviside.(u)

    h = @. hL * δ + hR * (1.0 - δ)
    b = @. bL * δ + bR * (1.0 - δ)
    sh = @. shL * δ + shR * (1.0 - δ)
    sb = @. sbL * δ + sbR * (1.0 - δ)

    #--- construct interface variables ---#
    #w = moments_conserve(h, b, u, v, ω)
    #prim = conserve_prim(w, γ)

    primL = conserve_prim(wL, γ)
    primR = conserve_prim(wR, γ)
    Mu1, Mxi1, MuL1, MuR1 = gauss_moments(primL, inK)
    Muv1 = moments_conserve(MuL1, Mxi1, 0, 0)
    Mu2, Mxi2, MuL2, MuR2 = gauss_moments(primR, inK)
    Muv2 = moments_conserve(MuR2, Mxi2, 0, 0)

    w = @. primL[1] * Muv1 + primR[1] * Muv2
    prim = conserve_prim(w, γ)
    aL = pdf_slope(prim, (w .- wL) ./ dxL, inK)
    aR = pdf_slope(prim, (wR .- w) ./ dxR, inK)

    Mu, Mxi, MuL, MuR = gauss_moments(prim, inK)
    MauL = moments_conserve_slope(aL, MuL, Mxi, 1)
    MauR = moments_conserve_slope(aR, MuR, Mxi, 1)
    aT = pdf_slope(prim, -prim[1] .* (MauL .+ MauR), inK)

    #--- calculate integral time constants ---#
    tau = vhs_collision_time(prim, visRef, visIdx)

    Mt = zeros(5)
    Mt[4] = tau * (1.0 - exp(-dt / tau)) # f0
    Mt[5] = -tau * dt * exp(-dt / tau) + tau * Mt[4]
    Mt[1] = dt - Mt[4] # M0
    Mt[2] = -tau * Mt[1] + Mt[5]
    Mt[3] = dt^2 / 2.0 - tau * Mt[1]

    #--- calculate flux from M0 ---#
    Muv = moments_conserve(Mu, Mxi, 1, 0)
    MauL = moments_conserve_slope(aL, MuL, Mxi, 2)
    MauR = moments_conserve_slope(aR, MuR, Mxi, 2)
    MauT = moments_conserve_slope(aT, Mu, Mxi, 1)

    @. fw = Mt[1] * prim[1] * Muv + Mt[2] * prim[1] * (MauL + MauR) + Mt[3] * prim[1] * MauT

    #--- calculate flux from f0 ---#
    H = maxwellian(u, prim)
    B = H .* inK ./ (2.0 * prim[end])

    fw[1] += Mt[4] * sum(ω .* u .* h) - Mt[5] * sum(ω .* u .^ 2 .* sh)
    fw[2] += Mt[4] * sum(ω .* u .^ 2 .* h) - Mt[5] * sum(ω .* u .^ 3 .* sh)
    fw[3] +=
        Mt[4] * 0.5 * (sum(ω .* u .^ 3 .* h) + sum(ω .* u .* b)) -
        Mt[5] * 0.5 * (sum(ω .* u .^ 4 .* sh) + sum(ω .* u .^ 2 .* sb))

    @. fh =
        Mt[1] * u * H +
        Mt[2] * u^2 * (aL[1] * H + aL[2] * u * H + 0.5 * aL[3] * (u^2 * H + B)) * δ +
        Mt[2] *
        u^2 *
        (aR[1] * H + aR[2] * u * H + 0.5 * aR[3] * (u^2 * H + B)) *
        (1.0 - δ) +
        Mt[3] * u * (aT[1] * H + aT[2] * u * H + 0.5 * aT[3] * (u^2 * H + B)) +
        Mt[4] * u * h - Mt[5] * u^2 * sh
    @. fb =
        Mt[1] * u * B +
        Mt[2] *
        u^2 *
        (aL[1] * B + aL[2] * u * B + 0.5 * aL[3] * (u^2 * B + Mxi[2] * H)) *
        δ +
        Mt[2] *
        u^2 *
        (aR[1] * B + aR[2] * u * B + 0.5 * aR[3] * (u^2 * B + Mxi[2] * H)) *
        (1.0 - δ) +
        Mt[3] * u * (aT[1] * B + aT[2] * u * B + 0.5 * aT[3] * (u^2 * B + Mxi[2] * H)) +
        Mt[4] * u * b - Mt[5] * u^2 * sb

    return nothing

end

function flux_ugks!(
    fw::T1,
    fh::T2,
    fb::T2,
    wL::T3,
    hL::T4,
    bL::T4,
    wR::T3,
    hR::T4,
    bR::T4,
    u::T5,
    v::T5,
    ω::T5,
    inK,
    γ,
    visRef,
    visIdx,
    pr,
    dt,
    dxL,
    dxR,
    len,
    shL = zeros(eltype(hL), axes(hL))::T4,
    sbL = zeros(eltype(bL), axes(bL))::T4,
    shR = zeros(eltype(hR), axes(hR))::T4,
    sbR = zeros(eltype(bR), axes(bR))::T4,
) where {
    T1<:AbstractArray{<:AbstractFloat,1},
    T2<:AbstractArray{<:AbstractFloat,2},
    T3<:AbstractArray{<:Real,1},
    T4<:AbstractArray{<:AbstractFloat,2},
    T5<:AbstractArray{<:AbstractFloat,2},
} # 2D2F flux

    #--- reconstruct initial distribution ---#
    δ = heaviside.(u)

    h = @. hL * δ + hR * (1.0 - δ)
    b = @. bL * δ + bR * (1.0 - δ)
    sh = @. shL * δ + shR * (1.0 - δ)
    sb = @. sbL * δ + sbR * (1.0 - δ)

    #--- construct interface variables ---#
    #w = moments_conserve(h, b, u, v, ω)
    #prim = conserve_prim(w, γ)

    primL = conserve_prim(wL, γ)
    primR = conserve_prim(wR, γ)
    Mu1, Mv1, Mxi1, MuL1, MuR1 = gauss_moments(primL, inK)
    Muv1 = moments_conserve(MuL1, Mv1, Mxi1, 0, 0, 0)
    Mu2, Mv2, Mxi2, MuL2, MuR2 = gauss_moments(primR, inK)
    Muv2 = moments_conserve(MuR2, Mv2, Mxi2, 0, 0, 0)
    w = @. primL[1] * Muv1 + primR[1] * Muv2
    prim = conserve_prim(w, γ)

    aL = pdf_slope(prim, (w .- wL) ./ dxL, inK)
    aR = pdf_slope(prim, (wR .- w) ./ dxR, inK)

    Mu, Mv, Mxi, MuL, MuR = gauss_moments(prim, inK)
    MauL = moments_conserve_slope(aL, MuL, Mv, Mxi, 1, 0)
    MauR = moments_conserve_slope(aR, MuR, Mv, Mxi, 1, 0)
    aT = pdf_slope(prim, -prim[1] .* (MauL .+ MauR), inK)

    #--- calculate integral time constants ---#
    tau = vhs_collision_time(prim, visRef, visIdx)

    Mt = zeros(5)
    Mt[4] = tau * (1.0 - exp(-dt / tau)) # f0
    Mt[5] = -tau * dt * exp(-dt / tau) + tau * Mt[4]
    Mt[1] = dt - Mt[4] # M0
    Mt[2] = -tau * Mt[1] + Mt[5]
    Mt[3] = dt^2 / 2.0 - tau * Mt[1]

    # --- calculate flux from M0 ---#
    Muv = moments_conserve(Mu, Mv, Mxi, 1, 0, 0)
    MauL = moments_conserve_slope(aL, MuL, Mv, Mxi, 2, 0)
    MauR = moments_conserve_slope(aR, MuR, Mv, Mxi, 2, 0)
    MauT = moments_conserve_slope(aT, Mu, Mv, Mxi, 1, 0)

    @. fw = Mt[1] * prim[1] * Muv + Mt[2] * prim[1] * (MauL + MauR) + Mt[3] * prim[1] * MauT

    # --- calculate flux from f0 ---#
    H = maxwellian(u, v, prim)
    B = H .* inK ./ (2.0 * prim[end])

    fw[1] += Mt[4] * sum(ω .* u .* h) - Mt[5] * sum(ω .* u .^ 2 .* sh)
    fw[2] += Mt[4] * sum(ω .* u .^ 2 .* h) - Mt[5] * sum(ω .* u .^ 3 .* sh)
    fw[3] += Mt[4] * sum(ω .* v .* u .* h) - Mt[5] * sum(ω .* v .* u .^ 2 .* sh)
    fw[4] +=
        Mt[4] * 0.5 * (sum(ω .* u .* (u .^ 2 .+ v .^ 2) .* h) + sum(ω .* u .* b)) -
        Mt[5] *
        0.5 *
        (sum(ω .* u .^ 2 .* (u .^ 2 .+ v .^ 2) .* sh) + sum(ω .* u .^ 2 .* sb))

    @. fh =
        Mt[1] * u * H +
        Mt[2] *
        u^2 *
        (aL[1] * H + aL[2] * u * H + aL[3] * v * H + 0.5 * aL[4] * ((u^2 + v^2) * H + B)) *
        δ +
        Mt[2] *
        u^2 *
        (aR[1] * H + aR[2] * u * H + aR[3] * v * H + 0.5 * aR[4] * ((u^2 + v^2) * H + B)) *
        (1.0 - δ) +
        Mt[3] *
        u *
        (aT[1] * H + aT[2] * u * H + aT[3] * v * H + 0.5 * aT[4] * ((u^2 + v^2) * H + B)) +
        Mt[4] * u * h - Mt[5] * u^2 * sh
    @. fb =
        Mt[1] * u * B +
        Mt[2] *
        u^2 *
        (
            aL[1] * B +
            aL[2] * u * B +
            aL[3] * v * B +
            0.5 * aL[4] * ((u^2 + v^2) * B + Mxi[2] * H)
        ) *
        δ +
        Mt[2] *
        u^2 *
        (
            aR[1] * B +
            aR[2] * u * B +
            aR[3] * v * B +
            0.5 * aR[4] * ((u^2 + v^2) * B + Mxi[2] * H)
        ) *
        (1.0 - δ) +
        Mt[3] *
        u *
        (
            aT[1] * B +
            aT[2] * u * B +
            aT[3] * v * B +
            0.5 * aT[4] * ((u^2 + v^2) * B + Mxi[2] * H)
        ) +
        Mt[4] * u * b - Mt[5] * u^2 * sb

    # multiply interface length
    fw .*= len
    fh .*= len
    fb .*= len

    return nothing

end

# ------------------------------------------------------------
# 3F2V with AAP model
# ------------------------------------------------------------
function flux_ugks!(
    fw::T1,
    fh0::T2,
    fh1::T2,
    fh2::T2,
    wL::T3,
    h0L::T4,
    h1L::T4,
    h2L::T4,
    wR::T3,
    h0R::T4,
    h1R::T4,
    h2R::T4,
    u::T5,
    v::T5,
    ω::T5,
    inK,
    γ,
    mi,
    ni,
    me,
    ne,
    Kn,
    dt,
    dxL,
    dxR,
    len,
    sh0L = zeros(eltype(h0L), axes(h0L))::T4,
    sh1L = zeros(eltype(h1L), axes(h1L))::T4,
    sh2L = zeros(eltype(h2L), axes(h2L))::T4,
    sh0R = zeros(eltype(h0R), axes(h0R))::T4,
    sh1R = zeros(eltype(h1R), axes(h1R))::T4,
    sh2R = zeros(eltype(h2R), axes(h2R))::T4,
) where {
    T1<:AbstractArray{<:AbstractFloat,2},
    T2<:AbstractArray{<:AbstractFloat,3},
    T3<:AbstractArray{<:Real,2},
    T4<:AbstractArray{<:AbstractFloat,3},
    T5<:AbstractArray{<:AbstractFloat,3},
}

    #--- reconstruct initial distribution ---#
    δ = heaviside.(u)

    h0 = @. h0L * δ + h0R * (1.0 - δ)
    h1 = @. h1L * δ + h1R * (1.0 - δ)
    h2 = @. h2L * δ + h2R * (1.0 - δ)

    sh0 = @. sh0L * δ + sh0R * (1.0 - δ)
    sh1 = @. sh1L * δ + sh1R * (1.0 - δ)
    sh2 = @. sh2L * δ + sh2R * (1.0 - δ)

    primL = mixture_conserve_prim(wL, γ)
    primR = mixture_conserve_prim(wR, γ)

    #--- construct interface distribution ---#
    Mu1, Mv1, Mxi1, MuL1, MuR1 = mixture_gauss_moments(primL, inK)
    Muv1 = mixture_moments_conserve(MuL1, Mv1, Mxi1, 0, 0, 0)
    Mu2, Mv2, Mxi2, MuL2, MuR2 = mixture_gauss_moments(primR, inK)
    Muv2 = mixture_moments_conserve(MuR2, Mv2, Mxi2, 0, 0, 0)

    w = similar(wL)
    for j in axes(w, 2)
        @. w[:, j] = primL[1, j] * Muv1[:, j] + primR[1, j] * Muv2[:, j]
    end
    prim = mixture_conserve_prim(w, γ)
    tau = aap_hs_collision_time(prim, mi, ni, me, ne, Kn)
    prim = aap_hs_prim(prim, tau, mi, ni, me, ne, Kn)

    a = mixture_pdf_slope(prim, (wR .- wL) ./ (dxR + dxL), inK)
    Mu, Mv, Mxi, MuL, MuR = mixture_gauss_moments(prim, inK)
    Mau = mixture_moments_conserve_slope(a, Mu, Mv, Mxi, 1, 0, 0)
    aT = mixture_pdf_slope(prim, -prim[1] .* Mau, inK)

    Mt = zeros(2, 2)
    @. Mt[2, :] = tau * (1.0 - exp(-dt / tau)) # f0
    @. Mt[1, :] = dt - Mt[2, :] # M0

    Mt = zeros(5, axes(w, 2))
    for j in axes(Mt, 2)
        Mt[4, j] = tau[j] * (1.0 - exp(-dt / tau[j])) # f0
        Mt[5, j] = -tau[j] * dt * exp(-dt / tau[j]) + tau[j] * Mt[4, j]
        Mt[1, j] = dt - Mt[4, j] # M0
        Mt[2, j] = -tau[j] * Mt[1, j] + Mt[5, j]
        Mt[3, j] = dt^2 / 2.0 - tau[j] * Mt[1, j]
    end

    #--- calculate interface flux ---#
    # flux from M0
    Muv = mixture_moments_conserve(Mu, Mv, Mxi, 1, 0, 0)
    Mau = mixture_moments_conserve_slope(a, Mu, Mv, Mxi, 2, 0, 0)
    MauT = mixture_moments_conserve_slope(aT, Mu, Mv, Mxi, 1, 0, 0)
    for j in axes(fw, 2)
        @. fw[:, j] =
            Mt[1, j] * prim[1, j] * Muv[:, j] +
            Mt[2, j] * prim[1, j] * Mau[:, j] +
            Mt[3, j] * prim[1, j] * MauT[:, j]
    end

    # flux from f0
    H0 = mixture_maxwellian(u, v, prim)
    H1 = similar(H0)
    H2 = similar(H0)
    for j in axes(H0, 3)
        H1[:, :, j] = H0[:, :, j] .* prim[4, j]
        H2[:, :, j] .= H0[:, :, j] .* (prim[4, j]^2 + 1.0 / (2.0 * prim[5, j]))
    end

    for j in axes(fw, 2)
        fw[1, j] +=
            Mt[4, j] * sum(ω[:, :, j] .* u[:, :, j] .* h0[:, :, j]) -
            Mt[5, j] * sum(ω[:, :, j] .* u[:, :, j] .^ 2 .* sh0[:, :, j])
        fw[2, j] +=
            Mt[4, j] * sum(ω[:, :, j] .* u[:, :, j] .^ 2 .* h0[:, :, j]) -
            Mt[5, j] * sum(ω[:, :, j] .* u[:, :, j] .^ 3 .* sh0[:, :, j])
        fw[3, j] +=
            Mt[4, j] * sum(ω[:, :, j] .* v[:, :, j] .* u[:, :, j] .* h0[:, :, j]) -
            Mt[5, j] * sum(ω[:, :, j] .* v[:, :, j] .* u[:, :, j] .^ 2 .* sh0[:, :, j])
        fw[4, j] +=
            Mt[4, j] * sum(ω[:, :, j] .* u[:, :, j] .* h1[:, :, j]) -
            Mt[5, j] * sum(ω[:, :, j] .* u[:, :, j] .^ 2 .* sh1[:, :, j])
        fw[5, j] +=
            Mt[4, j] *
            0.5 *
            (
                sum(
                    ω[:, :, j] .* u[:, :, j] .* (u[:, :, j] .^ 2 .+ v[:, :, j] .^ 2) .*
                    h0[:, :, j],
                ) + sum(ω[:, :, j] .* u[:, :, j] .* h2[:, :, j])
            ) -
            Mt[5, j] *
            0.5 *
            (
                sum(
                    ω[:, :, j] .* u[:, :, j] .^ 2 .* (u[:, :, j] .^ 2 .+ v[:, :, j] .^ 2) .*
                    sh0[:, :, j],
                ) + sum(ω[:, :, j] .* u[:, :, j] .^ 2 .* sh2[:, :, j])
            )

        @. fh0[:, :, j] =
            Mt[1, j] * u[:, :, j] * H0[:, :, j] +
            Mt[2, j] *
            u[:, :, j]^2 *
            (
                a[1, j] * H0[:, :, j] +
                a[2, j] * u[:, :, j] * H0[:, :, j] +
                a[3, j] * v[:, :, j] * H0[:, :, j] +
                a[4, j] * u[:, :, j] * H1[:, :, j] +
                0.5 * a[5, j] * ((u[:, :, j]^2 + v[:, :, j]^2) * H0[:, :, j] + H2[:, :, j])
            ) +
            Mt[3, j] *
            u[:, :, j] *
            (
                aT[1, j] * H0[:, :, j] +
                aT[2, j] * u[:, :, j] * H0[:, :, j] +
                aT[3, j] * v[:, :, j] * H0[:, :, j] +
                aT[4, j] * u[:, :, j] * H1[:, :, j] +
                0.5 *
                aT[5, j] *
                ((u[:, :, j]^2 + v[:, :, j]^2) * H0[:, :, j] + H2[:, :, j])
            ) +
            Mt[4, j] * u[:, :, j] * h0[:, :, j] - Mt[5, j] * u[:, :, j]^2 * sh0[:, :, j]
        @. fh1[:, :, j] =
            Mt[1, j] * u[:, :, j] * H1[:, :, j] +
            Mt[2, j] *
            u[:, :, j]^2 *
            (
                a[1, j] * H1[:, :, j] +
                a[2, j] * u[:, :, j] * H1[:, :, j] +
                a[3, j] * v[:, :, j] * H1[:, :, j] +
                a[4, j] * u[:, :, j] * H2[:, :, j] +
                0.5 *
                a[5, j] *
                ((u[:, :, j]^2 + v[:, :, j]^2) * H1[:, :, j] + Mxi[3, j] * H0[:, :, j])
            ) +
            Mt[3, j] *
            u[:, :, j] *
            (
                aT[1, j] * H1[:, :, j] +
                aT[2, j] * u[:, :, j] * H1[:, :, j] +
                aT[3, j] * v[:, :, j] * H1[:, :, j] +
                aT[4, j] * u[:, :, j] * H2[:, :, j] +
                0.5 *
                aT[5, j] *
                ((u[:, :, j]^2 + v[:, :, j]^2) * H1[:, :, j] + Mxi[3, j] * H0[:, :, j])
            ) +
            Mt[4, j] * u[:, :, j] * h1[:, :, j] - Mt[5, j] * u[:, :, j]^2 * sh1[:, :, j]
        @. fh2[:, :, j] =
            Mt[1, j] * u[:, :, j] * H2[:, :, j] +
            Mt[2, j] *
            u[:, :, j]^2 *
            (
                a[1, j] * H2[:, :, j] +
                a[2, j] * u[:, :, j] * H2[:, :, j] +
                a[3, j] * v[:, :, j] * H2[:, :, j] +
                a[4, j] * u[:, :, j] * Mxi[3, j] * H0[:, :, j] +
                0.5 *
                a[5, j] *
                ((u[:, :, j]^2 + v[:, :, j]^2) * H2[:, :, j] + Mxi[4, j] * H0[:, :, j])
            ) +
            Mt[3, j] *
            u[:, :, j] *
            (
                aT[1, j] * H2[:, :, j] +
                aT[2, j] * u[:, :, j] * H2[:, :, j] +
                aT[3, j] * v[:, :, j] * H2[:, :, j] +
                aT[4, j] * u[:, :, j] * Mxi[3, j] * H0[:, :, j] +
                0.5 *
                aT[5, j] *
                ((u[:, :, j]^2 + v[:, :, j]^2) * H2[:, :, j] + Mxi[4, j] * H0[:, :, j])
            ) +
            Mt[4, j] * u[:, :, j] * h2[:, :, j] - Mt[5, j] * u[:, :, j]^2 * sh2[:, :, j]
    end

    @. fw *= len
    @. fh0 *= len
    @. fh1 *= len
    @. fh2 *= len

    return nothing

end


text:
"""
    StructArray{T,N,C,I} <: AbstractArray{T, N}

A type that stores an `N`-dimensional array of structures of type `T` as a structure of arrays.

- `getindex` and `setindex!` are overloaded to get/set values of type `T`.
- `getproperty` is overloaded to return individual field arrays.

# Fields

- `components`: a `NamedTuple` or `Tuple` of the arrays used by each field. These can be accessed by [`components(x)`](@ref).
"""
struct StructArray{T, N, C<:Tup, I} <: AbstractArray{T, N}
    components::C

    function StructArray{T, N, C}(c) where {T, N, C<:Tup}
        if length(c) > 0
            ax = axes(c[1])
            length(ax) == N || error("wrong number of dimensions")
            for i = 2:length(c)
                axes(c[i]) == ax || error("all field arrays must have same shape")
            end
        end
        new{T, N, C, index_type(C)}(c)
    end
end

# common type used for indexing
index_type(::Type{NamedTuple{names, types}}) where {names, types} = index_type(types)
index_type(::Type{Tuple{}}) = Int
function index_type(::Type{T}) where {T<:Tuple}
    S, U = tuple_type_head(T), tuple_type_tail(T)
    IndexStyle(S) isa IndexCartesian ? CartesianIndex{ndims(S)} : index_type(U)
end

index_type(::Type{StructArray{T, N, C, I}}) where {T, N, C, I} = I

array_types(::Type{StructArray{T, N, C, I}}) where {T, N, C, I} = array_types(C)
array_types(::Type{NamedTuple{names, types}}) where {names, types} = types
array_types(::Type{TT}) where {TT<:Tuple} = TT

"""
    StructArray{T}((components...)::Union{Tuple, NamedTuple})
    StructArray{T}(name1=component1, name2=component2, ...)

Construct a `StructArray` of element type `T` from the specified field arrays.

    StructArray((components...)::Union{Tuple, NamedTuple})
    StructArray(name1=component1, name2=component2, ...)

Construct a `StructArray` with a `Tuple` or `NamedTuple` element type from the
specified field arrays.

# Examples

```julia-repl
julia> StructArray{ComplexF64}(([1.0, 2.0], [3.0, 4.0]))
2-element StructArray(::Array{Float64,1}, ::Array{Float64,1}) with eltype Complex{Float64}:
 1.0 + 3.0im
 2.0 + 4.0im

julia> StructArray{ComplexF64}(re=[1.0, 2.0], im=[3.0, 4.0])
2-element StructArray(::Array{Float64,1}, ::Array{Float64,1}) with eltype Complex{Float64}:
 1.0 + 3.0im
 2.0 + 4.0im
```

Any `AbstractArray` can be used as a field array
```julia-repl
julia> StructArray{Complex{Int64}}(([1, 2], 3:4))
2-element StructArray(::Array{Int64,1}, ::UnitRange{Int64}) with eltype Complex{Int64}:
 1 + 3im
 2 + 4im
```

If no element type `T` is provided, a `Tuple` or `NamedTuple` is used:
```julia-repl
julia> StructArray((zeros(2,2), ones(2,2)))
2×2 StructArray(::Array{Float64,2}, ::Array{Float64,2}) with eltype Tuple{Float64,Float64}:
 (0.0, 1.0)  (0.0, 1.0)
 (0.0, 1.0)  (0.0, 1.0)

julia> StructArray(a=zeros(2,2), b=ones(2,2))
2×2 StructArray(::Array{Float64,2}, ::Array{Float64,2}) with eltype NamedTuple{(:a, :b),Tuple{Float64,Float64}}:
 (a = 0.0, b = 1.0)  (a = 0.0, b = 1.0)
 (a = 0.0, b = 1.0)  (a = 0.0, b = 1.0)
```
"""
StructArray(tup::Union{Tuple,NamedTuple})

function StructArray{T}(c::C) where {T, C<:Tup}
    cols = strip_params(staticschema(T))(c)
    N = isempty(cols) ? 1 : ndims(cols[1])
    StructArray{T, N, typeof(cols)}(cols)
end

StructArray(c::C) where {C<:NamedTuple} = StructArray{eltypes(C)}(c)
StructArray(c::Tuple; names = nothing) = _structarray(c, names)

StructArray{T}(; kwargs...) where {T} = StructArray{T}(values(kwargs))
StructArray(; kwargs...) = StructArray(values(kwargs))

_structarray(args::T, ::Nothing) where {T<:Tuple} = StructArray{eltypes(T)}(args)
_structarray(args::Tuple, names) = _structarray(args, Tuple(names))
_structarray(args::Tuple, ::Tuple) = _structarray(args, nothing)
_structarray(args::NTuple{N, Any}, names::NTuple{N, Symbol}) where {N} = StructArray(NamedTuple{names}(args))

const StructVector{T, C<:Tup, I} = StructArray{T, 1, C, I}
StructVector{T}(args...; kwargs...) where {T} = StructArray{T}(args...; kwargs...)
StructVector(args...; kwargs...) = StructArray(args...; kwargs...)

"""
    StructArray{T}(A::AbstractArray; dims, unwrap=FT->FT!=eltype(A))

Construct a `StructArray` from slices of `A` along `dims`.

The `unwrap` keyword argument is a function that determines whether to
recursively convert fields of type `FT` to `StructArray`s.

!!! compat "Julia 1.1"
     This function requires at least Julia 1.1.

```julia-repl
julia> X = [1.0 2.0; 3.0 4.0]
2×2 Array{Float64,2}:
 1.0  2.0
 3.0  4.0

julia> StructArray{Complex{Float64}}(X; dims=1)
2-element StructArray(view(::Array{Float64,2}, 1, :), view(::Array{Float64,2}, 2, :)) with eltype Complex{Float64}:
 1.0 + 3.0im
 2.0 + 4.0im

julia> StructArray{Complex{Float64}}(X; dims=2)
2-element StructArray(view(::Array{Float64,2}, :, 1), view(::Array{Float64,2}, :, 2)) with eltype Complex{Float64}:
 1.0 + 2.0im
 3.0 + 4.0im
```

By default, fields will be unwrapped until they match the element type of the array:
```
julia> StructArray{Tuple{Float64,Complex{Float64}}}(rand(3,2); dims=1)
2-element StructArray(view(::Array{Float64,2}, 1, :), StructArray(view(::Array{Float64,2}, 2, :), view(::Array{Float64,2}, 3, :))) with eltype Tuple{Float64,Complex{Float64}}:
 (0.004767505234193781, 0.27949621887414566 + 0.9039320635041561im)
 (0.41853472213051335, 0.5760165160827859 + 0.9782723869433818im)
```
"""
StructArray(A::AbstractArray; dims, unwrap)
function StructArray{T}(A::AbstractArray; dims, unwrap=FT->FT!=eltype(A)) where {T}
    slices = Iterators.Stateful(eachslice(A; dims=dims))
    buildfromslices(T, unwrap, slices)
end
function buildfromslices(::Type{T}, unwrap::F, slices) where {T,F}
    if unwrap(T)
        buildfromschema(T) do FT
            buildfromslices(FT, unwrap, slices)
        end
    else
        return popfirst!(slices)
    end
end

function Base.IndexStyle(::Type{S}) where {S<:StructArray}
    index_type(S) === Int ? IndexLinear() : IndexCartesian()
end

function _undef_array(::Type{T}, sz; unwrap::F = alwaysfalse) where {T, F}
    if unwrap(T)
        return StructArray{T}(undef, sz; unwrap = unwrap)
    else
        return Array{T}(undef, sz)
    end
end

function _similar(v::AbstractArray, ::Type{Z}; unwrap::F = alwaysfalse) where {Z, F}
    if unwrap(Z)
        return buildfromschema(typ -> _similar(v, typ; unwrap = unwrap), Z)
    else
        return similar(v, Z)
    end
end

"""
    StructArray{T}(undef, dims; unwrap=T->false)

Construct an uninitialized `StructArray` with element type `T`, with `Array`
field arrays.

The `unwrap` keyword argument is a function that determines whether to
recursively convert arrays of element type `T` to `StructArray`s.

# Examples

```julia-repl
julia> StructArray{ComplexF64}(undef, (2,3))
2×3 StructArray(::Array{Float64,2}, ::Array{Float64,2}) with eltype Complex{Float64}:
  2.3166e-314+2.38405e-314im  2.39849e-314+2.38405e-314im  2.41529e-314+2.38405e-314im
 2.31596e-314+2.41529e-314im  2.31596e-314+2.41529e-314im  2.31596e-314+NaN*im
```
"""
StructArray(::Base.UndefInitializer, sz::Dims)

function StructArray{T}(::Base.UndefInitializer, sz::Dims; unwrap::F = alwaysfalse) where {T, F}
    buildfromschema(typ -> _undef_array(typ, sz; unwrap = unwrap), T)
end
StructArray{T}(u::Base.UndefInitializer, d::Integer...; unwrap::F = alwaysfalse) where {T, F} = StructArray{T}(u, convert(Dims, d); unwrap = unwrap)

function similar_structarray(v::AbstractArray, ::Type{Z}; unwrap::F = alwaysfalse) where {Z, F}
    buildfromschema(typ -> _similar(v, typ; unwrap = unwrap), Z)
end

"""
    StructArray(A; unwrap = T->false)

Construct a `StructArray` from an existing multidimensional array or iterator
`A`.

The `unwrap` keyword argument is a function that determines whether to
recursively convert arrays of element type `T` to `StructArray`s.

# Examples

## Basic usage

```julia-repl
julia> A = rand(ComplexF32, 2,2)
2×2 Array{Complex{Float32},2}:
 0.694399+0.94999im  0.422804+0.891131im
 0.101001+0.33644im  0.632468+0.811319im

julia> StructArray(A)
2×2 StructArray(::Array{Float32,2}, ::Array{Float32,2}) with eltype Complex{Float32}:
 0.694399+0.94999im  0.422804+0.891131im
 0.101001+0.33644im  0.632468+0.811319im
```

## From an iterator

```julia-repl
julia> StructArray((1, Complex(i, j)) for i = 1:3, j = 2:4)
3×3 StructArray(::Array{Int64,2}, ::Array{Complex{Int64},2}) with eltype Tuple{Int64,Complex{Int64}}:
 (1, 1+2im)  (1, 1+3im)  (1, 1+4im)
 (1, 2+2im)  (1, 2+3im)  (1, 2+4im)
 (1, 3+2im)  (1, 3+3im)  (1, 3+4im)
```

## Recursive unwrapping

```julia-repl
julia> StructArray((1, Complex(i, j)) for i = 1:3, j = 2:4; unwrap = T -> !(T<:Real))
3×3 StructArray(::Array{Int64,2}, StructArray(::Array{Int64,2}, ::Array{Int64,2})) with eltype Tuple{Int64,Complex{Int64}}:
 (1, 1+2im)  (1, 1+3im)  (1, 1+4im)
 (1, 2+2im)  (1, 2+3im)  (1, 2+4im)
 (1, 3+2im)  (1, 3+3im)  (1, 3+4im)
```
"""
function StructArray(v; unwrap::F = alwaysfalse)::StructArray where {F}
    collect_structarray(v; initializer = StructArrayInitializer(unwrap))
end

function StructArray(v::AbstractArray{T}; unwrap::F = alwaysfalse) where {T, F}
    s = similar_structarray(v, T; unwrap = unwrap)
    for i in eachindex(v)
        @inbounds s[i] = v[i]
    end
    s
end
StructArray(s::StructArray) = copy(s)

Base.convert(::Type{StructArray}, v::AbstractArray) = StructArray(v)
Base.convert(::Type{StructArray}, v::StructArray) = v

Base.convert(::Type{StructVector}, v::AbstractVector) = StructVector(v)
Base.convert(::Type{StructVector}, v::StructVector) = v

function Base.similar(::Type{<:StructArray{T, <:Any, C}}, sz::Dims) where {T, C}
    buildfromschema(typ -> similar(typ, sz), T, C)
end

Base.similar(s::StructArray, sz::Base.DimOrInd...) = similar(s, Base.to_shape(sz))
Base.similar(s::StructArray) = similar(s, Base.to_shape(axes(s)))
function Base.similar(s::StructArray{T}, sz::Tuple) where {T}
    StructArray{T}(map(typ -> similar(typ, sz), components(s)))
end

@deprecate fieldarrays(x) StructArrays.components(x)

"""
    components(s::StructArray)

Return the field arrays corresponding to the various entry of the struct as a `NamedTuple`, or a `Tuple` if the struct has no names.

# Examples

```julia-repl
julia> s = StructArray(rand(ComplexF64, 4));

julia> components(s)
(re = [0.396526, 0.486036, 0.459595, 0.0323561], im = [0.147702, 0.81043, 0.00993469, 0.487091])
```
"""
components(s::StructArray) = getfield(s, :components)

component(s::StructArray, key) = getfield(components(s), key)

Base.getproperty(s::StructArray, key::Symbol) = component(s, key)
Base.getproperty(s::StructArray, key::Int) = component(s, key)
Base.propertynames(s::StructArray) = propertynames(components(s))

staticschema(::Type{StructArray{T, N, C, I}}) where {T, N, C, I} = staticschema(C)
createinstance(::Type{<:StructArray{T}}, args...) where {T} = StructArray{T}(args)

Base.size(s::StructArray) = size(components(s)[1])
Base.size(s::StructArray{<:Any, <:Any, <:EmptyTup}) = (0,)
Base.axes(s::StructArray) = axes(components(s)[1])
Base.axes(s::StructArray{<:Any, <:Any, <:EmptyTup}) = (1:0,)

"""
    StructArrays.get_ith(cols::Union{Tuple,NamedTuple}, I...)

Form a `Tuple` of the `I`th index of each element of `cols`, i.e. is equivalent
to
```julia
map(c -> c[I...], Tuple(cols))
```
"""
@inline get_ith(cols::NamedTuple, I...) = get_ith(Tuple(cols), I...)
@inline function get_ith(cols::Tuple, I...)
    @inbounds r = first(cols)[I...]
    return (r, get_ith(Base.tail(cols), I...)...)
end
@inline get_ith(::Tuple{}, I...) = ()

Base.@propagate_inbounds function Base.getindex(x::StructArray{T, <:Any, <:Any, CartesianIndex{N}}, I::Vararg{Int, N}) where {T, N}
    cols = components(x)
    @boundscheck checkbounds(x, I...)
    return createinstance(T, get_ith(cols, I...)...)
end

Base.@propagate_inbounds function Base.getindex(x::StructArray{T, <:Any, <:Any, Int}, I::Int) where {T}
    cols = components(x)
    @boundscheck checkbounds(x, I)
    return createinstance(T, get_ith(cols, I)...)
end

@inline function Base.view(s::StructArray{T, N, C}, I...) where {T, N, C}
    StructArray{T}(map(v -> view(v, I...), components(s)))
end

Base.@propagate_inbounds function Base.setindex!(s::StructArray{<:Any, <:Any, <:Any, CartesianIndex{N}}, vals, I::Vararg{Int, N}) where {N}
    @boundscheck checkbounds(s, I...)
    foreachfield((col, val) -> (@inbounds col[I...] = val), s, vals)
    s
end

Base.@propagate_inbounds function Base.setindex!(s::StructArray{<:Any, <:Any, <:Any, Int}, vals, I::Int)
    @boundscheck checkbounds(s, I)
    foreachfield((col, val) -> (@inbounds col[I] = val), s, vals)
    s
end

function Base.push!(s::StructVector, vals)
    foreachfield(push!, s, vals)
    return s
end

function Base.append!(s::StructVector, vals::StructVector)
    foreachfield(append!, s, vals)
    return s
end

Base.copyto!(I::StructArray, J::StructArray) = (foreachfield(copyto!, I, J); I)

function Base.copyto!(I::StructArray, doffs::Integer, J::StructArray, soffs::Integer, n::Integer)
    foreachfield((dest, src) -> copyto!(dest, doffs, src, soffs, n), I, J)
    return I
end

function Base.resize!(s::StructArray, i::Integer)
    for a in components(s)
        resize!(a, i)
    end
    return s
end

function Base.empty!(s::StructArray)
    foreachfield(empty!, s)
end

for op in [:cat, :hcat, :vcat]
    @eval begin
        function Base.$op(args::StructArray...; kwargs...)
            f = key -> $op((getproperty(t, key) for t in args)...; kwargs...)
            T = mapreduce(eltype, promote_type, args)
            StructArray{T}(map(f, propertynames(args[1])))
        end
    end
end

Base.copy(s::StructArray{T}) where {T} = StructArray{T}(map(copy, components(s)))

function Base.reshape(s::StructArray{T}, d::Dims) where {T}
    StructArray{T}(map(x -> reshape(x, d), components(s)))
end

function showfields(io::IO, fields::NTuple{N, Any}) where N
    print(io, "(")
    for (i, field) in enumerate(fields)
        Base.showarg(io, fields[i], false)
        i < N && print(io, ", ")
    end
    print(io, ")")
end

function Base.showarg(io::IO, s::StructArray{T}, toplevel) where T
    print(io, "StructArray")
    showfields(io, Tuple(components(s)))
    toplevel && print(io, " with eltype ", T)
end

# broadcast
import Base.Broadcast: BroadcastStyle, ArrayStyle, AbstractArrayStyle, Broadcasted, DefaultArrayStyle

struct StructArrayStyle{Style} <: AbstractArrayStyle{Any} end

@inline combine_style_types(::Type{A}, args...) where A<:AbstractArray =
    combine_style_types(BroadcastStyle(A), args...)
@inline combine_style_types(s::BroadcastStyle, ::Type{A}, args...) where A<:AbstractArray =
    combine_style_types(Broadcast.result_style(s, BroadcastStyle(A)), args...)
combine_style_types(s::BroadcastStyle) = s

Base.@pure cst(::Type{SA}) where SA = combine_style_types(array_types(SA).parameters...)

BroadcastStyle(::Type{SA}) where SA<:StructArray = StructArrayStyle{typeof(cst(SA))}()

Base.similar(bc::Broadcasted{StructArrayStyle{S}}, ::Type{ElType}) where {S<:DefaultArrayStyle,N,ElType} =
    isstructtype(ElType) ? similar(StructArray{ElType}, axes(bc)) : similar(Array{ElType}, axes(bc))

ext: jl


text:
using SciMLTutorials
SciMLTutorials.tutorial_footer(WEAVE_ARGS[:folder],WEAVE_ARGS[:file])

max_stars_repo_name: SciML/SciMLTutorialsOutput

max_issues_repo_name: SciML/SciMLTutorialsOutput

max_forks_repo_name: SciML/SciMLTutorialsOutput


text:
import IntertemporalChoiceHeuristics:
    load,
    fit,
    models,
    make_closures,
    LogLoss,
    invlogit,
    stderrs

inputs, weights = load()

loss = LogLoss()
λ = 0.0

for model in models
    res = fit(model, inputs, weights)

    nll, nll_gr! = make_closures(model, loss, λ, inputs, weights)

    mle = Optim.minimizer(res)

    se = stderrs(model, inputs, weights)

    K = length(mle)

    n_points = 256
    for j in 1:K
        parameters = copy(mle)
        grid = linspace(
            mle[j] - 6 * se[j],
            mle[j] + 6 * se[j],
            n_points,
        )
        for i in 1:n_points
            parameters[j] = grid[i]
            @printf(
                "%s\t%d\t%s\t%s\t%s\t%s\n",
                model,
                j,
                grid[i],
                nll(parameters),
                mle[j],
                se[j],
            )
        end
    end
end


text:
using ClimateGNSS
using Test

@testset "ClimateGNSS.jl" begin
    # Write your own tests here.
end


text:
# this is testing retest(load=true), so not hijack stuff, but we already tested
# load=true for FakePackage, here we want to test another setting with two
# included modules

module HijackLoadTrueTests
using ReTest
@testset "HijackLoadTrueTests" begin end
end

module HijackLoadTrueTests2
using ReTest
@testset "HijackLoadTrueTests2" begin end
end

max_issues_repo_name: rfourquet/InlineTest

max_forks_repo_name: rfourquet/InlineTest


text:
# Burgers
import SDELab2
include("common.jl")
########################## problem parameters
T=1
diffusion=0.1
sig=sqrt(2*diffusion)

quad0=gauss_quadrature.delta_measure(0.)
#quad0=gauss_quadrature.Gauss_Hermite([0.,1e-8],4)
#########################################
function drift_ff(u)
  x=u[1]; mf=u[2]
  if x>mf
    return 0.
  elseif x==mf
    return 0.5
  else
    return 1.
  end
end
#
############################
function drift_fr(u,ell)
  x=u[1]; mf=u[2]
  return 0.5*erfc((x-mf)/ell)
end
#########################################
function diff_gg(u,sig)
  return sig
end
#########################################
function mf_r(x)
  return [x]
end
#########################################
fcn=SDELab2.set_fcn_mf(u->drift_fr(u,ell),
          u->diff_gg(u,sig),
          x->mf_r(x))
#########################################
function exact_cdf2(x,t,sig)
  sig_sq=sig^2
  scale=sig*sqrt(2*t)
 return 1- ( erfc(-x/scale)./
 (
 erfc(-x/scale)
 + exp((t-2*x)/(2*sig_sq)) .* (2-erfc((t-x)/scale))
 )    )
end
#
function exact_cdf3(x,t,sig)
  sig_sq=sig^2
  scale=sig*sqrt(2*t)
  T1=erfc(-x/scale)
  T2=exp((t-2*x)/(2*sig_sq))
  T3=erfc((t-x)/scale)
 return T2*(2-T3)./(T1+T2.*(2-T3))
end
######################################### method parameters
solver="MFEM0e"
dt=1e-2
params=SDELab2.set_opt(dt,solver)
params["const"]=1e0       # for no_Gauss_points
params["length_scale"]=1 # lam=1/length^2
params["cutoff"]=1  #
params["skiphalf"]=true
params["periodic"]=false
params[ "no_partition"]=20
#########################################
###################### computing reference value
print("Computing a reference value...")
phi=x->[x^2;x]
tic()
params["extrap"]=false
dtpts=0.1./(2.0).^(1:8 )
params["MaxStepSize"]=dtpts[end]
ell=sqrt(dtpts[end])

#ell=0.5#(dtpts[end])
ell=0.1
#ell=0.01
#ell=0.001
quad0=gauss_quadrature.delta_measure(0.)#-dtpts[end]/2)
fcn=SDELab2.set_fcn_mf(u->drift_fr(u,ell),
          u->diff_gg(u,sig),
          x->mf_r(x))
#quad0=gauss_quadrature.delta_measure(sqrt(dtpts[end]))
exact,cpu_time,stats,q1=get_soln(fcn,quad0,params,phi,true)
println(stats["no_pts"])
println("exact=", exact)
#########################################
#initial_final_cdf_pdf(q0,q1,[100,40],[0.08,0.01])
#

cum1=cumsum(q1.w)
x1=(q1.X[2:end]+q1.X[1:end-1])/2
#x1=q1.X
#x1=q1.X[1:end-1]
cum_exact=1-exact_cdf2(q1.X,1.,sig)
q2=gauss_quadrature.QUAD(length(x1),x1,q1.w)
println("l1 ",gauss_quadrature.l1cdf(q2,z->(1-exact_cdf2(z,1.,sig))))
using PyPlot
figure(2)
clf()
plot(q1.X,cum_exact,"b")
plot(x1,cum1[2:end],"g")#
axis([-1.5,1.5,0,1])
xlabel("X")
ylabel("cdf")
gcf()[:set_size_inches](width_inches,height_inches)
tight_layout()
  savefig(utf8("cdf_cmp.pdf"))
  clf()
  plot(q1.X,abs(cum_exact-cum1),"b")
  #plot(q1.X,cum1[2:end],"g")#
    axis([-1.5,1.5,0,0.04])
  xlabel("X")
  ylabel("error")
  gcf()[:set_size_inches](width_inches,height_inches)
  tight_layout()
  savefig(utf8("cdf_cmp_del.pdf"))
#########################################
dtpts=dtpts[1:end-1] # omit last value, used for reference
errors=zeros(length(dtpts),2); errors_no=deepcopy(errors); errors_nd=deepcopy(errors)
cpu_times=zeros(length(dtpts)); cpu_times_no=deepcopy(cpu_times);cpu_times_nd=deepcopy(cpu_times)
params["MaxStepSize"]=dtpts[1]
######################
params["Solver"]="MFEM0e"
dtpts_no=dtpts
for i=1:(length(dtpts))
    params["MaxStepSize"]=dtpts_no[i]
    qoi,cpu_time,stats,q=get_soln(fcn,quad0,params,phi,false)
    errors_no[i,:],cpu_times_no[i]=helper(qoi, exact, cpu_time, "normal")
end
########################
params["Solver"]="MFEM0e"
for i=1:(length(dtpts))
    params["MaxStepSize"]=dtpts[i]
    qoi,cpu_time,stats,q=get_soln(fcn,quad0,params,phi,true)
    errors[i,:],cpu_times[i]=helper(qoi, exact, cpu_time, "normal")
end
################
#########################################
figure(1);
# first functional
i=1

time_step_err(dtpts,errors[:,i],
  dtpts_no,errors_no[:,i],
  dtpts,errors_nd[:,i],[2,2],utf8("mom1.pdf"))
err_cpu(errors[:,i],cpu_times[:],
  errors_no[:,i],cpu_times_no[:],
  errors_nd[:,i],cpu_times_nd[:],[1,1], utf8("mom1_.pdf"))
# second functional
i=2
time_step_err(dtpts,errors[:,i],
      dtpts,errors_no[:,i],
      dtpts,errors_nd[:,i],[1,1],utf8("mom2.pdf"))
err_cpu(errors[:,i], cpu_times[:],
      errors_no[:,i],cpu_times_no[:],
      errors_nd[:,i],cpu_times_nd[:],[1,1], utf8("mom2_.pdf"))
#########################################
println("finished preparing plots")


text:
include("definitions.jl")
include("updateFunctions.jl")
include("utility.jl")

function sparseSCVB0_lda(documents::Array{SparseDocument,1},
                           numWords::Int64,
                           numDocuments::Int64,
                           numTopics::Int64,
                           numDocumentIterations::Int,
                           burnInPerDoc::Int,
                           minibatchSize::Int,
                           Alpha::Float64,
                           Beta::Float64,
                           dict,
                           tau::Float64,
                           kappa::Float64,
                           tau2::Float64,
                           kappa2::Float64,
                           scale::Float64,
                           sampleSize::Int,
                           time_limit::Float64,
                           miniEpoch::Int,
                           numDocuments_miniEpoch::Int64,
						   sparsity::Float64)

# sparse stochastic collapsed variational Bayes (sparseSCVB0) for LDA
    Alpha = ones(1,numTopics).*Alpha;
    Beta = ones(numWords,1).*Beta;

    # initialization of sparseSCVB0 specific model parameters
    #documentTopicCounts, wordTopicCounts, topicCounts, topicMem, sparseCounts = initialize(documents, numWords,numTopics,numDocuments);
    sparseVBparams=initialize(documents,Alpha, Beta, numTopics,numWords,numDocumentIterations,numDocuments);
    miniBatchesPerCorpus = numDocuments ./ minibatchSize;
    stepSize = scale ./ (tau^kappa); # step size for global expected counts
    #stepSize2 = 0.0; # initializing step size for local expected counts to make it soft local scope

    wordTopicCounts_hat = zeros(numWords,numTopics);
    topicCounts_hat = zeros(1,numTopics);
    iter = 0;

    saved_topics = Array{Array{Float64,2}}(0)
    saved_Totaltopics = Array{Array{Float64,2}}(0)
    saved_time = Float64[]
    saved_iters = Int64[]
    #saved_docTopics = Array{Array{Float64,2}}(0)

    prev_doc = 0; # to track miniEpoch docs
    doc_counter = 0;
    # main sparse stochastic collapsed variational Bayes (sparseSCVB0) loop
    while (sparseVBparams.wallClockTime < time_limit) && (doc_counter<numDocuments)
        prev_doc,documentTopicCounts,topicMem,sparseCounts = initialize_miniEpoch(documents,numTopics,numWords,numDocuments_miniEpoch,prev_doc);
        for epoch = 1:miniEpoch
            for epochDoc = 1:numDocuments_miniEpoch
                iter += 1;
                tic();
                #docInd_mini = mod(iter - 1, numDocuments_miniEpoch) + 1;
                #docInd = doc_counter + docInd_mini;
                docInd = doc_counter + epochDoc;
                docLength = documents[docInd].docLength;
                if docLength == 0
                    continue;
                end
                    # burn-in pass for a document
                sparseVBparams,documentTopicCounts,sparseCounts,topicMem = burnInUpdate(sparseVBparams,documentTopicCounts,documents,burnInPerDoc,docInd,epochDoc,docLength,tau2,kappa2,sparseCounts,topicMem,sampleSize,sparsity);

                    # main loop update for a word
                sparseVBparams,documentTopicCounts,wordTopicCounts_hat,topicCounts_hat,sparseCounts,topicMem = mainLoopUpdate(sparseVBparams,documentTopicCounts,documents,burnInPerDoc,docInd,epochDoc,docLength,tau2,kappa2,sparseCounts,
                                                                topicMem,wordTopicCounts_hat,topicCounts_hat,sampleSize);

                # mini-batch update of expected global counts
                if mod(iter, minibatchSize) == 0
                    # compute effective stepsize to account for minibatches
                    stepSize = 1 - (1-stepSize)^minibatchSize
                    # minibatch update of expected global counts
                    #wordTopicCounts,topicCounts,wordTopicCounts_hat,topicCounts_hat=minibatchUpdate(stepSize,miniBatchesPerCorpus,wordTopicCounts,topicCounts,wordTopicCounts_hat,topicCounts_hat);
                    topicCounts = sparseVBparams.topicCounts;
                    wordTopicCounts = sparseVBparams.wordTopicCounts;
                    for topic = 1:numTopics
                        for word = 1:numWords
                            wordTopicCounts[word,topic] = (1 - stepSize) .*  wordTopicCounts[word,topic] + stepSize .* miniBatchesPerCorpus .* wordTopicCounts_hat[word,topic];
                            wordTopicCounts_hat[word,topic] = wordTopicCounts_hat[word,topic]*0.0;
                        end
                        topicCounts[topic] = (1 - stepSize) .* topicCounts[topic] + stepSize .* miniBatchesPerCorpus .* topicCounts_hat[topic];
                        topicCounts_hat[topic] = topicCounts_hat[topic]*0.0;
                    end
                    stepSize = scale ./ (iter + tau)^kappa;
                end
                sparseVBparams.wallClockTime += toq();

                if iter == 1
                    push!(saved_topics,copy(sparseVBparams.wordTopicCounts))
                    push!(saved_Totaltopics,copy(sparseVBparams.topicCounts))
                    push!(saved_time,sparseVBparams.wallClockTime)
                    push!(saved_iters,iter)
                else
                    if mod(iter,200)==0
                        push!(saved_topics,copy(sparseVBparams.wordTopicCounts))
                        push!(saved_Totaltopics,copy(sparseVBparams.topicCounts))
                        push!(saved_time,sparseVBparams.wallClockTime)
                        push!(saved_iters,iter);
                    end
                end
            end
        end
        doc_counter = doc_counter + numDocuments_miniEpoch;
        #push!(saved_docTopics,copy(documentTopicCounts))
    end
    return sparseVBparams,saved_topics,saved_Totaltopics,saved_time,saved_iters,iter;
end

function initialize(documents,Alpha, Beta, numTopics,numWords,numDocumentIterations,numDocuments);
    # initialize sparseSCVB0 model parameters
    sparseVBparams = sparseSCVBO_params_type(Alpha, Beta, numTopics,numWords,numDocumentIterations,numDocuments);
    totalWordsInCorpus = 0;
    for i = 1:numDocuments
        docLength = documents[i].docLength;
        totalWordsInCorpus = totalWordsInCorpus + docLength;
    end
    # initialize expected global counts
    sparseVBparams.topicCounts.*= totalWordsInCorpus ./ numTopics;
    sparseVBparams.wordTopicCounts.*= totalWordsInCorpus ./ (numWords .* numTopics);

    # initializing alias samples from alias table for each word
    sparseVBparams.cachedSamples,sparseVBparams.cached_qw,sparseVBparams.cached_Qw = aliasMethod(sparseVBparams.numWords,sparseVBparams.Alpha,
                                       sparseVBparams.wordTopicCounts,sparseVBparams.topicCounts,sparseVBparams.Beta,
                                       sparseVBparams.BetaSum,sparseVBparams.numTopics,sparseVBparams.numCachedSamples,
                                       sparseVBparams.cachedSamples,sparseVBparams.cached_qw,sparseVBparams.cached_Qw);
    return sparseVBparams;
end

function initialize_miniEpoch(documents,numTopics,numWords,numDocuments_miniEpoch,prev_doc);
    # initialization of topic storage
    topicMem = Array{TopicArray}(0);
    # initialization for keeping track on non-zero sparse counts
    sparseCounts = Array{SparseCounts}(0);
    # initialize sparseSCVB0 miniEpoch parameters
    documentTopicCounts = zeros(numDocuments_miniEpoch,numTopics);
    j=0;
    for i = (prev_doc+1):(prev_doc+numDocuments_miniEpoch)
        j=j+1;
        docLength = documents[i].docLength;
        tokenLength = length(documents[i].tokenIndex); # storing topic for clumped documents
        Assigned = Array{Int64}(tokenLength); # array to store topic assignment for each token of clumped document
        nonzeroDocIndex = Array{Int64}(0); # array to keep track of non-zero document-topic counts
        documentTopicTmp = rand(1,numTopics);
        sum_documentTopicTmp = sum(documentTopicTmp);
        for topic = 1:numTopics
            documentTopicCounts[j,topic] = (documentTopicTmp[topic]./sum_documentTopicTmp).*docLength;
            push!(nonzeroDocIndex,topic);
        end
        Assigned = rand(1:numTopics,tokenLength);
        push!(topicMem,TopicArray(Assigned));
        push!(sparseCounts,SparseCounts(nonzeroDocIndex));
    end
    prev_doc = prev_doc+numDocuments_miniEpoch;
    return prev_doc,documentTopicCounts,topicMem,sparseCounts;
end

function aliasMethod(numWords,Alpha,wordTopicCounts,topicCounts,Beta,BetaSum,numTopics,numCachedSamples,cachedSamples,cached_qw,cached_Qw)
    qw = Array{Float64}(numTopics); # unnormalized discrete distribution of dense bucket
    for word = 1:numWords
        cached_Qw[word] = 0.0;
        for topic = 1:numTopics
            qw[topic] = Alpha[topic] .* (wordTopicCounts[word,topic] + Beta[word]) ./ (topicCounts[topic] + BetaSum);
            cached_qw[word,topic] = qw[topic];
            cached_Qw[word] += qw[topic];
        end
        aliasTable = generateAlias(qw,numTopics);
        for sample = 1:numCachedSamples
            cachedSamples[word,sample] = sampleAlias(aliasTable,numTopics);
        end
    end
    return cachedSamples,cached_qw,cached_Qw;
end


text:
group = addgroup!(SUITE, "kernel")

dummy_kernel() = nothing
group["launch"] = @benchmarkable @cuda dummy_kernel()

wanted_threads = 10000
function configurator(kernel)
    config = launch_configuration(kernel.fun)

    threads = Base.min(wanted_threads, config.threads)
    blocks = cld(wanted_threads, threads)

    return (threads=threads, blocks=blocks)
end
group["occupancy"] = @benchmarkable @cuda config=$configurator dummy_kernel()

src = CUDA.rand(Float32, 512, 1000)
dest = similar(src)
function indexing_kernel(dest, src)
    i = (blockIdx().x-1) * blockDim().x + threadIdx().x
    @inbounds dest[i] = src[i]
    return
end
group["indexing"] = @benchmarkable CUDA.@sync @cuda threads=size(src,1) blocks=size(src,2) $indexing_kernel($dest, $src)

function checked_indexing_kernel(dest, src)
    i = (blockIdx().x-1) * blockDim().x + threadIdx().x
    dest[i] = src[i]
    return
end
group["indexing_checked"] = @benchmarkable CUDA.@sync @cuda threads=size(src,1) blocks=size(src,2) $checked_indexing_kernel($dest, $src)


text:
using LightGraphs.SimpleGraphs

import LightGraphs.SimpleGraphs: fadj, badj, adj
import LightGraphs.edgetype, LightGraphs.has_edge
using Statistics: mean

struct DummySimpleGraph <: AbstractSimpleGraph{Int} end
struct DummySimpleEdge <: AbstractSimpleEdge{Int} end
DummySimpleEdge(x...) = DummySimpleEdge()
LightGraphs.edgetype(g::DummySimpleGraph) = DummySimpleEdge
has_edge(::DummySimpleGraph, ::DummySimpleEdge) = true

# function to check if the invariants for SimpleGraph and SimpleDiGraph holds
function isvalid_simplegraph(g::SimpleGraph{T}) where {T <: Integer}
    nf = length(g.fadjlist)
    n = T(nf)
    # checks it the adjacency lists are sorted, free of duplicates and in the correct range
    for u in one(T):n
        listu = g.fadjlist[u]
        isempty(listu) && continue
        issorted(listu) || return false
        allunique(listu) || return false
        issubset(extrema(listu), one(T):n) || return false
    end
    # checks if the edge count is correct
    edge_count = 0
    for u in one(T):n
        listu = g.fadjlist[u]
        for v in listu
            v > u && break
            edge_count += 1
        end
    end
    g.ne == edge_count || return false
    #  checks for backwards edge
    for u in one(T):n
        listu = g.fadjlist[u]
        for v in listu
            LightGraphs.insorted(u, g.fadjlist[v]) || return false
        end
    end
    return true
end

function isvalid_simplegraph(g::SimpleDiGraph{T}) where {T <: Integer}
    nf = length(g.fadjlist)
    nb = length(g.badjlist)
    nf == nb || return false
    n = T(nf)

    for u in one(T):n
        for listu in (g.fadjlist[u], g.badjlist[u])
            listu = g.fadjlist[u]
            isempty(listu) && continue
            issorted(listu) || return false
            allunique(listu) || return false
            issubset(extrema(listu), one(T):n) || return false
        end
    end
    # checks if the edge count is correct
    edge_count = 0
    for u in one(T):n
        edge_count += length(g.fadjlist[u])
    end
    g.ne == edge_count || return false
    edge_count = 0
    for u in one(T):n
        edge_count += length(g.badjlist[u])
    end
    g.ne == edge_count || return false
    #  checks for backwards edge
    for u in one(T):n
        listu = g.fadjlist[u]
        for v in listu
            LightGraphs.insorted(u, g.badjlist[v]) || return false
        end
    end
    return true
end

const simplegraphtestdir = dirname(@__FILE__)

const simple_tests = [
    "simplegraphs",
    "simpleedge",
    "simpleedgeiter",
    "generators/randgraphs",
    "generators/staticgraphs",
    "generators/smallgraphs",
    "generators/euclideangraphs"
]

@testset "LightGraphs.SimpleGraphs" begin
    for t in simple_tests
        tp = joinpath(simplegraphtestdir, "$(t).jl")
        include(tp)
    end
end


text:
module TestRectifier

println("\nRectifier: Demonstrating the ability to simulate models with state events")

using Modia
using Modia.Electric
using ModiaMath.plot

@model Rectifier begin
  R=Resistor(R=1)
  r=Resistor(R=0.01)
  C=Capacitor(C=1,start=0)
  D=IdealDiode()
  V=SineVoltage(V=5,freqHz=1.5, offset=0, startTime=0)
@equations begin
  connect(V.p, D.p)
  connect(D.n, R.p)
  connect(R.n, r.p)
  connect(r.n, V.n)
  connect(C.n, R.n)
  connect(C.p, R.p)
  end
end

result = checkSimulation(Rectifier, 2, "C.v", 0.4773911315322196)

plot(result, ("C.v", "V.v"), heading="Rectifier", figure=12)

end


text:
ENV["OMP_NUM_THREADS"] = 4

export PardisoShift

struct PardisoShift
    point::Float64
    verbose::Bool
end

"""
    PardisoShift(p::Number; verbose = false)

When used as the `point = PardisoShift(p)` keyword argument in `diagonalize`, it forces use
of the MKL Pardiso library for the shift-and-invert Lanczos method (it should be installed).
`verbose = true` makes Pardiso verbose.

# Example
```jldoctest
julia> using Pardiso, Arpack; h = rand(10,10); h = h' + h;

julia> d = diagonalizer(h, Arpack_IRAM, point = PardisoShift(0.0))
Diagonaliser{Arpack_IRAM{Float64}} for (10, 10) Hermitian matrix around point 0.0
```
"""
PardisoShift(p::Number; verbose = false) = PardisoShift(p, verbose)

getpoint(p::PardisoShift) = p.point

function linearmap(h::AbstractArray{Tv}, p::PardisoShift) where {Tv}
    ps = pardisosolver(;verbose = p.verbose)
    _set_matrixtype!(ps, Tv)
    hp = ensuresparse(Pardiso.get_matrix(ps, h - Tv(p.point) * I, :N))
    preparesolver!(ps, hp)
    lmap = let ps = ps, hp = hp
        LinearMap{Tv}((x, b) -> Pardiso.pardiso(ps, x, hp, b), size(h)...,
                      ismutating = true, ishermitian = true)
    end
    return lmap, ps
end

function pardisosolver(; verbose = false)
    ps = Pardiso.MKLPardisoSolver()
    finalizer(releasePardiso, ps)
    verbose && Pardiso.set_msglvl!(ps, Pardiso.MESSAGE_LEVEL_ON)
    Pardiso.pardisoinit(ps)
    Pardiso.set_iparm!(ps, 12, 2) # Pardiso expects CSR, Julia uses CSC
    return ps
end

_set_matrixtype!(ps, ::Type{<:AbstractFloat}) = Pardiso.set_matrixtype!(ps, Pardiso.REAL_SYM_INDEF)
_set_matrixtype!(ps, ::Type{<:Complex}) = Pardiso.set_matrixtype!(ps, Pardiso.COMPLEX_HERM_INDEF)

ensuresparse(a::SparseMatrixCSC) = a
ensuresparse(a::AbstractMatrix) = sparse(a)

function preparesolver!(ps, hp::AbstractArray{Tv}) where {Tv}
    b = Vector{Tv}(undef, size(hp, 1))
    Pardiso.set_phase!(ps, Pardiso.ANALYSIS_NUM_FACT)
    Pardiso.pardiso(ps, hp, b)
    Pardiso.set_phase!(ps, Pardiso.SOLVE_ITERATIVE_REFINE)
    return ps
end

function releasePardiso(ps)
    original_phase = Pardiso.get_phase(ps)
    Pardiso.set_phase!(ps, Pardiso.RELEASE_ALL);
    Pardiso.pardiso(ps)
    Pardiso.set_phase!(ps, original_phase)
    return ps
end


text:
module PostgreSQLLoader

using Octo.Repo: ExecuteResult

# https://github.com/invenia/LibPQ.jl
using LibPQ # v0.9.1
using .LibPQ.Tables

const current = Dict{Symbol, Any}(
    :conn => nothing,
)

current_conn() = current[:conn]

# db_connect
function db_connect(; kwargs...)
    if !isempty(kwargs)
        str = join(map(kv->join(kv, '='), collect(kwargs)), ' ')
        conn = LibPQ.Connection(str)
        current[:conn] = conn
    end
end

# db_disconnect
function db_disconnect()
    conn = current_conn()
    close(conn)
    current[:conn] = nothing
end

# query
function query(sql::String)
    conn = current_conn()
    stmt = LibPQ.prepare(conn, sql)
    result = LibPQ.execute(stmt)
    df = Tables.rowtable(result)
    LibPQ.close(result)
    df
end

function query(prepared::String, vals::Vector)
    conn = current_conn()
    stmt = LibPQ.prepare(conn, prepared)
    result = LibPQ.execute(stmt, vals)
    df = Tables.rowtable(result)
    LibPQ.close(result)
    df
end

# execute
function execute(sql::String)::ExecuteResult
    conn = current_conn()
    LibPQ.execute(conn, sql)
    ExecuteResult()
end

function execute(prepared::String, vals::Vector)::ExecuteResult
    conn = current_conn()
    stmt = LibPQ.prepare(conn, prepared)
    LibPQ.execute(stmt, vals)
    ExecuteResult()
end

function execute(prepared::String, nts::Vector{<:NamedTuple})::ExecuteResult
    conn = current_conn()
    stmt = LibPQ.prepare(conn, prepared)
    for tup in nts
        LibPQ.execute(stmt, collect(tup))
    end
    ExecuteResult()
end

end # module Octo.Backends.PostgreSQLLoader


text:
using Documenter, JOLI

makedocs(
    sitename = "JOLI Reference",
    modules = [JOLI]
)

deploydocs(
    repo = "github.com/slimgroup/JOLI.jl.git"
)


text:
function select_candidate(
    candidates::Vector{Candidate},
    trials::Vector{Trial},
    weights::Vector,
    scoring::Bool,
)
    if scoring
        _, i = findmax(score_candidates(candidates, trials, weights))
        i, candidates[i]
    else
        i = rand(1:length(candidates))
        i, candidates[i]
    end
end


text:
using LightGraphs

import Base: eltype, iterate, length

export MultipleEdgeIter

struct MultipleEdgeIter{G<:AbstractMultigraph} <:AbstractEdgeIter
    g::G
end

function iterate(eit::MultipleEdgeIter{G}, state=(one(eltype(eit.g)), one(eltype(eit.g)))) where {G <: AbstractMultigraph}
    g = eit.g
    n = nv(g)
    vs = vertices(g)
    sort!(vs)
    u, i = state

    @inbounds while u <= n
        list_u = outneighbors(g, vs[u])
        if i > length(list_u)
            u += 1
            i = one(u)
            continue
        end
        if is_directed(g)
            e = MultipleEdge(vs[u], list_u[i], mul(g, vs[u], list_u[i]))
            state = (u, i + 1)
            return e, state
        else
            if list_u[i] >= vs[u]
                e = MultipleEdge(vs[u], list_u[i], mul(g, vs[u], list_u[i]))
                state = (u, i + 1)
                return e, state
            else
                i += 1
            end
        end
    end

    if n == 0 || u > n
        return nothing
    end
end

length(eit::MultipleEdgeIter) = ne(eit.g)