# update_eta_Met.jl

## To do:
# - construct_Dh doesn't need h as an argument

export βδ_x_h_modify_γ, δ_x_h_modify_γ;

function update_η_h_Met!(model::Model_BNP_WMReg_Joint, h::Int, Λβ0star_ηy::Array{T,1}, βΛβ0star_ηy::T; update_ηy::Bool=true) where T <: Real
    if model.Σx_type == :full
        update_η_h_Met_Σx_full!(model::Model_BNP_WMReg_Joint, h::Int, Λβ0star_ηy::Array{T,1}, βΛβ0star_ηy::T, update_ηy=update_ηy)
    elseif model.Σx_type == :diag
        update_η_h_Met_Σx_diag!(model::Model_BNP_WMReg_Joint, h::Int, Λβ0star_ηy::Array{T,1}, βΛβ0star_ηy::T, update_ηy=update_ηy)
    end
end

# pre_compute Λ0star_ηy*β0star_ηy and β0star_ηy'Λ0star_ηy*β0star_ηy which are not indexed by h
function update_η_h_Met_Σx_full!(model::Model_BNP_WMReg_Joint, h::Int, Λβ0star_ηy::Array{T,1}, βΛβ0star_ηy::T; update_ηy::Bool=true) where T <: Real

    ## if you modify this function, you must modify its sisters.

    indx_h = findall(model.state.S .== h)
    n_h = length(indx_h)

    μ_x_h_old = model.state.μ_x[h,:]
    lδ_x_h_old = log.(model.state.δ_x[h,:])
    β_x_h_old = [ model.state.β_x[k][h,:] for k = 1:(model.K - 1) ]
    ηlδ_x_old = vcat(μ_x_h_old,
        vcat(β_x_h_old...),
        lδ_x_h_old)

    ## Draw candidate for η_x
    ηlδ_x_cand = ηlδ_x_old + rand(MvNormal(model.state.cSig_ηlδx[h]))
    μ_x_h_cand = ηlδ_x_cand[model.indx_ηx[:μ]]
    β_x_h_candvec = ηlδ_x_cand[model.indx_ηx[:β]]
    β_x_h_cand = [ β_x_h_candvec[model.indx_β_x[k]] for k = 1:(model.K - 1) ]
    lδ_x_h_cand = ηlδ_x_cand[model.indx_ηx[:δ]]
    δ_x_h_cand = exp.(lδ_x_h_cand)

    auto_reject = (!all(δ_x_h_cand .> 0.0)) || any(isinf.(δ_x_h_cand))
    not_auto_reject = !auto_reject

    lNX_mat_cand = deepcopy(model.state.lNX)

    ## bookkeeping for variable selection
    if typeof(model.state.γ) == BitArray{1}
        γ_h = deepcopy(model.state.γ)
    elseif typeof(model.state.γ) == BitArray{2}
        γ_h = deepcopy(model.state.γ[h,:])
    end
    γindx = findall(γ_h)
    nγ = length(γindx)

    if not_auto_reject
        if model.state.γδc == Inf # subset method
            if nγ == 0
                lNX_mat_cand[:,h] = zeros(Float64, model.n)
            elseif nγ == 1
                lNX_mat_cand[:,h] = logpdf.(Normal(μ_x_h_cand[γindx[1]], sqrt(δ_x_h_cand[γindx[1]])), vec(model.X[:,γindx]))
            elseif nγ > 1
                βγ_x_h_cand, δγ_x_h_cand = βδ_x_h_modify_γ(β_x_h_cand, δ_x_h_cand,
                    γ_h, model.state.γδc) # either variance-inflated or subset

                tmp = lNX_sqfChol(permutedims(model.X[:,γindx]),
                    μ_x_h_cand[γindx], βγ_x_h_cand, δγ_x_h_cand, false) # false means set to return nothing if not pos.def.

                if isnothing(tmp)
                    auto_reject = true
                    not_auto_reject = !auto_reject
                else
                    lNX_mat_cand[:,h] = deepcopy(tmp)
                end
            end

            lωNX_vec_cand = lωNXvec(model.state.lω, lNX_mat_cand) # n vector

        elseif isnothing(model.state.γδc) # integration method
            if nγ == 0
                lNX_mat_cand[:,h] = zeros(Float64, model.n)
            elseif nγ == 1
                tmp = sqfChol_to_Σ(β_x_h_cand, δ_x_h_cand, false) # false means don't throw error if not PosDef

                if isnothing(tmp)
                    auto_reject = true
                    not_auto_reject = !auto_reject
                else
                    σ2x = tmp.mat[γindx, γindx][1]
                    lNX_mat_cand[:,h] = logpdf.(Normal(μ_x_h_cand[γindx[1]], sqrt(σ2x)), vec(model.X[:,γindx]))
                end
            elseif nγ > 1
                tmp = sqfChol_to_Σ(β_x_h_cand, δ_x_h_cand, false) # false means don't throw error if not PosDef

                if isnothing(tmp)
                    auto_reject = true
                    not_auto_reject = !auto_reject
                else
                    Σx = PDMat_adj(tmp.mat[γindx, γindx])
                    lNX_mat_cand[:,h] = logpdf(MultivariateNormal(μ_x_h_cand[γindx], Σx), permutedims(model.X[:,γindx]))
                end
            end

            lωNX_vec_cand = lωNXvec(model.state.lω, lNX_mat_cand) # n vector

        else # variance-inflation method

            throw("Variance-inflation method no longer supported.")

            # βγ_x_h_cand, δγ_x_h_cand = βδ_x_h_modify_γ(β_x_h_cand, δ_x_h_cand,
            #     γ_h, model.state.γδc) # either variance-inflated or subset

            # tmp = lNX_sqfChol(permutedims(model.X), μ_x_h_cand, βγ_x_h_cand, δγ_x_h_cand, false) # false means don't throw error if not PosDef

            # if isnothing(tmp)
            #     auto_reject = true
            #     not_auto_reject = !auto_reject
            # else
            #     lNX_mat_cand[:,h] = deepcopy(tmp)
            #     lωNX_vec_cand = lωNXvec(model.state.lω, lNX_mat_cand) # n vector
            # end

        end
    end

    if n_h == 0

        ## Metropolis step for η_x

        ## Compute acceptance ratio
        if not_auto_reject
            lar = lG0_ηlδx(μ_x_h_cand, β_x_h_cand, lδ_x_h_cand, model.state) -
                    sum(lωNX_vec_cand) -
                    lG0_ηlδx(μ_x_h_old, β_x_h_old, lδ_x_h_old, model.state) +
                    sum(model.state.lωNX_vec)
        end

        ## Decision and update
        if not_auto_reject && log(rand()) < lar # accept

            model.state.μ_x[h,:] = deepcopy(μ_x_h_cand)
            for k = 1:(model.K - 1)
                model.state.β_x[k][h,:] = deepcopy(β_x_h_cand[k])
            end
            model.state.δ_x[h,:] = deepcopy(δ_x_h_cand)

            model.state.lNX[:,h] = deepcopy(lNX_mat_cand[:,h])
            model.state.lωNX_vec = deepcopy(lωNX_vec_cand)

            ηlδ_x_out = deepcopy(ηlδ_x_cand)
            model.state.accpt[h] += 1

        else # reject
            ηlδ_x_out = deepcopy(ηlδ_x_old)
        end

        ## Full conditional draw (from G0) for η_y

        if update_ηy
            model.state.δ_y[h] = rand(InverseGamma(0.5 * model.state.ν_δy,
                                        0.5 * model.state.ν_δy * model.state.s0_δy))
            βstar_ηy = model.state.β0star_ηy + (model.state.Λ0star_ηy.chol.U \ randn((model.K + 1)) .* sqrt(model.state.δ_y[h]))
            model.state.μ_y[h] = deepcopy(βstar_ηy[model.indx_ηy[:μ]])
            model.state.β_y[h,:] = deepcopy(βstar_ηy[model.indx_ηy[:β]])
        end

    else # some observations are assigned to cluster h (i.e., n_h > 0)

        ## Metropolis step for η_x

            ## Candidate, lNX, lωNXvec, previously computed
            ## Important quantities

        y_h = model.y[indx_h] # doesn't need deepcopy if indexing (result of call to getindex)
        X_h = model.X[indx_h,:]
        D_h_old = construct_Dh(h, X_h, model.state.μ_x[h,:], γ_h)

        Λ1star_ηy_h_old = get_Λ1star_ηy_h(D_h_old, model.state.Λ0star_ηy)
        β1star_ηy_h_old = get_β1star_ηy_h(Λ1star_ηy_h_old, Λβ0star_ηy, D_h_old, y_h)

        a1_δy_old = get_a1_δy_h(model.state.ν_δy, n_h) # posterior IG shape
        b1_δy_old = get_b1_δy_h(model.state.ν_δy, model.state.s0_δy, y_h, βΛβ0star_ηy, Λ1star_ηy_h_old, β1star_ηy_h_old) # posterior IG scale

        ## Important quantities for candidate
        if not_auto_reject
            D_h_cand = construct_Dh(h, X_h, μ_x_h_cand, γ_h)

            Λ1star_ηy_h_cand = get_Λ1star_ηy_h(D_h_cand, model.state.Λ0star_ηy)
            β1star_ηy_h_cand = get_β1star_ηy_h(Λ1star_ηy_h_cand, Λβ0star_ηy, D_h_cand, y_h)

            a1_δy_cand = get_a1_δy_h(model.state.ν_δy, n_h) # posterior IG shape
            b1_δy_cand = get_b1_δy_h(model.state.ν_δy, model.state.s0_δy, y_h, βΛβ0star_ηy, Λ1star_ηy_h_cand, β1star_ηy_h_cand) # posterior IG scale

            ## Compute acceptance ratio (lcc_ηlδx does not need γ-modified βx and δx)
            lar = lcc_ηlδx(h, indx_h, lNX_mat_cand, lωNX_vec_cand,
                            model.state, μ_x_h_cand, β_x_h_cand, lδ_x_h_cand,
                            Λ1star_ηy_h_cand, a1_δy_cand, b1_δy_cand) -
                lcc_ηlδx(h, indx_h, model.state.lNX, model.state.lωNX_vec,
                            model.state, μ_x_h_old, β_x_h_old, lδ_x_h_old,
                            Λ1star_ηy_h_old, a1_δy_old, b1_δy_old)

        end

        ## Decision and update
        if not_auto_reject && log(rand()) < lar # accept

            model.state.μ_x[h,:] = deepcopy(μ_x_h_cand)
            for k = 1:(model.K - 1)
                model.state.β_x[k][h,:] = deepcopy(β_x_h_cand[k])
            end
            model.state.δ_x[h,:] = deepcopy(δ_x_h_cand)

            model.state.lNX[:,h] = deepcopy(lNX_mat_cand[:,h])
            model.state.lωNX_vec = deepcopy(lωNX_vec_cand)

            ηlδ_x_out = deepcopy(ηlδ_x_cand)
            model.state.accpt[h] += 1

            ## Full conditional draw for η_y ( use prec. prameterization )
            if update_ηy
                model.state.δ_y[h] = rand(InverseGamma(a1_δy_cand, b1_δy_cand))
                βstar_ηy = β1star_ηy_h_cand + (Λ1star_ηy_h_cand.chol.U \ randn((model.K + 1)) .* sqrt(model.state.δ_y[h]))
                model.state.μ_y[h] = deepcopy(βstar_ηy[model.indx_ηy[:μ]])
                model.state.β_y[h,:] = deepcopy(βstar_ηy[model.indx_ηy[:β]])
            end

        else # reject

            ηlδ_x_out = deepcopy(ηlδ_x_old)

            ## Full conditional draw for η_y
            if update_ηy
                model.state.δ_y[h] = rand(InverseGamma(a1_δy_old, b1_δy_old))
                βstar_ηy = β1star_ηy_h_old + (Λ1star_ηy_h_old.chol.U \ randn((model.K + 1)) .* sqrt(model.state.δ_y[h]))
                model.state.μ_y[h] = deepcopy(βstar_ηy[model.indx_ηy[:μ]])
                model.state.β_y[h,:] = deepcopy(βstar_ηy[model.indx_ηy[:β]])
            end
        end
    end

    if model.state.adapt && (model.state.iter % model.state.adapt_thin == 0)
        model.state.adapt_iter += 1
        model.state.runningsum_ηlδx[h,:] += ηlδ_x_out
        runningmean = model.state.runningsum_ηlδx[h,:] / float(model.state.adapt_iter)
        runningdev = ( ηlδ_x_out - runningmean )
        model.state.runningSS_ηlδx[h,:,:] = runningdev * runningdev' # the mean is changing, but this approx. is fine.
    end

    return nothing
end

function update_η_h_Met_Σx_diag!(model::Model_BNP_WMReg_Joint, h::Int, Λβ0star_ηy::Array{T,1}, βΛβ0star_ηy::T; update_ηy::Bool=true) where T <: Real

    ## if you modify this function, you must modify its sisters.

    indx_h = findall(model.state.S .== h)
    n_h = length(indx_h)

    μ_x_h_old = model.state.μ_x[h,:]
    lδ_x_h_old = log.(model.state.δ_x[h,:])
    ηlδ_x_old = vcat(μ_x_h_old,
        # vcat(β_x_h_old...),
        lδ_x_h_old)

    ## Draw candidate for η_x
    ηlδ_x_cand = ηlδ_x_old + rand(MvNormal(model.state.cSig_ηlδx[h]))
    μ_x_h_cand = ηlδ_x_cand[model.indx_ηx[:μ]]
    lδ_x_h_cand = ηlδ_x_cand[model.indx_ηx[:δ]]
    δ_x_h_cand = exp.(lδ_x_h_cand)

    auto_reject = (!all(δ_x_h_cand .> 0.0)) || any(isinf.(δ_x_h_cand))
    not_auto_reject = !auto_reject

    lNX_mat_cand = deepcopy(model.state.lNX)

    ## bookkeeping for variable selection
    if typeof(model.state.γ) == BitArray{1}
        γ_h = deepcopy(model.state.γ)
    elseif typeof(model.state.γ) == BitArray{2}
        γ_h = deepcopy(model.state.γ[h,:])
    end
    γindx = findall(γ_h)
    nγ = length(γindx)

    if not_auto_reject
        if model.state.γδc == Inf || isnothing(model.state.γδc) # subset and integration methods are identical with SigX_diag
            if nγ == 0
                lNX_mat_cand[:,h] = zeros(Float64, model.n)
            elseif nγ == 1
                lNX_mat_cand[:,h] = logpdf.(Normal(μ_x_h_cand[γindx[1]], sqrt(δ_x_h_cand[γindx[1]])), vec(model.X[:,γindx]))
            elseif nγ > 1
                lNX_mat_cand[:,h] = lNX_Σdiag(permutedims(model.X[:,γindx]),
                    μ_x_h_cand[γindx], δ_x_h_cand[γindx])
            end
        else # variance-inflation method
            throw("Variance inflation method not supported when Sigma_x is diagonal.")
        end

        lωNX_vec_cand = lωNXvec(model.state.lω, lNX_mat_cand) # n vector
    end

    if n_h == 0

        ## Metropolis step for η_x

        ## Compute acceptance ratio
        if not_auto_reject
            lar = lG0_ηlδx(μ_x_h_cand, nothing, lδ_x_h_cand, model.state) -
                    sum(lωNX_vec_cand) -
                    lG0_ηlδx(μ_x_h_old, nothing, lδ_x_h_old, model.state) +
                    sum(model.state.lωNX_vec)
        end

        ## Decision and update
        if not_auto_reject && log(rand()) < lar # accept

            model.state.μ_x[h,:] = deepcopy(μ_x_h_cand)
            model.state.δ_x[h,:] = deepcopy(δ_x_h_cand)

            model.state.lNX[:,h] = deepcopy(lNX_mat_cand[:,h])
            model.state.lωNX_vec = deepcopy(lωNX_vec_cand)

            ηlδ_x_out = deepcopy(ηlδ_x_cand)
            model.state.accpt[h] += 1

        else # reject
            ηlδ_x_out = deepcopy(ηlδ_x_old)
        end

        ## Full conditional draw (from G0) for η_y
        if update_ηy
            model.state.δ_y[h] = rand(InverseGamma(0.5 * model.state.ν_δy,
                                        0.5 * model.state.ν_δy * model.state.s0_δy))
            βstar_ηy = model.state.β0star_ηy + (model.state.Λ0star_ηy.chol.U \ randn((model.K + 1)) .* sqrt(model.state.δ_y[h]))
            model.state.μ_y[h] = deepcopy(βstar_ηy[model.indx_ηy[:μ]])
            model.state.β_y[h,:] = deepcopy(βstar_ηy[model.indx_ηy[:β]])
        end
    else # some observations are assigned to cluster h (i.e., n_h > 0)

        ## Metropolis step for η_x

        ## Candidate, lNX, lωNXvec, previously computed
        ## Important quantities

        y_h = model.y[indx_h] # doesn't need deepcopy if indexing (result of call to getindex)
        X_h = model.X[indx_h,:]
        D_h_old = construct_Dh(h, X_h, model.state.μ_x[h,:], γ_h)

        Λ1star_ηy_h_old = get_Λ1star_ηy_h(D_h_old, model.state.Λ0star_ηy)
        β1star_ηy_h_old = get_β1star_ηy_h(Λ1star_ηy_h_old, Λβ0star_ηy, D_h_old, y_h)

        a1_δy_old = get_a1_δy_h(model.state.ν_δy, n_h) # posterior IG shape
        b1_δy_old = get_b1_δy_h(model.state.ν_δy, model.state.s0_δy, y_h, βΛβ0star_ηy, Λ1star_ηy_h_old, β1star_ηy_h_old) # posterior IG scale

        ## Important quantities for candidate
        if not_auto_reject
            D_h_cand = construct_Dh(h, X_h, μ_x_h_cand, γ_h)

            Λ1star_ηy_h_cand = get_Λ1star_ηy_h(D_h_cand, model.state.Λ0star_ηy)
            β1star_ηy_h_cand = get_β1star_ηy_h(Λ1star_ηy_h_cand, Λβ0star_ηy, D_h_cand, y_h)

            a1_δy_cand = get_a1_δy_h(model.state.ν_δy, n_h) # posterior IG shape
            b1_δy_cand = get_b1_δy_h(model.state.ν_δy, model.state.s0_δy, y_h, βΛβ0star_ηy, Λ1star_ηy_h_cand, β1star_ηy_h_cand) # posterior IG scale

            ## Compute acceptance ratio (lcc_ηlδx does not need γ-modified βx and δx)
            lar = lcc_ηlδx(h, indx_h, lNX_mat_cand, lωNX_vec_cand,
                            model.state, μ_x_h_cand, nothing, lδ_x_h_cand,
                            Λ1star_ηy_h_cand, a1_δy_cand, b1_δy_cand) -
                lcc_ηlδx(h, indx_h, model.state.lNX, model.state.lωNX_vec,
                            model.state, μ_x_h_old, nothing, lδ_x_h_old,
                            Λ1star_ηy_h_old, a1_δy_old, b1_δy_old)

        end

        ## Decision and update
        if not_auto_reject && log(rand()) < lar # accept

            model.state.μ_x[h,:] = deepcopy(μ_x_h_cand)
            model.state.δ_x[h,:] = deepcopy(δ_x_h_cand)

            model.state.lNX[:,h] = deepcopy(lNX_mat_cand[:,h])
            model.state.lωNX_vec = deepcopy(lωNX_vec_cand)

            ηlδ_x_out = deepcopy(ηlδ_x_cand)
            model.state.accpt[h] += 1

            ## Full conditional draw for η_y ( use prec. prameterization )
            if update_ηy
                model.state.δ_y[h] = rand(InverseGamma(a1_δy_cand, b1_δy_cand))
                βstar_ηy = β1star_ηy_h_cand + (Λ1star_ηy_h_cand.chol.U \ randn((model.K + 1)) .* sqrt(model.state.δ_y[h]))
                model.state.μ_y[h] = deepcopy(βstar_ηy[model.indx_ηy[:μ]])
                model.state.β_y[h,:] = deepcopy(βstar_ηy[model.indx_ηy[:β]])
            end
        else # reject

            ηlδ_x_out = deepcopy(ηlδ_x_old)

            ## Full conditional draw for η_y
            if update_ηy
                model.state.δ_y[h] = rand(InverseGamma(a1_δy_old, b1_δy_old))
                βstar_ηy = β1star_ηy_h_old + (Λ1star_ηy_h_old.chol.U \ randn((model.K + 1)) .* sqrt(model.state.δ_y[h]))
                model.state.μ_y[h] = deepcopy(βstar_ηy[model.indx_ηy[:μ]])
                model.state.β_y[h,:] = deepcopy(βstar_ηy[model.indx_ηy[:β]])
            end
        end
    end

    if model.state.adapt && (model.state.iter % model.state.adapt_thin == 0)
        model.state.adapt_iter += 1
        model.state.runningsum_ηlδx[h,:] += ηlδ_x_out
        runningmean = model.state.runningsum_ηlδx[h,:] / float(model.state.adapt_iter)
        runningdev = ( ηlδ_x_out - runningmean )
        model.state.runningSS_ηlδx[h,:,:] = runningdev * runningdev' # the mean is changing, but this approx. is fine.
    end

    return nothing
end

function update_η_h_Met_K1!(model::Model_BNP_WMReg_Joint, h::Int, Λβ0star_ηy::Array{T,1}, βΛβ0star_ηy::T; update_ηy::Bool=true) where T <: Real

    ## if you modify this function, you must modify the originals above.

    indx_h = findall(model.state.S .== h)
    n_h = length(indx_h)

    μ_x_h_old = model.state.μ_x[h,:]
    lδ_x_h_old = log.(model.state.δ_x[h,:])
    ηlδ_x_old = vcat(μ_x_h_old, lδ_x_h_old)

    ## Draw candidate for η_x
    ηlδ_x_cand = ηlδ_x_old + rand(MvNormal(model.state.cSig_ηlδx[h]))
    μ_x_h_cand = ηlδ_x_cand[model.indx_ηx[:μ]]
    lδ_x_h_cand = ηlδ_x_cand[model.indx_ηx[:δ]]
    δ_x_h_cand = exp.(lδ_x_h_cand)

    auto_reject = (!all(δ_x_h_cand .> 0.0)) || any(isinf.(δ_x_h_cand))
    not_auto_reject = !auto_reject

    lNX_mat_cand = deepcopy(model.state.lNX)

    if typeof(model.state.γ) == BitArray{1}
        γ_h = deepcopy(model.state.γ)
    elseif typeof(model.state.γ) == BitArray{2}
        γ_h = deepcopy(model.state.γ[h,:])
    end

    ## bookkeeping for variable selection
    if not_auto_reject
        if model.state.γδc == Inf || isnothing(model.state.γδc) # subset or integration method
            if γ_h[1]
                lNX_mat_cand[:,h] = logpdf.(Normal(μ_x_h_cand[1], sqrt(δ_x_h_cand[1])), vec(model.X))
            else
                lNX_mat_cand[:,h] = zeros(Float64, model.n)
            end
        else # variance-inflation method
            δγ_x_h_cand = δ_x_h_modify_γ(δ_x_h_cand, γ_h, model.state.γδc)
            lNX_mat_cand[:,h] = logpdf.(Normal(μ_x_h_cand[1], sqrt(δγ_x_h_cand[1])), vec(model.X))
        end

        lωNX_vec_cand = lωNXvec(model.state.lω, lNX_mat_cand) # n vector

    end

    if n_h == 0

        ## Metropolis step for η_x

        ## Compute acceptance ratio
        if not_auto_reject
            lar = lG0_ηlδx(μ_x_h_cand[1], lδ_x_h_cand[1], model.state) -
                    sum(lωNX_vec_cand) -
                    lG0_ηlδx(μ_x_h_old[1], lδ_x_h_old[1], model.state) +
                    sum(model.state.lωNX_vec)
        end

        ## Decision and update
        if not_auto_reject && log(rand()) < lar # accept

            model.state.μ_x[h,:] = deepcopy(μ_x_h_cand)
            model.state.δ_x[h,:] = deepcopy(δ_x_h_cand)

            model.state.lNX[:,h] = deepcopy(lNX_mat_cand[:,h])
            model.state.lωNX_vec = deepcopy(lωNX_vec_cand)

            ηlδ_x_out = deepcopy(ηlδ_x_cand)
            model.state.accpt[h] += 1

        else # reject
            ηlδ_x_out = deepcopy(ηlδ_x_old)
        end

        ## Full conditional draw (from G0) for η_y
        if update_ηy
            model.state.δ_y[h] = rand(InverseGamma(0.5 * model.state.ν_δy,
                                        0.5 * model.state.ν_δy * model.state.s0_δy))
            βstar_ηy = model.state.β0star_ηy + (model.state.Λ0star_ηy.chol.U \ randn((model.K + 1)) .* sqrt(model.state.δ_y[h]))
            model.state.μ_y[h] = deepcopy(βstar_ηy[model.indx_ηy[:μ]])
            model.state.β_y[h,:] = deepcopy(βstar_ηy[model.indx_ηy[:β]])
        end
    else # some observations are assigned to cluster h (i.e., n_h > 0)

        ## Metropolis step for η_x

        ## Candidate, lNX, lωNXvec, previously computed
        ## Important quantities

        y_h = model.y[indx_h] # doesn't need to copy if indexing (result of call to getindex)
        X_h = model.X[indx_h,:]
        D_h_old = construct_Dh(h, X_h, model.state.μ_x[h,:], γ_h)

        Λ1star_ηy_h_old = get_Λ1star_ηy_h(D_h_old, model.state.Λ0star_ηy)
        β1star_ηy_h_old = get_β1star_ηy_h(Λ1star_ηy_h_old, Λβ0star_ηy, D_h_old, y_h)

        a1_δy_old = get_a1_δy_h(model.state.ν_δy, n_h) # posterior IG shape
        b1_δy_old = get_b1_δy_h(model.state.ν_δy, model.state.s0_δy, y_h, βΛβ0star_ηy, Λ1star_ηy_h_old, β1star_ηy_h_old) # posterior IG scale

        if not_auto_reject
            ## Important quantities for candidate
            D_h_cand = construct_Dh(h, X_h, μ_x_h_cand, γ_h)

            Λ1star_ηy_h_cand = get_Λ1star_ηy_h(D_h_cand, model.state.Λ0star_ηy)
            β1star_ηy_h_cand = get_β1star_ηy_h(Λ1star_ηy_h_cand, Λβ0star_ηy, D_h_cand, y_h)

            a1_δy_cand = get_a1_δy_h(model.state.ν_δy, n_h) # posterior IG shape
            b1_δy_cand = get_b1_δy_h(model.state.ν_δy, model.state.s0_δy, y_h, βΛβ0star_ηy, Λ1star_ηy_h_cand, β1star_ηy_h_cand) # posterior IG scale

            ## Compute acceptance ratio (lcc_ηlδx does not need γ-modified δx)
            lar = lcc_ηlδx(h, indx_h, lNX_mat_cand, lωNX_vec_cand,
                            model.state, μ_x_h_cand[1], lδ_x_h_cand[1],
                            Λ1star_ηy_h_cand, a1_δy_cand, b1_δy_cand) -
                        lcc_ηlδx(h, indx_h, model.state.lNX, model.state.lωNX_vec,
                            model.state, μ_x_h_old[1], lδ_x_h_old[1],
                            Λ1star_ηy_h_old, a1_δy_old, b1_δy_old)
        end

        ## Decision and update
        if not_auto_reject && log(rand()) < lar  # accept

            model.state.μ_x[h,:] = deepcopy(μ_x_h_cand)
            model.state.δ_x[h,:] = deepcopy(δ_x_h_cand)

            model.state.lNX[:,h] = deepcopy(lNX_mat_cand[:,h])
            model.state.lωNX_vec = deepcopy(lωNX_vec_cand)

            ηlδ_x_out = deepcopy(ηlδ_x_cand)
            model.state.accpt[h] += 1

            ## Full conditional draw for η_y ( use prec. prameterization )
            if update_ηy
                model.state.δ_y[h] = rand(InverseGamma(a1_δy_cand, b1_δy_cand))
                βstar_ηy = β1star_ηy_h_cand + (Λ1star_ηy_h_cand.chol.U \ randn((model.K + 1)) .* sqrt(model.state.δ_y[h]))
                model.state.μ_y[h] = deepcopy(βstar_ηy[model.indx_ηy[:μ]])
                model.state.β_y[h,:] = deepcopy(βstar_ηy[model.indx_ηy[:β]])
            end
        else # reject

            ηlδ_x_out = deepcopy(ηlδ_x_old)

            ## Full conditional draw for η_y
            if update_ηy
                model.state.δ_y[h] = rand(InverseGamma(a1_δy_old, b1_δy_old))
                βstar_ηy = β1star_ηy_h_old + (Λ1star_ηy_h_old.chol.U \ randn((model.K + 1)) .* sqrt(model.state.δ_y[h]))
                model.state.μ_y[h] = deepcopy(βstar_ηy[model.indx_ηy[:μ]])
                model.state.β_y[h,:] = deepcopy(βstar_ηy[model.indx_ηy[:β]])
            end
        end
    end

    if model.state.adapt && (model.state.iter % model.state.adapt_thin == 0)
        model.state.adapt_iter += 1
        model.state.runningsum_ηlδx[h,:] += ηlδ_x_out
        runningmean = model.state.runningsum_ηlδx[h,:] / float(model.state.adapt_iter)
        runningdev = ( ηlδ_x_out - runningmean )
        model.state.runningSS_ηlδx[h,:,:] = runningdev * runningdev' # the mean is changing, but this approx. is fine.
    end

    return nothing
end

function construct_Dh(h::Int, Xh::Array{T,2}, μ_x_h::Array{T,1}) where T <: Real
    D0 = broadcast(-, permutedims(μ_x_h), Xh)
    return hcat(ones(size(Xh, 1)), D0)
end
function construct_Dh(h::Int, Xh::Array{T,2}, μ_x_h::Array{T,1},
                      γ::BitArray{1}) where T <: Real
    D0 = broadcast(-, permutedims(μ_x_h), Xh)
    D0[:, findall(.!(γ))] *= 0.0
    return hcat(ones(size(Xh, 1)), D0)
end

function lG0_ηlδx(μ_x::Array{T,1}, β_x::Union{Array{Array{T,1},1}, Nothing},
                  lδ_x::Array{T,1},
                  state::State_BNP_WMReg_Joint) where T <: Real

    K = length(μ_x)
    Q_μ = - 0.5 * PDMats.quad(state.Λ0_μx, (μ_x - state.μ0_μx))
    Q_β = 0.0
    Q_δ = 0.0
    if !isnothing(β_x)
        for k = 1:(K - 1)
            Q_β -= 0.5 * PDMats.quad(state.Λ0_βx[k], (β_x[k] - state.β0_βx[k]))
            # Q_δ -= ( (0.5*state.ν_δx[k] + 1.0)*lδ_x[k] +
            #           0.5*state.ν_δx[k]*state.s0_δx[k]/exp(lδ_x[k]) ) # build Jacobian into it
            # Q_δ -= 0.5*( state.ν_δx[k]*lδ_x[k] + state.ν_δx[k]*state.s0_δx[k]/exp(lδ_x[k]) ) # Jacobian built in
        end
    end
    # Q_δ -= ( (0.5*state.ν_δx[K] + 1.0)*lδ_x[K] +
    #           0.5*state.ν_δx[K]*state.s0_δx[K]/exp(lδ_x[K]) )
    for k = 1:K
        Q_δ -= 0.5 * ( state.ν_δx[k] * lδ_x[k] + state.ν_δx[k] * state.s0_δx[k] / exp(lδ_x[k]) ) # Jacobian built in
    end
    # J = sum(lδ_x) # Jacobian for log(δ) transformation
    # Jacobian now built in

    Q = Q_μ + Q_β + Q_δ # + J
    return Q
end
function lG0_ηlδx(μ_x::T, lδ_x::T,
                  state::State_BNP_WMReg_Joint) where T <: Real
    ## K = 1 case
    Q_μ = - 0.5 * Matrix(state.Λ0_μx)[1] * (μ_x - state.μ0_μx[1])^2.0
    Q_δ = - 0.5 * ( state.ν_δx[1] * lδ_x + state.ν_δx[1] * state.s0_δx[1] / exp(lδ_x) ) # Jacobian built in

    Q = Q_μ + Q_δ
    return Q
end

### log-collapsed conditional
function lcc_ηlδx(h::Int, indx_h::Array{Int,1}, lNX_mat::Array{T,2}, lωNX_vec::Array{T,1},
    state::State_BNP_WMReg_Joint, μ_x_h::Array{T,1}, β_x_h::Union{Array{Array{T,1},1}, Nothing},
    lδ_x_h::Array{T,1},
    Λ1star_ηy_h::PDMat{T}, a1_δy::T, b1_δy::T) where T <: Real

    # Q1
    Q1 = sum(lNX_mat[indx_h, h])

    # Q2
    Q2 = sum(lωNX_vec)

    # Q3 , G0 including Jacobian
    Q3 = lG0_ηlδx(μ_x_h, β_x_h, lδ_x_h, state)

    # Q4
    Q4 = logdet(Λ1star_ηy_h)

    # Q5
    Q5 = log(b1_δy)

    return Q1 - Q2 + Q3 - 0.5 * Q4 - (a1_δy) * Q5

end
function lcc_ηlδx(h::Int, indx_h::Array{Int,1}, lNX_mat::Array{T,2}, lωNX_vec::Array{T,1},
    state::State_BNP_WMReg_Joint, μ_x_h::T, lδ_x_h::T,
    Λ1star_ηy_h::PDMat{T}, a1_δy::T, b1_δy::T) where T <: Real

    ## K = 1 case

    # Q1
    Q1 = sum(lNX_mat[indx_h, h])

    # Q2
    Q2 = sum(lωNX_vec)

    # Q3 , G0 including Jacobian
    Q3 = lG0_ηlδx(μ_x_h, lδ_x_h, state)

    # Q4
    Q4 = logdet(Λ1star_ηy_h)

    # Q5
    Q5 = log(b1_δy)

    return Q1 - Q2 + Q3 - 0.5 * Q4 - (a1_δy) * Q5

end

## variance-inflation method
function βδ_x_h_modify_γ(β_x::Array{Array{T,1},1}, δ_x::Array{T,1},
                        γ::BitArray{1}, γδc::Array{T,1}) where T <: Real
    modify_indx = findall(.!(γ))
    K = length(γ)

    βout = deepcopy(β_x)
    δout = deepcopy(δ_x)

    for k in modify_indx # this works even if no modifications are necessary
        δout[k] += γδc[k]
    end

    for k = 1:(K - 1)
        if !γ[k]
            βout[k] *= 0.0
        else
            modify2 = intersect((k + 1):K, modify_indx)
            βout[k][(modify2 .- k)] *= 0.0
        end
    end

    return βout, δout
end
function δ_x_h_modify_γ(δ_x::Array{T,1},
                      γ::BitArray{1}, γδc::Array{T,1}) where T <: Real
    modify_indx = findall(.!(γ))
    δout = deepcopy(δ_x)

    for k in modify_indx # this works even if no modifications are necessary
        δout[k] += γδc[k]
    end
    return δout
end

## subset method
function βδ_x_h_modify_γ(β_x::Array{Array{T,1},1}, δ_x::Array{T,1},
                        γ::BitArray{1}, γδc::Float64) where T <: Real

    γδc == Inf || throw("A single variance inflation should be equal to Inf")

    γindx = findall(γ)
    nγ = length(γindx)
    nγ > 1 || throw("βδ_x_h_modify_γ requires more than one selected variable.")

    βout = [ deepcopy(β_x[γindx[k]][(γindx[(k + 1):nγ] .- γindx[k]) ])  for k = 1:(nγ - 1) ] # vector of (nγ-1):1 length vectors
    δout = deepcopy(δ_x[γindx]) # H nows and sum(gamma) cols

    return βout, δout
end
function δ_x_h_modify_γ(δ_x::Array{T,1},
                      γ::BitArray{1}, γδc::Float64) where T <: Real
    γδc == Inf || throw("A single variance inflation should be equal to Inf")
    γindx = findall(γ)
    δout = deepcopy(δ_x[γindx]) # H nows and sum(gamma) cols

    return δout
end

## helper functions for marginalization
function get_a1_δy_h(ν_δy::T, n_h::Int) where T <: Real
    return 0.5 * (ν_δy + n_h)
end

## must do no more with y_h than y'y because of its use in update_alloc!()
function get_b1_δy_h(ν_δy::T, s0_δy::T, y_h::Array{T,1}, βΛβ0star_ηy::T, Λ1star_ηy_h::PDMat, β1star_ηy_h::Array{T,1}) where T <: Real
    return 0.5 * (ν_δy * s0_δy + y_h'y_h + βΛβ0star_ηy - PDMats.quad(Λ1star_ηy_h, β1star_ηy_h))
end

function get_Λ1star_ηy_h(D_h::Array{T,2}, Λ0star_ηy::PDMat) where T <: Real
    return PDMat_adj(D_h'D_h + Λ0star_ηy)
end

function get_β1star_ηy_h(Λ1star_ηy_h::PDMat, Λβ0star_ηy::Array{T,1}, D_h::Array{T,2}, y_h::Array{T,1}) where T <: Real
    return Λ1star_ηy_h \ (Λβ0star_ηy + D_h'y_h)
end


state(prob::AbstractProblem, net, args...) = state(input_type(net), prob, net, args...)
state(T::Type{<:Number}, prob::AbstractProblem, net) = state(T, basis(prob), net)
#state(T::Type{<:Number}, prob::LRhoKLocalOpProblem, net) = state_lut(T, basis(prob), net)
state_lut(prob::AbstractProblem, net, args...) = state_lut(input_type(net), prob, net, args...)
state_lut(T::Type{<:Number}, prob::AbstractProblem, net) = state_lut(T, basis(prob), net)

function state(T::Type{<:Number}, hilb::Basis, net)
    !is_homogeneous(hilb) && error("Could not generate a state.")
    sys_state = _homogeneous_system_state(T, first(hilb.bases),
                                          length(hilb.bases), net)

    state = _network_state(sys_state, net)
    return state
end

function state_lut(T::Type{<:Number}, hilb::Basis, net)
    !is_homogeneous(hilb) && error("Could not generate a state.")
    sys_state = _homogeneous_system_state(T, first(hilb.bases),
                                          length(hilb.bases), net)

    state = _lut_state(sys_state, net)
    return state
end

function _lut_state(T::Type{<:Number}, hilb::Basis, net::MatrixNet)
    bare_state = state(T, basis(prob), net)
    state = DoubleState(ModifiedState(row(bare_state)))
    lut   = lookup(net)

    isnothing(lut) && return state
    return LUState(state, lut)
end

function _homogeneous_system_state(T::Type{<:Number}, hilb::Basis, nsites, net::Union{MatrixNet, KetNet})
    loc_size = first(hilb.shape)
    return NAryState(T, loc_size, nsites)
end

_network_state(sys_state::State, net::Union{MatrixNet, PureNet}) =
    DoubleState(sys_state)

_network_state(sys_state::State, net::KetNet) =
    return sys_state

function _lut_state(sys_state::State, net::Union{MatrixNet, PureNet})
    state = DoubleState(ModifiedState(sys_state))
    lut   = lookup(net)

    isnothing(lut) && return state
    return LUState(state, lut)
end


"""
    Contact{T,N}

    Abstract type contain contact information associated with a parent Body object.
"""
abstract type Contact{T,N} end

# constraints
constraint(model::Contact, body::Body, λ, timestep) = constraint(model, next_configuration(body.state, timestep)..., λ)

# constraint Jacobians
constraint_jacobian_configuration(model::Contact, body::Body, id, λ, timestep) = constraint_jacobian_configuration(model, next_configuration(body.state, timestep)..., current_configuration_velocity(body.state)..., λ, timestep)
constraint_jacobian_velocity(model::Contact, body::Body, id, λ, timestep) = constraint_jacobian_velocity(model, next_configuration(body.state, timestep)..., current_configuration_velocity(body.state)..., λ, timestep)

# impulses
impulse_map(model::Contact, body::Body, id, λ, timestep) = impulse_map(model, next_configuration(body.state, timestep)..., λ)

function impulse_map(model::Contact, x::AbstractVector, q::Quaternion, λ)
    X = force_mapping(model, x, q)
    # Q = - X * q * skew(model.contact_point - vector_rotate(model.offset, inv(q)))
    Q = - X * rotation_matrix(q) * skew(model.contact_point - vector_rotate(model.offset, inv(q)))
    return [X'; Q']
end

# force mapping
function force_mapping(model::Contact, x::AbstractVector, q::Quaternion)
    X = [model.surface_normal_projector;
         szeros(1,3);
         model.surface_projector]
    return X
end

# utilities
Base.length(model::Contact{T,N}) where {T,N} = N
neutral_vector(model::Contact{T,N}) where {T,N} = sones(T, Int(N/2))
cone_degree(model::Contact{T,N}) where {T,N} = Int(N/2)


# module containing patterns used for Metamaterial waveguide structure
module MetamaterialWaveguide
    using Devices, Clipper
    import Devices: Meta, GDSMeta
    export cross_polygon, mwg_unit_cell!, mwg_last_cell!, mwg_input_cap!, mwg_cell_with_claw!

    """
        cross_polygon(hline_width, hline_height, vline_width,
        vline_height; center=true)
    Create a cross-shaped polygon of specified widths and heights of horizontal
    line and vertical line. The keyword argument `center=true` determines whether
    the returned polygon is centered at the origin or at distance hline_width/2
    along the x-axis away from the origin.
    """
    function cross_polygon(hline_width, hline_height, vline_width,
        vline_height; center=true)

        if hline_width < vline_width
            error("The width of the horizontal line must be larger than that of
                the vertical line")
        elseif hline_height > vline_height
            error("The height of the vertical line must be larger than that of
                the horizontal line")
        end

        offset =  (~center) * Point(hline_width/2, typeof(hline_width)(0))
        hline = centered(Rectangle(hline_width, hline_height)) + offset
        vline = centered(Rectangle(vline_width, vline_height)) + offset

        poly = clip(Clipper.ClipTypeUnion, hline, vline)[1]
        return poly
    end

    """
        mwg_unit_cell!{T,S}(c::Cell{T,S}, Lr_width, Lr_height, Cr_overlap,
        res_ground_length, res_conductor_width, res_conductor_gap,
        Ck_terminal_width, Ck_overlap, wg_meander_length, wg_meander_nlegs::Int,
        wg_ground_gap, wg_conductor_width, wg_conductor_gap, centerline_width,
        centerline_length, meta=GDSMeta(0,0))
    Create a unit cell of the metamaterial waveguide.
    - `c` : a cell to render metamaterial waveguide in
    - `Lr_width` : width of the rectangular coil Lr
    - `Lr_height` : height of the rectangular coil Lr
    - `Cr_overlap` : overlap length of the interdigital capacitor Cr forming
        the resonator
    - `res_ground_length` : length of the line from the top of the resonator
        to the ground. This also sets the gap between the resonator
        and the ground.
    - `res_conductor_width` : width of the conductor forming a resonator
    - `res_conductor_gap` : gap between conductor traces forming a resonator
    - `Ck_terminal_width` : terminal width of the coupling capacitor Ck between
        the resonator and the centerline
    - `Ck_overlap` : overlap length of the coupling capacitor Ck
    - `wg_meander_length` : the length of meandering of waveguide. The trace of
        the meander goes back and forth vertically with this distance.
    - `wg_meander_nlegs` : number of legs in the meander forming the waveguide
    - `wg_ground_gap` : the distance between the meander trace of the waveguide
        and the ground
    - `wg_conductor_width` : width of the conductor forming the meander
    - `wg_conductor_gap` : gap between conductor traces forming the meander
    - `centerline_width` : width of the centerline that connects the resonator
        with waveguide.
    - `centerline_length` : length of the centerline. This has to be equal to or
        greater than `Lr_width`
    """
    function mwg_unit_cell!(c::Cell{T,S}, Lr_width, Lr_height, Cr_overlap,
        res_ground_length, res_conductor_width, res_conductor_gap, Ck_terminal_width,
        Ck_overlap, wg_meander_length, wg_meander_nlegs::Int, wg_ground_gap,
        wg_conductor_width, wg_conductor_gap, centerline_width, centerline_length,
        meta=GDSMeta(0,0)) where {T,S}
        if Lr_width > centerline_length
            error("The length of the centerline must be larger than the width of the coil")
        end
        wg_length = (wg_meander_nlegs - 1) * (wg_conductor_width + wg_conductor_gap) + wg_conductor_width / 2
        res_height = 2 * (Lr_height + res_conductor_gap) + Cr_overlap

        # define a mask for unit cell
        wg_mask_width = 2 * wg_length + centerline_length
        wg_mask_height = wg_meander_length + 2 * (wg_conductor_width + wg_ground_gap)

        res_mask_width = Lr_width + 2 * res_ground_length
        res_mask_height = centerline_width + 2 * (res_ground_length + res_height + Ck_terminal_width + Ck_overlap + 3 * res_conductor_gap)

        # global mask = res_mask ∪ wg_mask
        # mask = Cell{T,S}(uniquename("mask"))
        # mask_poly = cross_polygon(wg_mask_width, wg_mask_height, res_mask_width,
            # res_mask_height; center=false)
        # render!(mask, mask_poly)

        # lumped-element resonator with coupling capacitor
        res_with_Ck = Cell{T,S}(uniquename("lumped_elem_resonator_with_Ck"))
        lumped_elem_resonator_with_Ck!(res_with_Ck, Lr_width, Lr_height, Cr_overlap, Ck_overlap,
            Ck_terminal_width, res_conductor_width, res_conductor_gap, meta,
            ground_length=res_ground_length, Ck_terminal_width2=zero(T))

        res = Cell{T,S}(uniquename("resonator"))
        push!(res.refs, CellReference(res_with_Ck, Point(wg_mask_width/2, centerline_width/2)))
        push!(res.refs, CellReference(res_with_Ck,
            Point(wg_mask_width/2, -centerline_width/2), xrefl=true))

        # meander forming waveguide part of MWG
        meander = Cell{T,S}(uniquename("meander"))
        rectangular_meander!(meander, wg_meander_length, wg_meander_nlegs,
            wg_conductor_width, wg_conductor_gap, meta;
            first_length=wg_meander_length + wg_conductor_width,
            last_length=(wg_meander_length + centerline_width)/2,
            start_upward=true)

        wg = Cell{T,S}(uniquename("waveguide"))
        push!(wg.refs, CellReference(meander, Point(-wg_conductor_width/2,
            -wg_conductor_width-wg_meander_length/2)))
        push!(wg.refs, CellReference(meander,
            Point(wg_conductor_width/2+ wg_mask_width,
            wg_conductor_width + wg_meander_length/2), rot=π))

        # centerline
        centerline = Cell{T,S}(uniquename("centerline"))
        render!(centerline, centered(Rectangle(centerline_length, centerline_width)) +
            Point(wg_mask_width/2, zero(T)), meta)

        # # take the negative trace of each element
        # res_neg = negative(res, mask)
        # wg_neg = negative(wg, mask)
        # centerline_neg = negative(centerline, mask)

        # take the intersection res_neg ∩ wg_neg ∩ centerline_neg
        # unit_cell_poly = clip(Clipper.ClipTypeIntersection,
        #     polygon.(elements(res_neg)), polygon.(elements(wg_neg)))
        # unit_cell_poly = clip(Clipper.ClipTypeIntersection,
        #     unit_cell_poly, polygon.(elements(centerline_neg)))
        #
        # for p in unit_cell_poly
        #     render!(c, p, meta)
        # end
        for part in [res, wg, centerline]
            push!(c.refs, CellReference(part, Point(zero(T), zero(T))))
        end
        return c
    end

    """
    Metamaterial waveguide cell with claw coming out from the resonator node
    """
    function mwg_cell_with_claw!(c::Cell{T,S}, Lr_width, Lr_height, Cr_overlap,
        res_ground_length, res_conductor_width, res_conductor_gap, Ck_terminal_width,
        Ck_overlap, wg_meander_length, wg_meander_nlegs::Int, wg_ground_gap,
        wg_conductor_width, wg_conductor_gap, centerline_width, centerline_length,
        res_claw_gap, res_claw_conductor_width, claw_width, claw_length,
        claw_ground_gap, claw_dist, meta::Meta=GDSMeta(0,0); rot=0) where {T,S}

        wg_length = ((wg_meander_nlegs - 1) * (wg_conductor_width + wg_conductor_gap)
            + wg_conductor_width/2)
        res_height = 2 * (Lr_height + res_conductor_gap) + Cr_overlap

        # define a mask for unit cell
        wg_mask_width = 2 * wg_length + centerline_length
        wg_mask_height = wg_meander_length + 2 * (wg_conductor_width + wg_ground_gap)

        res_mask_width = Lr_width + 2 * res_ground_length
        res_mask_height = centerline_width + 2 * (res_ground_length + res_height
            + Ck_terminal_width + Ck_overlap + 3 * res_conductor_gap)

        claw_mask_left = wg_mask_width/2 + Lr_width/2
        claw_mask_right = (claw_mask_left + res_claw_gap + res_claw_conductor_width
            + claw_length + claw_ground_gap)
        claw_mask_vdist = claw_dist/2 - claw_ground_gap
        claw_mask_hdist = res_claw_gap + res_claw_conductor_width + claw_ground_gap

        wg_mask = Rectangle(Point(zero(T), -wg_mask_height/2),
            Point(wg_mask_width/2, wg_mask_height/2))
        res_mask = centered(Rectangle(res_mask_width, res_mask_height)) + Point(wg_mask_width/2, zero(T))

        claw_mask_bulk = Rectangle(Point(claw_mask_left, -res_mask_height/2),
            Point(claw_mask_right, res_mask_height/2))
        claw_mask_rm = Rectangle(Point(claw_mask_left + claw_mask_hdist,
            -claw_mask_vdist), Point(claw_mask_right, claw_mask_vdist))
        claw_mask = clip(Clipper.ClipTypeDifference, claw_mask_bulk, claw_mask_rm)[1]

        # global mask = res_mask ∪ wg_mask ∪ claw_mask
        # mask = Cell{T,S}(uniquename("mask"))
        # mask_poly = clip(Clipper.ClipTypeUnion, wg_mask, res_mask)[1]
        # mask_poly = clip(Clipper.ClipTypeUnion, mask_poly, claw_mask)[1]
        # render!(mask, mask_poly, meta_mask)

        # lumped-element resonator with coupling capacitor
        res_with_Ck = Cell{T,S}(uniquename("lumped_elem_resonator_with_Ck"))
        lumped_elem_resonator_with_Ck!(res_with_Ck, Lr_width, Lr_height, Cr_overlap, Ck_overlap,
        Ck_terminal_width, res_conductor_width, res_conductor_gap, meta,
        ground_length=res_ground_length, Ck_terminal_width2=zero(T))

        res = Cell{T,S}(uniquename("resonator"))
        push!(res.refs, CellReference(res_with_Ck,
            Point(wg_mask_width/2, centerline_width/2)))
        push!(res.refs, CellReference(res_with_Ck,
            Point(wg_mask_width/2, -centerline_width/2), xrefl=true))

            # meander forming waveguide part of MWG
        meander = Cell{T,S}(uniquename("meander"))
        rectangular_meander!(meander, wg_meander_length, wg_meander_nlegs,
            wg_conductor_width, wg_conductor_gap, meta;
            first_length=wg_meander_length + wg_conductor_width,
            last_length=(wg_meander_length + centerline_width)/2,
            start_upward=true)

        wg = Cell{T,S}(uniquename("waveguide"))
        push!(wg.refs, CellReference(meander, Point(-wg_conductor_width/2,
            -wg_conductor_width-wg_meander_length/2)))

        # centerline
        centerline = Cell{T,S}(uniquename("centerline"))
        render!(centerline, centered(Rectangle(centerline_length, centerline_width)) +
            Point(wg_mask_width/2, zero(T)), meta)

        # claw
        claw = Cell{T,S}(uniquename("claw"))
        res_claw_conn_vdist = (res_mask_height/2 - Lr_height - res_ground_length
            - 2 * res_conductor_gap - Cr_overlap)
        render!(claw, Rectangle(Point(claw_mask_left, res_claw_conn_vdist - 2*res_conductor_width),
            Point(claw_mask_left + res_claw_gap, res_claw_conn_vdist)), meta)
        render!(claw, Rectangle(Point(claw_mask_left, -res_claw_conn_vdist + 2*res_conductor_width),
            Point(claw_mask_left + res_claw_gap, -res_claw_conn_vdist)), meta)

            claw_left = claw_mask_left + res_claw_gap + res_claw_conductor_width
            claw_right = claw_left + claw_length

        render!(claw, Rectangle(Point(claw_left - res_claw_conductor_width,
            -res_claw_conn_vdist), Point(claw_left, res_claw_conn_vdist)), meta)

        render!(claw, Rectangle(Point(claw_left, claw_mask_vdist + claw_ground_gap),
            Point(claw_right, claw_mask_vdist + claw_width + claw_ground_gap)), meta)
        render!(claw, Rectangle(Point(claw_left, -claw_mask_vdist - claw_ground_gap),
            Point(claw_right, -claw_mask_vdist - claw_width - claw_ground_gap)), meta)

        # This part discarded due to instability of Clipper
        # take the negative trace of each element
        # res_neg = negative(res, mask)
        # wg_neg = negative(wg, mask)
        # centerline_neg = negative(centerline, mask)
        # claw_neg = negative(claw, mask)
        # take the intersection res_neg ∩ wg_neg ∩ centerline_neg
        # cell_poly = clip(Clipper.ClipTypeIntersection,
        #     polygon.(elements(res_neg)), polygon.(elements(wg_neg)))
        # cell_poly = clip(Clipper.ClipTypeIntersection,
        #     cell_poly, polygon.(elements(centerline_neg)))
        # cell_poly = clip(Clipper.ClipTypeIntersection,
        #     cell_poly, polygon.(elements(claw_neg)))
        # for p in cell_poly
            # render!(c, p, meta)
        # end

        for part in [res, wg, centerline, claw]# ,mask]
            push!(c.refs, CellReference(part, Point(zero(T), zero(T)), rot=rot))
        end

        # render!(c, mask, meta_mask)
        return c
    end

    #
    # """
    # First cell for the metamaterial waveguide including the input capacitor
    # """
    # function mwg_first_cell!(c::Cell{T,S}, Lr_width, Lr_height, Cr_overlap_first,
    #     res_ground_length, res_conductor_width, res_conductor_gap, Ck_terminal_width,
    #     Ck_overlap, wg_meander_length, wg_meander_nlegs::Int, wg_ground_gap,
    #     wg_conductor_width, wg_conductor_gap, centerline_width, centerline_length,
    #     res_claw_gap, res_claw_conductor_width_first, claw_width_first,
    #     claw_length_first, claw_ground_gap_first, input_cap_width, input_cap_gap, ground_width_first,
    #     meta::Meta=GDSMeta(0,0)) where {T,S}
    #     tmp = Cell{T,S}(uniquename("last_cell"))
    #     mwg_last_cell!(tmp, Lr_width, Lr_height, Cr_overlap_first,
    #         res_ground_length, res_conductor_width, res_conductor_gap,
    #         Ck_terminal_width, Ck_overlap, wg_meander_length, wg_meander_nlegs,
    #         wg_ground_gap, wg_conductor_width, wg_conductor_gap,
    #         centerline_width, centerline_length, res_claw_gap,
    #         res_claw_conductor_width_first, claw_width_first, claw_length_first,
    #         claw_ground_gap_first, input_cap_width, input_cap_gap,
    #         ground_width_first, meta)
    #     for p in polygon.(tmp.elements)
    #         render!(c, Rotation(π)(p)+Point(width(bounds(tmp)), zero(T)), meta)
    #     end
    #     return c
    # end

    # """
    # Input capacitor for the metamaterial waveguide
    # """
    # function mwg_input_cap!(c::Cell{T,S}, res_mask_width, res_mask_height,
    #     wg_mask_width, wg_mask_height, wg_meander_length, wg_meander_nlegs::Int,
    #     wg_ground_gap, wg_conductor_width, wg_conductor_gap, centerline_width,
    #     input_cap_terminal_width, input_cap_terminal_width2, input_cap_overlap,
    #     input_cap_finger_width, input_cap_finger_spacing, input_cap_end_gap,
    #     cpw_gap, cpw_trace, meta=GDSMeta(0,0)) where {T,S}
    #
    #     wg_length = (wg_meander_nlegs - 1) * (wg_conductor_width + wg_conductor_gap) + wg_conductor_width/2
    #
    #     input_cap_terminal_length = 2 * wg_conductor_width + wg_meander_length
    #
    #     # define a mask for input capacitor
    #     wg_mask = (centered(Rectangle(wg_mask_width/2, wg_mask_height))
    #         + Point((wg_mask_width/2+res_mask_width)/2, zero(T)))
    #     res_mask = (centered(Rectangle(res_mask_width, res_mask_height))
    #         + Point(res_mask_width/2, zero(T)))
    #
    #     # global mask = res_mask ∪ wg_mask
    #     mask = clip(Clipper.ClipTypeUnion, wg_mask, res_mask)[1]
    #
    #     # meander forming waveguide part of MWG
    #     meander = Cell{T,S}(uniquename("meander"))
    #     rectangular_meander!(meander, wg_meander_length, wg_meander_nlegs,
    #         wg_conductor_width, wg_conductor_gap;
    #         first_length=wg_meander_length + wg_conductor_width,
    #         last_length=(wg_meander_length + centerline_width)/2,
    #         start_upward=true)
    #     wg = Cell{T,S}(uniquename("waveguide"))
    #     push!(wg.refs, CellReference(meander,
    #         Point(wg_conductor_width/2 + wg_mask_width/2 + res_mask_width/2,
    #         wg_conductor_width + wg_meander_length/2), rot=π))
    #
    #     # input capacitor
    #     cap = Cell{T,S}("capacitor")
    #     interdigit = Cell{T,S}("interdigit")
    #     interdigit_capacitor!(interdigit, input_cap_terminal_width, input_cap_terminal_length,
    #         input_cap_overlap, input_cap_finger_width, input_cap_finger_spacing, input_cap_end_gap,
    #         terminal_width2=input_cap_terminal_width2)
    #     push!(cap.refs, CellReference(interdigit, Point(zero(T), -input_cap_terminal_length/2)))
    #
    #     # centerline
    #     conn_line = Cell{T,S}(uniquename("connecting_line"))
    #     render!(conn_line, Rectangle(Point(width(bounds(interdigit)), -centerline_width/2),
    #         Point(bounds(wg).ll.x, centerline_width/2)), meta)
    #     push!(cap.refs, CellReference(conn_line, Point(zero(T), zero(T))))
    #
    #     input_cap_neg = negative(cap, mask)
    #     input_wg_neg = negative(wg, mask)
    #
    #     # take the intersection res_neg ∩ wg_neg ∩ centerline_neg
    #     input_cap_poly = clip(Clipper.ClipTypeIntersection,
    #         polygon.(elements(input_cap_neg)), polygon.(elements(input_wg_neg)))
    #
    #     input_path = Path(Point(typeof(input_cap_terminal_length)(0),
    #         typeof(input_cap_terminal_length)(0)), α0 = 0)
    #     input_taper_len = (wg_mask_width - res_mask_width)/2
    #     input_trace = input_cap_terminal_length
    #     input_gap = input_trace * cpw_gap/cpw_trace
    #
    #     straight!(input_path, input_taper_len, Paths.CPW(
    #         t->(cpw_trace * (input_taper_len - t)/input_taper_len + t/input_taper_len * input_trace),
    #         t->(cpw_gap * (input_taper_len - t)/input_taper_len + t/input_taper_len * input_gap)))
    #
    #     render!(c, input_path, meta)
    #     trans = Translation(Point(input_taper_len, zero(T)))
    #     for p in input_cap_poly
    #         render!(c, trans(p), meta)
    #     end
    #     return c
    # end
end

# package code goes here


module NNHelferlein

import Base: iterate, length, summary
using LinearAlgebra
using Unicode
import ZipFile
import HDF5
import JLD2
using Statistics: mean, std
using ProgressMeter, Dates
using IterTools: ncycle, takenth
import CSV
# import OdsIO    # removed because of PyCall incompatibilities!
import DataFrames
import Random
using Printf
import CUDA
using Knet #: KnetArray, Param, @diff
import Images, Colors
import Augmentor
import MLDataUtils
using TensorBoardLogger, Logging

include("types.jl")
include("util.jl")
include("nets.jl")
include("layers.jl")
include("attn.jl")
include("funs.jl")
include("transformers.jl")
include("images.jl")
include("dataframes.jl")
include("texts.jl")
include("train.jl")
include("imagenet.jl")

DATA_DIR = normpath(joinpath(dirname(pathof(@__MODULE__)), "..", "data"))

export DNN, Classifier, Regressor, Chain,          # chains
       ImageLoader, DataLoader, preproc_imagenet,
       get_class_labels,
       iterate, length,
       Layer, Dense, Conv, Pool, Flat, PyFlat,         # layers
       DeConv, UnPool,
       Embed,
       RSeqClassifier, RSeqTagger,
       hidden_states, cell_states,
       Softmax, Dropout, BatchNorm, LayerNorm,
       Linear,
       AttentionMechanism, AttnBahdanau,
       AttnLuong, AttnDot, AttnLocation,
       AttnInFeed,
       leaky_sigm, leaky_relu, leaky_tanh,
       PositionalEncoding, positional_encoding_sincos,  # transformers
       mk_padding_mask, mk_peek_ahead_mask,
       dot_prod_attn,
       dataframe_read, dataframe_split,         # import data
       dataframe_minibatches, mk_class_ids,
       mk_image_minibatch,
       tb_train!,
       predict_top5, predict_imagenet,
       predict, hamming_dist,
       get_imagenet_classes,                    # images
       image2array, array2image, array2RGB,
       WordTokenizer,                           # texts
       get_tatoeba_corpus,
       seq_minibatch, seq2seq_minibatch,
       crop_array, init0, convert2KnetArray,             # utils
       blowup_array, recycle_array,
       de_embed,
       print_network,
       DATA_DIR

end # module


#
## Sidef base
#

import Base.-,
       Base.+,
       Base.*,
       Base./,
       Base.//,
       Base.%,
       Base.^,
       Base.$,
       Base.&,
       Base.|,
       Base.<=,
       Base.>=,
       Base.<,
       Base.>,
       #Base.==,
       #Base.!=,
       Base.!,
       Base.abs,
       Base.min,
       Base.map,
       Base.sqrt,
       Base.floor,
       Base.print,
       Base.ifelse,
       Base.string,
       Base.println,
       Base.setindex!,
       Base.getindex;

abstract Sidef_Object
abstract Sidef_Types_Number_Complex

immutable Sidef_Types_Nil_Nil <: Sidef_Object
    value::Void
end

immutable Sidef_Types_Bool_Bool <: Sidef_Object
    value::Bool
end

immutable Sidef_Types_Block_Block <: Sidef_Object
    value::Function
end

immutable Sidef_Types_Number_Number <: Sidef_Object
    value::Number
end

immutable Sidef_Types_String_String <: Sidef_Object
    value::AbstractString
end

immutable Sidef_Types_Range_RangeNumber <: Sidef_Object
    value::Range
end

immutable Sidef_Types_Hash_Hash <: Sidef_Object
    value::Dict{Any,Any}
end

immutable Sidef_Types_Array_Array <: Sidef_Object
    value::Array{Any}
end

const NIL = Sidef_Types_Nil_Nil(Void())

const TRUE = Sidef_Types_Bool_Bool(true)
const FALSE = Sidef_Types_Bool_Bool(false)

const MONE = Sidef_Types_Number_Number(-1)
const ZERO = Sidef_Types_Number_Number(0)
const ONE = Sidef_Types_Number_Number(1)

function interpolate(::Type{Sidef_Types_String_String}, a::Any...)
    str = ""
    for item in a
        str *= string(item)
    end
    Sidef_Types_String_String(str)
end

include("object.jl")
include("bool.jl")
include("nil.jl")
include("number.jl")
include("hash.jl")
include("complex.jl")
include("block.jl")
include("range.jl")
include("string.jl")
include("array.jl")
include("convert.jl")


using PackageCompiler
using Glob

using TurbulenceConvection
using CalibrateEDMF
using EnsembleKalmanProcesses
using OrdinaryDiffEq
using NCDatasets
using StochasticDiffEq
using ClimaCore

cedmf = pkgdir(CalibrateEDMF)

if isempty(Glob.glob("CEDMF.so"))
    create_sysimage(
        [
            "TurbulenceConvection",
            "EnsembleKalmanProcesses",
            "OrdinaryDiffEq",
            "NCDatasets",
            "StochasticDiffEq",
            "ClimaCore",
        ];
        sysimage_path = "CEDMF.so",
        precompile_execution_file = joinpath(cedmf, "test", "runtests.jl"),
    )
end


using Viznet, YaoBlocks

using YaoExtensions

struct CircuitCanvas
    frontier::Vector{Int}
end

nline(c::CircuitCanvas) = length(c)

function CircuitCanvas(nqubit::Int, depth::Int)
    CircuitCanvas(nqubit, depth)
end

function frontier(c::CircuitCanvas, start, stop)
    maximum(i->frontier[i], start:stop)
end

function Base.:>>(b::AbstractBlock, c::CircuitCanvas)
    error("block type $(typeof(b)) does not support visualization.")
end

function Base.:>>(p::PutBlock{N,1}, c::CircuitCanvas)
    i = frontier(c)
    for j in p.locs
        CircuitCanvas.lattice[i, j]
        tb >>
    end
end

function Base.:>>(cb::ControlBlock, c::CircuitCanvas)
    i = frontier(c)
end


# Useful Functions

"""
    max(vec::Vector{Atom}, para)

Do this will return the maximum data of `para` in `vec_atom`.

# Example
```julia-repl
max_data = max(vec_atom, "coord")
```
"""
function max(vec::Vector{Atom}, para)
    fields = fieldnames(typeof(vec[1]))
    para = Meta.parse(para)
    if !in(para, fields)
        error(join(["Data_Type Atom doesn't have field: ", string(para)]))
    end
    data_para = getfield(vec[1], para)
    len = length(data_para)
    if len != 1
        max = Vector{typeof(data_para).parameters[1]}(undef, len)
        max[:] = data_para[:]
        for i = 2 : length(vec)
            data_now = getfield(vec[i], para)
            for j = 1 : len
                if max[j] <= data_now[j]
                    max[j] = data_now[j]
                end
            end
        end
    else
        max = getfield(vec[1], para)
        for i = 2 : length(vec)
            data_now = getfield(vec[i], para)
            if max <= data_now
                max = data_now
            end
        end
    end
    max
end

function max(array_01::AbstractArray, array_02::AbstractArray)
    if typeof(array_01).parameters[2] == 1
        num_cols = 1
        num_rows = length(array_01)
    else
        num_rows, num_cols = size(array_01)
    end
    [array_01[row, col]>array_02[row, col] ? array_01[row, col] : array_02[row, col] for row = 1:num_rows, col = 1:num_cols]
end
"""
    min(vec::Vector{Atom}, para)

Do this will return the minimum data of `para` in `vec_atom`.

# Example
```julia-repl
in_data = min(vec_atom, "coord")
```
"""
function min(vec::Vector{Atom}, para)
    fields = fieldnames(typeof(vec[1]))
    para = Meta.parse(para)
    if !in(para, fields)
        error(join(["Data_Type Atom doesn't have field: ", string(para)]))
    end
    data_para = getfield(vec[1], para)
    len = length(data_para)
    if len != 1
        min = Vector{typeof(data_para).parameters[1]}(undef, len)
        min[:] = data_para[:]
        for i = 2 : length(vec)
            data_now = getfield(vec[i], para)
            for j = 1 : len
                if min[j] >= data_now[j]
                    min[j] = data_now[j]
                end
            end
        end
    else
        min = getfield(vec[1], para)
        for i = 2 : length(vec)
            data_now = getfield(vec[i], para)
            if min >= data_now
                min = data_now
            end
        end
    end
    min
end

function min(array_01::AbstractArray, array_02::AbstractArray)
    if typeof(array_01).parameters[2] == 1
        num_cols = 1
        num_rows = length(array_01)
    else
        num_rows, num_cols = size(array_01)
    end
    [array_01[row, col]<array_02[row, col] ? array_01[row, col] : array_02[row, col] for row = 1:num_rows, col = 1:num_cols]
end

"""
    get_data(vec_unit::Vector{T}, para::AbstractString) where T <: Data_Unit
Do this will return a vector of variable `para` for all elements in `vec_unit`.
# Example
```julia-repl
str = Si3N4()
cell = genr_cell([10, 10, 10])
data = genr(cell, str)
get_data(data.vec_angle, "atom")
```
"""
function get_data(vec_unit::Vector{T}, para::AbstractString) where T <: Unit
    fields = fieldnames(typeof(vec_unit[1]))
    para = Meta.parse(para)
    if !in(para, fields)
        error(join(["Data_Type Atom doesn't have field: ", string(para)]))
    end
    num_units = length(vec_unit)
    data = getfield(vec_unit[1], para)
    res = zeros(typeof(data[1]), num_units, length(data))

    for unit = 1 : num_units
        res[unit, :] .= getfield(vec_unit[unit], para) # try a=zeros(3, 1) a[1, :] = 1
    end
    return res
end

"""
    add(vec_unit::Vector{T}, tilt, para::AbstractString) where T <: Data_Unit

Do this will add `tilt` to variable `para` for all elements in `vec_unit`

# Example
```julia-repl
str = Si3N4()
cell = genr_cell([10, 10, 10])
data = genr(cell, str)
add!(data.vec_bond, 3, "typ")
```
"""
function add!(vec_unit::Vector{T}, tilt, para::AbstractString) where T <: Unit
    fields = fieldnames(typeof(vec_unit[1]))
    para = Meta.parse(para)
    if !in(para, fields)
        error(join(["Data_Type Atom doesn't have field: ", string(para)]))
    end
    num_units = length(vec_unit)

    for unit = 1 : num_units
        data = getfield(vec_unit[unit], para)
        len = length(data)
        if len == 1
            data += tilt
        else
            data .+= tilt
        end
        setfield!(vec_unit[unit], para, data)
    end
end

"""
    set!(vec_unit::Vector{T}, tilt, para::AbstractString) where T <: Unit

Do this will change variable `para` of all elements in `vec_unit` to `tilt`

# Example
```julia-repl
str = Si3N4()
cell = genr_cell([10, 10, 10])
data = genr(cell, str)
change(data.vec_bond, 3, "typ")
```
"""
function set!(vec_unit::Vector{T}, goal, para::AbstractString) where T <: Unit
    fields = fieldnames(typeof(vec_unit[1]))
    para = Meta.parse(para)
    if !in(para, fields)
        error(join(["Data_Type Atom doesn't have field: ", string(para)]))
    end
    num_units = length(vec_unit)

    for unit = 1 : num_units
        data = getfield(vec_unit[unit], para)
        data = goal
        setfield!(vec_unit[unit], para, data)
    end
end

function set!(vec_unit::Vector{T}, goal::AbstractArray, para::AbstractString) where T <: Unit
    fields = fieldnames(typeof(vec_unit[1]))
    para = Meta.parse(para)
    dims = size(goal)[2]
    if !in(para, fields)
        error(join(["Data_Type Atom doesn't have field: ", string(para)]))
    end
    num_units = length(vec_unit)

    for unit = 1 : num_units
        data = getfield(vec_unit[unit], para)
        if dims == 1
            data = goal[unit]
        else
            data = goal[unit, :]
        end
        setfield!(vec_unit[unit], para, data)
    end
end

function diag(vec::AbstractArray)
    len = length(vec)
    result = zeros(len, len)
    for i = 1 : len
        result[i, i] = vec[i]
    end
    result
end

function conv(vec::Vector, dim)
    len = length(vec)
    para = typeof(vec).parameters
    if dim == 1
        result = convert.(para[1], zeros(1, len))
    elseif dim == 2
        result = convert.(para[1], zeros(len, 1))
    else
        error("dim should be 1 or 2, representing Row or Column vector respectively")
    end
    for i in 1 : len
        result[i] = vec[i]
    end
    result
end

function conv(mat::Matrix, redu=0)
    para = typeof(mat).parameters
    if redu == 1
        (row, col) = size(mat)
        len = row>col ? row : col
        result = convert.(para[1], zeros(len))
    elseif redu == 0
        result = convert.(para[1], zeros(size(mat)[2], size(mat)[1]))
    else
        error("redu should be 0 or 1, representing wether Martix will be reduced to Vector.")
    end
    result[:] = mat[:]
end

"""
    dist(atom::Atom, org::Array)

Do this will return the distance between coord of `atom` and coord of `org`

# Example
```julia-repl
r = dist(atom, [0 0 0])
```
"""
function dist(atom::Atom, org::Array)
    r = 0
    for dim = 1 : 3
        r += (org[dim] - atom.coord[dim])^2
    end
    sqrt(r)
end

"""
    dist(atom::Atom, org::Array, dim)

Do this will return the 2-d distance between coord of `atom` and coord of `org` in plane vertical to axies `dim`.
`dim = 1, 2, 3` represents axies x, y, z respectively.

# Example
```julia-repl
r = dist(atom, [0 0 0], 1)
```
"""
function dist(atom::Atom, org::Array, dim)
    r = 0
    dim_range = [1 2 3]
    dim_range = dim_range[1:end .!= dim]
    for i in dim_range
        r += (org[i] - atom.coord[i])^2
    end
    sqrt(r)
end

"""
    dist(pos::Array, org::Array)

Do this will return the distance between coord of `pos` and coord of `org`

# Example
```julia-repl
r = dist([1 2 3], [0 0 0], 1)
```
"""
function dist(pos::Array, org::Array)
    r = 0
    for dim = 1 : 3
    r += (org[dim] - pos[dim])^2
    end
    sqrt(r)
end

"""
    dist(pos::Array, org::Array, dim)

Do this will return the 2-d distance between coord of `pos` and coord of `org` in plane vertical to axies `dim`.
`dim = 1, 2, 3` represents axies x, y, z respectively.

# Example
```julia-repl
r = dist([1 2 3], [0 0 0], 1)
```
"""
function dist(pos::Array, org::Array, dim)
    r = 0
    dim_range = [1 2 3]
    dim_range = dim_range[1:end .!= dim]
    for i in dim_range
        r += (org[i] - pos[i])^2
    end
    sqrt(r)
end

"""
    norm_vec(vec::Vector)
Do this will return a normalized vector along the direction of `vec`
"""
function norm_vec(vec::Vector)
    return vec ./ norm(vec)
end

"""
    function rot_mat(θx, θy, θz)
This will return a rotation matrix, coord times which will have a rotation along x, y, and z axies with angle θx, θy, θz respectively.

More details in [Wiki](https://zh.wikipedia.org/wiki/%E6%97%8B%E8%BD%AC%E7%9F%A9%E9%98%B5)
"""
function rot_mat(θx, θy, θz)
    x_mat = [
        1  0        0
        0  cos(θx)  -sin(θx)
        0  sin(θx)  cos(θx)
    ]

    y_mat = [
        cos(θy)  0  sin(θy)
        0        1  0
        -sin(θy) 0  cos(θy)
    ]

    z_mat = [
        cos(θz)  -sin(θz) 0
        sin(θz)  cos(θz)  0
        0        0        1
    ]

    return x_mat * y_mat * z_mat
end

"""
    function central_point_atom(data::Data)
This will return a vector of the coordinate of the center points, based on the atom infomation.
**Recommand** to use when model is small.
"""
function central_point_atom(data::Data)
	coord_vec = hcat([data.vec_atom[n].coord for n = 1:data.data_basic.num_atoms]...)
	return mean(coord_vec, dims=2)
end

"""
    function central_point_atom(data::Data)
This will return a vector of the coordinate of the center points, based on the box infomation.
**Recommand** to use when model is large.
"""
function central_point_box(data::Data)
	box_size = data.data_basic.box_size
	return (box_size[:, 2] .- box_size[:, 1])/2
end

"""
    genr_box_diag(data::Data)
This will return the matrix of box size information.
"""
function genr_box_diag(data::Data)
	box_info = data.data_basic.box_size
	box_info = box_info[:, 2] - box_info[:, 1]
	box_diag = diag(box_info)
	try
		data.data_basic.box_tilt
	catch
		return box_diag
	end
	box_diag[2, 1] = data.data_basic.box_tilt[1]
	box_diag[3, 1] = data.data_basic.box_tilt[2]
	box_diag[3, 2] = data.data_basic.box_tilt[3]
	return box_diag
end

"""
    genr_box_inv(data::Data)
This will return the inverse matrix of box size information.
"""
function genr_box_inv(data::Data)
	box_diag = genr_box_diag(data)
	return inv(box_diag)
end

function copy_array(goal)
    res = zeros(typeof(goal).parameters[1], size(goal))
    res[:] = goal[:]
    return res
end

"""
    function grid(x, y)
This will generate two array from two vector for contour.
"""
function grid(x, y)
    len_x = length(x)
    len_y = length(y)

    x_mat = zeros(len_y, len_x)
    y_mat = zeros(len_y, len_x)
    z_mat = zeros(len_y, len_x)

    for i = 1:len_y
        for j = 1:len_x
            x_mat[i, j] = x[j]
            y_mat[i, j] = y[i]
        end
    end

    return x_mat, y_mat
end

function in_str(data_str::Str, vec_str::Vector{T}) where T <: Str
    for str in vec_str
        if str.atom_name == data_str.atom_name
            return true
        end
    end
    return false
end


include("Args.jl") # Parsing the command-line
##################################################
# General parameters
##################################################
"""
    PARALLEL

Determining if the code is run in parallel.
"""
const PARALLEL = parsed_args["parallel"]
if ((PARALLEL != "yes") && (PARALLEL != "no"))
    error("ERROR: UNKNOWN PARALLEL") # Unknown parallel procedure
end

"""
    nbK_default

Default number of sampling points for the 3D integrals.
"""
const nbK_default = parsed_args["nbK"]

"""
    nbAvr_default

Default number of sampling points for the orbit-averaging integrals.
"""
const nbAvr_default = parsed_args["nbAvr"]

"""
    qCalc

Anisotropy parameter q for the Plummer model.
"""
const qCalc = parsed_args["q"]

##################################################
# Options to be removed in GitHub
##################################################

"""
    gAni_default

Anisotropy parameter gAni for the Isochrone model.
"""
const gAni_default = parsed_args["gAni"]

"""
    logJrcutoff

log10(Jr) lower cut-off used in the computation of dRc/dt
"""
const logJrcutoff = parsed_args["logJrcutoff"]

"""
    logLcutoff

log10(L) lower cut-off used in the computation of dRc/dt
"""
const logLcutoff = parsed_args["logLcutoff"]

##################################################
##################################################

include("Constants.jl")
include("Mean.jl")
include("Bath.jl")
include("LocalDeflection.jl")
include("OrbitAverage.jl")
include("ActionSpace.jl")
include("Flux.jl")


using InteractiveUtils
include("halfedge.jl")
include("catmull-clark.jl")

code_warntype(build, (Vector{Float32}, Vector{UInt32}, Vector{UInt32}, Vector{UInt32}))

#code_warntype(get_vertex, (Vector{Float32}, UInt32))

#code_warntype(set_vertex, (Vector{Float32}, UInt32, Vector{Float32}))

#code_warntype(subdivide, (Vector{Float32}, Vector{UInt32}, Vector{UInt32}, Vector{UInt32}))


# *********************************************************************************
# REopt, Copyright (c) 2019-2020, Alliance for Sustainable Energy, LLC.
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without modification,
# are permitted provided that the following conditions are met:
#
# Redistributions of source code must retain the above copyright notice, this list
# of conditions and the following disclaimer.
#
# Redistributions in binary form must reproduce the above copyright notice, this
# list of conditions and the following disclaimer in the documentation and/or other
# materials provided with the distribution.
#
# Neither the name of the copyright holder nor the names of its contributors may be
# used to endorse or promote products derived from this software without specific
# prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
# ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
# WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED.
# IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT,
# INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
# BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
# DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
# LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE
# OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED
# OF THE POSSIBILITY OF SUCH DAMAGE.
# *********************************************************************************
function add_dv_UnservedLoad_constraints(m,p)
    # effective load balance (with slack in dvUnservedLoad)
    @constraint(m, [s in p.elecutil.scenarios, tz in p.elecutil.outage_start_timesteps, ts in p.elecutil.outage_timesteps],
        m[:dvUnservedLoad][s, tz, ts] >= p.elec_load.critical_loads_kw[tz+ts]
        - sum(  m[:dvMGRatedProduction][t, s, tz, ts] * p.production_factor[t, tz+ts] * p.levelization_factor[t]
              - m[:dvMGProductionToStorage][t, s, tz, ts] - m[:dvMGCurtail][t, s, tz, ts]
            for t in p.techs
        )
        - m[:dvMGDischargeFromStorage][s, tz, ts]
    )
end

# constrain minimum hours that critical load is met
function add_min_hours_crit_ld_met_constraint(m,p)
    if p.min_resil_timesteps <= length(p.elecutil.outage_timesteps)
        @constraint(m, [s in p.elecutil.scenarios, tz in p.elecutil.outage_start_timesteps, ts in 1:p.min_resil_timesteps],
            m[:dvUnservedLoad][s, tz, ts] <= 0
        )
    end
end

function add_outage_cost_constraints(m,p)
    # TODO: fixed cost, account for outage_is_major_event
    @constraint(m, [s in p.elecutil.scenarios, tz in p.elecutil.outage_start_timesteps],
        m[:dvMaxOutageCost][s] >= p.pwf_e * sum(p.VoLL[tz+ts] * m[:dvUnservedLoad][s, tz, ts] for ts in 1:p.elecutil.outage_durations[s])
    )

    @expression(m, ExpectedOutageCost,
        sum(m[:dvMaxOutageCost][s] * p.elecutil.outage_probabilities[s] for s in p.elecutil.scenarios)
    )

    @constraint(m, [t in p.techs],
        m[:binMGTechUsed][t] => {m[:dvMGTechUpgradeCost][t] >= p.microgrid_premium_pct * p.two_party_factor *
		                         p.cap_cost_slope[t] * m[:dvMGsize][t]}
    )

    @constraint(m,
        m[:binMGStorageUsed] => {m[:dvMGStorageUpgradeCost] >= p.microgrid_premium_pct * m[:TotalStorageCapCosts]}
    )

    @expression(m, mgTotalTechUpgradeCost,
        sum( m[:dvMGTechUpgradeCost][t] for t in p.techs )
    )
end

function add_MG_size_constraints(m,p)
    @constraint(m, [t in p.techs],
        m[:binMGTechUsed][t] => {m[:dvMGsize][t] >= 1.0}  # 1 kW min size to prevent binaryMGTechUsed = 1 with zero cost
    )

    @constraint(m, [b in p.storage.types],
        m[:binMGStorageUsed] => {m[:dvStoragePower][b] >= 1.0} # 1 kW min size to prevent binaryMGStorageUsed = 1 with zero cost
    )

    if p.mg_tech_sizes_equal_grid_sizes
        @constraint(m, [t in p.techs],
            m[:dvMGsize][t] == m[:dvSize][t]
        )
    else
        @constraint(m, [t in p.techs],
            m[:dvMGsize][t] <= m[:dvSize][t]
        )
    end
end

function add_MG_production_constraints(m,p)

	# Electrical production sent to storage or export must be less than technology's rated production
	@constraint(m, [t in p.techs, s in p.elecutil.scenarios, tz in p.elecutil.outage_start_timesteps, ts in p.elecutil.outage_timesteps],
		m[:dvMGProductionToStorage][t, s, tz, ts] + m[:dvMGCurtail][t, s, tz, ts] <=
		p.production_factor[t, tz+ts] * p.levelization_factor[t] * m[:dvMGRatedProduction][t, s, tz, ts]
    )

    @constraint(m, [t in p.techs, s in p.elecutil.scenarios, tz in p.elecutil.outage_start_timesteps, ts in p.elecutil.outage_timesteps],
        m[:dvMGRatedProduction][t, s, tz, ts] >= 0
    )

    @constraint(m, [t in p.techs, s in p.elecutil.scenarios, tz in p.elecutil.outage_start_timesteps, ts in p.elecutil.outage_timesteps],
        m[:dvMGRatedProduction][t, s, tz, ts] <= m[:dvMGsize][t]
    )
end

function add_MG_fuel_burn_constraints(m,p)
    # Define dvMGFuelUsed by summing over outage timesteps.
    @constraint(m, [t in p.gentechs, s in p.elecutil.scenarios, tz in p.elecutil.outage_start_timesteps],
        m[:dvMGFuelUsed][t, s, tz] == p.generator.fuel_slope_gal_per_kwh * p.hours_per_timestep * p.levelization_factor[t] *
        sum( p.production_factor[t, tz+ts] * m[:dvMGRatedProduction][t, s, tz, ts] for ts in 1:p.elecutil.outage_durations[s])
        + p.generator.fuel_intercept_gal_per_hr * p.hours_per_timestep *
        sum( m[:binMGGenIsOnInTS][s, tz, ts] for ts in 1:p.elecutil.outage_durations[s])
    )

    # For each outage the fuel used is <= fuel_avail_gal
    @constraint(m, [t in p.gentechs, s in p.elecutil.scenarios, tz in p.elecutil.outage_start_timesteps],
        m[:dvMGFuelUsed][t, s, tz] <= p.generator.fuel_avail_gal
    )

    @constraint(m, [s in p.elecutil.scenarios, tz in p.elecutil.outage_start_timesteps],
        m[:dvMGMaxFuelUsage][s] >= sum( m[:dvMGFuelUsed][t, s, tz] for t in p.gentechs )
    )

    @expression(m, ExpectedMGFuelUsed,
        sum( m[:dvMGMaxFuelUsage][s] * p.elecutil.outage_probabilities[s] for s in p.elecutil.scenarios )
    )

    # fuel cost = gallons * $/gal for each tech, outage
    @expression(m, MGFuelCost[t in p.gentechs, s in p.elecutil.scenarios, tz in p.elecutil.outage_start_timesteps],
        m[:dvMGFuelUsed][t, s, tz] * p.generator.fuel_cost_per_gallon
    )

    @constraint(m, [s in p.elecutil.scenarios, tz in p.elecutil.outage_start_timesteps],
        m[:dvMGMaxFuelCost][s] >= sum( MGFuelCost[t, s, tz] for t in p.gentechs )
    )

    @expression(m, ExpectedMGFuelCost,
        sum( m[:dvMGMaxFuelCost][s] * p.elecutil.outage_probabilities[s] for s in p.elecutil.scenarios )
    )
end

function add_binMGGenIsOnInTS_constraints(m,p)
    # The following 2 constraints define binMGGenIsOnInTS to be the binary corollary to dvMGRatedProd for generator,
    # i.e. binMGGenIsOnInTS = 1 for dvMGRatedProd > min_turn_down_pct * dvMGsize, and binMGGenIsOnInTS = 0 for dvMGRatedProd = 0
    @constraint(m, [t in p.gentechs, s in p.elecutil.scenarios, tz in p.elecutil.outage_start_timesteps, ts in p.elecutil.outage_timesteps],
        !m[:binMGGenIsOnInTS][s, tz, ts] => { m[:dvMGRatedProduction][t, s, tz, ts] <= 0 }
    )
    @constraint(m, [t in p.gentechs, s in p.elecutil.scenarios, tz in p.elecutil.outage_start_timesteps, ts in p.elecutil.outage_timesteps],
        m[:binMGGenIsOnInTS][s, tz, ts] => {
            m[:dvMGRatedProduction][t, s, tz, ts] >= p.generator.min_turn_down_pct * m[:dvMGsize][t]
        }
    )
    @constraint(m, [t in p.gentechs, s in p.elecutil.scenarios, tz in p.elecutil.outage_start_timesteps, ts in p.elecutil.outage_timesteps],
        m[:binMGTechUsed][t] >= m[:binMGGenIsOnInTS][s, tz, ts]
    )
    # TODO? make binMGGenIsOnInTS indexed on p.gentechs
end

function add_MG_storage_dispatch_constraints(m,p)
    # initial SOC at start of each outage equals the grid-optimal SOC
    @constraint(m, [s in p.elecutil.scenarios, tz in p.elecutil.outage_start_timesteps],
        m[:dvMGStoredEnergy][s, tz, 0] <= m[:dvStoredEnergy][:elec, tz]
    )

    # state of charge
    @constraint(m, [s in p.elecutil.scenarios, tz in p.elecutil.outage_start_timesteps, ts in p.elecutil.outage_timesteps],
        m[:dvMGStoredEnergy][s, tz, ts] == m[:dvMGStoredEnergy][s, tz, ts-1] + p.hours_per_timestep * (
            p.storage.charge_efficiency[:elec] * sum(m[:dvMGProductionToStorage][t, s, tz, ts] for t in p.techs)
            - m[:dvMGDischargeFromStorage][s, tz, ts] / p.storage.discharge_efficiency[:elec]
        )
    )

    # Minimum state of charge
    @constraint(m, [s in p.elecutil.scenarios, tz in p.elecutil.outage_start_timesteps, ts in p.elecutil.outage_timesteps],
        m[:dvMGStoredEnergy][s, tz, ts] >=  p.storage.soc_min_pct[:elec] * m[:dvStorageEnergy][:elec]
    )

    # Dispatch to MG electrical storage is no greater than inverter capacity
    # and can't charge the battery unless binMGStorageUsed = 1
    @constraint(m, [s in p.elecutil.scenarios, tz in p.elecutil.outage_start_timesteps, ts in p.elecutil.outage_timesteps],
        m[:dvStoragePower][:elec] >= sum(m[:dvMGProductionToStorage][t, s, tz, ts] for t in p.techs)
    )

    # Dispatch from MG storage is no greater than inverter capacity
    # and can't discharge from storage unless binMGStorageUsed = 1
    @constraint(m, [s in p.elecutil.scenarios, tz in p.elecutil.outage_start_timesteps, ts in p.elecutil.outage_timesteps],
        m[:dvStoragePower][:elec] >= m[:dvMGDischargeFromStorage][s, tz, ts]
    )

    # Dispatch to and from electrical storage is no greater than power capacity
    @constraint(m, [s in p.elecutil.scenarios, tz in p.elecutil.outage_start_timesteps, ts in p.elecutil.outage_timesteps],
        m[:dvStoragePower][:elec] >= m[:dvMGDischargeFromStorage][s, tz, ts]
            + sum(m[:dvMGProductionToStorage][t, s, tz, ts] for t in p.techs)
    )

    # State of charge upper bound is storage system size
    @constraint(m, [s in p.elecutil.scenarios, tz in p.elecutil.outage_start_timesteps, ts in p.elecutil.outage_timesteps],
        m[:dvStorageEnergy][:elec] >= m[:dvMGStoredEnergy][s, tz, ts]
    )

    @constraint(m, [s in p.elecutil.scenarios, tz in p.elecutil.outage_start_timesteps, ts in p.elecutil.outage_timesteps],
        !m[:binMGStorageUsed] => { sum(m[:dvMGProductionToStorage][t, s, tz, ts] for t in p.techs) <= 0 }
    )

    @constraint(m, [s in p.elecutil.scenarios, tz in p.elecutil.outage_start_timesteps, ts in p.elecutil.outage_timesteps],
        !m[:binMGStorageUsed] => { m[:dvMGDischargeFromStorage][s, tz, ts] <= 0 }
    )
end

function add_cannot_have_MG_with_only_PVwind_constraints(m, p)
    renewable_techs = setdiff(p.techs, p.gentechs)
    # can't "turn down" renewable_techs
    if !isempty(renewable_techs)
        @constraint(m, [t in renewable_techs, s in p.elecutil.scenarios, tz in p.elecutil.outage_start_timesteps, ts in p.elecutil.outage_timesteps],
            m[:binMGTechUsed][t] => { m[:dvMGRatedProduction][t, s, tz, ts] >= m[:dvMGsize][t] }
        )
        @constraint(m, [t in renewable_techs, s in p.elecutil.scenarios, tz in p.elecutil.outage_start_timesteps, ts in p.elecutil.outage_timesteps],
            !m[:binMGTechUsed][t] => { m[:dvMGRatedProduction][t, s, tz, ts] <= 0 }
        )
        if !isempty(p.gentechs) # PV or Wind alone cannot be used for a MG
            @constraint(m, [t in renewable_techs, s in p.elecutil.scenarios, tz in p.elecutil.outage_start_timesteps, ts in p.elecutil.outage_timesteps],
                m[:binMGTechUsed][t] => { sum(m[:binMGTechUsed][tek] for tek in p.gentechs) + m[:binMGStorageUsed] >= 1 }
            )
        else
            @constraint(m, [t in renewable_techs, s in p.elecutil.scenarios, tz in p.elecutil.outage_start_timesteps, ts in p.elecutil.outage_timesteps],
                m[:binMGTechUsed][t] => { m[:binMGStorageUsed] >= 1 }
            )
        end
    end
end


function compute_ERI!(I::IntegralHelper{T, RIFIT, O}, aoints::IntegralHelper{T, RIFIT, AtomicOrbitals}) where {T<:AbstractFloat, O<:AbstractRestrictedOrbitals}
    Bμν = aoints["ERI"]

    C = I.orbitals.C
    if eltype(I.orbitals.C) !== T
        C = T.(C)
    end
    @tensoropt (Q=>50x, p=>x, q=>x, μ=>x, ν=>x) begin
        Bpq[Q,p,q] :=  Bμν[Q, μ, ν]*C[μ, p]*C[ν, q]
    end
    I["ERI"] = Bpq
end

function compute_BOO!(I::IntegralHelper{T, E, O}) where {T<:AbstractFloat, E<:AbstractDFERI, O<:AbstractRestrictedOrbitals}
    # Create AO integral object
    basis = I.orbitals.basis
    aoorbs = AtomicOrbitals(I.molecule, basis)
    aoints = IntegralHelper{T}(molecule=I.molecule, orbitals=aoorbs, basis=basis, eri_type=I.eri_type)
    compute_BOO!(I,aoints)
end

function compute_BVV!(I::IntegralHelper{T, E, O}) where {T<:AbstractFloat, E<:AbstractDFERI, O<:AbstractRestrictedOrbitals}
    # Create AO integral object
    basis = I.orbitals.basis
    aoorbs = AtomicOrbitals(I.molecule, basis)
    aoints = IntegralHelper{T}(molecule=I.molecule, orbitals=aoorbs, basis=basis, eri_type=I.eri_type)
    compute_BVV!(I,aoints)
end

function compute_BOV!(I::IntegralHelper{T, E, O}) where {T<:AbstractFloat, E<:AbstractDFERI, O<:AbstractRestrictedOrbitals}
    # Create AO integral object
    basis = I.orbitals.basis
    aoorbs = AtomicOrbitals(I.molecule, basis)
    aoints = IntegralHelper{T}(molecule=I.molecule, orbitals=aoorbs, basis=basis, eri_type=I.eri_type)
    compute_BOV!(I,aoints)
end

function compute_BOO!(I::IntegralHelper{T, E, O}, aoints::IntegralHelper{T, E, AtomicOrbitals}) where {T<:AbstractFloat, E<:AbstractDFERI, O<:AbstractRestrictedOrbitals}

    Bμν = aoints["ERI"]

    core = Options.get("drop_occ")
    o = (1+core):I.molecule.Nα
    Co = I.orbitals.C[:,o]
    if eltype(I.orbitals.C) !== T
        Co = T.(Co)
    end
    @tensoropt (Q=>50x, i=>x, j=>x, μ=>10x, ν=>10x) begin
        Bij[Q,i,j] :=  Bμν[Q, μ, ν]*Co[μ, i]*Co[ν, j]
    end
    I["BOO"] = Bij
end

function compute_BOV!(I::IntegralHelper{T, E, O}, aoints::IntegralHelper{T, E, AtomicOrbitals}) where {T<:AbstractFloat, E<:AbstractDFERI, O<:AbstractRestrictedOrbitals}

    Bμν = aoints["ERI"]

    inac = Options.get("drop_vir")
    core = Options.get("drop_occ")
    ndocc = I.molecule.Nα
    nbf = size(I.orbitals.C,1)
    o = (1+core):ndocc
    v = (ndocc+1):(nbf - inac)
    Cv = I.orbitals.C[:,v]
    Co = I.orbitals.C[:,o]
    if eltype(I.orbitals.C) !== T
        Cv = T.(Cv)
        Co = T.(Co)
    end
    @tensoropt (Q=>50x, a=>8x, i=>x, μ=>10x, ν=>10x) begin
        Bia[Q,i,a] :=  Bμν[Q, μ, ν]*Co[μ, i]*Cv[ν, a]
    end
    I["BOV"] = Bia
end

function compute_BVV!(I::IntegralHelper{T, E, O}, aoints::IntegralHelper{T, E, AtomicOrbitals}) where {T<:AbstractFloat, E<:AbstractDFERI, O<:AbstractRestrictedOrbitals}

    Bμν = aoints["ERI"]

    inac = Options.get("drop_vir")
    ndocc = I.molecule.Nα
    nbf = size(I.orbitals.C,1)
    v = (ndocc+1):(nbf - inac)
    Cv = I.orbitals.C[:,v]
    if eltype(I.orbitals.C) !== T
        Cv = T.(Cv)
    end
    @tensoropt (Q=>50x, a=>8x, b=>8x, μ=>10x, ν=>10x) begin
        Bab[Q,a,b] :=  Bμν[Q, μ, ν]*Cv[μ, a]*Cv[ν, b]
    end
    I["BVV"] = Bab
end

function compute_OOOO!(I::IntegralHelper{T,E,O}) where {T<:AbstractFloat, E<:AbstractDFERI, O<:AbstractRestrictedOrbitals}
    Boo = I["BOO"]
    @tensoropt (i=>x, j=>x, k=>x, l=>x, a=>10x, b=>10x, c=>10x, d=>10x, Q=>20x) begin
        OOOO[i,j,k,l] :=  Boo[Q, i, j]*Boo[Q, k, l]
    end
    I["OOOO"] = OOOO
end

function compute_OOOV!(I::IntegralHelper{T,E,O}) where {T<:AbstractFloat, E<:AbstractDFERI, O<:AbstractRestrictedOrbitals}
    Boo = I["BOO"]
    Bov = I["BOV"]
    @tensoropt (i=>x, j=>x, k=>x, l=>x, a=>10x, b=>10x, c=>10x, d=>10x, Q=>20x) begin
        OOOV[i,j,k,a] :=  Boo[Q, i, j]*Bov[Q, k, a]
    end
    I["OOOV"] = OOOV
end

function compute_OOVV!(I::IntegralHelper{T,E,O}) where {T<:AbstractFloat, E<:AbstractDFERI, O<:AbstractRestrictedOrbitals}
    Boo = I["BOO"]
    Bvv = I["BVV"]
    @tensoropt (i=>x, j=>x, k=>x, l=>x, a=>10x, b=>10x, c=>10x, d=>10x, Q=>20x) begin
        OOVV[i,j,a,b] :=  Boo[Q, i, j]*Bvv[Q, a, b]
    end
    I["OOVV"] = OOVV
end

function compute_OVOV!(I::IntegralHelper{T,E,O}) where {T<:AbstractFloat, E<:AbstractDFERI, O<:AbstractRestrictedOrbitals}
    Bov = I["BOV"]
    @tensoropt (i=>x, j=>x, k=>x, l=>x, a=>10x, b=>10x, c=>10x, d=>10x, Q=>20x) begin
        OVOV[i,a,j,b] :=  Bov[Q, i, a]*Bov[Q, j, b]
    end
    I["OVOV"] = OVOV
end

function compute_OVVV!(I::IntegralHelper{T,E,O}) where {T<:AbstractFloat, E<:AbstractDFERI, O<:AbstractRestrictedOrbitals}
    Bov = I["BOV"]
    Bvv = I["BVV"]
    @tensoropt (i=>x, j=>x, k=>x, l=>x, a=>10x, b=>10x, c=>10x, d=>10x, Q=>20x) begin
        OVVV[i,a,b,c] :=  Bov[Q, i, a]*Bvv[Q, b, c]
    end
    I["OVVV"] = OVVV
end

function compute_VVVV!(I::IntegralHelper{T,E,O}) where {T<:AbstractFloat, E<:AbstractDFERI, O<:AbstractRestrictedOrbitals}
    Bvv = I["BVV"]
    @tensoropt (i=>x, j=>x, k=>x, l=>x, a=>10x, b=>10x, c=>10x, d=>10x, Q=>20x) begin
        VVVV[a,b,c,d] :=  Bvv[Q, a, b]*Bvv[Q, c, d]
    end
    I["VVVV"] = VVVV
end


__precompile__()

"""
Module for dependent bootstrap procedures, by Colin T Bowers

Implemented bootstrap methods: \n
	- IID
	- Stationary
	- Moving Block
	- Circular Block
	- NoOverlapBlock

Implemented block length selection procedures: \n
	- Patton, Politis, and White (2009) Correction to Automatic Block Length Selection For The Dependent Bootstrap

Accepted input dataset types: \n
	- Vector{<:Number}
	- Matrix{<:Number} (where rows are observations and columns are variables)
	- Vector{Vector{<:Number}} (where elements of inner vectors are observations and outer vectors are variables)
	- DataFrame
	- TimeSeries.TimeArray{T,N} (only for N = 1 and N = 2)

Additional input dataset types are easily added. Please open an issue at https://github.com/colintbowers/DependentBootstrap.jl

The module has only a single exported type: \n
	- BootInput  <-- Core input type accepted by all exported functions. Typically constructed via keyword method. See ?BootInput for more detail.

All exported functions exhibit the following keyword signatures: \n
	- exported_func(data, bootinput::BootInput)
	- exported_func(data ; kwargs...)

Most users will be content to use the keyword argument method. In practice, this method wraps a
keyword argument BootInput constructor, which is then input to the exported function BootInput
method. For more detail on accepted keywords, see ?BootInput. All exported functions then use
the input dataset and bootstrap methodology described in BootInput in order to return the
appropriate statistics. A list of exported functions follows: \n
	- optblocklength <-- Estimate the optimal block length for the input dataset
	- dbootinds      <-- Get a vector of bootstrap resampling indices
	- dbootdata      <-- Get a vector of resampled datasets
	- dbootlevel1    <-- Get a vector of level 1 resampled bootstrap statistics
	- dboot          <-- Get the level 2 bootstrap statistic
	- dbootlevel2    <-- Identical to dboot. Included for completeness
	- dbootvar       <-- Wrapper on dboot that sets the level 2 statistic as the variance
	- dbootconf      <-- Wrapper on dboot that sets the level 2 statistic as a confidence interval

I use the phrases level 1 and level 2 statistics in this package in the same manner discussed in
Chapter 1 of Lahiri's textbook Resampling Methods for Dependent Data.

This package has an MIT license. Please see associated LICENSE.md file for more detail.
"""
module DependentBootstrap

using 	StatsBase, Distributions, DataFrames, TimeSeries

import 	Base: 	show

export 	BootInput,
		optblocklength,
		dbootinds,
		dbootdata,
		dbootlevel1,
		dbootlevel2,
		dboot,
		dbootvar,
		dbootconf

const NUM_RESAMPLE = 1000::Int #Default number of bootstrap resamples

#-----------------------------------------------------------------------------------------
# OLD CODE THAT USED Requires.jl TO LOAD DataFrames.jl and TimeSeries.jl
#-----------------------------------------------------------------------------------------
#In order to accommodate different types of datasets, we use the lazy loading
#features offered by the Requires package here, so that the modules in which these
#datasets are defined are not loaded unless they are actually needed. All functions
#that reference these modules are placed here inside __init__() (which is a
#requirement when using Requires package on Julia v0.7+). If you want to add a new
#dataset type to DependentBootstrap, it should be as simple as adding appropriate
#methods here. This module is structured so that nothing else should need to be done.
#The methods that need to be added are:
#    num_obs(data::T)::Int <- Number of observations in dataset
#	 num_var(data::T)::Int <- Number of variables in dataset
#    local_get_var(data::T, i::Int)::Vector <- Get the data associated with the ith variable and output as a Vector
#	 local_get_index(data::T, inds::Vector{Int})::T <- Resample data using resampling indices in inds
#For local_get_index, note that the output type should always match the input type of the dataset.
# function __init__()
# 	@require TimeSeries="9e3dc215-6440-5c97-bce1-76c03772f85e" begin
# 		(num_obs(data::TimeSeries.TimeArray{T,1})::Int) where {T} = size(data, 1)
# 		(num_obs(data::TimeSeries.TimeArray{T,2})::Int) where {T} = size(data, 1)
# 		(num_var(data::TimeSeries.TimeArray{T,1})::Int) where {T} = 1
# 		(num_var(data::TimeSeries.TimeArray{T,2})::Int) where {T} = size(data, 2)
# 		(local_get_var(data::TimeSeries.TimeArray{T,1}, i::Int)::Vector{T}) where {T} = i == 1 ? data.values[:] : error("Invalid index $(i) given data $(typeof(data)) with number of columns $(num_var(data))")
# 		(local_get_var(data::TimeSeries.TimeArray{T,2}, i::Int)::Vector{T}) where {T} = (1 <= i <= num_var(data)) ? data.values[:, i] : error("Invalid index $(i) given data $(typeof(data)) with number of columns $(num_var(data))")
# 		(local_get_index(data::TimeSeries.TimeArray{T,1}, inds::Vector{Int})::TimeSeries.TimeArray{T,1}) where {T} = TimeSeries.TimeArray(TimeSeries.timestamp(data), data.values[inds], TimeSeries.colnames(data) ; unchecked=true)
# 		(local_get_index(data::TimeSeries.TimeArray{T,2}, inds::Vector{Int})::TimeSeries.TimeArray{T,2}) where {T} = TimeSeries.TimeArray(TimeSeries.timestamp(data), data.values[inds, :], TimeSeries.colnames(data) ; unchecked=true)
# 	end
# 	@require DataFrames="a93c6f00-e57d-5684-b7b6-d8193f3e46c0" begin
# 		(num_obs(data::DataFrames.DataFrame)::Int) = size(data, 1)
# 		(num_var(data::DataFrames.DataFrame)::Int) = size(data, 2)
# 		(local_get_var(data::DataFrames.DataFrame, i::Int)) = (1 <= i <= num_var(data)) ? data[:, i] : error("Invalid index $(i) given data $(typeof(data)) with number of columns $(num_var(data))")
# 		(local_get_index(data::DataFrames.DataFrame, inds::Vector{Int})::DataFrames.DataFrame) = data[inds, :]
# 	end
# end
#-----------------------------------------------------------------------------------------

include("types.jl")
include("blocklength.jl")
include("bootinds.jl")
#include("tapered.jl")
include("core.jl")

end # module


# [src/vectors_to_states.jl]

"""
Convert a state vector history to a series of states given those states
as inputs. If the total number of numeric states is ns and the number of
time steps represented is nt, then yv is nt-by-ns. The remaining
arguments are example states. The outputs will have the same form as the
example states, but will be arrays of these states, each nt long, with
appropriate values from yv.

When a state is a column or row vector of numeric types, the output will
be an nt-by-x array where x is length(state)

When a state is an n-dimensional matrix (neither column nor vector), the
output will have the dimensions of the matrix concatenated on the first
unused dimension. That is, if the state is a-by-b, the output will be
a-by-b-by-nt.

When a state is a struct, the output will be a struct array of length nt.

Otherwise, the state will be an array of length nt, each element containing
the example type filled in with values from yv.

y1, y2, ... yn = vectors_to_states(yv, y1, y2, ... ,yn);

Inputs:

    yv              State vector history (nt-by-ns)
    y1, y2, ... yn  Example states

Outputs:

    y1, y2, ... yn  State histories as described above

Examples:

    # A history of 3 samples of a row vector:

    yv = [state_to_vector([1  2])';    # Sample 1
          state_to_vector([7  8])';    # Sample 2
          state_to_vector([13 14])'];  # Sample 3

    y  = zeros(1, 2);                  # Example state
    ys = vectors_to_states(yv, y)      # 3-by-2 output

    # A history of 3 samples of 3 different states: a column vector, a cell
    #    array, and a struct:

    # Make the example states.
    mutable struct TEST_STRUC
        a
        b
    end

    y1  = zeros(3, 1);
    y2  = [1 2; 3 4];
    y3  = TEST_STRUC(1, 2);

    # Make the history.
    y1v = [state_to_vector([1  2  3]')';
           state_to_vector([7  8  9]')';
           state_to_vector([13 14 15]')'];

    y2v = [state_to_vector([1 2; 7 8])';
           state_to_vector([3 4; 9 0])';
           state_to_vector([5 6; 1 2])'];

    y3v = [1 2; 3 4; 5 6];
    yv  = [y1v y2v y3v];

    # Get the state histories.
    y1s, y2s, y3s = vectors_to_states(yv, (y1, y2, y3) )  # NOTE that there are multiple reference states, input as a tuple

See also: vector_to_state, state_to_vector.

Online doc: http://www.anuncommonlab.com/doc/odehybrid/vectors_to_states.html

Copyright 2014 An Uncommon Lab
"""
# Called with there is only one reference state
function vectors_to_states(yv, varargin)

    # Number of samples
    ns = size(yv, 1);

    # Number of states
    nx = size(yv, 2) # There is only ONE reference state

    # For each type of state...
    count = 0;
    varargout = Array{Any}(undef, nx) # zeros(1, nx) #cell(1, nx);

    # If there are no states, return empties.
    if isempty(yv)
        return;
    end

    for x = 1:length(varargin)

        # See if it's just a bunch of numbers (easy!).
        if isa(varargin[x], Number) || isa(varargin[x], Char)
            nv  = length(varargin[x]);
            yvx = yv[:, count .+ (1:nv)];
            szs = size(varargin[x]);

            if length(szs) == 0   # Empty, meaning varargin[x] is a scalar (Which should ALWAYS be the case)
                varargout[x] = reshape(yvx, (1, ns)) # NOT casting to type of input...
            end

            count = count + nv;

        # Check for an array of numbers or chars
        elseif (isa(varargin[x], Array{<:Number}) && (any(size(varargin[x]) .== 1)) ) ||
                 (isa(varargin[x], Array{<:Char}) && (any(size(varargin[x]) .== 1)) )

            nv = length(varargin[x]);
            yvx = yv[:, count .+ (1:nv)]
            szs = size(varargin[x]);

            if length(szs) == 2 && any(szs[1:2] .== 1)
                varargout[x] = yvx
            else
                @warn "Potential problem in vecs to states (matrix not array)!"
            end
            count += nv

        # See if it is a struct
        elseif length(varargin[x]) == 1

            vargout[x] = vargin[x];
            dummy = state_to_vector(varargin[x]);
            yvx = yv[:, count + (1:maximum(size(dummy)))]
            for k = 1:ns
                varargout[x][k, 1] = vector_to_state(yvx[k, :], varargin[x]);
            end
            count += maximum(size(dummy));

        else

            varargout[x] = Array{Any}(undef, ns) #zeros(ns, 1);
            dummy = state_to_vector(varargin[x]);
            yvx = yv[:, count .+ (1:maximum(size(dummy)))];
            for k = 1:ns
                t, _ = vector_to_state(yvx[k, :], varargin[x]);
                varargout[x][k] = deepcopy(t)
            end
        end

    end

    return vec_to_mat(varargout)'
end

function vectors_to_states(yv, varargin::Tuple)
    """
        Alternate version run when length(varargin) > 1 (i.e., multiple reference states)
    """

    # Number of samples
    ns = size(yv, 1);

    # Number of states
    nx = length(varargin); # Elements in tuple

    # For each type of state...
    count = 0;
    varargout = Array{Any}(undef, nx)

    # If there are no states, return empties.
    if isempty(yv)
        return;
    end

    for x = 1:length(varargin)

        # See if it's just a bunch of numbers (easy!).
        if isa(varargin[x], Number) || isa(varargin[x], Char)
            nv  = length(varargin[x]);
            yvx = yv[:, count .+ (1:nv)];
            szs = size(varargin[x]);

            if length(szs) == 0   # Empty, meaning varargin[x] is a scalar (Which should ALWAYS be the case)
                varargout[x] = reshape(yvx, (1, ns)) # NOT casting to type of input...
            end

            count = count + nv;

        # Check for an array of numbers or chars
        elseif (isa(varargin[x], Array{<:Number}) && (any(size(varargin[x]) .== 1)) ) ||
                 (isa(varargin[x], Array{<:Char}) && (any(size(varargin[x]) .== 1)) )

            nv = length(varargin[x]);
            yvx = yv[:, count .+ (1:nv)]
            szs = size(varargin[x]);

            if length(szs) == 2 && any(szs[1:2] .== 1)
                varargout[x] = yvx
            else
                @warn "Potential problem in vecs to states (matrix not array?)!"
            end
            count += nv

        else
            varargout[x] = Array{Any}(undef, ns) #zeros(ns, 1);
            dummy = state_to_vector(varargin[x]);
            yvx = yv[:, count .+ (1:maximum(size(dummy)))];
            for k = 1:ns
                t, _ = vector_to_state(yvx[k, :], varargin[x]);
                varargout[x][k] = deepcopy(t)
            end

            count += maximum(size(dummy))
        end

    end

    return varargout
end


include("filtering.jl");
## Number of pokes during the session

Df = combine(groupby(streaks,:Phase)) do dd
    summarize(dd,:Trial,:Num_pokes)
end
Drug_colors!(Df)
@df Df plot(:Xaxis,:Mean,
    group = :Phase,
    color = :color,
    ribbon = :SEM,
    linewidth = 2,
    fillalpha = 0.1,
    xticks = 0:10:100)
savefig(joinpath(figs_loc,"Fig1/trials_selected.pdf"))


include("./topoprovision.jl")

ProvisionIPTopology("B4", scaling=5, density=0.98, linknum=52, ilp_planning=false, tofile=true)
ProvisionIPTopology("IBM", scaling=2, density=0.98, linknum=85, ilp_planning=false, tofile=true)