text:
# general.jl

export State_BNP_WMReg_Joint, init_state_BNP_WMReg_Joint,
    Prior_BNP_WMReg_Joint, Model_BNP_WMReg_Joint,
    Monitor_BNP_WMReg_Joint, Updatevars_BNP_WMReg_Joint,
    lNXmat, reset_adapt!,
    llik_BNP_WMReg_Joint;

mutable struct State_BNP_WMReg_Joint

    ### Parameters
    # component (kernel parameters) η
    μ_y::Array{Float64, 1}  # H vector
    β_y::Array{Float64, 2}  # H by K matrix
    δ_y::Array{Float64, 1}  # H vector
    μ_x::Array{Float64, 2}  # H by K matrix
    β_x::Union{Array{Array{Float64, 2}, 1}, Nothing}  # vector of H by (k in (K-1):1) matrices
    δ_x::Array{Float64,2}   # H by K matrix

    # variable selection
    γ::Union{BitArray{1}, BitArray{2}} # K vector (for global, H by K matrix for local)
    γδc::Union{Float64, Array{Float64, 1}, Nothing} # scale factors for δ_x
    π_γ::Array{Float64, 1} # K vector of probabilities that γ_k = 1 (true in code)
    ξ::BitArray{1} # mixture indicator for pi_gamma

    # allocation states, weights
    S::Array{Int, 1}        # n vector
    lω::Array{Float64, 1}   # H vector
    v::Array{Float64, 1}    # H-1 vector
    α::Float64

    # G0
    β0star_ηy::Array{Float64, 1}    # K+1 vector
    Λ0star_ηy::PDMat{Float64}       # K+1 by K+1 precision matrix

    ν_δy::Float64
    s0_δy::Float64

    μ0_μx::Array{Float64, 1}    # K vector
    Λ0_μx::PDMat{Float64}       # K by K precision matrix

    β0_βx::Union{Array{Array{Float64, 1}, 1}, Nothing} # vector of (k in (K-1):1) vectors
    Λ0_βx::Union{Array{PDMat{Float64}, 1}, Nothing}    # vector of k by k (for k in (K-1):1) precison matrices

    ν_δx::Array{Float64, 1}     # K vector
    s0_δx::Array{Float64, 1}    # K vector

    ### other state objects
    iter::Int

    accpt::Array{Int, 1} # H vector
    cSig_ηlδx::Array{PDMat{Float64}, 1}
    adapt::Bool
    adapt_iter::Union{Int, Nothing}
    adapt_thin::Union{Int, Nothing}
    runningsum_ηlδx::Union{Array{Float64, 2}, Nothing} # H by (K + K(K+1)/2) matrix (unless using Sigma_x_diag, in which case it has 2K columns)
    runningSS_ηlδx::Union{Array{Float64, 3}, Nothing}  # H by (K + K(K+1)/2) by (K + K(K+1)/2) matrix

    lNX::Array{Float64, 2} # n by H matrix of log(pdf(Normal(x_i))) under each obs i and allocation h
    lωNX_vec::Array{Float64, 1} # n vector with log( sum_j ωN(x_i) )

    n_occup::Int
    llik::Float64

end

### Outer constructors for State_BNP_WMReg_Joint
## for coninuing an adapt phase
function State_BNP_WMReg_Joint(μ_y, β_y, δ_y, μ_x, β_x, δ_x, γ, γδc, π_γ, ξ,
    S, lω::Vector{Float64}, α::Float64, β0star_ηy, Λ0star_ηy, ν_δy, s0_δy, μ0_μx, Λ0_μx,
    β0_βx, Λ0_βx, ν_δx, s0_δx,
    iter, accpt, cSig_ηlδx,
    adapt, adapt_iter, adapt_thin, runningsum_ηlδx, runningSS_ηlδx, lNX, lωNX_vec, llik)

    State_BNP_WMReg_Joint(deepcopy(μ_y), deepcopy(β_y), deepcopy(δ_y), deepcopy(μ_x),
        deepcopy(β_x), deepcopy(δ_x), deepcopy(γ), deepcopy(γδc), deepcopy(π_γ), deepcopy(ξ),
        deepcopy(S), deepcopy(lω), lω_to_v(lω), α, deepcopy(β0star_ηy),
        deepcopy(Λ0star_ηy), ν_δy, s0_δy, deepcopy(μ0_μx), deepcopy(Λ0_μx),
        deepcopy(β0_βx), deepcopy(Λ0_βx), deepcopy(ν_δx), deepcopy(s0_δx),
        iter, accpt, deepcopy(cSig_ηlδx),
        adapt, adapt_iter, adapt_thin, deepcopy(runningsum_ηlδx), deepcopy(runningSS_ηlδx),
        deepcopy(lNX), deepcopy(lωNX_vec), length(unique(S)), llik)
end
## same but includes v
function State_BNP_WMReg_Joint(μ_y, β_y, δ_y, μ_x, β_x, δ_x, γ, γδc, π_γ, ξ,
    S, lω::Vector{Float64}, v::Vector{Float64}, α::Float64, β0star_ηy, Λ0star_ηy, ν_δy, s0_δy, μ0_μx, Λ0_μx,
    β0_βx, Λ0_βx, ν_δx, s0_δx,
    iter, accpt, cSig_ηlδx,
    adapt, adapt_iter, adapt_thin, runningsum_ηlδx, runningSS_ηlδx, lNX, lωNX_vec, llik)

    State_BNP_WMReg_Joint(deepcopy(μ_y), deepcopy(β_y), deepcopy(δ_y), deepcopy(μ_x),
        deepcopy(β_x), deepcopy(δ_x), deepcopy(γ), deepcopy(γδc), deepcopy(π_γ), deepcopy(ξ),
        deepcopy(S), deepcopy(lω), v, α, deepcopy(β0star_ηy),
        deepcopy(Λ0star_ηy), ν_δy, s0_δy, deepcopy(μ0_μx), deepcopy(Λ0_μx),
        deepcopy(β0_βx), deepcopy(Λ0_βx), deepcopy(ν_δx), deepcopy(s0_δx),
        iter, accpt, deepcopy(cSig_ηlδx),
        adapt, adapt_iter, adapt_thin, deepcopy(runningsum_ηlδx), deepcopy(runningSS_ηlδx),
        deepcopy(lNX), deepcopy(lωNX_vec), length(unique(S)), llik)
end

## for starting new
function State_BNP_WMReg_Joint(μ_y, β_y, δ_y, μ_x, β_x, δ_x, γ, γδc, π_γ, ξ,
    S, lω::Vector{Float64}, α::Float64, β0star_ηy, Λ0star_ηy, ν_δy, s0_δy, μ0_μx, Λ0_μx,
    β0_βx, Λ0_βx, ν_δx, s0_δx,
    cSig_ηlδx, adapt)

    State_BNP_WMReg_Joint(deepcopy(μ_y), deepcopy(β_y), deepcopy(δ_y), deepcopy(μ_x),
        deepcopy(β_x), deepcopy(δ_x), deepcopy(γ), deepcopy(γδc), deepcopy(π_γ), deepcopy(ξ),
        deepcopy(S), deepcopy(lω), lω_to_v(lω), α, deepcopy(β0star_ηy), deepcopy(Λ0star_ηy),
        ν_δy, s0_δy, deepcopy(μ0_μx), deepcopy(Λ0_μx),
        deepcopy(β0_βx), deepcopy(Λ0_βx), deepcopy(ν_δx), deepcopy(s0_δx),
        0, zeros(Int, length(lω)), deepcopy(cSig_ηlδx),
        adapt, 0, 1,
        zeros( Float64, length(lω), Int(length(μ0_μx) + length(μ0_μx)*(length(μ0_μx)+1)/2) ),
        zeros( Float64, length(lω), Int(length(μ0_μx) + length(μ0_μx)*(length(μ0_μx)+1)/2),
                        Int(length(μ0_μx) + length(μ0_μx)*(length(μ0_μx)+1)/2) ),
        zeros(Float64, 1, 1), zeros(Float64, 1),
        length(unique(S)), 0.0 )
end

## for starting new ( given a value of v )
function State_BNP_WMReg_Joint(μ_y, β_y, δ_y, μ_x, β_x, δ_x, γ, γδc, π_γ, ξ,
    S, lω::Vector{Float64}, v::Vector{Float64}, α, β0star_ηy, Λ0star_ηy, ν_δy, s0_δy, μ0_μx, Λ0_μx,
    β0_βx, Λ0_βx, ν_δx, s0_δx,
    cSig_ηlδx, adapt)

    State_BNP_WMReg_Joint(deepcopy(μ_y), deepcopy(β_y), deepcopy(δ_y), deepcopy(μ_x),
        deepcopy(β_x), deepcopy(δ_x), deepcopy(γ), deepcopy(γδc), deepcopy(π_γ), deepcopy(ξ),
        deepcopy(S), deepcopy(lω), deepcopy(v), α, deepcopy(β0star_ηy), deepcopy(Λ0star_ηy),
        ν_δy, s0_δy, deepcopy(μ0_μx), deepcopy(Λ0_μx),
        deepcopy(β0_βx), deepcopy(Λ0_βx), deepcopy(ν_δx), deepcopy(s0_δx),
        0, zeros(Int, length(lω)), deepcopy(cSig_ηlδx),
        adapt, 0, 1,
        zeros( Float64, length(lω), Int(length(μ0_μx) + length(μ0_μx)*(length(μ0_μx)+1)/2) ),
        zeros( Float64, length(lω), Int(length(μ0_μx) + length(μ0_μx)*(length(μ0_μx)+1)/2),
                        Int(length(μ0_μx) + length(μ0_μx)*(length(μ0_μx)+1)/2) ),
        zeros(Float64, 1, 1), zeros(Float64, 1),
        length(unique(S)), 0.0 )
end

## for starting new but not adapting
function State_BNP_WMReg_Joint(μ_y, β_y, δ_y, μ_x, β_x, δ_x, γ, γδc, π_γ, ξ,
    S, lω::Vector{Float64}, α::Float64, β0star_ηy, Λ0star_ηy, ν_δy, s0_δy, μ0_μx, Λ0_μx,
    β0_βx, Λ0_βx, ν_δx, s0_δx,
    iter, accpt, cSig_ηlδx)

    State_BNP_WMReg_Joint(deepcopy(μ_y), deepcopy(β_y), deepcopy(δ_y), deepcopy(μ_x),
        deepcopy(β_x), deepcopy(δ_x), deepcopy(γ), deepcopy(γδc), deepcopy(π_γ), deepcopy(ξ),
        deepcopy(S), deepcopy(lω), lω_to_v(lω), α, deepcopy(β0star_ηy),
        deepcopy(Λ0star_ηy), ν_δy, s0_δy, deepcopy(μ0_μx), deepcopy(Λ0_μx),
        deepcopy(β0_βx), deepcopy(Λ0_βx), deepcopy(ν_δx), deepcopy(s0_δx),
        iter, accpt, deepcopy(cSig_ηlδx),
        false, nothing, nothing, nothing, nothing,
        zeros(Float64, 1, 1), zeros(Float64, 1),
        length(unique(S)), 0.0)
end

mutable struct Prior_BNP_WMReg_Joint
    α_sh::Float64   # gamma shape
    α_rate::Float64 # gamma rate

    π_sh::Array{Float64, 2} # Beta shape parameters
    π_ξ::Union{Float64, Array{Float64, 1}} # Mixture probability that pi_gamma comes from a beta (instead of point mass at 0)

    β0star_ηy_mean::Array{Float64, 1}   # MVN mean vector
    β0star_ηy_Cov::PDMat{Float64}       # MVN covariance matrix
    β0star_ηy_Prec::PDMat{Float64}      # MVN precision matrix

    Λ0star_ηy_df::Float64           # Wishart deg. of freedom
    Λ0star_ηy_S0::PDMat{Float64}    # Prior harmonic mean of Λ0inv; inverse scale of Wishart divided by deg. of freedom

    s0_δy_df::Float64   # scaled inv. chi-square deg. of freedom
    s0_δy_s0::Float64   # scaled inv. chi-square harmonic mean

    μ0_μx_mean::Array{Float64, 1}   # MVN mean vector
    μ0_μx_Cov::PDMat{Float64}       # MVN covariance matrix
    μ0_μx_Prec::PDMat{Float64}      # MVN precision matrix

    Λ0_μx_df::Float64           # Wishart deg. of freedom
    Λ0_μx_S0::PDMat{Float64}    # Prior harmonic mean of Λ0inv; inverse scale of Wishart divided by deg. of freedom

    β0_βx_mean::Union{Array{Array{Float64, 1}, 1}, Nothing} # vector of MVN mean vectors
    β0_βx_Cov::Union{Array{PDMat{Float64}, 1}, Nothing}     # vector of MVN covariance matrices
    β0_βx_Prec::Union{Array{PDMat{Float64}, 1}, Nothing}    # vector of MVN precision matrices

    Λ0_βx_df::Union{Array{Float64, 1}, Nothing}         # vector of Wishart deg. of freedom
    Λ0_βx_S0::Union{Array{PDMat{Float64}, 1}, Nothing}  # vector of Prior harmonic mean of Λ0inv; inverse scale of Wishart divided by deg. of freedom

    s0_δx_df::Array{Float64, 1}   # vector of scaled inv. chi-square deg. of freedom
    s0_δx_s0::Array{Float64, 1}   # vector of scaled inv. chi-square harmonic mean
end

### Outer constructors for Prior_BNP_WMReg_Joint
## automatic creation of precision matrices from covariance matrices
function Prior_BNP_WMReg_Joint(α_sh, α_rate, π_sh, π_ξ, β0star_ηy_mean, β0star_ηy_Cov,
    Λ0star_ηy_df, Λ0star_ηy_S0, s0_δy_df, s0_δy_s0,
    μ0_μx_mean, μ0_μx_Cov, Λ0_μx_df, Λ0_μx_S0,
    β0_βx_mean, β0_βx_Cov, Λ0_βx_df, Λ0_βx_S0,
    s0_δx_df, s0_δx_s0)

    Prior_BNP_WMReg_Joint(α_sh, α_rate, deepcopy(π_sh), deepcopy(π_ξ),
        deepcopy(β0star_ηy_mean),
        deepcopy(β0star_ηy_Cov), inv(β0star_ηy_Cov),
        deepcopy(Λ0star_ηy_df), deepcopy(Λ0star_ηy_S0), s0_δy_df, s0_δy_s0,
        deepcopy(μ0_μx_mean), deepcopy(μ0_μx_Cov), inv(μ0_μx_Cov), Λ0_μx_df,
        deepcopy(Λ0_μx_S0), deepcopy(β0_βx_mean), deepcopy(β0_βx_Cov),
        ( typeof(β0_βx_Cov) == Nothing ? nothing : [inv(β0_βx_Cov[k]) for k = 1:length(β0_βx_Cov)] ),
        deepcopy(Λ0_βx_df), deepcopy(Λ0_βx_S0),
        deepcopy(s0_δx_df), deepcopy(s0_δx_s0))
end

## default prior spec
function Prior_BNP_WMReg_Joint(K::Int, H::Int;
    center_y::Float64=0.0, center_X::Vector{Float64}=zeros(K),
    range_y::Float64=6.0, range_X::Vector{Float64}=fill(6.0, K), snr::Float64=5.0,
    Σx_type::Symbol=:full, γ_type::Symbol=:global) # the default 6.0 assumes data standardized

    s0_δy_s0 = (range_y/6.0)^2 / snr

    if Σx_type == :full
        s0_δx_s0_ranXfact = 8.0
        α_sh = 5.0
    elseif Σx_type == :diag
        s0_δx_s0_ranXfact = 12.0
        α_sh = 6.0
    end

    s0_δx_s0 = (range_X ./ s0_δx_s0_ranXfact ).^2

    Prior_BNP_WMReg_Joint(α_sh, # α_sh
    1.0, # α_rate
    hcat(fill(1.0, K), fill(0.5, K)), # π_sh
    fill(0.25, K), # π_ξ
    vcat(center_y, zeros(K)), # β0star_ηy_mean
    PDMat(Matrix(Diagonal(vcat((range_y/6.0)^2, fill(1.0, K))))), # β0star_ηy_Cov
    50.0*(K+1+2), # Λ0star_ηy_df # was multiplied by 100...
    PDMat(Matrix(Diagonal( vcat((range_y/2.0)^2, fill(16.0, K)) ./ s0_δy_s0 ))), # Λ0star_ηy_S0
    5.0, # s0_δy_df
    s0_δy_s0, # s0_δy_s0; the outer divide follows a SNR arguement
    center_X, # μ0_μx_mean
    PDMat(Matrix(Diagonal( (range_X ./ 6.0).^2 ))), # μ0_μx_Cov; needs to stay in close to center_X
    10.0*(K+2), # Λ0_μx_df; should be strong
    PDMat(Matrix(Diagonal( (range_X ./ 2.0).^2  ))), # Λ0_μx_S0; should give flexibility to mu_xs
    (K > 1 ? [zeros(k) for k = (K-1):-1:1] : nothing), # β0_βx_mean
    (K > 1 ? [ PDMat(Matrix(Diagonal(fill(1.0, k)))) for k = (K-1):-1:1 ] : nothing), # β0_βx_Cov
    (K > 1 ? fill(10.0*(K+2), K-1) : nothing), # Λ0_βx_df
    (K > 1 ? [ PDMat(Matrix(Diagonal(fill(2.0, k)))) for k = (K-1):-1:1 ] : nothing), # Λ0_βx_S0
    fill(5.0, K), # s0_δx_df
    s0_δx_s0 ) # s0_δx_s0
end

mutable struct Model_BNP_WMReg_Joint
    y::Array{Float64, 1}
    X::Array{Float64, 2}
    n::Int # length of data
    K::Int # number of predictors (columns in X)
    H::Int # DP truncation level
    prior::Prior_BNP_WMReg_Joint

    indx_ηy::Dict
    indx_ηx::Dict
    indx_β_x::Union{Dict, Nothing}

    Σx_type::Symbol # currently one of :full and :diag
    γ_type::Symbol # will one of :global, :local, or :fixed (for no variable selection)

    state::State_BNP_WMReg_Joint # this is the only thing that should change

    Model_BNP_WMReg_Joint(y, X, n, K, H, prior,
        indx_ηy, indx_ηx, indx_β_x, Σx_type, γ_type, state) = new(deepcopy(y), deepcopy(X), copy(n),
        deepcopy(K), deepcopy(H), deepcopy(prior),
        deepcopy(indx_ηy), deepcopy(indx_ηx), deepcopy(indx_β_x),
        deepcopy(Σx_type), deepcopy(γ_type),
        deepcopy(state))
end

function Model_BNP_WMReg_Joint(y::Array{Float64, 1}, X::Array{Float64, 2},
    H::Int, prior::Prior_BNP_WMReg_Joint, state::State_BNP_WMReg_Joint; Σx_type::Symbol=:full, γ_type::Symbol=:global)

    n = length(y)
    nx, K = size(X)

    n == nx || throw(error("X and y size mismatch."))

    indx_ηy = Dict( :μ => 1, :β => 2:(K+1), :δ => K+2 )

    if K > 1 && Σx_type == :full
        indx_ηx = Dict( :μ => 1:K, :β => (K+1):Int(K*(K+1)/2), :δ => Int(K*(K+1)/2 + 1):Int(K*(K+1)/2 + K) )
        indx_β_x = Dict()
        ii = 1
        jj = K - 1
        for k = 1:(K-1)
            indx_β_x[k] = Int(ii):Int(ii+jj-1)
            ii += jj
            jj -= 1
        end
    elseif K > 1 && Σx_type == :diag
        indx_ηx = Dict( :μ => 1:K, :δ => Int(K+1):Int(2*K) )
        indx_β_x = nothing
        state.β_x = nothing
        state.β0_βx = nothing
        state.Λ0_βx = nothing
        size(state.cSig_ηlδx[1], 1) == 2*K || throw("Initialize cSig_ηlδx with correct dimensions.")
        state.runningsum_ηlδx = zeros(Float64, H, 2*K)
        state.runningSS_ηlδx = zeros(Float64, H, 2*K, 2*K)
    else
        indx_ηx = Dict( :μ => 1:K, :δ => Int(K+1):Int(2*K) )
        indx_β_x = nothing
    end

    return Model_BNP_WMReg_Joint(deepcopy(y), deepcopy(X), n, K, deepcopy(H), deepcopy(prior),
        indx_ηy, indx_ηx, indx_β_x, Σx_type, γ_type, deepcopy(state))
end

mutable struct Monitor_BNP_WMReg_Joint
    ηlω::Bool
    γ::Bool
    S::Bool
    G0::Bool
end

mutable struct Updatevars_BNP_WMReg_Joint
    η::Bool
    lω::Bool
    γ::Bool
    α::Bool
    S::Bool
    G0::Bool
end

function postSimsInit_BNP_WMReg_Joint(m::Monitor_BNP_WMReg_Joint, n_keep::Int, init_state::State_BNP_WMReg_Joint)

    symb = [ :n_occup, :llik ]

    if m.ηlω
        push!(symb, [:μ_y, :β_y, :δ_y, :μ_x, :β_x, :δ_x, :lω, :α]...)
    end

    if m.γ
        push!(symb, [:γ, :π_γ, :ξ]...)
    end

    if m.S
        push!(symb, :S)
    end

    if m.G0
        push!(symb, [:β0star_ηy, :Λ0star_ηy, :s0_δy,
                     :μ0_μx, :Λ0_μx, :β0_βx, :Λ0_βx, :s0_δx]...)
    end

    sims = fill( BayesInference.deepcopyFields(init_state, symb), n_keep )

    return sims, symb
end

function lNXmat(X::Array{T, 2}, μ::Array{T, 2}, β::Array{Array{T, 2}, 1}, δ::Array{T, 2}) where T <: Real
    hcat( [lNX_sqfChol( Matrix(X'), μ[h,:], [ β[k][h,:] for k = 1:(size(X,2)-1) ], δ[h,:] )
            for h = 1:size(μ, 1) ]...)
end
function lNXmat(x::Union{Array{T,1}, Array{T,2}},
                μ::Union{Array{T,1}, Array{T,2}},
                δ::Union{Array{T,1}, Array{T,2}}) where T <: Real
    # K = 1 case

    x = vec(x)
    μ = vec(μ)
    δ = vec(δ)
    hcat( [ logpdf.(Normal(μ[h], sqrt(δ[h])), x) for h = 1:length(μ) ]...)
end

function lNXmat_Σdiag(X::Array{T,2}, μ::Array{T,2}, δ::Array{T,2}) where T <: Real
    n = size(X,1)
    H, KK = size(μ)
    out = zeros(n, H)

    for i = 1:n
        for h = 1:H
            for k = 1:KK
                out[i,h] += -0.5*log( 2.0π * δ[h,k] ) - 0.5*( X[i,k] - μ[h,k] )^2 / δ[h,k]
            end
        end
    end

    return out
end

function lNX_Σdiag(X::Union{Array{T, 1}, Array{T, 2}}, μ::Array{T, 1}, δ::Array{T, 1}) where T <: Real

    size(X,1) == length(μ) || throw("In lNX_sqfChol functions, the columns of X are the observations.")

    KK, n = size(X)

    out = zeros(n)

    for i = 1:n
        for k = 1:KK
            out[i] += -0.5*log( 2.0π * δ[k] ) - 0.5*( X[k,i] - μ[k] )^2 / δ[k]
        end
    end

    return out
end

function lωNXvec(lω::Array{T, 1}, lNX_mat::Array{T, 2}) where T <: Real
    lωNX_mat = broadcast(+, permutedims(lω), lNX_mat) # n by H
    vec( BayesInference.logsumexp(lωNX_mat, 2) ) # n vector
end

function reset_adapt!(model::Model_BNP_WMReg_Joint)
    model.state.adapt_iter = 0
    model.state.runningsum_ηlδx = zeros( Float64, size(model.state.runningsum_ηlδx) )
    model.state.runningSS_ηlδx = zeros( Float64, size(model.state.runningSS_ηlδx) )
    return nothing
end

function init_state_BNP_WMReg_Joint(n::Int, K::Int, H::Int,
    prior::Prior_BNP_WMReg_Joint; random::Int=1, Σx_type::Symbol=:full, γ_type::Symbol=:global)

    ν_δx = fill(5.0, K)
    ν_δy = 5.0

    γδc = Inf
    # γδc = fill(1.0e6, K)
    # γδc = nothing
    π_γ = fill(0.25, K)
    ξ = trues(K)

    if γ_type in (:fixed, :global)
        γ = trues(K) # always start with all variables
    elseif γ_type == :local
        γ = trues(H, K)
    end

    if random == 0
        s0_δx = [ prior.s0_δx_s0[k] for k = 1:K ]
        Λ0_βx = ( K > 1 ? [ inv(prior.Λ0_βx_S0[k]) for k = 1:(K-1) ] : nothing)
        β0_βx = ( K > 1 ? [ prior.β0_βx_mean[k] for k = 1:(K-1) ] : nothing)
        Λ0_μx = inv(prior.Λ0_μx_S0)
        μ0_μx = deepcopy(prior.μ0_μx_mean)
        s0_δy = deepcopy(prior.s0_δy_s0)
        Λ0star_ηy = inv(prior.Λ0star_ηy_S0)
        β0star_ηy = deepcopy(prior.β0star_ηy_mean)
        α = prior.α_sh / prior.α_rate
        lω = log.(fill(1.0 / H, H))
        v = lω_to_v(lω)
        S = [ sample(Weights(ones(H))) for i = 1:n ]
        δ_x = vcat([ deepcopy(s0_δx) for h = 1:H ]'...)
        β_x = ( K > 1 ? [ vcat( [deepcopy(β0_βx[k]) for h = 1:H]'... ) for k = 1:(K-1) ] : nothing)
        μ_x = vcat([ deepcopy(μ0_μx) for h = 1:H ]'...)
        δ_y = fill(s0_δy, H)
        β_y = vcat([ β0star_ηy[2:(K+1)] for h = 1:H ]'...)
        μ_y = fill(β0star_ηy[1], H)
    elseif random == 1 # G0 not random
        s0_δx = [ prior.s0_δx_s0[k] for k = 1:K ]
        Λ0_βx = ( K > 1 ? [ inv(prior.Λ0_βx_S0[k]) for k = 1:(K-1) ] : nothing)
        β0_βx = ( K > 1 ? [ prior.β0_βx_mean[k] for k = 1:(K-1) ] : nothing)
        Λ0_μx = inv(prior.Λ0_μx_S0)
        μ0_μx = deepcopy(prior.μ0_μx_mean)
        s0_δy = deepcopy(prior.s0_δy_s0)
        Λ0star_ηy = inv(prior.Λ0star_ηy_S0)
        β0star_ηy = deepcopy(prior.β0star_ηy_mean)
        α = rand(Gamma(prior.α_sh, 1.0/prior.α_rate))
        # lω = log.(fill(1.0 / H, H)) # start with equal weights
        lω, lv = rGenDirichlet(ones(H-1), fill(α, H-1); logout=true)
        v = exp.(lv)
        S = [ sample(Weights(ones(H))) for i = 1:n ] # random allocation
        δ_x = [ rand(InverseGamma(ν_δx[k]/2.0, ν_δx[k]*s0_δx[k]/2.0)) for h=1:H, k=1:K ]
        β_x = ( K > 1 ? [ vcat( [rand(MvNormal(β0_βx[k], inv(Λ0_βx[k]))) for h = 1:H]'... ) for k = 1:(K-1) ] : nothing)
        μ_x = vcat([ rand(MvNormal(μ0_μx, inv(Λ0_μx))) for h = 1:H ]'...)
        δ_y = rand( InverseGamma(ν_δy/2.0, ν_δy*s0_δy/2.0), H )
        β_y = vcat([ rand(MvNormal(β0star_ηy[2:(K+1)], δ_y[h]*inv(Λ0star_ηy).mat[2:(K+1), 2:(K+1)])) for h = 1:H ]'...)
        μ_y = [ rand( Normal(β0star_ηy[1], sqrt(δ_y[h]*inv(Λ0star_ηy).mat[1,1])) ) for h = 1:H ]
    elseif random == 2
        s0_δx = [ rand(Gamma(ν_δx[k]*prior.s0_δx_df[k]/2.0, 2.0*prior.s0_δx_s0[k]/(ν_δx[k]*prior.s0_δx_df[k]))) for k = 1:K ] # shape and scale
        Λ0_βx = ( K > 1 ? [ PDMat(rand(Wishart(prior.Λ0_βx_df[k], inv(prior.Λ0_βx_S0[k])/prior.Λ0_βx_df[k]))) for k = 1:(K-1) ] : nothing)
        β0_βx = ( K > 1 ? [ rand(MvNormal(prior.β0_βx_mean[k], prior.β0_βx_Cov[k])) for k = 1:(K-1) ] : nothing)
        Λ0_μx = PDMat( rand( Wishart(prior.Λ0_μx_df, inv(prior.Λ0_μx_S0)/prior.Λ0_μx_df) ) )
        μ0_μx = rand(MvNormal(prior.μ0_μx_mean, prior.μ0_μx_Cov))
        s0_δy = rand(Gamma(ν_δy*prior.s0_δy_df/2.0, 2.0*prior.s0_δy_s0/(ν_δy*prior.s0_δy_df))) # the prior here was actually gamma
        Λ0star_ηy = PDMat( rand( Wishart(prior.Λ0star_ηy_df, inv(prior.Λ0star_ηy_S0)/prior.Λ0star_ηy_df) ) )
        β0star_ηy = rand( MvNormal(prior.β0star_ηy_mean, prior.β0star_ηy_Cov) )
        α = rand(Gamma(prior.α_sh, 1.0/prior.α_rate))
        # lω = log.(fill(1.0 / H, H)) # start with equal weights
        lω, lv = rGenDirichlet(ones(H-1), fill(α, H-1); logout=true)
        v = exp.(lv)
        S = [ sample(Weights(ones(H))) for i = 1:n ] # random allocation
        δ_x = [ rand(InverseGamma(ν_δx[k]/2.0, ν_δx[k]*s0_δx[k]/2.0)) for h=1:H, k=1:K ]
        β_x = ( K > 1 ? [ vcat( [rand(MvNormal(β0_βx[k], inv(Λ0_βx[k]))) for h = 1:H]'... ) for k = 1:(K-1) ] : nothing)
        μ_x = vcat([ rand(MvNormal(μ0_μx, inv(Λ0_μx))) for h = 1:H ]'...)
        δ_y = rand( InverseGamma(ν_δy/2.0, ν_δy*s0_δy/2.0), H )
        β_y = vcat([ rand(MvNormal(β0star_ηy[2:(K+1)], δ_y[h]*inv(Λ0star_ηy).mat[2:(K+1), 2:(K+1)])) for h = 1:H ]'...)
        μ_y = [ rand( Normal(β0star_ηy[1], sqrt(δ_y[h]*inv(Λ0star_ηy).mat[1,1]))) for h = 1:H ]
    end

    if Σx_type == :full
        cSig_ηlδx = [ PDMat(Matrix(Diagonal(fill(0.1, Int(K+K*(K+1)/2))))) for h = 1:H ]
    elseif Σx_type == :diag
        cSig_ηlδx = [ PDMat(Matrix(Diagonal(fill(0.1, Int(2*K))))) for h = 1:H ]
        β_x = nothing
    end

    adapt = false

    State_BNP_WMReg_Joint(μ_y, β_y, δ_y, μ_x, β_x, δ_x, γ, γδc, π_γ, ξ,
        S, lω, v, α, β0star_ηy, Λ0star_ηy, ν_δy, s0_δy, μ0_μx, Λ0_μx,
        β0_βx, Λ0_βx, ν_δx, s0_δx,
        cSig_ηlδx, adapt)
end

function llik_BNP_WMReg_Joint(model::Model_BNP_WMReg_Joint)
    llik_num_mat = llik_numerator(model)
    return llik_BNP_WMReg_Joint(llik_num_mat, model.state.lωNX_vec)
end
function llik_BNP_WMReg_Joint(y::Array{T,1}, X::Array{T,2}, K::Int, H::Int,
    μ_y::Array{T, 1}, β_y::Array{T, 2}, δ_y::Array{T, 1},
    μ_x::Array{T, 2},
    β_x::Union{Array{Array{Float64, 2}, 1}, Nothing},
    δ_x::Array{Float64,2},
    γ::Union{BitArray{1}, BitArray{2}}, γδc::Union{Float64, Array{T, 1}},
    lω::Array{T, 1},
    lωNX_vec::Array{T, 1}) where T <: Real

    llik_num_mat = llik_numerator(y, X, K, H, μ_y, β_y, δ_y, μ_x, β_x, δ_x, γ, γδc, lω)
    return llik_BNP_WMReg_Joint(llik_num_mat, lωNX_vec)
end
function llik_BNP_WMReg_Joint(llik_numerator_mat::Array{T, 2},
    lωNX_vec::Array{T, 1}) where T <: Real

    llik_num_vec = vec( BayesInference.logsumexp(llik_numerator_mat, 2) ) # n vector
    llik_vec = llik_num_vec - lωNX_vec
    return sum(llik_vec)
end

"""
    rGenDirichlet(a, b[, logout])

  Single draw from the generalized Dirichlet distribution (Connor and Mosimann '69), option for log scale. From SparseProbVec package.

  ### Example
  ```julia
  a = ones(5)
  b = ones(5)
  rGenDirichlet(a, b)
  ```
"""
function rGenDirichlet(a::Vector{Float64}, b::Vector{Float64}; logout::Bool=false)
    n = length(a)
    K = n + 1
    length(b) == n || error("Dimension mismatch between a and b.")
    all(a .> 0.0) || error("All elements of a must be positive.")
    all(b .> 0.0) || error("All elements of b must be positive.")

    lz = Vector{Float64}(undef, n)
    loneminusz = Vector{Float64}(undef, n)
    for i = 1:n
        lx1 = log( rand( Distributions.Gamma( a[i] ) ) )
        lx2 = log( rand( Distributions.Gamma( b[i] ) ) )
        lxm = max(lx1, lx2)
        lxsum = lxm + log( exp(lx1 - lxm) + exp(lx2 - lxm) ) # logsumexp
        lz[i] = lx1 - lxsum
        loneminusz[i] = lx2 - lxsum
    end

    ## break the Stick
    lw = Vector{Float64}(undef, K)
    lwhatsleft = 0.0

    for i in 1:n
        lw[i] = lz[i] + lwhatsleft
        # lwhatsleft += log( 1.0 - exp(lw[i] - lwhatsleft) ) # logsumexp (not numerically stable)
        lwhatsleft += loneminusz[i]
    end
    lw[K] = copy(lwhatsleft)

    if logout
        out = (lw, lz)
    else
        out = (exp.(lw), exp.(lz))
    end
    return out
end


text:
using Random: seed!
cost(x) = x ^ 2
dcost(x) = 2 * x

seed!(1)
epochs = 2000

# momentum
x_momentum = rand(1)
opt = Momentum()

for i in 1:epochs
    g = dcost(x_momentum)
    δ = update(opt, g)
    global x_momentum -= δ
end

# adagrad
x_adagrad = rand(1)
opt = Adagrad()

for i in 1:epochs
    g = dcost(x_adagrad)

    δ = update(opt, g)
    global x_adagrad -= δ
end

# adadelta
x_adadelta = rand(1)
opt = Adadelta()

for i in 1:epochs
    g = dcost(x_adadelta)

    δ = update(opt, g)
    global x_adadelta -= δ
end

# rmsprop
x_rmsprop = rand(1)
opt = RMSprop()

for i in 1:epochs
    g = dcost(x_rmsprop)

    δ = update(opt, g)
    global x_rmsprop -= δ
end

# adam
x_adam = rand(1)
opt = Adam()

for i in 1:epochs
    g = dcost(x_adam)

    δ = update(opt, g)
    global x_adam -= δ
end

# adamax
x_adamax = rand(1)
opt = Adamax()

for i in 1:epochs
    g = dcost(x_adamax)

    δ = update(opt, g)
    global x_adamax -= δ
end

# nadam
x_nadam = rand(1)
opt = Nadam()

for i in 1:epochs
    g = dcost(x_nadam)

    δ = update(opt, g)
    global x_nadam -= δ
end

# vanilla
x_vanilla = rand(1)
opt = VanillaGradDescent()

for i in 1:epochs
    g = dcost(x_vanilla)

    δ = update(opt, g)
    global x_vanilla -= δ
end

# inversedecay
x_invdec = rand(1)
opt = Inversedecay()

for i in 1:epochs
    g = dcost(x_invdec)

    δ = update(opt, g)
    global x_invdec -= δ
end

@testset "Quadatric" begin
    @test 0.0 ≈ x_momentum[1] atol=5e-2
    @test 0.0 ≈ x_adagrad[1] atol=5e-2
    @test 0.0 ≈ x_adadelta[1] atol=5e-2
    @test 0.0 ≈ x_rmsprop[1] atol=5e-2
    @test 0.0 ≈ x_adam[1] atol=5e-2
    @test 0.0 ≈ x_adamax[1] atol=5e-2
    @test 0.0 ≈ x_nadam[1] atol=5e-2
    @test 0.0 ≈ x_vanilla[1] atol=5e-2
end

ext: jl


text:
# Test VectorTuples
@testset "VectorTuple" begin
    # Define useful data objects
    a = 42
    b = [4, -1, 2]
    c = [6 7; 5 1]
    c2 = JuMPC.DenseAxisArray(c, 3:4, [:q, :z])
    d = JuMPC.SparseAxisArray(Dict((2, 1) => 8, (2, 2) => 9, (1, 1) => -2))

    @testset "Definition" begin
        # Existence check
        @testset "DataType" begin
            @test IC.VectorTuple isa UnionAll
        end
        # _make_ordered (SparseAxisArray)
        @testset "_make_ordered (SparseAxisArray)" begin
            @test IC._make_ordered(d, IC.indices(d)) == [-2, 8, 9]
        end
        # _make_ordered (Fallback)
        @testset "_make_ordered (Fallback)" begin
            @test IC._make_ordered(a, nothing) == a
        end
        # VectorTuple constructor from Tuple
        @testset "Constructor (Tuple)" begin
            # test simple
            @test IC.VectorTuple((a,)) isa IC.VectorTuple
            @test IC.VectorTuple((a,)).values == [a]
            @test IC.VectorTuple((a,)).ranges == [1:1]
            @test IC.VectorTuple((a,)).indices == (nothing,)
            # test complicated
            tuple = (a, b, c, c2, d)
            @test sort(IC.VectorTuple(tuple).values) == sort([(tuple...)...])
            @test IC.VectorTuple(tuple).ranges == [1:1, 2:4, 5:8, 9:12, 13:15]
            @test IC.VectorTuple(tuple).indices == IC.indices.(tuple)
        end
        # VectorTuple constructor from args...
        @testset "Constructor (items...)" begin
            # test simple
            @test IC.VectorTuple(a) isa IC.VectorTuple
            @test IC.VectorTuple(a).values == [a]
            @test IC.VectorTuple(a).ranges == [1:1]
            @test IC.VectorTuple(a).indices == (nothing,)
            # test complicated
            tuple = (a, b, c, c2, d)
            @test sort(IC.VectorTuple(tuple...).values) == sort([(tuple...)...])
            @test IC.VectorTuple(tuple...).ranges == [1:1, 2:4, 5:8, 9:12, 13:15]
            @test IC.VectorTuple(tuple...).indices == IC.indices.(tuple)
        end
        # Base.:(==)
        @testset "Base.:(==)" begin
            @test IC.VectorTuple(a) == IC.VectorTuple(a)
            @test IC.VectorTuple(a, b, c, c2, d) == IC.VectorTuple(a, b, c, c2, d)
            @test IC.VectorTuple(a, b) != IC.VectorTuple(b, c)
        end
        # test same_structure
        @testset "same_structure" begin
            vt1 = IC.VectorTuple(a, b, c2)
            vt2 = IC.VectorTuple(a, b, c)
            vt3 = IC.VectorTuple(a, b, d)
            vt4 = IC.VectorTuple(a, b, b)
            vt5 = IC.VectorTuple('a', ['d', 't', 'q'], ['b' 'a'; 'd' 'f'])
            # test with indices
            @test IC.same_structure(vt2, vt5)
            @test !IC.same_structure(vt1, vt2)
            @test !IC.same_structure(vt2, vt3)
            @test !IC.same_structure(vt3, vt4)
            @test !IC.same_structure(vt1, vt5)
            # test without indices
            @test IC.same_structure(vt2, vt5, use_indices = false)
            @test IC.same_structure(vt1, vt2, use_indices = false)
            @test IC.same_structure(vt3, vt4, use_indices = false)
            @test !IC.same_structure(vt2, vt3, use_indices = false)
            @test !IC.same_structure(vt4, vt5, use_indices = false)
            @test IC.same_structure(vt1, vt5, use_indices = false)
        end
        # VectorTuple copying
        @testset "Base.copy" begin
            vt1 = IC.VectorTuple(a, b, c2)
            @test copy(vt1) == vt1
        end
    end

    @testset "Indexing" begin
        vt = IC.VectorTuple(b, a, c2, d)
        values = copy(vt.values)
        # test size
        @testset "Base.size" begin
            @test size(vt, 1) == 4
            @test_throws ArgumentError size(vt, 2)
        end
        # test firstindex
        @testset "Base.firstindex" begin
            @test firstindex(vt) == 1
            @test firstindex(vt, 1) == 1
            @test firstindex(vt, 2) == 1
        end
        # test lastindex
        @testset "Base.lastindex" begin
            @test lastindex(vt) == 11
            @test lastindex(vt, 1) == 4
            @test_throws ArgumentError lastindex(vt, 2)
        end
        # test getindex (linear indexing)
        @testset "Base.getindex (Linear)" begin
            # test scalar indexing
            @test vt[2] == -1
            @test vt[end] == 9
            # test slicing
            @test vt[1:4] == [b; a]
            @test vt[5:end] == [c2..., -2, 8, 9]
            @test vt[:] == vt.values
            # test out of bounds
            @test_throws BoundsError vt[1:end+1]
            @test_throws BoundsError vt[13]
        end
        # test getindex (tuple indexing)
        @testset "Base.getindex (Tuple)" begin
            # test scalar
            @test vt[1, 2] == b[2]
            @test vt[2, 1] == a
            @test vt[end, 3] == 9
            # test slicing over 2nd dimension
            @test vt[4, :] == [-2, 8, 9]
            @test vt[2, :] == [a]
            @test vt[1, 2:3] == b[2:3]
            @test vt[3, 1:4] == [c2...]
            # test slicing over first dimension
            @test vt[:, :] == [b, [a], [c2...], [-2, 8, 9]]
            @test vt[:, 1] == [first(b), a, first(c2), -2]
            @test vt[3:4, 2:3] == [c[2:3], [8, 9]]
            @test vt[[1, 4], 2:3] == [b[2:3], [8, 9]]
            # test out of bounds errors
            @test_throws BoundsError vt[:, 2]
            @test_throws BoundsError vt[5, 1]
            @test_throws BoundsError vt[2, 2]
            @test_throws BoundsError vt[3:4, 2:4]
        end
        # test setindex! (linear indexing)
        @testset "Base.setindex! (Linear)" begin
            # test scalar indexing
            @test (vt[2] = 1) == 1
            @test (vt[2] = -1) == -1
            @test (vt[end] = 2) == 2
            @test (vt[end] = 9) == 9
            # test slicing
            @test (vt[1:4] = collect(1:4)) == 1:4
            @test (vt[1:4] = [b; a]) == [b; a]
            @test (vt[6:7] = [3, 4]) == [3, 4]
            @test (vt[6:7] = c[2:3]) == c[2:3]
            # test everything reset
            @test vt.values == values
            # test out of bounds
            @test_throws BoundsError (vt[13] = 42)
            # test size mismatch
            @test_throws DimensionMismatch (vt[1:2] = [1])
        end
        # test setindex! (tuple indexing)
        @testset "Base.setindex! (Tuple)" begin
            # test scalar
            @test (vt[1, 2] = 42) == 42
            @test (vt[1, 2] = b[2]) == b[2]
            @test (vt[end, 3] = 3) == 3
            @test (vt[end, 3] = 9) == 9
            # test slicing over 2nd dimension
            @test (vt[4, :] = [1, 2, 5]) == [1, 2, 5]
            @test (vt[4, :] = [-2, 8, 9]) == [-2, 8, 9]
            @test (vt[1, 2:3] = [2, 6]) == [2, 6]
            @test (vt[1, 2:3] = b[2:3]) == b[2:3]
            # test slicing over first dimension
            @test (vt[:, 1] = [1, 3, 4, 5]) == [1, 3, 4, 5]
            @test (vt[:, 1] = [first(b), a, first(c2), -2]) == [first(b), a, first(c2), -2]
            @test (vt[3:4, 2:3] = [[1, 1], [2, 2]]) == [[1, 1], [2, 2]]
            @test (vt[3:4, 2:3] = [c[2:3], [8, 9]]) == [c[2:3], [8, 9]]
            @test (vt[[1, 4], 2:3] = [[1, 1], [2, 2]]) == [[1, 1], [2, 2]]
            @test (vt[[1, 4], 2:3] = [b[2:3], [8, 9]]) == [b[2:3], [8, 9]]
            # test more colon combos
            vt2 = IC.VectorTuple(b, d)
            values2 = copy(vt2.values)
            @test (vt2[:, 2:3] = [[1, 1], [3, 3]]) == [[1, 1], [3, 3]]
            @test (vt2[:, 2:3] = [b[2:3], [8, 9]]) == [b[2:3], [8, 9]]
            @test (vt2[:, :] = [[1, 1, 1], [3, 3, 3]]) == [[1, 1, 1], [3, 3, 3]]
            @test (vt2[:, :] = [b, [-2, 8, 9]]) == [b, [-2, 8, 9]]
            # test everything reset
            @test vt.values == values
            @test vt2.values == values2
            # test out of bounds errors
            @test_throws BoundsError (vt[2, 2] = 42)
            @test_throws BoundsError (vt[4, 1:4] = [1, 2, 3, 4])
            @test_throws BoundsError (vt[1:2, 1:2] = [[1, 1], [2, 2]])
            # test dimension errors
            @test_throws AssertionError (vt[[1, 3], 1] = 1)
            @test_throws DimensionMismatch (vt[3, 2:3] = [1, 4, 2])
            @test_throws DimensionMismatch (vt[:, 1] = [1, 4, 2])
        end
    end

    @testset "Basic Queries" begin
        vt = IC.VectorTuple(a, b, c2)
        # test length
        @testset "Base.length" begin
            @test length(vt) == length(a) + length(b) + length(c2)
        end
        # test isempty
        @testset "Base.isempty" begin
            @test !isempty(vt)
            @test isempty(IC.VectorTuple())
        end
        # test eachindex
        @testset "Base.eachindex" begin
            @test eachindex(vt) == eachindex(vt.values)
        end
        # test keys
        @testset "Base.keys" begin
            @test keys(vt) == keys(vt.values)
        end
        # test findfirst
        @testset "Base.findfirst" begin
            @test findfirst(isequal(-1), vt) == 3
        end
        # test findall
        @testset "Base.findall" begin
            @test findall(isequal(42), vt) == [1]
        end
        # test in
        @testset "Base.in" begin
            @test -1 in vt
            @test !(92 in vt)
        end
        # test iterate
        @testset "Base.iterate" begin
            @test [i for i in vt] == vt.values
        end
    end

    @testset "Copy Modification" begin
        # test _update_indices (SparseAxisArray)
        @testset "_update_indices (SparseAxisArray)" begin
            inds = IC.indices(d)
            new_axes = [(1, 1), (2, 2)]
            @test IC._update_indices(inds, 2).axes == new_axes
            inds = IC.ContainerIndices(CartesianIndices(new_axes), new_axes)
            @test IC._update_indices(inds, 2) isa Nothing
        end
        # test _update_indices (DenseAxisArray)
        @testset "_update_indices (DenseAxisArray)" begin
            inds = IC.indices(c2)
            @test IC._update_indices(inds, 2).axes == [(3, :q), (3, :z), (4, :z)]
        end
        # test _update_indices (Array)
        @testset "_update_indices (Array)" begin
            inds = IC.indices(c)
            @test IC._update_indices(inds, 2).axes == [(1, 1), (1, 2), (2, 2)]
        end
        # test restricted_copy
        @testset "restricted_copy" begin
            # test with linear indexing with simple results
            vt = IC.VectorTuple(b, a)
            @test isempty(IC.restricted_copy(vt, [false, false, false, false]))
            vt = IC.VectorTuple(b, a)
            @test IC.restricted_copy(vt, [false, false, false, true]) == IC.VectorTuple(a)
            vt = IC.VectorTuple(c2, a)
            @test IC.restricted_copy(vt, [true, true, true, true, false]) == IC.VectorTuple(c2)
            # test with linear indexing and array reduction
            vt = IC.VectorTuple(d, a, c)
            new_d = JuMPC.SparseAxisArray(delete!(copy(d.data), (1, 1)))
            inds = CartesianIndices(c)
            cs = JuMPC.SparseAxisArray(Dict(Tuple(inds[k]) => c[inds[k]] for k = [1, 4]))
            copy_inds = [false, true, true, true, true, false, false, true]
            @test IC.restricted_copy(vt, copy_inds) == IC.VectorTuple(new_d, a, cs)
            @test vt == IC.VectorTuple(d, a, c)
            # test with linear indexing and array reduction again
            vt = IC.VectorTuple(a, c, d)
            new_d = JuMPC.SparseAxisArray(delete!(copy(d.data), (2, 1)))
            inds = CartesianIndices(c)
            cs = JuMPC.SparseAxisArray(Dict(Tuple(inds[k]) => c[inds[k]] for k = [2, 4]))
            copy_inds = [false, false, true, false, true, true, false, true]
            @test IC.restricted_copy(vt, copy_inds) == IC.VectorTuple(cs, new_d)
            @test vt == IC.VectorTuple(a, c, d)
        end
    end

    @testset "Tuple Creation" begin
        # test _make_array (Array)
        @testset "_make_array (Array)" begin
            vt = IC.VectorTuple(b, a, c)
            @test IC._make_array(vt.values[1:3], vt.indices[1]) == b
            @test IC._make_array(vt.values[5:8], vt.indices[3]) == c
        end
        # test _make_array (DenseAxisArray)
        @testset "_make_array (DenseAxisArray)" begin
            vt = IC.VectorTuple(c2, a)
            @test IC._make_array(vt.values[1:4], vt.indices[1]) == c2
        end
        # test _make_array (SparseAxisArray)
        @testset "_make_array (SparseAxisArray)" begin
            vt = IC.VectorTuple(a, d)
            @test IC._make_array(vt.values[2:4], vt.indices[2]) == d
        end
        # test _make_array (Fallback)
        @testset "_make_array (Fallback)" begin
            vt = IC.VectorTuple(a, d)
            @test IC._make_array(vt.values[1:1], vt.indices[1]) == a
            vt = IC.VectorTuple(b, (1, 2))
            @test IC._make_array(vt.values[4:5], vt.indices[2]) == [1, 2]
        end
        # test Tuple
        @testset "Base.Tuple" begin
            # test reproducing original
            vt = IC.VectorTuple(a, d, c2, b)
            @test Tuple(vt) == (a, d, c2, b)
            # test vector mode
            vt = IC.VectorTuple(a, c)
            @test Tuple(vt, use_indices = false) == (a, [c...])
        end
        # test Tuple
        @testset "Base.Tuple (other values)" begin
            # test reproducing original
            vt = IC.VectorTuple(a, c, b)
            values = [:a for i in 1:length(vt)]
            @test Tuple(values, vt) == (:a, [:a :a; :a :a], [:a, :a, :a])
            # test vector mode
            vt = IC.VectorTuple(a, c)
            values = [:a for i in 1:length(vt)]
            @test Tuple(values, vt, use_indices = false) == (:a, [:a, :a, :a, :a])
        end
    end

    @testset "Printing Methods" begin
        vt = IC.VectorTuple(a, c)
        @testset "Base.string" begin
            @test string(vt) == "(42, [6, 5, 7, 1])"
        end
        # test Base.show (REPL)
        @testset "Base.show (REPL)" begin
            show_test(REPLMode, vt, "(42, [6, 5, 7, 1])")
        end
        # test Base.show (IJulia)
        @testset "Base.show (IJulia)" begin
            show_test(IJuliaMode, vt, "(42, [6, 5, 7, 1])")
        end
    end
end


text:
function matchPermutedCanonical(input_types::Vector{Type}, outbound_type::Type)
    # TODO: this implementation only works when the inbound types match the outbound type
    nothing_inputs = 0
    message_inputs = 0
    for input_type in input_types
        if input_type == Nothing
            nothing_inputs += 1
        elseif matches(input_type, outbound_type)
            message_inputs += 1
        end
    end

    return (nothing_inputs == 1) && (message_inputs == 2)
end

mutable struct SPEqualityGaussian <: SumProductRule{Equality} end
outboundType(::Type{SPEqualityGaussian}) = Message{GaussianWeightedMeanPrecision}
isApplicable(::Type{SPEqualityGaussian}, input_types::Vector{Type}) = matchPermutedCanonical(input_types, Message{Gaussian})

mutable struct SPEqualityGammaWishart <: SumProductRule{Equality} end
outboundType(::Type{SPEqualityGammaWishart}) = Message{Union{Gamma, Wishart}}
isApplicable(::Type{SPEqualityGammaWishart}, input_types::Vector{Type}) = matchPermutedCanonical(input_types, Message{Union{Gamma, Wishart}})

mutable struct SPEqualityBernoulli <: SumProductRule{Equality} end
outboundType(::Type{SPEqualityBernoulli}) = Message{Bernoulli}
isApplicable(::Type{SPEqualityBernoulli}, input_types::Vector{Type}) = matchPermutedCanonical(input_types, Message{Bernoulli})

mutable struct SPEqualityBeta <: SumProductRule{Equality} end
outboundType(::Type{SPEqualityBeta}) = Message{Beta}
isApplicable(::Type{SPEqualityBeta}, input_types::Vector{Type}) = matchPermutedCanonical(input_types, Message{Beta})

mutable struct SPEqualityCategorical <: SumProductRule{Equality} end
outboundType(::Type{SPEqualityCategorical}) = Message{Categorical}
isApplicable(::Type{SPEqualityCategorical}, input_types::Vector{Type}) = matchPermutedCanonical(input_types, Message{Categorical})

mutable struct SPEqualityDirichlet <: SumProductRule{Equality} end
outboundType(::Type{SPEqualityDirichlet}) = Message{Dirichlet}
isApplicable(::Type{SPEqualityDirichlet}, input_types::Vector{Type}) = matchPermutedCanonical(input_types, Message{Dirichlet})

mutable struct SPEqualityPointMass <: SumProductRule{Equality} end
outboundType(::Type{SPEqualityPointMass}) = Message{PointMass}
function isApplicable(::Type{SPEqualityPointMass}, input_types::Vector{Type})
    nothing_inputs = 0
    soft_inputs = 0
    point_mass_inputs = 0
    for input_type in input_types
        if input_type == Nothing
            nothing_inputs += 1
        elseif matches(input_type, Message{SoftFactor})
            soft_inputs += 1
        elseif matches(input_type, Message{PointMass})
            point_mass_inputs += 1
        end
    end

    return (nothing_inputs == 1) && (soft_inputs == 1) && (point_mass_inputs == 1)
end

mutable struct SPEqualityFn <: SumProductRule{Equality} end
outboundType(::Type{SPEqualityFn}) = Message{Function}
isApplicable(::Type{SPEqualityFn}, input_types::Vector{Type}) = matchPermutedCanonical(input_types, Message{Function})

mutable struct SPEqualityFnG <: SumProductRule{Equality} end
outboundType(::Type{SPEqualityFnG}) = Message{GaussianMeanVariance}
function isApplicable(::Type{SPEqualityFnG}, input_types::Vector{Type})
    nothing_inputs = 0
    function_inputs = 0
    gaussian_inputs = 0

    for input_type in input_types
        if input_type == Nothing
            nothing_inputs += 1
        elseif matches(input_type, Message{Function})
            function_inputs += 1
        elseif matches(input_type, Message{Gaussian})
            gaussian_inputs += 1
        end
    end

    return (nothing_inputs == 1) && (function_inputs == 1) && (gaussian_inputs == 1)
end

mutable struct SPEqualityFnFactor <: SumProductRule{Equality} end
outboundType(::Type{SPEqualityFnFactor}) = Message{SampleList}
function isApplicable(::Type{SPEqualityFnFactor}, input_types::Vector{Type})
    nothing_inputs = 0
    function_inputs = 0
    factor_inputs = 0

    for input_type in input_types
        if input_type == Nothing
            nothing_inputs += 1
        elseif matches(input_type, Message{Function})
            function_inputs += 1
        elseif matches(input_type, Message{Gaussian})
            return false
        elseif matches(input_type, Message{FactorNode})
            factor_inputs += 1
        end
    end

    return (nothing_inputs == 1) && (function_inputs == 1) && (factor_inputs == 1)
end

mutable struct SPEqualityGFactor <: SumProductRule{Equality} end
outboundType(::Type{SPEqualityGFactor}) = Message{GaussianMeanVariance}
function isApplicable(::Type{SPEqualityGFactor}, input_types::Vector{Type})
    nothing_inputs = 0
    factor_inputs = 0
    gaussian_inputs = 0

    for input_type in input_types
        if input_type == Nothing
            nothing_inputs += 1
        elseif matches(input_type, Message{PointMass})
            return false
        elseif matches(input_type, Message{Gaussian})
            gaussian_inputs += 1
        elseif matches(input_type, Message{FactorNode})
            factor_inputs += 1
        end
    end

    return (nothing_inputs == 1) && (factor_inputs == 1) && (gaussian_inputs == 1)
end

mutable struct SPEqualityFactor <: SumProductRule{Equality} end
outboundType(::Type{SPEqualityFactor}) = Message{SampleList}
function isApplicable(::Type{SPEqualityFactor}, input_types::Vector{Type})
    nothing_inputs = 0
    first_type, second_type = nothing, nothing

    for input_type in input_types
        if input_type == Nothing
            nothing_inputs += 1
        elseif matches(input_type, Message{PointMass})
            return false
        elseif matches(input_type, Message{Gaussian})
            return false
        elseif matches(input_type, Message{Function})
            return false
        elseif matches(input_type, Message{PointMass})
            return false
        elseif matches(input_type, Message{FactorNode})
            if first_type === nothing
                first_type = input_type
            else
                second_type = input_type
            end
        end
    end

    return (nothing_inputs == 1) && !matches(first_type, second_type) && !matches(second_type, first_type)
end


text:
### A Pluto.jl notebook ###
# v0.14.5

using Markdown
using InteractiveUtils

# ╔═╡ 831a3d08-b4f3-11eb-101b-a19a95d2b62e
begin
	using CSV
	using DataFrames
	using LightGraphs
	using GLM
	using GraphPlot
	using ParserCombinator.Parsers.DOT: nodes, edges, parse_dot, StringID
    using Statistics
end

# ╔═╡ 1d7b6d76-8974-4480-942d-35e895ac0c69

# ╔═╡ cb7406e4-ae75-4c99-b7d5-842055b52aee
df = CSV.read("./data/sodium.csv", DataFrame)
df = select!(df, Not(:Column1))   # drop index column

# ╔═╡ 3c9d639a-dea3-40f5-bec0-be46ecf332e3
begin
    ols = lm(@formula(sbp ~ sodium + age + proteinuria), df)

    xt = copy(df[:,Not(:sbp)])
    xt₁ = copy(xt)
    xt₁.sodium .= 1
    xt₂ = copy(xt)
    xt₂.sodium .= 0

    ate_est = mean(predict(ols, xt₁) - predict(ols, xt₂))
    println("ATE estimate: $ate_est")
end

begin
    ols_causal = lm(@formula(sbp ~ sodium + age), df)

    xt = copy(df[:,Not([:sbp, :proteinuria])])
    xt₁ = copy(xt)
    xt₁.sodium .= 1
    xt₂ = copy(xt)
    xt₂.sodium .= 0

    causal_ate_est = mean(predict(ols_causal, xt₁) - predict(ols_causal, xt₂))
    println("causal ATE estimate: $causal_ate_est")
end

# ╔═╡ 48203bec-8156-489a-b04d-9a613526c85e
function _dot_read_one_graph(pg)
    isdir = pg.directed
    nvg = length(nodes(pg))
    nodedict = Dict(zip(collect(nodes(pg)), 1:nvg))
    if isdir
        g = LightGraphs.DiGraph(nvg)
    else
        g = LightGraphs.Graph(nvg)
    end
    for es in edges(pg)
        s = nodedict[es[1]]
        d = nodedict[es[2]]
        add_edge!(g, s, d)
    end
    return g, collect(nodes(pg))
end

# ╔═╡ b1abedc3-647e-4f64-ad78-c4b0eeff8510
function loaddot(gdot::String, gname::String)
    p = parse_dot(gdot)
    for pg in p
        isdir = pg.directed
        possname = isdir ? StringID("digraph") : StringID("graph")
        name = get(pg.id, possname).id
        name == gname && return _dot_read_one_graph(pg)
    end
    error("Graph $gname not found")
end

# ╔═╡ 7e9da9c2-ac34-4ec9-a6ad-f81d67435c5b
gio = "digraph graphname {
    age -> sodium;
	age -> sbp;
	sodium -> proteinuria;
	sodium -> sbp;
    sbp -> proteinuria;
}"

# ╔═╡ abb14b1e-75fb-4a3f-9973-c6ec4fdb2722

# ╔═╡ 8302c22b-816c-43cf-9d06-60a9daaa56fb
graph, nodelabels = loaddot(gio, "graphname")

# ╔═╡ 4ae575f1-a224-4b98-8142-207a1c9b0094
gplot(graph, layout=circular_layout, nodelabel=nodelabels)

# ╔═╡ ae162157-5768-4935-a1ee-b720fb2d1dc7
begin
	Xt = df[:, ["sodium", "age", "proteinuria"]]
	y = df[!, "sbp"]
end

# ╔═╡ 2aa55dbe-4776-46df-9107-8bd64a545ee4

# ╔═╡ f7a19ddb-391e-49e0-9912-990edc76e92b

# ╔═╡ 41ed3b5e-cdb5-4478-86ce-138f08c09784

# ╔═╡ Cell order:
# ╠═831a3d08-b4f3-11eb-101b-a19a95d2b62e
# ╠═1d7b6d76-8974-4480-942d-35e895ac0c69
# ╠═cb7406e4-ae75-4c99-b7d5-842055b52aee
# ╠═3c9d639a-dea3-40f5-bec0-be46ecf332e3
# ╠═b1abedc3-647e-4f64-ad78-c4b0eeff8510
# ╠═48203bec-8156-489a-b04d-9a613526c85e
# ╠═7e9da9c2-ac34-4ec9-a6ad-f81d67435c5b
# ╠═abb14b1e-75fb-4a3f-9973-c6ec4fdb2722
# ╠═8302c22b-816c-43cf-9d06-60a9daaa56fb
# ╠═4ae575f1-a224-4b98-8142-207a1c9b0094
# ╠═ae162157-5768-4935-a1ee-b720fb2d1dc7
# ╠═2aa55dbe-4776-46df-9107-8bd64a545ee4
# ╠═f7a19ddb-391e-49e0-9912-990edc76e92b
# ╠═41ed3b5e-cdb5-4478-86ce-138f08c09784

max_stars_repo_name: Magier/Aetia

max_issues_repo_name: Magier/Aetia

max_forks_repo_name: Magier/Aetia


text:
using nanoJulia
using Test

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


text:
# struct Bounds{L,U} end
# Bounds() = Bounds{-Inf,Inf}()
# Base.min(::Bounds{L}) where {L} = L
# Base.max(::Bounds{L,U}) where {L,U} = U

struct RealScalar{L,U} <: AbstractParameter
    RealScalar{L,U}() where {L,U} = new{Float64(L),Float64(U)}()
end
struct RealArray{S<:Tuple,L,U,D} <: AbstractParameter
    dynamicsizes::NTuple{D,Int}
    RealArray{S,L,U,D}(d::NTuple{D,Int}) where {S,L,U,D} = ((@assert U > L); new{S,Float64(L),Float64(U),D}(d))
end
const RealVector{N,L,U,D} = RealArray{Tuple{N},L,U,D}
const RealMatrix{M,N,L,U,D} = RealArray{Tuple{M,N},L,U,D}
RealScalar() = RealScalar{-Inf,Inf}()
RealScalar{L}() where {L} = RealScalar{Float64(L),Inf}()
RealArray{S}() where {S} = RealArray{S,-Inf,Inf,0}(tuple())
RealArray{S,L}() where {S,L} = RealArray{S,Float64(L),Inf,0}(tuple())
RealArray{S,L,U}() where {S,L,U} = RealArray{S,Float64(L),Float64(U),0}(tuple())
RealArray{S,L,U,0}() where {S,L,U} = RealArray{S,Float64(L),Float64(U),0}(tuple())
RealArray{S}(ds::NTuple{D}) where {S<:Tuple,D} = RealArray{S,-Inf,Inf,D}(ds)
RealArray{S,L}(ds::NTuple{D}) where {S<:Tuple,L,D} = RealArray{S,Float64(L),Inf,D}(ds)
RealArray{S,L,U}(ds::NTuple{D}) where {S<:Tuple,L,U,D} = RealArray{S,Float64(L),Float64(U),D}(ds)
RealArray(ds::NTuple{1,Int}) = RealArray{Tuple{-1},-Inf,Inf,1}(ds)
RealArray(ds::NTuple{2,Int}) = RealArray{Tuple{-1,-1},-Inf,Inf,2}(ds)
RealArray{L}(ds::NTuple{1,Int}) where {L} = RealArray{Tuple{-1},Float64(L),Inf,1}(ds)
RealArray{L}(ds::NTuple{2,Int}) where {L} = RealArray{Tuple{-1,-1},Float64(L),Inf,2}(ds)
RealArray{L,U}(ds::NTuple{1,Int}) where {L,U} = RealArray{Tuple{-1},Float64(L),Float64(U),1}(ds)
RealArray{L,U}(ds::NTuple{2,Int}) where {L,U} = RealArray{Tuple{-1,-1},Float64(L),Float64(U),2}(ds)

using PaddedMatrices
import PaddedMatrices: memory_length
memory_length(::RealScalar) = 1
memory_length(::RealVector{M}) where {M} = M
memory_length(::RealMatrix{M,N}) where {M,N} = M*N
memory_length(::Type{<:RealScalar}) = 1
memory_length(::Type{<:RealVector{M}}) where {M} = M
memory_length(::Type{<:RealMatrix{M,N}}) where {M,N} = M*N
@generated PaddedMatrices.memory_length(::RealArray{S}) where {S} = PaddedMatrices.simple_vec_prod(S.parameters)
@generated PaddedMatrices.memory_length(::Type{<:RealArray{S}}) where {S} = PaddedMatrices.simple_vec_prod(S.parameters)

abstract type AbstractMissingDataArray{M,L,U,T,N,A <: AbstractArray{T,N}} <: AbstractParameter end
memory_length(::AbstractMissingDataArray{M}) where {M} = M
"""
Parameters are:
M: number Missing
B: Bounds
T: DataType
N: Dimensionality of Array (how many axis?)
"""
struct MissingDataArray{M,L,U,T,N,A} <: AbstractMissingDataArray{M,L,U,T,N,A}
    data::A
    inds::Vector{CartesianIndex{N}}
end
struct ThreadedMissingDataArray{M,L,U,T,N,A} <: AbstractMissingDataArray{M,L,U,T,N,A}
    data::Vector{A}
    inds::Vector{CartesianIndex{N}}
end

@enum ParameterType::Int8 begin
    RealScalarParameter
    RealArrayParameter
    LKJCorrCholeskyParameter
    SimplexParameter
end

"""
The basic plan is to use Val{Bounds}() for dispatch.
"""
struct Bounds
    lb::Float64
    ub::Float64
end

Bounds(::RealScalar{L,U}) where {L,U} = Bounds(L,U)
Bounds(::RealArray{S,L,U}) where {S,L,U} = Bounds(L,U)

Base.min(b::Bounds) = b.lb
Base.max(b::Bounds) = b.ub
ismin(x::T) where {T} = ((x == typemin(T)) | (x == -floatmax(T))) # floatmin
ismax(x::T) where {T} = ((x == typemax(T)) | (x ==  floatmax(T)))
function isunbounded(b::Bounds)
    ismin(b.lb) & ismax(b.ub)
end
function islowerbounded(b::Bounds)
    (!ismin(b.lb)) & ismax(b.ub)
end
function isupperbounded(b::Bounds)
    ismin(b.lb) & (!ismax(b.ub))
end
function isbounded(b::Bounds)
    (!ismin(b.lb)) & (!ismax(b.ub))
end

# struct ParamDescriptionCore
#     b::Bounds
#     ind::Int32
#     t::ParameterType
# end

# struct SizedParamDescription{N}
#     s::NTuple{N,Int}
#     d::ParamDescriptionCore
# end
struct LengthParamDescription
    l::Int
    ind::Int32
    r::Int8
    LengthParamDescription(l::Int, d) = new(l, d % Int32, (l & ((VectorizationBase.REGISTER_SIZE>>>3)-1) ) % Int8)
end
# LengthDescription(::Type{A}, ind) where {A} = LengthDescription(memory_length(A)::Int, ind % Int32)
LengthParamDescription(A, ind) = LengthParamDescription(memory_length(A)::Int, ind % Int32)

# LengthParamDescription(spd::SizedParamDescription{0}) = LengthParamDescription(1, spd.d)#, one(Int8))
# function LengthParamDescription(spd::SizedParamDescription)
#     s = spd.s
#     l = minimum(s) < 0 ? -1 : prod(s)
#     LengthParamDescription(l, spd.d)
# end
# function update_rem!(descript::AbstractVector{LengthParamDescription}, R)
#     Rm1 = R - 1
#     for j ∈ 1:length(descript)
#         d = descript[j]
#         d = LengthParamDescription(d.l, d.ind, (d.l & Rm1) % Int8)
#         descript[j] = d
#     end
# end

function Base.isless(lpd1::LengthParamDescription, lpd2::LengthParamDescription)
    l1 = lpd1.l; l2 = lpd2.l
    l1dR = iszero(lpd1.r)
    l2dR = iszero(lpd2.r)
    if l1dR ⊻ l2dR
        l1dR
    else#if l1dR & l2dR
        l1 > l2
    # else
        # isless(lpd1, lpd2, R >>> 1)
    end
end
function shift_entries_right!(x, by = 1, start = firstindex(x), stop = lastindex(x))
    @assert stop ≤ lastindex(x)
    @assert start ≥ firstindex(x)
    @inbounds for i ∈ stop-by:-1:start
        x[i+by] = x[i]
    end
end
# Search for parameter pairs that preserve alignemnt, and move them to the front of unaligned params.
# E.g., given [5,5,5,3,3,1] and 8-Float64 alignment, move the (5,3)s in front, reordering to [5,3,5,3,5,1]
function combine_param_pairs!(descript::AbstractVector{LengthParamDescription}, R, start = firstindex(descript), jstart = start)
    N = length(descript)
    start > N - 2 && return start
    for j ∈ jstart:N
        descriptⱼ = descript[j]
        rⱼ = descriptⱼ.r
        Rdiff = R - rⱼ
        for k ∈ j+1:N
            descriptₖ = descript[k]
            rₖ = descriptₖ.r
            if rₖ == Rdiff
                shift_entries_right!(descript, 1, j + 1, k)
                shift_entries_right!(descript, 2, start, j + 1)
                descript[start] = descriptⱼ
                descript[start+1] = descriptₖ
                return combine_param_pairs!(descript, R, start + 2, j + 2)
            end
        end
    end
    start
end
# Search for parameter tripples that preserve alignment, and move them in front of unaligned params.
# E.g., given [5,5,5,2,2,1] and 8-Float64 alignment, move (5,2,1) in front so that it is [5,2,1,5,5,2]
function combine_param_tripples!(descript::AbstractVector{LengthParamDescription}, R, start = firstindex(descript), istart = start)
    N = length(descript)
    start > N - 3 && return start
    for i ∈ start:N
        descriptᵢ = descript[i]
        rᵢ = descriptᵢ.r
        Rdiffᵢ = R - rᵢ
        for j ∈ i+1:N
            descriptⱼ = descript[j]
            rⱼ = descriptⱼ.r
            Rdiffⱼ = Rdiffᵢ - rⱼ
            for k ∈ j+1:N
                descriptₖ = descript[k]
                rₖ = descriptₖ.r
                if rₖ == Rdiffⱼ
                    shift_entries_right!(descript, 1, j + 1, k)
                    shift_entries_right!(descript, 2, i + 1, j + 1)
                    shift_entries_right!(descript, 3, start, i + 1)
                    descript[start  ] = descriptᵢ
                    descript[start+1] = descriptⱼ
                    descript[start+2] = descriptₖ
                    return combine_param_tripples!(descript, R, start + 3, i + 3)
                end
            end
        end
    end
    start
end
function sort_by_rem!(descript::Vector{LengthParamDescription}, R = VectorizationBase.REGISTER_SIZE >>> 3)
    i = 1
    N = length(descript)
    # Original version looped while reducing R
    # We only do this once, because that risks misaligning larger objects.
    # E.g., align to 8 doubles, if we have [(l = 10007, id=1, r = 8), (l = 4, id=2, r = 4)]
    # The algorithm would now choose (id=1,id=2) order, but with the loop it would favor (id=2,id=1)
    #
    # descriptv = @view(descript[1:end])
    # while R > 1
    # descriptv = @view(descriptv[i:end])
    descriptv = descript
    # update_rem!(descriptv, R)
    sort!(descriptv)
    i = findfirst(d -> !iszero(d.r), descriptv)
    i === nothing && return
    i = combine_param_pairs!(descriptv, R, i)
    i = combine_param_tripples!(descriptv, R, i)
    nothing
        # N < i || return
        # R >>>= 1
    # end
end

function parameter_offsets(descript::Vector{LengthParamDescription})
    sort_by_rem!(descript) # should hopefully align loads and stores reasonably well
    offsets = similar(descript, Int)
    offset = 0
    for d ∈ descript
        # @show d.d.ind, d.l, offset
        offsets[d.ind] = offset
        offset += d.l
    end
    offsets
end
@generated function parameter_offsets_noinline(::Val{descript}) where {descript}
    lpd = LengthParamDescription[LengthParamDescription(d)::LengthParamDescription for d ∈ descript]
    offsetvec = parameter_offsets(lpd)
    tup = Expr(:tuple)
    append!(tup.args, offsetvec)
    Expr(:block, Expr(:meta, :noinline), tup)
end
@generated parameter_offsets(::Val{descript}) where {descript} = parameter_offsets_noinline(Val{descript}())
@generated parameter_offset(::Val{descript}, ::Val{I}) where {descript,I} = parameter_offsets_noinline(Val{descript}())[I]


text:
module RecyclePermanent

using RNGPool
using Random

include("model.jl")
include("rpio.jl")
include("common.jl")
include("commonDiscrete.jl")
include("simpleEstimator.jl")
include("biasedEstimator.jl")
include("rpEstimator.jl")
include("kmEstimator.jl")

export Model, simulateModel, RPIO, simpleEstimate, biasedEstimate, rpEstimate,
  rpEstimateDiscrete, biasedEstimateDiscrete, Int64Particle, KMIO, kmEstimate

end # module


text:
using ArgParse
using Formatting
using JSON,TOML
using DataStructures

include("runCmd.jl")
include("tools.jl")

s = ArgParseSettings()
@add_arg_table s begin

    "--diag-file"
        help = "The TOML file containing the information to diagnose."
        arg_type = String
        required = true

    "--diag-overwrite"
        action = :store_true

    "--extra-overwrite"
        action = :store_true

    "--diagcase"
        arg_type = String
        default = ""

    "--convert-only"
        action = :store_true
end

parsed = DataStructures.OrderedDict(parse_args(ARGS, s))

JSON.print(parsed, 4)

include("analysis_dict.jl")

cfg = TOML.parsefile(parsed["diag-file"]) |> DataStructures.OrderedDict
println("Loaded diagnose configuration file ", parsed["diag-file"])
JSON.print(cfg, 4)

result_dir = "result"
diag_data_dir = joinpath(result_dir, cfg["diag-label"])
mkpath(result_dir)
mkpath(diag_data_dir)

lib_dir=joinpath(@__DIR__, "lib")

domain_atm = cfg["domains"]["atm"]
domain_ocn = cfg["domains"]["ocn"]
domain_ice = cfg["domains"]["ice"]

if parsed["diagcase"] != ""
    diagcases = Dict(
        parsed["diagcase"] => cfg["diagcases"][parsed["diagcase"]]
    )
else
    diagcases = cfg["diagcases"]
end

for (diagcase_name, diagcase) in diagcases

    println("Diagcase detail: ")
    JSON.print(diagcase, 4)

    extra_data_dir = joinpath(result_dir, "extra_data", diagcase["casename"])
    diagcase_data_dir = joinpath(diag_data_dir, diagcase_name)

    mkpath(extra_data_dir)
    mkpath(diagcase_data_dir)

    hist_dir_atm = joinpath(cfg["archive-root"], diagcase["casename"], "atm", "hist")
    hist_dir_ice = joinpath(cfg["archive-root"], diagcase["casename"], "ice", "hist")
    hist_dir_ocn = joinpath(cfg["archive-root"], diagcase["casename"], "ocn", "hist")

    casename = diagcase["casename"]
    diag_beg_year, diag_end_year = diagcase["year-rng"]

    if cfg["diagnose"]["atm"]

        println("Diagnosing atm...")

        pleaseRun(`julia $(@__DIR__)/make_extra_data_atm.jl
            --casename   $(casename)
            --input-dir  $(hist_dir_atm)
            --output-dir $(extra_data_dir)
            --year-rng   $(diag_beg_year) $(diag_end_year)
            --domain-file $(domain_atm)
        `)

        if ! parsed["convert-only"]

            output_file = "$(diagcase_data_dir)/atm_analysis_AHT.nc"
            if !isfile(output_file) || parsed["diag-overwrite"]
                pleaseRun(`julia $(lib_dir)/atm_heat_transport.jl
                    --data-file-prefix "$(hist_dir_atm)/$(casename).cam.h0."
                    --data-file-timestamp-form YEAR_MONTH
                    --domain-file $(domain_atm)
                    --output-file $output_file
                    --beg-year $diag_beg_year
                    --end-year $diag_end_year
                `)
            end

            output_file = "$(diagcase_data_dir)/atm_analysis_PDO.nc"
            if !isfile(output_file) || parsed["diag-overwrite"]
                pleaseRun(`julia $(lib_dir)/EOFs/PDO.jl
                    --data-file-prefix "$(hist_dir_atm)/$(casename).cam.h0."
                    --data-file-timestamp-form YEAR_MONTH
                    --domain-file $(domain_atm)
                    --output-file $output_file
                    --beg-year $diag_beg_year
                    --end-year $diag_end_year
                    --SST-varname SST
                    --dims XYT
                    --sparsity 1
                `)
            end

            output_file = "$(diagcase_data_dir)/atm_analysis_ENSO.nc"
            if !isfile(output_file) || parsed["diag-overwrite"]
                pleaseRun(`julia $(lib_dir)/EOFs/ENSO.jl
                    --data-file-prefix "$(hist_dir_atm)/$(casename).cam.h0."
                    --data-file-timestamp-form YEAR_MONTH
                    --domain-file $(domain_atm)
                    --output-file $output_file
                    --beg-year $diag_beg_year
                    --end-year $diag_end_year
                    --SST-varname SST
                    --dims XYT
                    --sparsity 1
                `)
            end
            output_file = "$(diagcase_data_dir)/atm_analysis_AO.nc"
            if !isfile(output_file) || parsed["diag-overwrite"]
                pleaseRun(`julia $(lib_dir)/EOFs/AO.jl
                    --data-file-prefix "$(hist_dir_atm)/$(casename).cam.h0."
                    --data-file-timestamp-form YEAR_MONTH
                    --domain-file $(domain_atm)
                    --output-file $output_file
                    --beg-year $diag_beg_year
                    --end-year $diag_end_year
                    --SLP-varname PSL
                    --dims XYT
                    --sparsity 3
                `)
            end

            output_file = "$(diagcase_data_dir)/atm_analysis_AAO.nc"
            if !isfile(output_file) || parsed["diag-overwrite"]
                pleaseRun(`julia $(lib_dir)/EOFs/AAO.jl
                    --data-file-prefix "$(hist_dir_atm)/$(casename).cam.h0."
                    --data-file-timestamp-form YEAR_MONTH
                    --domain-file $(domain_atm)
                    --output-file $output_file
                    --beg-year $diag_beg_year
                    --end-year $diag_end_year
                    --SLP-varname PSL
                    --dims XYT
                    --sparsity 3
                `)
            end

            output_file = "$(diagcase_data_dir)/atm_analysis_NAO.nc"
            if !isfile(output_file) || parsed["diag-overwrite"]
                pleaseRun(`julia $(lib_dir)/EOFs/NAO.jl
                    --data-file-prefix "$(hist_dir_atm)/$(casename).cam.h0."
                    --data-file-timestamp-form YEAR_MONTH
                    --domain-file $(domain_atm)
                    --output-file $output_file
                    --beg-year $diag_beg_year
                    --end-year $diag_end_year
                    --SLP-varname PSL
                    --dims XYT
                    --sparsity 1
                `)
            end

            output_file = "$(diagcase_data_dir)/atm_analysis_AHT_OHT.nc"
            if !isfile(output_file) || parsed["diag-overwrite"]
                pleaseRun(`julia $(lib_dir)/atmocn_heat_transport.jl
                    --data-file-prefix "$(hist_dir_atm)/$(casename).cam.h0."
                    --data-file-timestamp-form YEAR_MONTH
                    --domain-file $(domain_atm)
                    --output-file $output_file
                    --beg-year $diag_beg_year
                    --end-year $diag_end_year
                `)
            end

            # 3-dimensional fields
            for varname in ["U", "T", "ICEFRAC"]
                output_file = "$(diagcase_data_dir)/atm_analysis_mean_var_$(varname).nc"
                if !isfile(output_file) || parsed["diag-overwrite"]
                    pleaseRun(`julia $(lib_dir)/mean_var.jl
                         --data-file-prefix "$(hist_dir_atm)/$(casename).cam.h0."
                         --output-file $output_file
                         --beg-year $diag_beg_year
                         --end-year $diag_end_year
                         --varname  $(varname)
                    `)
                end
            end

            # Meridional averaged varables
            #for varname in ["U", "T", "VU", "VT", "VQ"]
            for varname in ["U", "T", "Z3"]
                output_file = "$(diagcase_data_dir)/atm_analysis_mean_anomaly_$(varname)_zm.nc"
                if !isfile(output_file) || parsed["diag-overwrite"]
                    pleaseRun(`julia $(lib_dir)/mean_anomaly.jl
                         --data-file-prefix "$(extra_data_dir)/$(casename).cam_extra2_zmean."
                         --data-file-timestamp-form YEAR_MONTH
                         --domain-file $(domain_atm)
                         --output-file $output_file
                         --beg-year $diag_beg_year
                         --end-year $diag_end_year
                         --varname  $(varname)
                         --dims     YZT
                    `)
                end
            end

            if !isfile(output_file) || parsed["diag-overwrite"]
                output_file = "$(diagcase_data_dir)/atm_analysis_mean_anomaly_streamfunction.nc"
                pleaseRun(`julia $(lib_dir)/mean_anomaly.jl
                     --data-file-prefix "$(extra_data_dir)/$(casename).cam_extra3_streamfunction."
                     --data-file-timestamp-form YEAR_MONTH
                     --domain-file $(domain_atm)
                     --output-file $output_file
                     --beg-year $diag_beg_year
                     --end-year $diag_end_year
                     --varname  psi
                     --dims     YZ
                `)
            end

            begin
                varname = "PREC_TOTAL"
                output_file = "$(diagcase_data_dir)/atm_analysis_mean_anomaly_$(varname).nc"
                output_file_zm = "$(diagcase_data_dir)/atm_analysis_mean_anomaly_$(varname)_zm.nc"

                if !isfile(output_file) || parsed["diag-overwrite"]
                    pleaseRun(`julia $(lib_dir)/mean_anomaly.jl
                         --data-file-prefix "$(extra_data_dir)/$(casename).cam_extra1."
                         --data-file-timestamp-form YEAR_MONTH
                         --domain-file $(domain_atm)
                         --output-file $output_file
                         --beg-year $diag_beg_year
                         --end-year $diag_end_year
                         --varname  $(varname)
                         --dims     XYT
                    `)

                    pleaseRun(`ncwa -h -O -a Nx $output_file $output_file_zm`)
                end
            end

            for varname in ["TREFHT", "SST", "TAUX", "TAUY", "PSL", "ICEFRAC", "LHFLX", "SHFLX", "SWCF", "LWCF", "FSNT", "FLNT"]
                output_file = "$(diagcase_data_dir)/atm_analysis_mean_anomaly_$(varname).nc"
                output_file_zm = "$(diagcase_data_dir)/atm_analysis_mean_anomaly_$(varname)_zm.nc"

                if varname in ["SST", "PSL", "TREFHT"]
                    output_monthly_anomalies = true
                else
                    output_monthly_anomalies = false
                end

                if !isfile(output_file) || parsed["diag-overwrite"]
                    pleaseRun(`julia $(lib_dir)/mean_anomaly.jl
                         --data-file-prefix "$(hist_dir_atm)/$(casename).cam.h0."
                         --data-file-timestamp-form YEAR_MONTH
                         --domain-file $(domain_atm)
                         --output-file $output_file
                         --beg-year $diag_beg_year
                         --end-year $diag_end_year
                         --varname  $(varname)
                         --dims     XYT
                        --output-monthly-anomalies $output_monthly_anomalies
                    `)

                    #pleaseRun(`ncwa -h -O -a Nx $output_file $output_file_zm`)
                end

            end

            input_file = "$(diagcase_data_dir)/atm_analysis_mean_anomaly_SST.nc"
            output_file = "$(diagcase_data_dir)/atm_analysis_SST_CORR.nc"
            if !isfile(output_file) || parsed["diag-overwrite"]
                pleaseRun(`julia $(lib_dir)/SST_correlation.jl
                     --data-file $input_file
                     --domain-file $domain_atm
                     --output-file $output_file
                     --SSTA  SST_MA
                `)

            end

    # Atmosphere tropical precipitation asymmetry index
    # julia $script_analysis_dir/atm_PAI.jl --data-file-prefix="$atm_hist_dir/$casename.cam.h0." --data-file-timestamp-form=YEAR_MONTH --domain-file=$atm_domain --output-file=$atm_analysis31 --beg-year=$diag_beg_year --end-year=$diag_end_year

    # Downstream data. No need to specify --beg-year --end-year
    #julia $script_analysis_dir/AO.jl --data-file=$atm_analysis1 --domain-file=$atm_domain --output-file=$atm_analysis5 --sparsity=$PCA_sparsity
        end

    end

    if cfg["diagnose"]["ice"]
        pleaseRun(`julia $(@__DIR__)/make_extra_data_ice.jl
            --casename   $(casename)
            --input-dir  $(hist_dir_ice)
            --output-dir $(extra_data_dir)
            --year-rng   $(diag_beg_year) $(diag_end_year)
            --domain-file $(domain_ice)
            --remap-file-nn $(cfg["remap-files"]["ice2atm"]["nn"])
        `)

        if ! parsed["convert-only"]
            for varname in ["aice", "vice"]
                output_file = "$(diagcase_data_dir)/ice_analysis_mean_anomaly_$(varname).nc"
                output_file_zm = "$(diagcase_data_dir)/ice_analysis_mean_anomaly_$(varname)_zm.nc"

                if !isfile(output_file) || parsed["diag-overwrite"]
                    pleaseRun(`julia $(lib_dir)/mean_anomaly.jl
                         --data-file-prefix "$(extra_data_dir)/$(casename).cice_extra2_rg."
                         --data-file-timestamp-form YEAR_MONTH
                         --domain-file $(domain_atm)
                         --output-file $output_file
                         --beg-year $diag_beg_year
                         --end-year $diag_end_year
                         --varname  $(varname)
                         --dims     XYT
                    `)

                    pleaseRun(`ncwa -h -O -a Nx $output_file $output_file_zm`)
                end

            end

            begin

                output_file = "$(diagcase_data_dir)/ice_analysis_total_seaice.nc"

                if !isfile(output_file) || parsed["diag-overwrite"]

                    pleaseRun(`julia $(lib_dir)/seaice.jl
                        --data-file-prefix $(extra_data_dir)/$(casename).cice_extra1.
                        --data-file-timestamp-form YEAR_MONTH
                        --domain-file $(domain_ice)
                        --output-file $(output_file)
                         --beg-year $diag_beg_year
                         --end-year $diag_end_year

                    `)
    #julia $script_analysis_dir/seaice.jl --data-file-prefix="$concat_dir/$casename.cice_extra.h." --data-file-timestamp-form=YEAR_MONTH --domain-file=$ocn_domain --output-file=$ice_analysis1 --beg-year=$diag_beg_year --end-year=$diag_end_year

                end

            end

        end

    end

    if cfg["diagnose"]["ocn"]

        if haskey(diagcase, "pop2") && diagcase["pop2"] == true
            println("IMPORTANT: This ocean case is produced with POP2. Turn on --pop2 option")
            pop2_option = true
        else
            pop2_option = false
        end

        pleaseRun(`julia $(@__DIR__)/make_extra_data_ocn.jl
            --casename   $(casename)
            --input-dir  $(hist_dir_ocn)
            --output-dir $(extra_data_dir)
            --year-rng   $(diag_beg_year) $(diag_end_year)
            --domain-file $(domain_ocn)
            --remap-file-nn $(cfg["remap-files"]["ocn2atm"]["nn"])
            --pop2     $(pop2_option)
        `)

        if ! parsed["convert-only"]
            #for varname in ["SST", "HMXL"]

            if ! pop2_option
                output_file = "$(diagcase_data_dir)/ocn_analysis_OHT.nc"
                if !isfile(output_file) || parsed["diag-overwrite"]
                    pleaseRun(`julia $(lib_dir)/ocn_heat_transport.jl
                         --data-file-prefix "$(extra_data_dir)/$(casename).EMOM_extra4_OHT."
                         --data-file-timestamp-form YEAR_MONTH
                         --domain-file $(domain_ocn)
                         --output-file $output_file
                         --beg-year $diag_beg_year
                         --end-year $diag_end_year
                    `)
                end
            end

            for varname in ["SST",]
                output_file = "$(diagcase_data_dir)/ocn_analysis_mean_anomaly_$(varname).nc"
                output_file_zm = "$(diagcase_data_dir)/ocn_analysis_mean_anomaly_$(varname)_zm.nc"

                if !isfile(output_file) || parsed["diag-overwrite"]
                    pleaseRun(`julia $(lib_dir)/mean_anomaly.jl
                         --data-file-prefix "$(extra_data_dir)/$(casename).EMOM_extra1_rg."
                         --data-file-timestamp-form YEAR_MONTH
                         --domain-file $(domain_atm)
                         --output-file $output_file
                         --beg-year $diag_beg_year
                         --end-year $diag_end_year
                         --varname  $(varname)
                         --dims     XYT
                    `; igs=true)

                    pleaseRun(`ncwa -h -O -a Nx $output_file $output_file_zm`; igs=true)
                end

            end

            for varname in ["STRAT",]
                output_file = "$(diagcase_data_dir)/ocn_analysis_mean_anomaly_$(varname).nc"
                output_file_zm = "$(diagcase_data_dir)/ocn_analysis_mean_anomaly_$(varname)_zm.nc"

                if !isfile(output_file) || parsed["diag-overwrite"]
                    pleaseRun(`julia $(lib_dir)/mean_anomaly.jl
                         --data-file-prefix "$(extra_data_dir)/$(casename).EMOM_extra2_rg."
                         --data-file-timestamp-form YEAR_MONTH
                         --domain-file $(domain_atm)
                         --output-file $output_file
                         --beg-year $diag_beg_year
                         --end-year $diag_end_year
                         --varname  $(varname)
                         --dims     XYT
                    `; igs=true)

                    pleaseRun(`ncwa -h -O -a Nx $output_file $output_file_zm`; igs=true)
                end

            end
        end
    end

end

max_stars_repo_name: meteorologytoday/CESM-diagnostic

max_issues_repo_name: meteorologytoday/CESM-diagnostic

max_forks_repo_name: meteorologytoday/CESM-diagnostic


text:
let magic
    magic(x) = false
    sentinel = gensym()
    @test magic(sentinel) == false

    # Getting closers to work means having a function created in the current scope
    patch = @patch magic(x) = x == sentinel
    apply(patch) do
        @test (@mock magic(sentinel)) == true
    end
end


text:
@testset "sphericalQuadrature" begin
  ɛ = 1000*eps(Float64);
  @polyvar x y z;

  C₂ = SphericalHarmonicCoefficients(randn(9)); # L = 2
  C₄ = SphericalHarmonicCoefficients(randn(25));# L = 4

  p₂ = sphericalHarmonicsExpansion(C₂,x,y,z);
  p_func₂ = @fastfunc p₂;
  p₄ = sphericalHarmonicsExpansion(C₄,x,y,z);
  p_func₄ = @fastfunc p₄;

  # 4-design with 14 points:
  coordinates₄ = reshape([1.0e0 0.0e0 0.0e0
        5.947189772040725e-1 1.776539926025823e-1 7.840589244857197e-1
        5.947189772040725e-1 -7.678419429698292e-1 -2.381765915652909e-1
        5.947189772040725e-1 5.90187950367247e-1 -5.458823329204288e-1
        -5.947189772040725e-1 1.776539926025823e-1 -7.840589244857197e-1
        -5.947189772040725e-1 5.90187950367247e-1 5.458823329204288e-1
        -5.947189772040725e-1 -7.678419429698292e-1 2.381765915652909e-1
        3.012536847870683e-1 8.79474443923065e-1 3.684710570566285e-1
        3.012536847870683e-1 -7.588425179318781e-1 5.774116818882528e-1
        3.012536847870683e-1 -1.206319259911869e-1 -9.458827389448813e-1
        -3.012536847870683e-1 8.79474443923065e-1 -3.684710570566285e-1
        -3.012536847870683e-1 -1.206319259911869e-1 9.458827389448813e-1
        -3.012536847870683e-1 -7.588425179318781e-1 -5.774116818882528e-1
        -1.0e0 0.0e0 0.0e0],14,3)

  # 8-design with 36 points:
  coordinates₈ = reshape([
          .5074754464108170E+00 -.3062000132395710E+00 .8054254920116630E+00
        -.3062000132395690E+00 .8054254920116630E+00 .5074754464108170E+00
        -.5074754464108170E+00 .3062000132395700E+00 .8054254920116630E+00
        .8054254920116630E+00 .5074754464108170E+00 -.3062000132395690E+00
        .3062000132395690E+00 .8054254920116640E+00 -.5074754464108170E+00
        .8054254920116630E+00 -.5074754464108170E+00 .3062000132395690E+00
        .3062000132395690E+00 -.8054254920116630E+00 .5074754464108160E+00
        -.8054254920116630E+00 -.5074754464108170E+00 -.3062000132395690E+00
        -.3062000132395700E+00 -.8054254920116640E+00 -.5074754464108160E+00
        -.8054254920116630E+00 .5074754464108180E+00 .3062000132395690E+00
        .5074754464108170E+00 .3062000132395700E+00 -.8054254920116630E+00
        -.5074754464108170E+00 -.3062000132395700E+00 -.8054254920116630E+00
        .6263636702652710E+00 -.2435277754091940E+00 -.7405152092807200E+00
        -.2435277754091950E+00 -.7405152092807200E+00 .6263636702652710E+00
        -.6263636702652710E+00 .2435277754091940E+00 -.7405152092807200E+00
        -.7405152092807200E+00 .6263636702652700E+00 -.2435277754091950E+00
        .2435277754091950E+00 -.7405152092807190E+00 -.6263636702652710E+00
        -.7405152092807200E+00 -.6263636702652700E+00 .2435277754091950E+00
        .2435277754091950E+00 .7405152092807190E+00 .6263636702652710E+00
        .7405152092807200E+00 -.6263636702652700E+00 -.2435277754091950E+00
        -.2435277754091950E+00 .7405152092807200E+00 -.6263636702652710E+00
        .7405152092807200E+00 .6263636702652700E+00 .2435277754091950E+00
        .6263636702652710E+00 .2435277754091940E+00 .7405152092807200E+00
        -.6263636702652710E+00 -.2435277754091940E+00 .7405152092807200E+00
        -.2862487234260350E+00 .9571203270924580E+00 -.4452356458542100E-01
        .9571203270924580E+00 -.4452356458542000E-01 -.2862487234260350E+00
        .2862487234260350E+00 -.9571203270924580E+00 -.4452356458542100E-01
        -.4452356458542000E-01 -.2862487234260350E+00 .9571203270924580E+00
        -.9571203270924580E+00 -.4452356458541900E-01 .2862487234260350E+00
        -.4452356458542100E-01 .2862487234260340E+00 -.9571203270924580E+00
        -.9571203270924580E+00 .4452356458542000E-01 -.2862487234260340E+00
        .4452356458542100E-01 .2862487234260340E+00 .9571203270924580E+00
        .9571203270924580E+00 .4452356458542000E-01 .2862487234260340E+00
        .4452356458542100E-01 -.2862487234260340E+00 -.9571203270924580E+00
        -.2862487234260340E+00 -.9571203270924580E+00 .4452356458542100E-01
        .2862487234260350E+00 .9571203270924580E+00 .4452356458542100E-01
          ],36,3)

  values₂ = Array{Float64,1}(undef, size(coordinates₄,1));
  for i=1:size(coordinates₄,1)
      values₂[i] = p_func₂(coordinates₄[i,1],coordinates₄[i,2],coordinates₄[i,3]);
  end
  values₄ = Array{Float64,1}(undef, size(coordinates₈,1));
  for i=1:size(coordinates₈,1)
      values₄[i] = p_func₄(coordinates₈[i,1],coordinates₈[i,2],coordinates₈[i,3]);
  end

  Cnew₂ = sphericalQuadrature(values₂,coordinates₄,2);
  Cnew₄ = sphericalQuadrature(values₄,coordinates₈,4);

  @test isapprox(C₂,Cnew₂,atol=ε);
  @test isapprox(C₄,Cnew₄,atol=ε);

end


text:
using DifferentialEquations
using Plots
plotlyjs()

#constants
l = 1.0                             # length [m]
m = 1.0                             # mass[Kg]
g = 9.81                            # gravitational acceleration [m/s²]

#function to compute ODE
function pendulum!(du,u,p,t)
    du[1] = u[2]                    # θ'(t) = ω(t)
    du[2] = -(g\l)*sin(u[1])  # ω'(t) = -(g/l) sin θ(t)
end

#initial conditions
θ₀ = π/6                         # initial angular deflection [rad]
ω₀ = 0.0                            # initial angular velocity [rad/s]
u₀ = [θ₀, ω₀]                       # initial state vector
tspan = (0.0,100.0)                  # time interval

#defining the problem and solving
prob = ODEProblem(pendulum!,u₀,tspan)
sol = solve(prob,saveat=0.1)

#plotting the solution
display(plot(sol,linewidth=2,xaxis="t (s)",
    color=["green" "blue"],label=["θ [rad]" "ω [rad/s]"],layout=(2,1)))
savefig("./plots/simple pendulum-plot.png")

# trajectory
t=sol.t
U=sol[1:end,:]
(θ,ω)=[U[x,:] for x in 1:size(U,1)]
x=l.*sin.(π.-θ)
y=zeros(size(t))
z=l.*cos.(π.-θ)
display(plot(x,y,z,linewidth=2,label="trajectory",
    xaxis="x (m)",yaxis="y (m)",zaxis="z (m)",grid=(:on,:black)))
savefig("./plots/simple pendulum-trajectory.png")

max_stars_repo_path: Modeling the Dragon Module as a Pendulum/Scripts/pendulum.jl


text:
@testset "MHA" begin
    using NeuralAttentionlib: CausalMask, BiLengthMask, BatchedMask
    head = 4
    input_dims = 7
    head_dims = 5
    output_dims = 8

    function atten(mha, x, mask=nothing)
        q = mha.iqproj(x)
        k = mha.ikproj(x)
        v = mha.ivproj(x)
        a = NeuralAttentionlib.multihead_qkv_attention(4, q, k, v, mask)
        return mha.oproj(a)
    end

    @testset "w/o mask" begin
        x = randn(input_dims, 3, 2)
        mha = MultiheadAttention(head, input_dims, head_dims, output_dims;
                                 future = true, pdrop = 0) # disable dropout

        @test mha(x,x,x) ≈ atten(mha, x)

        gradN = Flux.gradient(mha, x) do mha, x
            sum(sin.(atten(mha, x)))
        end
        gradT = Flux.gradient(mha, x) do mha, x
            sum(sin.(mha(x,x,x)))
        end

        @test isapprox(gradN[1].iqproj.weight, gradT[1].iqproj.weight; atol = 1e-12)
        @test isapprox(gradN[1].ikproj.weight, gradT[1].ikproj.weight; atol = 1e-12)
        @test isapprox(gradN[1].ivproj.weight, gradT[1].ivproj.weight; atol = 1e-12)
        @test isapprox(gradN[1].oproj.weight, gradT[1].oproj.weight; atol = 1e-12)
        @test isapprox(gradN[1].iqproj.bias, gradT[1].iqproj.bias; atol = 1e-12)
        @test isapprox(gradN[1].ikproj.bias, gradT[1].ikproj.bias; atol = 1e-12)
        @test isapprox(gradN[1].ivproj.bias, gradT[1].ivproj.bias; atol = 1e-12)
        @test isapprox(gradN[1].oproj.bias, gradT[1].oproj.bias; atol = 1e-12)
        @test gradN[2] ≈ gradT[2]
    end

    @testset "w/ future mask" begin
        x = randn(input_dims, 3, 2)
        mha = MultiheadAttention(head, input_dims, head_dims, output_dims;
                                 future = false, pdrop = 0) # disable dropout

        @test mha(x,x,x) ≈ atten(mha, x, CausalMask())

        gradN = Flux.gradient(mha, x) do mha, x
            sum(sin.(atten(mha, x, CausalMask())))
        end
        gradT = Flux.gradient(mha, x) do mha, x
            sum(sin.(mha(x,x,x)))
        end

        @test isapprox(gradN[1].iqproj.weight, gradT[1].iqproj.weight; atol = 1e-12)
        @test isapprox(gradN[1].ikproj.weight, gradT[1].ikproj.weight; atol = 1e-12)
        @test isapprox(gradN[1].ivproj.weight, gradT[1].ivproj.weight; atol = 1e-12)
        @test isapprox(gradN[1].oproj.weight, gradT[1].oproj.weight; atol = 1e-12)
        @test isapprox(gradN[1].iqproj.bias, gradT[1].iqproj.bias; atol = 1e-12)
        @test isapprox(gradN[1].ikproj.bias, gradT[1].ikproj.bias; atol = 1e-12)
        @test isapprox(gradN[1].ivproj.bias, gradT[1].ivproj.bias; atol = 1e-12)
        @test isapprox(gradN[1].oproj.bias, gradT[1].oproj.bias; atol = 1e-12)
        @test gradN[2] ≈ gradT[2]
    end

    @testset "w/ future + length mask" begin
        x = randn(input_dims, 6, 2)
        mha = MultiheadAttention(head, input_dims, head_dims, output_dims;
                                 future = false, pdrop = 0) # disable dropout

        q_len = [3, 5]
        k_len = [4, 4]
        lmask = BatchedMask(BiLengthMask(q_len, k_len))
        mask = CausalMask() & lmask

        old_q_len = Old_Impl.getmask(map(l->ones(l), q_len))
        old_k_len = Old_Impl.getmask(map(l->ones(l), k_len))
        old_mask = zeros(6, 6, 2)
        old_mask[1:4, 1:5, :] .= Old_Impl.getmask(old_k_len, old_q_len)

        @test mha(x,x,x, mask=old_mask) ≈ atten(mha, x, mask)

        gradN = Flux.gradient(mha, x, mask) do mha, x, m
            sum(sin.(atten(mha, x, m)))
        end
        gradT = Flux.gradient(mha, x, old_mask) do mha, x, m
            sum(sin.(mha(x,x,x, mask=m)))
        end

        @test isapprox(gradN[1].iqproj.weight, gradT[1].iqproj.weight; atol = 1e-12)
        @test isapprox(gradN[1].ikproj.weight, gradT[1].ikproj.weight; atol = 1e-12)
        @test isapprox(gradN[1].ivproj.weight, gradT[1].ivproj.weight; atol = 1e-12)
        @test isapprox(gradN[1].oproj.weight, gradT[1].oproj.weight; atol = 1e-12)
        @test isapprox(gradN[1].iqproj.bias, gradT[1].iqproj.bias; atol = 1e-12)
        @test isapprox(gradN[1].ikproj.bias, gradT[1].ikproj.bias; atol = 1e-12)
        @test isapprox(gradN[1].ivproj.bias, gradT[1].ivproj.bias; atol = 1e-12)
        @test isapprox(gradN[1].oproj.bias, gradT[1].oproj.bias; atol = 1e-12)
        @test gradN[2] ≈ gradT[2]
        @test gradN[3] == gradT[3]
    end

end


text:
import GraphUtils: get_graph, get_vtx

export
    AbstractGraphState,
    GraphState,
    # get_vtx,
    get_t,
    AbstractGraphAction,
    GraphAction,
    get_e,
    get_dt

abstract type AbstractGraphState end
@with_kw_noshow struct GraphState <: AbstractGraphState
    vtx::Int        = -1 # vertex of graph
    t::Int          = -1
end
get_vtx(s::AbstractGraphState)  = s.vtx
get_t(s::AbstractGraphState)    = s.t
is_valid(s::AbstractGraphState) = get_vtx(s) >= 1
Base.string(s::AbstractGraphState) = "(v=$(get_vtx(s)),t=$(get_t(s)))"
Base.show(io::IO,s::AbstractGraphState) = print(io,typeof(s),string(s)) #"(v=$(get_vtx(s)),t=$(get_t(s)))"
# allow paths to start at non-zero start times
get_start_index(path::Path{S,A,C}) where {S<:AbstractGraphState,A,C} = get_t(get_initial_state(path))

# AbstractGraphAction
export AbstractGraphAction

abstract type AbstractGraphAction end
@with_kw_noshow struct GraphAction <: AbstractGraphAction
    e::Edge{Int}    = Edge(-1,-1)
    dt::Int         = 1
end
get_e(a::AbstractGraphAction)   = a.e
get_dt(a::AbstractGraphAction)  = a.dt
Base.reverse(a::AbstractGraphAction) = typeof(a)(a,e=reverse(a.e))
Base.string(a::AbstractGraphAction) = "(e=$(get_e(a).src) → $(get_e(a).dst))"
Base.show(io::IO,a::AbstractGraphAction) = print(io,typeof(a),string(a)) #"(v=$(get_vtx(s)),t=$(get_t(s)))"

export
    GraphEnv

"""
    GraphEnv{S,A,C} <: AbstractLowLevelEnv{S,A,C}

An abstract environment type whose concrete subtypes comply with a standard
interface.
"""
abstract type GraphEnv{S,A,C} <: AbstractLowLevelEnv{S,A,C} end

export
    # get_graph,
    get_cost_model,
    get_agent_id,
    get_constraints,
    get_goal,
    get_heuristic_model

"""
    GraphUtils.get_graph(env::GraphEnv)

Must be implemented for all concrete subtypes of `GraphEnv`
"""
get_graph(env::GraphEnv)            = env.graph
get_cost_model(env::GraphEnv)       = env.cost_model
get_agent_id(env::GraphEnv)         = env.agent_idx
get_constraints(env::GraphEnv)      = env.constraints
get_goal(env::GraphEnv)             = env.goal
get_heuristic_model(env::GraphEnv)  = env.heuristic
function GraphUtils.get_distance(env::GraphEnv,s1::AbstractGraphState,s2::AbstractGraphState)
    get_distance(get_graph(env),get_vtx(s1),get_vtx(s2))
end

set_agent_idx!(env::E,i) where {E<:GraphEnv} = E(env,agent_idx=i)

cost_type(env::GraphEnv)            = cost_type(get_cost_model(env))

is_valid(env::GraphEnv,s::AbstractGraphState) = has_vertex(get_graph(env),get_vtx(s))
is_valid(env::GraphEnv,a::AbstractGraphAction) = has_edge(get_graph(env),get_e(a))
# get_possible_actions(env::GraphEnv,s) = map(v->GraphAction(e=Edge(get_vtx(s),v)),outneighbors(get_graph(env),get_vtx(s)))
get_next_state(s::AbstractGraphState,a::AbstractGraphAction) = GraphState(get_e(a).dst,get_t(s)+get_dt(a))
# get_next_state(env::GraphEnv,s,a) = get_next_state(s,a)
wait(s::AbstractGraphState) = GraphAction(e=Edge(get_vtx(s),get_vtx(s)))
# wait(env::GraphEnv,s) = GraphAction(e=Edge(get_vtx(s),get_vtx(s)))
function get_transition_cost(c::TravelTime,env::GraphEnv,s,a,sp)
    return cost_type(c)(get_dt(a))
end
function get_conflict_value(c::C,agent_id::Int,env::GraphEnv,s::AbstractGraphState,t=get_t(s)) where {C<:SoftConflictTable}
    idx,t = serialize_jointly(env,s,t)
    get_conflict_value(c,agent_id,idx,t)
end
function get_conflict_value(c::C,agent_id::Int,env::GraphEnv,a::AbstractGraphAction,t) where {C<:SoftConflictTable}
    idx,t = serialize_jointly(env,a,t)
    get_conflict_value(c,agent_id,idx,t)
end
function get_transition_cost(c::C,env::GraphEnv,s,a,sp) where {C<:ConflictCostModel}
    state_conflict_value = get_conflict_value(c, get_agent_id(env), get_vtx(sp), get_t(sp))
    edge_conflict_value = min(
        get_conflict_value(c,get_agent_id(env), get_vtx(sp), get_t(s)),
        get_conflict_value(c,get_agent_id(env), get_vtx(s), get_t(sp)),
        )
    if edge_conflict_value > 0
        # println("Possible Edge Conflict")
        edge_conflict_value = 0
        for (i,p) in enumerate(c.table.paths)
            if (get_planned_vtx(c.table, i, get_t(s)) == get_vtx(sp)) && (get_planned_vtx(c.table, i, get_t(sp)) == get_vtx(s))
                edge_conflict_value += 1
            end
        end
    end
    return state_conflict_value + edge_conflict_value
end
function get_transition_cost(c::ConflictCostModel{S},env::GraphEnv,s,a,sp) where {S<:SoftConflictTable}
    state_conflict_value = get_conflict_value(c.table,get_agent_id(env),env,sp)
    action_conflict_value = get_conflict_value(c.table,get_agent_id(env),env,a,get_t(sp))
    return state_conflict_value + action_conflict_value
end
function get_transition_cost(c::TravelDistance,env::GraphEnv,s,a,sp)
    return (get_vtx(s) == get_vtx(sp)) ? 0.0 : 1.0
end
compute_path_cost(c::TravelTime,env::GraphEnv,path,i) = cost_type(c)(get_t(get_final_state(path)))
function compute_path_cost(c::TravelDistance,env::GraphEnv,path,i)
    cost_type(c)(sum(map(
        n->get_e(get_a(n)).src != get_e(get_a(n)).dst,
        path.path_nodes)))
end
# get_heuristic_cost(env::GraphEnv,s) = get_heuristic_cost(env,get_heuristic_model(env),s)
get_heuristic_cost(env::GraphEnv,s) = get_heuristic_cost(get_heuristic_model(env),env,s)
# function get_heuristic_cost(env::GraphEnv,h::H,s) where {H<:Union{PerfectHeuristic,DefaultPerfectHeuristic}}
function get_heuristic_cost(h::H,env::GraphEnv,s) where {H<:Union{PerfectHeuristic,DefaultPerfectHeuristic}}
    get_heuristic_cost(h, get_vtx(get_goal(env)), get_vtx(s))
end
# function get_heuristic_cost(env::GraphEnv,h::EnvDistanceHeuristic,s)
function get_heuristic_cost(h::EnvDistanceHeuristic,env::GraphEnv,s)
    get_distance(env, s, get_goal(env))
end
# function get_heuristic_cost(env::GraphEnv,h::H,s) where {E<:GraphEnv, H<:ConflictTableHeuristic}
function get_heuristic_cost(h::H,env::GraphEnv,s) where {E<:GraphEnv, H<:ConflictTableHeuristic}
    get_heuristic_cost(h, get_agent_id(env), get_vtx(s), get_t(s))
end
# states_match
states_match(s1::AbstractGraphState,s2::AbstractGraphState) = (get_vtx(s1) == get_vtx(s2))
states_match(env::GraphEnv,s1,s2) = (get_vtx(s1) == get_vtx(s2))
################################################################################
######################## Low-Level (Independent) Search ########################
################################################################################
num_states(env::GraphEnv)                           = nv(get_graph(env))
num_actions(env::GraphEnv)                          = num_states(env)^2 # NOT actually true, but necessary for O(1) serialization
state_space_trait(env::GraphEnv)                    = DiscreteSpace()
action_space_trait(env::GraphEnv)                   = DiscreteSpace()
serialize(env::GraphEnv,s::AbstractGraphState,t=get_t(s))            = get_vtx(s), t
# NOTE the return type of deserialize doesn't need to match the state/action type of env, since it is only used for constraint checking
deserialize(env::GraphEnv,s::AbstractGraphState,idx::Int,t=-1) = GraphState(vtx=idx,t=t), t
function serialize(env::GraphEnv,a::AbstractGraphAction,t=-1)
    (get_e(a).src-1)*num_states(env)+get_e(a).dst, t
end
function deserialize(env::GraphEnv,s::AbstractGraphAction,idx::Int,t=get_t(s))
    GraphAction(e = Edge(
        div(idx-1,num_states(env))+1,
        mod(idx-1,num_states(env))+1)
        ), t
end
serialize_jointly(env::GraphEnv,s::AbstractGraphState,t=get_t(s)) = serialize(env,s,t)
function serialize_jointly(env::GraphEnv,a::A,t=-1) where {A<:AbstractGraphAction}
    a_ = A(a,e=Edge(sort([get_e(a).src, get_e(a).dst])...))
    idx, t = serialize(env,a_,t)
    return idx+num_states(env), t
end

"""
    create_reservations(env::GraphEnv,s,a,sp,t=-1)

Generates three reservations as follows
```
t0 = get_t(s)
tF = get_t(sp)
t_mid = (t0+tF)/2
reservations = [
    ResourceReservation{Int}(s_idx,get_agent_id(env),   (t0,    t_mid)),
    ResourceReservation{Int}(a_idx,get_agent_id(env),   (t0,    tF)),
    ResourceReservation{Int}(sp_idx,get_agent_id(env),  (t_mid, tF)),
]
```
In this way, the reservations for one path node will not interfere with those
for the next path node.
"""
function create_reservations(env::GraphEnv,s,a,sp,t=-1)
    s_idx,t0    = serialize(env,s)
    a_idx,_     = serialize_jointly(env,a) # serialize action to not overlap with state serialization space?
    sp_idx,tF   = serialize(env,sp)
    t_mid = (t0+tF)/2
    reservations = [
        ResourceReservation{Float64}(s_idx,get_agent_id(env),   (t0,    t_mid)),
        ResourceReservation{Float64}(a_idx,get_agent_id(env),   (t0,    tF)),
        ResourceReservation{Float64}(sp_idx,get_agent_id(env),  (t_mid, tF)),
    ]
end
function create_reservations(env::GraphEnv,n::PathNode,t=-1)
    create_reservations(env,get_s(n),get_a(n),get_sp(n),t)
end
# is_goal
function is_goal(env::GraphEnv,s)
    if get_t(s) >= get_t(get_goal(env))
        if states_match(s, get_goal(env)) || !is_valid(get_goal(env))
            return true
        end
    end
    return false
end
function violates_constraints(env::GraphEnv, s, a, sp)
    t = get_t(sp)
    if has_constraint(env,get_constraints(env),
        state_constraint(get_agent_id(get_constraints(env)),PathNode(s,a,sp),t)
        )
        return true
    elseif has_constraint(env,get_constraints(env),
        action_constraint(get_agent_id(get_constraints(env)),PathNode(s,a,sp),t)
        )
        return true
    end
    return false
end

################################################################################
###################### Conflict-Based Search (High-Level) ######################
################################################################################
function detect_state_conflict(n1::N,n2::N) where {S<:AbstractGraphState,A<:AbstractGraphAction,N<:PathNode{S,A}}
    if get_vtx(n1.sp) == get_vtx(n2.sp) && get_t(n1.sp) == get_t(n2.sp)
        return true
    end
    return false
end
function detect_action_conflict(n1::N,n2::N) where {S<:AbstractGraphState,A<:AbstractGraphAction,N<:PathNode{S,A}}
    if (get_e(n1.a).src == get_e(n2.a).dst) && (get_e(n1.a).dst == get_e(n2.a).src) && (get_t(n1.sp) == get_t(n2.sp))
        return true
    end
    return false
end

################################################################################
############################### HELPER FUNCTIONS ###############################
################################################################################
""" Helper for displaying Paths """
function convert_to_vertex_lists(path::Path{S,A,C}) where {S<:AbstractGraphState,A<:AbstractGraphAction,C}
    vtx_list = [get_vtx(n.sp) for n in path.path_nodes]
    # if length(path) > 0
    #     vtx_list = [get_s(get_path_node(path,1)).vtx, vtx_list...]
    # else
        vtx_list = [get_vtx(get_initial_state(path)), vtx_list...]
    # end
    vtx_list
end