text:
"""
    Tuple(perm::Permutation)

Extract tuple representation of Permutation.

The result can be passed to `permutedims` and `PermutedDimsArray`.

Returns `nothing` if `perm` is a `NoPermutation`.

# Examples

```jldoctest
julia> Tuple(Permutation(3, 2, 1))
(3, 2, 1)

julia> Tuple(NoPermutation()) === nothing
true
```
"""
Base.Tuple(::Permutation{p}) where {p} = p
Base.Tuple(::NoPermutation) = nothing

"""
    length(perm::AbstractPermutation)

Returns length of permutation.

For `NoPermutation`, returns `nothing`.

# Examples

```jldoctest
julia> length(Permutation(3, 2, 1))
3

julia> length(NoPermutation()) === nothing
true
```
"""
Base.length(::Permutation{p}) where {p} = length(p)
Base.length(::NoPermutation) = nothing

"""
    isperm(perm::AbstractPermutation) -> Bool

Returns `true` if `perm` is a valid permutation, `false` otherwise.

The result is known at compile time.

# Examples

```jldoctest
julia> isperm(Permutation(3, 1, 2))
true

julia> isperm(Permutation(4, 1, 2))
false

julia> isperm(NoPermutation())
true
```
"""
Base.isperm(::NoPermutation) = true

function Base.isperm(::Permutation{P}) where {P}
    P :: Tuple
    if @generated
        # Call isperm tuple implementation in base Julia
        pp = isperm(P)
        :( $pp )
    else
        isperm(P)
    end
end

"""
    *(p::AbstractPermutation, collection)

Apply permutation to the given collection.

The collection may be a `Tuple` or a `CartesianIndex` to be reordered. If `p` is
a [`Permutation`](@ref), the collection must have the same length as `p`.

# Examples

```jldoctest
julia> p = Permutation(2, 3, 1);

julia> p * (36, 42, 14)
(42, 14, 36)

julia> p * CartesianIndex(36, 42, 14)
CartesianIndex(42, 14, 36)
```
"""
*(::NoPermutation, t::Tuple) = t

@inline function *(::Permutation{p,N}, t::Tuple{Vararg{Any,N}}) where {N,p}
    @inbounds ntuple(i -> t[p[i]], Val(N))
end

@inline *(p::AbstractPermutation, I::CartesianIndex) = CartesianIndex(p * Tuple(I))

"""
    *(p::AbstractPermutation, q::AbstractPermutation)

Compose two permutations: apply permutation `p` to permutation `q`.

Note that permutation composition is non-commutative.

# Examples

```jldoctest
julia> p = Permutation(2, 3, 1);

julia> q = Permutation(1, 3, 2);

julia> p * q
Permutation(3, 2, 1)

julia> q * p
Permutation(2, 1, 3)

julia> p * inv(p)
Permutation(1, 2, 3)

julia> inv(p) * p
Permutation(1, 2, 3)
```
"""
*(p::Permutation, q::Permutation) = Permutation(p * Tuple(q))
*(::NoPermutation, q::AbstractPermutation) = q
*(p::AbstractPermutation, ::NoPermutation) = p
*(p::NoPermutation, ::NoPermutation) = p

"""
    /(y::AbstractPermutation, x::AbstractPermutation)

Get relative permutation needed to get from `x` to `y`. That is, the permutation
`p` such that `p * x == y`.

# Examples

```jldoctest
julia> x = Permutation(3, 1, 2);

julia> y = Permutation(2, 1, 3);

julia> p = y / x
Permutation(3, 2, 1)

julia> p * x == y
true
```
"""
function /(::Permutation{q,N}, ::Permutation{p,N}) where {p, q, N}
    if @generated
        perm = map(v -> findfirst(==(v), p)::Int, q)
        @assert Permutation(perm) * p === q
        :( Permutation($perm) )
    else
        perm = map(v -> findfirst(==(v), p)::Int, q)
        @assert Permutation(perm) * p === q
        Permutation(perm)
    end
end

/(y::AbstractPermutation, ::NoPermutation) = y

# In this case, the result is the inverse permutation of `x`, such that
# `perm * x == (1, 2, 3, ...)`.
/(::NoPermutation, x::Permutation{p,N}) where {p,N} =
    identity_permutation(Val(N)) / x

"""
    \\(p::AbstractPermutation, x)

Undo permutation `p` from permuted collection `x`.

In other words, apply inverse of permutation `p` to `x`. This is effectively
equivalent to `inv(p) * x`.
"""
\(p::AbstractPermutation, x) = inv(p) * x

"""
    inv(p::Permutation)
    invperm(p::Permutation)

Returns the inverse permutation of `p`.

Functionally equivalent to Julia's `invperm`, with the advantage that the result
is a compile time constant.

See also [`/`](@ref).

# Examples

```jldoctest
julia> p = Permutation(2, 3, 1);

julia> q = inv(p)
Permutation(3, 1, 2)

julia> t_orig = (36, 42, 14);

julia> t_perm = p * t_orig
(42, 14, 36)

julia> q * t_perm === t_orig
true

```
"""
Base.inv(x::AbstractPermutation) = NoPermutation() / x
Base.invperm(x::AbstractPermutation) = inv(x)

"""
    identity_permutation(::Val{N})
    identity_permutation(A::AbstractArray{T,N})

Returns the identity permutation `Permutation(1, 2, ..., N)`.
"""
identity_permutation(::Val{N}) where {N} = Permutation(ntuple(identity, Val(N)))
identity_permutation(A::AbstractArray) = identity_permutation(Val(ndims(A)))

"""
    isidentity(p::Permutation)

Returns `true` if `p` is an identity permutation, i.e. if it is equivalent to
`(1, 2, 3, ...)`.

```jldoctest
julia> isidentity(Permutation(1, 2, 3))
true

julia> isidentity(Permutation(1, 3, 2))
false

julia> isidentity(NoPermutation())
true
```
"""
isidentity(::NoPermutation) = true

function isidentity(perm::Permutation)
    N = length(perm)
    perm === identity_permutation(Val(N))
end

# Comparisons: (1, 2, ..., N) is considered equal to NoPermutation, for any N.
==(::Permutation{p}, ::Permutation{q}) where {p,q} = p === q
==(::NoPermutation, ::NoPermutation) = true
==(p::Permutation, ::NoPermutation) = isidentity(p)
==(np::NoPermutation, p::Permutation) = p == np

"""
    append(p::Permutation, ::Val{M})

Append `M` non-permuted dimensions to the given permutation.

# Examples

```jldoctest
julia> append(Permutation(2, 3, 1), Val(2))
Permutation(2, 3, 1, 4, 5)

julia> append(NoPermutation(), Val(2))
NoPermutation()
```
"""
function append(::Permutation{p}, ::Val{M}) where {p, M}
    N = length(p)
    Permutation(p..., ntuple(i -> N + i, Val(M))...)
end

append(np::NoPermutation, ::Val) = np

"""
    prepend(p::Permutation, ::Val{M})

Prepend `M` non-permuted dimensions to the given permutation.

# Examples

```jldoctest
julia> prepend(Permutation(2, 3, 1), Val(2))
Permutation(1, 2, 4, 5, 3)

julia> prepend(NoPermutation(), Val(2))
NoPermutation()
```
"""
function prepend(::Permutation{p}, ::Val{M}) where {p, M}
    Permutation(ntuple(identity, Val(M))..., (M .+ p)...)
end

prepend(np::NoPermutation, ::Val) = np


text:
using ModelBaseEcon
using SparseArrays
using Test

@testset "Options" begin
    o = Options(tol=1e-7, maxiter=25)
    @test getoption(o, tol=1e7) == 1e-7
    @test getoption(o, abstol=1e-10) == 1e-10
    @test "abstol" ∉ o
    @test getoption!(o, abstol=1e-11) == 1e-11
    @test :abstol ∈ o
    @test setoption!(o, reltol=1e-3, linear=false) isa Options
    @test "reltol" ∈ o
    @test :linear ∈ o
    @test getoption!(o, tol=nothing, linear=true, name="Zoro") == (1e-7, false, "Zoro")
    @test "name" ∈ o && o.name == "Zoro"
    z = Options(o)
    z.name = "Oro"
    @test o.name == "Zoro"
end

@testset "Vars" begin
    y1 = :y
    y2 = ModelSymbol(:y)
    y3 = ModelSymbol("y3", :y)
    y4 = ModelSymbol(quote "y4" y end)
    @test hash(y1) == hash(:y)
    @test hash(y2) == hash(:y)
    @test hash(y3) == hash(:y)
    @test hash(y4) == hash(:y)
    @test_throws ArgumentError ModelSymbol(:(x + 5))
    @test y1 == y2
    @test y3 == y1
    @test y1 == y4
    @test y2 == y3
    @test y2 == y4
    @test y3 == y4
    ally = Symbol[y1,y2,y3,y4]
    @test y1 in ally
    @test y2 in ally
    @test y3 in ally
    @test y4 in ally
    @test indexin([y1,y2,y3,y4], ally) == [1,1,1,1]
    ally = ModelSymbol[y1,y2,y3,y4,:y,quote "y5" y end]
    @test indexin([y1,y2,y3,y4], ally) == [1,1,1,1]
    @test length(unique(hash.(ally))) == 1
    ally = Dict{Symbol,Any}()
    get!(ally, y1, "y1")
    get!(ally, y2, "y2")
    @test length(ally) == 1
    @test ally[y3] == "y1"
    ally = Dict{ModelSymbol,Any}()
    get!(ally, y1, "y1")
    get!(ally, y2, "y2")
    @test length(ally) == 1
    @test ally[y3] == "y1"
    @test sprint(print, y2, context=IOContext(stdout, :compact => true)) == "y"
    @test sprint(print, y2, context=IOContext(stdout, :compact => false)) == "y"
    @test sprint(print, y3, context=IOContext(stdout, :compact => true)) == "y"
    @test sprint(print, y3, context=IOContext(stdout, :compact => false)) == "\"y3\" y"

end

@testset "VarTypes" begin
    lvars = ModelSymbol[]
    push!(lvars, :ly)
    push!(lvars, quote "ly" ly end)
    push!(lvars, quote @log ly end)
    push!(lvars, quote "ly" @log ly end)
    push!(lvars, quote @lin ly end)
    push!(lvars, quote "ly" @lin ly end)
    push!(lvars, quote @steady ly end)
    push!(lvars, quote "ly" @steady ly end)
    push!(lvars, ModelSymbol(:ly, :lin))
    for i = 1:length(lvars)
        for j = i + 1:length(lvars)
            @test lvars[i] == lvars[j]
        end
        @test lvars[i] == :ly
    end
    @test lvars[1].type == :lin
    @test lvars[2].type == :lin
    @test lvars[3].type == :log
    @test lvars[4].type == :log
    @test lvars[5].type == :lin
    @test lvars[6].type == :lin
    @test lvars[7].type == :steady
    @test lvars[8].type == :steady
    @test lvars[9].type == :lin
    for i = 1:length(lvars)
        @test sprint(print, lvars[i], context=IOContext(stdout, :compact => true)) == "ly"
    end
    @test sprint(print, lvars[1], context=IOContext(stdout, :compact => false)) == "ly"
    @test sprint(print, lvars[2], context=IOContext(stdout, :compact => false)) == "\"ly\" ly"
    @test sprint(print, lvars[3], context=IOContext(stdout, :compact => false)) == "@log ly"
    @test sprint(print, lvars[4], context=IOContext(stdout, :compact => false)) == "\"ly\" @log ly"
    @test sprint(print, lvars[5], context=IOContext(stdout, :compact => false)) == "ly"
    @test sprint(print, lvars[6], context=IOContext(stdout, :compact => false)) == "\"ly\" ly"
    @test sprint(print, lvars[7], context=IOContext(stdout, :compact => false)) == "@steady ly"
    @test sprint(print, lvars[8], context=IOContext(stdout, :compact => false)) == "\"ly\" @steady ly"

    let m = Model()
        @variables m p q r
        @variables m begin
            x; @log y; @steady z;
        end
        @test [v.type for v in m.allvars] == [:lin, :lin, :lin, :lin, :log, :steady]
    end
    let m = Model()
        @shocks m p q r
        @shocks m begin
            x; @log y; @steady z;
        end
        @test [v.type for v in m.allvars] == [:shock, :shock, :shock, :shock, :shock, :shock]
    end
    let m = Model()
        @logvariables m p q r
        @logvariables m begin
            x; @log y; @steady z;
        end
        @test [v.type for v in m.allvars] == [:log, :log, :log, :log, :log, :log]
    end
    let m = Model()
        @steadyvariables m p q r
        @steadyvariables m begin
            x; @log y; @steady z;
        end
        @test [v.type for v in m.allvars] == [:steady, :steady, :steady, :steady, :steady, :steady]
    end
end

module E
    using ModelBaseEcon
end
@testset "Evaluations" begin
    ModelBaseEcon.initfuncs(E)
    E.eval(ModelBaseEcon.makefuncs(:(x + 3 * y), [:x, :y], mod=E))
    @test :resid_1 ∈ names(E, all=true)
    @test :RJ_1 ∈ names(E, all=true)
    @test E.resid_1([1.1, 2.3]) == 8.0
    @test E.RJ_1([1.1, 2.3]) == (8.0, [1.0, 3.0])
end

@testset "Misc" begin
    m = Model()
    out = let io = IOBuffer()
        print(io, m.flags)
        readlines(seek(io, 0))
    end
    @test length(out) == 3
    for line in out[2:end]
        sline = strip(line)
        @test isempty(sline) || length(split(sline, "=")) == 2
    end
    @test_throws ModelBaseEcon.ModelErrorBase ModelBaseEcon.modelerror()
    @variables m x y z
    @logvariables m k l m
    @steadyvariables m p q r
    @shocks m a b c
    for s in (:a, :b, :c)
        @test m.:($s) isa ModelSymbol && isshock(m.:($s))
    end
    for s in (:x, :y, :z)
        @test m.:($s) isa ModelSymbol && islin(m.:($s))
    end
    for s in (:k, :l, :m)
        @test m.:($s) isa ModelSymbol && islog(m.:($s))
    end
    for s in (:p, :q, :r)
        @test m.:($s) isa ModelSymbol && issteady(m.:($s))
    end
    @equations m begin
        p[t] = 0
    end
    @initialize m
    @test Symbol(m.variables[1]) == m.variables[1]
    m.sstate.values .= rand(length(m.sstate.values))
    @test begin (l, s) = m.sstate.x.data; l == m.sstate.x.level && s == m.sstate.x.slope end
    @test begin (l, s) = m.sstate.k.data; exp(l) == m.sstate.k.level && exp(s) == m.sstate.k.slope end

    xdata = m.sstate.x[1:8, ref=3]
    @test xdata[3] ≈ m.sstate.x.level
    @test xdata ≈ m.sstate.x.level .+ ((1:8) .- 3) .* m.sstate.x.slope
    kdata = m.sstate.k[1:8, ref=3]
    @test kdata[3] ≈ m.sstate.k.level
    @test kdata ≈ m.sstate.k.level .* m.sstate.k.slope .^ ((1:8) .- 3)

    @test_throws Exception m.sstate.x.data = [1,2]
    @test_throws ArgumentError m.sstate.nosuchvariable

    @steadystate m m = l
    @steadystate m slope m = l
    @test length(m.sstate.constraints) == 2

    let io = IOBuffer()
        show(io, m.sstate.x)
        lines = split(String(take!(io)), '\n')
        @test length(lines) == 1 && occursin('+', lines[1])

        show(io, m.sstate.k)
        lines = split(String(take!(io)), '\n')
        @test length(lines) == 1 && !occursin('+', lines[1]) && occursin('*', lines[1])

        m.sstate.y.slope = 0
        show(io, m.sstate.y)
        lines = split(String(take!(io)), '\n')
        @test length(lines) == 1 && !occursin('+', lines[1]) && !occursin('*', lines[1])

        m.sstate.l.slope = 1
        show(io, m.sstate.l)
        lines = split(String(take!(io)), '\n')
        @test length(lines) == 1 && !occursin('+', lines[1]) && !occursin('*', lines[1])

        show(io, m.sstate.p)
        lines = split(String(take!(io)), '\n')
        @test length(lines) == 1 && !occursin('+', lines[1]) && !occursin('*', lines[1])

        ModelBaseEcon.show_aligned5(io, m.sstate.x, mask=[true, false])
        lines = split(String(take!(io)), '\n')
        @test length(lines) == 1 && length(split(lines[1], '?')) == 2

        ModelBaseEcon.show_aligned5(io, m.sstate.k, mask=[false, false])
        lines = split(String(take!(io)), '\n')
        @test length(lines) == 1 && length(split(lines[1], '?')) == 3

        ModelBaseEcon.show_aligned5(io, m.sstate.l, mask=[false, false])
        println(io)
        ModelBaseEcon.show_aligned5(io, m.sstate.y, mask=[false, false])
        println(io)
        ModelBaseEcon.show_aligned5(io, m.sstate.p, mask=[false, true])
        lines = split(String(take!(io)), '\n')
        @test length(lines) == 3
        for line in lines
            @test length(split(line, '?')) == 2
        end
    end
end

module MetaTest
using ModelBaseEcon
    params = @parameters
    custom(x) = x + one(x)
    const val = 12.0
    params.b = custom(val)
    params.a = @link custom(val)
end

@testset "Parameters" begin
    params = Parameters()
    push!(params, :a => 1.0)
    push!(params, :b => @link 1.0 - a)
    push!(params, :c => @alias b)
    push!(params, :e => [1,2,3])
    push!(params, :d => @link (sin(2π / e[3])))
    @test length(params) == 5
    # dot notation evaluates
    @test params.a isa Number
    @test params.b isa Number
    @test params.c isa Number
    @test params.d isa Number
    @test params.e isa Vector{<:Number}
    # [] notation returns the holding structure
    a = params[:a]
    b = params[:b]
    c = params[:c]
    d = params[:d]
    e = params[:e]
    @test a isa ModelParam
    @test b isa ModelParam
    @test c isa ModelParam
    @test d isa ModelParam
    @test e isa ModelParam
    @test a.depends == Set([:b])
    @test b.depends == Set([:c])
    @test c.depends == Set([])
    @test d.depends == Set([])
    @test e.depends == Set([:d])
    # circular dependencies not allowed
    @test_throws ArgumentError push!(params, :a => @alias b)
    # even deep ones
    @test_throws ArgumentError push!(params, :a => @alias c)
    # even when it is in an expr
    @test_throws ArgumentError push!(params, :a => @link 5 + b^2)
    @test_throws ArgumentError push!(params, :a => @link 3 - c)

    @test params.d ≈ √3 / 2.0
    params.e[3] = 2
    update_links!(params)
    @test 1.0 + params.d ≈ 1.0

    params.d = @link cos(2π / e[2])
    @test params.d ≈ -1.0

    @test_throws ArgumentError @alias a + 5
    @test_throws ArgumentError @link 28

    @test MetaTest.params.a ≈ 13.0
    @test MetaTest.params.b ≈ 13.0
    MetaTest.eval(quote custom(x) = 2x + one(x) end)
    update_links!(MetaTest.params)
    @test MetaTest.params.a ≈ 25.0
    @test MetaTest.params.b ≈ 13.0

    @test @alias(c) == ModelParam(Set(), :c, nothing)
    @test @link(c) == ModelParam(Set(), :c, nothing)
    @test @link(c + 1) == ModelParam(Set(), :(c + 1), nothing)

    @test_throws ArgumentError params[:contents] = 5
    @test_throws ArgumentError params.abc
end

@testset "ifelse" begin
    m = Model()
    @variables m x
    @equations m begin
        x[t] = 0
    end
    @initialize m
    @test_throws ArgumentError ModelBaseEcon.process_equation(m, :(y[t] = 0))
    @test_throws ArgumentError ModelBaseEcon.process_equation(m, :(x[t] = p))
    @test_throws ArgumentError ModelBaseEcon.process_equation(m, :(x[t] = if false 2 end))
    @test ModelBaseEcon.process_equation(m, :(x[t] = if false 2 else 0 end)) isa Equation
    @test ModelBaseEcon.process_equation(m, :(x[t] = ifelse(false, 2, 0))) isa Equation
end

@testset "Meta" begin
    mod = Model()
    @parameters mod a = 0.1 b = @link(1.0 - a)
    @variables mod x
    @shocks mod sx
    @equations mod begin
    x[t - 1] = sx[t + 1]
        @lag(x[t]) = @lag(sx[t + 2])
        #
        x[t - 1] + a = sx[t + 1] + 3
        @lag(x[t] + a) = @lag(sx[t + 2] + 3)
        #
        x[t - 2] = sx[t]
        @lag(x[t], 2) = @lead(sx[t - 2], 2)
        #
        x[t] - x[t - 1] = x[t + 1] - x[t] + sx[t]
        @d(x[t]) = @d(x[t + 1]) + sx[t]
        #
        (x[t] - x[t + 1]) - (x[t - 1] - x[t]) = sx[t]
        @d(x[t] - x[t + 1]) = sx[t]
        #
        x[t] - x[t - 2] = sx[t]
        @d(x[t],0,2) = sx[t]
        #
        (x[t] - x[t - 1]) - (x[t - 1] - x[t - 2]) = sx[t]
        @d(x[t],2) = sx[t]
        #
        (x[t] - x[t - 2]) - (x[t - 1] - x[t - 3]) = sx[t]
        @d(x[t],1,2) = sx[t]
        #
        log(x[t] - x[t - 2]) - log(x[t - 1] - x[t - 3]) = sx[t]
        @dlog(@d(x[t],0,2)) = sx[t]
        #
        (x[t] + 0.3x[t + 2]) + (x[t - 1] + 0.3x[t + 1]) + (x[t - 2] + 0.3x[t]) = 0
        @movsum(x[t] + 0.3x[t + 2],3) = 0
        #
        ((x[t] + 0.3x[t + 2]) + (x[t - 1] + 0.3x[t + 1]) + (x[t - 2] + 0.3x[t])) / 3 = 0
        @movav(x[t] + 0.3x[t + 2],3) = 0
    end
    @initialize mod

    compare_resids(e1, e2) = (
        e1.resid.head == e2.resid.head  && (
            (length(e1.resid.args) == length(e2.resid.args) == 2 && e1.resid.args[2] == e2.resid.args[2]) ||
            (length(e1.resid.args) == length(e2.resid.args) == 1 && e1.resid.args[1] == e2.resid.args[1])
        )
    )

    for i = 2:2:length(mod.equations)
        @test compare_resids(mod.equations[i - 1], mod.equations[i])
    end
    # test errors and warnings
    mod.warn.no_t = false
    @test  add_equation!(mod, :(x = sx[t])) isa Model
    @test  add_equation!(mod, :(x[t] = sx)) isa Model
    @test  add_equation!(mod, :(x[t] = sx[t])) isa Model
    @test compare_resids(mod.equations[end], mod.equations[end - 1])
    @test compare_resids(mod.equations[end], mod.equations[end - 2])
    @test_throws ArgumentError add_equation!(mod, :(@notametafunction(x[t]) = 7))
    @test_throws ArgumentError add_equation!(mod, :(x[t] = unknownsymbol))
    @test_throws ArgumentError add_equation!(mod, :(x[t] = unknownseries[t]))
    @test_throws ArgumentError add_equation!(mod, :(x[t] = let c = 5; sx[t + c]; end))

end

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

@testset "export" begin
    let m = Model()
        m.warn.no_t = false
        @parameters m begin
            a = 0.3
            b = @link 1 - a
            d = [1,2,3]
            c = @link sin(2π / d[3])
        end
        @variables m begin
            "variable x" x
        end
        @shocks m sx
        @equations m begin
            "This equation is super cool"
            a * @d(x) = b * @d(x[t + 1]) + sx
        end
        @initialize m
        @steadystate m x = a + 1

        export_model(m, "TestModel", "../examples/")

        @test isfile("../examples/TestModel.jl")
        @using_example TestModel

        @test parameters(TestModel.model) == parameters(m)
        @test variables(TestModel.model) == variables(m)
        @test shocks(TestModel.model) == shocks(m)
        @test equations(TestModel.model) == equations(m)
        @test sstate(TestModel.model).constraints == sstate(m).constraints

        @test_throws ArgumentError TestModel.model.parameters.d = @alias c

        rm("../examples/TestModel.jl")
    end
end

@testset "@log eqn" begin
    let m = Model()
        @parameters m rho = 0.1
        @variables m X
        @shocks m EX
        @equations m begin
            @log X[t] = rho * X[t - 1] + EX[t]
        end
        @initialize m
        @test length(m.equations) == 1 && islog(m.equations[1])
    end
end

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

function test_eval_RJ(m::Model, known_R, known_J)
    nrows = 1 + m.maxlag + m.maxlead
    ncols = length(m.allvars)
    R, J = eval_RJ(zeros(nrows, ncols), m)
    @test R ≈ known_R
    @test J ≈ known_J
end

function compare_RJ_R!_(m::Model)
    nrows = 1 + m.maxlag + m.maxlead
    ncols = length(m.variables) + length(m.shocks) + length(m.auxvars)
    point = rand(nrows, ncols)
    R, J = eval_RJ(point, m)
    S = similar(R)
    eval_R!(S, point, m)
    @test R ≈ S
end

@using_example E1
@testset "E1" begin
    @test length(E1.model.parameters) == 2
    @test length(E1.model.variables) == 1
    @test length(E1.model.shocks) == 1
    @test length(E1.model.equations) == 1
    @test E1.model.maxlag == 1
    @test E1.model.maxlead == 1
    test_eval_RJ(E1.model, [0.0], [-0.5 1.0 -0.5 0.0 -1.0 0.0])
    compare_RJ_R!_(E1.model)
    @test E1.model.tol == E1.model.options.tol
    tol = E1.model.tol
    E1.model.tol = tol * 10
    @test E1.model.options.tol == E1.model.tol
    E1.model.tol = tol
    @test E1.model.linear == E1.model.flags.linear
    E1.model.linear = true
    @test E1.model.linear
end

@testset "E1.sstate" begin
    let io = IOBuffer(), m = E1.model
        @test issssolved(m) == false
        E1.model.sstate.mask .= true
        @test issssolved(m) == true
        @test neqns(m.sstate) == 2
        @steadystate m y = 5
        @test_throws ErrorException @steadystate m sin(y + 7)
        @test length(m.sstate.constraints) == 1
        @test neqns(m.sstate) == 3
        @test length(alleqns(m.sstate)) == 3
        printsstate(io, m)
        lines = split(String(take!(io)), '\n')
        @test length(lines) == 2 + length(m.allvars)
    end
end

@testset "E1.lin" begin
    m = deepcopy(E1.model)
    with_linearized(m) do lm
        @test islinearized(lm)
        test_eval_RJ(lm, [0.0], [-0.5 1.0 -0.5 0.0 -1.0 0.0])
        compare_RJ_R!_(lm)
    end
    @test !islinearized(m)
    lm = linearized(m)
    test_eval_RJ(lm, [0.0], [-0.5 1.0 -0.5 0.0 -1.0 0.0])
    compare_RJ_R!_(lm)
    @test islinearized(lm)
    @test !islinearized(m)
    linearize!(m)
    @test islinearized(m)
end

@testset "E1.params" begin
    let m = E1.model
        @test propertynames(m.parameters) == (:α, :β)
        m.β = @link 1.0 - α
        m.parameters.beta = @alias β
        for α = 0.0:0.1:1.0
            m.α = α
            test_eval_RJ(m, [0.0], [-α 1.0 -m.beta 0.0 -1.0 0.0;])
        end
    end
    let io = IOBuffer(), m = E1.model
        show(io, m.parameters)
        @test length(split(String(take!(io)), '\n')) == 1
        show(io, MIME"text/plain"(), m.parameters)
        @test length(split(String(take!(io)), '\n')) == 4
    end
    end

module AUX
using ModelBaseEcon
model = Model()
model.substitutions = true
@variables model x y
@equations model begin
    x[t + 1] = log(x[t] - x[t - 1])
    y[t + 1] = y[t] + log(y[t - 1])
end
@initialize model
end
@testset "AUX" begin
    let m = AUX.model
    @test m.nvars == 2
        @test m.nshks == 0
        @test m.nauxs == 2
        @test length(m.auxeqns) == 2
        x = ones(2, 2)
        @test_throws ErrorException ModelBaseEcon.update_auxvars(x, m)
        x = 2 .* ones(4, 2)
        ax = ModelBaseEcon.update_auxvars(x, m; default=0.1)
        @test size(ax) == (4, 4)
        @test x == ax[:, 1:2]
        @test ax[:, 3:4] == [0.0 0.0; 0.1 log(2.0); 0.1 log(2.0); 0.0 0.0]
    end
    end

@using_example E2
@testset "E2" begin
    @test length(E2.model.parameters) == 3
    @test length(E2.model.variables) == 3
    @test length(E2.model.shocks) == 3
    @test length(E2.model.equations) == 3
    @test E2.model.maxlag == 1
    @test E2.model.maxlead == 1
    test_eval_RJ(E2.model, [0.0, 0.0, 0.0],
        [-.5      1  -.48     0    0  0    0   -.02     0  0  -1  0  0  0 0 0  0 0;
           0  -.375     0  -.75    1  0    0  -.125     0  0   0  0  0 -1 0 0  0 0;
           0      0  -.02     0  .02  0  -.5      1  -.48  0   0  0  0  0 0 0 -1 0])
    compare_RJ_R!_(E2.model)
end

@testset "E2.sstate" begin
    m = E2.model
    ss = m.sstate
    empty!(ss.constraints)
    out = let io = IOBuffer()
        print(io, ss)
        readlines(seek(io, 0))
    end
    @test length(out) == 2
    @steadystate m pinf = rate + 1
    out = let io = IOBuffer()
        print(io, ss)
        readlines(seek(io, 0))
    end
    @test length(out) == 3
    @test length(split(out[end], "=")) == 2
    #
    @test propertynames(ss) == tuple(m.allvars...)
    @test ss.pinf.level == ss.pinf.data[1]
    @test ss.pinf.slope == ss.pinf.data[2]
    ss.pinf.data .= [2.3, 0.7]
    @test ss.values[1:2] == [2.3, 0.7]
    ss.rate.level = 21
    ss.rate.slope = 0.21
    @test ss.rate.level == 21 && ss.rate.slope == 0.21
    @test ss.rate.data == [21, 0.21]
end

@using_example E3
@testset "E3" begin
    @test length(E3.model.parameters) == 3
    @test length(E3.model.variables) == 3
    @test length(E3.model.shocks) == 3
    @test length(E3.model.equations) == 3
    @test E3.model.maxlag == 2
    @test E3.model.maxlead == 3
    compare_RJ_R!_(E3.model)
    test_eval_RJ(E3.model, [0.0, 0.0, 0.0],
        sparse(
            [1, 1, 2, 1, 3, 1, 1, 2, 2, 3,  3,  3,  1,  2,  3,  3,  1,  2,  3],
            [2, 3, 3, 4, 4, 5, 6, 8, 9, 9, 13, 14, 15, 15, 15, 16, 21, 27, 33],
            [-0.5, 1.0, -0.375, -0.3, -0.02, -0.05, -0.05, -0.75, 1.0, 0.02, -0.25,
             -0.25, -0.02, -0.125, 1.0, -0.48, -1.0, -1.0, -1.0],
            3, 36,
        )
    )
end

@using_example E6
@testset "E6" begin
    @test length(E6.model.parameters) == 2
    @test length(E6.model.variables) == 6
    @test length(E6.model.shocks) == 2
    @test length(E6.model.equations) == 6
    @test E6.model.maxlag == 2
    @test E6.model.maxlead == 3
    compare_RJ_R!_(E6.model)
    nt = 1 + E6.model.maxlag + E6.model.maxlead
    test_eval_RJ(E6.model, [-0.0027, -0.0025, 0.0, 0.0, 0.0, 0.0],
        sparse(
            [2, 2, 2, 3, 5, 2, 2, 2, 1, 1, 3, 4, 1, 3, 6, 5, 5, 4, 4, 6, 6, 2, 1],
            [1, 2, 3, 3, 3, 4, 5, 6, 8, 9, 9, 9, 10, 15, 15, 20, 21, 26, 27, 32, 33, 39, 45],
            [-0.1, -0.1, 1.0, -1.0, -1.0, -0.1, -0.1, -0.1, -0.2, 1.0, -1.0, -1.0, -0.2, 1.0,
             -1.0, -1.0, 1.0, -1.0, 1.0, -1.0, 1.0, -1.0, -1.0],
            6, 6 * 8,
        ))
end

@testset "VarTypesSS" begin
    let m = Model()

        m.verbose = !true

        @variables m begin
            p
            @log q
        end
        @equations m begin
            2p[t] = p[t + 1] + 0.1
            q[t] = p[t] + 1
        end
        @initialize m

        # clear_sstate!(m)
        # ret = sssolve!(m)
        # @test ret ≈ [0.1, 0.0, log(1.1), 0.0]

        eq1, eq2, eq3, eq4 = m.sstate.equations
        x = rand(Float64, (4,))
        R, J = eq1.eval_RJ(x[eq1.vinds])
        @test R ≈ x[1] - x[2] - 0.1
        @test J ≈ [1.0, -1.0, 0, 0][eq1.vinds]

        for sh = 1:5
            m.shift = sh
            R, J = eq3.eval_RJ(x[eq3.vinds])
            @test R ≈ x[1] + (sh - 1) * x[2] - 0.1
            @test J ≈ [1.0, sh - 1.0, 0, 0][eq3.vinds]
        end

        R, J = eq2.eval_RJ(x[eq2.vinds])
        @test R ≈ exp(x[3]) - x[1] - 1
        @test J ≈ [-1, 0.0, exp(x[3]), 0.0][eq2.vinds]

        for sh = 1:5
            m.shift = sh
            R, J = eq4.eval_RJ(x[eq4.vinds])
            @test R ≈ exp(x[3] + sh * x[4]) - x[1] - sh * x[2] - 1
            @test J ≈ [-1.0, -sh, exp(x[3] + sh * x[4]), exp(x[3] + sh * x[4]) * sh][eq4.vinds]
        end

    end

    let m = Model()
        @variables m begin
            lx
            @log x
        end
        @shocks m s1 s2
        @equations m begin
            "linear growth with slope 0.2"
            lx[t] = lx[t - 1] + 0.2 + s1[t]
            "exponential with the same rate as the slope of lx"
            log(x[t]) = lx[t] + s2[t + 1]
        end
        @initialize m
        #
        @test nvariables(m) == 2
        @test nshocks(m) == 2
        @test nequations(m) == 2
        ss = sstate(m)
        @test neqns(ss) == 4
        eq1, eq2, eq3, eq4 = ss.equations
        @test length(ss.values) == 2 * length(m.allvars)
        #
        # test with eq1
        ss.lx.data .= [1.5, 0.2]
        ss.x.data .= [0.0, 0.2]
        ss.s1.data .= [0.0, 0.0]
        ss.s2.data .= [0.0, 0.0]
        for s1 = -2:0.1:2
            ss.s1.level = s1
            @test eq1.eval_resid(ss.values[eq1.vinds]) ≈ -s1
        end
        ss.s1.level = 0.0
        for lxslp = -2:0.1:2
            ss.lx.slope = lxslp
            @test eq1.eval_resid(ss.values[eq1.vinds]) ≈ lxslp - 0.2
        end
        ss.lx.slope = 0.2
        R, J = eq1.eval_RJ(ss.values[eq1.vinds])
        TMP = fill!(similar(ss.values), 0.0)
        TMP[eq1.vinds] .= J
        @test R == 0
        @test TMP[[1,2,5]] ≈ [0.0, 1.0, -1.0]
        # test with eq4
        ss.lx.data .= [1.5, 0.2]
        ss.x.data .= [1.5, 0.2]
        ss.s1.data .= [0.0, 0.0]
        ss.s2.data .= [0.0, 0.0]
        for s2 = -2:0.1:2
            ss.s2.level = s2
            @test eq4.eval_resid(ss.values[eq4.vinds]) ≈ -s2
        end
        ss.s2.level = 0.0
        for lxslp = -2:0.1:2
            ss.lx.slope = lxslp
            @test eq4.eval_resid(ss.values[eq4.vinds]) ≈ m.shift * (0.2 - lxslp)
        end
        ss.lx.slope = 0.2
        for xslp = -2:0.1:2
            ss.x.data[2] = xslp
            @test eq4.eval_resid(ss.values[eq4.vinds]) ≈ m.shift * (xslp - 0.2)
        end
        ss.x.slope = exp(0.2)
        R, J = eq4.eval_RJ(ss.values[eq4.vinds])
        TMP = fill!(similar(ss.values), 0.0)
        TMP[eq4.vinds] .= J
        @test R+1.0 ≈ 0.0+1.0
        @test TMP[[1,2,3,4,7]] ≈ [-1.0, -m.shift, 1.0, m.shift, -1.0]
        for xlvl = 0.1:0.1:2
            ss.x.level = exp(xlvl)
            R, J = eq4.eval_RJ(ss.values[eq4.vinds])
            @test R ≈ xlvl - 1.5
            TMP[eq4.vinds] .= J
            @test TMP[[1,2,3,4,7]] ≈ [-1.0, -m.shift, 1.0, m.shift, -1.0]
        end
    end
end

ext: jl


text:
# We're currently reusing this type name to minimize the number of changes from pdmat.jl
CholType{T, S<:AbstractMatrix} = LinearAlgebra.CholeskyPivoted{T, S}

#=
The code below is just a slight modification of the PDMat code in PDMats.jl
https://github.com/JuliaStats/PDMats.jl/blob/master/src/pdmat.jl

TODO: Integrate this with PDMats.jl

References for discussion on supporting degenerate mvnormal distributions:
- https://discourse.julialang.org/t/multivariate-normal-with-positive-semi-definite-covariance-matrix/3029
- https://github.com/JuliaStats/Distributions.jl/issues/366
=#

"""
Positive semi-definite matrix together with a CholeskyPivoted factorization object.
"""
struct PSDMat{T<:Real,S<:AbstractMatrix} <: AbstractPDMat{T}
    dim::Int
    mat::S
    chol::CholType{T, S}

    PSDMat{T,S}(d::Int,m::AbstractMatrix{T},c::CholType{T,S}) where {T,S} = new{T, S}(d, m, c)
end

function PSDMat(mat::AbstractMatrix, chol::CholType)
    d = size(mat, 1)
    size(chol, 1) == d ||
        throw(DimensionMismatch("Dimensions of mat and chol are inconsistent."))
    PSDMat{eltype(mat),typeof(mat)}(d, mat, chol)
end

PSDMat(mat::Matrix) = PSDMat(mat, cholesky(mat, Val(true); check=false))
PSDMat(mat::Symmetric) = PSDMat(Matrix(mat))
PSDMat(fac::CholType) = PSDMat(Matrix(fac), fac)

### Conversion
Base.convert(::Type{PSDMat{T}}, a::PSDMat) where {T<:Real} = PSDMat(convert(AbstractArray{T}, a.mat))
Base.convert(::Type{AbstractArray{T}}, a::PSDMat) where {T<:Real} = convert(PSDMat{T}, a)

### Basics

PDMats.dim(a::PSDMat) = a.dim
Base.Matrix(a::PSDMat) = copy(a.mat)
Base.getindex(a::PSDMat, i::Int) = getindex(a.mat, i)
Base.getindex(a::PSDMat, I::Vararg{Int, N}) where N = getindex(a.mat, I...)
LinearAlgebra.diag(a::PSDMat) = diag(a.mat)

### Arithmetics

function PDMats.pdadd!(r::Matrix, a::Matrix, b::PSDMat, c)
    PDMats.@check_argdims size(r) == size(a) == size(b)
    PDMats._addscal!(r, a, b.mat, c)
end

*(a::PSDMat{S}, c::T) where {S<:Real, T<:Real} = PSDMat(a.mat * c)
*(a::PSDMat, x::StridedVecOrMat) = a.mat * x
\(a::PSDMat, x::StridedVecOrMat) = a.chol \ x

### Algebra

Base.inv(a::PSDMat) = PSDMat(inv(a.chol))
LinearAlgebra.logdet(a::PSDMat) = logdet(a.chol)
LinearAlgebra.eigmax(a::PSDMat) = eigmax(a.mat)
LinearAlgebra.eigmin(a::PSDMat) = eigmin(a.mat)
LinearAlgebra.kron(a::PSDMat, b::PSDMat) = PSDMat(kron(a.mat, b.mat))
LinearAlgebra.kron(a::AbstractPDMat, b::PSDMat) = PSDMat(kron(Matrix(a), b.mat))
LinearAlgebra.kron(a::PSDMat, b::AbstractPDMat) = PSDMat(kron(a.mat, Matrix(b)))

### whiten and unwhiten

function PDMats.whiten!(r::StridedVecOrMat, a::PSDMat, x::StridedVecOrMat)
    cf = a.chol.U
    v = PDMats._rcopy!(r, x)
    istriu(cf) ? ldiv!(transpose(cf), v) : ldiv!(cf, v)
end

function PDMats.unwhiten!(r::StridedVecOrMat, a::PSDMat, x::StridedVecOrMat)
    cf = a.chol.U
    v = PDMats._rcopy!(r, x)
    istriu(cf) ? lmul!(transpose(cf), v) : lmul!(cf, v)
end

### quadratic forms

PDMats.quad(a::PSDMat, x::StridedVector) = dot(x, a * x)
PDMats.invquad(a::PSDMat, x::StridedVector) = dot(x, a \ x)

"""
    quad!(r::AbstractArray, a::AbstractPDMat, x::StridedMatrix)
Overwrite `r` with the value of the quadratic form defined by `a` applied columnwise to `x`
"""
PDMats.quad!(r::AbstractArray, a::PSDMat, x::StridedMatrix) = PDMats.colwise_dot!(r, x, a.mat * x)

"""
    invquad!(r::AbstractArray, a::AbstractPDMat, x::StridedMatrix)
Overwrite `r` with the value of the quadratic form defined by `inv(a)` applied columnwise to `x`
"""
PDMats.invquad!(r::AbstractArray, a::PSDMat, x::StridedMatrix) = PDMats.colwise_dot!(r, x, a.mat \ x)

### tri products

function X_A_Xt(a::PSDMat, x::StridedMatrix)
    z = copy(x)
    cf = a.chol.UL
    rmul!(z, istriu(cf) ? transpose(cf) : cf)
    z * transpose(z)
end

function Xt_A_X(a::PSDMat, x::StridedMatrix)
    cf = a.chol.UL
    z = lmul!(istriu(cf) ? cf : transpose(cf), copy(x))
    transpose(z) * z
end

function X_invA_Xt(a::PSDMat, x::StridedMatrix)
    cf = a.chol.UL
    z = rdiv!(copy(x), istriu(cf) ? cf : transpose(cf))
    z * transpose(z)
end

function Xt_invA_X(a::PSDMat, x::StridedMatrix)
    cf = a.chol.UL
    z = ldiv!(istriu(cf) ? transpose(cf) : cf, copy(x))
    transpose(z) * z
end


text:
using DataStructures: MutableBinaryMaxHeap, top_with_handle, delete!

function proj_simplex0(v::Array{T, 1}, s::T=one(T)) where {T<:Real}
    n = length(v)
    if sum(v) == s && all(v .≥ 0)
        w = v
    else
        u = sort(v, rev=true)
        cssv = cumsum(u)
        ρ = findall(>(0), u .* Array((1:n)) .> (cssv .- s))[end]
        θ = (cssv[ρ] - s) / (ρ)
        w = max.(v .- θ, zero(T))
    end
    return w
end

""" Algorithm from Held, M., Wolfe, P., Crowder, H.: 'Validation of
subgradient optimization'. The description is from the Condat L.:
'Fast Projection onto the Simplex and the l1 Ball'. (Algorithm 1)"""

function proj_simplex1(y::Array{T, 1}, a::T=one(T)) where {T<:Real}
    n = length(y)
    if sum(y) == a && all(y .≥ 0)
        x = y
    else
        let τ
            u = sort(y, rev=true)
            cumsum_u = zero(T)
            for k in 1:n
                if cumsum_u + u[k] < k * u[k] + a
                    cumsum_u += u[k]
                    τ = (cumsum_u - a) / k
                else
                    break
                end
            end
            x = max.(y .- τ, zero(T))
        end
    end
    return x
end

function proj_simplex12(y::Array{T, 1}, a::T=one(T)) where {T<:Real}
    n = length(y)
    if sum(y) == a && all(y .≥ 0)
        return y
    else
        τ = 0.0
        u = sort(y, rev=true)
        cumsum_u = zero(T)
        for k in 1:n
            if cumsum_u + u[k] < k * u[k] + a
                cumsum_u += u[k]
                τ = (cumsum_u - a) / k
            else
                break
            end
        end
        y = max.(y .- τ, zero(T))
    end
    return y
end

""" Algorithm from van den Berg, E., Friedlander, M.P.: 'Probing the
Pareto frontier for basis pursuit solution'. The description is from
Condat L: 'Fast Projection onto the Simplex and the l1
Ball'. (Algorithm 2)"""

function proj_simplex2(y::Array{T, 1}, a::T=one(T)) where {T<:Real}
    N = length(y)
    if sum(y) == a && all(y .≥ 0)
        x = y
    else
        τ = zero(T)
        v = MutableBinaryMaxHeap(y)
        cumsum_u = zero(T)
        for k in 1:N
            u = first(v)
            if cumsum_u + u < k * u + a
                cumsum_u += u
                i = top_with_handle(v)[2]
                delete!(v, i)
                τ = (cumsum_u - a) / k
            else
                break
            end
        end
        x = max.(y .- τ, zero(T))
    end
    return x
end

function proj_simplex22(y::Array{Float64, 1}, a=1.0)
    N = length(y)
    if sum(y) == a && all(y .≥ 0)
        x = y
    else
        τ = 0.
        v = MutableBinaryMaxHeap(y)
        cumsum_u = 0.
        for k in 1:N
            u = first(v)
            if cumsum_u + u < k * u + a
                cumsum_u += u
                i = top_with_handle(v)[2]
                delete!(v, i)
                τ = (cumsum_u - a) / k
            else
                break
            end
        end
        x = max.(y .- τ, 0.)
    end
    return x
end

"""
From Condat “Fast projection onto the simplex and the l_1 ball”. In: Mathematical Programming 158.1
(2016), pp. 575–585
"""
function proj_simplex_condat(y::Array{Float64, 1}, a=Float64(1.0))
    N = length(y)
    v = [y[1]]
    v_tilde = Float64[]
    ρ = y[1] - a
    for n in 2:N
        yn = y[n]
        if yn > ρ
            ρ += (yn - ρ) / (length(v) + 1)
            if ρ > yn - a
                append!(v, y[n])
            else
                append!(v_tilde, v)
                v = [yn]
                ρ = yn - a
            end
        end
    end

    if !isempty(v_tilde)
        for yi in v_tilde
            if yi > ρ
                append!(v, yi)
                ρ += (yi - ρ) / length(v)
            end
        end
    end
    # if during the loop rho is increases at least once, then the flag is true. Otherwise we stop
    flag = true
    while flag
        flag = false
        for (i, yi) in enumerate(v)
            if yi ≤ ρ
                deleteat!(v, i)
                ρ += (ρ - yi) / length(v)
                flag = true
            end
        end
    end
    τ = ρ
    x = max.(y .- τ, zero(y))
    return x
end

function proj_simplex3(v, z=1.)
    n = length(v)
    U = Array((1:n))
    s = 0
    ρ = 0
    while length(U) > 0
        G = []
        L = []
        k = U[rand(1:length(U))]
        ds = v[k]
        for j in U
            if v[j] >= v[k]
                if j != k
                    ds += v[j]
                    append!(G, j)
                end
            elseif v[j] < v[k]
                append!(L, j)
            end
        end
        drho = length(G) + 1
        if s + ds - (ρ + drho) * v[k] < z
            s += ds
            ρ += drho
            U = L
        else
            U = G
        end
    end
    theta = (s - z) / ρ
    return max.(v .- theta, 0.)

end

function proj_simplex4(v, z=1., τ=1e-7, max_iter=1000)
    lower = 0
    upper = maximum(v)
    current = Inf
    w = zeros(n)
    for it in 1:max_iter
        if abs(current) / z < τ && current < 0.
            break
        end
        theta = (upper + lower) / 2.0
        w = max.(v .- theta, 0.)
        current = sum(w) - z
        if current <= 0.
            upper = theta
        else
            lower = theta
        end
    end
    return w
end

function softmax(x::Array{T, 1}) where {T<:Real}
    # compare with NNlib implementation
    res = exp.(x .- maximum(x))
    res ./= sum(res)
    return res
end


text:
"""
Solve any non-conjugate likelihood using Variational Inference
by making a numerical approximation (quadrature or MC integration)
of the expected log-likelihood ad its gradients
Gradients are computed as in "The Variational Gaussian Approximation
Revisited" by Opper and Archambeau 2009
"""
abstract type NumericalVI{T<:Real} <: VariationalInference{T} end

include("quadratureVI.jl")
include("MCVI.jl")

""" `NumericalVI(integration_technique::Symbol=:quad;ϵ::T=1e-5,nMC::Integer=1000,nGaussHermite::Integer=20,optimizer::Optimizer=Adam(α=0.1))`

General constructor for Variational Inference via numerical approximation.

**Argument**

    -`integration_technique::Symbol` : Method of approximation can be `:quad` for quadrature see [QuadratureVI](@ref) or `:mc` for MC integration see [MCIntegrationVI](@ref)

**Keyword arguments**

    - `ϵ::T` : convergence criteria, which can be user defined
    - `nMC::Int` : Number of samples per data point for the integral evaluation (for the MCIntegrationVI)
    - `nGaussHermite::Int` : Number of points for the integral estimation (for the QuadratureVI)
    - `optimizer::Optimizer` : Optimizer used for the variational updates. Should be an Optimizer object from the [GradDescent.jl]() package. Default is `Adam()`
"""
function NumericalVI(integration_technique::Symbol=:quad;ϵ::T=1e-5,nMC::Integer=1000,nGaussHermite::Integer=20,optimizer::Optimizer=Momentum(η=1e-5)) where {T<:Real}
    if integration_technique == :quad
        QuadratureVI{T}(ϵ,nGaussHermite,optimizer,false,0)
    elseif integration_technique == :mc
        MCIntegrationVI{T}(ϵ,nMC,optimizer,false,0)
    else
        @error "Only possible integration techniques are quadrature : :quad or mcmc integration :mcmc"
    end
end

""" `NumericalSVI(integration_technique::Symbol=:quad;ϵ::T=1e-5,nMC::Integer=1000,nGaussHermite::Integer=20,optimizer::Optimizer=Adam(α=0.1))`

General constructor for Stochastic Variational Inference via numerical approximation.

**Argument**

    -`nMinibatch::Integer` : Number of samples per mini-batches
    -`integration_technique::Symbol` : Method of approximation can be `:quad` for quadrature see [QuadratureVI](@ref) or `:mc` for MC integration see [MCIntegrationVI](@ref)

**Keyword arguments**

    - `ϵ::T` : convergence criteria, which can be user defined
    - `nMC::Int` : Number of samples per data point for the integral evaluation (for the MCIntegrationVI)
    - `nGaussHermite::Int` : Number of points for the integral estimation (for the QuadratureVI)
    - `optimizer::Optimizer` : Optimizer used for the variational updates. Should be an Optimizer object from the [GradDescent.jl]() package. Default is `Adam()`
"""
function NumericalSVI(nMinibatch::Integer,integration_technique::Symbol=:quad;ϵ::T=1e-5,nMC::Integer=200,nGaussHermite::Integer=20,optimizer::Optimizer=Momentum(η=1e-5)) where {T<:Real}
    if integration_technique == :quad
        QuadratureVI{T}(ϵ,nGaussHermite,optimizer,true,nMinibatch)
    elseif integration_technique == :mc
        MCIntegrationVI{T}(ϵ,nMC,optimizer,true,nMinibatch)
    else
        @error "Only possible integration techniques are quadrature : :quad or mcmc integration :mc"
    end
end

function Base.show(io::IO,inference::NumericalVI{T}) where T
    print(io,"$(inference.Stochastic ? "Stochastic numerical" : "Numerical") inference by $(isa(inference,MCIntegrationVI) ? "Monte Carlo Integration" : "Quadrature")")
end

∇E_μ(::Likelihood,i::NVIOptimizer,::AbstractVector) = -i.ν
∇E_Σ(::Likelihood,i::NVIOptimizer,::AbstractVector) = 0.5.*i.λ

function variational_updates!(model::AbstractGP{T,L,<:NumericalVI}) where {T,L}
    compute_grad_expectations!(model)
    natural_gradient!.(model.likelihood,model.inference,model.inference.vi_opt,[get_y(model)],model.f)
    global_update!(model)
end

function natural_gradient!(l::Likelihood,i::NumericalVI,opt::NVIOptimizer,y::AbstractVector,X::AbstractMatrix,gp::_VGP{T}) where {T,L}
    opt.∇η₂ .= Symmetric(Diagonal(∇E_Σ(l,opt,y)) - 0.5*inv(gp.K) - gp.η₂)
    opt.∇η₁ .= ∇E_μ(l,opt,y) - gp.K \ (gp.μ - gp.μ₀(X)) - 2 * opt.∇η₂ * gp.μ
end

function natural_gradient!(l::Likelihood,i::NumericalVI,opt::NVIOptimizer,y::AbstractVector,Z::AbstractMatrix,gp::_SVGP{T}) where {T,L}
    opt.∇η₂ .= Symmetric(i.ρ*transpose(gp.κ)*Diagonal(∇E_Σ(l,opt,y))*gp.κ - 0.5*inv(gp.K) - gp.η₂)
    opt.∇η₁ .= i.ρ * transpose(gp.κ) * ∇E_μ(l,opt,y) - gp.K \ (gp.μ - gp.μ₀(Z)) - 2 * opt.∇η₂ * gp.μ
end

function global_update!(model::AbstractGP{T,L,<:NumericalVI}) where {T,L}
    for (gp,opt) in zip(model.f,model.inference.vi_opt)
        Δ = update(opt.optimizer,vcat(opt.∇η₁,vec(LowerTriangular(-2*opt.∇η₂*opt.L))))
        Δ₁ = Δ[1:model.nFeatures]
        Δ₂ = LowerTriangular(reshape(Δ[model.nFeatures+1:end],model.nFeatures,model.nFeatures))
        display(ELBO(model))
        α = 1.0
        # # Loop to verify update keeps positive definiteness
        while !isposdef(Symmetric((opt.L+α*Δ₂)*(opt.L+α*Δ₂)')) &&  α > 1e-7
            α *= 0.1
        end
        if α < 1e-7
            @error "α too small, positive definiteness could not be achieved"
        end
        opt.L = LowerTriangular(opt.L+α*Δ₂)
        gp.η₂ .= Symmetric(-opt.L*opt.L')
        display(gp.η₁)
        # display(opt.optimizer.η)
        display(Δ₁)

        gp.η₁ .+= α*Δ₁
        ## Passed the pos. def. test, now update parameters
        # gp.η₂ = Symmetric(gp.η₂+α*Δ₂)
        # gp.η₁ .+= α*Δ₁

        ## Save the new scaling on the optimizer
        if isa(opt.optimizer,Adam)
            opt.optimizer.α = min(opt.optimizer.α * α* 2.0,0.1)
        elseif isa(opt.optimizer,Union{VanillaGradDescent,Momentum,RMSprop})
            opt.optimizer.η = min(opt.optimizer.η*α*1.1,0.1)
        elseif isa(opt.optimizer,ALRSVI)
        elseif isa(opt.optimizer,InverseDecay)
        end

        ## Reparametrize to the normal distribution
        global_update!.(model.f)
    end
end

function ELBO(model::AbstractGP{T,L,<:NumericalVI}) where {T,L}
    return model.inference.ρ*expec_log_likelihood(model.likelihood,model.inference,get_y(model),mean_f(model),diag_cov_f(model)) - GaussianKL(model)
end


text:
function get_component_to_pm(ix::Int, load::PSY.StaticLoad)
    PM_load = Dict{String, Any}(
        "source_id" => ["bus", PSY.get_name(PSY.get_bus(load))],
        "load_bus" => PSY.get_number(PSY.get_bus(load)),
        "status" => Int(PSY.get_available(load)),
        "qd" => PSY.get_reactive_power(load), # TODO: get Q load from time series
        "pd" => PSY.get_active_power(load),
        "index" => ix,
    )
    return PM_load
end

function get_component_to_pm(ix::Int, shunt::PSY.FixedAdmittance)
    PM_shunt = Dict{String, Any}(
        "source_id" => ["bus", PSY.get_name(PSY.get_bus(shunt))],
        "shunt_bus" => PSY.get_number(PSY.get_bus(shunt)),
        "status" => Int(PSY.get_available(shunt)),
        "gs" => real(PSY.get_Y(shunt)),
        "bs" => imag(PSY.get_Y(shunt)),
        "index" => ix,
    )
    return PM_shunt
end

function get_time_series_to_pm!(
    pm_data::Dict{String, Any},
    pm_category::String,
    pm_id::String,
    device::PSY.StaticLoad,
    start_time::Dates.DateTime,
    time_periods::UnitRange{Int},
)
    psy_forecast_name = "max_active_power" # change this line for different forecasts
    pm_field_name = "pd" # change this line to apply forecast to different fields

    ts_data = PSY.get_time_series_values(
        PSY.AbstractDeterministic,
        device,
        psy_forecast_name,
        start_time = start_time,
    )
    if haskey(pm_data, "nw")
        pm_update = Dict{String, Any}("nw" => Dict{String, Any}())
        for (t, nw) in pm_data["nw"]
            pm_update["nw"][t] = Dict(
                pm_category => Dict(
                    pm_id =>
                        Dict(pm_field_name => ts_data[time_periods[parse(Int, t)]]),
                ),
            )
        end
    else
        throw(@error "cannot apply time series to single period data")
    end

    PM.update_data!(pm_data, pm_update)
end

function get_time_series_to_pm!(
    pm_data::Dict{String, Any},
    pm_category::String,
    pm_id::String,
    device::PSY.StaticLoad,
    start_time::Dates.DateTime,
    time_period::Int,
)
    psy_forecast_name = "max_active_power" # change this line for different forecasts
    pm_field_name = "pd" # change this line to apply forecast to different fields

    ts_data = PSY.get_time_series_values(
        PSY.AbstractDeterministic,
        device,
        psy_forecast_name,
        start_time = start_time,
    )
    if haskey(pm_data, "nw")
        throw(@error "cannot apply single time period data to multi-network")
    else
        pm_update = Dict{String, Any}(
            pm_category => Dict(pm_id => Dict(pm_field_name => ts_data[time_period])),
        )
    end

    PM.update_data!(pm_data, pm_update)
end

function get_time_series_to_pm!(
    pm_data::Dict{String, Any},
    pm_category::String,
    pm_id::String,
    device::PSY.FixedAdmittance,
    start_time::Dates.DateTime,
    time_periods::UnitRange{Int},
)
    psy_forecast_name = "max_active_power" # change this line for different forecasts
    pm_field_name = "qd" # change this line to apply forecast to different fields

    ts_data = PSY.get_time_series_values(
        PSY.AbstractDeterministic,
        device,
        psy_forecast_name,
        start_time = start_time,
    )
    if haskey(pm_data, "nw")
        pm_update = Dict{String, Any}("nw" => Dict{String, Any}())
        for (t, nw) in pm_data["nw"]
            pm_update["nw"][t] = Dict(
                pm_category => Dict(
                    pm_id =>
                        Dict(pm_field_name => ts_data[time_periods[parse(Int, t)]]),
                ),
            )
        end
    else
        throw(@error "cannot apply time series to single period data")
    end

    PM.update_data!(pm_data, pm_update)
end

function get_time_series_to_pm!(
    pm_data::Dict{String, Any},
    pm_category::String,
    pm_id::String,
    device::PSY.FixedAdmittance,
    start_time::Dates.DateTime,
    time_period::Int,
)
    psy_forecast_name = "max_active_power" # change this line for different forecasts
    pm_field_name = "qd" # change this line to apply forecast to different fields

    ts_data = PSY.get_time_series_values(
        PSY.AbstractDeterministic,
        device,
        psy_forecast_name,
        start_time = start_time,
    )
    if haskey(pm_data, "nw")
        throw(@error "cannot apply single time period data to multi-network")
    else
        pm_update = Dict{String, Any}(
            pm_category => Dict(pm_id => Dict(pm_field_name => ts_data[time_period])),
        )
    end

    PM.update_data!(pm_data, pm_update)
end

max_stars_repo_name: NREL-SIIP/PowerModelsInterface

max_issues_repo_name: NREL-SIIP/PowerModelsInterface

max_forks_repo_name: NREL-SIIP/PowerModelsInterface


text:
using MoireHamiltonian
using Test

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


text:
function worker_bits()
  ws = nextpow2(num_threads())
  ifelse(Static.lt(ws,StaticInt{8}()), StaticInt{8}(), ws)
end
function worker_mask_count()
  bits = worker_bits()
  (bits + StaticInt{63}()) ÷ StaticInt{64}() # cld not defined on `StaticInt`
end
worker_size() = worker_bits() ÷ worker_mask_count()

_mask_type(::StaticInt{8}) = UInt8
_mask_type(::StaticInt{16}) = UInt16
_mask_type(::StaticInt{32}) = UInt32
_mask_type(::StaticInt{64}) = UInt64
worker_type() = _mask_type(worker_size())
worker_pointer_type() = Ptr{worker_type()}

worker_pointer() = Base.unsafe_convert(worker_pointer_type(), pointer_from_objref(WORKERS))

function free_threads!(freed_threads)
  _atomic_or!(worker_pointer(), freed_threads)
  nothing
end
@inline _remaining(x::Tuple) = Base.tail(x)
@inline _remaining(@nospecialize(x)) = nothing
@inline _first(x::Tuple{}) = nothing
@inline _first(x::Tuple{X,Vararg}) where {X<:Unsigned} = getfield(x,1)
@inline _first(x::Union{Unsigned,Nothing}) = x
@inline function _request_threads(num_requested::UInt32, wp::Ptr, ::StaticInt{N}, threadmask) where {N}
  ui, ft, num_requested, wp = __request_threads(num_requested, wp, _first(threadmask))
  uit, ftt = _request_threads(num_requested, wp, StaticInt{N}()-StaticInt{1}(), _remaining(threadmask))
  (ui, uit...), (ft, ftt...)
end
@inline function _request_threads(num_requested::UInt32, wp::Ptr, ::StaticInt{1}, threadmask)
  ui, ft, num_requested, wp = __request_threads(num_requested, wp, _first(threadmask))
  (ui, ), (ft, )
end
@inline function _exchange_mask!(wp, ::Nothing)
  all_threads = _atomic_xchg!(wp, zero(worker_type()))
  all_threads, all_threads
end
@inline function _exchange_mask!(wp, threadmask::Unsigned)
  all_threads = _atomic_xchg!(wp, zero(worker_type()))
  saved = all_threads & (~(threadmask%worker_type()))
  _atomic_store!(wp, saved)
  all_threads | saved, all_threads & threadmask
end
@inline function __request_threads(num_requested::UInt32, wp::Ptr, threadmask)
  no_threads = zero(worker_type())
  if (num_requested ≢ StaticInt{-1}()) && (num_requested % Int32 ≤ zero(Int32))
    return UnsignedIteratorEarlyStop(zero(worker_type()), 0x00000000), no_threads, 0x00000000, wp
  end
  # to get more, we xchng, setting all to `0`
  # then see which we need, and free those we aren't using.
  wpret = wp + 8 # (worker_type() === UInt64) | (worker_mask_count() === StaticInt(1)) #, so adding 8 is fine.
  # _all_threads = all_threads = _apply_mask(_atomic_xchg!(wp, no_threads), threadmask)
  _all_threads, all_threads = _exchange_mask!(wp, threadmask)
  additional_threads = count_ones(all_threads) % UInt32
  # num_requested === StaticInt{-1}() && return reserved_threads, all_threads
  if num_requested === StaticInt{-1}()
    return UnsignedIteratorEarlyStop(all_threads), all_threads, num_requested, wpret
  end
  nexcess = num_requested - additional_threads
  if signed(nexcess) ≥ 0
    return UnsignedIteratorEarlyStop(all_threads), all_threads, nexcess, wpret
  end
  # we need to return the `excess` to the pool.
  lz = leading_zeros(all_threads) % UInt32
  while true
    # start by trying to trim off excess from lz
    lz += (-nexcess)%UInt32
    m = (one(worker_type()) << (UInt32(last(worker_size())) - lz)) - one(worker_type())
    masked = (all_threads & m) ⊻ all_threads
    nexcess += count_ones(masked) % UInt32
    all_threads &= (~masked)
    nexcess == zero(nexcess) && break
  end
  _atomic_store!(wp, _all_threads & (~all_threads))
  return UnsignedIteratorEarlyStop(all_threads, num_requested), all_threads, 0x00000000, wpret
end

@inline function request_threads(num_requested, threadmask)
  _request_threads(num_requested % UInt32, worker_pointer(), worker_mask_count(), threadmask)
end
@inline request_threads(num_requested) = request_threads(num_requested, nothing)


text:
cd(@__DIR__); import Pkg; Pkg.activate(".")
coords = (left = -145.0, right = -50.0, bottom = 20.0, top = 75.0)
Pkg.develop(path="$(homedir())/github/SimpleSDMLayers.jl")
using Revise
using SimpleSDMLayers
using Test

# Import QGIS-cropped layer as example
ref = geotiff(SimpleSDMPredictor, "../../rasters/wc1_cropped.tif")

## Step 1: Fix getindex ####

layer = SimpleSDMPredictor(WorldClim, BioClim, 1)
layer_subset = layer[coords]
size(layer_subset) # should be 330, 570

left, right, bottom, top = coords

import SimpleSDMLayers: _match_latitude, _match_longitude

# Inside getindex call
imax = _match_longitude(layer, isnothing(right) ? layer.right : right)
imin = _match_longitude(layer, isnothing(left) ? layer.left : left)
jmax = _match_latitude(layer, isnothing(top) ? layer.top : top)
jmin = _match_latitude(layer, isnothing(bottom) ? layer.bottom : bottom)
test1a = layer[jmin:jmax, imin:imax]
test1b = layer[jmin+1:jmax-1, imin+1:imax] # almost same as ref
ref

isequal(ref.grid, test1b.grid) # equal

stride(ref)
stride(layer) # original stride ok
stride(test1a) # stride is wrong...
stride(test1b) # still wrong

longitudes(ref)
longitudes(test1a)
longitudes(test1b)
isequal(collect(longitudes(ref)), collect(longitudes(test1b)))
isequal(collect(longitudes(ref)), collect(longitudes(test1b)))

latitudes(test1)
latitudes(ref)

# Check functions 1 by 1
# getindex(layer; coords...) # should be ok in theory
# _match_longitude # could be the problem
_match_longitude(layer, isnothing(right) ? layer.right : right) # ok
_match_longitude(layer, isnothing(left) ? layer.left : left) # should be 211
lon = left
last(findmin(abs.(lon .- longitudes(layer))))
@time lon .- longitudes(layer)
@time abs.(lon .- longitudes(layer))
@time findmin(abs.(lon .- longitudes(layer)))
@time last(findmin(abs.(lon .- longitudes(layer))))

isapprox(abs.(lon .- longitudes(layer))[[210, 211]]...)
findall(x -> x ≈ stride(layer,1), abs.(lon .- longitudes(layer)))

@time filter(x -> lon-stride(layer,1) <= x <= lon+stride(layer, 1), longitudes(layer))
i1, i2 = findall(x -> lon-stride(layer,1) <= x <= lon+stride(layer, 1), longitudes(layer))
(lon .- longitudes(layer))[[i1, i2]] .|> abs
isapprox.(stride(layer,1), abs.((lon .- longitudes(layer))[[i1, i2]]))

@time any(x -> isapprox(x, left), longitudes(layer) .- stride(layer,1))
findall(x -> isapprox(x, left), longitudes(layer) .- stride(layer,1))
findall(x -> isapprox(x, left), longitudes(layer) .+ stride(layer,1))
findall(x -> isapprox(x, left+0.001), longitudes(layer) .- stride(layer,1))
findall(x -> isapprox(x, right), longitudes(layer) .- stride(layer,1))

findall(x -> isapprox(x, bottom), latitudes(layer) .- stride(layer,1))
findall(x -> isapprox(x, top), latitudes(layer) .+ stride(layer,1))

@time LinRange(layer.bottom+stride(layer, 2), layer.top-stride(layer, 2), size(layer,1))
@time range(layer.bottom+stride(layer, 2), layer.top-stride(layer, 2), length=size(layer,1))

# As previously
@time _match_longitude(layer, isnothing(left) ? layer.left : left) # 25 alloc
# with return l
@time _match_longitude(layer, isnothing(left) ? layer.left : left) # 25 alloc
# with if else
@time _match_longitude(layer, isnothing(left) ? layer.left : left) # 25 alloc
@time _match_longitude(layer, isnothing(left) ? layer.left : left; side=:none) # 25 alloc
@time _match_longitude(layer, isnothing(left) ? layer.left : left; side=:left) # 25 alloc
@time _match_longitude(layer, isnothing(left) ? layer.left : left; side=:youppi) # 25 alloc
# with any statement
ldiff = abs.(lon .- longitudes(layer))
@time any(x -> isapprox(x, stride(layer,1)), ldiff)
@time findlast(x -> isapprox(x, stride(layer,1)), ldiff)
@time _match_longitude(layer, isnothing(left) ? layer.left : left) # 26 alloc
@time _match_longitude(layer, isnothing(left) ? layer.left : left; side=:none) # 26 alloc
@time _match_longitude(layer, isnothing(left) ? layer.left : left; side=:left) # 15.15k alloc
@time _match_longitude(layer, isnothing(right) ? layer.right : right; side=:right) # 10.95k alloc
@time _match_longitude(layer, isnothing(left+0.001) ? layer.left : left+0.001; side=:left) # 15.15k alloc
@time _match_longitude(layer, isnothing(left-0.001) ? layer.left : left-0.001) # 28 alloc
@time _match_longitude(layer, layer.right) # 28 alloc
@time _match_longitude(layer, layer.right; side=:right) # 28 alloc
@time _match_longitude(layer, layer.right; side=:left) # 28 alloc
# Test full fix
_match_longitude(layer, isnothing(right) ? layer.right : right; side=:right)
_match_longitude(layer, isnothing(left) ? layer.left : left; side=:left)
_match_latitude(layer, isnothing(top) ? layer.top : top; side=:top)
_match_latitude(layer, isnothing(bottom) ? layer.bottom : bottom; side=:bottom)
@time layer[coords] # 0.002510 seconds (47.69 k allocations: 1.676 MiB) (now)
@time layer[coords] # 0.001570 seconds (222 allocations: 981.422 KiB) (before)

ref
isequal(layer[coords].grid, ref.grid) # true!
isequal(layer[coords], ref) # false, as bounding coordinates are not exactly the same

tmp = collect(range(layer.left+stride(layer, 1), layer.right-stride(layer, 1), length=size(layer,2)))
145.0 in tmp .+ stride(layer,1)
# longitudes # should be ok in theory
# stride # should be ok in theory
# getindex(layer, irange, jrange) # could be the problem too
i = jmin:jmax # need to be inversed
j = imin:imax
layer.grid[i, j] # wrong dimensions
i_fix = i
j_fix = j
layer.grid[i_fix,j_fix]
minimum(longitudes(layer)[j_fix])-stride(layer,1)
maximum(longitudes(layer)[j_fix])+stride(layer,1)
minimum(latitudes(layer)[i_fix])-stride(layer,2)
maximum(latitudes(layer)[i_fix])+stride(layer,2)


text:
export ConcurrentCircularBuffer

"""

Thread-safe `CircularBuffer{T}`. Can be used to implement a fixed-capacity FIFO/FILO message queue
with preemption, i.e., pushing a message to a full queue causes a messages already in the queue to
be discarded.
"""
struct ConcurrentCircularBuffer{T}
    cb::CircularBuffer{T}           # underlying circular buffer for storage
    lk::Threads.SpinLock            # for protecting access to the buffer
    c::Threads.Condition            # used to notify threads waiting to pop from the buffer when elements are added
    open::Threads.Atomic{Bool}      # set to false when calling close(ch) on the ConcurrentCircularBuffer
end

ConcurrentCircularBuffer{T}(n::Integer) where T = ConcurrentCircularBuffer{T}(CircularBuffer{T}(n), Threads.SpinLock(), Threads.Condition(), Threads.Atomic{Bool}(true))

function DataStructures.isfull(ch::ConcurrentCircularBuffer)
    lock(ch.lk) do
        return isfull(ch.cb)
    end
end

function Base.isopen(ch::ConcurrentCircularBuffer)
    lock(ch.lk) do
        return ch.open[]
    end
end

"""

Close the `ConcurrentCircularBuffer`. An `InvalidStateException` is thrown by `put!` on a closed
`ConcurrentCircularBuffer` and `pop!` or `popfirst!` on an empty closed `ConcurrentCircularBuffer`.
"""
function Base.close(ch::ConcurrentCircularBuffer)
    lock(ch.lk) do
        ch.open[] = false
    end
    lock(ch.c) do
        notify(ch.c)
    end
    return
end

function DataStructures.isempty(ch::ConcurrentCircularBuffer)
    lock(ch.lk) do
        return isempty(ch.cb)
    end
end

function DataStructures.empty!(ch::ConcurrentCircularBuffer)
    lock(ch.lk) do
        return empty!(ch.cb)
    end
end

DataStructures.capacity(ch::ConcurrentCircularBuffer) = capacity(ch.cb)
Base.eltype(ch::ConcurrentCircularBuffer) = eltype(ch.cb)

function Base.length(ch::ConcurrentCircularBuffer)
    lock(ch.lk) do
        return length(ch.cb)
    end
end
Base.size(ch::ConcurrentCircularBuffer) = length(ch)

"""

Return `true` if there is at least one element in the buffer.
"""
Base.isready(ch::ConcurrentCircularBuffer) = length(ch) > 0

"""

Add an element to the back and overwrite front if full.
"""
function Base.push!(ch::ConcurrentCircularBuffer, v)
    lock(ch.lk) do
        if !ch.open[]
            throw(InvalidStateException("ConcurrentCircularBuffer is closed.", :closed))
        end
        push!(ch.cb, v)
    end
    lock(ch.c) do
        notify(ch.c)
    end
    ch
end

"""

Add an element to the front and overwrite back if full.
"""
function DataStructures.pushfirst!(ch::ConcurrentCircularBuffer, v)
    lock(ch.lk) do
        if !ch.open
            throw(InvalidStateException("ConcurrentCircularBuffer is closed.", :closed))
        end
        pushfirst!(ch.cb, v)
    end
    lock(ch.c) do
        notify(ch.c)
    end
    ch
end

"""

Add all elements in `vs` to the back and overwrite front if full.
"""
function Base.append!(ch::ConcurrentCircularBuffer, vs::AbstractVector)
    lock(ch.lk) do
        if !ch.open[]
            throw(InvalidStateException("ConcurrentCircularBuffer is closed.", :closed))
        end
        append!(ch.cb, vs)
    end
    lock(ch.c) do
        notify(ch.c)
    end
    ch
end

"""

Remove the element at the back.
"""
function DataStructures.pop!(ch::ConcurrentCircularBuffer{T})::T where T
    local rv::T
    success = false
    lock(ch.c) do
        while !success
            lock(ch.lk) do
                if length(ch.cb) > 0
                    success = true
                    rv = pop!(ch.cb)
                elseif !ch.open[]
                    throw(InvalidStateException("ConcurrentCircularBuffer is closed.", :closed))
                end
            end
            if !success
                wait(ch.c)
            end
        end
    end
    rv
end

"""

Remove the element at the front.
"""
function DataStructures.popfirst!(ch::ConcurrentCircularBuffer{T})::T where T
    local rv::T
    success = false
    lock(ch.c) do
        while !success
            lock(ch.lk) do
                if length(ch.cb) > 0
                    success = true
                    rv = popfirst!(ch.cb)
                elseif !ch.open[]
                    throw(InvalidStateException("ConcurrentCircularBuffer is closed.", :closed))
                end
            end
            if !success
                wait(ch.c)
            end
        end
    end
    rv
end

"""

Same as `popfirst!`.
"""
Base.take!(ch::ConcurrentCircularBuffer) = popfirst!(ch)

function _pop!(vs::AbstractVector, ch::ConcurrentCircularBuffer)
    rv = 0
    for i in 1:length(vs)
        if length(ch.cb) == 0
            break
        end
        vs[i] = pop!(ch.cb)
        rv += 1
    end
    rv
end

function _popfirst!(vs::AbstractVector, ch::ConcurrentCircularBuffer)
    rv = 0
    for i in 1:length(vs)
        if length(ch.cb) == 0
            break
        end
        vs[i] = popfirst!(ch.cb)
        rv += 1
    end
    rv
end

"""

Remove up to `length(vs)` element at the back, place them in `vs`, and return the number of
elements removed.
"""
function DataStructures.pop!(vs::AbstractVector, ch::ConcurrentCircularBuffer)
    rv = 0
    success = false
    lock(ch.c) do
        while !success
            lock(ch.lk) do
                if length(ch.cb) > 0
                    success = true
                    rv = _pop!(vs, ch)
                elseif !ch.open[]
                    throw(InvalidStateException("ConcurrentCircularBuffer is closed.", :closed))
                end
            end
            if !success
                wait(ch.c)
            end
        end
    end
    rv
end

"""

Remove up to `length(vs)` element at the front, place them in `vs`, and return the number of
elements removed.
"""
function DataStructures.popfirst!(vs::AbstractVector, ch::ConcurrentCircularBuffer)
    rv = 0
    success = false
    lock(ch.c) do
        while !success
            lock(ch.lk) do
                if length(ch.cb) > 0
                    success = true
                    rv = _popfirst!(vs, ch)
                elseif !ch.open[]
                    throw(InvalidStateException("ConcurrentCircularBuffer is closed.", :closed))
                end
            end
            if !success
                wait(ch.c)
            end
        end
    end
    rv
end


text:
using Dispersal, Test, Statistics
using Dispersal: precalc_human_dispersal!, upsample_index, initdownsample, downsample,
                 CellGravity, CellInterval, downsample!, populate, populate!

@testset "CellGravity sums and compares gravity" begin
    g = CellGravity(5.0, (1, 1))
    @test g + 3.7 == 8.7
    @test g > 4.0
    @test g < 6.0
    @test g < CellGravity(6.0, (1, 1))
end

@testset "CellInterval sums and compares interval" begin
    i = CellInterval(0.5, (1, 1))
    @test i > 0.4
    @test i < 0.6
    @test i < CellInterval(0.6, (1, 1))
end

# Setup precalc
init =      [0.0f0 0.0f0  0.0f0 0.0f0 0.0f0
             0.0f0 0.0f0  0.0f0 0.0f0 0.0f0
             0.0f0 0.0f0 10.0f0 0.0f0 0.0f0
             0.0f0 0.0f0  0.0f0 0.0f0 0.0f0
             0.0f0 0.0f0  0.0f0 0.0f0 0.0f0]
human_pop = [1.0f0   2.0f0   0.1f0 0.0f0 1.0f0
             1.0f0   0.0f0   0.3f0 0.4f0 1.0f0
             missing missing 1.0f0 1.0f0 1.0f0
             missing missing 1.0f0 1.0f0 1.0f0
             0.3f0   0.3f0   1.0f0 1.0f0 missing]

cellsize = 1
scale = 2
aggregator = mean
human_exponent = 2
dist_exponent = 1
nshortlisted = 4
dispersalperactivity = 0.01
max_dispersers = 4

# Allocate memory
human_buffer = downsample(human_pop, aggregator, scale)
distances = initdownsample(human_pop, scale)
gravities, gravity_vector = Dispersal.alloc_gravities(human_buffer, nshortlisted)
dest_shortlists = Dispersal.alloc_dest_shortlist(nshortlisted, size(human_buffer))
gravity_vector

@testset "downsample works" begin
    test_downsample = [1.0f0   0.2f0 1.0f0
                       missing 1.0f0 1.0f0
                       0.3f0   1.0f0 missing]
    for i in 1:length(human_buffer)
        @test human_buffer[i] === test_downsample[i]
    end
    Dispersal.precalc_human_exponent!(human_buffer, human_exponent)
    test_exponent =  [1.0f0    0.2f0^2 1.0f0
                      missing  1.0f0   1.0f0
                      0.09f0   1.0f0   missing]
    for i in 1:length(human_buffer)
        @test human_buffer[i] === test_exponent[i]
    end
end

@testset "distances are correct" begin
    Dispersal.precalc_distances!(distances, dist_exponent, cellsize, scale)
    @test distances[3, 1] == sqrt(2.0f0^2 + 0.0f0^2) * 2
    @test distances[3, 2] == sqrt(2.0f0^2 + 1.0f0^2) * 2
    @test distances[2, 3] == sqrt(1.0f0^2 + 2.0f0^2) * 2
end

@testset "gravities" begin
    Dispersal.precalc_distances!(distances, dist_exponent, cellsize, scale)
    Dispersal.assign_gravities!(gravities, human_buffer, distances, 2, 2)
    @test map(g -> g.gravity, gravities) ==
        [0.35355338f0  0.020000001f0 0.35355338f0
         0.0f0         1.3068552f0     0.5f0
         0.031819806f0 0.5f0         0.0f0]
    @test map(g -> g.index, gravities) ==
        [(1, 1) (1, 2) (1, 3)
         (2, 1) (2, 2) (2, 3)
         (3, 1) (3, 2) (3, 3)]
    gravity_vector .= vec(gravities)
    partialsort!(gravity_vector, 1:nshortlisted, rev=true)
    gravity_shortlist = view(gravity_vector, 1:nshortlisted)
    dest = dest_shortlists[2, 2]
    Dispersal.gravity2inverval!(dest, gravity_shortlist)
    @test map(d -> d.index, dest) == [(1, 1), (2, 3), (3, 2), (2, 2)]
    @test map(d -> d.cumprop, dest) == [0.1328944f0, 0.32083547f0, 0.50877655f0, 1.0f0]
end

@testset "cell" begin
    Dispersal.precalc_cell!(dest_shortlists, nshortlisted, human_buffer, distances, (gravities, gravity_vector), 2, 2)
    dest = dest_shortlists[2, 2]
    @test map(d -> d.index, dest) == [(1, 1), (2, 3), (3, 2), (2, 2)]
    @test map(d -> d.cumprop, dest) == [0.1328944f0, 0.32083547f0, 0.50877655f0, 1.0f0]
end

@testset "shortlists are correctly ordered" begin
    precalc_human_dispersal!(dest_shortlists, human_pop, cellsize, scale, aggregator,
                             human_exponent, dist_exponent, nshortlisted, human_buffer,
                             distances)
    @test length(dest_shortlists) == 9
    @test length(dest_shortlists[1, 1]) == 4
    @test searchsortedfirst(dest_shortlists[1, 1], 1.0) == 4
    @test dest_shortlists[1, 1][4].index == (1, 1)
    @test dest_shortlists[1, 1][3].index == (2, 2)
end

@testset "populate and HumanDispersal construction" begin
    # Setup precalc
    human_pop = Union{Float32,Missing}[1 for i = 1:5, j = 1:5]
    init = zeros(Float32, size(human_pop)...)
    init[3, 3] = 10.0
    cellsize = 1
    scale = 1
    aggregator = mean
    human_exponent = 1
    dist_exponent = 1
    nshortlisted = 9
    dispersalperpop = 0.1
    mode = BatchGroups()

    humandisp = HumanDispersal(mode=mode, human_pop=human_pop, dispersalperpop=dispersalperpop, cellsize=cellsize, scale=scale,
       aggregator=aggregator, human_exponent=human_exponent, max_dispersers=max_dispersers,
       dist_exponent=dist_exponent, nshortlisted=nshortlisted);
    dest_shortlists = humandisp.dest_shortlists

    @testset "the highest value is the cell itself" begin
        @test (x->x[end].index).(humandisp.dest_shortlists) == Tuple.(collect(CartesianIndices(human_pop)))
    end

    @testset "precalc is simmetrical with simmetrical inputs" begin
        a = zeros(5, 5)
        b = zeros(5, 5)
        # There is a bug in julia 1.6.0 partialsort!
        @test (x->x.cumprop).(dest_shortlists[5, 5]) == (x->x.cumprop).(dest_shortlists[1, 1])
        @test (x->x.cumprop).(dest_shortlists[2, 2]) == (x->x.cumprop).(dest_shortlists[4, 4])
        @test a == reverse(reverse(b, dims=2), dims=1)
    end

    a = populate(humandisp, 1, 1)
    b = populate(humandisp, 5, 5)
    c = populate(humandisp[5, 3], size(humandisp), scale)
    d = populate(humandisp[3, 3], size(humandisp), scale)

    @test all(sum.((a, b, c, d)) .≈ 1)

    @test (a .> 0) ==
               [1 1 1 0 0
                1 1 1 0 0
                1 1 1 0 0
                0 0 0 0 0
                0 0 0 0 0]

    @test (b .> 0) ==
               [0 0 0 0 0
                0 0 0 0 0
                0 0 1 1 1
                0 0 1 1 1
                0 0 1 1 1]

    @test (c .> 0) ==
               [0 0 0 0 0
                0 0 0 0 0
                0 0 1 0 0
                0 1 1 1 0
                1 1 1 1 1]

    @test (d .> 0) ==
               [0 0 0 0 0
                0 1 1 1 0
                0 1 1 1 0
                0 1 1 1 0
                0 0 0 0 0]

    @test populate(humandisp, 1, 1) ≈
        reverse(populate(humandisp, 1, 5); dims=2) ≈
        reverse(populate(humandisp, 5, 1); dims=1) ≈
        reverse(reverse(populate(humandisp, 5, 5); dims=1); dims=2)
    @test populate(humandisp, 2, 2) ≈
        reverse(populate(humandisp, 2, 4); dims=2) ≈
        reverse(populate(humandisp, 4, 2); dims=1) ≈
        reverse(reverse(populate(humandisp, 4, 4); dims=1); dims=2)
    @test populate(humandisp, 2, 3) ≈
        rotr90(populate(humandisp, 3, 2)) ≈
        rot180(populate(humandisp, 4, 3)) ≈
        rotl90(populate(humandisp, 3, 4))

    # Symmetry is close but not exactly identical due to truncation
    # of the shortlist and row/column order. This causes small differences
    # choosing the last item in the shortlist when the last index +1
    # has the same gravity. It should make little real difference
    # with larger shortlists as intervals at the end of the shortlist
    # should be small.
    @test (populate(humandisp, 1, 2) ≈
        rotl90(populate(humandisp, 2, 5)) ≈
        rot180(populate(humandisp, 5, 4)) ≈
        rotl90(populate(humandisp, 2, 5))) == false

    combined = populate(humandisp)
    sum(combined) ≈ 25
    @test (combined .> 0) ==
               [1 1 1 1 1
                1 1 1 1 1
                1 1 1 1 1
                1 1 1 1 1
                1 1 1 1 1]

    # Again symmetry is not exact, there are small errors from truncation.
    reverse(combined; dims=1) ≈ combined == false
    reverse(combined; dims=2) ≈ combined == false

    @testset "populate handles missing" begin
        @test populate!(init, missing, scale) === missing
    end
end

@testset "human dispersal simulation maintains total population" begin
    dispersalperpop = 0.1
    mode = BatchGroups()

    humandisp = HumanDispersal(human_pop=human_pop, dispersalperpop=dispersalperpop, cellsize=cellsize, scale=scale,
       aggregator=aggregator, human_exponent=human_exponent, max_dispersers=max_dispersers,
       dist_exponent=dist_exponent, nshortlisted=nshortlisted);

    ruleset = Ruleset(humandisp)
    output = ArrayOutput(init; tspan=1:10)

    for i = 1:100
        sim!(output, ruleset)
        @test init == output[1]
        @test init != output[10]
        # Population is allways maintained
        @test sum(init) == sum(output[2])
        @test sum(init) == sum(output[3])
        @test sum(init) == sum(output[5])
        @test sum(init) == sum(output[10])
    end

end


text:
const INDENT=4

"julia Expr printing color schema"
module Color
using Crayons.Box

kw(x) = LIGHT_MAGENTA_FG(Base.string(x))
fn(x) = LIGHT_BLUE_FG(Base.string(x))
line(x) = DARK_GRAY_FG(Base.string(x))
literal(x) = YELLOW_FG(Base.string(x))
type(x) = LIGHT_GREEN_FG(Base.string(x))
string(x::String) = Box.CYAN_FG(x)
string(x) = Box.CYAN_FG(Base.string(x))

end

no_indent(io::IO) = IOContext(io, :indent=>0)
no_indent_first_line(io::IO) = IOContext(io, :no_indent_first_line=>true)

# 1.0 compatibility
function indent_print(io::IO, ::Nothing)
    indent = get(io, :indent, 0)
    tab = get(io, :tab, " ")
    print(io, tab^indent, "nothing")
end

function indent_print(io::IO, xs...)
    indent = get(io, :indent, 0)
    tab = get(io, :tab, " ")
    Base.print(io, tab^indent, xs...)
end

function indent_println(io::IO, xs...)
    if get(io, :no_indent_first_line, false)
        indent_print(no_indent(io), xs..., "\n")
    else
        indent_print(io, xs..., "\n")
    end
end

function within_line(f, io)
    indent_print(io)
    f(no_indent(io))
end

function within_indent(f, io)
    f(indent(io))
end

function indent(io)
    IOContext(io, :indent => get(io, :indent, 0) + INDENT)
end

"""
    with_marks(f, io, lhs, rhs)

Print using `f` with marks specified on LHS and RHS by `lhs` and `rhs`.
See also [`with_parathesis`](@ref), [`with_curly`](@ref), [`with_brackets`](@ref),
[`with_begin_end`](@ref).
"""
function with_marks(f, io::IO, lhs, rhs)
    indent_print(io, lhs)
    f()
    indent_print(io, rhs)
end

"""
    with_parathesis(f, io::IO)

Print with parathesis. See also [`with_marks`](@ref),
[`with_curly`](@ref), [`with_brackets`](@ref),
[`with_begin_end`](@ref).

# Example

```julia
julia> with_parathesis(stdout) do
        print(1, ", ", 2)
    end
(1, 2)
```
"""
with_parathesis(f, io::IO) = with_marks(f, io, "(", ")")

"""
    with_curly(f, io::IO)

Print with curly parathesis. See also [`with_marks`](@ref), [`with_parathesis`](@ref),
[`with_brackets`](@ref), [`with_begin_end`](@ref).
"""
with_curly(f, io::IO) = with_marks(f, io, "{", "}")

"""
    with_brackets(f, io::IO)

Print with brackets. See also [`with_marks`](@ref), [`with_parathesis`](@ref),
[`with_curly`](@ref), [`with_begin_end`](@ref).
"""
with_brackets(f, io::IO) = with_marks(f, io, "[", "]")

"""
    with_triple_quotes(f, io::IO)

Print with triple quotes.
"""
with_triple_quotes(f, io::IO) = with_marks(f, io, Color.string("\"\"\"\n"), Color.string("\"\"\""))

"""
    with_double_quotes(f, io::IO)

Print with double quotes.
"""
with_double_quotes(f, io::IO) = with_marks(f, io, Color.string("\""), Color.string("\""))

"""
    with_begin_end(f, io::IO)

Print with begin ... end. See also [`with_marks`](@ref), [`with_parathesis`](@ref),
[`with_curly`](@ref), [`with_brackets`](@ref).
"""
with_begin_end(f, io::IO) = with_marks(f, io, "begin", "end")

"""
    print_collection(io, xs; delim=",")

Print a collection `xs` with deliminator `delim`, default is `","`.
"""
function print_collection(io, xs; delim=", ")
    for i in 1:length(xs)
        print_ast(io, xs[i])
        if i !== length(xs)
            indent_print(io, delim)
        end
    end
end

"""
    print_ast(io::IO, xs...)

Print Julia AST. This is a custom implementation of
`Base.show(io, ::Expr)`.
"""
function print_ast(io::IO, xs...)
    foreach(xs) do x
        print_ast(io, x)
    end
end

function print_ast(io::IO, ex)
    tab = get(io, :tab, " ")
    first_line_io = get(io, :no_indent_first_line, false) ? no_indent(io) : io
    @match ex begin
        ::Union{Number} => indent_print(first_line_io, Color.literal(ex))
        ::String => indent_print(first_line_io, Color.string(repr(ex)))
        ::Symbol => indent_print(first_line_io, ex)

        Expr(:tuple, xs...) => begin
            with_parathesis(io) do
                print_collection(no_indent(io), xs)
            end
        end

        Expr(:(::), name, type) => begin
            within_line(io) do io
                print_ast(io, name)
                indent_print(io, "::")
                print_ast(io, Color.type(type))
            end
        end

        Expr(:kw, name, value) => begin
            within_line(io) do io
                print_ast(io, name)
                indent_print(io, tab, "=", tab)
                print_ast(io, value)
            end
        end

        Expr(:(=), l, r) => begin
            print_ast(io, l)
            print(io, tab, Color.kw("="), tab)
            print_ast(no_indent_first_line(io), r)
        end

        Expr(:call, name, args...) => begin
            if !get(io, :no_indent_first_line, false)
                indent_print(io)
            end
            if name in [:+, :-, :*, :/, :\, :(===), :(==), :(:)]
                print_collection(no_indent(io), args; delim=Color.fn(string(tab, name, tab)))
            else
                indent_print(no_indent(io), Color.fn(name))
                with_parathesis(no_indent(io)) do
                    print_collection(no_indent(io), args)
                end
            end
        end

        Expr(:block, stmts...) => begin
            indent_println(io, Color.kw("begin"))
            io = IOContext(io, :no_indent_first_line=>false)
            within_indent(io) do io
                for i in 1:length(stmts)
                    print_ast(io, stmts[i])
                    indent_println(io)
                end
            end
            indent_print(io, Color.kw("end"))
        end

        Expr(:let, vars, body) => begin
            if get(io, :no_indent_first_line, false)
                print(io, Color.kw("let"))
            else
                indent_print(io, Color.kw("let"))
            end

            isempty(vars.args) || print_collection(no_indent(io), vars.args)
            println(io)
            io = IOContext(io, :no_indent_first_line=>false)
            within_indent(io) do io
                if body isa Expr && body.head === :block
                    stmts = body.args
                    for i in 1:length(stmts)
                        print_ast(io, stmts[i])
                        indent_println(io)
                    end
                else
                    print_ast(io, body)
                    println(io)
                end
            end
            indent_print(io, Color.kw("end"))
        end

        Expr(:return, xs...) => begin
            within_line(io) do io
                print_ast(io, Color.kw("return"))
                indent_print(io, tab)
                print_ast(io, xs...)
            end
        end

        Expr(:if, xs...) => begin
            print_ast(io, JLIfElse(ex))
        end

        ::LineNumberNode => indent_print(io, Color.line(ex))
        # fallback to default printing
        _ => begin
            if get(io, :no_indent_first_line, false)
                print(io, ex)
            else
                indent_print(io, ex)
            end
        end
    end
end

function print_ast(io::IO, def::JLFor)
    def.kernel === nothing && return
    tab = get(io, :tab, " ")
    indent_print(io, Color.kw("for"), tab)

    within_indent(io) do io
        for i in 1:length(def.vars)
            print_ast(i==1 ? no_indent(io) : io, def.vars[i])
            print(io, tab, Color.kw("in"), tab)
            print_ast(no_indent(io), def.iterators[i])
            i < length(def.vars) && print(io, ",")
            println(io)
        end
        indent_println(io, Color.line("#= loop body =#"))
        print_ast(io, def.kernel)
        println(io)
    end

    indent_print(io, Color.kw("end"))
end

function print_ast(io::IO, def::JLIfElse)
    isempty(def.map) && return print_ast(io, def.otherwise)
    tab = get(io, :tab, " ")
    indent_print(io, Color.kw("if"), tab)
    for (k, (cond, action)) in enumerate(def.map)
        print_ast(no_indent(io), cond)
        indent_println(io)
        print_ast(indent(io), action)
        indent_println(io)

        if k !== length(def.map)
            indent_print(io, Color.kw("elseif"), tab)
        end
    end
    if def.otherwise !== nothing
        indent_print(io, Color.kw("else"), "\n")
        print_ast(indent(io), def.otherwise)
        indent_println(io)
    end
    indent_print(io, Color.kw("end"))
end

function print_ast(io::IO, def::JLMatch)
    print_ast(io, def.line)
    println(io)
    isempty(def.map) && return print_ast(io, def.fallthrough)
    tab = get(io, :tab, " ")
    indent_println(io, Color.kw("@match"), tab, def.item, tab, Color.kw("begin"))
    within_indent(io) do io
        for (k, (pattern, action)) in enumerate(def.map)
            within_line(io) do io
                print_ast(io, pattern)
                print(io, tab, Color.kw("=>"), tab)
                print_ast(io, action)
            end
            println(io)
        end

        # match must have fallthrough
        indent_print(io, "_")
        print(io, tab, Color.kw("=>"), tab)
        print_ast(io, def.fallthrough)
    end
    println(io)
    indent_print(io, Color.kw("end"))
end

function print_ast(io::IO, def::JLFunction)
    tab = get(io, :tab, " ")
    within_line(io) do io
        def.head === :function && indent_print(io, Color.kw("function"), tab)
        # print calls
        def.name === nothing || indent_print(io, Color.fn(def.name))
        with_parathesis(io) do
            print_collection(io, def.args)
            if def.kwargs !== nothing
                indent_print(io, "; ")
                print_collection(io, def.kwargs)
            end
        end

        if def.whereparams !== nothing
            indent_print(io, tab, Color.kw("where"), tab)
            with_curly(io) do
                print_collection(io, def.whereparams)
            end
        end

        def.head === :(=) && indent_print(io, tab, "=", tab)
        def.head === :(->) && indent_print(io, tab, "->", tab)
    end

    # print body
    if def.head === :function
        println(io)
        within_indent(io) do io
            @match def.body begin
                Expr(:block, stmts...) => begin
                    for i in 1:length(stmts)
                        print_ast(io, stmts[i])
                        indent_println(io)
                    end
                end

                _ => begin
                    print_ast(io, def.body)
                    println(io)
                end
            end
        end
        indent_print(io, Color.kw("end"))
    else
        print_ast(no_indent_first_line(io), def.body)
    end
end

function print_ast(io::IO, def::JLStruct)
    print_ast_struct(io, def)
end

function print_ast(io::IO, def::JLKwStruct)
    print_ast_struct(io, def)
end

function print_ast(io::IO, def::JLField)
    print_ast_struct_field(io, def)
end

function print_ast(io::IO, def::JLKwField)
    print_ast_struct_field(io, def)
    def.default === no_default || indent_print(no_indent(io), " = ", def.default)
end

function print_ast_doc(io::IO, def)
    def.doc === nothing && return
    doc = def.doc
    with_triple_quotes(io) do
        indent_print(io, Color.string(doc))
    end
    indent_println(io)
end

function print_ast_struct(io::IO, def)
    def.line === nothing || indent_println(io, Color.line(def.line))
    print_ast_doc(io, def)
    print_ast_struct_head(io, def)
    for each in def.fields
        indent_println(no_indent(io))
        print_ast(indent(io), each)
    end
    indent_println(no_indent(io))

    for each in def.constructors
        print_ast(indent(io), each)
        indent_println(io)
    end

    indent_print(io, Color.kw("end"))
end

function print_ast_struct_field(io::IO, def)
    def.line === nothing || indent_println(io, Color.line(def.line))
    if def.doc !== nothing
        indent_print(io)
        with_double_quotes(no_indent(io)) do
            indent_print(no_indent(io), Color.string(def.doc))
        end
        indent_println(io)
    end
    indent_print(io, def.name)
    def.type === Any || indent_print(no_indent(io), "::", Color.type(def.type))
end

function print_ast_struct_head(io::IO, def)
    tab = get(io, :tab, " ")
    # make sure there is only one indent in the same line
    printed_indent = false
    if def isa JLKwStruct
        indent_print(io, Color.line("#= kw =#"), tab)
        printed_indent = true
    end

    if def.ismutable
        indent_print(printed_indent ? no_indent(io) : io, Color.kw("mutable"), tab)
    end

    indent_print(printed_indent ? no_indent(io) : io, Color.kw("struct"))
    indent_print(io, tab, def.name)

    isempty(def.typevars) || with_curly(no_indent(io)) do
        print_collection(no_indent(io), def.typevars)
    end

    if def.supertype !== nothing
        indent_print(no_indent(io), tab, "<:", tab, Color.type(def.supertype))
    end
end

print_ast(::IO, def::JLExpr) = error("Printings.print_ast is not defined for $(typeof(def))")
Base.show(io::IO, def::JLExpr) = print_ast(io, def)


text:
using RecipesBase

################################################################################
# Destriping results

@doc raw"""
    mutable struct DestripingData{T <: Real, O <: Healpix.Order}

Structure containing the result of a destriping operation. It contains
the following fields:

- `skymap`: a `PolarizedMap` containing the maximum-likelihood map
- `hitmap`: a map containing the weights of each pixel (this is not
  the hit count, as it is normalized over the value of ``σ^2`` for
  each sample in the TOD)
- `nobs_matrix`: an array of [`NobsMatrixelement`](@ref) objects,
  which can be used to determine which pixels were the most
  troublesome in the reconstruction of tehe I/Q/U components
- `baselines`: the computed baselines. These must be kept as a list of
  `RunLengthArray` objects
- `stopping_factors`: sequence of stopping factors, one for
  each iteration of the CG algorithm
- `threshold`: upper limit on the tolerable error for the Conjugated
  Gradient method (optional, default is `1e-9`). See
  [`calc_stopping_factor`](@ref).
- `max_iterations`: maximum number of allowed iterations for the
  Conjugated Gradient algorithm (optional, default is `1000`).
- `use_preconditioner`: a Boolean flag telling whether a
  preconditioner should be used or not (optional, default is
  `true`). Usually you want this to be `true`.

The structure provides only one constructor, which accepts the
following parameters:

- `nside`: resolution of `skymap` and `hitmap`

- `obs_list`: vector of [`Observation`](@ref); it is used to determine
  which pixels in the map are going to be observed

- `runs`: list of vectors (one for each element in `obs_list`),
  containing the number of samples to be included in each baseline for
  each observation

The constructor accepts the following keywords as well:

- `threshold` (default: `1e-9`)

- `max_iterations` (default: 1000)

- `use_preconditioner` (default: `true`)

The function [`destripe!`](@ref) use `DestripingData` to return the
result of its computation.

"""
mutable struct DestripingData{T <: Real, O <: Healpix.Order}
    skymap::Healpix.PolarizedMap{T, O}
    hitmap::Healpix.Map{T, O}
    nobs_matrix::Vector{NobsMatrixElement{T}}
    baselines::Vector{RunLengthArray{Int, T}}
    stopping_factors::Vector{T}
    threshold::Float64
    max_iterations::Int
    use_preconditioner::Bool

    function DestripingData{T,O}(
        nside,
        obs_list::Vector{Observation{T}},
        runs;
        threshold = 1e-9,
        max_iterations = 1000,
        use_preconditioner = true,
        comm = nothing,
    ) where {T <: Real, O <: Healpix.Order}

        @assert length(obs_list) == length(runs)
        @assert threshold > 0.0
        @assert max_iterations > 0

        npix = Healpix.nside2npix(nside)
        nobs_matrix = [NobsMatrixElement{T}() for i in 1:npix]
        compute_nobs_matrix!(nobs_matrix, obs_list, comm = comm)

        new(
            Healpix.PolarizedMap{T, O}(
                zeros(npix),
                zeros(npix),
                zeros(npix),
            ),
            Healpix.Map{T,O}(nside),
            nobs_matrix,
            [RunLengthArray{Int,T}(x, zeros(length(x))) for x in runs],
            T[],  # List of stopping factors
            threshold,
            max_iterations,
            use_preconditioner,
        )
    end

end

function Base.show(io::IO, d::DestripingData)
    print(io, "DestripingData(NSIDE=$(d.hitmap.resolution.nside))")
end

function Base.show(io::IO, ::MIME"text/plain", d::DestripingData)
    # Count how many pixels have been observed in the map
    obspix = 0
    for elem in d.nobs_matrix
        elem.neglected || (obspix += 1)
    end

    # Put all the baselines of each observation into one long list
    # (without allocating memory!)
    flattened_baselines = Iterators.flatten([values(x) for x in d.baselines])

    println(
        io,
        @sprintf(
            """Destriping data:
NSIDE of the map........................ %d
Number of observed pixels............... %d (%.1f%%)
Number of iterations.................... %d/%d
Best convergence factor................. %s (threshold: %e)
Number of observations.................. %d
Average value of the baselines.......... %e
RMS of the baselines.................... %e
Preconditioner.......................... %s
""",
            d.hitmap.resolution.nside,
            obspix,
            100.0 * obspix / length(d.nobs_matrix),
            length(d.stopping_factors),
            d.max_iterations,
            isempty(d.stopping_factors) ? "none" : @sprintf("%e", minimum(d.stopping_factors)),
            d.threshold,
            length(d.baselines),
            mean(flattened_baselines),
            std(flattened_baselines),
            "$(d.use_preconditioner)",
        )
    )
end

@doc raw"""
    reset_maps!(d::DestripingData{T, O}) where {T <: Real, O <: Healpix.Order}

Set the skymap and the hitmap in `d` to zero. Nothing is done on the
other fields.

"""
function reset_maps!(d::DestripingData{T, O}) where {T <: Real, O <: Healpix.Order}
    d.skymap.i .= 0.0
    d.skymap.q .= 0.0
    d.skymap.u .= 0.0

    d.hitmap .= 0.0
end

@recipe function plot_destriping_data(d::DestripingData{T, O}) where {T <: Real, O <: Healpix.Order}
    seriestype := :path
    xlabel := "Iteration number"
    ylabel := "Stopping factor"
    linecolor := :black

    isempty(d.stopping_factors) && return ([0.0], [0.0])
    num_of_steps = length(d.stopping_factors)

    @series begin
        # Horizontal line showing the threshold
        seriestype := :path
        linecolor := :red
        ([1, num_of_steps], [d.threshold, d.threshold])
    end

    ((1:num_of_steps), d.stopping_factors)
end

@doc raw"""
    plot(d::DestripingData{T, O}) where {T <: Real, O <: Healpix.Order}

Plot the convergence plot for a destriping solution.
"""
plot_destriping_data