text:
# Practical Machine learning
# Decision Trees based learning - Random Forests example
# Chapter 5

type Regressor
    n_samples::Int
    n_features::Int
    n_max_features::Int
    improvements::Vector{Float64}
    oob_error::Float64
    trees::Vector{Tree}

    function Regressor(rf, x, y)
        n_samples, n_features = size(x)

        if n_samples != length(y)
            throw(DimensionMismatch(""))
        end

        n_max_features = resolve_max_features(rf.max_features, n_features)
        @assert 0 < n_max_features <= n_features

        improvements = zeros(Float64, n_features)
        trees = Array(Tree, rf.n_estimators)
        new(n_samples, n_features, n_max_features, improvements, NaN, trees)
    end
end

typealias RandomForestRegressor RandomForest{Regressor}

function RandomForestRegressor(;n_estimators::Int=10, max_features::Union{Integer,AbstractFloat,Symbol}=:third, max_depth=nothing, min_samples_split::Int=2)
    RandomForest{Regressor}(n_estimators, max_features, max_depth, min_samples_split, :mse)
end

function fit{T<:TabularData}(rf::RandomForestRegressor, x::T, y::AbstractVector)
    learner = Regressor(rf, x, y)
    n_samples = learner.n_samples

    # pre-allocation
    bootstrap = Array(Int, n_samples)
    sample_weight = Array(Float64, n_samples)
    oob_predict = zeros(n_samples)
    oob_count = zeros(Int, n_samples)

    for b in 1:rf.n_estimators
        rand!(bootstrap, 1:n_samples)
        set_weight!(bootstrap, sample_weight)
        example = Trees.Example{T}(x, y, sample_weight)
        tree = Trees.Tree()
        Trees.fit(tree, example, rf.criterion, learner.n_max_features, rf.max_depth, rf.min_samples_split)
        learner.trees[b] = tree

        for s in 1:n_samples
            if sample_weight[s] == 0.0
                oob_predict[s] += Trees.predict(tree, vec(x[s, :]))
                oob_count[s] += 1
            end
        end
    end

    oob_error = 0.

    for s in 1:n_samples
        if oob_count[s] == 0
            continue
        end

        avg = oob_predict[s] / oob_count[s]
        d = y[s] - avg
        oob_error += d * d
    end

    set_improvements!(learner)
    learner.oob_error = sqrt(oob_error / countnz(oob_count))
    rf.learner = learner
    return
end

function predict{T<:TabularData}(rf::RandomForestRegressor, x::T)
    if is(rf.learner, nothing)
        error("not yet trained")
    end

    n_samples = size(x, 1)
    output = Array(Float64, n_samples)
    vs = Array(Float64, rf.n_estimators)

    for i in 1:n_samples
        for b in 1:rf.n_estimators
            tree = rf.learner.trees[b]
            vs[b] = Trees.predict(tree, vec(x[i, :]))
        end
        output[i] = mean(vs)
    end

    output
end

function oob_error(rf::RandomForestRegressor)
    if is(rf.learner, nothing)
        error("not yet trained")
    end

    rf.learner.oob_error
end


text:
@testset "mask" begin
    using LinearAlgebra
    using NeuralAttentionlib.Masks
    using NeuralAttentionlib: getmask

    causal(x) = batched_triu!(copy(x), 0)
    trilu(x, d) = batched_tril!(batched_triu!(copy(x), -d), d)
    bandpart(x, l, u) = batched_tril!(batched_triu!(copy(x), -l), u)

    test_random(x, p) = isapprox(sum(x .* RandomMask(p)) / length(x), 1-p, atol=1e-1)

    @testset "dataless mask" begin
        a = ones(Int, 10, 10)
        b = ones(Int, 10, 10, 2)
        c = ones(Int, 10, 10, 2, 2)

        @test a .* CausalMask() == causal(a)
        @test b .* CausalMask() == causal(b)
        @test c .* CausalMask() == causal(c)

        @test a .* LocalMask(5) == trilu(a, 4)
        @test b .* LocalMask(5) == trilu(b, 4)
        @test c .* LocalMask(5) == trilu(c, 4)

        @test test_random(a, 0.1)
        @test test_random(b, 0.1)
        @test test_random(c, 0.1)
        @test test_random(a, 0.9)
        @test test_random(b, 0.9)
        @test test_random(c, 0.9)

        @test a .* BandPartMask(3, 5) == bandpart(a, 3, 5)
        @test b .* BandPartMask(3, 5) == bandpart(b, 3, 5)
        @test c .* BandPartMask(3, 5) == bandpart(c, 3, 5)

    end

    function grow_length(len1, len2, n)
        @assert size(len1) == size(len2)
        y = zeros(Int, n, n, size(len1)...)
        for i = 1:length(len1)
            idx = Tuple(CartesianIndices(len1)[i])
            l1 = len1[idx...]
            l2 = len2[idx...]
            y[1:l1, 1:l2, idx...] .= 1
        end
        return y
    end

    @testset "array mask" begin
        a = ones(Int, 10, 10)
        b = ones(Int, 10, 10, 2)
        c = ones(Int, 10, 10, 2, 2)

        gmask_a = rand(Bool, 10, 10)
        gmask_b = rand(Bool, 10, 10, 2)
        gmask_c = rand(Bool, 10, 10, 2, 2)

        @test a .* GenericMask(gmask_a) == convert(AbstractArray{Int}, gmask_a)
        @test b .* GenericMask(gmask_b) == convert(AbstractArray{Int}, gmask_b)
        @test c .* GenericMask(gmask_c) == convert(AbstractArray{Int}, gmask_c)

        smask_a = [5]
        smask_b = [3, 5]
        smask_c = rand(1:10, 2, 2)

        @test a .* SymLengthMask(smask_a) == dropdims(grow_length(smask_a, smask_a, 10), dims=3)
        @test b .* SymLengthMask(smask_b) == grow_length(smask_b, smask_b, 10)
        @test c .* SymLengthMask(smask_c) == grow_length(smask_c, smask_c, 10)

        bmaskq_a, bmaskk_a = [4], [5]
        bmaskq_b, bmaskk_b = [3, 5], [4, 6]
        bmaskq_c, bmaskk_c = rand(1:10, 2, 2), rand(1:10, 2, 2)

        @test a .* BiLengthMask(bmaskq_a, bmaskk_a) == dropdims(grow_length(bmaskk_a, bmaskq_a, 10), dims=3)
        @test b .* BiLengthMask(bmaskq_b, bmaskk_b) == grow_length(bmaskk_b, bmaskq_b, 10)
        @test c .* BiLengthMask(bmaskq_c, bmaskk_c) == grow_length(bmaskk_c, bmaskq_c, 10)

    end

    @testset "wrapper mask" begin
        a = ones(Int, 10, 10)
        b = ones(Int, 10, 10, 2)
        c = ones(Int, 10, 10, 2, 2)

        gmask_a = rand(Bool, 10, 10)
        gmask_b = rand(Bool, 10, 10, 2)
        gmask_c = rand(Bool, 10, 10, 2, 2)
        smask_a = [5]
        smask_b = [3, 5]
        smask_c = rand(1:10, 2, 2)
        bmaskq_a, bmaskk_a = [4], [5]
        bmaskq_b, bmaskk_b = [3, 5], [4, 6]
        bmaskq_c, bmaskk_c = rand(1:10, 2, 2), rand(1:10, 2, 2)

        @test a .* !CausalMask() == 1 .- causal(a)
        @test b .* !CausalMask() == 1 .- causal(b)
        @test c .* !CausalMask() == 1 .- causal(c)
        @test a .* !LocalMask(5) == 1 .- trilu(a, 4)
        @test b .* !LocalMask(5) == 1 .- trilu(b, 4)
        @test c .* !LocalMask(5) == 1 .- trilu(c, 4)
        @test a .* !BandPartMask(3, 5) == 1 .- bandpart(a, 3, 5)
        @test b .* !BandPartMask(3, 5) == 1 .- bandpart(b, 3, 5)
        @test c .* !BandPartMask(3, 5) == 1 .- bandpart(c, 3, 5)
        @test a .* !GenericMask(gmask_a) == 1 .- convert(AbstractArray{Int}, gmask_a)
        @test b .* !GenericMask(gmask_b) == 1 .- convert(AbstractArray{Int}, gmask_b)
        @test c .* !GenericMask(gmask_c) == 1 .- convert(AbstractArray{Int}, gmask_c)
        @test a .* !SymLengthMask(smask_a) == 1 .- dropdims(grow_length(smask_a, smask_a, 10), dims=3)
        @test b .* !SymLengthMask(smask_b) == 1 .- grow_length(smask_b, smask_b, 10)
        @test c .* !SymLengthMask(smask_c) == 1 .- grow_length(smask_c, smask_c, 10)
        @test a .* !BiLengthMask(bmaskq_a, bmaskk_a) == 1 .- dropdims(grow_length(bmaskk_a, bmaskq_a, 10), dims=3)
        @test b .* !BiLengthMask(bmaskq_b, bmaskk_b) == 1 .- grow_length(bmaskk_b, bmaskq_b, 10)
        @test c .* !BiLengthMask(bmaskq_c, bmaskk_c) == 1 .- grow_length(bmaskk_c, bmaskq_c, 10)

        @test a .* (!BandPartMask(4, 7) | BiLengthMask(bmaskq_a, bmaskk_a) | LocalMask(2)) == min.((1 .- bandpart(a, 4, 7) .+ dropdims(grow_length(bmaskk_a, bmaskq_a, 10), dims=3) .+ trilu(a, 1)), 1)
        @test b .* (!BandPartMask(4, 7) | BiLengthMask(bmaskq_b, bmaskk_b) | LocalMask(2)) == min.((1 .- bandpart(b, 4, 7) .+ grow_length(bmaskk_b, bmaskq_b, 10) .+ trilu(b, 1)), 1)
        @test c .* (!BandPartMask(4, 7) | BiLengthMask(bmaskq_c, bmaskk_c) | LocalMask(2)) == min.((1 .- bandpart(c, 4, 7) .+ grow_length(bmaskk_c, bmaskq_c, 10) .+ trilu(c, 1)), 1)

        @test a .* (!CausalMask() & BiLengthMask(bmaskq_a, bmaskk_a) & LocalMask(2)) == (1 .- causal(a)) .* dropdims(grow_length(bmaskk_a, bmaskq_a, 10), dims=3) .* trilu(a, 1)
        @test b .* (!CausalMask() & BiLengthMask(bmaskq_b, bmaskk_b) & LocalMask(2)) == (1 .- causal(b)) .* grow_length(bmaskk_b, bmaskq_b, 10) .* trilu(b, 1)
        @test c .* (!CausalMask() & BiLengthMask(bmaskq_c, bmaskk_c) & LocalMask(2)) == (1 .- causal(c)) .* grow_length(bmaskk_c, bmaskq_c, 10) .* trilu(c, 1)

        @test a .* (!CausalMask() | BiLengthMask(bmaskq_a, bmaskk_a) & LocalMask(2) | BandPartMask(4, 7)) == min.((1 .- causal(a)) .+ (dropdims(grow_length(bmaskk_a, bmaskq_a, 10), dims=3) .* trilu(a, 1)) .+ bandpart(a, 4, 7), 1)
        @test b .* (!CausalMask() | BiLengthMask(bmaskq_b, bmaskk_b) & LocalMask(2) | BandPartMask(4, 7)) == min.((1 .- causal(b)) .+ (grow_length(bmaskk_b, bmaskq_b, 10) .* trilu(b, 1)) .+ bandpart(a, 4, 7), 1)
        @test c .* (!CausalMask() | BiLengthMask(bmaskq_c, bmaskk_c) & LocalMask(2) | BandPartMask(4, 7)) == min.((1 .- causal(c)) .+ (grow_length(bmaskk_c, bmaskq_c, 10) .* trilu(c, 1)) .+ bandpart(a, 4, 7), 1)

        @test a .* (!(CausalMask() | BiLengthMask(bmaskq_a, bmaskk_a)) & LocalMask(2) | BandPartMask(4, 7)) ==
            min.((1 .- min.(causal(a) .+ (dropdims(grow_length(bmaskk_a, bmaskq_a, 10), dims=3)), 1)) .* trilu(a, 1) .+ bandpart(a, 4, 7), 1)
        @test b .* (!(CausalMask() | BiLengthMask(bmaskq_b, bmaskk_b)) & LocalMask(2) | BandPartMask(4, 7)) ==
            min.((1 .- min.(causal(b) .+ grow_length(bmaskk_b, bmaskq_b, 10), 1)) .* trilu(b, 1) .+ bandpart(a, 4, 7), 1)
        @test c .* (!(CausalMask() | BiLengthMask(bmaskq_c, bmaskk_c)) & LocalMask(2) | BandPartMask(4, 7)) ==
            min.((1 .- min.(causal(c) .+ grow_length(bmaskk_c, bmaskq_c, 10), 1)) .* trilu(c, 1) .+ bandpart(a, 4, 7), 1)

        @test c .* BatchedMask(BiLengthMask(bmaskq_b, bmaskk_b)) == begin
            m = reshape(grow_length(bmaskk_b, bmaskq_b, 10), 10, 10, 1, 2)
            cat(m, m; dims=3)
        end

        @test c .* BatchedMask(SymLengthMask(smask_b)) == begin
            m = reshape(grow_length(smask_b, smask_b, 10), 10, 10, 1, 2)
            cat(m, m; dims=3)
        end

        d = ones(Int, 10, 10, 2, 2, 2)

        @test d .* BatchedMask(BiLengthMask(bmaskq_c, bmaskk_c)) == begin
            m = reshape(grow_length(bmaskk_c, bmaskq_c, 10), 10, 10, 1, 2, 2)
            cat(m, m; dims=3)
        end
        @test d .* BatchedMask(BiLengthMask(bmaskq_b, bmaskk_b)) == begin
            m = reshape(grow_length(bmaskk_b, bmaskq_b, 10), 10, 10, 1, 1, 2)
            tmp = cat(m, m; dims=3)
            cat(tmp, tmp; dims=4)
        end

        e = ones(Int, 10, 10, 4)
        f = ones(Int, 10, 10, 2, 6)
        @test e .* RepeatMask(BiLengthMask(bmaskq_b, bmaskk_b), 2) == repeat(grow_length(bmaskk_b, bmaskq_b, 10), inner=(1,1,2))
        @test f .* RepeatMask(BiLengthMask(bmaskq_c, bmaskk_c), 3) == repeat(grow_length(bmaskk_c, bmaskq_c, 10), inner=(1,1,1,3))

    end

    @testset "Op" begin
        c = ones(10, 10, 2, 2)

        @test apply_mask(CausalMask(), c) == c .* CausalMask()
        @test apply_mask(NaiveAttenMaskOp(), CausalMask(), c) == c .* CausalMask()
        @test apply_mask(GenericAttenMaskOp(.*, true, 2), CausalMask(), c) == 2 .* c .* !CausalMask()
        @test apply_mask(GenericAttenMaskOp(.*, false, 2), CausalMask(), c) == 2 .* c .* CausalMask()
        @test apply_mask(GenericAttenMaskOp(.+, false, 2), CausalMask(), c) == c .+ 2 .* CausalMask()
        @test apply_mask(GenericAttenMaskOp(+, false, 2), CausalMask(), c) == c .+ 2 .* CausalMask()
        @test apply_mask(GenericAttenMaskOp(-, false, 2), CausalMask(), c) == c .- 2 .* CausalMask()
        @test apply_mask(GenericAttenMaskOp(-, true, 2), CausalMask(), c) == c .- 2 .* !CausalMask()
        @test apply_mask(GenericAttenMaskOp(./, true, 2), CausalMask(), c) == inv(2) .* c .* !CausalMask()
        @test apply_mask(GenericAttenMaskOp(./, false, 2), CausalMask(), c) == inv(2) .* c .* CausalMask()

    end

    @testset "Broadcast" begin
        @test trues(5, 5, 1, 2, 1) .* BatchedMask(SymLengthMask([2])) .* CausalMask() ==
            trues(5, 5, 1, 2, 1) .* (BatchedMask(SymLengthMask([2])) & CausalMask())
        @test trues(5, 5, 1, 2, 1) .* BatchedMask(SymLengthMask([2])) .| CausalMask() ==
            trues(5, 5, 1, 2, 1) .* (BatchedMask(SymLengthMask([2])) | CausalMask())
        @test trues(5, 5, 1, 2, 1) .* (BatchedMask(SymLengthMask([2])) .* CausalMask()) ==
            trues(5, 5, 1, 2, 1) .* (BatchedMask(SymLengthMask([2])) & CausalMask())
        @test trues(5,5,1, 2,1) .* (BatchedMask(SymLengthMask([2])) .| CausalMask() .* LocalMask(1)) ==
            trues(5,5,1, 2,1) .* (BatchedMask(SymLengthMask([2])) | (CausalMask() & LocalMask(1)))
        @test trues(5,5,1, 2,1) .* .!(BatchedMask(SymLengthMask([2])) .| CausalMask() .* LocalMask(1)) ==
            trues(5,5,1, 2,1) .* !(BatchedMask(SymLengthMask([2])) | (CausalMask() & LocalMask(1)))

        @test_throws DimensionMismatch randn(5) .* CausalMask()
        @test_throws DimensionMismatch randn(5, 5, 2) .* SymLengthMask([2])
        @test_throws DimensionMismatch randn(5, 5, 2, 1) .* BiLengthMask([2,3], [2,2])
        @test_throws DimensionMismatch randn(5, 5, 3) .* BiLengthMask([2,3], [2,2])
        @test_throws DimensionMismatch randn(5, 5, 2) .* RepeatMask(BiLengthMask([2,3], [2,2]), 3)
        @test_throws DimensionMismatch randn(5, 5, 3) .* RepeatMask(SymLengthMask([2]), 2)
        @test_throws DimensionMismatch randn(5, 5, 2, 3) .* BatchedMask(BiLengthMask([2,3], [2,2]))
        @test_throws DimensionMismatch randn(5, 4) .* GenericMask(rand(Bool, 3, 4))
        @test_throws DimensionMismatch randn(5, 4, 2) .* (GenericMask(rand(Bool, 3, 4)) | BiLengthMask([2,3], [2,2]))
        @test_throws DimensionMismatch randn(5, 4) .* (GenericMask(rand(Bool, 3, 4)) | SymLengthMask([2]))
    end

    @testset "AD" begin
        m = (LocalMask(1) | CausalMask() & !(BandPartMask(5,5)) | BiLengthMask([2,3], [3, 7]))

        test_rrule(getindex, m ⊢ NoTangent(), 1, 4, 1)
        test_rrule(getindex, NeuralAttentionlib.GetIndexer(m) ⊢ NoTangent(), 5, 4, 2)
        test_rrule(getindex, m ⊢ NoTangent(), (1, 4, 1))
        test_rrule(getindex, NeuralAttentionlib.GetIndexer(m) ⊢ NoTangent(), (5, 4, 2))
        test_rrule(LocalMask, 2)

        test_rrule(getmask, m ⊢ NoTangent(), randn(10, 10, 2), 0.5 ⊢ NoTangent())
        test_rrule(NeuralAttentionlib.apply_mask, NaiveAttenMaskOp(), m ⊢ NoTangent(), randn(10, 10, 2))

        test_rrule(NeuralAttentionlib.apply_broadcast_mask, (*) ⊢ NoTangent(), m ⊢ NoTangent(), randn(10, 10, 2), 3 ⊢ NoTangent())
        test_rrule(NeuralAttentionlib.apply_broadcast_mask, (+) ⊢ NoTangent(), m ⊢ NoTangent(), randn(10, 10, 2), -1e9 ⊢ NoTangent(); atol=1e-2)

    end

end

ext: jl


text:
#=

lucas_stokey.jl

@author: Shunsuke Hori

=#

module LS

using QuantEcon
using Dierckx

using NLsolve
using NLopt

type Para{TAF<:AbstractFloat}
    beta::TAF
    Pi::Array{TAF,2}
    G::Vector{TAF}
    Theta::Vector{TAF}
    transfers::Bool
    U::Function
    Uc::Function
    Ucc::Function
    Un::Function
    Unn::Function
    n_less_than_one::Bool
end

"""
Class returns planner's allocation as a function of the multiplier
on the implementability constraint mu
"""
type Planners_Allocation_Sequential{TI<:Integer, TAF<:AbstractFloat}
    para::Para{TAF}
    mc::MarkovChain
    S::TI
    cFB::Vector{TAF}
    nFB::Vector{TAF}
    XiFB::Vector{TAF}
    zFB::Vector{TAF}
end

"""
Initializes the class from the calibration Para
"""
function Planners_Allocation_Sequential(para::Para)
    beta, Pi, G, Theta =
        para.beta, para.Pi, para.G, para.Theta
    mc = MarkovChain(Pi)
    S = size(Pi,1) # number of states
    #now find the first best allocation
    cFB, nFB, XiFB, zFB = find_first_best(para,S,1)

    return Planners_Allocation_Sequential(
        para, mc,S, cFB, nFB, XiFB, zFB)
end

"""
Find the first best allocation
"""
function find_first_best(para::Para,S::Integer,
            version::Integer)
    if version != 1 && version != 2
        throw(ArgumentError("version must be 1 or 2"))
    end
    beta, Theta, Uc, Un, G, Pi =
        para.beta, para.Theta, para.Uc, para.Un, para.G, para.Pi
    function res!(z, out)
        c = z[1:S]
        n = z[S+1:end]
        out[1:S] = Theta.*Uc(c,n)+Un(c,n)
        out[S+1:end] = Theta.*n - c - G
    end
    res = nlsolve(res!, 0.5*ones(2*S))

    if converged(res) == false
        error("Could not find first best")
    end

    if version == 1
        cFB = res.zero[1:S]
        nFB = res.zero[S+1:end]
        XiFB = Uc(cFB,nFB) #multiplier on the resource constraint.
        zFB = vcat(cFB,nFB,XiFB)
        return cFB, nFB, XiFB, zFB
    elseif version == 2
        cFB = res.zero[1:S]
        nFB = res.zero[S+1:end]
        IFB = Uc(cFB, nFB).*cFB + Un(cFB, nFB).*nFB
        xFB = \(eye(S) - beta*Pi, IFB)
        zFB = Array{Array}(S)
        for s in 1:S
            zFB[s] = vcat(cFB[s], nFB[s], xFB)
        end
        return cFB, nFB, IFB, xFB, zFB
    end
end

"""
Computes optimal allocation for time t\geq 1 for a given \mu
"""
function time1_allocation(pas::Planners_Allocation_Sequential, mu::Real)
    para, S = pas.para, pas.S
    Theta, beta, Pi, G, Uc, Ucc, Un, Unn =
        para.Theta, para.beta, para.Pi, para.G,
        para.Uc, para.Ucc, para.Un, para.Unn
    function FOC!(z::Vector, out)
        c = z[1:S]
        n = z[S+1:2S]
        Xi = z[2S+1:end]
        out[1:S] = Uc(c,n) - mu*(Ucc(c,n).*c+Uc(c,n)) -Xi #foc c
        out[S+1:2S] = Un(c,n) - mu*(Unn(c,n).*n+Un(c,n)) + Theta.*Xi #foc n
        out[2S+1:end] = Theta.*n - c - G #resource constraint
        return out
    end
    #find the root of the FOC
    res = nlsolve(FOC!, pas.zFB)
    if res.f_converged == false
        error("Could not find LS allocation.")
    end
    z = res.zero
    c, n, Xi = z[1:S], z[S+1:2S], z[2S+1:end]
    #now compute x
    I  = Uc(c,n).*c +  Un(c,n).*n
    x = \(eye(S) - beta*para.Pi, I)
    return c, n, x, Xi
end

"""
Finds the optimal allocation given initial government debt B_ and state s_0
"""
function time0_allocation(pas::Planners_Allocation_Sequential,
                            B_::AbstractFloat, s_0::Integer)
    para = pas.para
    Pi, Theta, G, beta =
        para.Pi, para.Theta, para.G, para.beta
    Uc, Ucc, Un, Unn =
        para.Uc, para.Ucc, para.Un, para.Unn
    #first order conditions of planner's problem
    function FOC!(z, out)
        mu, c, n, Xi = z[1], z[2], z[3], z[4]
        xprime = time1_allocation(pas, mu)[3]
        out .= vcat(
        Uc(c,n).*(c-B_) + Un(c,n).*n + beta*dot(Pi[s_0,:],xprime),
        Uc(c,n) - mu*(Ucc(c,n).*(c-B_) + Uc(c,n)) - Xi,
        Un(c,n) - mu*(Unn(c,n).*n+Un(c,n)) + Theta[s_0].*Xi,
        (Theta.*n - c - G)[s_0]
        )
    end
    #find root
    res = nlsolve(FOC!,
        [0.0, pas.cFB[s_0], pas.nFB[s_0], pas.XiFB[s_0]])
    if res.f_converged == false
        error("Could not find time 0 LS allocation.")
    end
    return (res.zero...)
end

"""
Find the value associated with multiplier mu
"""
function time1_value(pas::Planners_Allocation_Sequential,mu::Real)
    para= pas.para
    c, n, x, Xi = time1_allocation(pas, mu)
    U_val = para.U(c,n)
    V = \(eye(pas.S) - para.beta*para.Pi, U_val)
    return c, n, x, V
end

"""
Computes Tau given c,n
"""
function Tau(para::Para,c::Union{Real,Vector},n::Union{Real,Vector})
    Uc, Un = para.Uc.(c,n), para.Un.(c,n)
    return 1+Un./(para.Theta .* Uc)
end

"""
Simulates planners policies for T periods
"""
function simulate(pas::Planners_Allocation_Sequential,
                    B_::AbstractFloat, s_0::Integer, T::Integer, sHist=nothing)
    para = pas.para
    Pi, beta, Uc =
        para.Pi, para.beta, para.Uc
    if sHist == nothing
        sHist = QuantEcon.simulate(pas.mc, T, init=s_0)
    end
    cHist = zeros(T)
    nHist = zeros(T)
    Bhist = zeros(T)
    TauHist = zeros(T)
    muHist = zeros(T)
    RHist = zeros(T-1)
    #time0
    mu, cHist[1], nHist[1], _  =
        time0_allocation(pas,B_,s_0)
    TauHist[1] = Tau(pas.para, cHist[1],nHist[1])[s_0]
    Bhist[1] = B_
    muHist[1] = mu
    #time 1 onward
    for t in 2:T
        c, n, x, Xi = time1_allocation(pas,mu)
        u_c = Uc(c,n)
        s = sHist[t]
        TauHist[t] = Tau(pas.para,c,n)[s]
        Eu_c = dot(Pi[sHist[t-1],:],u_c)
        cHist[t], nHist[t], Bhist[t] = c[s], n[s], x[s]/u_c[s]
        RHist[t-1] = Uc(cHist[t-1],nHist[t-1])/(beta*Eu_c)
        muHist[t] = mu
    end
    return cHist, nHist, Bhist, TauHist, sHist, muHist, RHist
end

"""
Bellman equation for the continuation of the Lucas-Stokey Problem
"""
type BellmanEquation{TI<:Integer,TF<:AbstractFloat}
    para::Para
    S::TI
    xbar::Array{TF}
    time_0::Bool
    z0::Array{Array}
    cFB::Vector{TF}
    nFB::Vector{TF}
    xFB::Vector{TF}
    zFB::Vector{Array}
end

"""
Initializes the class from the calibration Para
"""
function BellmanEquation(para::Para,xgrid::AbstractVector,
                        policies0)
    S = size(para.Pi,1) # number of states
    xbar = [minimum(xgrid),maximum(xgrid)]
    time_0 = false
    z0 = Array{Array}(length(xgrid),S)
    cf, nf, xprimef = policies0
    for s in 1:S
        for (i_x, x) in enumerate(xgrid)
            xprime0 = Array{typeof(para.beta)}(S)
            for sprime in 1:S
                xprime0[sprime] = xprimef[s,sprime](x)
            end
            z0[i_x,s] = vcat(cf[s](x),nf[s](x),xprime0)
        end
    end
    cFB, nFB, IFB, xFB, zFB = find_first_best(para, S, 2)
    return BellmanEquation(para,S,xbar,time_0,z0,cFB,nFB,xFB,zFB)
end

"""
Finds the optimal policies
"""
function get_policies_time1(T::BellmanEquation,
                        i_x::Integer, x::AbstractFloat,
                        s::Integer, Vf::Array)
    para, S = T.para, T.S
    beta, Theta, G, Pi = para.beta, para.Theta, para.G, para.Pi
    U, Uc, Un = para.U, para.Uc, para.Un

    function objf(z::Vector,grad)
        c,xprime = z[1], z[2:end]
        n=c+G[s]

        Vprime = Array{typeof(beta)}(S)
        for sprime in 1:S
            Vprime[sprime] = Vf[sprime](xprime[sprime])
        end
        return -(U(c,n)+beta*dot(Pi[s,:],Vprime))
    end
    function cons(z::Vector,grad)
        c, xprime = z[1], z[2:end]
        n=c+G[s]
        return x - Uc(c,n)*c-Un(c,n)*n - beta*dot(Pi[s,:],xprime)
    end
    lb = vcat(0, T.xbar[1]*ones(S))
    ub = vcat(1-G[s], T.xbar[2]*ones(S))
    opt = Opt(:LN_COBYLA, length(T.z0[i_x,s])-1)
    min_objective!(opt, objf)
    equality_constraint!(opt, cons)
    lower_bounds!(opt, lb)
    upper_bounds!(opt, ub)
    maxeval!(opt, 300)
    maxtime!(opt, 10)
    init = vcat(T.z0[i_x,s][1], T.z0[i_x,s][3:end])
    for (i, val) in enumerate(init)
        if val > ub[i]
            init[i] = ub[i]
        elseif val < lb[i]
            init[i] = lb[i]
        end
    end
    (minf, minx, ret) = optimize(opt, init)
    T.z0[i_x,s] = vcat(minx[1], minx[1]+G[s], minx[2:end])
    return vcat(-minf, T.z0[i_x, s])
end
"""
Finds the optimal policies
"""
function get_policies_time0(T::BellmanEquation,
                        B_::AbstractFloat, s0::Integer, Vf::Array)
    para, S = T.para, T.S
    beta, Theta, G, Pi = para.beta, para.Theta, para.G, para.Pi
    U, Uc, Un = para.U, para.Uc, para.Un
    function objf(z,grad)
        c, xprime = z[1], z[2:end]
        n = c+G[s0]
        Vprime = Array{typeof(beta)}(S)
        for sprime in 1:S
            Vprime[sprime] = Vf[sprime](xprime[sprime])
        end
        return -(U(c,n)+beta*dot(Pi[s0,:],Vprime))
    end
    function cons(z::Vector, grad)
        c, xprime = z[1], z[2:end]
        n = c+G[s0]
        return -Uc(c,n)*(c-B_)-Un(c,n)*n - beta*dot(Pi[s0,:],xprime)
    end
    lb = vcat(0, T.xbar[1]*ones(S))
    ub = vcat(1-G[s0], T.xbar[2]*ones(S))
    opt = Opt(:LN_COBYLA, length(T.zFB[s0])-1)
    min_objective!(opt,objf)
    equality_constraint!(opt,cons)
    lower_bounds!(opt,lb)
    upper_bounds!(opt,ub)
    maxeval!(opt, 300)
    maxtime!(opt, 10)
    init = vcat(T.zFB[s0][1],T.zFB[s0][3:end])
    for (i, val) in enumerate(init)
        if val > ub[i]
            init[i] = ub[i]
        elseif val < lb[i]
            init[i] = lb[i]
        end
    end
    (minf,minx,ret) = optimize(opt,init)
    return vcat(-minf, vcat(minx[1], minx[1]+G[s0], minx[2:end]))
end

"""
Compute the planner's allocation by solving Bellman
equation.
"""
type Planners_Allocation_Bellman
    para::Para
    mc::MarkovChain
    S::Integer
    T::BellmanEquation
    mugrid::Vector
    xgrid::Vector
    Vf::Array
    policies::Array
end

"""
Initializes the class from the calibration Para
"""
function Planners_Allocation_Bellman(para::Para,mugrid::AbstractArray)
    mc = MarkovChain(para.Pi)
    G = para.G
    S = size(para.Pi,1) # number of states
    mugrid = collect(mugrid)
    #now find the first best allocation
    Vf, policies, T, xgrid = solve_time1_bellman(para, mugrid)
    T.time_0 = true #Bellman equation now solves time 0 problem
    return Planners_Allocation_Bellman(
        para,mc,S,T,mugrid,xgrid, Vf,policies)
end

"""
Solve the time 1 Bellman equation for calibration Para and initial grid mugrid0
"""
function solve_time1_bellman(para::Para, mugrid::AbstractArray)
    mugrid0 = mugrid
    S = size(para.Pi,1)
    #First get initial fit
    PP = Planners_Allocation_Sequential(para)
    Tbeta=typeof(PP.para.beta)
    c = Array{Tbeta}(2,length(mugrid))
    n = Array{Tbeta}(2,length(mugrid))
    x = Array{Tbeta}(2,length(mugrid))
    V = Array{Tbeta}(2,length(mugrid))
    for (i,mu) in enumerate(mugrid0)
        c[:,i], n[:,i], x[:,i], V[:,i] = time1_value(PP,mu)
    end
    c=c'
    n=n'
    x=x'
    V=V'
    Vf = Vector{Function}(2)
    cf = Vector{Function}(2)
    nf = Vector{Function}(2)
    xprimef = Array{Function}(2, S)
    for s in 1:2
        splc = Spline1D(vec(x[:,s])[end:-1:1], vec(c[:,s])[end:-1:1], k=3)
        spln = Spline1D(vec(x[:,s])[end:-1:1], vec(n[:,s])[end:-1:1], k=3)
        splV = Spline1D(vec(x[:,s])[end:-1:1], vec(V[:,s])[end:-1:1], k=3)
        cf[s] = x -> evaluate(splc, x)
        nf[s] = x -> evaluate(spln, x)
        Vf[s] = x -> evaluate(splV, x)
        for sprime in 1:S
            splxp = Spline1D(vec(x[:,s])[end:-1:1], vec(x[:,s])[end:-1:1], k=3)
            xprimef[s,sprime] = x -> evaluate(splxp, x)
        end
    end
    policies = [cf,nf,xprimef]
    #create xgrid
    xbar = [maximum(minimum(x,1)),minimum(maximum(x,1))]
    xgrid = linspace(xbar[1],xbar[2],length(mugrid0))
    #Now iterate on bellman equation
    T = BellmanEquation(para, xgrid, policies)
    diff = 1.0
    while diff > 1e-6
        if T.time_0 == false
            Vfnew, policies =
                fit_policy_function(PP,
                (i_x, x, s) -> get_policies_time1(T, i_x, x, s, Vf), xgrid)
        elseif T.time_0 == true
            Vfnew, policies =
                fit_policy_function(PP,
                (i_x, B_, s0) -> get_policies_time0(T, i_x, B_, s0, Vf), xgrid)
        else
            error("T.time_0 is $(T.time_0), which is invalid")
        end
        diff = 0.0
        for s in 1:S
            diff = max(diff,
                    maxabs((Vf[s](xgrid)-Vfnew[s](xgrid))/Vf[s](xgrid))
                        )
        end
        print("diff = $diff \n")
        Vf = Vfnew
    end
    # store value function policies and Bellman Equations
    return Vf, policies, T, xgrid
end

"""
Fits the policy functions PF using the points xgrid using UnivariateSpline
"""
function fit_policy_function(PP::Planners_Allocation_Sequential,
                            PF::Function,xgrid::AbstractArray)
    S = PP.S
    Vf = Vector{Function}(S)
    cf = Vector{Function}(S)
    nf = Vector{Function}(S)
    xprimef = Array{Function}(S, S)
    for s in 1:S
        PFvec = Array{typeof(PP.para.beta)}(length(xgrid), 3+S)
        for (i_x, x) in enumerate(xgrid)
            PFvec[i_x,:] = PF(i_x, x, s)
        end
        splV = Spline1D(xgrid, PFvec[:,1], s=0, k=1)
        splc = Spline1D(xgrid, PFvec[:,2], s=0, k=1)
        spln = Spline1D(xgrid, PFvec[:,3], s=0, k=1)
        Vf[s] = x -> evaluate(splV, x)
        cf[s] = x -> evaluate(splc, x)
        nf[s] = x -> evaluate(spln, x)
        for sprime in 1:S
            splxp = Spline1D(xgrid, PFvec[:,3+sprime], k=1)
            xprimef[s,sprime] = x -> evaluate(splxp, x)
        end
    end
    return Vf, [cf, nf, xprimef]
end

"""
Finds the optimal allocation given initial government debt B_ and state s_0
"""
function time0_allocation(pab::Planners_Allocation_Bellman,
                            B_::AbstractFloat, s0::Integer)
    xgrid = pab.xgrid
    if pab.T.time_0 == false
        z0 = get_policies_time1(pab.T, i_x, x, s, pab.Vf)
    elseif pab.T.time_0 == true
        z0 = get_policies_time0(pab.T, B_, s0, pab.Vf)
    else
        error("T.time_0 is $(T.time_0), which is invalid")
    end
    c0, n0, xprime0 = z0[2], z0[3], z0[4:end]
    return c0, n0, xprime0
end

"""
Simulates Ramsey plan for T periods
"""
function simulate(pab::Planners_Allocation_Bellman,
                B_::AbstractFloat, s_0::Integer, T::Integer, sHist=nothing)
    para, S, policies = pab.para, pab.S, pab.policies
    beta, Pi, Uc = para.beta, para.Pi, para.Uc
    cf, nf, xprimef = policies[1], policies[2], policies[3]
    if sHist == nothing
        sHist = QuantEcon.simulate(mc, s_0, T)
    end
    cHist = zeros(T)
    nHist = zeros(T)
    Bhist = zeros(T)
    TauHist = zeros(T)
    muHist = zeros(T)
    RHist = zeros(T-1)
    #time0
    cHist[1], nHist[1], xprime = time0_allocation(pab,B_,s_0)
    TauHist[1] = Tau(pab.para,cHist[1], nHist[1])[s_0]
    Bhist[1] = B_
    muHist[1] = 0.0
    #time 1 onward
    for t in 2:T
        s, x = sHist[t], xprime[sHist[t]]
        c, n, xprime = Vector{typeof(beta)}(S), nf[s](x), Vector{typeof(beta)}(S)
        for shat in 1:S
            c[shat] = cf[shat](x)
        end
        for sprime in 1:S
            xprime[sprime] = xprimef[s,sprime](x)
        end
        TauHist[t] = Tau(pab.para,c,n)[s]
        u_c = Uc(c,n)
        Eu_c = dot(Pi[sHist[t-1],:], u_c)
        muHist[t] = pab.Vf[s](x)
        RHist[t-1] = Uc(cHist[t-1],nHist[t-1])/(beta*Eu_c)
        cHist[t], nHist[t], Bhist[t] = c[s], n, x/u_c[s]
    end
    return cHist, nHist, Bhist, TauHist, sHist, muHist, RHist
end
end


text:
include("FTypesModule.jl")
include("BoxModule.jl")
include("MatrixUtilityModule.jl")
include("PhysicalUnitModule.jl")
include("RotationUtilModule.jl")
include("CSysModule.jl")
include("FESetModule.jl")
include("FENodeSetModule.jl")
include("FENodeToFEMapModule.jl")
include("FieldModule.jl")
include("GeneralFieldModule.jl")
include("NodalFieldModule.jl")
include("ElementalFieldModule.jl")
include("MeshUtilModule.jl")
include("MeshSelectionModule.jl")
include("MeshExportModule.jl")
include("MeshModificationModule.jl")
include("MeshImportModule.jl")
include("MeshQuadrilateralModule.jl")
include("MeshLineModule.jl")
include("MeshTriangleModule.jl")
include("MeshHexahedronModule.jl")
include("MeshTetrahedronModule.jl")
include("VectorCacheModule.jl")
include("ForceIntensityModule.jl")
include("SurfaceNormalModule.jl")
include("MatModule.jl")
include("AssemblyModule.jl")
include("IntegRuleModule.jl")
include("IntegDomainModule.jl")
include("FEMMBaseModule.jl")
include("AlgoBaseModule.jl")


text:
# Notation:
# λ is the initial step length
# α current trial step length
# β next trial step length
# d is the search direction
# x is the current iterate
# f is the objective
# φ is the line search objective and is a function of the step length only

# This file contains several implementation of what we might call "Backtracking".
# The AbstractBacktracking line searches try to satisfy the Amijo(-Goldstein)
# condition:
#     |f(x + α*d)| < (1-c_1*α)*|f(x)|
# That is: the function should

# As per [Nocedal & Wright, pp. 37] we don't have to think about the curvature
# condition as long as we use backtracking.

abstract type AbstractBacktracking end
abstract type BacktrackingInterp end

"""
  _safe_α(α_cand, α_curr, c, ratio)

Returns the safeguarded value of α in a Amijo
backtracking line search.

σ restriction 0 < c < ratio < 1
"""
function _safe_α(α_cand, α_curr, decrease=0.1, ratio=0.5)
  α_cand < decrease*α_curr && return decrease*α_curr
  α_cand > ratio*α_curr    && return ratio*α_curr

  α_cand # if the candidate is in the interval, just return it
end

"""
  Backtracking(;)

Constructs an object to control the Backtracking line search algorithm. The
algorithm tries out an initial point and then it iteratively tries to find
a good enough step length as measured by the improvement compared to a first
order Taylor approximation around a step length of zero.

The Backtracking constructor takes the following keyword arguments
 - ratio: the ratio between the current trial step length and the next
 - decrease: a parameter that controls when descent is sufficient
 - maxiter: maximum number of times to search for a better step length
 - interp: which type of interpolation to use, if any
 - steprange: the allowed range for steps lengths to be in given as a tuple
 - verbose: the verbosity level

`ratio` controls how much time to spend looking for a large step size. It
is very common to chose 1/2, but others can be chosen. It's

When chosing parameters, it's important to note that it must be true that
  0 < decrease < ratio

`maxiter` defaults to 26. With no interpolation and with a ratio of 1/2, this
means that we quit when the step size reaches sqrt(eps(Float64)). Provide another
value as appropriate.

"""
struct Backtracking{T1, T2, T3, TR} <: LineSearcher
    ratio::T1
	decrease::T1
	maxiter::T2
	interp::T3
	steprange::TR
	verbose::Bool
end
summary(bt::Backtracking) = "backtracking ("*summary(bt.interp)*")"
Backtracking(; ratio=0.5, decrease=1e-4, maxiter=26,
	           steprange=(lower=0, upper=Inf), interp=FixedInterp(),
	           verbose=false) =
 Backtracking(ratio, decrease, maxiter, interp, steprange, verbose)

struct FixedInterp <: BacktrackingInterp end
summary(fi::FixedInterp) = ("no interp")
struct FFQuadInterp <: BacktrackingInterp end
summary(ffq::FFQuadInterp) = ("quadratic interp")
struct FFFQuadInterp <: BacktrackingInterp end

"""
  interpolate(itp::BacktrackingInterp, ...)

Calculates the minimizer of a polynomial approximation to the line restricted
objective function.
"""
function interpolate(itp::FixedInterp, φ, φ0, dφ0, α, f_α, ratio)
	β = α
	α = ratio*α
	φ_α = φ(α)
	β, α, φ_α
end

# There are many ways to come up with the next step length trial value.
# If we want to interpolate, then we generally need to have some combination
# of values and derivatives of the function we're building an approximation
# for. In the non-linear equations case, a line search algorithm implies that
# some sort of mert function is used. When the squared two norm is used we
# can calculate a cubic interpolation using the "two point" method, and we
# have the two values and the derivative that we need from
#
# f(α) = ||F(xₙ+αdₙ)||²₂
# f(0) = ||F(xₙ)||²₂
# f'(0) = 2*(F'(xₙ)'dₙ)'F(xₙ) = 2*F(xₙ)'*(F'(xₙ)*dₙ) < 0
#
# if f'(0) >= 0, then dₙ  is not a descent direction for the merit function.
# This can happen for broyden. Notice, that when calculating dₙ using an
# inexact Newton's method, F'(xₙ)*dₙ can be obtained from the final residual.
# In either case, this is stored in φ.dφ0 as should be handled using the LineObjective.

# two-point parabolic
# at α = 0 we know f and f' from F and J as written above
# at αc define
function twopoint(f, f_0, df_0, α, f_α, ratio)
	ρ_lo, ρ_hi = 0.1, ratio
    # get the minimum (requires df0 < 0)
	c = (f_α - f_0 - df_0*α)/α^2

	# p(α) = f0 + df_0*α + c*α^2 is the function
	# we have df_0 < 0. Then if  f_α > f(0) then c > 0
	# by the expression above, and p is convex. Then,
	# we have a minimum between 0 and α at

	γ = -df_0/(2*c) # > 0 by df0 < 0 and c > 0
    # safeguard α
    return max(min(γ, α*ρ_hi), α*ρ_lo) # σs
end

function interpolate(itp::FFQuadInterp, φ, φ0, dφ0, α, f_α, ratio)
	β = α
	α = twopoint(φ, φ0, dφ0, α, f_α, ratio)
	φ_α = φ(α)
	β, α, φ_α
end

"""
    find_steplength(---)

Returns a step length, (merit) function value at step length and success flag.
"""
function find_steplength(mstyle, ls::Backtracking, φ::T, λ) where T
 	#== unpack ==#
	φ0, dφ0 = φ.φ0, φ.dφ0
	Tf = typeof(φ0)
	ratio, decrease, maxiter, verbose = Tf(ls.ratio), Tf(ls.decrease), ls.maxiter, ls.verbose

	#== factor in Armijo condition ==#
    t = -decrease*dφ0

    iter, α, β = 0, λ, λ # iteration variables
	f_α = φ(α) # initial function value

	if verbose
		println("Entering line search with step size: ", λ)
		println("Initial value: ", φ0)
		println("Value at first step: ", f_α)
	end
	is_solved = isfinite(f_α) && f_α <= φ0 + α*t
    while !is_solved && iter <= maxiter
        iter += 1
        β, α, f_α = interpolate(ls.interp, φ, φ0, dφ0, α, f_α, ratio)
		is_solved = isfinite(f_α) && f_α <= φ0 + α*t
    end

	ls_success = iter >= maxiter ? false : true

    if verbose
		!ls_success && println("maxiter exceeded in backtracking")
        println("Exiting line search with step size: ", α)
        println("Exiting line search with value: ", f_α)
    end
    return α, f_α, ls_success
end


text:
module ElectrostaticPIC1D

using Base.Threads
using ConcreteStructs
using ForwardDiff
using FFTW
using IterativeSolvers
using JLD2
using LinearAlgebra
using LoopVectorization
using Memoization
using Preconditioners
using QuadGK
using SparseArrays
using SpecialFunctions
using ThreadsX
using ToeplitzMatrices

include("BoundaryConditions.jl")

include("BasisFunctions.jl")
export BSpline, GaussianShape, TentShape, TopHatShape, DeltaFunctionShape
export BasisFunction, lower, upper, width, limits, centre, weight
include("Fields.jl")
export EquispacedValueGrid, PeriodicGridBC
export FourierField, FiniteDifferenceField, LSFEMField, GalerkinFEMField, FEMGrid
export cellcentres, solve!, update!, cell, cells, cellsize, energydensity
include("Particles.jl")
export AbstractParticle, Nuclide, Particle, pushposition!, pushvelocity!
export velocity, charge, mass, deposit!, antideposit, integral, basis, electricfield
export momentum, momentumdensity, energy, energydensity, charge, chargedensity
include("Plasmas.jl")
export Plasma, Species, positions, velocities
include("TimeIntegrators.jl")
export LeapFrogTimeIntegrator, LeapFrogTimeIntegrator
include("Simulations.jl")
export Simulation, SimulationParameters, run!, init!

end # module


text:
#=
Implement a URL shortener with the following methods:

    shorten(url), which shortens the url into a six-character alphanumeric string, such as zLg6wl.
    restore(short), which expands the shortened string into the original url. If no such shortened string exists, return null.

Hint: What if we enter the same URL twice?
=#

using Test
include("Solutions/problem55_URLshortener.jl")

Random.seed!(123)
@test shorten("google.ca") == "4UkInf"
@test shorten("google.ca") == "4UkInf"
@test shorten("amazon.com") == "edi7B7"

@test restore("4UkInf") == "google.ca"
@test isnothing(restore("df211sdd"))

max_stars_repo_name: DominiqueCaron/daily-coding-problem

max_issues_repo_name: DominiqueCaron/daily-coding-problem

max_forks_repo_name: DominiqueCaron/daily-coding-problem


text:
#=
Test Resynthesizer by comparing out images to reference good result image.

Using testing framework.
=#

# See "Code Loading"
# Either include, or import package
include("../src/apps/inpaint.jl")
# using Resynthesizer

# the Julia testing framework
using ReferenceTests
using Test

# profiling framework
using Profile

# for "testimage" of packaged, remote  test files ?
using TestImages

# For "load" local  files
using FileIO

using ColorTypes

#using Pkg
#Pkg.instantiate()

# This is executable, not a definition
@testset "Resynthesizer" begin

  # lighthouse image

  # Here the in image is the same as the reference
  # The in image has already been processed, we mask the result off,
  # and process it again.

  # input

  tiny = true
  donkey = false

  # a mask, that covers the lighthouse
  # height dimension first for Julia
  if tiny
      # 8w x 5h

      img_path = "/home/bootch/git/resynthesizerJulia/test/data/in/tinylighthouse.png"
      img = load(img_path)
      mask = falses(size(img))
      mask[3:4, 2:3] .= true
  elseif donkey
      img_path = "/home/bootch/git/resynthesizerJulia/test/data/indonkey.png"
      img = load(img_path)
      mask = falses(size(img))
      # square mask of the donkey itself
      mask[175:275,100:230] .= true
  else
      # 768x x 512h

      # The file is remote, in the  TestImage  pkg ??
      #img  = Float64.(Gray.(testimage("lighthouse")))
      img  = testimage("lighthouse")
      mask = falses(size(img))

      # Usual size mask
      #mask[50:350,300:400] .= true
      # Smaller masks, for dev
      mask[50:60,300:320] .= true
      # Tiny mask
      # mask[50:52,300:302] .= true
  end

  # time or profile or none
  #out = @time inpaint(img, mask)
  # @profile
  # or better to do this in the REPL?
  # out = inpaint(img, mask)
  out = @time inpaint(img, mask)

  img_path = "/home/bootch/git/resynthesizerJulia/test/out.png"
  save(img_path, out)

  # Julia macro that compares file and expression
  @test_reference "data/out/lighthouse.png" out
end

max_stars_repo_name: bootchk/resynthesizerJulia

max_issues_repo_name: bootchk/resynthesizerJulia

max_forks_repo_name: bootchk/resynthesizerJulia


text:
# Helper code

"""
  spdiagm_nonsquare(m, n, args...)

Construct a sparse diagonal matrix from Pairs of vectors and diagonals. Each
vector arg.second will be placed on the arg.first diagonal. By default (if
size=nothing), the matrix is square and its size is inferred from kv, but a
non-square size m×n (padded with zeros as needed) can be specified by passing
m,n as the first arguments.

# Arguments
  - `m::Int64`: First dimension of the output matrix
  - `n::Int64`: Second dimension of the output matrix
  - `args::Tuple{T} where T<:Pair{<:Integer,<:AbstractVector}`

# Outputs

  - sparse matrix of size mxn containing the values in args

"""
function spdiagm_nonsquare(m, n, args...)
    I, J, V = SparseArrays.spdiagm_internal(args...)
    return sparse(I, J, V, m, n)
end

# One-dimensional codes

"""
    nablasq_grid(n, h)

Laplacian matrix on a 1D grid

# Arguments:
    - `n`: Number of points
    - `h`: Aspect ratio

# Outputs:
    - discrete Laplacian matrix
"""
function nablasq_grid(n::Int64, h::Float64 = 1.0)
  o = ones(n) / h
  del = spdiagm_nonsquare(n + 1, n, -1 => -o, 0 => o)
  A1D = del' * del
  A1D[1, 1] = 1.0 / h ^ 2
  A1D[n, n] = 1.0 / h ^ 2
  return A1D
end

"""

   matern_1d_grid(y, idx, m, eps, h)

Matern Interpolation in one dimension

# Arguments:
  - `y`: the vector of y's for which the values are known
  - `idx`: the vector of indices for which values are to be interpolated
  - `m::Int64 = 1`: Matern parameter m
  - `eps = 0.0`: Matern parameter eps
  - `h = 1.0`: aspect ratio

# Outputs:
  - vector of interpolated data

# Example:
```<julia-repl>
x = 1:100
h = x[2] - x[1]
y = sin.(2 * pi * x * 0.2)
discard = randperm(100)[1:50]
# Laplace interpolation
y_lap = matern_1d_grid(y, discard, 1, 0.0, h)
# Matern interpolation
y_mat = matern_1d_grid(y, discard, 2, 0.1, h)
```

"""
function matern_1d_grid(y::Vector{T}, idx::Vector{Int64},
                        m::Int64 = 1, eps = 0.0, h::Float64 = 1.0) where{T<:Number}
  n = length(y)
  A1D = nablasq_grid(n, h)
  C = sparse(I, n, n)
  for i in idx
    C[i, i] = 0.0
  end
  if ((eps == 0)||(eps == 0.0)) && (m == 1)
    # Laplace Interpolation
    return ((C - (sparse(I, n, n) - C) * A1D)) \ (C * y)
  else
    # Matern Interpolation
    for i = 1:size(A1D,1)
      A1D[i, i] = A1D[i, i] + eps^2
    end
    A1DM = A1D ^ m
    return ((C - (sparse(I, n, n) - C) * A1DM)) \ (C * y)
  end
end

# Two dimensional codes

"""
  bdy_nodes(Nx, Ny)

...
# Arguments

  - `Nx::Int64`: the number of points in the first dimension
  - `Ny::Int64`: the number of points in the second dimension
...

...
# Outputs
  - vector containing the indices of coordinates on the boundary of the 2D rectangle
...

"""
function bdy_nodes(Nx, Ny)
  bdy = []
  xnb = []
  ynb = []
  counter = 0
      for j = 1:Ny
          for i = 1:Nx
              counter += 1
              if (j == 1 || j == Ny || i == 1 || i == Nx)
                  bdy = push!(bdy, counter)
                  if (j == 1 || j == Ny)
                      push!(ynb, 1)
                  else
                      push!(ynb, 2)
                  end
                  if (i == 1 || i == Nx)
                      push!(xnb, 1)
                  else
                      push!(xnb, 2)
                  end
              end
          end
      end
  return bdy, xnb, ynb
end

"""
    nablasq_grid(Nx, Ny, h, k)

Laplacian matrix on a 2D grid

# Arguments:
    - `Nx::Int64`: Number of points in first dimension
    - `Ny::Int64`: Number of points in second dimension
    - `h::Float64`: Aspect ratio in first dimension
    - `k::Float64`: Aspect ratio in second dimension

# Outputs:
    - discrete Laplacian matrix in 2D
"""
function nablasq_grid(Nx::Int64, Ny::Int64, h::Float64, k::Float64)
  o1 = ones(Nx) / h
  del1 = spdiagm_nonsquare(Nx + 1, Nx, -1 => -o1, 0 => o1)
  o2 = ones(Ny) / k
  del2 = spdiagm_nonsquare(Ny + 1, Ny, -1 => -o2,0 => o2)
  A2D = (kron(sparse(I, Ny, Ny), del1' * del1) +
          kron(del2' * del2, sparse(I, Nx, Nx)))
  bdy, xnb, ynb = bdy_nodes(Nx, Ny)
  count = 1
  for i in bdy
      A2D[i, i] = 0.0
      A2D[i, i] = A2D[i, i] + xnb[count] / h ^ 2 + ynb[count] / k ^ 2
      count += 1
  end
  return A2D
end

"""

  matern_2d_grid(mat, discard, m, eps, h, k)

...
# Arguments
  - `mat`: the matrix containing the image
  - `idx`: the linear indices of the nodes to be discarded
  - `eps`: Matern parameter eps
  - `m`: The Matern exponent (integer)
  - `h`: The aspect ratio in the first dimension
  - `k`: The aspect ratio in the second dimension

# Outputs
  - matrix containing the interpolated image

# Example:

```<julia-repl>
x = y = 1:30
h = k = x[2] - x[1]
y = sin.(2 * pi * x * 0.2) * cos.(2 * pi * y * 0.3)
discard = randperm(900)[1:450]
# Laplace interpolation
y_lap = matern_2d_grid(y, discard, 1, 0.0, h, k)
# Matern interpolation
y_mat = matern_2d_grid(y, discard, 2, 0.1, h, k)
...

"""
function matern_2d_grid(mat::Matrix, discard::Vector, m::Int64 = 1, eps = 0.0,
                        h::Float64 = 1.0, k::Float64 = 1.0)
    rows, columns = size(mat)
    A2D = nablasq_grid(rows, columns, h, k)
    sizeA = size(A2D, 1)
    C = sparse(I, sizeA, sizeA)
    for i in discard
        C[i, i] = 0.0
    end
    if ((m == 1)||(m == 1.0)) && (eps == 0.0)
        # Laplace interpolation
        u = ((C - (sparse(I, sizeA, sizeA) - C) * A2D)) \ (C * mat[:])
        return reshape(u, rows, columns)
    else
        # Matern interpolation
        for i = 1:sizeA
            A2D[i,i] = A2D[i,i] + eps^2
        end
        A2D = A2D^m
        u = ((C - (sparse(I, sizeA, sizeA) - C) * A2D)) \ (C * mat[:])
        return reshape(u, rows, columns)
    end
end

max_stars_repo_name: rayosborn/nxpeaks

max_issues_repo_name: rayosborn/nxpeaks

max_forks_repo_name: axmas-anl/nxrefine


text:
function test_fixedLoad()

    T = 2
    load = [1.0, 3.0]

    l = DR.DER.FixedLoad(index=0, T=T, load=load)
    model = GLPK.Optimizer()

    # Instantiate initial model
    h = DR.DER.House(
        0, T,
        [0.0, 0.0], [10.0, 10.0],
        [1.0, 1.0],
        [l]
    )

    model, var2idx, con2idx = DR.DER.build_model!(h, model)

    MOI.optimize!(model)

    # Test solution
    @test MOI.get(model, MOI.ObjectiveValue()) ≈ 4.0
    for t in 1:T
        pnet = var2idx[(:house, h.index, :pnet, t)]
        pfix = var2idx[(:fixed, l.index, :pnet, t)]
        @test MOI.get(model, MOI.VariablePrimal(), pnet) ≈ load[t]
        @test MOI.get(model, MOI.VariablePrimal(), pfix) ≈ load[t]
    end

    return nothing
end

@testset "FixedLoad" begin
    test_fixedLoad()
end


text:
#------------------------------
#   Solve for the equilibrium interest rate in the Hugget economy.
#
#   First, solve the HJB equation to get the A ("transition") matrix, then
#   solve the discretized KF equations. Then iterate over interest rates to
#   get the asset supply function. Update interest rate until convergence.
#
#   This code uses an implicit method for the solution: use
#   an 'upwind' scheme consisting of forwards and backwards difference
#   approximation to the derivative of the value function.
#
#   Assume: CRRA utility function u(c) = c^(1-s)/(1-s)
#   Note that r < rho
#
#   Written by Ben Moll et al (2015). Optimized for speed #   by SeHyoun Ahn. Translation into Julia by James Graham
#------------------------------

# Define types for workspace variables
type Huggett_model
  sigma::Float64            # risk aversion coefficient
  rho::Float64              # discount rate.
  z1::Float64               # Income in state 1
  z2::Float64               # Income in state 2
  z::Array{Float64}        # Vector for the two incomes
  la1::Float64              # State 1 switching intensity
  la2::Float64              # State 2 switching intensity
  la::Array{Float64}       # Vector for the two intensities
  I::Int64                  # Number of grids for assets
  Ir::Int64                   # Number of interest rate iterations
  amin::Float64             # Min assets = borrowing constraint
  amax::Float64
  rmin::Float64             # Min interest rate
  rmax::Float64             # Max interest rate
  agrid::Vector{Float64}    # Grid space for assets
  aa::Array{Float64}       # Matrix for asset grid space
  zz::Array{Float64}       # Matrix for income grid space
  da::Float64               # This is the implied grid step size
  maxit::Int64
  crit::Float64
  crit_S::Float64             # Tolerance for asset mkt clearing
  Delta::Float64
  Aswitch::SparseMatrixCSC{Float64,Int64}    # Transition matrix
  v0::Array{Float64}       # Initial guess for value function
  r0::Float64                 # Intial guess for interest rate
end

# Create model function
function Huggett_model(;sigma::Float64=1.5,
  rho::Float64=0.05,
  I::Int64=1000,
  Ir::Int64=40,
  amin::Float64=-0.15,
  amax::Float64=5.0,
  rmin::Float64=0.01,
  rmax::Float64=0.04,
  z1::Float64=0.1,
  z2::Float64=0.2,
  la1::Float64=1.2,
  la2::Float64=1.5,
  maxit::Int64=100,
  crit::Float64=1e-6,
  crit_S::Float64=1e-5,
  Delta::Float64=1000.0,
  r0::Float64=0.03)

  z       = hcat(z1,z2)
  la      = hcat(la1,la2)
  agrid   = collect(linspace(amin,amax,I))
  aa      = hcat(agrid, agrid)
  zz      = ones(I,1)*z
  da      = (amax-amin)/(I-1)
  v0 = hcat( (z[1] + r0.*agrid).^(1-sigma)/(1-sigma)/rho, (z[2] + r0.*agrid).^(1-sigma)/(1-sigma)/rho )

  Aswitch =  vcat( hcat(-speye(I)*la[1], speye(I)*la[1]), hcat(speye(I)*la[2],-speye(I)*la[2]) )

  Huggett_model(sigma, rho, z1, z2, z, la1, la2, la, I, Ir, amin, amax, rmin, rmax, agrid, aa, zz, da, maxit, crit, crit_S, Delta, Aswitch,  v0, r0)

end

# Creates an instance of the model
mod = Huggett_model()

# Define Utility funciton
function util(x; sigma::Float64=mod.sigma)
  u = x.^(1-sigma)/(1-sigma)
  return u
end

function utilprime(x; sigma::Float64=mod.sigma)
  uprime = x.^(-sigma)
  return uprime
end

# Define inverse of derivative of utility function
function utilprimeinv(x; sigma::Float64=mod.sigma)
  uprimeinv = x.^(-1/sigma)
  return uprimeinv
end

# Function to compute the KFE
function KFE(A::SparseMatrixCSC{Float64,Int64}; I::Int64=mod.I, da::Float64=mod.da)
  AT          = A'
  b           = zeros(2*I,1)
  i_fix       = 1    # fix one value so that matrix is not singular
  b[i_fix]    = 0.1
  row         = hcat(zeros(1,i_fix-1), 1, zeros(1,2*I-i_fix))
  AT[i_fix,:] = row
  #Solve linear system
  gg          = AT\b          # A'*gg = b --> gg = (A')^(-1)*b
  g_sum       = gg'*ones(2*I,1)*da
  gg          = gg./g_sum     # Normalize so sum(gg)*da) = 1
  g           = hcat(gg[1:I], gg[I+1:2*I])
  #check1      = g[:,1]'*ones(mod.I,1)*mod.da
  #check2      = g[:,2]'*ones(mod.I,1)*mod.da
  return g
end

# Initialization
dVf         = zeros(mod.I,2)      # Forward difference in each state
dVb         = zeros(mod.I,2)      # Backward difference in each state
c           = zeros(mod.I,2)      # Consumption in each state
A           = sparse(zeros(2*mod.I,2*mod.I))   # Transition matrix
r_r         = zeros(mod.Ir,1)     # Interest rate guesses
rmin_r      = zeros(mod.Ir,1)     # Lower bound for bisection
rmax_r      = zeros(mod.Ir,1)     # Upper bound for bisection
g_r         = zeros(mod.I,2,mod.Ir)   # Stationary distribution
adot        = zeros(mod.I,2,mod.Ir)   # Policy function
V_r         = zeros(mod.I,2,mod.Ir)   # Value function
V           = zeros(mod.I,2)        # Value function
S           = zeros(mod.Ir,1)     # Aggregate asset supply
r           = mod.r0
dist        = zeros(mod.maxit,1)
r_end       = 0                   # Last iteration value

for ir in 1:mod.Ir
  r_r[ir]     = r
  rmin_r[ir]  = mod.rmin
  rmax_r[ir]  = mod.rmax

  if ir>1
    mod.v0 = V_r[:,:,ir-1]
  end

  v = mod.v0

  for n = 1:mod.maxit
    V               = v
    # forward difference
    dVf[1:mod.I-1,:]    = (V[2:mod.I,:]-V[1:mod.I-1,:])/mod.da
    dVf[mod.I,:]        = utilprime(mod.z + r.*mod.amax)
    # backward difference
    dVb[1,:]            = utilprime(mod.z + r.*mod.amin)  #state constraint boundary condition
    dVb[2:mod.I,:]      = (V[2:mod.I,:]-V[1:mod.I-1,:])/mod.da

    # I_concave       = dVb .> dVf #indicator whether value function is concave (problems arise if this is not the case)

    #consumption and savings with forward difference
    cf              = utilprimeinv(dVf)
    ssf             = mod.zz + r.*mod.aa - cf
    #consumption and savings with backward difference
    cb              = utilprimeinv(dVb)
    ssb             = mod.zz + r.*mod.aa - cb
    #consumption and derivative of value function at steady state
    c0              = mod.zz + r.*mod.aa
    dV0             =utilprime(c0)

    # dV_upwind makes a choice of forward or backward differences based on the sign of the drift
    If              = ssf .> 0    # positive drift --> forward difference
    Ib              = ssb .< 0    # negative drift --> backward difference
    I0              = (1-If-Ib)   # at steady state
    #make sure backward difference is used at amax
    #Ib(I,:) = 1; If(I,:) = 0;
    #STATE CONSTRAINT at amin: USE BOUNDARY CONDITION UNLESS muf > 0:
    #already taken care of automatically

    dV_Upwind       = dVf.*If + dVb.*Ib + dV0.*I0    #important to include third term
    c               = utilprimeinv(dV_Upwind)
    u               = util(c)

    #CONSTRUCT MATRIX
    X               = - min(ssb,0)/mod.da
    Y               = - max(ssf,0)/mod.da + min(ssb,0)/mod.da
    Z               = max(ssf,0)/mod.da

    A1              = spdiagm(Y[:,1],0,mod.I,mod.I) + spdiagm(X[2:mod.I,1],-1,mod.I,mod.I) + spdiagm(Z[1:mod.I-1,1],1,mod.I,mod.I)
    A2              = spdiagm(Y[:,2],0,mod.I,mod.I) + spdiagm(X[2:mod.I,2],-1,mod.I,mod.I) + spdiagm(Z[1:mod.I-1,2],1,mod.I,mod.I)
    A               = vcat(hcat(A1, sparse(zeros(mod.I,mod.I))), hcat(sparse(zeros(mod.I,mod.I)), A2)) + mod.Aswitch
    B               = (mod.rho + 1/mod.Delta)*speye(2*mod.I) - A

    u_stacked       = vcat(u[:,1], u[:,2])
    V_stacked       = vcat(V[:,1], V[:,2])

    b               = u_stacked + V_stacked/mod.Delta
    V_stacked       = B\b             # Solve linear system of eqns

    V               = hcat(V_stacked[1:mod.I], V_stacked[mod.I+1:2*mod.I])

    Vchange         = V - v
    v               = V

    dist[n] = maximum(maximum(abs(Vchange)))
    if dist[n]<mod.crit
      @printf("-----------------------\n")
      @printf("Interest rate = %f\n", r)
      @printf("Value Function Converged, Iteration = %i\n", n)
      break
    end

  end

  ## Solve the Kolmogorov Forward (Fokker-Planck) equations
  g = KFE(A)

  # Pack up output
  g_r[:,:,ir] = g
  adot[:,:,ir] = mod.zz + r.*mod.aa - c
  V_r[:,:,ir] = V
  S[ir] = (g[:,1]'*mod.agrid*mod.da + g[:,2]'*mod.agrid*mod.da)[1,1]

  #UPDATE INTEREST RATE
  r_end = ir

  if S[ir]>mod.crit_S
    @printf("Excess Supply\n")
    mod.rmax = r
    r = 0.5*(r+mod.rmin)
  elseif S[ir]<-mod.crit_S
    @printf("Excess Demand\n")
    mod.rmin = r
    r = 0.5*(r+mod.rmax)
  elseif abs(S[ir])<mod.crit_S
    @printf("Equilibrium Found, Interest rate = %f/n", r)
    break
  end

end

## EXPORT DATA FOR GRAPHING IN PYTHON
using DataFrames
df = DataFrame(adot[:,:,r_end])
writetable("adot.csv", df)

df = DataFrame(c)
writetable("cons.csv", df)

df = DataFrame(g_r[:,:,r_end])
writetable("dist.csv", df)

## GRAPHS
using Gadfly

# Savings Policy function
myplot1 = plot(layer(x=mod.agrid,y=adot[:,1,r_end],Geom.line,Theme(line_width=1.5mm)), layer(x=mod.agrid,y=adot[:,2,r_end],Geom.line, Theme(line_width=1.5mm,default_color=colorant"red")),
xintercept=[mod.amin], yintercept=[0], Geom.hline(color=colorant"black", size=0.5mm), Geom.vline(color=colorant"black", size=0.5mm),
Guide.title("Savings policy function"),
Guide.xlabel("Wealth, a"),
Guide.ylabel("Savings, si(a)"),
Coord.cartesian(xmin=mod.amin, xmax=mod.amax), Theme(major_label_font_size=6mm,
minor_label_font_size=4mm))

using Cairo
using Fontconfig

# Distributions
plot(layer(x=mod.agrid,y=g_r[:,1,r_end],Geom.line,Theme(line_width=1.5mm)), layer(x=mod.agrid,y=g_r[:,2,r_end],Geom.line, Theme(line_width=1.5mm,default_color=colorant"red")),
xintercept=[mod.amin], Geom.vline(color=colorant"black", size=0.5mm), yintercept=[0], Geom.hline(color=colorant"black", size=0.5mm),
Guide.title("Stationary distributions"),
Guide.xlabel("Wealth, a"),
Guide.ylabel("Density, gi(a)"),
Coord.cartesian(xmin=mod.amin, xmax=1), Theme(major_label_font_size=6mm,
minor_label_font_size=4mm))


text:
map(calc_first,grammar)

function isProduction(p)
	# if we just do return typeof(p) == Symbol it might bug when p is an Enum
	#Dont aske me why
	if typeof(p) == Symbol
		return true
	end

	return false
end

function calc_follow(p::Union{Production,Int})
	subprods = -1

	if typeof(p) == Int
		subprods = getProd(p).subprods
	else
		subprods = p.subprods
	end

	for sprod in subprods
		for elemsprod_idx=1:(length(sprod)-1)

			if isProduction(sprod[elemsprod_idx]) #Symbols are non-terminals. Enums are tokens
				father_id=getProd_idx(sprod[elemsprod_idx])
				if isProduction(sprod[elemsprod_idx+1])

					for fst in getProd(elemsprod_idx+1).firsts
						#@show sprod[elemsprod_idx]
						#println("----")
						grammar[father_id].follows =
														unique(vcat(grammar[father_id].follows,fst))
					end
				else
					nxtToken = Int(sprod[elemsprod_idx+1])
					if nxtToken != Int(EPS)
						grammar[father_id].follows = unique(vcat(grammar[father_id].follows,nxtToken))
					end
				end

			else
				#If it's not a symbol, then it's a token, alas, there's no follows for it
			end

		end

		if typeof(last(sprod)) == Symbol
			son_id = getProd_idx(last(sprod))
			father_id = p.enum
			#@show p
			for fl in grammar[father_id].follows
				grammar[son_id].follows =
											 unique(vcat(grammar[son_id].follows,fl))
			end
		end
	end

end

map(calc_follow,grammar)

max_stars_repo_name: NaelsonDouglas/namelesscompiler

max_issues_repo_name: NaelsonDouglas/namelesscompiler

max_forks_repo_name: NaelsonDouglas/namelesscompiler


text:
@testset "424.longest-repeating-character-replacement.jl" begin
    @test character_replacement("ABAB", 2) == 4
    @test character_replacement("AABABBA", 1) == 4
    @test character_replacement("AABCABBB", 2) == 6
end


text:
"""
default template.
"""
const DEFAULT_TEMPLATE = Ref{String}("default")

"""
default user name.
"""
const DEFAULT_USERNAME = Ref{String}("")

function set_default_template(name::String="default")
    DEFAULT_TEMPLATE[] = name
    return
end

function set_default_username(name::String="")
    DEFAULT_USERNAME[] = name
    return
end

struct PDTN{name} end

PDTN(name::String) = PDTN{Symbol(name)}()

macro PDTN_str(name::String)
    return PDTN{Symbol(name)}
end

Base.show(io::IO, ::PDTN{template}) where {template} = print(io, "Pre-defined Template Name Type ", string(template))

"""
create a project or package.

# Arguments

- `path`: path of the project you want to create

# Options

- `--user <name>`: your GitHub user name for this package.
- `--template <template name>`: template name.

# Flags

- `-i, --interactive`: enable to start interactive configuration interface.
"""
@cast function create(path; user=DEFAULT_USERNAME[], interactive::Bool=false, template=DEFAULT_TEMPLATE[])
    if !isabspath(path)
        fullpath = joinpath(pwd(), path)
    else
        fullpath = path
    end

    if ispath(fullpath)
        error("$path exists, remove it or use a new path")
    end

    # if not specified, check if user.name is set in git
    if isempty(user)
        user = readchomp(`git config user.name`)
    end
    # TODO: use .ionrc to save user configuration
    # and reuse it next time

    t = create_template(PDTN(template), dirname(fullpath), user, interactive)
    t(basename(fullpath))
    return
end


text:
__precompile__(true)
""""
Add C, Python and type-based formatting to Str string literals

Copyright 2016-2018 Gandalf Software, Inc., Scott P. Jones
Licensed under MIT License, see LICENSE.md
"""
module StrFormat

using ModuleInterfaceTools

@api extend! Format, StrLiterals

function _parse_format(str, pos, fun)
    ex, j = parse(Expr, str, pos; greedy=false)
    check_expr(ex)
    ex, k = parse(Expr, string("a", str[pos:j-1]), 1, greedy=true)
    check_expr(ex)
    isa(ex, Symbol) && (println(string("a", str[pos:j-1])) ; dump(ex))
    ex.args[1] = fun
    ex, j
end

function _parse_fmt(sx::Vector{Any}, s::AbstractString, unescape::Function,
                    i::Integer, j::Integer, k::Integer)
    # Move past \\, k should point to '%'
    c, k = str_next(s, k)
    check_done(s, k, "Incomplete % expression")
    # Handle interpolation
    isempty(s[i:j-1]) || push!(sx, unescape(s[i:j-1]))
    if s[k] == '('
        # Need to find end to parse to
        ex, j = _parse_format(s, k, Format.fmt)
    else
        # Move past %, c should point to letter
        beg = k
        while true
            c, k = str_next(s, k)
            check_done(s, k, "Incomplete % expression")
            s[k] == '(' && break
        end
        ex, j = _parse_format(s, k, Format.cfmt)
        insert!(ex.args, 2, s[beg-1:k-1])
    end
    push!(sx, esc(ex))
    j
end

function _parse_pyfmt(sx::Vector{Any}, s::AbstractString, unescape::Function,
                      i::Integer, j::Integer, k::Integer)
    # Move past \\, k should point to '{'
    c, k = str_next(s, k)
    check_done(s, k, "Incomplete {...} Python format expression")
    # Handle interpolation
    isempty(s[i:j-1]) || push!(sx, unescape(s[i:j-1]))
    beg = k # start location
    c, k = str_next(s, k)
    while c != '}'
        check_done(s, k, string("\\{ missing closing } in ", c))
        c, k = str_next(s, k)
    end
    check_done(s, k, "Missing (expr) in Python format expression")
    c, k = str_next(s, k)
    c == '(' || parse_error(string("Missing (expr) in Python format expression: ", c))
    # Need to find end to parse to
    ex, j = _parse_format(s, k-1, Format.pyfmt)
    insert!(ex.args, 2, s[beg:k-3])
    push!(sx, esc(ex))
    j
end

function __init__()
    StrLiterals.interpolate['%'] = _parse_fmt
    StrLiterals.interpolate['{'] = _parse_pyfmt
end

end # module StrFormat


text:
__precompile__(false)
module PFKernels
    using SparseArrays
    using Printf
    include("utils.jl")

    include("indexes.jl")
    using .IndexSet
    include("parsers/parse_mat.jl")
    using .ParseMAT
    include("parsers/parse_psse.jl")
    using .ParsePSSE
    include("PowerSystem/PowerSystem.jl")
    using .PowerSystem

    const PS = PowerSystem

    using ForwardDiff
    abstract type AbstractBackend end
    struct CPUBackend <: AbstractBackend end
    struct CUDABackend <: AbstractBackend end
    struct AMDGPUBackend <: AbstractBackend end
    struct oneAPIBackend <: AbstractBackend end
    backend = nothing
    AT = Array
    try
        using CUDA
        @info("CUDA loaded")
        global backend = CUDABackend()
        global AT = CuVector

    catch
        try
            using AMDGPU
            @info("AMDGPU loaded")
            global backend = AMDGPUBackend()
            global AT = ROCVector
        catch
            try
                using oneAPI
                @info("oneAPI loaded")
                global backend = oneAPIBackend()
                global AT = oneArray
            catch
                @info("No GPU backend could be loaded, switching to CPU")
                global backend = CPUBackend()
            end
        end
    end

    # Defining dummy macros for the code to run
    if !(@isdefined CUDA)
        macro cuda(ex)
            return :( error("CUDA module not loaded") )
        end
    end

    if !(@isdefined AMDGPU)
        macro roc(ex)
            return :( error("AMDGPU module not loaded") )
        end
    end

    if !(@isdefined oneAPI)
        macro oneapi(ex...)
            return :( error("oneAPI module not loaded") )
        end
    end

    function kernel!(::AMDGPUBackend, F, args...)
        F_ = copy(F)
        io = open("code_lowered.txt", "w")
        print(io, @device_code_lowered @roc kernel_roc!(F_, args...))
        close(io)
        F_ = copy(F)
        io = open("code_typed.txt", "w")
        print(io, @device_code_lowered @roc kernel_roc!(F_, args...))
        close(io)
        F_ = copy(F)
        io = open("code_llvm.txt", "w")
        print(io, @device_code_lowered @roc kernel_roc!(F_, args...))
        close(io)
        wait(@roc kernel_roc!(F, args...))
    end

    function kernel!(::CUDABackend, F, args...)
    F_ = copy(F)
    io = open("code_lowered.txt", "w")
        @sync begin
            print(io, @device_code_lowered @cuda kernel_cuda!(F_, args...))
        end
    close(io)
    F_ = copy(F)
    io = open("code_typed.txt", "w")
        @sync begin
            print(io, @device_code_lowered @cuda kernel_cuda!(F_, args...))
        end
    close(io)
    io = open("code_llvm.txt", "w")
        @sync begin
            print(io, @device_code_lowered @cuda kernel_cuda!(F, args...))
        end
    close(io)
    end

    function kernel!(F, v_m, v_a,
                                  ybus_re_nzval, ybus_re_colptr, ybus_re_rowval,
                                  ybus_im_nzval, ybus_im_colptr, ybus_im_rowval,
                                  pinj, qinj, pv, pq, nbus, ::oneAPIBackend)
        n = length(pv) + length(pq)
        @oneapi items=n kernel_oneapi!(F, v_m, v_a,
                    ybus_re_nzval, ybus_re_colptr, ybus_re_rowval,
                    ybus_im_nzval, ybus_im_colptr, ybus_im_rowval,
                    pinj, qinj, pv, pq, nbus)
    end

    function loaddata(case)
        datafile = joinpath(dirname(@__FILE__), "..", "data", case)
        data_raw = ParsePSSE.parse_raw(datafile)
        data, bus_to_indexes = ParsePSSE.raw_to_exapf(data_raw)

        BUS_I, BUS_TYPE, PD, QD, GS, BS, BUS_AREA, VM, VA, BASE_KV, ZONE, VMAX, VMIN,
        LAM_P, LAM_Q, MU_VMAX, MU_VMIN = IndexSet.idx_bus()
        bus = data["bus"]
        gen = data["gen"]
        SBASE = data["baseMVA"][1]
        nbus = size(bus, 1)
        Ybus = PS.makeYbus(data, bus_to_indexes)
        V = Array{Complex{Float64}}(undef, nbus)
        for i in 1:nbus
            V[i] = bus[i, VM]*exp(1im * pi/180 * bus[i, VA])
        end

        Vm = abs.(V)
        Va = angle.(V)
        bus = data["bus"]
        gen = data["gen"]
        nbus = size(bus, 1)
        ngen = size(gen, 1)

        SBASE = data["baseMVA"][1]
        Sbus, Sload = PS.assembleSbus(gen, bus, SBASE, bus_to_indexes)
        ref, pv, pq = PS.bustypeindex(bus, gen, bus_to_indexes)
        return V, Ybus, Sbus, Sload, ref, pv, pq
    end
end


text:
# a type to hold data related to coloring the mesh
#TODO: reverse dimensions of revadj
mutable struct ColoringData
  adj_dict::Dict{Int, Array{Int, 1}}  # map local element to non-locals
  revadj::Array{Int32, 2}  # map non-local elements to local elements

  nonlocal_colors::Array{Int32, 1}
  # put other fields here

  function ColoringData(dict, revadj)
    obj = new()
    obj.adj_dict = dict
    obj.revadj = revadj
    return obj
  end
end

#------------------------------------------------------------------------------
# Distance-0 coloring
#------------------------------------------------------------------------------

# perform a distance-0 coloring of the mesh (ie. all elements same color)
function colorMesh0(mesh::PumiMesh)

  # make every element color 1
  for i=1:mesh.numEl
    apf.numberJ(mesh.coloring_Nptr, mesh.elements[i], 0, 0, 1)
  end

  return 1  # return number of colors
end

#------------------------------------------------------------------------------
# Distance- coloring
#------------------------------------------------------------------------------

# perform distance-1 coloring of mesh
# not sure if this works correctly
#=
function colorMesh1(mesh::PumiMesh, masks::Array{BitArray{1}})
# each element must have a different color than its neighbors with which it
# shares and edge

for i=1:mesh.numEl
  apf.numberJ(mesh.coloring_Nptr, mesh.elements[i], 0, 0, 0)
end

adj_size = 6  # guess number of neighboring faces
numc = 4  # guess number of colors
adj = Array{Ptr{Void}}( adj_size)  # hold neighboring faces
adj_color =zeros(Int32, adj_size)  # colors of neighboring faces
cnt_colors = zeros(Int32, numc)  # count how many of each color

for i=1:mesh.numEl
  el_i = mesh.elements[i]
  # get faces that share a vert with el_i
  # this is actually more restrictive than we want because
  # edge stabilization only causes interaction between
  # elements that share an edge
  # However, this is much easier to write an algorithm for

  # This is more restrictive than we want, but it ensure
  # that we void assigning the same color to two elements
  # bordering the same element
  # ie. we must avoid vertex neighbor non uniqueness,
  # which is not the same thing as having uniqueness because
  # there are unassigned values remaining
  # thus, this is a logical system where a double negative
  # is not a positive

  num_adj = apf.countBridgeAdjacent(mesh.m_ptr, el_i, 0, 2)

  if num_adj > adj_size
    println("resizing adj")
    println("element number = ", i)
    resize!(adj, num_adj)
    resize!(adj_color, num_adj)
    adj_size = num_adj
  end

  # need to verify this works in parallel (proper ghosting)
  apf.getBridgeAdjacent(adj)

  for j=1:num_adj
    adj_color[j] = apf.getNumberJ(mesh.coloring_Nptr, adj[j], 0, 0)
  end

  min_color = getMinColor2(adj_color, numc)

  if min_color > numc
    resize!(cnt_colors, min_color)
    cnt_colors[min_color] = 0  # initialize new value to zero
    numc = min_color
  end

  apf.numberJ(mesh.coloring_Nptr, el_i, 0, 0, min_color)

  cnt_colors[min_color] += 1  # update counts
#  masks[min_color][i] = true  # update mask

  fill!(adj_color, 0)

end
println("maximum number of neighbor faces = ", adj_size)
println("number of colors = ", numc)
println("number of each color = ", cnt_colors)
mesh.color_cnt = cnt_colors

return nothing

end
=#

#------------------------------------------------------------------------------
# Distance-2 coloring
#------------------------------------------------------------------------------

# perform distance-1 coloring of mesh
# this can be generalized to 3D via mesh.dim field
function colorMesh2(mesh::PumiMeshDG, colordata::ColoringData)
# each element must have a different color than its neighbors with which it
# shares and edge

println(mesh.f, "----- Entered colorMesh2 -----")

for i=1:mesh.numEl
  apf.numberJ(mesh.coloring_Nptr, mesh.elements[i], 0, 0, 0)
end

nfaces = mesh.numFacesPerElement
numc = nfaces + 1  # guess number of colors

adj = Array{Ptr{Void}}(nfaces)  # distance-1 edge neighbors
adj2 = Array{Array{Ptr{Void}, 1}}(nfaces)  # distance-2 edge neighbors + distance-1

for i=1:nfaces
  adj2[i] = Array{Ptr{Void}}(nfaces+1)
end

el_faces = Array{Ptr{Void}}(mesh.numFacesPerElement)
part_nums = Array{Cint}(1)
matched_entities = Array{Ptr{Void}}(1)
adj_entities = Array{Ptr{Void}}(1)
matchdata = LocalNeighborMatches(el_faces, part_nums, matched_entities, adj_entities)

colors = zeros(Int32, mesh.numFacesPerElement*(mesh.numFacesPerElement+1))

cnt_colors = zeros(Int32, numc)  # count how many of each color
adj_dict = colordata.adj_dict

for i=1:mesh.numEl
#  println("Processing element ", i)
  if haskey(adj_dict, i)  # skip any element on a parallel boundary
    continue
  end

  el_i = mesh.elements[i]

  # get faces that share an edge with el_i, and all elements
  # those elements share an edge with
  # this is actually more restrictive than we want because
  # a unique distance-1 coloring, not distance-2
  # However, this is much easier to write an algorithm for

  # This is more restrictive than we want, but it ensure
  # that we avoid assigning the same color to two elements
  # bordering the same element
  # ie. we must avoid distance-2 neighbor non uniqueness,
  # which is not the same thing as having uniqueness because
  # there are unassigned values remaining
  # thus, this is a logical system where a double negative
  # is not a positive

  getDistance2Colors(mesh, i, adj, adj2, colors, matchdata)

  min_color = getMinColor2(colors, numc)

  if min_color > numc
    println(mesh.f, "adding color ", min_color)
    resize!(cnt_colors, min_color)
    cnt_colors[min_color] = 0  # initialize new value to zero
    numc = min_color
  end

  apf.numberJ(mesh.coloring_Nptr, el_i, 0, 0, min_color)

  cnt_colors[min_color] += 1  # update counts

  fill!(colors, 0)

end

# now color the parallel boundary elements and their neighbors
numc = colorMeshBoundary2(mesh, colordata, numc, cnt_colors)

println(mesh.f, "number of colors = ", numc)
println(mesh.f, "number of each color = ", cnt_colors)
mesh.color_cnt = cnt_colors

return numc

end

# perform distance-2 coloring of mesh
function colorMesh2(mesh::PumiMesh2)
# each element must have a different color than its neighbors with which it
# shares and edge

# figure out the number of colors
# this is a lot of random memory access

for i=1:mesh.numEl
  apf.numberJ(mesh.coloring_Nptr, mesh.elements[i], 0, 0, 0)
end

nfaces = mesh.numFacesPerElement

#adj_size = 6  # guess number of neighboring faces
numc = nfaces + 1  # guess number of colors
#adj = Array{Ptr{Void}}(adj_size)  # hold neighboring faces
#adj_color =zeros(Int32, adj_size)  # colors of neighboring faces

adj = Array{Ptr{Void}}(nfaces)  # distance-1 edge neighbors
adj2 = Array{Array{Ptr{Void}, 1}}(nfaces)  # distance-2 edge neighbors + distance-1

for i=1:3
  adj2[i] = Array{Ptr{Void}}(nfaces+1)
end

el_faces = Array{Ptr{Void}}(mesh.numFacesPerElement)
part_nums = Array{Cint}(1)
matched_entities = Array{Ptr{Void}}(1)
adj_entities = Array{Ptr{Void}}(1)
matchdata = LocalNeighborMatches(el_faces, part_nums, matched_entities, adj_entities)

# space for distance-2 neighbors + self
colors = zeros(Int32, nfaces*(nfaces+1))

cnt_colors = zeros(Int32, numc)  # count how many of each color
for i=1:mesh.numEl
#  println("Processing element ", i)
  el_i = mesh.elements[i]
  # get faces that share an edge with el_i, and all elements
  # those elements share an edge with
  # this is actually more restrictive than we want because
  # a unique distance-1 coloring, not distance-2
  # However, this is much easier to write an algorithm for

  # This is more restrictive than we want, but it ensure
  # that we avoid assigning the same color to two elements
  # bordering the same element
  # ie. we must avoid distance-2 neighbor non uniqueness,
  # which is not the same thing as having uniqueness because
  # there are unassigned values remaining
  # thus, this is a logical system where a double negative
  # is not a positive

  getDistance2Colors(mesh, i, adj, adj2, colors, matchdata)

  min_color = getMinColor2(colors, numc)

  if min_color > numc
    resize!(cnt_colors, min_color)
    cnt_colors[min_color] = 0  # initialize new value to zero
    numc = min_color
  end

  apf.numberJ(mesh.coloring_Nptr, el_i, 0, 0, min_color)

  cnt_colors[min_color] += 1  # update counts

  fill!(colors, 0)

end
mesh.color_cnt = cnt_colors

return numc

end

#TODO: investigate whether duplicating any non-local elements that share a face
#      with more than one local elements produces better coloring
function colorMeshBoundary2(mesh::PumiMeshDG, colordata::ColoringData, numc, cnt_colors)

  println(mesh.f, "----- Entered colorMeshBoundary2 -----")
  nfaces = mesh.numFacesPerElement
  colordata.nonlocal_colors = zeros(Int32, mesh.numSharedEl)
  adj = Array{Ptr{Void}}(nfaces)  # distance-1 edge neighbors
  adj2 = Array{Array{Ptr{Void}, 1}}(nfaces)  # distance-2 edge neighbors + distance-1
  self = Array{Ptr{Void}}(1)  # the current element pointer
  for i=1:nfaces
    adj2[i] = Array{Ptr{Void}}(nfaces+1)
  end

  el_faces = Array{Ptr{Void}}(mesh.numFacesPerElement)
  part_nums = Array{Cint}(1)
  matched_entities = Array{Ptr{Void}}(1)
  adj_entities = Array{Ptr{Void}}(1)
  matchdata = LocalNeighborMatches(el_faces, part_nums, matched_entities, adj_entities)

  const local_start = 1
  const nonlocal_start = nfaces*(nfaces+1) + local_start
  const nonlocal_d2_start = (nfaces-1)*nfaces + nonlocal_start
  const final_start = (nfaces-2)*(nfaces-1) + nonlocal_d2_start

  colors = zeros(Int32, final_start-1)

  # these calculations are for a 2D mesh (triangles), but are similar for
  # 3d
  # the first 3 sets of four elements are for the colors of the distance-2
  # neighbors + their common distance-2 neighbor
  # the last 2*3 entries are for the non-local neighbors of the distance-1
  # neighbors
  # the last 2 entries are for the local distance-2 neighbors connected by a
  # non-local element
  # (for the case where a single element has 2 non local neighbors)
  local_colors = sview(colors, local_start:(nonlocal_start-1))
  nonlocal_d1neighborcolors = sview(colors, nonlocal_start:(nonlocal_d2_start-1))
  nonlocal_neighborcolors = sview(colors, nonlocal_d2_start:(final_start-1))

  for i in keys(colordata.adj_dict)
    el_i = mesh.elements[i]
    self[1] = el_i

    num_adj = getDistance2Colors(mesh, i, adj, adj2, local_colors, matchdata)
    getNonLocalColors(mesh, sview(adj, 1:num_adj), colordata, nonlocal_d1neighborcolors)
    getNonlocalDistance2Colors(mesh, i, colordata, nonlocal_neighborcolors)
    # the non-local elements have not been colored yet, so no need to get their colors
#    getNonLocalColors(mesh, self, colordata, nonlocal_neighborcolors)

    min_color = getMinColor2(colors, numc)
    if min_color > numc
      resize!(cnt_colors, min_color)
      cnt_colors[min_color] = 0  # initialize new value to zero
      numc = min_color
    end
    # record the color
    apf.numberJ(mesh.coloring_Nptr, el_i, 0, 0, min_color)
    cnt_colors[min_color] += 1  # update counts

    fill!(colors, 0)
  end

  # now do non-local elements
  # it is preferable to do these last because they have strictly fewer
  # adjacent elements than other elemenets, so the odds are better they can
  # get a lower color by doing coloring all the boundary elements first

  # store local d1 adjacencies in adj (there are nfaces-1 maximum)
  # store colors of their local d1 neighbors + self color in next
  # (nfaces-1)(nfaces) entries of colors
  # store their nonlocal d1 neighbors in next (nfaces-1)(nfaces-1) entries of
  # colors
  # It is possible to have a fully occluded element, so the two extreme cases
  # are nfaces local neighbors and 0 nonlocal neighbor or 0 local neighbors and
  # nfaces nonlocal neighbors.
  # There layout of the colors array is:
  # nfaces (for d1 neighbors), then
  # nfaces*nfaces (for the d2 local neighbors)
  # nfaces*nfaces (for the d2 nonlocal neighbors)

  revadj = colordata.revadj
  colors = zeros(Int32, nfaces + nfaces*nfaces + nfaces*nfaces)

  d1_neighbors = aview(colors, 1:nfaces)
  local_d2_neighbors = Array{ContiguousView{Int32,1, Array{Int32, 1}}}(nfaces)
  nonlocal_d2_neighbors = Array{ContiguousView{Int32,1, Array{Int32, 1}}}(nfaces)
  # starting indices for the local and nonlocal sections
  pos_local = nfaces + 1
  pos_nonlocal = nfaces + nfaces*nfaces + 1

  for i=1:nfaces
    # local d2 neighbors
    local_d2_neighbors[i] = aview(colors, pos_local:(pos_local + nfaces - 1))
    pos_local += nfaces

    # nonlocal d2 neighbors
    nonlocal_d2_neighbors[i] = aview(colors, pos_nonlocal:(pos_nonlocal + nfaces - 1))
    pos_nonlocal += nfaces
  end
  flush(mesh.f)

  # check the sections don't overlab
  @assert pos_local == (nfaces + nfaces*nfaces + 1)
  @assert pos_nonlocal == length(colors) + 1

  d1_ptr = Array{Ptr{Void}}(1)
  for i=1:mesh.numSharedEl
    for j=1:nfaces
      neighbor = revadj[i, j]
      if neighbor != 0
        # get d1 neighbor color
        d1_ptr[1] = mesh.elements[neighbor]
        d1_neighbors[j] = apf.getNumberJ(mesh.coloring_Nptr, d1_ptr[1], 0, 0)

        # get d2 local and nonlocal neighbor colors (by getting the d1
        # neighbors of neighbor of the d1 neighbor)
        d2_local = local_d2_neighbors[j]
        d2_nonlocal = nonlocal_d2_neighbors[j]
        getDistance1Colors(mesh, neighbor, adj, d2_local)
        getNonLocalColors(mesh, d1_ptr, colordata, d2_nonlocal)
      end
    end

    min_color = getMinColor2(colors, numc)

    if min_color > numc
      println(mesh.f, "adding color ", min_color)
      resize!(cnt_colors, min_color)
      cnt_colors[min_color] = 0  # initialize new value to zero
      numc = min_color
    end
    colordata.nonlocal_colors[i] = min_color
    cnt_colors[min_color] += 1
  end

  fill!(colors, 0)

  return numc
end

struct LocalNeighborMatches
  faces::Array{Ptr{Void}, 1}
  part_nums::Array{Cint, 1}
  matched_entities::Array{Ptr{Void}, 1}
  adj_entities::Array{Ptr{Void}, 1}
end

function getNeighborMatches(mesh::PumiMesh, el::Ptr{Void}, adj::AbstractArray{Ptr{Void}}, num_used::Integer, thunk::LocalNeighborMatches)
# check to see if any faces of the element are matched, and get their parent
# elements
# adj is the partially populated array of neighbor elements
# num_used is the number of entries in adj already used

  faces = Array{Ptr{Void}}(mesh.numFacesPerElement)
  n = apf.getDownward(mesh.m_ptr, el, mesh.dim-1, faces)

  part_nums = Array{Cint}(1)
  matched_entities = Array{Ptr{Void}}(1)
  adj_entities = Array{Ptr{Void}}(1)
  for i=1:n
    nmatches = apf.countMatches(mesh.m_ptr, faces[i])
    apf.getMatches(part_nums, matched_entities)
    @assert nmatches <= 1

    if nmatches == 1 && part_nums[1] == mesh.myrank
      other_entity = matched_entities[1]
      nel = apf.countAdjacent(mesh.m_ptr, other_entity, mesh.dim)
      @assert nel == 1
      apf.getAdjacent(adj_entities)

      # put it in the next spot in the array
      num_used += 1
      adj[num_used] = adj_entities[1]
    end
  end

  return num_used
end

function getDistance2Colors(mesh::PumiMesh, elnum::Integer, adj::AbstractArray{Ptr{Void}}, adj2::AbstractArray{Array{Ptr{Void}, 1}}, colors, matchdata::LocalNeighborMatches)
# get the distance-2 neighbors of a given element
# repeats included
# adj : array to be populated with distance-1 edge neighbors
# adj2 : array of arrays to be populated with sum of distance-1 and distance-2 edge
# neighbors
# colors : array to be populated with colors of elements in adj2
# getting the distance-2 colors (instead of the distance-1 colors) is necessary
# to avoid a problem where 2 elements that have a common neighbor element
# will be visited without visiting the common neighbor first, coloring
# the 2 elements the same

  el_i = mesh.elements[elnum]

  num_adj = apf.countBridgeAdjacent(mesh.m_ptr, el_i, mesh.dim-1, mesh.dim)
#  adj = Array{Ptr{Void}}(3)
  apf.getBridgeAdjacent(adj)

  # check for matches
  if num_adj < mesh.numFacesPerElement
    num_adj = getNeighborMatches(mesh, el_i, adj, num_adj, matchdata)
  end

#  adj2 = Array{Array{Ptr{Void}, 1}}(num_adj)  # hold distance 2 neighbors

  adj_cnt = zeros(Int, num_adj)
  for j=1:num_adj
    num_adj_j = apf.countBridgeAdjacent(mesh.m_ptr, adj[j], mesh.dim-1, mesh.dim)
 #   adj2[j] = Array{Ptr{Void}}(num_adj_j + 1)
    apf.getBridgeAdjacent(adj2[j])

    if num_adj_j < mesh.numFacesPerElement
      num_adj_j = getNeighborMatches(mesh, adj[j], adj2[j], num_adj_j, matchdata)
    end

    adj2[j][num_adj_j + 1] = adj[j]  # include the distance-1 neighbor
    adj_cnt[j] = num_adj_j + 1
  end

  num_adj_total = sum(adj_cnt)  # total number of neighbors, including repeats

#  colors = zeros(Int32, num_adj_total)

  # get the colors of all the elements
  # also flatten adj2 array of arrays into a single array
  index = 1
  for i=1:num_adj
    for j=1:adj_cnt[i]
      colors[index] = apf.getNumberJ(mesh.coloring_Nptr, adj2[i][j], 0, 0)
      index += 1
    end
  end

  return num_adj
end

function getNonlocalDistance2Colors(mesh::PumiMeshDG, elnum::Integer, colordata::ColoringData, colors)
# get the distance 2 neighbors of a local element that are connected via a non-local element, not including the current element

  el = mesh.elements[elnum]
  nfaces = mesh.numFacesPerElement
  pos = 1
  if haskey(colordata.adj_dict, elnum)
    vals = colordata.adj_dict[elnum]
    for i=1:length(vals)
      if (vals[i] != 0)
        for j=1:(nfaces-1)
          neighbor = colordata.revadj[vals[i] - mesh.numEl, j]
          if neighbor != 0 && neighbor != elnum
            neighbor_ptr = mesh.elements[neighbor]
            colors[pos] = apf.getNumberJ(mesh.coloring_Nptr, neighbor_ptr, 0, 0)
            pos += 1
          end
        end
      end
    end
  end

  return nothing
end

function getDistance1Colors(mesh::PumiMeshDG, elnum::Integer, adj, colors)
# get the distance1 colors of a specified element

  el = mesh.elements[elnum]
  num_adj = apf.countBridgeAdjacent(mesh.m_ptr, el, mesh.dim-1, mesh.dim)
  apf.getBridgeAdjacent(adj)

  for i=1:num_adj
    colors[i] = apf.getNumberJ(mesh.coloring_Nptr, adj[i], 0, 0)
  end

  return num_adj
end

function getNonLocalColors(mesh, adj::AbstractArray{Ptr{Void}}, colordata::ColoringData, colors::AbstractArray)
# gets the colors of the non-local neighbors of the specified elements
# adj holds the pointers to the elements

  adj_dict = colordata.adj_dict
  pos = 1
  for i=1:length(adj)
    elnum = apf.getNumberJ(mesh.el_Nptr, adj[i], 0, 0) + 1
    if haskey(adj_dict, elnum)
      vals = adj_dict[elnum]
      for j=1:length(vals)
        if vals[j] != 0
          val_i = vals[j]
          colors[pos] = colordata.nonlocal_colors[val_i - mesh.numEl]
          pos += 1
        end
      end
    end
  end

  return nothing
end

function getMinColor(adj::AbstractArray{T}) where T
# adj contains colors of adjacent elements
  min_color = 1
  sort!(adj)  # adj must be in increasing order for this to work
  for i=1:length(adj)
    if min_color < adj[i]
      continue
    elseif min_color == adj[i]
	min_color = adj[i] + 1
    end  # if min_color > adj[i] do nothing
  end  # end for loop

  return min_color
end

#TODO: see if output can be removed
function getMinColor2(adj::AbstractArray{T}, numc::Integer) where T
# ensure uniqueness of neighboring colors
# adj is array of colors of adjacent faces
# numc is the current number of colors

  mask = zeros(Bool, numc)
  sort!(adj)
  min_color = 0
  for i=1:length(adj) # identify already used colors
    if adj[i] != 0
      mask[adj[i]] = true
    end
  end

  mask_sum = sum(mask)

  if mask_sum == numc  # all existing colors used, so add another
    min_color = numc + 1
  elseif mask_sum == (numc - 1)  # there is exactly 1 color remaining
    # find out which color is missing and use it
    for i=1:numc
      if !mask[i]  # if mask is false
        min_color = i
      end
    end
  else  # some colors are missing
    min_color = getMinColor(adj)  # get the minimum
  end

  @assert min_color != 0

  return min_color

end

#------------------------------------------------------------------------------
# functions that get the bookkeeping information after coloring is done
#------------------------------------------------------------------------------

function getColors0(mesh, masks::AbstractArray{BitArray{1}, 1})
# populate the masks that indicate which elements are perturbed by which
# colors
# for a distance-0 coloring, every element is perturbed and there is only
# 1 color

  masks[1] = trues(mesh.numEl)
  return nothing
end

function getColors1(mesh, masks::AbstractArray{BitArray{1}, 1}, neighbor_colors::AbstractArray{T, 2}, neighbor_nums::AbstractArray{T2, 2}; verify=true ) where {T, T2}
# verify edge neighbor faces have different colors (ie. distance-1 coloring
# of graph where elements are the vertices connected by edges if
# there is a mesh edge connecting them
# also construct the BitArray masks that describe which elements to perturb
# for which colors, and the list of neighbor + self colors
# masks is an array of bitarrays, number of arrays = number of colors
# neighbor_colors is 4 by numEl array of integers holding the colors
# of the nieghbors + own color
# neighbor_nums is 4 by numEl array of integers holding the element numbers
# of the neighbors + self

nfaces = mesh.numFacesPerElement
adj = Array{Ptr{Void}}(nfaces + 1) # pointers to element i + 3 neighbors
adj_color = zeros(Int32, nfaces + 1)  # element colors

el_faces = Array{Ptr{Void}}(mesh.numFacesPerElement)
part_nums = Array{Cint}(1)
matched_entities = Array{Ptr{Void}}(1)
adj_entities = Array{Ptr{Void}}(1)
matchdata = LocalNeighborMatches(el_faces, part_nums, matched_entities, adj_entities)

# initialize masks to 0 (false)
for i=1:length(masks)
  masks[i] = falses(mesh.numEl)
end

#if verify
#  println("verifying distance-1 coloring")
#end

cnt = 0
for i=1:mesh.numEl
  el_i = mesh.elements[i]

  # check edge neighbors only
  num_adj = apf.countBridgeAdjacent(mesh.m_ptr, el_i, mesh.dim-1, mesh.dim)
  @assert num_adj <= nfaces
  apf.getBridgeAdjacent(adj)

  if num_adj < mesh.numFacesPerElement
    num_adj = getNeighborMatches(mesh, el_i, adj, num_adj, matchdata)
  end

  adj[num_adj + 1] = el_i  # insert the current element
  # get color, element numbers
  for j=1:(num_adj + 1)
    adj_color[j] = apf.getNumberJ(mesh.coloring_Nptr, adj[j], 0, 0)
    neighbor_colors[j, i] = adj_color[j]
    neighbor_nums[j, i] = apf.getNumberJ(mesh.el_Nptr, adj[j], 0, 0) + 1
  end

  color_i = adj_color[num_adj + 1]  # color of current element
  masks[color_i][i] = true  # indicate element i gets perturbed by color_i

  if verify
    sort!(adj_color)

    # remove leading zeros
    nnz = countnz(adj_color)
    nz_arr = zeros(eltype(adj_color), nnz)
    start_idx = length(adj_color) - nnz + 1
    for k=1:nnz
      nz_arr[k] = adj_color[start_idx]
      start_idx += 1
    end

    #TODO: do this by verifying adjacent values are not equal, rather than
    #      calling unique which allocates memory
    if nz_arr != unique(nz_arr)
      println("element ", i, " has non unique colors")
      println("adj_color = ", nz_arr)
      cnt += 1
    end
  end   # end if verify
  fill!(adj_color, 0)
end

if verify
#  println("color-1 verification finished")
  if cnt != 0
    println(STDERR, "number of element with non unique coloring = ", cnt)
    throw(ErrorException("non unique element coloring"))
  end
end

return cnt

end

function getColors1(mesh, colordata::ColoringData, masks::AbstractArray{BitArray{1}, 1}, neighbor_colors::AbstractArray{T, 2}, neighbor_nums::AbstractArray{T2, 2}; verify=true ) where {T, T2}
# verify edge neighbor faces have different colors (ie. distance-1 coloring
# of graph where elements are the vertices connected by edges if
# there is a mesh edge connecting them
# also construct the BitArray masks that describe which elements to perturb
# for which colors, and the list of neighbor + self colors
# masks is an array of bitarrays, number of arrays = number of colors
# neighbor_colors is 4 by numEl array of integers holding the colors
# of the nieghbors + own color
# neighbor_nums is 4 by numEl array of integers holding the element numbers
# of the neighbors + self

nfaces = mesh.numFacesPerElement
adj = Array{Ptr{Void}}(nfaces + 1)   # pointers to element i + 3 neighbors
adj_color = zeros(Int32, nfaces + 1)  # element colors
adj_elnum = zeros(Int32, nfaces + 1)  # element numbers
adj_dict = colordata.adj_dict
nonlocal_colors = colordata.nonlocal_colors

# initialize masks to 0 (false)
for i=1:length(masks)
  masks[i] = falses(mesh.numEl)
end

if verify
  println(mesh.f, "verifying distance-1 coloring")
end

cnt = 0
for i=1:mesh.numEl
  el_i = mesh.elements[i]
  elnum = apf.getNumberJ(mesh.el_Nptr, el_i, 0, 0) + 1

  # check edge neighbors only
  num_adj = apf.countBridgeAdjacent(mesh.m_ptr, el_i, mesh.dim-1, mesh.dim)
  @assert num_adj <= nfaces
  apf.getBridgeAdjacent(adj)
  adj[num_adj + 1] = el_i  # insert the current elementa

  # get color, element numbers for non-local elements
  pos = 1  #current index in the adj_color, adj_elnum arrays
  if haskey(adj_dict, elnum)
    nonlocal_els = adj_dict[elnum]
    for j=1:(nfaces-1)
      nonlocal_elnum = nonlocal_els[j]
      if nonlocal_elnum != 0  # if this is a real ghost neighbor
        elcolor = nonlocal_colors[nonlocal_elnum - mesh.numEl]
        adj_color[pos] = elcolor
        adj_elnum[pos] = nonlocal_elnum
        pos += 1
      end
    end
  end

  # get color, element numbers for local elements + self
  for j=1:(num_adj + 1)
    adj_color[pos] = apf.getNumberJ(mesh.coloring_Nptr, adj[j], 0, 0)
    adj_elnum[pos] = apf.getNumberJ(mesh.el_Nptr, adj[j], 0, 0) + 1
    pos += 1
  end

  # copy them into the global arrays
  for j=1:(nfaces + 1)
    neighbor_colors[j, i] = adj_color[j]
    neighbor_nums[j, i] = adj_elnum[j]
  end

  color_i = adj_color[pos - 1]  # color of current element
  masks[color_i][i] = true  # indicate element i gets perturbed by color_i

  if verify
    sort!(adj_color)

    # remove leading zeros
    nnz = countnz(adj_color)
    nz_arr = zeros(eltype(adj_color), nnz)
    start_idx = length(adj_color) - nnz + 1
    for k=1:nnz
      nz_arr[k] = adj_color[start_idx]
      start_idx += 1
    end

    #TODO: do this by comparing adjacent entries, without allocating memory
    if nz_arr != unique(nz_arr)
      println(mesh.f, "element ", i, " has non unique colors")
      println(mesh.f, "adj_color = ", nz_arr)
      cnt += 1
    end
  end   # end if verify
  fill!(adj_color, 0)
  fill!(adj_elnum, 0)
end

if verify
  println(mesh.f, "color-1 verification finished"); flush(mesh.f)
  if cnt != 0
    throw(ErrorException("Non unique distance-1 coloring: $cnt elements have bad colors"))
  end
  println(mesh.f, "number of element with non unique coloring = ", cnt)
  println(mesh.f, "abc")
end

# now create color masks for non-local elements
println(mesh.f, "getting non-local color masks")
nonlocal_masks = Array{Array{BitArray{1}, 1}}(mesh.npeers)
start_elnum = 1
for i=1:mesh.npeers
  nonlocal_masks[i] = Array{BitArray{1}}(mesh.numColors)
  masks_i = nonlocal_masks[i]  # starting element number for this peer
  numel_i = mesh.shared_element_offsets[i+1] - mesh.shared_element_offsets[i]

  # create the masks BitArrays
  for j=1:mesh.numColors
    masks_i[j] = falses(numel_i)
  end

  # set flag to true if the element has the color
  for k=1:numel_i
    global_elnum = start_elnum + k - 1
    color_k = nonlocal_colors[global_elnum]
    masks_i[color_k][k] = true
  end

  start_elnum += numel_i
end

return cnt, nonlocal_masks

end

function getPertNeighbors0(mesh::PumiMesh)
# get the element that is perturbed for each element for each color
# of a distance-0 coloring
# for a distance-0 coloring, each element is only perturbed by
# other dofs on the element

  pertNeighborEls = zeros(Int32, mesh.numEl, mesh.numColors)

  for i=1:mesh.numEl
    pertNeighborEls[i, 1] = i
  end

  return pertNeighborEls
end

function getPertNeighbors1(mesh::PumiMesh)
# populate the array with the element that is perturbed for each element
# for each color for a distance-1 coloring
# element number == 0 if no perturbation

#  println("getting neighbor list")
  pertNeighborEls = zeros(Int32, mesh.numEl, mesh.numColors)

for color = 1:mesh.numColors
  num_neigh = size(mesh.neighbor_colors, 1)
#  fill!(arr, 0)
  for i=1:mesh.numEl
    # find out if current element or its neighbors have the current color
    pos = 0
    for j=1:num_neigh
      if color == mesh.neighbor_colors[j, i]
	pos = j
	break
      end
    end

    if pos != 0
      pertNeighborEls[i, color] = mesh.neighbor_nums[pos, i]
    else
       pertNeighborEls[i, color] = 0
     end

   end  # end loop over elements

end  # end loop over colors

  return pertNeighborEls
end

# perterubed neighbors for edge-based residual
# TODO: make this 3d
function getPertEdgeNeighbors(mesh::PumiMesh)

  neighbor_nums = zeros(Int32, mesh.numEl, 3)

  edges = Array{Ptr{Void}}(3)  # hold element edges
  adj = Array{Ptr{Void}}(2)  # hold adjacent elements

  for i=1:mesh.numEl
    el_i = mesh.elements[i]

    apf.getDownward(mesh.m_ptr, el_i, 1, edges)

    #=
    # check edge neighbors only
    num_adj = apf.countBridgeAdjacent(mesh.m_ptr, el_i, 1, 2)
    @assert num_adj <= 3
    apf.getBridgeAdjacent(adj)
    =#

    # get color, element numbers
    for j=1:3 # loop over edges
      # in use, this array is traveres by numEl first, so we have to
      # populate it by rows here
      num_adj = apf.countAdjacent(mesh.m_ptr, edges[j], 2)
      @assert num_adj <= 2
      if num_adj == 2  # if there is another adjacnet element
	apf.getAdjacent(adj)

	# figure out which adjacent element is the *other* one
	if adj[1] == el_i
	  other_el = adj[2]
	else
	  other_el = adj[1]
	end

        neighbor_nums[i, j] = apf.getNumberJ(mesh.el_Nptr, other_el, 0, 0) + 1
      end  # end if num_adj == 2
    end  # end loop j=1:3
  end  # end i=1:mesh.numEl

  return neighbor_nums

end  # end getPertEdgeNeighbors