text:
                                            ################################
                                            ### Problem Set 0 - ECON 603 ###
                                            ################################

                                                #Sebastian Melo-Martin

# Call some useful packages
using Distributions
using Random
using QuantEcon
using Plots
using NLopt
using SpecialFunctions
using StatsBase
using DataFrames
using DelimitedFiles
using CSV

Random.seed!(641993)

############
# Question 1
############

std_norm_cdf(x::T) where {T <: Real} = 0.5 * erfc(-x/sqrt(2))
std_norm_cdf(x::Array{T}) where {T <: Real} = 0.5 .* erfc(-x./sqrt(2))

function tauchen(N::Integer, ρ::T1, σ::T2, μ=zero(promote_type(T1, T2)), n_std::T3=3) where {T1 <: Real, T2 <: Real, T3 <: Real}
    # Discretized space
    a_bar = n_std * sqrt(σ^2 / (1 - ρ^2))
    y = range(-a_bar, stop=a_bar, length=N)
    d = y[2] - y[1]

    # Transition probabilities
    Π = zeros(promote_type(T1, T2), N, N)
    for row = 1:N
        # Do end points first
        Π[row, 1] = std_norm_cdf((y[1] - ρ*y[row] + d/2) / σ)
        Π[row, N] = 1 - std_norm_cdf((y[N] - ρ*y[row] - d/2) / σ)

        # Fill in columns
        for col = 2:N-1
            Π[row, col] = (std_norm_cdf((y[col] - ρ*y[row] + d/2) / σ) -
                           std_norm_cdf((y[col] - ρ*y[row] - d/2) / σ))
        end
    end

    yy = y .+ μ / (1 - ρ)
    Π = Π./sum(Π, dims = 2)

    return Π, yy

end

vec_matrix = [zeros(3,3), 1:1:length(1), zeros(3,3), 1:1:length(1), zeros(3,3), 1:1:length(1), zeros(3,3), 1:1:length(1)]

for i in [0.2, 0.7, 0.9, 0.98]
    z = tauchen(3,i,0.4)
    if i==0.2
        vec_matrix[1] = z[1]
        vec_matrix[2] = z[2]
    elseif i==0.7
        vec_matrix[3] = z[1]
        vec_matrix[4] = z[2]
    elseif i==0.9
        vec_matrix[5] = z[1]
        vec_matrix[6] = z[2]
    else
        vec_matrix[7] = z[1]
        vec_matrix[8] = z[2]
    end
end

y_tau = [zeros(1000,1), zeros(1000,1), zeros(1000,1), zeros(1000,1)]
pos = [zeros(1000,1), zeros(1000,1), zeros(1000,1), zeros(1000,1)]

for k=1:4

    l = 2*k - 1
    m = 2*k
    y_tau[k][1] = vec_matrix[m][rand(1:3)]

    for i=2:1000
        x=rand()
        for j=1:3
            if y_tau[k][i-1]==vec_matrix[m][j]
                pos[k][i] = j
            end
        end
        prob_row = vec_matrix[l][floor(Int, pos[k][i]),:]
        thr_1 = prob_row[1]
        thr_2 = thr_1 + prob_row[2]
        if x<=thr_1
            y_tau[k][i]=vec_matrix[m][1]
        elseif x>thr_1 && x<=thr_2
            y_tau[k][i] = vec_matrix[m][2]
        else
            y_tau[k][i] = vec_matrix[m][3]
        end
    end
end

t = collect(range(1,stop=1000,length=1000))
plot(t, y_tau[1], title=" ", legend = :false, xlabel = "Time")
savefig("./PS0_tau3_1.png")
plot(t, y_tau[2], title=" ", legend = :false, xlabel = "Time")
savefig("./PS0_tau3_2.png")
plot(t, y_tau[3], title=" ", legend = :false, xlabel = "Time")
savefig("./PS0_tau3_3.png")
plot(t, y_tau[4], title=" ", legend = :false, xlabel = "Time")
savefig("./PS0_tau3_4.png")

table1_1 = zeros(4,6)

for j=1:4
    table1_1[j,1] = mean(y_tau[j])
    table1_1[j,2] = var(y_tau[j])
    table1_1[j,3] = percentile(vec(y_tau[j]),25)
    table1_1[j,4] = percentile(vec(y_tau[j]),50)
    table1_1[j,5] = percentile(vec(y_tau[j]),75)
    table1_1[j,6] = percentile(vec(y_tau[j]),90)
end

CSV.write("table1_1.csv",  Tables.table(table1_1), writeheader=false)

##############
# Question 2 #
##############

#Getting initial vector and probability matrix
function rouwenhorst(N::Integer, ρ::Real, σ::Real, μ::Real)

    σ_y = σ / sqrt(1-ρ^2)
    p  = (1+ρ)/2
    Θ = [p 1-p; 1-p p]
    ψ = sqrt(N-1) * σ_y
    m = μ / (1 - ρ)

    state_values, p = _rouwenhorst(p, p, m, ψ, N)

    return p, state_values
end

function _rouwenhorst(p::Real, q::Real, m::Real, Δ::Real, n::Integer)
    if n == 2
        return [m-Δ, m+Δ],  [p 1-p; 1-q q]
    else
        _, θ_nm1 = _rouwenhorst(p, q, m, Δ, n-1)
        θN = p    *[θ_nm1 zeros(n-1, 1); zeros(1, n)] +
             (1-p)*[zeros(n-1, 1) θ_nm1; zeros(1, n)] +
             q    *[zeros(1, n); zeros(n-1, 1) θ_nm1] +
             (1-q)*[zeros(1, n); θ_nm1 zeros(n-1, 1)]

        θN[2:end-1, :] ./= 2

        return range(m-Δ, stop=m+Δ, length=n), θN
    end
end

vec_matrix = [zeros(3,3), 1:1:length(1), zeros(3,3), 1:1:length(1), zeros(3,3), 1:1:length(1), zeros(3,3), 1:1:length(1)]

for i in [0.2, 0.7, 0.9, 0.98]
    z = rouwenhorst(3,i,0.4,0)
    if i==0.2
        vec_matrix[1] = z[1]
        vec_matrix[2] = z[2]
    elseif i==0.7
        vec_matrix[3] = z[1]
        vec_matrix[4] = z[2]
    elseif i==0.9
        vec_matrix[5] = z[1]
        vec_matrix[6] = z[2]
    else
        vec_matrix[7] = z[1]
        vec_matrix[8] = z[2]
    end
end

y_rou = [zeros(1000,1), zeros(1000,1), zeros(1000,1), zeros(1000,1)]
pos = [zeros(1000,1), zeros(1000,1), zeros(1000,1), zeros(1000,1)]

#Simulating the time series
for k=1:4

    l = 2*k - 1
    m = 2*k
    y_rou[k][1] = vec_matrix[m][rand(1:3)]

    for i=2:1000
        x=rand()
        for j=1:3
            if y_rou[k][i-1]==vec_matrix[m][j]
                pos[k][i] = j
            end
        end
        prob_row = vec_matrix[l][floor(Int, pos[k][i]),:]
        thr_1 = prob_row[1]
        thr_2 = thr_1 + prob_row[2]
        if x<=thr_1
            y_rou[k][i]=vec_matrix[m][1]
        elseif x>thr_1 && x<=thr_2
            y_rou[k][i] = vec_matrix[m][2]
        else
            y_rou[k][i] = vec_matrix[m][3]
        end
    end
end

t = collect(range(1,stop=1000,length=1000))
plot(t, y_rou[1], title=" ", legend = :false, xlabel = "Time", color = :deepskyblue3)
savefig("./PS0_rou3_1.png")
plot(t, y_rou[2], title=" ", legend = :false, xlabel = "Time", color = :deepskyblue3)
savefig("./PS0_rou3_2.png")
plot(t, y_rou[3], title=" ", legend = :false, xlabel = "Time", color = :deepskyblue3)
savefig("./PS0_rou3_3.png")
plot(t, y_rou[4], title=" ", legend = :false, xlabel = "Time", color = :deepskyblue3)
savefig("./PS0_rou3_4.png")

table2_2 = zeros(4,6)

for j=1:4
    table2_2[j,1] = mean(y_rou[j])
    table2_2[j,2] = var(y_rou[j])
    table2_2[j,3] = percentile(vec(y_rou[j]),25)
    table2_2[j,4] = percentile(vec(y_rou[j]),50)
    table2_2[j,5] = percentile(vec(y_rou[j]),75)
    table2_2[j,6] = percentile(vec(y_rou[j]),90)
end

CSV.write("table2_2.csv",  Tables.table(table2_2), writeheader=false)

##############
# Question 3 #
##############

function normal_ar(ρ, N)
    d = Normal(0, 0.16)
    y = [0.0 for i = 1:N]
    ϵ = rand(d, N)
    for i in 1:(N-1)
        y[i+1] = ρ*y[i] + ϵ[i]
    end
    return y
end

y_ar = [zeros(1000), zeros(1000), zeros(1000), zeros(1000)]

N=1000
for i in [0.2, 0.7, 0.9, 0.98]
    ρ=i
    if i==0.2
        y_ar[1] = normal_ar(ρ,N)
    elseif i==0.7
        y_ar[2] = normal_ar(ρ,N)
    elseif i==0.9
        y_ar[3] = normal_ar(ρ,N)
    else
        y_ar[4] = normal_ar(ρ,N)
    end
end

t = collect(range(1,stop=1000,length=1000))
plot(t, y_ar[1], title=" ", legend = :false, xlabel = "Time", color = :tan2)
savefig("./PS0_ar_1.png")
plot(t, y_ar[2], title=" ", legend = :false, xlabel = "Time", color = :tan2)
savefig("./PS0_ar_2.png")
plot(t, y_ar[3], title=" ", legend = :false, xlabel = "Time", color = :tan2)
savefig("./PS0_ar_3.png")
plot(t, y_ar[4], title=" ", legend = :false, xlabel = "Time", color = :tan2)
savefig("./PS0_ar_4.png")

##############
# Question 4 #
##############

# Tauchen
vec_matrix = [zeros(10,10), zeros(10), zeros(10,10), zeros(10), zeros(10,10), zeros(10), zeros(10,10), zeros(10)]

for i in [0.2, 0.7, 0.9, 0.98]
    z = tauchen(10,i,0.4)
    if i==0.2
        vec_matrix[1] = z[1]
        vec_matrix[2] = Vector(z[2])
    elseif i==0.7
        vec_matrix[3] = z[1]
        vec_matrix[4] = Vector(z[2])
    elseif i==0.9
        vec_matrix[5] = z[1]
        vec_matrix[6] = Vector(z[2])
    else
        vec_matrix[7] = z[1]
        vec_matrix[8] = Vector(z[2])
    end
end

y_tau = [zeros(1000,1), zeros(1000,1), zeros(1000,1), zeros(1000,1)]
pos = [zeros(1000,1), zeros(1000,1), zeros(1000,1), zeros(1000,1)]

for k=1:4

    l = 2*k - 1
    m = 2*k
    y_tau[k][1] = vec_matrix[m][rand(1:3)]

    for i=2:1000
        x=rand()
        for j=1:10
            if y_tau[k][i-1]==vec_matrix[m][j]
                pos[k][i] = j
            end
        end
        prob_row = vec_matrix[l][floor(Int, pos[k][i]),:]
        thr_1 = prob_row[1]
        thr_2 = thr_1 + prob_row[2]
        thr_3 = thr_2 + prob_row[3]
        thr_4 = thr_3 + prob_row[4]
        thr_5 = thr_4 + prob_row[5]
        thr_6 = thr_5 + prob_row[6]
        thr_7 = thr_6 + prob_row[7]
        thr_8 = thr_7 + prob_row[8]
        thr_9 = thr_8 + prob_row[9]
        if x<=thr_1
            y_tau[k][i]= vec_matrix[m][1]
        elseif x>thr_1 && x<=thr_2
            y_tau[k][i] = vec_matrix[m][2]
        elseif x>thr_2 && x<=thr_3
            y_tau[k][i] = vec_matrix[m][3]
        elseif x>thr_3 && x<=thr_4
            y_tau[k][i] = vec_matrix[m][4]
        elseif x>thr_4 && x<=thr_5
            y_tau[k][i] = vec_matrix[m][5]
        elseif x>thr_5 && x<=thr_6
            y_tau[k][i] = vec_matrix[m][6]
        elseif x>thr_6 && x<=thr_7
            y_tau[k][i] = vec_matrix[m][7]
        elseif x>thr_7 && x<=thr_8
            y_tau[k][i] = vec_matrix[m][8]
        elseif x>thr_8 && x<=thr_9
            y_tau[k][i] = vec_matrix[m][9]
        else
            y_tau[k][i] = vec_matrix[m][10]
        end
    end
end

t = collect(range(1,stop=1000,length=1000))
plot(t, y_tau[1], title=" ", legend = :false, xlabel = "Time")
savefig("./PS0_tau10_1.png")
plot(t, y_tau[2], title=" ", legend = :false, xlabel = "Time")
savefig("./PS0_tau10_2.png")
plot(t, y_tau[3], title=" ", legend = :false, xlabel = "Time")
savefig("./PS0_tau10_3.png")
plot(t, y_tau[4], title=" ", legend = :false, xlabel = "Time")
savefig("./PS0_tau10_4.png")

# Rouwenhorst
vec_matrix = [zeros(10,10), zeros(10), zeros(10,10), zeros(10), zeros(10,10), zeros(10), zeros(10,10), zeros(10)]

for i in [0.2, 0.7, 0.9, 0.98]
    z = rouwenhorst(10,i,0.4,0)
    if i==0.2
        vec_matrix[1] = z[1]
        vec_matrix[2] = Vector(z[2])
    elseif i==0.7
        vec_matrix[3] = z[1]
        vec_matrix[4] = Vector(z[2])
    elseif i==0.9
        vec_matrix[5] = z[1]
        vec_matrix[6] = Vector(z[2])
    else
        vec_matrix[7] = z[1]
        vec_matrix[8] = Vector(z[2])
    end
end

y_rou = [zeros(1000,1), zeros(1000,1), zeros(1000,1), zeros(1000,1)]
pos = [zeros(1000,1), zeros(1000,1), zeros(1000,1), zeros(1000,1)]

for k=1:4

    l = 2*k - 1
    m = 2*k
    y_rou[k][1] = vec_matrix[m][rand(1:3)]

    for i=2:1000
        x=rand()
        for j=1:10
            if y_rou[k][i-1]==vec_matrix[m][j]
                pos[k][i] = j
            end
        end
        prob_row = vec_matrix[l][floor(Int, pos[k][i]),:]
        thr_1 = prob_row[1]
        thr_2 = thr_1 + prob_row[2]
        thr_3 = thr_2 + prob_row[3]
        thr_4 = thr_3 + prob_row[4]
        thr_5 = thr_4 + prob_row[5]
        thr_6 = thr_5 + prob_row[6]
        thr_7 = thr_6 + prob_row[7]
        thr_8 = thr_7 + prob_row[8]
        thr_9 = thr_8 + prob_row[9]
        if x<=thr_1
            y_rou[k][i]=vec_matrix[m][1]
        elseif x>thr_1 && x<=thr_2
            y_rou[k][i] = vec_matrix[m][2]
        elseif x>thr_2 && x<=thr_3
            y_rou[k][i] = vec_matrix[m][3]
        elseif x>thr_3 && x<=thr_4
            y_rou[k][i] = vec_matrix[m][4]
        elseif x>thr_4 && x<=thr_5
            y_rou[k][i] = vec_matrix[m][5]
        elseif x>thr_5 && x<=thr_6
            y_rou[k][i] = vec_matrix[m][6]
        elseif x>thr_6 && x<=thr_7
            y_rou[k][i] = vec_matrix[m][7]
        elseif x>thr_7 && x<=thr_8
            y_rou[k][i] = vec_matrix[m][8]
        elseif x>thr_8 && x<=thr_9
            y_rou[k][i] = vec_matrix[m][9]
        else
            y_rou[k][i] = vec_matrix[m][10]
        end
    end
end

t = collect(range(1,stop=1000,length=1000))
plot(t, y_rou[1], title=" ", legend = :false, xlabel = "Time", color = :deepskyblue3)
savefig("./PS0_rou10_1.png")
plot(t, y_rou[2], title=" ", legend = :false, xlabel = "Time", color = :deepskyblue3)
savefig("./PS0_rou10_2.png")
plot(t, y_rou[3], title=" ", legend = :false, xlabel = "Time", color = :deepskyblue3)
savefig("./PS0_rou10_3.png")
plot(t, y_rou[4], title=" ", legend = :false, xlabel = "Time", color = :deepskyblue3)
savefig("./PS0_rou10_4.png")

table1 = zeros(4,6)
table2 = zeros(4,6)
table3 = zeros(4,6)

for j=1:4
    table1[j,1] = mean(y_tau[j])
    table1[j,2] = var(y_tau[j])
    table1[j,3] = percentile(vec(y_tau[j]),25)
    table1[j,4] = percentile(vec(y_tau[j]),50)
    table1[j,5] = percentile(vec(y_tau[j]),75)
    table1[j,6] = percentile(vec(y_tau[j]),90)
end

CSV.write("table1.csv",  Tables.table(table1), writeheader=false)

for j=1:4
    table2[j,1] = mean(y_rou[j])
    table2[j,2] = var(y_rou[j])
    table2[j,3] = percentile(vec(y_rou[j]),25)
    table2[j,4] = percentile(vec(y_rou[j]),50)
    table2[j,5] = percentile(vec(y_rou[j]),75)
    table2[j,6] = percentile(vec(y_rou[j]),90)
end

CSV.write("table2.csv",  Tables.table(table2), writeheader=false)

for j=1:4
    table3[j,1] = mean(y_ar[j])
    table3[j,2] = var(y_ar[j])
    table3[j,3] = percentile(vec(y_ar[j]),25)
    table3[j,4] = percentile(vec(y_ar[j]),50)
    table3[j,5] = percentile(vec(y_ar[j]),75)
    table3[j,6] = percentile(vec(y_ar[j]),90)
end

CSV.write("table3.csv",  Tables.table(table3), writeheader=false)

##############
# Question 5 #
##############

#= The coding for this question is already included previously in loops
    =#


text:
using HSL
using LinearOperators

K = diagm(1 ./ [1.0; 2.0; 3.0; 4.0; 5.0])
rhs = ones(5)
x1=K\rhs
spK = sparse(K)

K[3,3]=eps()

LBL = Ma97(K)
ma97_factorize(LBL)
x2 = ma97_solve(LBL, rhs)  # or x = LBL \ b
ret = ma97_inquire(LBL)
d=ret[2]
piv=ret[1]
show(piv);println()
show(d);println()
d[1,3] = 3.0
show(d);println()

ma97_alter(LBL, d)

x3 = ma97_solve(LBL, rhs)

show(x3)

#for i=1:5
    K = sprand(15,15,0.4)
    K = K + K'
    LBL = Ma97(K)
    ma97_factorize(LBL)
    ret = ma97_inquire(LBL)
    d=ret[2]
    piv=ret[1]
    println(" pivots : $piv")
    apiv = abs(piv)
    D1=spdiagm((d[1,apiv],d[2,apiv][1:end-1],d[2,apiv][1:end-1]),[0 -1 1])
        D2=spdiagm((d[1,:],d[2,[1:end-1]],d[2,[1:end-1]]),[0 -1 1])
D=D2[apiv,apiv]
    b=ones(15)
    x2 = ma97_solve(LBL, b)
    p1=copy(b)
    ma97_solve!(LBL, p1, job=:PL)
    p2 = D*p1
    p3 = copy(p2)
    ma97_solve!(LBL, p3, job=:LPS)

show([x2-p3])
println("\n")

    b=ones(15)
    x2 = ma97_solve(LBL, b)
    p1=copy(b)
    ma97_solve!(LBL, p1, job=:PL)
    p2 = (D*p1)[apiv]
    p3 = copy(p2)
    ma97_solve!(LBL, p3, job=:LPS)

show([x2-p3])
println("\n")

    b=ones(15)
    x2 = ma97_solve(LBL, b)
    p1=copy(b)
    ma97_solve!(LBL, p1, job=:PL)
    p2[apiv] = D*p1[apiv]
    p3 = copy(p2)
    ma97_solve!(LBL, p3, job=:LPS)

show([x2-p3])
println("\n")
    b=ones(15)
    x2 = ma97_solve(LBL, b)
    p1=copy(b)
    ma97_solve!(LBL, p1, job=:PL)
    p2[apiv] = D*p1
    p3 = copy(p2)
    ma97_solve!(LBL, p3, job=:LPS)

[x2-p3]

#end

ext: jl


text:
#
# Generate BayesNets
#
# Create (random) Bayesian Networks
#
# Also, Bernoulli does not count as a Categoical RV, so they can't be used in
# a DiscreteBayesNet

"""
    rand_discrete_bn(num_nodes16, max_num_parents=3,
            max_num_states=5, connected=true)

Generate a random DiscreteBayesNet.

Creates DiscreteBayesNet with `num_nodes` nodes, with each node having
a random number of states and parents, up to `max_num_parents` and
`max_num_parents`, respectively.
If `connected`, each node (except the first) will be guaranteed at least one
parent, making the graph connected.
"""
function rand_discrete_bn(num_nodes::Int=16,
        max_num_parents::Int=3,
        max_num_states::Int=5,
        connected::Bool=true)
    num_nodes > 0 || throw(ArgumentError("`num_nodes` must be greater than 0"))
    max_num_parents > 0 || throw(ArgumentError("`max_num_parents` must be " *
                "greater than 0"))
    max_num_states > 1  || throw(ArgumentError("`max_num_states` must be " *
                "greater than 1"))

    min_parents = connected ? 1 : 0

    bn = DiscreteBayesNet();

    for i in range(1, num_nodes)
        # the symbol to add, NX, where X is its number
        s = Symbol("N", i)
        n_states = rand(2:max_num_states)

        # keep trying different parents until it isn't cyclic
        # it should be impossible for a cycle to appear based on the
        #  algorithm but ...
        while true
            # how many parents we can possibly add
            n_par = rand(0:min(length(bn), max_num_parents))
            parents = sample(names(bn),
                    min(length(bn), max(min_parents, n_par)); replace=false)

            # add the new random cpd with random parents to the network
            push!(bn, rand_cpd(bn, n_states, s, parents))

            if is_cyclic(bn.dag)
                delete!(bn, s)
            else
                break
            end
        end
    end

    return bn
end

"""
    rand_bn_inference(bn, num_query=2, num_evidence=3)

Generate a random inference state for a Bayesian Network with an evidence
assignment sample uniformly over the chosen nodes' domain.
"""
function rand_bn_inference(bn::BayesNet, num_query::Int=2, num_evidence::Int=3)
    (num_query + num_evidence) <= length(bn) ||
        throw(ArgumentError("Number of qurey and evidence nodes " *
                    "($(num_query + num_evidence)) is greater than number " *
                    "of nodes ($(length(bn)))"))

    # yay for convoluted
    non_hidden = sample(names(bn), num_query + num_evidence; replace=false)
    query = non_hidden[1:num_query]
    evidence_nodes = non_hidden[(num_query+1):end]
    evidence = Assignment()

    for ev in evidence_nodes
        ncat = ncategories(bn, ev)
        evidence[ev] = sample(1:ncat)
    end

    return InferenceState(bn, query, evidence)
end

"""
The usual sprinkler problem
"""
function get_sprinkler_bn()
    sprinkler = DiscreteBayesNet()

    push!(sprinkler, rand_cpd(sprinkler, 2, :Rain, Symbol[]))
    push!(sprinkler, rand_cpd(sprinkler, 2, :Sprinkler, Symbol[:Rain]))
    push!(sprinkler, rand_cpd(sprinkler, 2, :WetGrass,
        Symbol[:Sprinkler, :Rain]))

    return sprinkler
end

"""
Satellite failure network from DMU, pg 17
"""
function get_sat_fail_bn()
    sat_fail = DiscreteBayesNet()

    # DiscreteBayes don't support Bernouli, has to be Categorical
    push!(sat_fail, rand_cpd(sat_fail, 2, :B, Symbol[]))
    push!(sat_fail, rand_cpd(sat_fail, 2, :S, Symbol[]))
    push!(sat_fail, rand_cpd(sat_fail, 2, :E, [:B, :S]))
    push!(sat_fail, rand_cpd(sat_fail, 2, :D, [:E]))
    push!(sat_fail, rand_cpd(sat_fail, 2, :C, [:E]))

    return sat_fail
end

"""
An ergodic version of the asia network, with the E variable removed

Orignal network: Lauritzen, Steffen L. and David J. Spiegelhalter, 1988
"""
function get_asia_bn()
    asia = DiscreteBayesNet()

    push!(asia, CategoricalCPD(:A, Symbol[], Int64[],
        [Categorical([0.99, 0.01])]))
    push!(asia, CategoricalCPD(:S, Symbol[], Int64[],
        [Categorical([0.5, 0.5])]))
    push!(asia, CategoricalCPD(:T, [:A], [2],
        [Categorical([0.99, 0.01]), Categorical([0.95, 0.05])]))
    push!(asia, CategoricalCPD(:L, [:S], [2],
        [Categorical([0.99, 0.01]), Categorical([0.9, 0.1])]))
    push!(asia, CategoricalCPD(:B, [:S], [2],
        [Categorical([0.97, 0.03]), Categorical([0.4, 0.6])]))

    push!(asia, CategoricalCPD(:X, [:T, :L], [2, 2], [
        Categorical([0.95, 0.05]), # false, false
        Categorical([0.02, 0.98]), # true, false
        Categorical([0.02, 0.98]), # false, true
        Categorical([0.02, 0.98]), # true, true
        ]))

    push!(asia, CategoricalCPD(:D, [:T, :L, :B], [2, 2, 2],[
        Categorical([0.9, 0.1]), # false, false, false
        Categorical([0.2, 0.8]), # true, false, false
        Categorical([0.3, 0.7]), # false, true, false
        Categorical([0.1, 0.9]), # true, true, false

        Categorical([0.3, 0.7]), # false, false, true
        Categorical([0.1, 0.9]), # true, false, true
        Categorical([0.3, 0.7]), # false, true, true
        Categorical([0.1, 0.9]), # true, true, true
        ]))

    return asia
end


text:
# This shows how to solve the 2D incompressible Stokes equations using SNES solvers,
# using a staggered grid discretization and a Velocity-Pressure formulation
#
#   Governing equations:
#
#              dVx/dx + dVz/dz = 0                              |   Mass balance
#   -dP/dx + dTxx/dx + dTxz/dz = 0                              |   Horizontal force balance
#   -dP/dz + dTxz/dx + dTzz/dz = rho*g                          |   Vertical force balance
#
#   with:
#       Exx = dVx/dx, Ezz = dVz/dz, Exz = 0.5*(dVx/dz + dVz/dx) |   Strain rate definition
#
#      | Txx |   | 2*eta   0     0   |  | Exx |
#      | Tzz | = |   0   2*eta   0   |  | Ezz |                 |   Linear viscous (isotropic) rheology
#      | Txz |   |   0     0   2*eta |  | Exz |
#
# This example also uses the Automatic differentiation package ForwardDiff

using PETSc, MPI
using ForwardDiff, SparseArrays, SparseDiffTools

petsclib = PETSc.petsclibs[1];
PETSc.initialize(petsclib)

function FormRes!(ptr_fx_g, ptr_x_g, user_ctx)

    # Copy global to local vectors
    PETSc.update!(user_ctx.x_l, ptr_x_g,   PETSc.INSERT_VALUES)
    PETSc.update!(user_ctx.f_l, ptr_fx_g,  PETSc.INSERT_VALUES)

    # Retrieve arrays from the local vectors
    ArrayLocal_x     =   PETSc.DMStagVecGetArrayRead(user_ctx.dm,   user_ctx.x_l);      # array with all local x-data (solution array)
    ArrayLocal_f     =   PETSc.DMStagVecGetArray(user_ctx.dm,       user_ctx.f_l);      # array with all local residual

    # Compute local residual
    ComputeLocalResidual(user_ctx.dm, ArrayLocal_x, ArrayLocal_f, user_ctx)

    # Finalize local arrays
    Base.finalize(ArrayLocal_x)
    Base.finalize(ArrayLocal_f)

    # Copy local into global residual vector
    PETSc.update!( ptr_fx_g, user_ctx.f_l, PETSc.INSERT_VALUES)

end

function  ForwardDiff_res(x, user_ctx)
    f   = zero(x)               # vector of zeros, of same type as x (local vector)

    ArrayLocal_x     =   PETSc.DMStagVecGetArray(user_ctx.dm, x);        # array with all local x-data
    ArrayLocal_f     =   PETSc.DMStagVecGetArray(user_ctx.dm, f);        # array with all local residual

    ComputeLocalResidual(user_ctx.dm, ArrayLocal_x, ArrayLocal_f, user_ctx);

    # As the residual vector f is linked with ArrayLocal_f, we don't need to
    # pass ArrayLocal_f back to f

    return f;
end

function ComputeSparsityPatternJacobian_automatic(x_l, user_ctx)
    # This computes the sparsity pattern and coloring of the jacobian automatically
    # This will works for any equation but is slow @ high resolutions

    f_Residual  =   (x -> ForwardDiff_res(x, user_ctx));        # pass additional arguments into the routine
    J_julia     =   ForwardDiff.jacobian(f_Residual,x_l*0 .+ 1);# Compute a "blank" jacobian

    # employ sparse structure to compute jacobian - to be moved inside routines
    jac         =   sparse(J_julia);
    colors      =   matrix_colors(jac)          # the number of nonzeros per row

    return jac, colors
end

function  func(out, x, user_ctx)

    ArrayLocal_x     =   PETSc.DMStagVecGetArray(user_ctx.dm, x);          # array with all local x-data (solition array)
    ArrayLocal_f     =   PETSc.DMStagVecGetArray(user_ctx.dm, out);        # array with all local residual

    ComputeLocalResidual(user_ctx.dm, ArrayLocal_x, ArrayLocal_f, user_ctx);

    return nothing
end

function FormJacobian!(ptr_x_g, J, P, user_ctx)
    # This requires several steps:
    #
    #   1) Extract local vector from global solution (x) vector
    #   2) Compute local jacobian from the residual routine (note that
    #       this routine requires julia vectors as input)
    #   3) Store the matrix inside the global Petsc matrix J

    # Extract the local vector
    #PETSc.DMGlobalToLocal(user_ctx.dm, cx_g,  PETSc.INSERT_VALUES,  user_ctx.x_l)
    PETSc.update!(user_ctx.x_l, ptr_x_g, PETSc.INSERT_VALUES)
    x               =   PETSc.unsafe_localarray(Float64, user_ctx.x_l.ptr;  write=false, read=true);

    # Check the sparsity pattern
    if isnothing(user_ctx.jac)
        error("You first have to define the sparsity pattern of the jacobian")
    else
        jac     =   user_ctx.jac;
        colors  =   user_ctx.colors;
    end
    out         =   similar(x);

    f_Res           =   ((out,x)->func(out, x, user_ctx));  # pass additional arguments into the routine

    if isnothing(user_ctx.jac_cache)
        # Allocate data required for jacobian computations
        user_ctx.jac_cache = ForwardColorJacCache(f_Res,x; colorvec=colors, sparsity = jac);
    end

    # Compute Jacobian using automatic differenciation
    forwarddiff_color_jacobian!(jac, f_Res, x, user_ctx.jac_cache);

    # Store global Jacobian inside Petsc Matrix
    ind             =   PETSc.LocalInGlobalIndices(user_ctx.dm);
    if PETSc.assembled(J) == false
        println("Assembling J")
        J           =   PETSc.MatSeqAIJ(sparse(jac[ind,ind]));
    else
        J           .=   sparse(jac[ind,ind]);
    end

    # Store Julia matrix and coloring
    user_ctx.jac    =   jac;
    user_ctx.colors =   colors;

   return jac, ind
end

# Define a struct that holds data we need in the local SNES routines below
mutable struct Data_Stokes2D
    # DMs and vectors
    dm
    dmCoeff
    coeff_l
    x_l
    f_l
    # physical parameters
    eta1
    eta2
    rho1
    rho2
    gz
    kappa
    # dimensions
    dx
    dz
    xlim
    zlim
    # jacobian and sparsity pattern
    jac
    jac_cache
    colors
end
user_ctx = Data_Stokes2D(nothing, nothing, nothing, nothing, nothing, nothing, nothing, nothing, nothing, nothing, nothing, nothing, nothing, nothing, nothing, nothing, nothing, nothing);  # holds data we need in the local

function ComputeLocalResidual(dm, ArrayLocal_x, ArrayLocal_f, user_ctx)
    # Compute the local residual. The vectors include ghost points

    # Compute shear stresses
    Txx,Tzz,Txz = ComputeStresses!(user_ctx, ArrayLocal_x);

    # Extracting arrays from the residual and solution vectors
    P       = PETSc.DMStagGetGhostArrayLocationSlot(user_ctx.dm,ArrayLocal_x,     PETSc.DMSTAG_ELEMENT, 0);
    Vx      = PETSc.DMStagGetGhostArrayLocationSlot(user_ctx.dm,ArrayLocal_x,     PETSc.DMSTAG_LEFT, 0);
    Vz      = PETSc.DMStagGetGhostArrayLocationSlot(user_ctx.dm,ArrayLocal_x,     PETSc.DMSTAG_DOWN, 0);

    f_p     = PETSc.DMStagGetGhostArrayLocationSlot(user_ctx.dm,ArrayLocal_f,     PETSc.DMSTAG_ELEMENT, 0);
    f_x     = PETSc.DMStagGetGhostArrayLocationSlot(user_ctx.dm,ArrayLocal_f,     PETSc.DMSTAG_LEFT, 0);
    f_z     = PETSc.DMStagGetGhostArrayLocationSlot(user_ctx.dm,ArrayLocal_f,     PETSc.DMSTAG_DOWN, 0);

    # Extracting physical parameters
    RhoC    = PETSc.DMStagGetGhostArrayLocationSlot(user_ctx.dmCoeff,user_ctx.coeff_l,     PETSc.DMSTAG_DOWN_LEFT, 1);

    # Gettins sizes
    dx = user_ctx.dx;
    dz = user_ctx.dz;

    # Getting global indices for the local vectors
    ind                 =   PETSc.DMStagGetIndices(dm);          # indices of (center/element) points, not including ghost values.
    #sx, sn  = PETSc.DMStagGetCentralNodes(user_ctx.dm); #indices of central points

    ix      =   ind.center[1];
    iz      =   ind.center[1];

    # Horizontal force balance f(x)

    # Boundary conditions (Dirichlet with no normal velocity at boundary)
    f_x[ix[1]    ,:] .= Vx[ix[1]    ,:] .- 0.0;
    f_x[ix[end]+1,:] .= Vx[ix[end]+1,:] .+ 0.0;

    # -dP/dx + dTxx/dx + dTxz/dz = 0
    f_x[ix[2]:ix[end],iz] .= .-(P[ix[2]:ix[end],iz] .- P[ix[1]:ix[end-1],iz])./dx .+
                               (Txx[2:end,:]        .- Txx[1:end-1,:])       ./dx .+
                               (Txz[2:end-1,2:end]  .- Txz[2:end-1,1:end-1]) ./dz;

    # Vertical force balance f(z)

    # Boundary conditions (Dirichlet with no normal velocity at boundary)
    f_z[:,iz[1]    ] .= Vz[:,iz[1]    ] .- 0.0;
    f_z[:,iz[end]+1] .= Vz[:,iz[end]+1] .+ 0.0;

    # -dP/dz + dTxz/dx + dTzz/dz = rho*g
    f_z[ix,iz[2]:iz[end]] .= .-(P[ix,iz[2]:iz[end]] .- P[ix,iz[1]:iz[end-1]])./dz .+
                               (Tzz[:,2:end]        .- Tzz[:,1:end-1])       ./dz .+
                               (Txz[ 2:end,2:end-1] .- Txz[1:end-1,2:end-1]) ./dx .+
                               (RhoC[ix.+1,iz[2]:iz[end]] .+ RhoC[ix,iz[2]:iz[end]  ]) .* 0.5 .* user_ctx.gz;

    # Mass balance f(p)
    # dVx/dx + dVz/dz = 0
    kappa = user_ctx.kappa; # penalty term to help solver
    f_p[ix,iz] .= (Vx[ix.+1,iz].-Vx[ix,iz])./dx .+ (Vz[ix,iz.+1].-Vz[ix,iz])./dz .+ 1/kappa*P[ix,iz];

    # Cleanup
    Base.finalize(Vx);    Base.finalize(Vz);
    Base.finalize(P);     Base.finalize(RhoC);
    Base.finalize(f_x);   Base.finalize(f_z);
    Base.finalize(Txx);   Base.finalize(Tzz);
    Base.finalize(Txz);
end

function ComputeStresses!(user_ctx, ArrayLocal_x)

    # Compute strain rates
    Exx,Ezz,Exz = ComputeStrainRates!(user_ctx, ArrayLocal_x);

    # Getting Eta at the center and corner of the cells
    EtaE    = PETSc.DMStagGetGhostArrayLocationSlot(user_ctx.dmCoeff,user_ctx.coeff_l,     PETSc.DMSTAG_ELEMENT, 0);
    EtaC    = PETSc.DMStagGetGhostArrayLocationSlot(user_ctx.dmCoeff,user_ctx.coeff_l,     PETSc.DMSTAG_DOWN_LEFT, 0);

    # Getting indices for center nodes (not ghost)
    ind     = PETSc.DMStagGetIndices(user_ctx.dm); #indices of central points
    #ix      =   sx[1]:sn[1];
    #iz      =   sx[2]:sn[2];
    ix  = ind.center[1];
    iz  = ind.center[1];

    # Compute shear stresses
    Txx    = 2 .* EtaE[ix,iz] .* Exx;
    Tzz    = 2 .* EtaE[ix,iz] .* Ezz;
    Txz    = 2 .* EtaC[ix[1]:ix[end]+1,iz[1]:iz[end]+1] .* Exz;

    return Txx,Tzz,Txz

    Base.finalize(Exx);   Base.finalize(Ezz);
    Base.finalize(Exz);

end

function ComputeStrainRates!(user_ctx, ArrayLocal_x)

    # Getting velocity vectors
    Vx      = PETSc.DMStagGetGhostArrayLocationSlot(user_ctx.dm,ArrayLocal_x,     PETSc.DMSTAG_LEFT, 0);
    Vz      = PETSc.DMStagGetGhostArrayLocationSlot(user_ctx.dm,ArrayLocal_x,     PETSc.DMSTAG_DOWN, 0);

    # Gettind dimensions
    dx = user_ctx.dx;
    dz = user_ctx.dz;

    # Getting indices for center nodes (not ghost)
    ind  = PETSc.DMStagGetIndices(user_ctx.dm); #indices of central points
    #ix      =   sx[1]:sn[1];
    #iz      =   sx[2]:sn[2];
    ix      =   ind.center[1];
    iz      =   ind.center[2];

    # Set ghost points (free slip)
    Vx[:,iz[1]-1]   .=  Vx[:,iz[1]];                # ghost points on left and right size (here: free slip)
    Vx[:,iz[end]+1] .=  Vx[:,iz[end]];
    Vz[ix[1]-1,:]   .=  Vz[ix[1],:];
    Vz[ix[end]+1,:] .=  Vz[ix[end],:];

    # Compute deviatoric Strain rates
    DivV  = (Vx[ix.+1,iz].-Vx[ix,iz])./dx .+ (Vz[ix,iz.+1].-Vz[ix,iz])./dz;
    Exx   = (Vx[ix.+1,iz].-Vx[ix,iz])./dx .- 1/3 .* DivV;
    Ezz   = (Vz[ix,iz.+1].-Vz[ix,iz])./dz .- 1/3 .* DivV;
    Exz   = 0.5.*((Vx[ix[1]:ix[end].+1,iz[1]:iz[end].+1].-Vx[ix[1]:ix[end].+1,iz[1].-1:iz[end]])./dz .+
                  (Vz[ix[1]:ix[end].+1,iz[1]:iz[end].+1].-Vz[ix[1].-1:ix[end],iz[1]:iz[end].+1])./dx);

    return Exx,Ezz,Exz

end

function PopulateCoefficientData!(user_ctx)

    # Create coefficient local vector and array (to store eta and rho)
    user_ctx.coeff_l    =   PETSc.createlocalvector(user_ctx.dmCoeff);
    coeff_array         =   PETSc.DMStagVecGetArray(user_ctx.dmCoeff,user_ctx.coeff_l);

    # Create Coordinates dm, local vector and array
    dm_coord  = PETSc.getcoordinateDM(user_ctx.dmCoeff);
    vec_coord = PETSc.getcoordinateslocal(user_ctx.dmCoeff);
    coord     = PETSc.DMStagVecGetArray(dm_coord, vec_coord);
   # start,n,nExtra = PETSc.DMStagGetCorners(user_ctx.dmCoeff);

    # Get the correct entries for each of our variables in local element-wise storage
    iec = PETSc.DMStagGetLocationSlot(user_ctx.dmCoeff, PETSc.DMSTAG_DOWN_LEFT, 0);  # location eta corner
    irc = PETSc.DMStagGetLocationSlot(user_ctx.dmCoeff, PETSc.DMSTAG_DOWN_LEFT, 1);  # location rho corner
    iee = PETSc.DMStagGetLocationSlot(user_ctx.dmCoeff, PETSc.DMSTAG_ELEMENT, 0);    # location eta element
    ixc = PETSc.DMStagGetLocationSlot(dm_coord, PETSc.DMSTAG_DOWN_LEFT, 0);          # location coord corner
    ixe = PETSc.DMStagGetLocationSlot(dm_coord, PETSc.DMSTAG_ELEMENT, 0);            # location coord element

    # Fill coefficient vectors with suitable value (determined by x-coordinate in function GetPhase())

    # Element nodes (eta)

    Coeff          = coeff_array[:,:,iee+1];
    X_coord        = coord[:,:,ixe+1];
    index1         = findall(x -> GetPhase(user_ctx,x,1),X_coord);
    Coeff[index1] .= user_ctx.eta1;

    index2         = findall(x -> GetPhase(user_ctx,x,2),X_coord);
    Coeff[index2] .= user_ctx.eta2;

    coeff_array[:,:,iee+1] .= Coeff;

    # Corner nodes (rho and eta)

    CoeffE         = coeff_array[:,:,iec+1];
    CoeffR         = coeff_array[:,:,irc+1];
    X_coord        = coord[:,:,ixc+1];
    index1         = findall(x -> GetPhase(user_ctx,x,1),X_coord);
    CoeffE[index1].= user_ctx.eta1;
    CoeffR[index1].= user_ctx.rho1;

    index2         = findall(x -> GetPhase(user_ctx,x,2),X_coord);
    CoeffE[index2].= user_ctx.eta2;
    CoeffR[index2].= user_ctx.rho2;

    coeff_array[:,:,iec+1] .= CoeffE;
    coeff_array[:,:,irc+1] .= CoeffR;

end

function GetPhase(ctx,x,n)
    # Divide domain in phases according to x-coordinate (similar to SolCx benchmark)
    if x < (ctx.xlim[2]-ctx.xlim[1])/2
        if n == 1 return true else return false end
    else
        if n == 1 return false else return true end
    end
end

function SetVecX!(user_ctx,x_g)
    # Set first guess for vector x to 0
    user_ctx.x_l .= 0;
    #PETSc.DMLocalToGlobal(user_ctx.dm,user_ctx.x_l, PETSc.INSERT_VALUES, x_g);
    PETSc.update!(x_g,user_ctx.x_l, PETSc.INSERT_VALUES);

end

# Main Solver
nx, nz           =   32, 32;                  # number of nodes is x and z direction
#nx, nz           =   8, 8;                  # number of nodes is x and z direction

user_ctx.xlim    =   [0,2];                   # x and z dimensions
user_ctx.zlim    =   [0,1];
xlim             =   user_ctx.xlim;
zlim             =   user_ctx.zlim;
user_ctx.dx      =   (xlim[2]-xlim[1])/nx;   # x-resolution
user_ctx.dz      =   (zlim[2]-zlim[1])/nz;   # z-resolution
user_ctx.eta1    =   1;                      # viscosity phase 1
user_ctx.eta2    =   2;                      # viscosity phase 2
user_ctx.rho1    =   3;                      # density phase 1
user_ctx.rho2    =   4;                      # density phase 2
user_ctx.gz      =   -1;                     # gravity magnitude
user_ctx.kappa   =   1e2;                    # incompressible penaly term

# Create Solution and coefficient DMs
user_ctx.dm           =   PETSc.DMStagCreate2d(petsclib,MPI.COMM_SELF,PETSc.DM_BOUNDARY_GHOSTED,PETSc.DM_BOUNDARY_GHOSTED,nx,nz,1,1,0,1,1,PETSc.DMSTAG_STENCIL_BOX,1);  # V edge, P Element
PETSc.DMStagSetUniformCoordinatesExplicit(user_ctx.dm, xlim[1], xlim[2], zlim[1], zlim[2]);            # set coordinates
user_ctx.dmCoeff      =   PETSc.DMStagCreateCompatibleDMStag(user_ctx.dm,2,0,1,0);   # rho and eta on VERTEX, eta on ELEMENT
PETSc.DMStagSetUniformCoordinatesExplicit(user_ctx.dmCoeff, xlim[1], xlim[2], zlim[1], zlim[2]);

# Populate phases
PopulateCoefficientData!(user_ctx);

# Create solution and residual vectors
x_g             =   PETSc.createglobalvector(user_ctx.dm);
f_g             =   PETSc.createglobalvector(user_ctx.dm);
user_ctx.x_l    =   PETSc.createlocalvector(user_ctx.dm);
user_ctx.f_l    =   PETSc.createlocalvector(user_ctx.dm);

# Create Petsc matrix for Jacobian
J               =   PETSc.creatematrix(user_ctx.dm);

# Set first guess values for solution vector
SetVecX!(user_ctx,x_g);

# Compute jacobian sparcity pattern
x_l               =   PETSc.unsafe_localarray(Float64, user_ctx.x_l.ptr;  write=false, read=true)

user_ctx.jac, user_ctx.colors = ComputeSparsityPatternJacobian_automatic(user_ctx.x_l, user_ctx);
ind             =   PETSc.LocalInGlobalIndices(user_ctx.dm);    # extract indices
J           =   PETSc.MatSeqAIJ(sparse(user_ctx.jac[ind,ind]));

#Setting up SNES
S = PETSc.SNES{Float64}(petsclib,MPI.COMM_SELF;
        snes_rtol=1e-12,
        snes_monitor=true,
        snes_max_it = 500,
        snes_monitor_true_residual=true,
        snes_converged_reason=true);
S.user_ctx  =       user_ctx;

PETSc.setfunction!(S, FormRes!, f_g)
PETSc.setjacobian!(S, FormJacobian!, J, J)

#Solve
PETSc.solve!(x_g, S);

# Copy solution to local vector
PETSc.update!(user_ctx.x_l, x_g,  PETSc.INSERT_VALUES);

# Extract solution
Vx  =   PETSc.DMStagGetGhostArrayLocationSlot(user_ctx.dm, user_ctx.x_l,   PETSc.DMSTAG_LEFT,      0);
Vz  =   PETSc.DMStagGetGhostArrayLocationSlot(user_ctx.dm, user_ctx.x_l,   PETSc.DMSTAG_DOWN,      0);
P   =   PETSc.DMStagGetGhostArrayLocationSlot(user_ctx.dm, user_ctx.x_l,   PETSc.DMSTAG_ELEMENT,   0);
Eta =   PETSc.DMStagGetGhostArrayLocationSlot(user_ctx.dmCoeff, user_ctx.coeff_l,   PETSc.DMSTAG_ELEMENT,   0);

# Getting indices for center nodes (not ghost)
ind     = PETSc.DMStagGetIndices(user_ctx.dm);
ix      =   ind.center[1];
iz      =   ind.center[2];

# Getting rid of ghost points
Vx  =   Vx[ix[1]:ix[end]+1,iz];
Vz  =   Vz[ix,iz[1]:iz[end]+1];
P   =    P[ix,iz];
Eta =  Eta[ix,iz];

# Compute central average velocities
Vx_cen = (Vx[2:end,:]  + Vx[1:end-1,:])/2;
Vz_cen = (Vz[:,2:end]  + Vz[:,1:end-1])/2;

#get coordinates
dm_coord  = PETSc.getcoordinateDM(user_ctx.dmCoeff);
vec_coord = PETSc.getcoordinateslocal(user_ctx.dmCoeff);

XCoord_e = PETSc.DMStagGetGhostArrayLocationSlot(dm_coord, vec_coord,   PETSc.DMSTAG_ELEMENT  ,      0); # location coord corner
ZCoord_e = PETSc.DMStagGetGhostArrayLocationSlot(dm_coord, vec_coord,   PETSc.DMSTAG_ELEMENT  ,      1); # location coord corner
XCoord_c = PETSc.DMStagGetGhostArrayLocationSlot(dm_coord, vec_coord,   PETSc.DMSTAG_DOWN_LEFT,      0);   # location coord element
ZCoord_c = PETSc.DMStagGetGhostArrayLocationSlot(dm_coord, vec_coord,   PETSc.DMSTAG_DOWN_LEFT,      1);   # location coord element

XCoord_e = XCoord_e[ix,iz];
ZCoord_e = ZCoord_e[ix,iz];
XCoord_c = XCoord_c[ix[1]:ix[end]+1,iz[1]:iz[end]+1];
ZCoord_c = ZCoord_c[ix[1]:ix[end]+1,iz[1]:iz[end]+1];

xe_1D    = XCoord_e[:,1];
ze_1D    = ZCoord_e[1,:];
xc_1D    = XCoord_c[:,1];
zc_1D    = ZCoord_c[1,:];

# Plot
CreatePlots = false;        # set to false by default such that testing does not require Plots
if CreatePlots==true
    using Plots

    p1 = heatmap(xe_1D,ze_1D, P', xlabel="Width", ylabel="Depth", title="Pressure",aspect_ratio = 1)
    p2 = heatmap(xe_1D,zc_1D, Vz', xlabel="Width", ylabel="Depth", title="Vz",aspect_ratio = 1)
    p3 = heatmap(xc_1D,ze_1D, Vx', xlabel="Width", ylabel="Depth", title="Vx",aspect_ratio = 1)
    quiver!(XCoord_e[:],ZCoord_e[:],quiver=(Vx_cen[:]*0.02,Vz_cen[:]*0.02), color=:white,aspect_ratio = 1)
    p4 = heatmap(xe_1D,ze_1D, Eta', xlabel="Width", ylabel="Depth", title="Eta",aspect_ratio = 1)
    plot(p1, p2, p3, p4, layout = (2, 2), legend = false)

end

# check if the solution makes sense
using LinearAlgebra, Test
check =  norm(user_ctx.x_l)
@test check ≈ 31.54688642064483


text:
typealias IndexList Vector{Int}
typealias ComplexMatrix AbstractMatrix{Complex}
#typealias Vector Vector

type PowerFlowSystem
  buses::Vector # Vector of bus names/indices

  pq::IndexList # pq bus indices
  pv::IndexList # pv bus indices
  sl::IndexList # slack bus indices

  from::IndexList # list of pair "from" bus indices
  to::IndexList   # list of pair "to"  bus indices

  Yᵈ::ComplexMatrix # diagonal of Y matrix
  Yᶠ::ComplexMatrix # admittance of edges exiting the bus
  Yᵗ::ComplexMatrix # admittance of edges entering the bus
  Y::ComplexMatrix  # admittance matrix

  v⁰::Vector # base solution voltage magnitudes
  ϑ⁰::Vector # base solution voltage angles

  p⁰::Vector # base solution active powers
  q⁰::Vector # base solution reactive powers

  nb::Integer # total number of buses
  np::Integer # total number of pairs

  npq::Integer # total number of PQ buses
  npv::Integer # total number of PV buses
  nsl::Integer # total number of Slack buses

  function PowerFlowSystem(buses, pq, pv, slack, from, to, Yᵈ, Yᶠ, Yᵗ, Y, v⁰, ϑ⁰ , p⁰, q⁰)
    nb = length(buses)
    np = length(from)
    npq = length(pq)
    npv = length(pv)
    nsl = length(slack)
    new(buses, pq, pv, slack, from, to, Yᵈ, Yᶠ, Yᵗ, Y, v⁰, ϑ⁰, p⁰, q⁰, nb, np, npq, npv, nsl)
  end
end

function PowerFlowSystem(pm::GenericPowerModel)
  net = Network(pm)
  bt = get_bus_types(pm, net)
  YYY = calc_YYY(pm, net)
  Y = calc_Y(YYY..., net)
  vϑ = get_vϑ(pm, net)
  pq = get_pq(pm, net)
  return PowerFlowSystem(net.buses, bt..., net.from, net.to, YYY..., Y, vϑ..., pq...)
end

function PowerFlowSystem(pm::GenericPowerModel, sol::Dict)
  net = Network(pm)
  bt = get_bus_types(pm, net)
  YYY = calc_YYY(pm, net)
  Y = calc_Y(YYY..., net)
  vϑ = get_sol_vϑ(sol, net)
  pq = get_pq(pm, net)
  return PowerFlowSystem(net.buses, bt..., net.from, net.to, YYY..., Y, vϑ..., pq...)
end

function calc_pq(Y::AbstractMatrix, v::Vector, ϑ::Vector)
  V = v.*exp(im*ϑ)
  return reim(V.*conj(Y*V))
end

function calc_Y(Yᵈ::AbstractMatrix, Yᶠ::AbstractMatrix, Yᵗ::AbstractMatrix, net::Network)
  @import_fields net
  Cf = sparse(1:np,   to, 1.0 + 0im, np, nb)
  Ct = sparse(1:np, from, 1.0 + 0im, np, nb)
  return Yᵈ + Yᶠ*Cf + Yᵗ*Ct
end

export PowerFlowSystem, calc_pq, calc_Y

max_stars_repo_name: hunghtd/PowerFlowAnalysis-tutorial

max_issues_repo_name: hunghtd/PowerFlowAnalysis-tutorial

max_forks_repo_name: hunghtd/PowerFlowAnalysis-tutorial


text:
using Distributions
using Base.LinAlg.BLAS
using Devectorize
import Base.getindex

typealias Gaussian Normal
abstract AbstractRBM

"""
    # Boltzmann.RBM{V,H} (RBMBase.jl)
    ## Description
        A structure for containing all of the restricted Boltzmann Machine (RBM)
        model parameters. Besides just the model parameters (couplings, biases),
        the structure also contains variables which are pertinent to the RBM training
        procedure.

    ## Structure
        - `W::Matrix{Float64}`:       The matrix of coupling parameters (RBM model parameter)
        - `W2::Matrix{Float64}`:      The square of `W` (used for EMF learning)
        - `W3::Matrix{Float64}`:      The cube of `W` (used for EMF learning)
        - `vbias::Vector{Float64}`:   The visible unit biases (RBM model parameter)
        - `hbias::Vector{Float64}`:   The hidden unit biases (RBM model parameter)
        - `dW::Matrix{Float64}`:      The current gradient on the coupling parameters (used for RBM training)
        - `dW_prev::Matrix{Float64}`: The last gradient on the coupling parmaters (used for RBM training)
        - `persistent_chain_vis::Matrix{Float64}`: Visible fantasy particles (used for RBM persistent mode training)
        - `persistent_chain_hid::Matrix{Float64}`: Hidden fantasy particles (used for RBM persistent mode training)
        - `momentum::Float64`:        Amount of last gradient to add to the current gradient (used for RBM training)
        - `VisShape::Tuple{Int,Int}`: Final output shape of the visible units
"""
type RBM{V,H} <: AbstractRBM
    W::Matrix{Float64}
    W2::Matrix{Float64}
    W3::Matrix{Float64}
    vbias::Vector{Float64}
    hbias::Vector{Float64}
    dW::Matrix{Float64}
    dW_prev::Matrix{Float64}
    persistent_chain_vis::Matrix{Float64}
    persistent_chain_hid::Matrix{Float64}
    momentum::Float64
    VisShape::Tuple{Int,Int}
end

function RBM(V::Type, H::Type, n_vis::Int, n_hid::Int, visshape::Tuple{Int,Int}; sigma=0.01, momentum=0.0, TrainData=[])
    # Some "tiny" value, used to enforce min/max boundary conditions
    eps = 1e-8

    # Initialize the weighting matrix by drawing from an iid Gaussian
    # of the specified standard deviation.
    W = rand(Normal(0, sigma), (n_hid, n_vis))

    # If the user specifies the training dataset, it can be useful to
    # initialize the visibile biases according to the empirical expected
    # feature values of the training data.
    #
    # TODO: Generalize this biasing. Currently, the biasing is only written for
    #       the case of binary RBMs.
    InitialVisBias = zeros(n_vis,1)
    if !isempty(TrainData)
        ProbVis = mean(TrainData,2)             # Mean across  samples
        ProbVis = max(ProbVis,eps)              # Some regularization (avoid Inf/NaN)
        ProbVis = min(ProbVis,1 - eps)          # ''
        @devec InitialVisBias = log(ProbVis ./ (1-ProbVis)) # Biasing as the log-proportion
    end

    RBM{V,H}(W,                                             # W
             W.*W,                                          # W2
             W.*W.*W,                                       # W3
             vec(InitialVisBias),                           # vbias
             zeros(n_hid),                                  # hbias
             zeros(n_hid, n_vis),                           # dW
             zeros(n_hid, n_vis),                           # dW_prev
             Array(Float64, 0, 0),                          # persistent_chain_vis
             Array(Float64, 0, 0),                          # persistent_chain_hid
             momentum,                                      # momentum
             visshape)                                      # Shape of the visible units (for display)
end

function Base.show{V,H}(io::IO, rbm::RBM{V,H})
    n_vis = size(rbm.vbias, 1)
    n_hid = size(rbm.hbias, 1)
    print(io, "RBM{$V,$H}($n_vis, $n_hid)")
end

typealias BernoulliRBM RBM{Bernoulli, Bernoulli}
BernoulliRBM(n_vis::Int, n_hid::Int, visshape::Tuple{Int,Int}; sigma=0.1, momentum=0.0, TrainData=[]) =
    RBM(Bernoulli, Bernoulli, n_vis, n_hid, visshape; sigma=sigma, momentum=momentum, TrainData=TrainData)

typealias GRBM RBM{Gaussian, Bernoulli}
GRBM(n_vis::Int, n_hid::Int, visshape::Tuple{Int,Int}; sigma=0.1, momentum=0.0, TrainData=[]) =
    RBM(Gaussian, Bernoulli, n_vis, n_hid, visshape; sigma=sigma, momentum=momentum, TrainData=TrainData)

"""
    # Boltzmann.PassHidToVis  (RBMBase.jl)
"""
function PassHidToVis(rbm::RBM, hid::Mat{Float64})
    return rbm.W' * hid .+ rbm.vbias
end

"""
    # Boltzmann.PassVisToHid  (RBMBase.jl)
"""
function PassVisToHid(rbm::RBM, vis::Mat{Float64})
    return rbm.W * vis .+ rbm.hbias
end

### These functions need to be generalized to detect the Distribution on
### the hidden and visible variables.
"""
    # Boltzmann.ProbHidCondOnVis  (RBMBase.jl)
"""
function ProbHidCondOnVis(rbm::RBM, vis::Mat{Float64})
    return logsig(PassVisToHid(rbm,vis))
end

"""
    # Boltzmann.ProbVisCondOnHid  (RBMBase.jl)
"""
function ProbVisCondOnHid(rbm::RBM, hid::Mat{Float64})
    return logsig(PassHidToVis(rbm,hid))
end

"""
    # Boltzmann.PinField! (RBMBase.jl)
"""
function PinField!(rbm::RBM,pinning_field::Vec{Float64})
    pos_inf_locations = pinning_field > 0
    neg_inf_locations = pinning_field < 0

    rbm.vbias(pos_inf_locations) =  Inf
    rbm.vbias(neg_inf_locations) = -Inf
end


text:
### acquisition optimization setup

defaultoptions(::Type{<:GPE}, ::Type{<:AbstractAcquisition}) =
    (method = :LD_LBFGS, restarts = 10, maxeval = 2000)
defaultoptions(::Type{<:GPE}, ::Type{ThompsonSamplingSimple}) =
    (method = :GN_DIRECT_L, restarts = 1, maxeval = 2000)

function wrap_gradient(f)
    (x, g) -> begin
        res = DiffResults.DiffResult(0., g)
        ForwardDiff.gradient!(res, f, x)
        res.value
    end
end
wrap_dummygradient(f) = (x, g) -> f(x)

function nlopt_setup(a::AbstractAcquisition, model, lowerbounds, upperbounds,
                     options)
    D = length(lowerbounds)
    opt = NLopt.Opt(options.method, D)
    NLopt.maxeval!(opt, options.maxeval)
    NLopt.lower_bounds!(opt, lowerbounds)
    NLopt.upper_bounds!(opt, upperbounds)
    setparams!(a, model)
    if string(options.method)[2] == 'D'
        f = wrap_gradient(acquisitionfunction(a, model))
        NLopt.max_objective!(opt, f)
    else
        NLopt.max_objective!(opt, wrap_dummygradient(acquisitionfunction(a, model)))
    end
    opt
end

### optimize acquisition function

acquire_max(o) = acquire_max(o.acquisition, o.model, o.lowerbounds, o.upperbounds, o.acquisitionoptions)
acquire_model_max(o; options = o.acquisitionoptions) = acquire_max(MaxMean(), o.model, o.lowerbounds, o.upperbounds, options)
# TODO: stochastic acquisition like Gutmann and Corander (2016), p. 20
function acquire_max(a::AbstractAcquisition, model, lowerbounds, upperbounds, options)
    opt = nlopt_setup(a, model, lowerbounds, upperbounds, options)
    acquire_max(opt, lowerbounds, upperbounds, options.restarts)
end

function acquire_max(opt, lowerbounds, upperbounds, restarts)
    maxf = -Inf
    maxx = lowerbounds
    x0s = latin_hypercube_sampling(lowerbounds, upperbounds, restarts)
    for i in 1:restarts
        f, x, ret = NLopt.optimize(opt, x0s[:, i])
#         println("$x0, $x, $f, $ret")
        ret == NLopt.FORCED_STOP && throw(InterruptException())
        if f > maxf
            maxf = f
            maxx = x
        end
    end
    maxf, maxx
end


text:
@recipe function f{Tx,To}(md::MultiData{Tx,To})
    seriestype --> :scatter
    xguide --> L"$f_1$"
    yguide --> L"$f_2$"
    zguide --> L"$f_3$"
    label --> ""
    numobjectives = length(md.objectives)
    if numobjectives > 3 || numobjectives < 2
        Base.error("Only plotting 2d and 3d fronts")
    end

    f1arr = convert(Array{Float64},
                    [val[1] for val in md.paretofront])
    f2arr = convert(Array{Float64},
                    [val[2] for val in md.paretofront])

    title --> "Pareto front with $(length(f1arr)) points"

    if numobjectives == 2
        xyz = (f1arr, f2arr)
    else
        f3arr = convert(Array{Float64},
                        [val[3] for val in md.paretofront])
        xyz = (f1arr, f2arr, f3arr)
    end
    (xyz...,)
end

@recipe function f(m::JuMP.Model)
    plot(get_multidata(m))
end


text:
@kwdef mutable struct BiodiversityObservationNetwork{IT<:Integer,M<:AbstractMatrix}
    numobservatories::IT = 50
    coordinates::M = missing
end

abstract type SpatialSampler end

struct SpatialSimulatedAnnealing <: SpatialSampler end


text:
module MapViews

using LinearAlgebra

export MapView

"""
MapView is similar to a regular Julia view. However, the data it accesses
is determined by an internal map. This allows for a reorganization
of the array in question. Used in cases where it is convenient to have
two different orderings of the same array at the same time.
"""
struct MapView{T, N} <: AbstractArray{T, N}
    indices::Array{Int, N}
    original::Array{T, N}
end

######################################
# AbstractArray interfaces for MapView
Base.IndexStyle(::Type{MapView}) = IndexLinear()
Base.getindex(X::MapView, i::Int) = X.original[X.indices[i]]

function Base.setindex!(X::MapView{T}, v::V, i::Int) where {T, V}
    X.original[X.indices[i]] = v
end

Base.firstindex(X::MapView) = 1
Base.lastindex(X::MapView) = length(X)

######################################
# Iterator interfaces for MapView
function Base.iterate(iter::MapView)
    if length(iter) == 0
        nothing
    else
        iter[1], 1
    end
end

function Base.iterate(iter::MapView, state)
    if state >= length(iter)
        nothing
    else
        iter[state + 1], state + 1
    end
end

Base.eltype(::Type{MapView{T}}) where T = T
Base.size(iter::MapView) = size(iter.indices)

##############################
# Show interface for MapView
Base.show(io::IO, x::MapView) = Base.show(io, collect(x))

end


text:
module problem123

description = """
Prime square remainders

Let pn be the nth prime: 2, 3, 5, 7, 11, ..., and let r be the remainder when (pn−1)^n + (pn+1)^n is divided by pn^2.
For example, when n = 3, p3 = 5, and 4^3 + 6^3 = 280 ≡ 5 mod 25.
The least value of n for which the remainder first exceeds 10^9 is 7037.

Find the least value of n for which the remainder first exceeds 10^10.
"""

# Computes a^b % n
# Definition from Wikipedia on Modular exponentiation: http://en.wikipedia.org/wiki/Modular_exponentiation
# BigInt is required, as a*a will overflow for large a and n
function powmod(a :: BigInt, b :: Integer, n :: Integer)
  x = BigInt(1)
  a = a % n
  while b > 0
    if (b % 2) == 1
      x = (x * a) % n
    end
    b = b >> 1
    a = (a * a) % n
  end
  return x
end

function solve(rmin = 10^10, maxprime = 10^6)
  ps = primes(maxprime)
  for n = 1:length(ps)
    p = BigInt(ps[n])
    psq = p * p
    if psq < rmin
      continue
    end
    r = (powmod(p-1, n, psq) + powmod(p+1, n, psq)) % psq
    if r > rmin
      return n
    end
  end
  error("No solution for rmin=$rmin maxprime=$maxprime")
end

using Base.Test

@test (powmod(BigInt(4), 3, 25) + powmod(BigInt(6), 3, 25)) % 25 == 5
@test solve(10^9) == 7037

end

max_stars_repo_name: mbuhot/mbuhot-euler-solutions

max_issues_repo_name: mbuhot/mbuhot-euler-solutions

max_forks_repo_name: mbuhot/mbuhot-euler-solutions


text:
module jInv

include("Utils/Utils.jl")
include("Mesh/Mesh.jl")
include("LinearSolvers/LinearSolvers.jl")
include("ForwardShare/ForwardShare.jl")
include("InverseSolve/InverseSolve.jl")
end


text:
# AUTO GENERATED FILE - DO NOT EDIT

export ''_antdscatter

"""
    ''_antdscatter(;kwargs...)

An AntdScatter component.

Keyword arguments:
- `id` (String; optional)
- `annotations` (optional)
- `appendPadding` (Real | Array of Reals | String; optional)
- `autoFit` (Bool; optional)
- `className` (String; optional)
- `color` (optional): . color has the following type: String | Array of Strings | lists containing elements 'func'.
Those elements have the following types:
  - `func` (String; optional)
- `colorField` (String; optional)
- `data` (Array of Dicts; required)
- `height` (Real; optional)
- `label` (optional)
- `legend` (optional)
- `loading_state` (optional): . loading_state has the following type: lists containing elements 'is_loading', 'prop_name', 'component_name'.
Those elements have the following types:
  - `is_loading` (Bool; optional): Determines if the component is loading or not
  - `prop_name` (String; optional): Holds which property is loading
  - `component_name` (String; optional): Holds the name of the component that is loading
- `locale` (String; optional)
- `meta` (optional)
- `padding` (Real | Array of Reals | String; optional)
- `pointStyle` (optional): . pointStyle has the following type: lists containing elements 'func'.
Those elements have the following types:
  - `func` (String; optional)
- `quadrant` (optional): . quadrant has the following type: lists containing elements 'xBaseline', 'yBaseline', 'lineStyle', 'regionStyle', 'labels'.
Those elements have the following types:
  - `xBaseline` (Real; optional)
  - `yBaseline` (Real; optional)
  - `lineStyle` (optional)
  - `regionStyle` (Array; optional)
  - `labels` (optional): . labels has the following type: Array of lists containing elements 'content', 'position', 'style'.
Those elements have the following types:
  - `content` (String; optional)
  - `position` (Array of Reals; optional)
  - `style` (optional)s
- `regressionLine` (optional): . regressionLine has the following type: lists containing elements 'type', 'style', 'algorithm', 'top'.
Those elements have the following types:
  - `type` (a value equal to: 'exp', 'linear', 'loess', 'log', 'poly', 'pow', 'quad'; optional)
  - `style` (optional)
  - `algorithm` (optional): . algorithm has the following type: lists containing elements 'func'.
Those elements have the following types:
  - `func` (String; optional)
  - `top` (Bool; optional)
- `renderer` (String; optional)
- `shape` (optional): . shape has the following type: String | Array of Strings | lists containing elements 'func'.
Those elements have the following types:
  - `func` (String; optional)
- `shapeField` (String; optional)
- `shapeLegend` (optional)
- `size` (optional): . size has the following type: Real | Array of Reals | lists containing elements 'func'.
Those elements have the following types:
  - `func` (String; optional)
- `sizeField` (String; optional)
- `sizeLegend` (optional)
- `style` (Dict; optional)
- `tooltip` (optional)
- `width` (Real; optional)
- `xAxis` (optional)
- `xField` (String; required)
- `yAxis` (optional)
- `yField` (String; required)
"""
function ''_antdscatter(; kwargs...)
        available_props = Symbol[:id, :annotations, :appendPadding, :autoFit, :className, :color, :colorField, :data, :height, :label, :legend, :loading_state, :locale, :meta, :padding, :pointStyle, :quadrant, :regressionLine, :renderer, :shape, :shapeField, :shapeLegend, :size, :sizeField, :sizeLegend, :style, :tooltip, :width, :xAxis, :xField, :yAxis, :yField]
        wild_props = Symbol[]
        return Component("''_antdscatter", "AntdScatter", "feffery_antd_charts", available_props, wild_props; kwargs...)
end

max_stars_repo_name: CNFeffery/feffery-antd-charts

max_issues_repo_name: CNFeffery/feffery-antd-charts

max_forks_repo_name: CNFeffery/feffery-antd-charts


text:
module TestUI

using Base.Iterators: partition
using SplittablesBase: halve
using SplittablesBase.Implementations: shape
using Test

macro expect(T, ex)
    quote
        local err
        err = try
            $(esc(ex))
            nothing
        catch err
            err isa $(esc(T)) || rethrow()
            Some(err)
        end
        err === nothing && error("exception not thrown")
        something(err)
    end
end

@testset "zip" begin
    err = @expect ArgumentError halve(zip(1:3, 1:4))
    msg = sprint(showerror, err)
    @test occursin("`halve(zip(...))` requires collections with identical `size`", msg)

    err = @expect MethodError halve(zip(1:3, "abc"))
    @test err.f === shape

    if !isdefined(Iterators, :Zip1)  # VERSION >= v"1.1-"
        err = @expect ArgumentError halve(zip())
        msg = sprint(showerror, err)
        @test occursin("empty `zip` does not have `size`", msg)
    end
end

end  # module


text:
import Temporal.acf  # used for running autocorrelation function

using Statistics

"""
    diffn(x::Array{Float64}; n::Int64=1)::Array{Float64}

Lagged differencing
"""
function diffn(x::Array{Float64}; n::Int64=1)::Array{Float64}
    @assert n<size(x,1) && n>0 "Argument n out of bounds."
    dx = zeros(size(x))
    dx[1:n] .= NaN
    @inbounds for i=n+1:size(x,1)
        dx[i] = x[i] - x[i-n]
    end
    return dx
end

"""
    diffn(X::Array{Float64,2}; n::Int64=1)::Array{Float64}

Lagged differencing
"""
function diffn(X::Array{Float64,2}; n::Int64=1)::Matrix{Float64}
    @assert n<size(X,1) && n>0 "Argument n out of bounds."
    dx = zeros(size(X))
    @inbounds for j = 1:size(X,2)
        dx[:,j] = diffn(X[:,j], n=n)
    end
    return dx
end

"""
(Adapted from StatsBase: https://raw.githubusercontent.com/JuliaStats/StatsBase.jl/master/src/scalarstats.jl)

`Compute the mode of an arbitrary array::Array{Float64}`
"""
function mode(a::AbstractArray{Float64})
    isempty(a) && error("mode: input array cannot be empty.")
    cnts = Dict{Float64,Int64}()
    # first element
    mc = 1
    mv = a[1]
    cnts[mv] = 1
    # find the mode along with table construction
    @inbounds for i = 2 : length(a)
        x = a[i]
        if haskey(cnts, x)
            c = (cnts[x] += 1)
            if c > mc
                mc = c
                mv = x
            end
        else
            cnts[x] = 1
            # in this case: c = 1, and thus c > mc won't happen
        end
    end
    return mv
end

"""
    runmean(x::Array{Float64}; n::Int64=10, cumulative::Bool=true)::Array{Float64}

Compute a running or rolling arithmetic mean of an array.
"""
function runmean(x::Array{Float64}; n::Int64=10, cumulative::Bool=true)::Array{Float64}
    @assert n<size(x,1) && n>1 "Argument n is out of bounds."
    out = zeros(size(x))
    out[1:n-1] .= NaN
    if cumulative
        fi = 1.0:size(x,1)
        @inbounds for i = n:size(x,1)
            out[i] = sum(x[1:i])/fi[i]
        end
    else
        @inbounds for i = n:size(x,1)
            out[i] = mean(x[i-n+1:i])
        end
    end
    return out
end

"""
    runsum(x::Array{Float64}; n::Int64=10, cumulative::Bool=true)::Array{Float64}

Compute a running or rolling summation of an array.
"""
function runsum(x::Array{Float64}; n::Int64=10, cumulative::Bool=true)::Array{Float64}
    @assert n<size(x,1) && n>1 "Argument n is out of bounds."
    if cumulative
        out = cumsum(x, dims=1)
        out[1:n-1] .= NaN
    else
        out = zeros(size(x))
        out[1:n-1] .= NaN
        @inbounds for i = n:size(x,1)
            out[i] = sum(x[i-n+1:i])
        end
    end
    return out
end

"""
    wilder_sum(x::Array{Float64}; n::Int64=10)::Array{Float64}

Welles Wilder summation of an array
"""
function wilder_sum(x::Array{Float64}; n::Int64=10)::Array{Float64}
    @assert n<size(x,1) && n>0 "Argument n is out of bounds."
    nf = float(n)  # type stability -- all arithmetic done on floats
    out = zeros(size(x))
    out[1] = x[1]
    @inbounds for i = 2:size(x,1)
        out[i] = x[i] + out[i-1]*(nf-1.0)/nf
    end
    return out
end

"""
    runmad(x::Array{Float64}; n::Int64=10, cumulative::Bool=true, fun::Function=median)::Array{Float64}

Compute the running or rolling mean absolute deviation of an array
"""
function runmad(x::Array{Float64}; n::Int64=10, cumulative::Bool=true, fun::Function=median)::Array{Float64}
    @assert n<size(x,1) && n>1 "Argument n is out of bounds."
    out = zeros(size(x))
    out[1:n-1] .= NaN
    center = 0.0
    if cumulative
        fi = collect(1.0:size(x,1))
        @inbounds for i = n:size(x,1)
            center = fun(x[1:i])
            out[i] = sum(abs.(x[1:i].-center)) / fi[i]
        end
    else
        fn = float(n)
        @inbounds for i = n:size(x,1)
            center = fun(x[i-n+1:i])
            out[i] = sum(abs.(x[i-n+1:i].-center)) / fn
        end
    end
    return out
end

"""
    runvar(x::Array{Float64}; n::Int64=10, cumulative=true)::Array{Float64}

Compute the running or rolling variance of an array
"""
function runvar(x::Array{Float64}; n::Int64=10, cumulative=true)::Array{Float64}
    @assert n<size(x,1) && n>1 "Argument n is out of bounds."
    out = zeros(size(x))
    out[1:n-1] .= NaN
    if cumulative
        @inbounds for i = n:size(x,1)
            out[i] = var(x[1:i])
        end
    else
        @inbounds for i = n:size(x,1)
            out[i] = var(x[i-n+1:i])
        end
    end
    return out
end

"""
    runsd(x::Array{Float64}; n::Int64=10, cumulative::Bool=true)::Array{Float64}

Compute the running or rolling standard deviation of an array
"""
function runsd(x::Array{Float64}; n::Int64=10, cumulative::Bool=true)::Array{Float64}
    return sqrt.(runvar(x, n=n, cumulative=cumulative))
end

"""
    runcov(x::Array{Float64}, y::Array{Float64}; n::Int64=10, cumulative::Bool=true)::Array{Float64}

Compute the running or rolling covariance of two arrays
"""
function runcov(x::Array{Float64}, y::Array{Float64}; n::Int64=10, cumulative::Bool=true)::Array{Float64}
    @assert length(x) == length(y) "Dimension mismatch: length of `x` not equal to length of `y`."
    @assert n<size(x,1) && n>1 "Argument n is out of bounds."
    out = zeros(size(x))
    out[1:n-1] .= NaN
    if cumulative
        @inbounds for i = n:length(x)
            out[i] = cov(x[1:i], y[1:i])
        end
    else
        @inbounds for i = n:length(x)
            out[i] = cov(x[i-n+1:i], y[i-n+1:i])
        end
    end
    return out
end

"""
    runcor(x::Array{Float64}, y::Array{Float64}; n::Int64=10, cumulative::Bool=true)::Array{Float64}

Compute the running or rolling correlation of two arrays
"""
function runcor(x::Array{Float64}, y::Array{Float64}; n::Int64=10, cumulative::Bool=true)::Array{Float64}
    @assert length(x) == length(y) "Dimension mismatch: length of `x` not equal to length of `y`."
    @assert n<size(x,1) && n>1 "Argument n is out of bounds."
    out = zeros(size(x))
    out[1:n-1] .= NaN
    if cumulative
        @inbounds for i = n:length(x)
            out[i] = cor(x[1:i], y[1:i])
        end
    else
        @inbounds for i = n:length(x)
            out[i] = cor(x[i-n+1:i], y[i-n+1:i])
        end
    end
    return out
end

"""
    runmax(x::Array{Float64}; n::Int64=10, cumulative::Bool=true, inclusive::Bool=true)::Array{Float64}

Compute the running or rolling maximum of an array
"""
function runmax(x::Array{Float64}; n::Int64=10, cumulative::Bool=true, inclusive::Bool=true)::Array{Float64}
    @assert n<size(x,1) && n>1 "Argument n is out of bounds."
    out = zeros(size(x))
    if inclusive
        if cumulative
            out[n] = maximum(x[1:n])
            @inbounds for i = n+1:size(x,1)
                out[i] = max(out[i-1], x[i])
            end
        else
            @inbounds for i = n:size(x,1)
                out[i] = maximum(x[i-n+1:i])
            end
        end
        out[1:n-1] .= NaN
        return out
    else
        if cumulative
            out[n+1] = maximum(x[1:n])
            @inbounds for i = n+1:size(x,1)-1
                out[i+1] = max(out[i-1], x[i-1])
            end
        else
            @inbounds for i = n:size(x,1)-1
                out[i+1] = maximum(x[i-n+1:i])
            end
        end
        out[1:n] .= NaN
        return out
    end
end

"""
    runmin(x::Array{Float64}; n::Int64=10, cumulative::Bool=true, inclusive::Bool=true)::Array{Float64}

Compute the running or rolling minimum of an array
"""
function runmin(x::Array{Float64}; n::Int64=10, cumulative::Bool=true, inclusive::Bool=true)::Array{Float64}
    @assert n<size(x,1) && n>1 "Argument n is out of bounds."
    out = zeros(size(x))
    if inclusive
        if cumulative
            out[n] = minimum(x[1:n])
            @inbounds for i = n+1:size(x,1)
                out[i] = min(out[i-1], x[i])
            end
        else
            @inbounds for i = n:size(x,1)
                out[i] = minimum(x[i-n+1:i])
            end
        end
        out[1:n-1] .= NaN
        return out
    else
        if cumulative
            out[n+1] = minimum(x[1:n])
            @inbounds for i = n+1:size(x,1)-1
                out[i+1] = min(out[i-1], x[i-1])
            end
        else
            @inbounds for i = n:size(x,1)-1
                out[i+1] = minimum(x[i-n+1:i])
            end
        end
        out[1:n] .= NaN
        return out
    end
end

"""
    runquantile(x::Vector{Float64}; p::Float64=0.05, n::Int=10, cumulative::Bool=true)::Vector{Float64}

Compute the running/rolling quantile of an array
"""
function runquantile(x::Array{Float64}; p::Float64=0.05, n::Int=10, cumulative::Bool=true)::Array{Float64}
    @assert n<size(x,1) && n>1 "Argument n is out of bounds."
    out = zeros(Float64, (size(x,1), size(x,2)))
    if cumulative
        @inbounds for j in 1:size(x,2), i in 2:size(x,1)
            out[i,j] = quantile(x[1:i,j], p)
        end
        out[1,:] .= NaN
    else
        @inbounds for j in 1:size(x,2), i in n:size(x,1)
            out[i,j] = quantile(x[i-n+1:i,j], p)
        end
        out[1:n-1,:] .= NaN
    end
    return out
end

"""
```
function runacf(x::Array{Float64};
                n::Int = 10,
                maxlag::Int = n-3,
                lags::AbstractArray{Int,1} = 0:maxlag,
                cumulative::Bool = true)::Matrix{Float64}
```

Compute the running/rolling autocorrelation of a vector.
"""
function runacf(x::Array{Float64};
                n::Int = 10,
                maxlag::Int = n-3,
                lags::AbstractArray{Int,1} = 0:maxlag,
                cumulative::Bool = true)::Matrix{Float64}
    @assert size(x, 2) == 1 "Autocorrelation input array must be one-dimensional"
    N = size(x, 1)
    @assert n < N && n > 0
    if length(lags) == 1 && lags[1] == 0
        return ones((N, 1))
    end
    out = zeros((N, length(lags))) * NaN
    if cumulative
        @inbounds for i in n:N
            out[i,:] = acf(x[1:i], lags=lags)
        end
    else
        @inbounds for i in n:N
            out[i,:] = acf(x[i-n+1:i], lags=lags)
        end
    end
    return out
end


text:
# By default, Julia/LLVM does not use fused multiply-add operations (FMAs).
# Since these FMAs can increase the performance of many numerical algorithms,
# we need to opt-in explicitly.
# See https://ranocha.de/blog/Optimizing_EC_Trixi for further details.
@muladd begin

# Container data structure (structure-of-arrays style) for DG elements
mutable struct ElementContainer1D{RealT<:Real, uEltype<:Real} <: AbstractContainer
  inverse_jacobian::Vector{RealT}        # [elements]
  node_coordinates::Array{RealT, 3}      # [orientation, i, elements]
  surface_flux_values::Array{uEltype, 3} # [variables, direction, elements]
  cell_ids::Vector{Int}                  # [elements]
  # internal `resize!`able storage
  _node_coordinates::Vector{RealT}
  _surface_flux_values::Vector{uEltype}
end

nvariables(elements::ElementContainer1D) = size(elements.surface_flux_values, 1)
nnodes(elements::ElementContainer1D) = size(elements.node_coordinates, 2)
Base.eltype(elements::ElementContainer1D) = eltype(elements.surface_flux_values)

# Only one-dimensional `Array`s are `resize!`able in Julia.
# Hence, we use `Vector`s as internal storage and `resize!`
# them whenever needed. Then, we reuse the same memory by
# `unsafe_wrap`ping multi-dimensional `Array`s around the
# internal storage.
function Base.resize!(elements::ElementContainer1D, capacity)
  n_nodes = nnodes(elements)
  n_variables = nvariables(elements)
  @unpack _node_coordinates, _surface_flux_values,
          inverse_jacobian, cell_ids = elements

  resize!(inverse_jacobian, capacity)

  resize!(_node_coordinates, 1 * n_nodes * capacity)
  elements.node_coordinates = unsafe_wrap(Array, pointer(_node_coordinates),
                                          (1, n_nodes, capacity))

  resize!(_surface_flux_values, n_variables * 2 * 1 * capacity)
  elements.surface_flux_values = unsafe_wrap(Array, pointer(_surface_flux_values),
                                             (n_variables, 2 * 1, capacity))

  resize!(cell_ids, capacity)

  return nothing
end

function ElementContainer1D{RealT, uEltype}(capacity::Integer, n_variables, n_nodes) where {RealT<:Real, uEltype<:Real}
  nan_RealT = convert(RealT, NaN)
  nan_uEltype = convert(uEltype, NaN)

  # Initialize fields with defaults
  inverse_jacobian = fill(nan_RealT, capacity)

  _node_coordinates = fill(nan_RealT, 1 * n_nodes * capacity)
  node_coordinates = unsafe_wrap(Array, pointer(_node_coordinates),
                                 (1, n_nodes, capacity))

  _surface_flux_values = fill(nan_uEltype, n_variables * 2 * 1 * capacity)
  surface_flux_values = unsafe_wrap(Array, pointer(_surface_flux_values),
                                    (n_variables, 2 * 1, capacity))

  cell_ids = fill(typemin(Int), capacity)

  return ElementContainer1D{RealT, uEltype}(
    inverse_jacobian, node_coordinates, surface_flux_values, cell_ids,
    _node_coordinates, _surface_flux_values)
end

# Return number of elements
@inline nelements(elements::ElementContainer1D) = length(elements.cell_ids)
# TODO: Taal performance, 1:nelements(elements) vs. Base.OneTo(nelements(elements))
@inline eachelement(elements::ElementContainer1D) = Base.OneTo(nelements(elements))
@inline Base.real(elements::ElementContainer1D) = eltype(elements.node_coordinates)

# Create element container and initialize element data
function init_elements(cell_ids, mesh::TreeMesh1D,
                       equations::AbstractEquations{1},
                       basis, ::Type{RealT}, ::Type{uEltype}) where {RealT<:Real, uEltype<:Real}
  # Initialize container
  n_elements = length(cell_ids)
  elements = ElementContainer1D{RealT, uEltype}(
    n_elements, nvariables(equations), nnodes(basis))

  init_elements!(elements, cell_ids, mesh, basis)
  return elements
end

function init_elements!(elements, cell_ids, mesh::TreeMesh1D, basis)
  nodes = get_nodes(basis)
  # Compute the length of the 1D reference interval by integrating
  # the function with constant value unity on the corresponding
  # element data type (using \circ)
  reference_length = integrate(one ∘ eltype, nodes, basis)
  # Compute the offset of the midpoint of the 1D reference interval
  # (its difference from zero)
  reference_offset = first(nodes) + reference_length / 2

  # Store cell ids
  elements.cell_ids .= cell_ids

  # Calculate inverse Jacobian and node coordinates
  for element in eachelement(elements)
    # Get cell id
    cell_id = cell_ids[element]

    # Get cell length
    dx = length_at_cell(mesh.tree, cell_id)

    # Calculate inverse Jacobian
    jacobian = dx / reference_length
    elements.inverse_jacobian[element] = inv(jacobian)

    # Calculate node coordinates
    # Note that the `tree_coordinates` are the midpoints of the cells.
    # Hence, we need to add an offset for `nodes` with a midpoint
    # different from zero.
      for i in eachnode(basis)
        elements.node_coordinates[1, i, element] = (
            mesh.tree.coordinates[1, cell_id] + jacobian * (nodes[i] - reference_offset))
      end
  end

  return elements
end

# Container data structure (structure-of-arrays style) for DG interfaces
mutable struct InterfaceContainer1D{uEltype<:Real} <: AbstractContainer
  u::Array{uEltype, 3}      # [leftright, variables, interfaces]
  neighbor_ids::Matrix{Int} # [leftright, interfaces]
  orientations::Vector{Int} # [interfaces]
  # internal `resize!`able storage
  _u::Vector{uEltype}
  _neighbor_ids::Vector{Int}
end

nvariables(interfaces::InterfaceContainer1D) = size(interfaces.u, 2)
Base.eltype(interfaces::InterfaceContainer1D) = eltype(interfaces.u)

# See explanation of Base.resize! for the element container
function Base.resize!(interfaces::InterfaceContainer1D, capacity)
  n_variables = nvariables(interfaces)
  @unpack _u, _neighbor_ids, orientations = interfaces

  resize!(_u, 2 * n_variables * capacity)
  interfaces.u = unsafe_wrap(Array, pointer(_u),
                             (2, n_variables, capacity))

  resize!(_neighbor_ids, 2 * capacity)
  interfaces.neighbor_ids = unsafe_wrap(Array, pointer(_neighbor_ids),
                                        (2, capacity))

  resize!(orientations, capacity)

  return nothing
end

function InterfaceContainer1D{uEltype}(capacity::Integer, n_variables, n_nodes) where {uEltype<:Real}
  nan = convert(uEltype, NaN)

  # Initialize fields with defaults
  _u = fill(nan, 2 * n_variables * capacity)
  u = unsafe_wrap(Array, pointer(_u),
                  (2, n_variables, capacity))

  _neighbor_ids = fill(typemin(Int), 2 * capacity)
  neighbor_ids = unsafe_wrap(Array, pointer(_neighbor_ids),
                             (2, capacity))

  orientations = fill(typemin(Int), capacity)

  return InterfaceContainer1D{uEltype}(
    u, neighbor_ids, orientations,
    _u, _neighbor_ids)
end

# Return number of interfaces
@inline ninterfaces(interfaces::InterfaceContainer1D) = length(interfaces.orientations)

# Create interface container and initialize interface data in `elements`.
function init_interfaces(cell_ids, mesh::TreeMesh1D,
                         elements::ElementContainer1D)
  # Initialize container
  n_interfaces = count_required_interfaces(mesh, cell_ids)
  interfaces = InterfaceContainer1D{eltype(elements)}(
    n_interfaces, nvariables(elements), nnodes(elements))

  # Connect elements with interfaces
  init_interfaces!(interfaces, elements, mesh)
  return interfaces
end

# Count the number of interfaces that need to be created
function count_required_interfaces(mesh::TreeMesh1D, cell_ids)
  count = 0

  # Iterate over all cells
  for cell_id in cell_ids
    for direction in eachdirection(mesh.tree)
      # Only count interfaces in positive direction to avoid double counting
      if direction == 1
        continue
      end

      # Skip if no neighbor exists
      if !has_any_neighbor(mesh.tree, cell_id, direction)
        continue
      end

      count += 1
    end
  end

  return count
end

# Initialize connectivity between elements and interfaces
function init_interfaces!(interfaces, elements, mesh::TreeMesh1D)
  # Construct cell -> element mapping for easier algorithm implementation
  tree = mesh.tree
  c2e = zeros(Int, length(tree))
  for element in eachelement(elements)
    c2e[elements.cell_ids[element]] = element
  end

  # Reset interface count
  count = 0

  # Iterate over all elements to find neighbors and to connect via interfaces
  for element in eachelement(elements)
    # Get cell id
    cell_id = elements.cell_ids[element]

    # Loop over directions
    for direction in eachdirection(mesh.tree)
      # Only create interfaces in positive direction
      if direction == 1
        continue
      end

      # Skip if no neighbor exists and current cell is not small
      if !has_any_neighbor(mesh.tree, cell_id, direction)
        continue
      end

      count += 1

      if has_neighbor(mesh.tree, cell_id, direction)
        neighbor_cell_id = mesh.tree.neighbor_ids[direction, cell_id]
        if has_children(mesh.tree, neighbor_cell_id) # Cell has small neighbor
          interfaces.neighbor_ids[2, count] = c2e[mesh.tree.child_ids[1, neighbor_cell_id]]
        else # Cell has same refinement level neighbor
          interfaces.neighbor_ids[2, count] = c2e[neighbor_cell_id]
        end
      else # Cell is small and has large neighbor
        parent_id = mesh.tree.parent_ids[cell_id]
        neighbor_cell_id = mesh.tree.neighbor_ids[direction, parent_id]
        interfaces.neighbor_ids[2, count] = c2e[neighbor_cell_id]
      end

      interfaces.neighbor_ids[1, count] = element
      # Set orientation (x -> 1)
      interfaces.orientations[count] = 1
    end
  end

  @assert count == ninterfaces(interfaces) ("Actual interface count ($count) does not match " *
                                            "expectations $(ninterfaces(interfaces))")
end

# Container data structure (structure-of-arrays style) for DG boundaries
mutable struct BoundaryContainer1D{RealT<:Real, uEltype<:Real} <: AbstractContainer
  u::Array{uEltype, 3}              # [leftright, variables, boundaries]
  neighbor_ids::Vector{Int}         # [boundaries]
  orientations::Vector{Int}         # [boundaries]
  neighbor_sides::Vector{Int}       # [boundaries]
  node_coordinates::Array{RealT, 2} # [orientation, elements]
  n_boundaries_per_direction::SVector{2, Int} # [direction]
  # internal `resize!`able storage
  _u::Vector{uEltype}
  _node_coordinates::Vector{RealT}
end

nvariables(boundaries::BoundaryContainer1D) = size(boundaries.u, 2)
Base.eltype(boundaries::BoundaryContainer1D) = eltype(boundaries.u)

# See explanation of Base.resize! for the element container
function Base.resize!(boundaries::BoundaryContainer1D, capacity)
  n_variables = nvariables(boundaries)
  @unpack _u, _node_coordinates,
          neighbor_ids, orientations, neighbor_sides = boundaries

  resize!(_u, 2 * n_variables * capacity)
  boundaries.u = unsafe_wrap(Array, pointer(_u),
                             (2, n_variables, capacity))

  resize!(_node_coordinates, 1 * capacity)
  boundaries.node_coordinates = unsafe_wrap(Array, pointer(_node_coordinates),
                                            (1, capacity))

  resize!(neighbor_ids, capacity)

  resize!(orientations, capacity)

  resize!(neighbor_sides, capacity)

  return nothing
end

function BoundaryContainer1D{RealT, uEltype}(capacity::Integer, n_variables, n_nodes) where {RealT<:Real, uEltype<:Real}
  nan_RealT = convert(RealT, NaN)
  nan_uEltype = convert(uEltype, NaN)

  # Initialize fields with defaults
  _u = fill(nan_uEltype, 2 * n_variables * capacity)
  u = unsafe_wrap(Array, pointer(_u),
                  (2, n_variables, capacity))

  neighbor_ids = fill(typemin(Int), capacity)

  orientations = fill(typemin(Int), capacity)

  neighbor_sides = fill(typemin(Int), capacity)

  _node_coordinates = fill(nan_RealT, 1 * capacity)
  node_coordinates = unsafe_wrap(Array, pointer(_node_coordinates),
                                 (1, capacity))

  n_boundaries_per_direction = SVector(0, 0)

  return BoundaryContainer1D{RealT, uEltype}(
    u, neighbor_ids, orientations, neighbor_sides,
    node_coordinates, n_boundaries_per_direction,
    _u, _node_coordinates)
end

# Return number of boundaries
nboundaries(boundaries::BoundaryContainer1D) = length(boundaries.orientations)

# Create boundaries container and initialize boundary data in `elements`.
function init_boundaries(cell_ids, mesh::TreeMesh1D,
                         elements::ElementContainer1D)
  # Initialize container
  n_boundaries = count_required_boundaries(mesh, cell_ids)
  boundaries = BoundaryContainer1D{real(elements), eltype(elements)}(
    n_boundaries, nvariables(elements), nnodes(elements))

  # Connect elements with boundaries
  init_boundaries!(boundaries, elements, mesh)
  return boundaries
end

# Count the number of boundaries that need to be created
function count_required_boundaries(mesh::TreeMesh1D, cell_ids)
  count = 0

  # Iterate over all cells
  for cell_id in cell_ids
    for direction in eachdirection(mesh.tree)
      # If neighbor exists, current cell is not at a boundary
      if has_neighbor(mesh.tree, cell_id, direction)
        continue
      end

      # If coarse neighbor exists, current cell is not at a boundary
      if has_coarse_neighbor(mesh.tree, cell_id, direction)
        continue
      end

      # No neighbor exists in this direction -> must be a boundary
      count += 1
    end
  end

  return count
end

# Initialize connectivity between elements and boundaries
function init_boundaries!(boundaries, elements, mesh::TreeMesh1D)
  # Reset boundaries count
  count = 0

  # Initialize boundary counts
  counts_per_direction = MVector(0, 0)

  # OBS! Iterate over directions first, then over elements, and count boundaries in each direction
  # Rationale: This way the boundaries are internally sorted by the directions -x, +x, -y etc.,
  #            obviating the need to store the boundary condition to be applied explicitly.
  # Loop over directions
  for direction in eachdirection(mesh.tree)
    # Iterate over all elements to find missing neighbors and to connect to boundaries
    for element in eachelement(elements)
      # Get cell id
      cell_id = elements.cell_ids[element]

      # If neighbor exists, current cell is not at a boundary
      if has_neighbor(mesh.tree, cell_id, direction)
        continue
      end

      # If coarse neighbor exists, current cell is not at a boundary
      if has_coarse_neighbor(mesh.tree, cell_id, direction)
        continue
      end

      # Create boundary
      count += 1
      counts_per_direction[direction] += 1

      # Set neighbor element id
      boundaries.neighbor_ids[count] = element

      # Set neighbor side, which denotes the direction (1 -> negative, 2 -> positive) of the element
      if direction == 2
        boundaries.neighbor_sides[count] = 1
      else
        boundaries.neighbor_sides[count] = 2
      end

      # Set orientation (x -> 1)
      boundaries.orientations[count] = 1

      # Store node coordinates
      enc = elements.node_coordinates
      if direction == 1 # -x direction
        boundaries.node_coordinates[:, count] .= enc[:, 1,  element]
      elseif direction == 2 # +x direction
        boundaries.node_coordinates[:, count] .= enc[:, end, element]
      else
        error("should not happen")
      end
    end
  end

  @assert count == nboundaries(boundaries) ("Actual boundaries count ($count) does not match " *
                                            "expectations $(nboundaries(boundaries))")
  @assert sum(counts_per_direction) == count

  boundaries.n_boundaries_per_direction = SVector(counts_per_direction)

  return boundaries.n_boundaries_per_direction
end

end # @muladd


text:
module Metal

using Reexport
using GPUArrays
using Adapt
using GPUCompiler
using LLVM
using LLVM.Interop

# core library
include("../lib/core/MTL.jl")
@reexport using .MTL

# device functionality
include("device/pointer_abstract.jl")
include("device/pointer_ptr.jl")
include("device/pointer_buf.jl")
include("device/array.jl")
include("device/metal.jl")

# compiler
include("execution/mtl_type_conversion.jl")
include("execution/device_type_conversion.jl")
include("execution/kernel.jl")

include("state.jl")

# array abstraction
include("host/array.jl")
include("host/memory.jl")

end # module


text:
using PolyaGammaSamplers
using Test

@testset "PolyaGammaSamplers.jl" begin
    include("polyagammapswsampler.jl")
end


text:
######################################################################
# sensors.jl
######################################################################

abstract type Sensor end

# Returns a possible observation to be received by sensor s from pose p
#  assuming the target is located according to m
function observe(m::SearchDomain, s::Sensor, p::Pose)
    observe(m.theta, s, p)
end

function observe(tx::TargetTuple, s::Sensor, p::Pose)
    error("`observe` not defined for this sensor.")
end

# Returns probability of measuring o with sensor s if vehicle pose is p
#  and target location is theta
function O(s::Sensor, theta::LocTuple, p::Pose, o)
    error("`O` (density version) not defined for this sensor.")
end

include("bearing.jl")
include("binned_bearing.jl")
include("range.jl")
include("diromni.jl")
include("binned_diromni.jl")
include("fov.jl")
#include("fov2.jl")
include("fovb.jl")
include("fovn.jl")
include("fov3.jl")


text:
#GracePlot: Publication-quality plots through Grace/xmgrace
#-------------------------------------------------------------------------------
#__precompile__()

module GracePlot

const rootpath = realpath(joinpath(dirname(realpath(@__FILE__)),"../."))

#Convenient accessor for sample GracePlot template (parameter) files:
template(name::AbstractString) =
	joinpath(GracePlot.rootpath, "sample", "template", "$name.par")

#==Ensure interface (similar to assert)
===============================================================================#
#=Similar to assert.  However, unlike assert, "ensure" is not meant for
debugging.  Thus, ensure is never meant to be compiled out.
=#
function _ensure(cond::Bool, err)
	if !cond; throw(err); end
end

include("graceconst.jl")
include("codegen.jl")
include("units.jl")
include("base.jl")
include("plotmanip.jl")
include("io.jl")

#==Exported symbols
===============================================================================#
export redraw #Whole plot
export autofit #Re-compute axes to fit data
export arrange #Re-tile plot with specified number of rows/cols
export clearall #clearall(Plot): Graphs or clearall(Graph): datasets

export graph #Obtain reference to an individual graph
export add #Add new dataset/graph

export set #Set Plot/Graph/Dataset properties
#   set(::Plot, arg1, arg2, ..., kwarg1=v1, kwarg2=v1, ...)
#      kwargs: active, focus,
#   set(::GraphRef, arg1, arg2, ..., kwarg1=v1, kwarg2=v1, ...)
#      args: axes()
#      kwargs: title, subtitle, xlabel, ylabel, frameline
#   set(::DatasetRef, arg1, arg2, ..., kwarg1=v1, kwarg2=v1, ...)
#      args: line(), glyph()
export defaults #Creates DefaultAttributes to set default plot attributes
export canvas #Creates CanvasAttributes to resize plot, etc
export limits #Creates CartesianLimAttributes to set view, world, ...
export text #Creates TextAttributes to set titles, etc
export axes #Creates AxesAttributes to modify axis
export line #Creates LineAttributes to modify line
export glyph #Creates GlyphAttributes to modify glyph
export legend #Creates LegendAttributes to edit legend
export addannotation

#==
Other interface tools (symbols not exported to avoid collisions):
	Inch, Meter, TPoint... TODO: move to external module
	Plot: Main plot object.
	new(): Creates a new Plot object.
	kill(graph): Kill already in Base.
	template("<GracePlot-provided template name>")
	Base.get.... get is already part of base, so can't export it...
==#

end #GracePlot

#Last line


text:
#!/usr/bin/env julia

#Start Test Script
using VREP
using Base.Test

# Run tests

tic()
println("Test 1")
@time @test include("test1.jl")
println("Test 2")
@time @test include("test2.jl")
toc()


text:
using Mamba, Stan

currdir = (@__FILE__() == nothing) ? "." : dirname(@__FILE__())
const BayesianSlopeModel = open(fh -> readall(fh), joinpath(currdir, "BayesianSlope.stan"))

# Create fake data to test
NumCoeffs = 100
nonzero_coeffs = [2, 3, -4, 7]
nx = length(nonzero_coeffs)
NumSpurious = NumCoeffs - nx
coeffs = zeros(NumCoeffs)
idxs_nonzero = shuffle(collect(1:NumCoeffs))[1:nx]
coeffs[idxs_nonzero] = nonzero_coeffs
p = length(coeffs)
s = 0.20
n = 60
X = randn(n, p)
y = X * coeffs .+ s * randn(n)

# Set lambdas according to Benjamini-Hochberg
using Distributions
const NormalDistr = Normal(0, 1)
function lambdas_bh(p, q)
    lambdas = zeros(Float64, p)
    for i in 1:p
        lambdas[p+1-i] = quantile(NormalDistr, 1-0.5*i*q/p)
    end
    lambdas
end

# Prep and then call stan
data = Dict("p" => p, "n" => n, "y" => y, "X" => X,
    "lambda" => lambdas_bh(p, 0.01), "a" => 1.0, "g" => 1.0);
bayes_slope_data = [data];

num_betas = length(coeffs);
beta_names = map(i -> "beta.$i", 1:num_betas);
convergence_stats = ["lp__", "accept_stat__"];

stanmodel = Stanmodel(name="bayesian_slope", model=BayesianSlopeModel,
    monitors = vcat(convergence_stats, beta_names), update=2000)
@time sim1 = stan(stanmodel, bayes_slope_data, currdir, CmdStanDir=CMDSTAN_HOME);

convstats_chain1 = sim1[1:size(sim1, 1), convergence_stats, 1];

betas_chain1 = sim1[1:size(sim1, 1), beta_names, 1].value;
est_betas_chain1 = median(betas_chain1, 1)[:];
#std_betas_chain1 = std(betas_chain1, 1)[:];
#includes_zero(val, stdest) = (val - stdest) < 0.0 < (val + stdest)
#nonzero_idxs = filter(i -> !includes_zero(est_betas_chain1[i], std_betas_chain1[i]), 1:length(est_betas_chain1));
includes_zeroes_hdr(vals, level=0.05) = (qs = quantile(vals, [level, 1.0-level]); qs[1] < 0.0 < qs[2])
nonzero_idxs = filter(i -> !includes_zeroes_hdr(betas_chain1[:,i,:][:]), 1:length(est_betas_chain1));
println("Non-zero coeffs:")
map(i -> println("  $i: ", join(map(v -> round(v,2), quantile(betas_chain1[:,i,:][:], [0.05, 0.50, 0.95])), " - ")), sort(nonzero_idxs));
println("Actual coeffs:")
map(i -> println("  $i: ", coeffs[i]), sort(idxs_nonzero));

# Now compare to Optim-based SLOPE regression
include("optim_based_slope.jl")
#betas_optim = optim_based_slope_regression(y, X; absmax = 1e2)
betas_optim = optim_based_slope_regression(y, X);
largest_abs = sortperm(abs(betas_optim), rev=true);
println("Largest 10: ", largest_abs[1:10])
println("Coeffs:     ", map(i -> round(betas_optim[i], 3), largest_abs[1:10]))

# Without extensive testing it seems that Bayesian SLOPE regression is more powerful:
#  - if the number of cases (n) is low and, in particular, if n<p
#  - in that it can give confidence intervals and not only point estimates
# However, it does take a lot longer time so if n > p it might not be the way to go...
# Optim-based SLOPE does not seem to work for n/p under 1.2 or so...