text:
# ============================================================
# Methods of Particle Velocity Space
# ============================================================

"""
    struct VSpace1D{TR<:Real,TI<:Integer,TA<:AbstractArray,TB<:AbstractArray{<:Real,1}} <: AbstractVelocitySpace
        u0::TR
        u1::TR
        nu::TI
        u::TA
        du::TA
        weights::TB
    end

1D velocity space

"""
struct VSpace1D{TR<:Real,TI<:Integer,TA<:AbstractArray,TB<:AbstractArray{<:Real,1}} <:
       AbstractVelocitySpace1D
    u0::TR
    u1::TR
    nu::TI
    u::TA
    du::TA
    weights::TB
end

function VSpace1D(
    U0,
    U1,
    NU::TI,
    TYPE = "rectangle",
    NG = zero(NU)::TI,
    PRECISION = Float64,
) where {TI<:Integer}

    δ = (U1 - U0) / NU
    u = OffsetArray{PRECISION}(undef, 1-NG:NU+NG)
    du = similar(u)
    weights = similar(u)

    if TYPE == "rectangle" # rectangular
        for i in eachindex(u)
            u[i] = U0 + (i - 0.5) * δ
            du[i] = δ
            weights[i] = δ
        end
    elseif TYPE == "newton" # newton-cotes
        for i in eachindex(u)
            u[i] = U0 + (i - 0.5) * δ
            du[i] = δ
            weights[i] = newton_cotes(i + NG, NU + NG * 2) * δ
        end
    elseif TYPE == "algebra" # algebraic
        _nu = NU + 1
        _u = [U1 / (_nu - 1)^3 * (-_nu + 1 + 2 * (i - 1))^3 for i = 1:_nu]
        u .= (_u[1:end-1] .+ _u[2:end]) ./ 2
        du .= _u[2:end] - _u[1:end-1]
        weights .= du
    else
        throw("No velocity quadrature available")
    end

    return VSpace1D{PRECISION,TI,typeof(u),typeof(weights)}(U0, U1, NU, u, du, weights)

end

VSpace1D() = VSpace1D(-5, 5, 50)
VSpace1D(U0::T, U1::T) where {T<:Real} = VSpace1D(U0, U1, 50)

"""
    struct VSpace2D{TR<:Real,TI<:Integer,TA<:AbstractArray{<:Real,2}} <: AbstractVelocitySpace
        u0::TR
        u1::TR
        nu::TI
        v0::TR
        v1::TR
        nv::TI
        u::TA
        v::TA
        du::TA
        dv::TA
        weights::TA
    end

2D velocity space

"""
struct VSpace2D{TR<:Real,TI<:Integer,TA<:AbstractArray{<:Real,2}} <: AbstractVelocitySpace2D
    u0::TR
    u1::TR
    nu::TI
    v0::TR
    v1::TR
    nv::TI
    u::TA
    v::TA
    du::TA
    dv::TA
    weights::TA
end

function VSpace2D(
    U0,
    U1,
    NU::TI,
    V0,
    V1,
    NV::TI,
    TYPE = "rectangle",
    NGU = zero(NU)::TI,
    NGV = zero(NV)::TI,
    PRECISION = Float64,
) where {TI<:Integer}

    δu = (U1 - U0) / NU
    δv = (V1 - V0) / NV
    u = OffsetArray{PRECISION}(undef, 1-NGU:NU+NGU, 1-NGV:NV+NGV)
    v = similar(u)
    du = similar(u)
    dv = similar(u)
    weights = similar(u)

    if TYPE == "rectangle" # rectangular
        for j in axes(u, 2)
            for i in axes(u, 1)
                u[i, j] = U0 + (i - 0.5) * δu
                v[i, j] = V0 + (j - 0.5) * δv
                du[i, j] = δu
                dv[i, j] = δv
                weights[i, j] = δu * δv
            end
        end
    elseif TYPE == "newton" # newton-cotes
        for j in axes(u, 2)
            for i in axes(u, 1)
                u[i, j] = U0 + (i - 0.5) * δu
                v[i, j] = V0 + (j - 0.5) * δv
                du[i, j] = δu
                dv[i, j] = δv
                weights[i, j] =
                    newton_cotes(i + NGU, NU + NGU * 2) *
                    δu *
                    newton_cotes(j + NGV, NV + NGV * 2) *
                    δv
            end
        end
    elseif TYPE == "algebra"
        _nu = NU + 1
        _nv = NV + 1
        _u = [U1 / (_nu - 1)^3 * (-_nu + 1 + 2 * (i - 1))^3 for i = 1:_nu]
        _v = [V1 / (_nv - 1)^3 * (-_nv + 1 + 2 * (j - 1))^3 for j = 1:_nv]
        __u = (_u[1:end-1] .+ _u[2:end]) ./ 2
        __v = (_v[1:end-1] .+ _v[2:end]) ./ 2
        u1, v1 = meshgrid(__u, __v)
        u .= u1 |> permutedims
        v .= v1 |> permutedims

        #_wx = [ 3.0 * (-_nu/2 + i - 0.5)^2 / (_nu/2 - 0.5)^3 for i = 1:_nu]
        #_wy = [ 3.0 * (-_nv/2 + j - 0.5)^2 / (_nv/2 - 0.5)^3 for j = 1:_nv]
        #wx = (_wx[1:end-1] .+ _wx[2:end]) ./ 2
        #wy = (_wy[1:end-1] .+ _wy[2:end]) ./ 2

        _du = _u[2:end] - _u[1:end-1]
        _dv = _v[2:end] - _v[1:end-1]
        du1, dv1 = meshgrid(_du, _dv)
        du .= du1 |> permutedims
        dv .= dv1 |> permutedims

        for j in axes(u, 2)
            for i in axes(u, 1)
                weights[i, j] = du[i, j] * dv[i, j]
                #weights[i, j] = wx[i] * wy[j]
            end
        end
    elseif TYPE == "maxwell"
        _u, _uw = maxwell_quadrature(NU, U1 / 5.76)
        _v, _vw = maxwell_quadrature(NV, V1 / 5.76)
        u1, v1 = meshgrid(_u, _v)
        u .= u1 |> permutedims
        v .= v1 |> permutedims
        for j in axes(weights, 2), i in axes(weights, 1)
            weights[i, j] = _uw[i] * _vw[j]
        end

        _du = zero(_u)
        for i in eachindex(_u)
            if i == 1
                _du[i] = _u[2] - _u[1]
            elseif i == length(_u)
                _du[i] = _u[end] - _u[end-1]
            else
                _du[i] = (_u[i+1] - _u[i-1]) / 2
            end
        end
        _dv = zero(_v)
        for i in eachindex(_v)
            if i == 1
                _dv[i] = _v[2] - _v[1]
            elseif i == length(_v)
                _dv[i] = _v[end] - _v[end-1]
            else
                _dv[i] = (_v[i+1] - _v[i-1]) / 2
            end
        end
        du1, dv1 = meshgrid(_du, _dv)
        du .= du1 |> permutedims
        dv .= dv1 |> permutedims
    else
        throw("No velocity quadrature available")
    end

    return VSpace2D{PRECISION,TI,typeof(u)}(U0, U1, NU, V0, V1, NV, u, v, du, dv, weights)

end

VSpace2D() = VSpace2D(-5, 5, 28, -5, 5, 28)
VSpace2D(U0::T, U1::T, V0::T, V1::T) where {T<:Real} = VSpace2D(U0, U1, 28, V0, V1, 28)

"""
    struct VSpace3D{TR<:Real,TI<:Integer,TA<:AbstractArray{<:Real,3}} <: AbstractVelocitySpace
        u0::TR
        u1::TR
        nu::TI
        v0::TR
        v1::TR
        nv::TI
        w0::TR
        w1::TR
        nw::TI
        u::TA
        v::TA
        w::TA
        du::TA
        dv::TA
        dw::TA
        weights::TA
    end

3D velocity space

"""
struct VSpace3D{TR<:Real,TI<:Integer,TA<:AbstractArray{<:Real,3}} <: AbstractVelocitySpace3D
    u0::TR
    u1::TR
    nu::TI
    v0::TR
    v1::TR
    nv::TI
    w0::TR
    w1::TR
    nw::TI
    u::TA
    v::TA
    w::TA
    du::TA
    dv::TA
    dw::TA
    weights::TA
end

function VSpace3D(
    U0,
    U1,
    NU::TI,
    V0,
    V1,
    NV::TI,
    W0,
    W1,
    NW::TI,
    TYPE = "rectangle",
    NGU = zero(NU)::TI,
    NGV = zero(NV)::TI,
    NGW = zero(NW)::TI,
    PRECISION = Float64,
) where {TI<:Integer}

    δu = (U1 - U0) / NU
    δv = (V1 - V0) / NV
    δw = (W1 - W0) / NW
    u = OffsetArray{PRECISION}(undef, 1-NGU:NU+NGU, 1-NGV:NV+NGV, 1-NGW:NW+NGW)
    v = similar(u)
    w = similar(u)
    du = similar(u)
    dv = similar(u)
    dw = similar(u)
    weights = similar(u)

    if TYPE == "rectangle" # rectangular
        for k in axes(u, 3), j in axes(u, 2), i in axes(u, 1)
            u[i, j, k] = U0 + (i - 0.5) * δu
            v[i, j, k] = V0 + (j - 0.5) * δv
            w[i, j, k] = W0 + (k - 0.5) * δw
            du[i, j, k] = δu
            dv[i, j, k] = δv
            dw[i, j, k] = δw
            weights[i, j, k] = δu * δv * δw
        end
    elseif TYPE == "newton" # newton-cotes
        for k in axes(u, 3), j in axes(u, 2), i in axes(u, 1)
            u[i, j, k] = U0 + (i - 0.5) * δu
            v[i, j, k] = V0 + (j - 0.5) * δv
            w[i, j, k] = W0 + (k - 0.5) * δw
            du[i, j, k] = δu
            dv[i, j, k] = δv
            dw[i, j, k] = δw
            weights[i, j, k] =
                newton_cotes(i + NGU, NU + NGU * 2) *
                δu *
                newton_cotes(j + NGV, NV + NGV * 2) *
                δv *
                newton_cotes(k + NGW, NW + NGW * 2) *
                δw
        end
    elseif TYPE == "algebra"
        _nu = NU + 1
        _nv = NV + 1
        _nw = NW + 1
        _u = [U1 / (_nu - 1)^3 * (-_nu + 1 + 2 * (i - 1))^3 for i = 1:_nu]
        _v = [V1 / (_nv - 1)^3 * (-_nv + 1 + 2 * (j - 1))^3 for j = 1:_nv]
        _w = [W1 / (_nw - 1)^3 * (-_nw + 1 + 2 * (k - 1))^3 for k = 1:_nw]
        __u = (_u[1:end-1] .+ _u[2:end]) ./ 2
        __v = (_v[1:end-1] .+ _v[2:end]) ./ 2
        __w = (_w[1:end-1] .+ _w[2:end]) ./ 2
        u1, v1, w1 = meshgrid(__u, __v, __w)
        u .= permutedims(u1, [3, 2, 1])
        v .= permutedims(v1, [3, 2, 1])
        w .= permutedims(w1, [3, 2, 1])

        _du = _u[2:end] - _u[1:end-1]
        _dv = _v[2:end] - _v[1:end-1]
        _dw = _w[2:end] - _w[1:end-1]
        du1, dv1, dw1 = meshgrid(_du, _dv, _dw)
        du .= permutedims(du1, [3, 2, 1])
        dv .= permutedims(dv1, [3, 2, 1])
        dw .= permutedims(dw1, [3, 2, 1])

        for k in axes(u, 3), j in axes(u, 2), i in axes(u, 1)
            weights[i, j, k] = du[i, j, k] * dv[i, j, k] * dw[i, j, k]
        end
    else
        throw("No velocity quadrature available")
    end

    return VSpace3D{PRECISION,TI,typeof(u)}(
        U0,
        U1,
        NU,
        V0,
        V1,
        NV,
        W0,
        W1,
        NW,
        u,
        v,
        w,
        du,
        dv,
        dw,
        weights,
    )

end

VSpace3D() = VSpace3D(-5, 5, 28, -5, 5, 28, -5, 5, 28)
VSpace3D(U0::T, U1::T, V0::T, V1::T, W0::T, W1::T) where {T<:Real} =
    VSpace3D(U0, U1, 28, V0, V1, 28, W0, W1, 28)

"""
    struct MVSpace1D{TR<:AbstractArray{<:Real,1},TI<:Integer,TA<:AbstractArray{<:Real,2}} <: AbstractVelocitySpace
        u0::TR
        u1::TR
        nu::TI
        u::TA
        du::TA
        weights::TA
    end

1D multi-component velocity space

"""
struct MVSpace1D{TR<:AbstractArray{<:Real,1},TI<:Integer,TA<:AbstractArray{<:Real,2}} <:
       AbstractVelocitySpace1D
    u0::TR
    u1::TR
    nu::TI
    u::TA
    du::TA
    weights::TA
end

function MVSpace1D(
    Ui0,
    Ui1,
    Ue0,
    Ue1,
    NU::TI,
    TYPE = "rectangle",
    NG = zero(NU)::TI,
    PRECISION = Float64,
) where {TI<:Integer}

    u0 = PRECISION.([Ui0, Ue0])
    u1 = PRECISION.([Ui1, Ue1])
    δ = (u1 .- u0) ./ NU
    u = OffsetArray{PRECISION}(undef, 1-NG:NU+NG, 1:2)
    du = similar(u)
    weights = similar(u)

    if TYPE == "rectangle" # rectangular
        for j in axes(u, 2), i in axes(u, 1)
            u[i, j] = u0[j] + (i - 0.5) * δ[j]
            du[i, j] = δ[j]
            weights[i, j] = δ[j]
        end
    elseif TYPE == "newton" # newton-cotes
        for j in axes(u, 2), i in axes(u, 1)
            u[i, j] = u0[j] + (i - 0.5) * δ[j]
            du[i, j] = δ[j]
            weights[i, j] = newton_cotes(i + NG, NU + NG * 2) * δ[j]
        end
    else
        throw("No velocity quadrature available")
    end

    return MVSpace1D{typeof(u0),TI,typeof(u)}(u0, u1, NU, u, du, weights)

end

MVSpace1D() = MVSpace1D(-5, 5, -10, 10, 28)
MVSpace1D(U0::T, U1::T, V0::T, V1::T) where {T<:Real} = MVSpace1D(U0, U1, V0, V1, 28)

"""
    struct MVSpace2D{TR<:AbstractArray{<:Real,1},TI<:Integer,TA<:AbstractArray{Float64,3}} <: AbstractVelocitySpace
        u0::TR
        u1::TR
        nu::TI
        v0::TR
        v1::TR
        nv::TI
        u::TA
        v::TA
        du::TA
        dv::TA
        weights::TA
    end

2D multi-component velocity space

"""
struct MVSpace2D{TR<:AbstractArray{<:Real,1},TI<:Integer,TA<:AbstractArray{Float64,3}} <:
       AbstractVelocitySpace2D
    u0::TR
    u1::TR
    nu::TI
    v0::TR
    v1::TR
    nv::TI
    u::TA
    v::TA
    du::TA
    dv::TA
    weights::TA
end

function MVSpace2D(
    Ui0,
    Ui1,
    Ue0,
    Ue1,
    NU::TI,
    Vi0,
    Vi1,
    Ve0,
    Ve1,
    NV::TI,
    TYPE = "rectangle",
    NGU = zero(NU)::TI,
    NGV = zero(NV)::TI,
    PRECISION = Float64,
) where {TI<:Integer}

    u0 = PRECISION.([Ui0, Ue0])
    u1 = PRECISION.([Ui1, Ue1])
    δu = (u1 .- u0) ./ NU
    v0 = PRECISION.([Vi0, Ve0])
    v1 = PRECISION.([Vi1, Ve1])
    δv = (v1 .- v0) ./ NV
    u = OffsetArray{PRECISION}(undef, 1-NGU:NU+NGU, 1-NGV:NV+NGV, 1:2)
    v = similar(u)
    du = similar(u)
    dv = similar(u)
    weights = similar(u)

    if TYPE == "rectangle" # rectangular
        for k in axes(u, 3), j in axes(u, 2), i in axes(u, 1)
            u[i, j, k] = u0[k] + (i - 0.5) * δu[k]
            v[i, j, k] = v0[k] + (j - 0.5) * δv[k]
            du[i, j, k] = δu[k]
            dv[i, j, k] = δv[k]
            weights[i, j, k] = δu[k] * δv[k]
        end
    elseif TYPE == "newton" # newton-cotes
        for k in axes(u, 3), j in axes(u, 2), i in axes(u, 1)
            u[i, j, k] = u0[k] + (i - 0.5) * δu[k]
            v[i, j, k] = v0[k] + (j - 0.5) * δv[k]
            du[i, j, k] = δu[k]
            dv[i, j, k] = δv[k]
            weights[i, j, k] =
                newton_cotes(i + NGU, NU + NGU * 2) *
                δu[k] *
                newton_cotes(j + NGV, NV + NGV * 2) *
                δv[k]
        end
    else
        throw("No velocity quadrature available")
    end

    return MVSpace2D{typeof(u0),TI,typeof(u)}(u0, u1, NU, v0, v1, NV, u, v, du, dv, weights)

end

MVSpace2D() = MVSpace2D(-5, 5, -10, 10, 28, -5, 5, -10, 10, 28)
MVSpace2D(U0::T, U1::T, V0::T, V1::T) where {T<:Real} =
    MVSpace2D(U0, U1, U0, U1, 28, V0, V1, V0, V1, 28)

"""
    struct MVSpace3D{TR<:AbstractArray{<:Real,1},TI<:Integer,TA<:AbstractArray{<:Real,4}} <: AbstractVelocitySpace
        u0::TR
        u1::TR
        nu::TI
        v0::TR
        v1::TR
        nv::TI
        w0::TR
        w1::TR
        nw::TI
        u::TA
        v::TA
        w::TA
        du::TA
        dv::TA
        dw::TA
        weights::TA
    end

3D multi-component velocity space

"""
struct MVSpace3D{TR<:AbstractArray{<:Real,1},TI<:Integer,TA<:AbstractArray{<:Real,4}} <:
       AbstractVelocitySpace3D
    u0::TR
    u1::TR
    nu::TI
    v0::TR
    v1::TR
    nv::TI
    w0::TR
    w1::TR
    nw::TI
    u::TA
    v::TA
    w::TA
    du::TA
    dv::TA
    dw::TA
    weights::TA
end

function MVSpace3D(
    Ui0,
    Ui1,
    Ue0,
    Ue1,
    NU::TI,
    Vi0,
    Vi1,
    Ve0,
    Ve1,
    NV::TI,
    Wi0,
    Wi1,
    We0,
    We1,
    NW::TI,
    TYPE = "rectangle",
    NGU = zero(NU)::TI,
    NGV = zero(NV)::TI,
    NGW = zero(NW)::TI,
    PRECISION = Float64,
) where {TI<:Integer}

    u0 = PRECISION.([Ui0, Ue0])
    u1 = PRECISION.([Ui1, Ue1])
    δu = (u1 .- u0) ./ NU
    v0 = PRECISION.([Vi0, Ve0])
    v1 = PRECISION.([Vi1, Ve1])
    δv = (v1 .- v0) ./ NV
    w0 = PRECISION.([Wi0, We0])
    w1 = PRECISION.([Wi1, We1])
    δw = (w1 .- w0) ./ NW

    u = OffsetArray{PRECISION}(undef, 1-NGU:NU+NGU, 1-NGV:NV+NGV, 1-NGW:NW+NGW, 1:2)
    v = similar(u)
    w = similar(u)
    du = similar(u)
    dv = similar(u)
    dw = similar(u)
    weights = similar(u)

    if TYPE == "rectangle" # rectangular
        for l in axes(u, 4), k in axes(u, 3), j in axes(u, 2), i in axes(u, 1)
            u[i, j, k, l] = u0[l] + (i - 0.5) * δu[l]
            v[i, j, k, l] = v0[l] + (j - 0.5) * δv[l]
            w[i, j, k, l] = w0[l] + (k - 0.5) * δw[l]
            du[i, j, k, l] = δu[l]
            dv[i, j, k, l] = δv[l]
            dw[i, j, k, l] = δw[l]
            weights[i, j, k, l] = δu[l] * δv[l] * δw[l]
        end
    elseif TYPE == "newton" # newton-cotes
        for l in axes(u, 4), k in axes(u, 3), j in axes(u, 2), i in axes(u, 1)
            u[i, j, k, l] = u0[l] + (i - 0.5) * δu[l]
            v[i, j, k, l] = v0[l] + (j - 0.5) * δv[l]
            w[i, j, k, l] = w0[l] + (k - 0.5) * δw[l]
            du[i, j, k, l] = δu[l]
            dv[i, j, k, l] = δv[l]
            dw[i, j, k, l] = δw[l]
            weights[i, j, k, l] =
                newton_cotes(i + NGU, NU + NGU * 2) *
                δu[l] *
                newton_cotes(j + NGV, NV + NGV * 2) *
                δv[l] *
                newton_cotes(k + NGW, NW + NGW * 2) *
                δw[l]
        end
    else
        throw("No velocity quadrature available")
    end

    return MVSpace3D{typeof(u0),TI,typeof(u)}(
        u0,
        u1,
        NU,
        v0,
        v1,
        NV,
        w0,
        w1,
        NW,
        u,
        v,
        v,
        du,
        dv,
        dw,
        weights,
    )

end

MVSpace3D() = MVSpace3D(-5, 5, -10, 10, 20, -5, 5, -10, 10, 20, -5, 5, -10, 10, 20)
MVSpace3D(U0::T, U1::T, V0::T, V1::T, W0::T, W1::T) where {T<:Real} =
    MVSpace3D(U0, U1, U0, U1, 20, V0, V1, V0, V1, 20, W0, W1, W0, W1, 20)

"""
    struct UnstructVSpace{TR<:Real,TI<:Integer,TA<:AbstractArray,TB<:AbstractArray{<:Real,1}} <: AbstractVelocitySpace
        u0::TR
        u1::TR
        nu::TI
        u::TA
        weights::TB
    end

Unstructured velocity space

"""
struct UnstructVSpace{
    TR<:Union{Real,AbstractVector},
    TI<:Union{Integer,AbstractVector},
    TA<:AbstractArray,
    TB<:AbstractVector{<:Real},
} <: AbstractVelocitySpace
    u0::TR
    u1::TR
    nu::TI
    u::TA
    weights::TB
end

"""
    newton_cotes(idx::T, num::T) where {T<:Integer}

Evaluate quadrature weight from Newton-Cotes rule
"""
function newton_cotes(idx::T, num::T) where {T<:Integer}
    if idx == 1 || idx == num
        nc_coeff = 14.0 / 45.0
    elseif (idx - 5) % 4 == 0
        nc_coeff = 28.0 / 45.0
    elseif (idx - 3) % 4 == 0
        nc_coeff = 24.0 / 45.0
    else
        nc_coeff = 64.0 / 45.0
    end

    return nc_coeff
end

# ------------------------------------------------------------
# Extended Base.show()
# ------------------------------------------------------------
function Base.show(io::IO, vs::VSpace1D{TR,TI,TA,TB}) where {TR,TI,TA,TB}
    print(
        io,
        "VelocitySpace1D{$TR,$TI,$TA,$TB}\n",
        "domain: ($(vs.u0),$(vs.u1))\n",
        "resolution: $(vs.nu)\n",
        "ghost: $(1-firstindex(vs.u))\n",
    )
end

function Base.show(io::IO, vs::VSpace2D{TR,TI,TA}) where {TR,TI,TA}
    print(
        io,
        "VelocitySpace2D{$TR,$TI,$TA}\n",
        "domain: ($(vs.u0),$(vs.u1)) × ($(vs.v0),$(vs.v1))\n",
        "resolution: $(vs.nu) × $(vs.nv)\n",
        "ghost in u: $(1-firstindex(vs.u[:, 1]))\n",
        "ghost in v: $(1-firstindex(vs.v[1, :]))\n",
    )
end

function Base.show(io::IO, vs::VSpace3D{TR,TI,TA}) where {TR,TI,TA}
    print(
        io,
        "VelocitySpace3D{$TR,$TI,$TA}\n",
        "domain: ($(vs.u0),$(vs.u1)) × ($(vs.v0),$(vs.v1)) × ($(vs.w0),$(vs.w1))\n",
        "resolution: $(vs.nu) × $(vs.nv) × $(vs.nw)\n",
        "ghost in u: $(1-firstindex(vs.u[:, 1, 1]))\n",
        "ghost in v: $(1-firstindex(vs.v[1, :, 1]))\n",
        "ghost in w: $(1-firstindex(vs.w[1, 1, :]))\n",
    )
end

function maxwell_quadrature(N::Integer, C = 1::Real)
    @assert N <= 33

    py"""
    import numpy as np
    from numpy import linalg as LA

    def dvGH(N2,C):
        N = N2//2

        a = np.zeros(N)
        b = np.zeros(N)
        a[0] = 1.0/np.sqrt(np.pi)
        a[1] = 2.0/np.sqrt(np.pi)/(np.pi-2.0)
        b[1] = a[0]/( a[0] + a[1])/2.0

        for i in range(2,N):
            b[i] = (i-1)+1.0/2.0-b[i-1]-a[i-1]**2
            a[i] = (i**2/4.0/b[i]-b[i-1]-1.0/2)/a[i-1]-a[i-1]

        J = np.diag(a) + np.diag(np.sqrt(b[1:N]),1) \
        + np.diag(np.sqrt(b[1:N]),-1)

        v,V = LA.eig(J)

        w = V[0,:]*V[0,:]*np.sqrt(np.pi)/2.0

        vw = np.transpose(np.vstack((v,w)))
        vw = vw[vw[:,0].argsort()]
        v = vw[:,0]
        w = vw[:,1]

        Xis = np.hstack((-np.flipud(v),v))
        weights = np.hstack((np.flipud(w),w))
        weights = weights*np.exp(Xis**2)*C
        Xis = Xis*C
        return (Xis, weights)
    """

    p, w = py"dvGH"(N, C)
end


text:
export add, add_sym, add_to!, mul, mul_sym, mul_to!, sqr, sqr_sym, sqr_to!

function add(a::SparseMatrix{T}, b::SparseMatrix{T}) where {T}
    @assert a.row_count == b.row_count
    @assert a.column_count == b.column_count

    local c = add_sym(a, b)
    add_to!(a, b, c)
    return c
end

function add_sym(a::SparseMatrix{T}, b::SparseMatrix{T}) where {T}
    @assert a.row_count == b.row_count
    @assert a.column_count == b.column_count

    local c = DirectedGraph{T}(a.row_count)

    clear!(a.list)

    for j = 1:a.column_count
        for i = a.columns[j]:(a.columns[j+1]-1)
            push!(a.list, a.columns_rows[i])
        end
        for i = b.columns[j]:(b.columns[j+1]-1)
            push!(a.list, b.columns_rows[i])
        end
        while !is_empty(a.list)
            connect!(c, T(0), pop!(a.list), j)
        end
    end

    return from_graph(c, a.row_count, a.column_count)
end

function add_to!(
    a::SparseMatrix{T},
    b::SparseMatrix{T},
    c::SparseMatrix{T},
) where {T}
    @assert a.row_count == b.row_count
    @assert a.column_count == b.column_count
    @assert a.row_count == c.row_count
    @assert a.column_count == c.column_count

    @inbounds begin
        local acc = zeros(T, a.column_count)

        for i = 1:a.row_count
            for j = a.columns[i]:(a.columns[i+1]-1)
                acc[a.columns_rows[j]] = a.values[j]
            end

            for j = b.columns[i]:(b.columns[i+1]-1)
                acc[b.columns_rows[j]] += b.values[j]
            end

            for j = c.columns[i]:(c.columns[i+1]-1)
                local jj = c.columns_rows[j]
                c.values[j] = acc[jj]
                acc[jj] = T(0)
            end
        end
    end
end

function mul(a::SparseMatrix{T}, b::SparseMatrix{T}) where {T}
    @assert a.column_count == b.row_count

    local c = mul_sym(a, b)
    mul_to!(a, b, c)
    return c
end

function mul_sym(a::SparseMatrix{T}, b::SparseMatrix{T}) where {T}
    @assert a.column_count == b.row_count

    clear!(a.list)

    local c = DirectedGraph{T}(a.row_count)

    for j = 1:b.column_count
        for i = b.columns[j]:(b.columns[j+1]-1)
            local ii = b.columns_rows[i]
            for k = a.columns[ii]:(a.columns[ii+1]-1)
                push!(a.list, a.columns_rows[k])
            end
        end

        while !is_empty(a.list)
            connect!(c, T(0), pop!(a.list), j)
        end
    end

    return from_graph(c, a.row_count, b.column_count)
end

function mul_to!(
    a::SparseMatrix{T},
    b::SparseMatrix{T},
    c::SparseMatrix{T},
) where {T}
    @assert a.row_count == c.row_count
    @assert a.column_count == b.row_count
    @assert b.column_count == c.column_count

    @inbounds begin
        local acc = zeros(T, c.row_count)

        for j = 1:b.column_count
            for i = b.columns[j]:(b.columns[j+1]-1)
                local ii = b.columns_rows[i]
                for k = a.columns[ii]:(a.columns[ii+1]-1)
                    acc[a.columns_rows[k]] += b.values[i] * a.values[k]
                end
            end

            for i = c.columns[j]:(c.columns[j+1]-1)
                local jj = c.columns_rows[i]
                c.values[i] = acc[jj]
                acc[jj] = T(0)
            end
        end
    end
end

function mul(a::SparseMatrix{T}, b::Vector{T}) where {T}
    local c = zeros(T, a.row_count)
    mul_to!(a, b, c)
    return c
end

function mul_to!(a::SparseMatrix{T}, b::Vector{T}, c::Vector{T}) where {T}
    @assert a.column_count == length(b)
    @assert a.row_count == length(c)

    for i = 1:a.row_count
        for j = a.rows[i]:(a.rows[i+1]-1)
            c[i] += b[a.rows_columns[j]] * a.values[a.positions[j]]
        end
    end
end

function mul(a::SparseMatrix{T}, b::SparseArray{T}) where {T}
    local c = SparseArray{T}(a.row_count)
    mul_to!(a, b, c)
    return c
end

function mul_to!(
    a::SparseMatrix{T},
    b::SparseArray{T},
    c::SparseArray{T},
) where {T}
    @assert a.column_count == b.size
    @assert a.row_count == c.size

    local zero = T(0)
    for i = 1:a.row_count
        if a.rows[i+1] == a.rows[i]
            continue
        end

        local x = zero

        for j = a.rows[i]:(a.rows[i+1]-1)
            local r = a.rows_columns[j]
            if !ismissing(b[r])
                x += b[r] * a.values[a.positions[j]]
            end
        end

        if zero != x
            c[i] = x
        end
    end
end

function sqr(a::SparseMatrix{T}) where {T}
    local s = sqr_sym(a)
    sqr_to!(a, s)
    return s
end

function sqr_sym(a::SparseMatrix{T}) where {T}
    local g = DirectedGraph{T}(a.column_count)
    clear!(a.list)

    for j = 1:a.column_count
        for i = a.columns[j]:(a.columns[j+1]-1)
            local ii = a.columns_rows[i]
            for k = a.rows[ii]:(a.rows[ii+1]-1)
                local r = a.rows_columns[k]
                push!(a.list, r)
            end
        end

        push!(a.list, j)

        while !is_empty(a.list)
            connect!(g, T(0), pop!(a.list), j)
        end
    end

    local ata = from_graph(g, a.column_count, a.column_count)

    return ata
end

function sqr_to!(a::SparseMatrix{T}, ata::SparseMatrix{T}) where {T}
    @assert default_layout == ata.layout

    @inbounds begin
        local acc = zeros(T, a.column_count)

        for j = 1:a.column_count
            for i = a.columns[j]:(a.columns[j+1]-1)
                local ii = a.columns_rows[i]

                for k = a.rows[ii]:(a.rows[ii+1]-1)
                    local r = a.rows_columns[k]

                    acc[r] += a.values[i] * a.values[a.positions[k]]
                end
            end

            for i = ata.columns[j]:(ata.columns[j+1]-1)
                local r = ata.columns_rows[i]
                ata.values[i] = acc[r]
                acc[r] = T(0)
            end
        end
    end
end

function sqr_sym(a::SparseMatrix{T}, pt::PermuteTable) where {T}
    @assert default_layout == a.layout

    local g = DirectedGraph{T}(a.column_count)
    clear!(a.list)

    for j = 1:a.column_count
        local jj = pt.permuted[j]

        for i = a.columns[jj]:(a.columns[jj+1]-1)
            local ii = a.columns_rows[i]
            for k = a.rows[ii]:(a.rows[ii+1]-1)
                local r = pt.primary[a.rows_columns[k]]

                if r < j
                    continue
                end

                push!(a.list, r)
            end
        end

        push!(a.list, j)

        while !is_empty(a.list)
            connect!(g, T(0), pop!(a.list), j)
        end
    end

    local ata = from_graph(g, a.column_count, a.column_count)
    ata.layout = lower_symmetric

    return ata
end

function sqr_to!(
    a::SparseMatrix{T},
    ata::SparseMatrix{T},
    pt::PermuteTable,
) where {T}
    @assert default_layout == a.layout
    @assert lower_symmetric == ata.layout

    @inbounds begin
        local acc = zeros(T, a.column_count)

        for j = 1:a.column_count
            local jj = pt.permuted[j]

            for i = a.columns[jj]:(a.columns[jj+1]-1)
                local ii = a.columns_rows[i]

                for k = a.rows[ii]:(a.rows[ii+1]-1)
                    local r = pt.primary[a.rows_columns[k]]

                    if r < j
                        continue
                    end

                    acc[r] += a.values[i] * a.values[a.positions[k]]
                end
            end

            for i = ata.columns[j]:(ata.columns[j+1]-1)
                local r = ata.columns_rows[i]
                ata.values[i] = acc[r]
                acc[r] = T(0)
            end
        end
    end
end

ext: jl


text:
abstract type AtomIterator end

Base.eltype(::AtomIterator) = Int64
Base.length(x::AtomIterator) = x.fullPass ? ProximalBase.numCoordinates(x.iterate) : nnz(x.iterate)

# T = SparseIterate, SymmetricSparseIterate, AtomIterate
# last argument represents full pass
mutable struct OrderedIterator{T} <: AtomIterator
  iterate::T
  fullPass::Bool

  OrderedIterator{T}(iterate::Union{SparseIterate,SymmetricSparseIterate,AtomIterate}, fullPass) where {T} = new(iterate, fullPass)
end

OrderedIterator(iterate) = OrderedIterator{typeof(iterate)}(iterate, true)

Base.iterate(x::OrderedIterator) = iterate(x, 1)
function Base.iterate(x::OrderedIterator, i)
    # check if done
    done = x.fullPass ? i > ProximalBase.numCoordinates(x.iterate) : i > nnz(x.iterate)
    if done
        return nothing
    end
    x.fullPass ? (i, i + 1) : (x.iterate.nzval2ind[i], i + 1)
end

# Base.start(x::OrderedIterator) = 1
# # current item, next state
# Base.next(x::OrderedIterator, i) = x.fullPass ? (i, i + 1) : (x.iterate.nzval2ind[i], i + 1)
# Base.done(x::OrderedIterator, i) = x.fullPass ? i > ProximalBase.numCoordinates(x.iterate) : i > nnz(x.iterate)
#
function reset!(x::OrderedIterator, fullPass::Bool)
  x.fullPass = fullPass
  x
end

mutable struct RandomIterator{T} <: AtomIterator
  iterate::T
  order::Vector{Int64}
  fullPass::Bool

  RandomIterator{T}(iterate::Union{SparseIterate,SymmetricSparseIterate,AtomIterate}, order, fullPass) where {T} =
    new(iterate, order, fullPass)
end

RandomIterator(iterate::Union{SparseIterate,SymmetricSparseIterate,AtomIterate}) =
  RandomIterator{typeof(iterate)}(iterate, collect(1:ProximalBase.numCoordinates(iterate)), true)

function reset!(x::RandomIterator, fullPass::Bool)
  x.fullPass = fullPass
  newLength = fullPass ? ProximalBase.numCoordinates(x.iterate) : nnz(x.iterate)
  @inbounds for i=1:newLength
    x.order[i] = i
  end
  @inbounds for i=1:newLength-1
    j = rand(i:newLength)
    x.order[i], x.order[j] = x.order[j], x.order[i]
  end
  x
end

Base.iterate(x::RandomIterator) = iterate(x, 1)
function Base.iterate(x::RandomIterator, i)
    # check if done
    done = x.fullPass ? i > ProximalBase.numCoordinates(x.iterate) : i > nnz(x.iterate)
    if done
        return nothing
    end
    x.fullPass ? (x.order[i], i + 1) : (x.iterate.nzval2ind[x.order[i]], i + 1)
end

# Base.start(x::RandomIterator) = 1
# Base.next(x::RandomIterator, i) = x.fullPass ? (x.order[i], i + 1) : (x.iterate.nzval2ind[x.order[i]], i + 1)
# Base.done(x::RandomIterator, i) = x.fullPass ? i > ProximalBase.numCoordinates(x.iterate) : i > nnz(x.iterate)


text:
# The code to generate Fig 3
using PyPlot
using SurNoR_2020
using Statistics

PyPlot.svg(true)
rcParams = PyPlot.PyDict(PyPlot.matplotlib."rcParams")
rcParams["svg.fonttype"] = "none"
rcParams["pdf.fonttype"] = 42

# ------------------------------------------------------------------------------
# Training Algorithms
# ------------------------------------------------------------------------------
Agents_SurNoR = SurNoR_2020.Func_SurNoR_Overall_train_opt_ind_i(3)

# ------------------------------------------------------------------------------
# ------------------------------------------------------------------------------
# Fig 3A: Count and Novelty Time Series
# Only the data for subject 9 is shown in the paper
# ------------------------------------------------------------------------------
# ------------------------------------------------------------------------------
Epi = 1
Block = 1
Colors = ["#E56969","#C14545","#B71616","#45A2C1","#167CB7","#B159E0"]
for Sub = 1:12
    Agent = Agents_SurNoR[Sub,Block,Epi]

    State_Set = [9,10,11,5,8,1]
    Order = [1,2,6,4,9,5,10,11,3,7,8]
    Names = ["State 8 (Trap)", "State 9 (Trap)", "State 10 (Trap)", "State 4", "State 7", "Goal State"]

    fig = figure(figsize=(7,6)); ax = gca()
    for j = 1:length(State_Set)
       i = State_Set[j]
       ax.plot(Agent.Counter[Order[i],:],color=Colors[j])
    end
    title(string("Counter over Time, 1st Block, 1st Epi., Sub. ",string(Agent.Sub)))
    ax.legend(Names)
    ax.set_xlabel("Time")
    ax.set_ylabel("Counter");

    fig = figure(figsize=(7,6)); ax = gca()
    for j = 1:length(State_Set)
       i = State_Set[j]
       ax.plot(Agent.Novelty[Order[i],:],color=Colors[j])
    end
    title(string("Novelty over Time, 1st Block, 1st Epi., Sub. ",string(Agent.Sub)))
    ax.legend(Names)
    ax.set_xlabel("Time")
    ax.set_ylabel("Novelty");
end

# ------------------------------------------------------------------------------
# ------------------------------------------------------------------------------
# Fig 3B: Novelty Map
# ------------------------------------------------------------------------------
# ------------------------------------------------------------------------------
Novelty = zeros(11,12)
Epi = 1
Block = 1
for Sub = 1:12
    Ind = 1:length(Agents_SurNoR[Sub,Block,Epi].Input.obs)
    Ind = Ind[Agents_SurNoR[Sub,Block,Epi].Input.obs .== 10]
    Novelty[:,Sub] = Agents_SurNoR[Sub,Block,Epi].Novelty[:,findmin(Ind)[1]]
end

Novelty_Mean = mean(Novelty,dims=2)

Order = [1,5,3,8,4,9,10,0,2,6,7]
Names = ["1","2","3","4","5","6","7","G","8","9","10"]

N_Ordered = zeros(size(Novelty_Mean))
for i=1:length(Novelty_Mean)
    N_Ordered[i] = Novelty_Mean[Order[i]+1]
end

x = [1,2,3,4,5,6,7,8, 3.5, 4.5, 5.5]
y = [1,1,1,1,1,1,1,1, 0.7, 0.7, 0.7]

fig = figure(figsize=(8,6)); ax = gca()
sc = ax.scatter(x, y, s=ones(size(x)).*1200, c=N_Ordered[:])

for i =1:length(Names)
 ax.text(x[i],y[i],Names[i],fontsize=12,
          horizontalalignment="center", verticalalignment="center")
end
ax.grid(false)
ax.set_xticklabels([]); ax.set_yticklabels([])
ax.set_xticks([]); ax.set_yticks([])
ax.set_xlim(-0.1,9.1); ax.set_ylim(0.2,1.4)
ax.set_title("Average of Novelty of States - 1st Episode of the 1st Block")

cbar= fig.colorbar(sc)
cbar.set_label("Novelty Value")


text:
function assemble(problem::LinearElasticityProblem{dim,T}, elementinfo::ElementFEAInfo{dim, T}) where {dim,T}
    globalinfo = GlobalFEAInfo(problem)
    assemble!(globalinfo, problem, elementinfo)
    return globalinfo
end

function assemble!(globalinfo::GlobalFEAInfo{T}, problem::LinearElasticityProblem{dim,T}, elementinfo::ElementFEAInfo{dim, T, TK}) where {dim, T, TK}
    ch = problem.ch
    dh = ch.dh
    K, f = globalinfo.K, globalinfo.f
    f .= elementinfo.fixedload
    Kes, fes = elementinfo.Kes, elementinfo.fes
    if K isa Symmetric
        K.data.nzval .= 0
        assembler = JuAFEM.AssemblerSparsityPattern(K.data, f, Int[], Int[])
    else
        K.nzval .= 0
        assembler = JuAFEM.AssemblerSparsityPattern(K, f, Int[], Int[])
    end
    global_dofs = zeros(Int, ndofs_per_cell(dh))
    fe = zeros(fes[1])
    Ke = zeros(T, size(Kes[1]))
    celliteratortype = CellIterator{typeof(dh).parameters...}
    _celliterator::celliteratortype = CellIterator(dh)
    for (i,cell) in enumerate(_celliterator)
        celldofs!(global_dofs, dh, i)
        if TK <: Symmetric
            JuAFEM.assemble!(assembler, global_dofs, Kes[i].data, fes[i])
        else
            JuAFEM.assemble!(assembler, global_dofs, Kes[i], fes[i])
        end
    end
    if TK <: Symmetric
        apply!(K.data, f, ch)
    else
        apply!(K, f, ch)
    end

    return
end

function assemble_f(problem::LinearElasticityProblem{dim,T}, elementinfo::ElementFEAInfo{dim, T}) where {dim, T}
    f = zeros(T, ndofs(problem.ch.dh))
    assemble_f!(f, problem, elementinfo)
    return f
end

function assemble_f!(f::AbstractVector, problem::LinearElasticityProblem, elementinfo::ElementFEAInfo)
    fes = elementinfo.fes
    f .= elementinfo.fixedload
    dof_cells = problem.metadata.dof_cells
    dof_cells_offset = problem.metadata.dof_cells_offset
    for dofidx in 1:ndofs(problem.ch.dh)
        r = dof_cells_offset[dofidx] : dof_cells_offset[dofidx+1]-1
        for i in r
            cellidx, localidx = dof_cells[i]
            f[dofidx] += fes[cellidx][localidx]
        end
    end
    return f
end

function assemble_f!(f::AbstractVector, problem, dloads)
    metadata = problem.metadata
    dof_cells = metadata.dof_cells
    dof_cells_offset = metadata.dof_cells_offset
    for dofidx in 1:ndofs(problem.ch.dh)
        r = dof_cells_offset[dofidx] : dof_cells_offset[dofidx+1]-1
        for i in r
            cellidx, localidx = dof_cells[i]
            f[dofidx] += dloads[cellidx][localidx]
        end
    end
    return f
end


text:
using StaticArrays

function Harmonic(X)
    return 0.5 * X[1]^2
end

function DoubleWell(X)
    return (X[1]^2 -1)^2
end

function Muller(x)

    aa = @SVector [-1, -1, -6.5, 0.7];
    bb = @SVector [0., 0., 11., 0.6];
    cc = @SVector [-10., -10., -6.5, 0.7];
    AA = @SVector [-200., -100., -170., 15.];
    XX = @SVector [1., 0., -0.5, -1.];
    YY = @SVector [0., 0.5, 1.5, 1.];

    return ( AA[1]*exp(aa[1]*(x[1]-XX[1])^2+bb[1]*(x[1]-XX[1])*(x[2]-YY[1])+cc[1]*(x[2]-YY[1])^2)
                 +AA[2]*exp(aa[2]*(x[1]-XX[2])^2+bb[2]*(x[1]-XX[2])*(x[2]-YY[2])+cc[2]*(x[2]-YY[2])^2)
                 +AA[3]*exp(aa[3]*(x[1]-XX[3])^2+bb[3]*(x[1]-XX[3])*(x[2]-YY[3])+cc[3]*(x[2]-YY[3])^2)
                 +AA[4]*exp(aa[4]*(x[1]-XX[4])^2+bb[4]*(x[1]-XX[4])*(x[2]-YY[4])+cc[4]*(x[2]-YY[4])^2));
end

max_stars_repo_name: cortner/JuBasicMD


text:
export with, Parameter, OptionParameter, LazyParameter

mutable struct Parameter{T}
  value::T
end

(p::Parameter{T})() where T = p.value::T
(p::Parameter{T})(newvalue::T) where T = p.value = newvalue::T

function with(f, p, newvalue)
  oldvalue = p()
  p(newvalue)
  try
    f()
  finally
    p(oldvalue)
  end
end

with(f, p, newvalue, others...) =
  with(p, newvalue) do
    with(f, others...)
  end

mutable struct OptionParameter{T}
  value::Union{Missing,T}
end
OptionParameter{T}() where T = OptionParameter{T}(missing)

(p::OptionParameter{T})() where T =
  ismissing(p.value) ?
    error("Parameter was not initialized") :
    p.value
(p::OptionParameter{T})(newvalue::T) where T = p.value = newvalue

mutable struct LazyParameter{T}
  initializer::Union{DataType,Function} #This should be a more specific type: None->T
  value::Union{T,Nothing}
end

LazyParameter(T::DataType, initializer::Union{DataType,Function}) =
  LazyParameter{T}(initializer, nothing)

(p::LazyParameter{T})() where T = isnothing(p.value) ? (p.value = p.initializer()) : p.value
(p::LazyParameter{T})(newvalue::T) where T = p.value = newvalue

Base.reset(p::LazyParameter{T}) where T = p.value = nothing


text:
export ℎ, ħ, ℏ, 𝑒, Φ₀, ϕ₀, ln2,
    sigma0, σ0, sigmax, σx, sigmay, σy, sigmaz, σz, sigmaplus, σ₊, sigmaminus, σ₋,
    ⊗

# Physical and mathematical constants
const ℎ = 6.626069934E-34 # Planck constant, in Joules*seconds, NIST 2017
const ħ = 1.054571783E-34 # Reduced Planck constant
const ℏ = ħ # Actual unicode symbol
const 𝑒 = 1.6021766208E-19 # Elementary charge, in Coulomb, 2014 CODATA
const Φ₀ = 2.067833831E-15 # Magnetic flux quantum (ℎ/2𝑒), in Weber, 2014 CODATA
const ϕ₀ = 3.291059757E-16 # Reduced flux quantum (ℏ/2𝑒), in Weber, 2014 CODATA
const ln2 = 0.693147180559945309 # Natural logarithm of 2
const sqrtln16 = 1.66510922231539551 # Square root of ln(16)

# Pauli matrices (spin-1/2 operators)
const sigma0 = qeye(2)
const sigmax = Operator(sparse([1,2],[2,1],[1.0,1.0],2,2),(2,))
const sigmay = Operator(sparse([1,2],[2,1],[-1.0im,1.0im],2,2),(2,))
const sigmaz = Operator(sparse([1,2],[1,2],[1.0,-1.0],2,2),(2,))
const σ0 = sigma0
const σx = sigmax
const σy = sigmay
const σz = sigmaz
const sigmaplus  = Operator(sparse([2],[1],[1.0],2,2),(2,)) # σ₊*[1,0] = [0,1]
const sigmaminus = Operator(sparse([1],[2],[1.0],2,2),(2,))
const σ₊ = sigmaplus
const σ₋ = sigmaminus

# Function aliases
const ⊗ = kron
const tensor = kron
const fock = basis
const qzeros = qzero

# Table of square root of factorials √(n!) for n ∈ 0:127 (1KiB)
const sqrtfact_table = Float64[
    1.0000000000000000000, 1.0000000000000000000, 1.4142135623730950488,
    2.4494897427831780982, 4.8989794855663561964, 10.954451150103322269,
    26.832815729997476357, 70.992957397195392511, 200.79840636817813151,
    602.39521910453439454, 1904.9409439665052252, 6317.9743589223278835,
    21886.105181141755630, 78911.474450804681438, 295259.70128007648582,
    1.1435359058639129588E6, 4.5741436234556518350E6,
    1.8859677306253148084E7, 8.0014834285449844953E7,
    3.4877657663442939413E8, 1.5597762686284978865E9,
    7.1477928181858656891E9, 3.3526120082371710076E10,
    1.6078562354540587669E11, 7.8768547132293829008E11,
    3.9384273566146914504E12, 2.0082117944245961319E13,
    1.0434974580907397764E14, 5.5216695356722848749E14,
    2.9735100460129106405E15, 1.6286585271694955843E16,
    9.0679869067935485290E16, 5.1296280268036349125E17,
    2.9467469553410734788E18, 1.7182339742875652406E19,
    1.0165209277917570217E20, 6.0991255667505421303E20,
    3.7099532465014090857E21, 2.2869687743093501008E22,
    1.4282115417961529469E23, 9.0328029052332240864E23,
    5.7838159214452708158E24, 3.7483411234209726053E25,
    2.4579516484946125896E26, 1.6304206741784307882E27,
    1.0937194378152021970E28, 7.4179661362209580728E28,
    5.0855013667402369566E29, 3.5233386996620226535E30,
    2.4663370897634158575E31, 1.7439636808636061170E32,
    1.2454391808865586995E33, 8.9809896543167155890E33,
    6.5382591597917144388E34, 4.8046196242703894246E35,
    3.5632012788584194610E36, 2.6664556771205919519E37,
    2.0131298891248228833E38, 1.5331540468207617694E39,
    1.1776379687564843276E40, 9.1219444817107882595E40,
    7.1244663931920178495E41, 5.6098104478126475754E42,
    4.4526490041372451123E43, 3.5621192033097960898E44,
    2.8718723147247460219E45, 2.3331200978034608324E46,
    1.9097411059666879697E47, 1.5748128594969087944E48,
    1.3081378078327271991E49, 1.0944666130115569586E50,
    9.2221396029764281460E50, 7.8252449403763769254E51,
    6.6858922078602825325E52, 5.7514219472399922436E53,
    4.9808775141931966697E54, 4.3422283469044442401E55,
    3.8102899106011063425E56, 3.3651569321810681095E57,
    2.9910169058002623210E58, 2.6752468492881886262E59,
    2.4077221643593697636E60, 2.1802851503903891306E61,
    1.9863343046226278804E62, 1.8205054612841328324E63,
    1.6784231035053557904E64, 1.5565055535934567420E65,
    1.4518117296604018405E66, 1.3619201234191322394E67,
    1.2848328747704294744E68, 1.2188994890809338170E69,
    1.1627560052213890684E70, 1.1152763807523813626E71,
    1.0755335917960171153E72, 1.0427685057848376926E73,
    1.0163650175128549387E74, 9.9583027412855333880E74,
    9.8077907646157562109E75, 9.7092175013660332284E76,
    9.6605494379949297313E77, 9.6605494379949297313E78,
    9.7087320283538361861E79, 9.8053387065593642395E80,
    9.9513319291872585027E81, 1.0148407138632952789E83,
    1.0399022830248479171E84, 1.0706449288791818466E85,
    1.1074837259283487726E86, 1.1509308491181513962E87,
    1.2016070835354182374E88, 1.2602561412358770729E89,
    1.3277622343967479571E90, 1.4051714689748984594E91,
    1.4937177607097713179E92, 1.5948541417547209534E93,
    1.7102905289724946877E94, 1.8420392733196266468E95,
    1.9924701154117019517E96, 2.1643765498993678241E97,
    2.3610560690520684628E98, 2.5864073371085860453E99,
    2.8450480708194446498E100, 3.1424583053452915085E101,
    3.4851548555301426772E102, 3.8809042007129483894E103,
    4.3389828034793201975E104, 4.8704961173170509835E105,
    5.4887703709097355030E106
]


text:
include("neural_networks.jl")

using IterTools
using Functors

using Random
using Optim
using Flux

Random.seed!(1)
function normalize_density(density, total)
    # return density
    s = sum(density)
    if s == 0
        return density
    end
    density * total / s
end
struct DensityPredictor
    # dx
    # rmax
    L::EquivConv
    Q::EquivProd
    D1::LocalDense
    D::LocalDense
    function DensityPredictor(dx, rmax)
        L = EquivConv(Props((0,)), Props((0, 2),), dx, rmax) |> mygpu
        Q = EquivProd(Props((0, 2, 0, 1)); spectra = true)
        n = length(Q.out.ranks)
        D = LocalDense(4n, 1) |> mygpu
        D1 = LocalDense(n, 4n, leakyrelu) |> mygpu
        # D2 = LocalDense(2n, 4n, leakyrelu) |> mygpu
        # D3 = LocalDense(4n, 1) |> mygpu
        new(L, Q, D1, D)
    end
end
@functor DensityPredictor

function Flux.params(m::DensityPredictor)
    @unpack L, Q, D1, D = m
    [L, D1, D]
end

function f_(x, ψ, E,L, Q, D1, D)
    (D ∘ D1 ∘ Q)(cat(L(x), ψ, E, dims = 4))
end
function (m::DensityPredictor)(x, ψ, E)
    @unpack L, Q, D1, D = m

    sz = size(x)
    x0 = ψ0 = zeros(sz)
    E0 = zeros(sz[1:3]..., 3)
    normalize_density(abs.(f_(x, ψ, E,L, Q, D1, D) - f_(x0, ψ0, E0,L, Q, D1, D)), sum(x))
end

mutable struct Frame
    atoms
    positions
    grid
    ops
    ρp
    ρe
    cache
    more
    function Frame(atoms, positions, grid;ρe=nothing)
        @unpack dV, dx, sz = grid
        total = sum(atoms)
        ρp = zeros(sz..., 1)
        put_point_source!(
            ρp,
            grid,
            positions,
            reshape(atoms, (1, length(atoms))),
        )
        @assert total ≈ sum(ρp) * dV

        rmax = dx * norm(sz)
        rmin = 2dx
        name = :inverse_squared_field
        E = LinearOperator(name; dx, rmin, rmax)
        name = :potential
        ψ = LinearOperator(name; dx, rmin, rmax)

        ops = (; ψ, E)
        cache = ()
        more = ()
        # x = ρp |> mygpu
        # y = electronic_density |> mygpu
        new(
            atoms,
            positions,
            grid,
            ops,
            ρp,
            ρe,
            cache,
            more,
        )
    end
end

function (m::DensityPredictor)(x::Frame)
    @unpack ρp, ops,cache = x
    if isempty(cache)
    ψp = ops.ψ(ρp)
    Ep = ops.E(ρp)
    x.cache=(;ψp,Ep)
else
    @unpack ψp,Ep=cache
end
m(ρp, ψp, Ep)
    # m.ρe=m(ρp, ψp, Ep)
end

function calc!(x, p)
    @unpack ρe,ρp, grid, atoms, positions, ops = x
    if p == :forces
        # ψ=ops.ψ(ρp)
        E = ops.E(ρe)
        F = hcat([
            atom * get(E, grid, position)
            for (atom, position) in zip(atoms, eachcol(positions))
        ]...)
        more = (; E, F)
        x.more = merge(x.more, more)
        return F
    end
    # more
end
##
# function fg!(F, G, x)
#     electronic_density_hat, ρp = predict_density(x, atoms, grid)
#     ρhat = -electronic_density_hat + ρp
#     @show forces_hat = calcforces(x, atoms, ρhat, grid)
#     if G != nothing
#         G .= -forces_hat
#     end
#     norm(forces_hat) / sum(atoms)
# end;
#
# x0 = positions
#
# xres = optimize(
#     Optim.only_fg!(fg!),
#     x0,
#     LBFGS(),
#     Optim.Options(show_trace = true, f_tol = tol),
# )
# xmin = Optim.minimizer(xres)
# dmin = norm(xmin[:, 1] - xmin[:, 2])
# @printf "\nOptimal bond length for Ecut=%.2f: %.3f Bohr\n" Ecut dmin

max_stars_repo_name: aced-differentiate/equiv-docs

max_issues_repo_name: aced-differentiate/equiv-docs

max_forks_repo_name: aced-differentiate/equiv-docs


text:
export ContinuousSpace

struct ContinuousSpace{D,P,T<:AbstractFloat, F} <: AbstractSpace
    grid::GridSpace{D,P}
    update_vel!::F
    dims::NTuple{D,Int}
    spacing::T
    extent::NTuple{D,T}
end
Base.eltype(s::ContinuousSpace{D,P,T,F}) where {D,P,T,F} = T
defvel(a, m) = nothing

"""
    ContinuousSpace(extent::NTuple{D, <:Real}, spacing = min(extent)/10; kwargs...)
Create a `D`-dimensional `ContinuousSpace` in range 0 to (but not including) `extent`.
`spacing` configures the compartment spacing that the space is divided in, in order to
accelerate nearest neighbor functions like [`nearby_ids`](@ref).
All dimensions in `extent` must be completely divisible by `spacing` (i.e. no
fractional remainder).
Your agent positions (field `pos`) must be of type `NTuple{D, <:Real}`.

The keyword argument `update_vel!` is a **function**, `update_vel!(agent, model)` that updates
the agent's velocity **before** the agent has been moved, see [`move_agent!`](@ref).
You can of course change the agents' velocities
during the agent interaction, the `update_vel!` functionality targets spatial force
fields acting on the agents individually (e.g. some magnetic field).
By default no update is done this way.
If you use `update_vel!`, the agent type must have a field `vel::NTuple{D, <:Real}`.

The keyword `periodic = true` configures whether the space is periodic or not. If set to
`false` an error will occur if an agent's position exceeds the boundary.

There is no "best" choice for the value of `spacing`. If you need optimal performance it's
advised to set up a benchmark over a range of choices. The value matters most when searching
for neighbors. In [`Models.flocking`](@ref) for example, an optimal value for `spacing` is
66% of the search distance.

**Note:** if your model requires linear algebra operations for which tuples are not supported,
a performant solution is to convert between Tuple and SVector using
[StaticArrays.jl](https://github.com/JuliaArrays/StaticArrays.jl)
as follows: `s = SVector(t)` and back with `t = Tuple(s)`.
"""
function ContinuousSpace(
        extent::NTuple{D,X},
        spacing = min(extent)/10.0;
        update_vel! = defvel,
        periodic = true,
    ) where {D, X<:Real}
    @assert extent./spacing == floor.(extent./spacing) "All dimensions in `extent` must be completely divisible by `spacing`"
    s = GridSpace(floor.(Int, extent ./ spacing), periodic = periodic, metric = :euclidean)
    Z = X <: AbstractFloat ? X : Float64
    return ContinuousSpace(s, update_vel!, size(s), Z(spacing), Z.(extent))
end

function random_position(model::ABM{A,<:ContinuousSpace}) where {A}
    map(dim->rand()*dim, model.space.extent)
end

pos2cell(pos::Tuple, model::ABM) = floor.(Int, pos ./ model.space.spacing) .+ 1
pos2cell(a::AbstractAgent, model::ABM) = pos2cell(a.pos, model)
function cell_center(pos::NTuple{D, F}, model) where {D, F}
    ε = model.space.spacing
    (pos2cell(pos, model) .- 1) .* ε .+ ε/2
end
distance_from_cell_center(pos, model::ABM) =
distance_from_cell_center(pos, cell_center(pos, model))
function distance_from_cell_center(pos::Tuple, center::Tuple)
    sqrt(sum(abs2.(pos .- center)))
end

function add_agent_to_space!(
        a::A,
        model::ABM{A,<:ContinuousSpace},
    ) where {A<:AbstractAgent}
    push!(model.space.grid.s[pos2cell(a, model)...], a.id)
    return a
end

function remove_agent_from_space!(
        a::A,
        model::ABM{A,<:ContinuousSpace},
    ) where {A<:AbstractAgent}
    prev = model.space.grid.s[pos2cell(a, model)...]
    ai = findfirst(i -> i == a.id, prev)
    deleteat!(prev, ai)
    return a
end

function move_agent!(
        a::A,
        pos::Tuple,
        model::ABM{A,<:ContinuousSpace{D,periodic}},
    ) where {A<:AbstractAgent,D,periodic}
    remove_agent_from_space!(a, model)
    if periodic
        pos = mod.(pos, model.space.extent)
    end
    a.pos = pos
    add_agent_to_space!(a, model)
end

"""
    move_agent!(agent::A, model::ABM{A, ContinuousSpace}, dt::Real = 1.0)
Propagate the agent forwards one step according to its velocity,
_after_ updating the agent's velocity (if configured, see [`ContinuousSpace`](@ref)).
Also take care of periodic boundary conditions.

For this continuous space version of `move_agent!`, the "evolution algorithm"
is a trivial Euler scheme with `dt` the step size, i.e. the agent position is updated
as `agent.pos += agent.vel * dt`. If you want to move the agent to a specified position, do
`move_agent!(agent, pos, model)`.
"""
function move_agent!(
        agent::A,
        model::ABM{A,<:ContinuousSpace},
        dt::Real = 1.0,
    ) where {A<:AbstractAgent}
    model.space.update_vel!(agent, model)
    pos = agent.pos .+ dt .* agent.vel
    move_agent!(agent, pos, model)
    return agent.pos
end

#######################################################################################
# %% Neighbors and stuff
#######################################################################################
function nearby_ids(
        pos::ValidPos,
        model::ABM{<:AbstractAgent,<:ContinuousSpace{D,P,T}},
        r = 1;
        exact = false,
    ) where {D,P,T}
    if exact
        grid_r_max = r < model.space.spacing ? T(1) : r / model.space.spacing + T(1)
        grid_r_certain = grid_r_max - T(1.2)*sqrt(D)
        focal_cell = CartesianIndex(pos2cell(pos, model))
        allcells = grid_space_neighborhood(focal_cell, model, grid_r_max)
        if grid_r_max >= 1
            certain_cells = grid_space_neighborhood(focal_cell, model, grid_r_certain)
            certain_ids = Iterators.flatten(ids_in_position(cell, model) for cell in certain_cells)

            uncertain_cells = setdiff(allcells, certain_cells) # This allocates, but not sure if there's a better way.
            uncertain_ids = Iterators.flatten(ids_in_position(cell, model) for cell in uncertain_cells)

            additional_ids = Iterators.filter(i->edistance(pos, model[i].pos, model) ≤ r, uncertain_ids)

            return Iterators.flatten((certain_ids, additional_ids))
        else
            all_ids = Iterators.flatten(ids_in_position(cell, model) for cell in allcells)
            return Iterators.filter(i->edistance(pos, model[i].pos, model) ≤ r, all_ids)
        end
    else
        δ = distance_from_cell_center(pos, cell_center(pos, model))
        grid_r = (r+δ)/model.space.spacing
        return nearby_ids_cell(pos, model, grid_r)
    end
end

grid_space_neighborhood(α, model::ABM{<:AbstractAgent,<:ContinuousSpace}, r) =
grid_space_neighborhood(α, model.space.grid, r)

function nearby_ids_cell(
        pos::ValidPos,
        model::ABM{<:AbstractAgent,<:ContinuousSpace},
        r = 1,
    )
    nn = grid_space_neighborhood(CartesianIndex(pos2cell(pos, model)), model, r)
    s = model.space.grid.s
    Iterators.flatten((s[i...] for i in nn))
end

function nearby_positions(
        pos::ValidPos,
        model::ABM{<:AbstractAgent,<:ContinuousSpace},
        r = 1,
    )
    nn = grid_space_neighborhood(CartesianIndex(pos2cell(pos, model)), model, r)
    Iterators.filter(!isequal(pos), nn)
end

function positions(model::ABM{<:AbstractAgent,<:ContinuousSpace})
    x = CartesianIndices(model.space.grid.s)
    return (Tuple(y) for y in x)
end

function ids_in_position(pos::ValidPos, model::ABM{<:AbstractAgent,<:ContinuousSpace})
    return model.space.grid.s[pos...]
end

################################################################################
### Pretty printing
################################################################################
function Base.show(io::IO, space::ContinuousSpace{D,P}) where {D,P}
    s = "$(P ? "periodic" : "") continuous space with $(join(space.dims, "×")) divisions"
    space.update_vel! ≠ defvel && (s *= " with velocity updates")
    print(io, s)
end

#######################################################################################
# Continuous space exclusive
#######################################################################################
export nearest_neighbor, elastic_collision!, interacting_pairs

"""
    nearest_neighbor(agent, model::ABM{<: ContinuousSpace}, r) → nearest
Return the agent that has the closest distance to given `agent`.
Return `nothing` if no agent is within distance `r`.
"""
function nearest_neighbor(
        agent::A,
        model::ABM{A,<:ContinuousSpace},
        r;
        exact = false,
    ) where {A}
    n = collect(nearby_ids(agent, model, r; exact = exact))
    length(n) == 0 && return nothing
    d, j = Inf, 1
    for i in 1:length(n)
        @inbounds dnew = edistance(agent.pos, model[n[i]].pos, model)
        if dnew < d
            d, j = dnew, i
        end
    end
    return @inbounds model[n[j]]
end

using LinearAlgebra

"""
    elastic_collision!(a, b, f = nothing)
Resolve a (hypothetical) elastic collision between the two agents `a, b`.
They are assumed to be disks of equal size touching tangentially.
Their velocities (field `vel`) are adjusted for an elastic collision happening between them.
This function works only for two dimensions.
Notice that collision only happens if both disks face each other, to avoid
collision-after-collision.

If `f` is a `Symbol`, then the agent property `f`, e.g. `:mass`, is taken as a mass
to weight the two agents for the collision. By default no weighting happens.

One of the two agents can have infinite "mass", and then acts as an immovable object
that specularly reflects the other agent. In this case of course momentum is not
conserved, but kinetic energy is still conserved.
"""
function elastic_collision!(a, b, f = nothing)
    # Do elastic collision according to
    # https://en.wikipedia.org/wiki/Elastic_collision#Two-dimensional_collision_with_two_moving_objects
    v1, v2, x1, x2 = a.vel, b.vel, a.pos, b.pos
    length(v1) ≠ 2 && error("This function works only for two dimensions.")
    r1 = x1 .- x2
    r2 = x2 .- x1
    m1, m2 = f == nothing ? (1.0, 1.0) : (getfield(a, f), getfield(b, f))
    # mass weights
    m1 == m2 == Inf && return false
    if m1 == Inf
        @assert v1 == (0, 0) "An agent with ∞ mass cannot have nonzero velocity"
        dot(r1, v2) ≤ 0 && return false
        v1 = ntuple(x -> zero(eltype(v1)), length(v1))
        f1, f2 = 0.0, 2.0
    elseif m2 == Inf
        @assert v2 == (0, 0) "An agent with ∞ mass cannot have nonzero velocity"
        dot(r2, v1) ≤ 0 && return false
        v2 = ntuple(x -> zero(eltype(v1)), length(v1))
        f1, f2 = 2.0, 0.0
    else
        # Check if disks face each other, to avoid double collisions
        !(dot(r2, v1) > 0 && dot(r2, v1) > 0) && return false
        f1 = (2m2 / (m1 + m2))
        f2 = (2m1 / (m1 + m2))
    end
    ken = norm(v1)^2 + norm(v2)^2
    dx = a.pos .- b.pos
    dv = a.vel .- b.vel
    n = norm(dx)^2
    n == 0 && return false # do nothing if they are at the same position
    a.vel = v1 .- f1 .* (dot(v1 .- v2, r1) / n) .* (r1)
    b.vel = v2 .- f2 .* (dot(v2 .- v1, r2) / n) .* (r2)
    return true
end

"""
    interacting_pairs(model, r, method; scheduler = model.scheduler)
Return an iterator that yields unique pairs of agents `(a1, a2)` that are close
neighbors to each other, within some interaction radius `r`.

This function is usefully combined with `model_step!`, when one wants to perform
some pairwise interaction across all pairs of close agents once
(and does not want to trigger the event twice, both with `a1` and with `a2`, which
is unavoidable when using `agent_step!`).

The argument `method` provides three pairing scenarios
- `:all`: return every pair of agents that are within radius `r` of each other,
  not only the nearest ones.
- `:nearest`: agents are only paired with their true nearest neighbor
  (existing within radius `r`).
  Each agent can only belong to one pair, therefore if two agents share the same nearest
  neighbor only one of them (sorted by id) will be paired.
- `:scheduler`: agents are scanned according to the given keyword `scheduler`
  (by default the model's scheduler), and each scanned
  agent is paired to its nearest neighbor. Similar to `:nearest`, each agent can belong
  to only one pair. This functionality is useful e.g. when you want some agents to be
  paired "guaranteed", even if some other agents might be nearest to each other.
- `:types`: For mixed agent models only. Return every pair of agents within radius `r`
  (similar to `:all`), only capturing pairs of differing types. For example, a model of
  `Union{Sheep,Wolf}` will only return pairs of `(Sheep, Wolf)`. In the case of multiple
  agent types, *e.g.* `Union{Sheep, Wolf, Grass}`, skipping pairings that involve
  `Grass`, can be achived by a [`scheduler`](@ref Schedulers) that doesn't schedule `Grass`
  types, *i.e.*: `scheduler = [a.id for a in allagents(model) of !(a isa Grass)]`.
"""
function interacting_pairs(
        model::ABM{A,<:ContinuousSpace},
        r::Real,
        method;
        scheduler = model.scheduler,
        exact = true,
    ) where {A}
    @assert method ∈ (:scheduler, :nearest, :all, :types)
    pairs = Tuple{Int,Int}[]
    if method == :nearest
        true_pairs!(pairs, model, r)
    elseif method == :scheduler
        scheduler_pairs!(pairs, model, r, scheduler)
    elseif method == :all
        all_pairs!(pairs, model, r, exact = exact)
    elseif method == :types
        type_pairs!(pairs, model, r, scheduler, exact = exact)
    end
    return PairIterator(pairs, model.agents)
end

function scheduler_pairs!(
        pairs::Vector{Tuple{Int,Int}},
        model::ABM{A,<:ContinuousSpace},
        r::Real,
        scheduler,
    ) where {A}
    #TODO: This can be optimized further I assume
    for id in scheduler(model)
        # Skip already checked agents
        any(isequal(id), p[2] for p in pairs) && continue
        a1 = model[id]
        a2 = nearest_neighbor(a1, model, r)
        # This line ensures each neighbor exists in only one pair:
        if a2 ≠ nothing && !any(isequal(a2.id), p[2] for p in pairs)
            push!(pairs, (id, a2.id))
        end
    end
end

function all_pairs!(
        pairs::Vector{Tuple{Int,Int}},
        model::ABM{A,<:ContinuousSpace},
        r::Real;
        exact = true,
    ) where {A}
    for a in allagents(model)
        for nid in nearby_ids(a, model, r; exact = exact)
            # Sort the pair to overcome any uniqueness issues
            new_pair = isless(a.id, nid) ? (a.id, nid) : (nid, a.id)
            new_pair ∉ pairs && push!(pairs, new_pair)
        end
    end
end

function true_pairs!(
        pairs::Vector{Tuple{Int,Int}},
        model::ABM{A,<:ContinuousSpace},
        r::Real,
    ) where {A}
    distances = Vector{Float64}(undef, 0)
    for a in allagents(model)
        nn = nearest_neighbor(a, model, r)
        nn == nothing && continue
        # Sort the pair to overcome any uniqueness issues
        new_pair = isless(a.id, nn.id) ? (a.id, nn.id) : (nn.id, a.id)
        if new_pair ∉ pairs
            # We also need to check if our current pair is closer to each
            # other than any pair using our first id already in the list,
            # so we keep track of nn distances.
            dist = edistance(a.pos, nn.pos, model)

            idx = findfirst(x -> first(new_pair) == x, first.(pairs))
            if idx == nothing
                push!(pairs, new_pair)
                push!(distances, dist)
            elseif idx != nothing && distances[idx] > dist
                # Replace this pair, it is not the true neighbor
                pairs[idx] = new_pair
                distances[idx] = dist
            end
        end
    end
end

function type_pairs!(
        pairs::Vector{Tuple{Int,Int}},
        model::ABM{A,<:ContinuousSpace},
        r::Real,
        scheduler;
        exact = true,
    ) where {A}
    # We don't know ahead of time what types the scheduler will provide. Get a list.
    available_types = unique(typeof(model[id]) for id in scheduler(model))
    for id in scheduler(model)
        for nid in nearby_ids(model[id], model, r, exact = exact)
            neigbor_type = typeof(model[nid])
            if neigbor_type ∈ available_types && neigbor_type !== typeof(model[id])
                # Sort the pair to overcome any uniqueness issues
                new_pair = isless(id, nid) ? (id, nid) : (nid, id)
                new_pair ∉ pairs && push!(pairs, new_pair)
            end
        end
    end
end

struct PairIterator{A}
    pairs::Vector{Tuple{Int,Int}}
    agents::Dict{Int,A}
end

Base.length(iter::PairIterator) = length(iter.pairs)
function Base.iterate(iter::PairIterator, i = 1)
    i > length(iter) && return nothing
    p = iter.pairs[i]
    id1, id2 = p
    return (iter.agents[id1], iter.agents[id2]), i + 1
end