tomlpath = joinpath(@__DIR__, "sbm_config.toml")

@testset "BMI" begin

    @testset "BMI functions" begin

        model = BMI.initialize(Wflow.Model, tomlpath)

        @testset "initialization and time functions" begin
            @test BMI.get_time_units(Wflow.Model) == "seconds since 1970-01-01T00:00:00"
            @test BMI.get_time_step(model) == 86400.0
            @test BMI.get_start_time(model) == 9.467712e8
            @test BMI.get_current_time(model) == 9.467712e8
            @test BMI.get_end_time(model) == 9.493632e8
        end

        @testset "model information functions" begin
            @test BMI.get_component_name(model) == "sbm"
            @test BMI.get_input_item_count(model) == 186
            @test BMI.get_output_item_count(model) == 186
            @test BMI.get_input_var_names(model)[[1, 5, 120, 186]] == [
                "vertical.Δt",
                "vertical.n_unsatlayers",
                "lateral.land.q_av",
                "lateral.river.reservoir.evaporation",
            ]
        end

        @testset "variable information functions" begin
            @test BMI.get_var_grid(model, "vertical.θₛ") == 4
            @test BMI.get_var_grid(model, "lateral.river.h") == 2
            @test BMI.get_var_grid(model, "lateral.river.reservoir.inflow") == 0
            @test_throws ErrorException BMI.get_var_grid(model, "lateral.river.lake.volume")
            # Vector{Float64} printing on Julia 1.6+
            @test BMI.get_var_type(model, "lateral.river.reservoir.inflow") in
                  ["Array{$Float,1}", "Vector{$Float}"]
            @test BMI.get_var_type(model, "vertical.n") == "Int64"
            @test BMI.get_var_units(model, "vertical.θₛ") == "-"
            @test BMI.get_var_itemsize(model, "lateral.subsurface.ssf") == sizeof(Float)
            @test BMI.get_var_nbytes(model, "vertical.n") == 8
            @test BMI.get_var_nbytes(model, "lateral.river.q") ==
                  length(model.lateral.river.q) * sizeof(Float)
            @test BMI.get_var_location(model, "lateral.river.q") == "node"
        end

        model = BMI.update(model)

        @testset "update and get and set functions" begin
            @test BMI.get_current_time(model) == 9.468576e8
            @test mean(BMI.get_value(model, "vertical.zi")) ≈ 276.3767651555451
            @test BMI.get_value_at_indices(model, "lateral.river.q", [1, 100, 5617]) ≈
                  [0.6211503865184697, 5.219305686635002, 0.026163746306482282]
            BMI.set_value(model, "vertical.zi", fill(300.0, length(model.vertical.zi)))
            @test mean(BMI.get_value(model, "vertical.zi")) == 300.0
            BMI.set_value_at_indices(model, "vertical.zi", [1], [250.0])
            @test BMI.get_value_at_indices(model, "vertical.zi", [1, 2]) == [250.0, 300.0]
        end

        @testset "model grid functions" begin
            @test BMI.get_grid_type(model, 0) == "unstructured"
            @test BMI.get_grid_rank(model, 0) == 2
            @test BMI.get_grid_size(model, 0) == 2
            @test BMI.get_grid_size(model, 2) == 5809
            @test BMI.get_grid_size(model, 4) == 50063
            @test BMI.get_grid_shape(model, 4) == 50063
            @test minimum(BMI.get_grid_x(model, 4)) ≈ 5.426666666666667f0
            @test maximum(BMI.get_grid_x(model, 4)) ≈ 7.843333333333344f0
            @test BMI.get_grid_x(model, 0) ≈ [5.760000000000002f0, 5.918333333333336f0]
            @test BMI.get_grid_y(model, 0) ≈ [48.92583333333333f0, 49.909166666666664f0]
            @test BMI.get_grid_node_count(model, 0) == 2
        end

        @testset "update until and finalize" begin
            time = BMI.get_current_time(model) + 2 * BMI.get_time_step(model)
            @test_logs (:info, "update model until 9.470304e8")
            model = BMI.update_until(model, time)
            BMI.finalize(model)
        end

    end

    @testset "BMI run SBM in parts" begin
        tomlpath = joinpath(@__DIR__, "sbm_gw.toml")
        model = BMI.initialize(Wflow.Model, tomlpath)

        # update the recharge part of the SBM model
        model = BMI.update(model, run = "sbm_until_recharge")

        @testset "recharge part of SBM" begin
            sbm = model.vertical
            @test sbm.interception[1] ≈ 0.6299999952316284f0
            @test sbm.ustorelayerdepth[1][1] ≈ 0.0f0
            @test sbm.snow[1] ≈ 3.1912317735997524f0
            @test sbm.recharge[5] ≈ -0.0727941579914808f0
            @test sbm.zi[5] ≈ 300.0f0
        end

        # set zi and exfiltwater from external source (e.g. a groundwater model)
        BMI.set_value(
            model,
            "lateral.subsurface.zi",
            fill(0.25, BMI.get_grid_size(model, 4)),
        )
        BMI.set_value(
            model,
            "lateral.subsurface.exfiltwater",
            fill(1.0e-5, BMI.get_grid_size(model, 4)),
        )
        # update SBM after subsurface flow
        model = BMI.update(model, run = "sbm_after_subsurfaceflow")

        @testset "SBM after subsurface flow" begin
            sbm = model.vertical
            sub = model.lateral.subsurface
            @test sbm.interception[1] ≈ 0.6299999952316284
            @test sbm.ustorelayerdepth[1][1] ≈ 0.0f0
            @test sbm.snow[1] ≈ 3.1912317735997524f0
            @test sbm.recharge[5] ≈ -0.0727941579914808f0
            @test sbm.zi[5] ≈ 250.0f0
            @test sub.zi[5] ≈ 0.25f0
            @test sub.exfiltwater[1] ≈ 1.0f-5
            @test sub.ssf[1] ≈ 0.0f0
        end

        BMI.finalize(model)
    end
end


# ## Animation
z = get_maximal_state(storage)
z = [[z[1] for t = 1:100]..., z..., [z[end] for t = 1:100]...]
T = length(z)
anim = MeshCat.Animation(convert(Int, floor(1.0 / 0.01)))
build_robot(mech, vis=vis)
for t = 1:T
    MeshCat.atframe(anim, t) do
        set_robot(vis, mech, z[t])
    end
end
MeshCat.setanimation!(vis, anim)
set_camera!(vis, cam_pos=[3,-0,0], zoom=0.84)

set_floor!(vis, x=0.2, y=10, z=0.02, color=RGBA(0,0,0,1))

# ## HMTL scene
render_static(vis)
open(joinpath(@__DIR__, "atlas_drop.html"), "w") do file
    write(file, static_html(vis))
end


# test various training parameters

bigdata = joinpath("..", "data", "big.txt")

vector_size = rand(100:300)
window = rand(1:5)
sample = [0, 1e-5][rand(1:2)]
hs = rand(0:1)
negative = rand(1:5)
threads = rand(1:16)
iter = rand(1:5)
min_count = rand(1:5)
alpha = 0.025
debug = 2
threshold = rand(50:100)
binary = 1
cbow = rand(0:1)
classes = rand(30:100)
save_vocab = "vocab.txt"
read_vocab = "vocab.txt"

word2phrase(bigdata, "bigphrase", min_count=min_count, threshold=threshold,
            debug=debug)

word2vec("bigphrase", "bigvecs.txt", size=vector_size, window=window, sample=sample,
         hs=hs, negative=negative, threads=threads, iter=iter,
         min_count=min_count, alpha=alpha, debug=debug, cbow=cbow,
         save_vocab=save_vocab)
word2vec(bigdata, "bigvecs.bin", size=vector_size, window=window,
         sample=sample, hs=hs, negative=negative, threads=threads,
         iter=iter, binary=binary, read_vocab = read_vocab)

rm("bigphrase")

model = wordvectors("bigvecs.txt")
len_vecs, num_words = size(model)
@test len_vecs == vector_size

word2clusters(bigdata, "bigclus.txt", classes, size=vector_size, window=window,
              sample=sample, hs=hs, negative=negative, threads=threads,
              iter=iter, min_count=min_count, alpha=alpha, debug=debug,
              cbow=cbow, read_vocab=read_vocab, save_vocab=save_vocab)

println("training passed test...")


"""
    $(TYPEDEF)

A struct used to reference the `equilibrator_api` object.

$(FIELDS)
"""
struct Equilibrator
    eq::PyObject
    cc::PyObject
end

function Base.show(io::IO, ::MIME"text/plain", equilibrator::Equilibrator)
    v = string(equilibrator.eq.__version__)
    println(crayon"blue bold", "Interface to equilibrator_api v$v", crayon"default !bold")
    println(
        crayon"yellow",
        "Temperature:\t",
        crayon"default",
        get_temperature(equilibrator),
    )
    println(
        crayon"yellow",
        "Ionic strength: ",
        crayon"default",
        get_ionic_strength(equilibrator),
    )
    println(crayon"yellow", "pH:\t\t", crayon"default", get_pH(equilibrator))
    println(crayon"yellow", "pMg:\t\t", crayon"default", get_pMg(equilibrator))
end

"""
    $(SIGNATURES)

A constructor for `Equilibrator`. Keyword arguments `ionic_strength` and `temperature` have units.
Any appropriate unit may be used, the conversion is handled internally. This is typically
slow to load, ~ 30 seconds, since all the thermodynamic data needs to be loaded into memory.

# Example
```
using eQuilibrator
using Unitful

equilibrator = Equilibrator(pH=6.0, pMg = 2.5, ionic_strength=250.0u"mM", temperature=30.0u"°C")
```
"""
function Equilibrator(;
    pH = 7.5,
    pMg = 3.0,
    ionic_strength = 0.25u"M",
    temperature = 298.15u"K",
)
    eq = pyimport("equilibrator_api")
    cc = eq.ComponentContribution()

    cc.p_h = eq.Q_(pH)
    cc.p_mg = eq.Q_(pMg)

    _ionic_strength = uconvert(u"M", float(ionic_strength))
    cc.ionic_strength = eq.Q_(string(_ionic_strength))

    _temperature = uconvert(u"K", float(temperature))
    cc.temperature = eq.Q_(string(_temperature))
    return Equilibrator(eq, cc)
end

"""
    $(SIGNATURES)

Reset `equilibrator` back to the default starting state with `pH = 7.5`,
`pMg = 3.0`, `temperature = 298 K` and `ionic_strength = 0.25 M`.
"""
function reset_equilibrator(equilibrator::Equilibrator)
    set_temperature(equilibrator, 298.15u"K")
    set_pH(equilibrator, 7.5)
    set_pMg(equilibrator, 3.0)
    set_ionic_strength(equilibrator, 0.25u"M")
end

"""
    $(SIGNATURES)

Get the temperature of `equilibrator`.

# Example
```
get_temperature(equilibrator)
```
"""
function get_temperature(equilibrator::Equilibrator)
    units = uparse(_parse_units(py"str"(equilibrator.cc.temperature.units)))
    (equilibrator.cc.temperature.m)units
end

"""
    $(SIGNATURES)

Set the temperature of `equilibrator` to `x`.

# Example
```
set_temperature(equilibrator, 40u"°C")
```
"""
function set_temperature(equilibrator::Equilibrator, x)
    _temperature = uconvert(u"K", float(x))
    equilibrator.cc.temperature = equilibrator.eq.Q_(string(_temperature))
    return nothing
end

"""
    $(SIGNATURES)

Get the ionic strength of `equilibrator`.

# Example
```
get_ionic_strength(equilibrator)
```
"""
function get_ionic_strength(equilibrator::Equilibrator)
    units = uparse(_parse_units(py"str"(equilibrator.cc.ionic_strength.units)))
    (equilibrator.cc.ionic_strength.m)units
end

"""
    $(SIGNATURES)

Set the ionic strength of `equilibrator` to `x`.

# Example
```
set_ionic_strength(equilibrator, 100u"mM")
```
"""
function set_ionic_strength(equilibrator::Equilibrator, x)
    _ionic_strength = uconvert(u"M", float(x))
    equilibrator.cc.ionic_strength = equilibrator.eq.Q_(string(_ionic_strength))
    return nothing
end

"""
    $(SIGNATURES)

Get the pH from `equilibrator`.

# Example
```
get_pH(equilibrator)
```
"""
function get_pH(equilibrator::Equilibrator)
    equilibrator.cc.p_h.m
end

"""
    $(SIGNATURES)

Set pH of `equilibrator` to `x`.

# Example
```
set_pH(equilibrator, 3.5)
```
"""
function set_pH(equilibrator::Equilibrator, x)
    equilibrator.cc.p_h = equilibrator.eq.Q_(float(x))
    return nothing
end

"""
    $(SIGNATURES)

Get pMg from `equilibrator`.

# Example
```
get_pMg(equilibrator)
```
"""
function get_pMg(equilibrator::Equilibrator)
    equilibrator.cc.p_mg.m
end

"""
    $(SIGNATURES)

Set pMg of `equilibrator` to `x`.

# Example
```
set_pMg(equilibrator, 1.2)
```
"""
function set_pMg(equilibrator::Equilibrator, x)
    equilibrator.cc.p_mg = equilibrator.eq.Q_(float(x))
    return nothing
end


fileid = parse(Int64, ARGS[1])
simregexp = Regex("postsim.*_.*$(fileid).*.bson")

using GPMTD
using SparseProbVec

using BSON
using BSON: @load

using RCall
using StatsBase

## for compatibility with the imported types
using Random
using Random123
using PDMats
using LinearAlgebra
using Statistics
using StatsBase
using Dates

files = filter(x -> occursin(simregexp, x), readdir("./postsim"))
nfiles = length(files)

i = 1
for i in 1:nfiles
    file_now = files[i]
    println("starting file: $(file_now)")

    BSON.@load "./postsim/$(file_now)" model sims init mesg
    nsim = length(sims)

    ### Send to R
    R"rm(list=ls())"
    iter = model.state.iter
    accptr =  model.state.accpt ./ model.state.iter
    L = model.L
    n = model.n
    priorinfo = "$(model.prior)"
    initinfo = "$(init)"

    sims_llik = [ sims[i][:llik] for i=1:nsim ]
    # sims_Scounts = permutedims( hcat( [ StatsBase.counts( sims[i][:ζ], 0:L ) for i = 1:nsim ]... ) )

    sims_lam = permutedims(hcat([ exp.(sims[i][:lλ]) for i=1:nsim ]...))
    sims_int_mu = [ sims[i][:intercept].μ for i=1:nsim ]
    sims_int_sig2 = [ sims[i][:intercept].σ2 for i=1:nsim ]

    @rput mesg nsim iter accptr L n priorinfo initinfo sims_llik sims_lam sims_int_mu sims_int_sig2

    R"ls()"
    R"save.image(file=paste0('postsim/mcmc_', $(mesg), '.rda'))"

    cp("./postsim/$(file_now)", "./postsimB/$(file_now)")
    rm("./postsim/$(file_now)")

    println("completed file number: $(i) of $(nfiles) \n\n")
end


@testset "Constraints" begin
    using ONNXNaiveNASflux.NaiveNASflux
    using ONNXNaiveNASflux: SizePseudoTransparent, MeasureNout
    using NaiveNASflux: actdim

    dv(name, invertex, outsize) = dv(name, invertex, nout(invertex), outsize)
    dv(name, invertex, insize, outsize) = fluxvertex(name, Dense(insize, outsize), invertex)
    cv(name, invertex, outsize) = cv(name, invertex, nout(invertex), outsize)
    cv(name, invertex, insize, outsize) =fluxvertex(name, Conv((3,3), insize => outsize, pad=(1,1)), invertex)

    @testset "Reshape" begin
        using ONNXNaiveNASflux: Reshape

        rv(name, invertex, dims) = absorbvertex(name, MeasureNout(Reshape(dims)), invertex; traitdecoration=SizePseudoTransparent)

        @testset "Reshape 2D -> 4D variable batch" begin
            v0 = inputvertex("in", 3)
            v1 = dv("v1", v0, 4)
            v2 = rv("v2", v1, (5,2,3,Colon()))
            v3 = cv("v3", v2, 4)

            @test actdim(v2) == [actdim(layer(v2))] == [3]
            @test layer(v2) isa Reshape

            g = CompGraph(v0, v3)

            @test size(g(ones(Float32, 3, 15))) == (5,2,4,2)

            @test_logs (:warn, r"Could not change nout") Δnout!(v1, 3)

            @test nout.(vertices(g)) == [3, 6, 3, 4]
            @test size(g(ones(Float32, 3, 15))) == (5,2,4,3)
        end

        @testset "Reshape 2D -> 4D variable shape" begin
            v0 = inputvertex("in", 3)
            v1 = dv("v1", v0, 4)
            v2 = rv("v2", v1, (5, Colon(),3, 2))
            v3 = cv("v3", v2, 4)

            @test actdim(v2) == [actdim(layer(v2))] == [3]
            @test layer(v2) isa Reshape

            g = CompGraph(v0, v3)

            @test size(g(ones(Float32, 3, 15))) == (5,2,4,2)

            @test_logs (:warn, r"Could not change nout") Δnout!(v1, 3)

            @test nout.(vertices(g)) == [3, 6, 3, 4]
            @test size(g(ones(Float32, 3, 15))) == (5,3,4,2)
        end

        @testset "Reshape 2D -> 4D variable nout" begin
            v0 = inputvertex("in", 3)
            v1 = dv("v1", v0, 4)
            v2 = rv("v2", v1, (5, 3, Colon(), 2))
            v3 = cv("v3", v2, 2, 4)

            @test actdim(v2) == [actdim(layer(v2))] == [3]
            @test layer(v2) isa Reshape

            g = CompGraph(v0, v3)

            @test size(g(ones(Float32, 3, 15))) == (5,3,4,2)

            # note that this works only because MeasureNout got the size when we evaluated the graph absorbvertex
            # In other cases one might need try_infer_sizes!
            @test_logs (:warn, r"Could not change nout") Δnout!(v1, 3)

            @test nout.(vertices(g)) == [3, 6, 3, 4]
            @test size(g(ones(Float32, 3, 15))) == (5,3,4,2)
        end

        @testset "Reshape 2D -> 4D variable nout 0-dim" begin
            v0 = inputvertex("in", 3)
            v1 = dv("v1", v0, 20)
            v2 = rv("v2", v1, (2, 5, Colon(), 0))
            v3 = cv("v3", v2, 2, 4)

            @test actdim(v2) == [actdim(layer(v2))] == [3]
            @test layer(v2) isa Reshape

            g = CompGraph(v0, v3)

            @test size(g(ones(Float32, 3, 2))) == (2,5,4,2)

            # note that this works only because MeasureNout got the size when we evaluated the graph absorbvertex
            # In other cases one might need try_infer_sizes!
            @test_logs (:warn, r"Could not change nout") Δnout!(v1, 3)

            @test nout.(vertices(g)) == [3, 30, 3, 4]
            @test size(g(ones(Float32, 3, 2))) == (2,5,4,2)
        end

        @testset "Reshape 2D -> 4D all fixed" begin
            v0 = inputvertex("in", 3)
            v1 = dv("v1", v0, 4)
            v2 = rv("v2", v1, (5, 2, 3, 2))
            v3 = cv("v3", v2, 4)

            @test actdim(v2) == [actdim(layer(v2))] == [3]
            @test layer(v2) isa Reshape

            g = CompGraph(v0, v3)

            @test size(g(ones(Float32, 3, 15))) == (5,2,4,2)

            @test @test_logs((:warn, r"Could not change nout"), match_mode=:any, Δnout!(v1, 3)) == false

            @test nout.(vertices(g)) == [3, 4, 3, 4]
            @test size(g(ones(Float32, 3, 15))) == (5,2,4,2)
        end

        @testset "Reshape 2D -> 4D all fixed 0-dim" begin
            v0 = inputvertex("in", 3)
            v1 = dv("v1", v0, 4)
            v2 = rv("v2", v1, (5, 2, 0, 2))
            v3 = cv("v3", v2, 4, 4)

            @test actdim(v2) == [actdim(layer(v2))] == [3]
            @test layer(v2) isa Reshape

            g = CompGraph(v0, v3)

            @test size(g(ones(Float32, 3, 20))) == (5,2,4,2)

            @test_logs (:warn, r"Could not change nout") Δnout!(v1, 3)

            @test nout.(vertices(g)) == [3, 6, 6, 4]
            @test size(g(ones(Float32, 3, 20))) == (5,2,4,2)
        end

        @testset "Reshape 4D -> 2D variable batch" begin
            v0 = inputvertex("in", 3)
            v1 = cv("v1", v0, 4)
            v2 = rv("v2", v1, (5 ,Colon()))
            v3 = dv("v3", v2, 4)

            @test actdim(v2) == [actdim(layer(v2))] == [1]
            @test layer(v2) isa Reshape

            g = CompGraph(v0, v3)

            @test size(g(ones(Float32, 5,2,3,2))) == (4, 16)

            @test_logs (:warn, r"Could not change nout") Δnout!(v1, 3)

            @test nout.(vertices(g)) == [3, 8, 10, 4]
            @test size(g(ones(Float32, 5,2,3,2))) == (4, 16)
        end

        @testset "Reshape 4D -> 2D variable nout" begin
            v0 = inputvertex("in", 3)
            v1 = cv("v1", v0, 4)
            v2 = rv("v2", v1, (Colon(), 16))
            v3 = dv("v3", v2, 5, 4)

            @test actdim(v2) == [actdim(layer(v2))] == [1]
            @test layer(v2) isa Reshape

            g = CompGraph(v0, v3)

            @test size(g(ones(Float32, 5,2,3,2))) == (4, 16)

            @test_logs (:warn, r"Could not change nout") Δnout!(v1, 3)

            @test nout.(vertices(g)) == [3, 16, 20, 4]
            @test size(g(ones(Float32, 5,2,3,2))) == (4, 16)
        end

        @testset "Reshape 4D -> 2D variable nout 0-dim" begin
            v0 = inputvertex("in", 3)
            v1 = cv("v1", v0, 4)
            v2 = rv("v2", v1, (Colon(), 0))
            v3 = dv("v3", v2, 40, 4)

            @test actdim(v2) == [actdim(layer(v2))] == [1]
            @test layer(v2) isa Reshape

            g = CompGraph(v0, v3)

            @test size(g(ones(Float32, 5,2,3,2))) == (4, 2)

            Δnout!(v1, -2)

            @test nout.(vertices(g)) == [3, 2, 20, 4]
            @test size(g(ones(Float32, 5,2,3,2))) == (4, 2)
        end
    end

    @testset "Flatten" begin
        import ONNXNaiveNASflux: Flatten
        fv(name, invertex, dim) = absorbvertex(name, MeasureNout(Flatten(dim)), invertex; traitdecoration=SizePseudoTransparent)

        function tg(outsize, dim)
            v0 = inputvertex("in", 3)
            v1 = cv("v1", v0, 4)
            v2 = fv("v2", v1, dim)
            v3 = dv("v3", v2, outsize, 4)

            @test actdim(v2) == [actdim(layer(v2))] == [1]
            @test layer(v2) isa Flatten

            g = CompGraph(v0, v3)
        end

        @testset "Flatten 4D dim 0" begin
            g = tg(80, 4)

            @test size(g(ones(Float32, 5,2,3,2))) == (4, 1)

            Δnout!(vertices(g)[2], -3)

            @test nout.(vertices(g)) == [3, 1, 20, 4]
            @test size(g(ones(Float32, 5,2,3,2))) == (4, 1)
        end

        @testset "Flatten 4D dim 1" begin
            g = tg(5, 1)

            @test size(g(ones(Float32, 5,2,3,2))) == (4, 16)

            Δnout!(vertices(g)[2], 3)

            @test nout.(vertices(g)) == [3, 7, 5, 4]
            @test size(g(ones(Float32, 5,2,3,2))) == (4, 28)
        end

        @testset "Flatten 4D dim 2" begin
            g = tg(10, 2)

            @test size(g(ones(Float32, 5,2,3,2))) == (4, 8)

            Δnout!(vertices(g)[2], -1)

            @test nout.(vertices(g)) == [3, 3, 10, 4]
            @test size(g(ones(Float32, 5,2,3,2))) == (4, 6)
        end

        @testset "Flatten 4D dim 3" begin
            g = tg(40, 3)

            @test size(g(ones(Float32, 5,2,3,2))) == (4, 2)

            Δnout!(vertices(g)[2], -2)

            @test nout.(vertices(g)) == [3, 2, 20, 4]
            @test size(g(ones(Float32, 5,2,3,2))) == (4, 2)
        end

        @testset "Flatten 4D dim 4" begin
            g = tg(80, 4)

            @test size(g(ones(Float32, 5,2,3,2))) == (4, 1)

            Δnout!(vertices(g)[2], -3)

            @test nout.(vertices(g)) == [3, 1, 20, 4]
            @test size(g(ones(Float32, 5,2,3,2))) == (4, 1)
        end

        @testset "Flatten 4D dim -4" begin
            g = tg(1, -4)

            @test size(g(ones(Float32, 5,2,3,2))) == (4, 80)

            Δnout!(vertices(g)[2], -2)

            @test nout.(vertices(g)) == [3, 2, 1, 4]
            @test size(g(ones(Float32, 5,2,3,2))) == (4, 40)
        end

        @testset "Flatten 4D dim -3" begin
            g = tg(5, -3)

            @test size(g(ones(Float32, 5,2,3,2))) == (4, 16)

            Δnout!(vertices(g)[2], 3)

            @test nout.(vertices(g)) == [3, 7, 5, 4]
            @test size(g(ones(Float32, 5,2,3,2))) == (4, 28)
        end

        @testset "Flatten 4D dim -2" begin
            g = tg(10, -2)

            @test size(g(ones(Float32, 5,2,3,2))) == (4, 8)

            Δnout!(vertices(g)[2], -1)

            @test nout.(vertices(g)) == [3, 3, 10, 4]
            @test size(g(ones(Float32, 5,2,3,2))) == (4, 6)
        end

        @testset "Flatten 4D dim -1" begin
            g = tg(40, -1)

            @test size(g(ones(Float32, 5,2,3,2))) == (4, 2)

            Δnout!(vertices(g)[2], -2)

            @test nout.(vertices(g)) == [3, 2, 20, 4]
            @test size(g(ones(Float32, 5,2,3,2))) == (4, 2)
        end
    end
end


# TODO add quantile estimation as indices
function buildarray_annual(C::ClimateTools.ClimGrid, dataout, numYears)
    lonsymbol = Symbol(C.dimension_dict["lon"])
    latsymbol = Symbol(C.dimension_dict["lat"])
    FD = AxisArray(dataout, Axis{lonsymbol}(C[1][Axis{lonsymbol}][:]), Axis{latsymbol}(C[1][Axis{latsymbol}][:]), Axis{:time}(Dates.year.(DateTime.(numYears))))
    return FD
end

function buildarray_climato(C::ClimateTools.ClimGrid, dataout)
    lonsymbol = Symbol(C.dimension_dict["lon"])
    latsymbol = Symbol(C.dimension_dict["lat"])
    FD = AxisArray(dataout, Axis{lonsymbol}(C[1][Axis{lonsymbol}][:]), Axis{latsymbol}(C[1][Axis{latsymbol}][:]))
    return FD
end

function buildarray_resample(C::ClimateTools.ClimGrid, dataout, newtime)
    lonsymbol = Symbol(C.dimension_dict["lon"])
    latsymbol = Symbol(C.dimension_dict["lat"])
    FD = AxisArray(dataout, Axis{lonsymbol}(C[1][Axis{lonsymbol}][:]), Axis{latsymbol}(C[1][Axis{latsymbol}][:]), Axis{:time}(newtime))
    return FD
end

# function buildarray_orig(C::ClimateTools.ClimGrid, dataout)
#     lonsymbol = Symbol(C.dimension_dict["lon"])
#     latsymbol = Symbol(C.dimension_dict["lat"])
#     FD = AxisArray(dataout, Axis{lonsymbol}(C[1][Axis{lonsymbol}][:]), Axis{latsymbol}(C[1][Axis{latsymbol}][:]), Axis{:time}(get_timevec(C)))
#     return FD
# end

"""
    prcp1(C::ClimGrid)

Annual number with preciptation >= 1 mm. This function returns a ClimGrid. Input data should be in mm.
"""
function prcp1(C::ClimGrid)
  @argcheck C[9] == "pr"
  years    = Dates.year.(C.data[Axis{:time}][:])
  numYears = unique(years)
  dataout  = zeros(Int64, (size(C.data, 1), size(C.data, 2), length(numYears)))
  datain   = (C.data.data)

  # Indice calculation
  Threads.@threads for i in 1:length(numYears)
    idx = searchsortedfirst(years, numYears[i]):searchsortedlast(years, numYears[i])
    Base.mapreducedim!(t -> t >= 1, +, view(dataout, :, :, i:i), view(datain, :,:, idx))
  end

  # Apply mask
  dataout = applymask(dataout, C.msk)

  # Build output AxisArray
  FD = buildarray_annual(C, dataout, numYears)

  # Return climGrid type containing the indice
  return ClimGrid(FD, longrid=C.longrid, latgrid=C.latgrid, msk=C.msk, grid_mapping=C.grid_mapping, dimension_dict=C.dimension_dict, timeattrib=C.timeattrib, model=C.model, frequency="year", experiment=C.experiment, run=C.run, project=C.project, institute=C.institute, filename=C.filename, dataunits="days", latunits=C.latunits, lonunits=C.lonunits, variable="prcp1", typeofvar=C.typeofvar, typeofcal=C.typeofcal, varattribs=C.varattribs, globalattribs=C.globalattribs)
end

"""
    frostdays(C::ClimGrid)

FD, Number of frost days: Annual count of days when TN (daily minimum temperature) < 0 Celsius.

Let TN[i,j] be daily minimum temperature on day i in year j. Count the number of days where:

  TN[i,j] < 0 Celsius.
"""
function frostdays(C::ClimGrid)
  @argcheck C[9] == "tasmin"
  years    = Dates.year.(C.data[Axis{:time}][:])
  numYears = unique(years)
  dataout  = zeros(Int64, (size(C.data, 1), size(C.data, 2), length(numYears)))
  datain   = C.data.data

  # Indice calculation
  Threads.@threads for i in 1:length(numYears)
    idx = searchsortedfirst(years, numYears[i]):searchsortedlast(years, numYears[i])
    Base.mapreducedim!(t -> t < 0, +, view(dataout, :, :, i:i), view(datain, :,:, idx))
  end

  # Apply mask
  dataout = applymask(dataout, C.msk)

  # Build output AxisArray
  FD = buildarray_annual(C, dataout, numYears)

  # Return climGrid type containing the indice
  return ClimGrid(FD, longrid=C.longrid, latgrid=C.latgrid, msk=C.msk, grid_mapping=C.grid_mapping, dimension_dict=C.dimension_dict, timeattrib=C.timeattrib, model=C.model, frequency="year", experiment=C.experiment, run=C.run, project=C.project, institute=C.institute, filename=C.filename, dataunits="days", latunits=C.latunits, lonunits=C.lonunits, variable="frostdays", typeofvar=C.typeofvar, typeofcal=C.typeofcal, varattribs=C.varattribs, globalattribs=C.globalattribs)
end

"""
    summerdays(C::ClimGrid)

SD, Number of summer days: Annual count of days when TX (daily maximum temperature) > 25 degree Celsius.

Let TX[i,j] be daily maximum temperature on day i in year j. Count the number of days where:

  TX[i,j] > 25 Celsius.
"""
function summerdays(C::ClimGrid)
  @argcheck C[9] == "tasmax"
  years    = Dates.year.(C.data[Axis{:time}][:])
  numYears = unique(years)
  dataout  = zeros(Int64, (size(C.data, 1), size(C.data, 2), length(numYears)))
  datain   = C.data.data

  # Indice calculation
  Threads.@threads for i in 1:length(numYears)
    idx = searchsortedfirst(years, numYears[i]):searchsortedlast(years, numYears[i])
    Base.mapreducedim!(t -> t > 25, +, view(dataout, :, :, i:i), view(datain, :,:, idx))
  end

  # Apply mask
  dataout = applymask(dataout, C.msk)

  # Build output AxisArray
  FD = buildarray_annual(C, dataout, numYears)

  # Return climGrid type containing the indice
  return ClimGrid(FD, longrid=C.longrid, latgrid=C.latgrid, msk=C.msk, grid_mapping=C.grid_mapping, dimension_dict=C.dimension_dict, timeattrib=C.timeattrib, model=C.model, frequency="year", experiment=C.experiment, run=C.run, project=C.project, institute=C.institute, filename=C.filename, dataunits="days", latunits=C.latunits, lonunits=C.lonunits, variable="summerdays", typeofvar=C.typeofvar, typeofcal=C.typeofcal, varattribs=C.varattribs, globalattribs=C.globalattribs)
end

"""
    icingdays(C::ClimGrid)

ID, Number of summer days: Annual count of days when TX (daily maximum temperature) < 0 degree Celsius.

Let TX[i,j] be daily maximum temperature on day i in year j. Count the number of days where:

  TX[i,j] < 0 Celsius.
"""
function icingdays(C::ClimGrid)
  @argcheck C[9] == "tasmax"
  years    = Dates.year.(C.data[Axis{:time}][:])
  numYears = unique(years)
  dataout  = zeros(Int64, (size(C.data, 1), size(C.data, 2), length(numYears)))
  datain   = C.data.data

  # Indice calculation
  Threads.@threads for i in 1:length(numYears)
    idx = searchsortedfirst(years, numYears[i]):searchsortedlast(years, numYears[i])
    Base.mapreducedim!(t -> t < 0, +, view(dataout, :, :, i:i), view(datain, :,:, idx))
  end

  # Apply mask
  dataout = applymask(dataout, C.msk)

  # Build output AxisArray
  FD = buildarray_annual(C, dataout, numYears)

  # Return climGrid type containing the indice
  return ClimGrid(FD, longrid=C.longrid, latgrid=C.latgrid, msk=C.msk, grid_mapping=C.grid_mapping, dimension_dict=C.dimension_dict, timeattrib=C.timeattrib, model=C.model, frequency="year", experiment=C.experiment, run=C.run, project=C.project, institute=C.institute, filename=C.filename, dataunits="days", latunits=C.latunits, lonunits=C.lonunits, variable="icingdays", typeofvar=C.typeofvar, typeofcal=C.typeofcal, varattribs=C.varattribs, globalattribs=C.globalattribs)
end

"""
    tropicalnights(C::ClimGrid)

TropicalNights, Number of tropical nights: Annual count of days when TN (daily maximum temperature) > 20 degree Celsius.

Let TN[i,j] be daily minimum temperature on day i in year j. Count the number of days where:

  TN[i,j] > 20 Celsius.
"""
function tropicalnights(C::ClimGrid)
  @argcheck C[9] == "tasmin"
  years    = Dates.year.(C.data[Axis{:time}][:])
  numYears = unique(years)
  dataout  = zeros(Int64, (size(C.data, 1), size(C.data, 2), length(numYears)))
  datain   = C.data.data

  # Indice calculation
  Threads.@threads for i in 1:length(numYears)
    idx = searchsortedfirst(years, numYears[i]):searchsortedlast(years, numYears[i])
    Base.mapreducedim!(t -> t > 20, +, view(dataout, :, :, i:i), view(datain, :,:, idx))
  end

  # Apply mask
  dataout = applymask(dataout, C.msk)

  # Build output AxisArray
  FD = buildarray_annual(C, dataout, numYears)

  # Return climGrid type containing the indice
  return ClimGrid(FD, longrid=C.longrid, latgrid=C.latgrid, msk=C.msk, grid_mapping=C.grid_mapping, dimension_dict=C.dimension_dict, timeattrib=C.timeattrib, model=C.model, frequency="year", experiment=C.experiment, run=C.run, project=C.project, institute=C.institute, filename=C.filename, dataunits="days", latunits=C.latunits, lonunits=C.lonunits, variable="tropicalnights", typeofvar=C.typeofvar, typeofcal=C.typeofcal, varattribs=C.varattribs, globalattribs=C.globalattribs)
end

"""
    customthresover(C::ClimGrid, thres)

customthresover, annual number of days over a specified threshold.

Let TS[i,j] be a daily time serie value on day i in year j. Count the number of days where:

  TS[i,j] > thres.

"""
function customthresover(C::ClimGrid, thres)
  years    = Dates.year.(C.data[Axis{:time}][:])
  numYears = unique(years)
  dataout  = fill(NaN, (size(C.data, 1), size(C.data, 2), length(numYears)))
  datain   = C.data.data

  dataout_rshp = reshape(dataout, (size(dataout, 1)*size(dataout, 2), size(dataout, 3)))
  datain_rshp = reshape(datain, (size(datain, 1)*size(datain, 2), size(datain, 3)))

  # Indice calculation
  Threads.@threads for k in 1:size(datain_rshp, 1)
    for i in 1:length(numYears)
      idx = searchsortedfirst(years, numYears[i]):searchsortedlast(years, numYears[i])
      dataout_rshp[k, i] = sum(datain_rshp[k, idx] .> thres)
      # Base.mapreducedim!(t -> t > thres, +, view(dataout, :, :, i:i), view(datain, :,:, idx))
    end
  end

  # Apply mask
  dataout = applymask(dataout, C.msk)

  # Build output AxisArray
  FD = buildarray_annual(C, dataout, numYears)

  # Return climGrid type containing the indice
  return ClimGrid(FD, longrid=C.longrid, latgrid=C.latgrid, msk=C.msk, grid_mapping=C.grid_mapping, dimension_dict=C.dimension_dict, timeattrib=C.timeattrib, model=C.model, frequency="year", experiment=C.experiment, run=C.run, project=C.project, institute=C.institute, filename=C.filename, dataunits="days", latunits=C.latunits, lonunits=C.lonunits, variable=string("Days over ", thres), typeofvar=C.typeofvar, typeofcal=C.typeofcal, varattribs=C.varattribs, globalattribs=C.globalattribs)
end

"""
    customthresunder(C::ClimGrid, thres)

customthresunder, annual number of days under a specified threshold.

Let TS[i,j] be a daily time serie value on day i in year j. Count the number of days where:

  TS[i,j] < thres.

"""
function customthresunder(C::ClimGrid, thres)
  years    = Dates.year.(C.data[Axis{:time}][:])
  numYears = unique(years)
  datain   = C.data.data
  dataout  = zeros(Int64, (size(C.data, 1), size(C.data, 2), length(numYears)))

  # Indice calculation
  Threads.@threads for i in 1:length(numYears)
    idx = searchsortedfirst(years, numYears[i]):searchsortedlast(years, numYears[i])
    Base.mapreducedim!(t -> t < thres, +, view(dataout, :, :, i:i), view(datain, :,:, idx))
  end

  # Apply mask
  dataout = applymask(dataout, C.msk)

  # Build output AxisArray
  FD = buildarray_annual(C, dataout, numYears)

  # Return climGrid type containing the indice
  return ClimGrid(FD, longrid=C.longrid, latgrid=C.latgrid, msk=C.msk, grid_mapping=C.grid_mapping, dimension_dict=C.dimension_dict, timeattrib=C.timeattrib, model=C.model, frequency="year", experiment=C.experiment, run=C.run, project=C.project, institute=C.institute, filename=C.filename, dataunits="days", latunits=C.latunits, lonunits=C.lonunits, variable=string("Days under ", thres), typeofvar=C.typeofvar, typeofcal=C.typeofcal, varattribs=C.varattribs, globalattribs=C.globalattribs)
end

"""
    annualmax(C::ClimGrid)

Annual maximum of array data.

Let data[i,j] be daily time serie on day i in year j. Extract the highest value for year j.
"""
function annualmax(C::ClimGrid)
  years    = Dates.year.(C.data[Axis{:time}][:])
  numYears = unique(years)
  datain   = C.data.data
  dataout  = zeros(typeof(datain[1]), (size(C.data, 1), size(C.data, 2), length(numYears)))

  # Indice calculation
  Threads.@threads for i in 1:length(numYears)
    idx = searchsortedfirst(years, numYears[i]):searchsortedlast(years, numYears[i])
    Statistics.maximum!(view(dataout, :, :, i:i), view(datain, :,:, idx))
  end

  # Apply mask
  dataout = applymask(dataout, C.msk)

  # Build output AxisArray
  FD = buildarray_annual(C, dataout, numYears)

  # Return climGrid type containing the indice
  return ClimGrid(FD, longrid=C.longrid, latgrid=C.latgrid, msk=C.msk, grid_mapping=C.grid_mapping, dimension_dict=C.dimension_dict, timeattrib=C.timeattrib, model=C.model, frequency="year", experiment=C.experiment, run=C.run, project=C.project, institute=C.institute, filename=C.filename, dataunits=C.dataunits, latunits=C.latunits, lonunits=C.lonunits, variable="annualmax", typeofvar=C.typeofvar, typeofcal=C.typeofcal, varattribs=C.varattribs, globalattribs=C.globalattribs)
end

"""
    annualmin(C::ClimGrid)

Annual minimum of array data.

Let data[i,j] be daily time serie on day i in year j. Extract the lowest value for year j.
"""
function annualmin(C::ClimGrid)
  years    = Dates.year.(C.data[Axis{:time}][:])
  numYears = unique(years)
  datain   = C.data.data
  dataout  = zeros(typeof(datain[1]), (size(C.data, 1), size(C.data, 2), length(numYears)))

  # Indice calculation
  Threads.@threads for i in 1:length(numYears)
    idx = searchsortedfirst(years, numYears[i]):searchsortedlast(years, numYears[i])
    Statistics.minimum!(view(dataout, :, :, i:i), view(datain, :,:, idx))
  end

  # Apply mask
  dataout = applymask(dataout, C.msk)

  # Build output AxisArray
  FD = buildarray_annual(C, dataout, numYears)

  # Return climGrid type containing the indice
  return ClimGrid(FD, longrid=C.longrid, latgrid=C.latgrid, msk=C.msk, grid_mapping=C.grid_mapping, dimension_dict=C.dimension_dict, timeattrib=C.timeattrib, model=C.model, frequency="year", experiment=C.experiment, run=C.run, project=C.project, institute=C.institute, filename=C.filename, dataunits=C.dataunits, latunits=C.latunits, lonunits=C.lonunits, variable="annualmin", typeofvar=C.typeofvar, typeofcal=C.typeofcal, varattribs=C.varattribs, globalattribs=C.globalattribs)
end

"""
    annualmean(C::ClimGrid)

Annual mean of array data.

Let data[i,j] be daily time serie on day i in year j. Calculate the mean value for year j.
"""
function annualmean(C::ClimGrid)
  years    = Dates.year.(C.data[Axis{:time}][:])
  numYears = unique(years)
  datain   = C.data.data
  dataout  = zeros(typeof(datain[1]), (size(C.data, 1), size(C.data, 2), length(numYears)))

  # Indice calculation
  Threads.@threads for i in 1:length(numYears)
    idx = searchsortedfirst(years, numYears[i]):searchsortedlast(years, numYears[i])
    Statistics.mean!(view(dataout, :, :, i:i), view(datain, :,:, idx))
  end

  # Apply mask
  dataout = applymask(dataout, C.msk)

  # Build output AxisArray
  FD = buildarray_annual(C, dataout, numYears)

  # Return climGrid type containing the indice
  return ClimGrid(FD, longrid=C.longrid, latgrid=C.latgrid, msk=C.msk, grid_mapping=C.grid_mapping, dimension_dict=C.dimension_dict, timeattrib=C.timeattrib, model=C.model, frequency="year", experiment=C.experiment, run=C.run, project=C.project, institute=C.institute, filename=C.filename, dataunits=C.dataunits, latunits=C.latunits, lonunits=C.lonunits, variable="annualmean", typeofvar=C.typeofvar, typeofcal=C.typeofcal, varattribs=C.varattribs, globalattribs=C.globalattribs)
end

"""
    annualsum(C::ClimGrid)

Annual sum of array data.

Let data[i,j] be daily time serie on day i in year j. Sums daily values for year j.
"""
function annualsum(C::ClimGrid)
  years    = Dates.year.(C.data[Axis{:time}][:])
  numYears = unique(years)
  datain   = C.data.data
  dataout  = zeros(typeof(datain[1]), (size(C.data, 1), size(C.data, 2), length(numYears)))

  # Indice calculation
  Threads.@threads for i in 1:length(numYears)
    idx = searchsortedfirst(years, numYears[i]):searchsortedlast(years, numYears[i])
    Base.sum!(view(dataout, :, :, i:i), view(datain, :,:, idx))
  end

  # Apply mask
  dataout = applymask(dataout, C.msk)

  # Build output AxisArray
  FD = buildarray_annual(C, dataout, numYears)

  # Return climGrid type containing the indice
  return ClimGrid(FD, longrid=C.longrid, latgrid=C.latgrid, msk=C.msk, grid_mapping=C.grid_mapping, dimension_dict=C.dimension_dict, timeattrib=C.timeattrib, model=C.model, frequency="year", experiment=C.experiment, run=C.run, project=C.project, institute=C.institute, filename=C.filename, dataunits=C.dataunits, latunits=C.latunits, lonunits=C.lonunits, variable="annualsum", typeofvar=C.typeofvar, typeofcal=C.typeofcal, varattribs=C.varattribs, globalattribs=C.globalattribs)
end


# API (Two-stage) #
# ========================== #
"""
    instantiate(stochasticmodel::StochasticModel{2},
                scenarios::Vector{<:AbstractScenario};
                instantiation::StochasticInstantiation = UnspecifiedInstantiation(),
                optimizer = nothing;
                defer::Bool = false,
                kw...)

Instantiate a two-stage stochastic program using the model definition stored in the two-stage `stochasticmodel`, and the given collection of `scenarios`. Optionally, supply an `optimizer`. If no explicit `instantiation` is provided, the structure is induced by the optimizer. The structure is `Deterministic` by default.
"""
function instantiate(sm::StochasticModel{2},
                     scenarios::Vector{<:AbstractScenario};
                     instantiation::StochasticInstantiation = UnspecifiedInstantiation(),
                     optimizer = nothing,
                     defer::Bool = false,
                     direct_model::Bool = false,
                     kw...)
    sp = StochasticProgram(parameters(sm.parameters[1]; kw...),
                           parameters(sm.parameters[2]; kw...),
                           scenarios,
                           instantiation,
                           optimizer)
    # Generate model recipes
    sm.generator(sp)
    # Return deterministic equivalent as JuMP model if direct model has been requested
    if direct_model
        optimizer(sp) isa MOI.AbstractOptimizer || error("MOI optimizer required to create direct model.")
        model = direct_model(optimizer_constructor(sp))
        model.ext[:decisions] = IgnoreDecisions()
        _generate_deterministic_equivalent!(sp, model)
        return model
    end
    # Generate stochastic program if not deferred
    if !defer
        generate!(sp)
    end
    return sp
end
"""
    instantiate(stochasticmodel::StochasticModel{2};
                instantiation::StochasticInstantiation = UnspecifiedInstantiation(),
                optimizer = nothing,
                scenario_type::Type{S} = Scenario,
                defer::Bool = false,
                kw...) where S <: AbstractScenario

Instantiate a deferred two-stage stochastic program using the model definition stored in the two-stage `stochasticmodel` over the scenario type `S`. Optionally, supply an `optimizer`. If no explicit `instantiation` is provided, the structure is induced by the optimizer. The structure is `Deterministic` by default.
"""
function instantiate(sm::StochasticModel{2};
                     instantiation::StochasticInstantiation = UnspecifiedInstantiation(),
                     optimizer = nothing,
                     scenario_type::Type{S} = Scenario,
                     kw...) where S <: AbstractScenario
    sp = StochasticProgram(parameters(sm.parameters[1]; kw...),
                           parameters(sm.parameters[2]; kw...),
                           scenario_type,
                           instantiation,
                           optimizer)
    # Generate model recipes
    sm.generator(sp)
    return sp
end
"""
    instantiate(stochasticmodel::StochasticModel,
                scenarios::Vector{<:AbstractScenario};
                instantiation::StochasticInstantiation = UnspecifiedInstantiation(),
                optimizer = nothing,
                defer::Bool = false,
                kw...)

Instantiate a stochastic program using the model definition stored in `stochasticmodel`, and the given collection of `scenarios`. Optionally, supply an `optimizer`. If no explicit `instantiation` is provided, the structure is induced by the optimizer. The structure is `Deterministic` by default.
"""
function instantiate(sm::StochasticModel{N},
                     scenarios::NTuple{M,Vector{<:AbstractScenario}};
                     instantiation::StochasticInstantiation = UnspecifiedInstantiation(),
                     optimizer = nothing,
                     defer::Bool = false,
                     kw...) where {N,M}
    M == N - 1 || error("Inconsistent number of stages $N and number of scenario types $M")
    stages = ntuple(Val(N)) do i
        Stage(parameters(sm.parameters[i]; kw...))
    end
    sp = StochasticProgram(stages,
                           scenarios,
                           instantiation,
                           optimizer)
    # Generate model recipes
    sm.generator(sp)
    # Generate stochastic program if not deferred
    if !defer
        generate!(sp)
    end
    return sp
end
"""
    instantiate(stochasticmodel::StochasticModel,
                sampler::AbstractSampler,
                n::Integer;
                instantiation::StochasticInstantiation = UnspecifiedInstantiation(),
                optimizer = nothing,
                defer::Bool = false,
                kw...)

Generate a sampled instance of size `n` using the model stored in the two-stage `stochasticmodel`, and the provided `sampler`. Optionally, supply an `optimizer`. If no explicit `instantiation` is provided, the structure is induced by the optimizer. The structure is `Deterministic` by default.

"""
function instantiate(stochasticmodel::StochasticModel{2},
                     sampler::AbstractSampler{S},
                     n::Integer;
                     instantiation::StochasticInstantiation = UnspecifiedInstantiation(),
                     optimizer = nothing,
                     defer::Bool = false,
                     kw...) where S <: AbstractScenario
    # Create new stochastic program instance
    sp = StochasticProgram(parameters(stochasticmodel.parameters[1]; kw...),
                           parameters(stochasticmodel.parameters[2]; kw...),
                           S,
                           instantiation,
                           optimizer)
    # Generate model recipes
    stochasticmodel.generator(sp)
    # Sample n scenarios
    add_scenarios!(sp, n) do
        return sample(sampler, 1/n)
    end
    # Generate stochastic program if not deferred
    if !defer
        generate!(sp)
    end
    # Return the sample instance
    return sp
end
function instantiate(stochasticmodel::StochasticModel{2},
                     sampler::AbstractSampler{Scenario},
                     n::Integer;
                     instantiation::StochasticInstantiation = UnspecifiedInstantiation(),
                     optimizer = nothing,
                     defer::Bool = false,
                     kw...)
    # Get concrete scenario type
    S = typeof(sampler())
    # Create new stochastic program instance
    sp = StochasticProgram(parameters(stochasticmodel.parameters[1]; kw...),
                           parameters(stochasticmodel.parameters[2]; kw...),
                           S,
                           instantiation,
                           optimizer)
    # Generate model recipes
    stochasticmodel.generator(sp)
    # Sample n scenarios
    add_scenarios!(sp, n) do
        return sample(sampler, 1/n)
    end
    # Generate stochastic program if not deferred
    if !defer
        generate!(sp)
    end
    # Return the sample instance
    return sp
end
"""
    optimize!(stochasticprogram::StochasticProgram; crash::AbstractCrash = Crash.None(); kw...)

Optimize the `stochasticprogram` in expectation. If an optimizer has not been set yet (see [`set_optimizer`](@ref)), a `NoOptimizer` error is thrown. An optional crash procedure can be set to warm-start.

## Examples

The following solves the stochastic program `sp` using the L-shaped algorithm.

```julia
set_optimizer(sp, LShaped.Optimizer)
set_optimizer_attribute(sp, MasterOptimizer(), GLPK.Optimizer)
set_optimizer_attribute(sp, SubproblemOptimizer(), GLPK.Optimizer)
optimize!(sp);

# output

L-Shaped Gap  Time: 0:00:02 (6 iterations)
  Objective:       -855.8333333333358
  Gap:             0.0
  Number of cuts:  8
  Iterations:      6
```

The following solves the stochastic program `sp` using GLPK on the extended form.

```julia
using GLPK

set_optimizer(sp, GLPK.Optimizer)
optimize!(sp)

```

See also: [`VRP`](@ref)
"""
function JuMP.optimize!(stochasticprogram::TwoStageStochasticProgram; crash::AbstractCrash = Crash.None(), cache::Bool = false, kw...)
    # Throw NoOptimizer error if no recognized optimizer has been provided
    check_provided_optimizer(stochasticprogram.optimizer)
    # Ensure stochastic program has been initialized at this point
    if deferred(stochasticprogram)
        generate!(stochasticprogram)
    end
    # Crash initial decision
    x₀ = crash(stochasticprogram)
    # Switch on structure and solver type
    optimize!(structure(stochasticprogram), optimizer(stochasticprogram), x₀; kw...)
    # Cache solution (if requested)
    if cache
        cache_solution!(stochasticprogram, structure(stochasticprogram), optimizer(stochasticprogram))
    end
    return nothing
end
"""
    termination_status(stochasticprogram::StochasticProgram)

Return the reason why the solver of the `stochasticprogram` stopped.
"""
function JuMP.termination_status(stochasticprogram::StochasticProgram)
    return MOI.get(stochasticprogram, MOI.TerminationStatus())::MOI.TerminationStatusCode
end
"""
    termination_status(stochasticprogram::StochasticProgram, stage::Integer, scenario_index::Integer)

Return the reason why the solver of the node at stage `stage` and scenario `scenario_index`
stopped.
"""
function JuMP.termination_status(stochasticprogram::StochasticProgram, stage::Integer, scenario_index::Integer)
    attr = ScenarioDependentModelAttribute(stage, scenario_index, MOI.TerminationStatus())
    return MOI.get(stochasticprogram, attr)::MOI.TerminationStatusCode
end
"""
    termination_status(stochasticprogram::TwoStageStochasticProgram, scenario_index::Integer)

Return the reason why the solver of scenario `scenario_index` stopped.
"""
function JuMP.termination_status(stochasticprogram::TwoStageStochasticProgram, scenario_index::Integer)
    return terminationstatus(stochasticprogram, 2, scenario_index)
end
"""
    raw_status(stochasticprogram::StochasticProgram)

Return the reason why the solver of the `stochasticprogram` stopped in its own words.
"""
function JuMP.raw_status(stochasticprogram::StochasticProgram)
    return MOI.get(stochasticprogram, MOI.RawStatusString())
end
"""
    raw_status(stochasticprogram::StochasticProgram, stage::Integer, scenario_index::Integer)

Return the reason why the solver of the node at stage `stage` and scenario `scenario_index`
stopped in its own words.
"""
function JuMP.raw_status(stochasticprogram::StochasticProgram, stage::Integer, scenario_index::Integer)
    attr = ScenarioDependentModelAttribute(stage, scenario_index, MOI.RawStatusString())
    return MOI.get(stochasticprogram, attr)
end
"""
    raw_status(stochasticprogram::TwoStageStochasticProgram, scenario_index::Integer)

Return the reason why the solver of scenario `scenario_index` stopped in its own words.
"""
function JuMP.raw_status(stochasticprogram::TwoStageStochasticProgram, scenario_index::Integer)
    return raw_status(stochasticprogram, 2, scenario_index)
end
"""
    primal_status(stochasticprogram::StochasticProgram; result::Int = 1)

Return the status of the most recent primal solution of the solver of
the `stochasticprogram`.
"""
function JuMP.primal_status(stochasticprogram::StochasticProgram; result::Int = 1)
    return MOI.get(stochasticprogram, MOI.PrimalStatus(result))::MOI.ResultStatusCode
end
"""
    primal_status(stochasticprogram::StochasticProgram, stage::Integer, scenario_index::Integer; result::Int = 1)

Return the status of the most recent primal solution of the solver of
the node at stage `stage` and scenario `scenario_index`.
"""
function JuMP.primal_status(stochasticprogram::StochasticProgram, stage::Integer, scenario_index::Integer; result::Int = 1)
    attr = ScenarioDependentModelAttribute(stage, scenario_index, MOI.PrimalStatus(result))
    return MOI.get(stochasticprogram, attr)::MOI.ResultStatusCode
end
"""
    primal_status(stochasticprogram::TwoStageStochasticProgram, scenario_index::Integer; result::Int = 1)

Return the status of the most recent primal solution of the solver of
the scenario `scenario_index`.
"""
function JuMP.primal_status(stochasticprogram::TwoStageStochasticProgram, scenario_index::Integer; result::Int = 1)
    return primal_status(stochasticprogram, 2, scenario_index; result = result)
end
"""
    dual_status(stochasticprogram::StochasticProgram; result::Int = 1)

Return the status of the most recent dual solution of the solver of
the `stochasticprogram`.
"""
function JuMP.dual_status(stochasticprogram::StochasticProgram; result::Int = 1)
    return MOI.get(stochasticprogram, MOI.DualStatus(result))::MOI.ResultStatusCode
end
"""
    dual_status(stochasticprogram::StochasticProgram, stage::Integer, scenario_index::Integer; result::Int = 1)

Return the status of the most recent dual solution of the solver of
the node at stage `stage` and scenario `scenario_index`.
"""
function JuMP.dual_status(stochasticprogram::StochasticProgram, stage::Integer, scenario_index::Integer; result::Int = 1)
    attr = ScenarioDependentModelAttribute(stage, scenario_index, MOI.DualStatus(result))
    return MOI.get(stochasticprogram, attr)::MOI.ResultStatusCode
end
"""
    dual_status(stochasticprogram::TwoStageStochasticProgram, scenario_index::Integer; result::Int = 1)

Return the status of the most recent dual solution of the solver of
scenario `scenario_index`.
"""
function JuMP.dual_status(stochasticprogram::TwoStageStochasticProgram, scenario_index::Integer; result::Int = 1)
    return dual_status(stochasticprogram, 2, scenario_index; result = result)
end
"""
    solve_time(stochasticprogram::StochasticProgram, stage::Integer, scenario_index::Integer)

If available, returns the solve time reported by the solver of
the `stochasticprogram`.
"""
function JuMP.solve_time(stochasticprogram::StochasticProgram)
    return MOI.get(stochasticprogram, MOI.SolveTime())
end
"""
    solve_time(stochasticprogram::StochasticProgram, stage::Integer, scenario_index::Integer)

If available, returns the solve time reported by the solver of
the node at stage `stage` and scenario `scenario_index`.
"""
function JuMP.solve_time(stochasticprogram::StochasticProgram, stage::Integer, scenario_index::Integer)
    attr = ScenarioDependentModelAttribute(stage, scenario_index, MOI.SolveTime())
    return MOI.get(stochasticprogram, attr)
end
"""
    solve_time(stochasticprogram::TwoStageStochasticProgram, scenario_index::Integer)

If available, returns the solve time reported by the solver of
scenario `scenario_index`.
"""
function JuMP.solve_time(stochasticprogram::TwoStageStochasticProgram, scenario_index::Integer)
    return solve_time(stochasticprogram, 2, scenario_index)
end
"""
    optimal_decision(stochasticprogram::StochasticProgram)

Return the optimal first stage decision of `stochasticprogram`, after a call to `optimize!(stochasticprogram)`.
"""
function optimal_decision(stochasticprogram::StochasticProgram)
    if termination_status(stochasticprogram) == MOI.OPTIMIZE_NOT_CALLED
        throw(OptimizeNotCalled())
    end
    return JuMP.value.(all_decision_variables(stochasticprogram, 1))
end
"""
    optimal_decision(stochasticprogram::StochasticProgram, stage::Integer, scenario_index::Integer)

Return the optimal decision in the node at stage `stage` and scenario `scenario_index`
of `stochasticprogram`, after a call to `optimize!(stochasticprogram)`.
"""
function optimal_decision(stochasticprogram::StochasticProgram, stage::Integer, scenario_index::Integer)
    if termination_status(stochasticprogram, stage, scenario_index) == MOI.OPTIMIZE_NOT_CALLED
        throw(OptimizeNotCalled())
    end
    return JuMP.value.(all_decision_variables(stochasticprogram, stage), scenario_index)
end
"""
    optimal_decision(stochasticprogram::TwoStageStochasticProgram, scenario_index::Integer)

Return the optimal decision in scenario `scenario_index` of the two-stage `stochasticprogram`,
after a call to `optimize!(stochasticprogram)`.
"""
function optimal_decision(stochasticprogram::TwoStageStochasticProgram, scenario_index::Integer)
    return optimal_decision(stochasticprogram, 2, scenario_index)
end
"""
    optimal_recourse_decision(stochasticprogram::StochasticProgram, scenario_index::Integer)

Return the optimal recourse decision in the final stage of `stochasticprogram` in
the scenario `scenario_index`, after a call to `optimize!(stochasticprogram)`.
"""
function optimal_recourse_decision(stochasticprogram::StochasticProgram{N}, scenario_index::Integer) where N
    return optimal_decision(stochasticprogram, N, scenario_index)
end
"""
    optimize!(stochasticmodel::StochasticModel, sampler::AbstractSampler; crash::AbstractCrash = Crash.None(), kw...)

Approximately optimize the `stochasticmodel` when the underlying scenario distribution is inferred by `sampler`. If an optimizer has not been set yet (see [`set_optimizer`](@ref)), a `NoOptimizer` error is thrown.
"""
function JuMP.optimize!(stochasticmodel::StochasticModel,
                        sampler::AbstractSampler;
                        crash::AbstractCrash = Crash.None(),
                        kw...)
    # Throw NoOptimizer error if no recognized optimizer has been provided
    check_provided_optimizer(stochasticmodel.optimizer)
    # Crash initial decision
    x₀ = crash(stochasticmodel, sampler)
    # Switch on solver type
    return _optimize!(stochasticmodel, sampler, optimizer(stochasticmodel), x₀; kw...)
end
function _optimize!(stochasticmodel::StochasticModel, sampler::AbstractSampler, optimizer::AbstractSampledOptimizer, x₀; kw...)
    # Load optimizer
    load_model!(optimizer, stochasticmodel, sampler, x₀)
    # Run sample-based optimization procedure
    MOI.optimize!(optimizer)
    return nothing
end
function JuMP.termination_status(stochasticmodel::StochasticModel)
    return MOI.get(stochasticmodel, MOI.TerminationStatus())::MOI.TerminationStatusCode
end
function JuMP.objective_sense(stochasticmodel::StochasticModel)
    return MOI.get(stochasticmodel, MOI.ObjectiveSense())::MOI.OptimizationSense
end
"""
    optimal_decision(stochasticmodel::StochasticModel)

Return the optimal first stage decision of `stochasticmodel`, after a call to `optimize!(stochasticmodel)`.
"""
function optimal_decision(stochasticmodel::StochasticModel)
    if termination_status(stochasticmodel) == MOI.OPTIMIZE_NOT_CALLED
        throw(OptimizeNotCalled())
    end
    return optimal_decision(instance(stochasticmodel))
end
"""
    objective_value(stochasticmodel::StochasticModel; result::Int = 1)

Returns the value of the recourse problem after a call to `optimize!(stochasticmodel)`.
"""
function JuMP.objective_value(stochasticmodel::StochasticModel; result::Int = 1)
    # Throw NoOptimizer error if no recognized optimizer has been provided
    check_provided_optimizer(stochasticmodel.optimizer)
    # Switch on termination status
    status = JuMP.termination_status(stochasticmodel)
    if status in AcceptableTermination
        return MOI.get(stochasticmodel, MOI.ObjectiveValue(result))
    else
        if status == MOI.OPTIMIZE_NOT_CALLED
            throw(OptimizeNotCalled())
        elseif status == MOI.INFEASIBLE
            @warn "Stochastic program is infeasible"
            return objective_sense(stochasticmodel) == MOI.MAX_SENSE ? -Inf : Inf
        elseif status == MOI.DUAL_INFEASIBLE
            @warn "Stochastic program is unbounded"
            return objective_sense(stochasticmodel) == MOI.MAX_SENSE ? Inf : -Inf
        else
            @warn("Stochastic program could not be solved, returned status: $status")
            return NaN
        end
    end
end
"""
    stage_parameters(stochasticprogram::StochasticProgram, stage::Integer)

Return the parameters at `stage` in `stochasticprogram`.
"""
function stage_parameters(stochasticprogram::StochasticProgram{N}, stage::Integer) where N
    1 <= stage <= N || error("Stage $stage not in range 1 to $N.")
    return stochasticprogram.stages[stage].parameters
end
"""
    decision(stochasticprogram::StochasticProgram, index::MOI.VariableIndex)

Return the current value of the first-stage decision at `index` of `stochasticprogram`.
"""
function decision(stochasticprogram::StochasticProgram, index::MOI.VariableIndex)
    return decision(structure(stochasticprogram), 1, index)
end
"""
    decision(stochasticprogram::StochasticProgram)

Return the current first-stage decision values of `stochasticprogram`.
"""
function decision(stochasticprogram::StochasticProgram)
    return JuMP.value(all_decision_variables(stochasticprogram))
end
"""
    all_decision_variables(stochasticprogram::StochasticProgram{N}) where N

Returns a stage-wise list of all decisions currently in the `stochasticprogram`. The decisions are
ordered by creation time. Defaults to the first stage.
"""
function all_decision_variables(stochasticprogram::StochasticProgram{N}) where N
    return ntuple(Val{N}()) do s
        return all_decision_variables(stochasticprogram, s)
    end
end
"""
    all_decision_variables(stochasticprogram::StochasticProgram{N}, stage::Integer = 1) where N

Returns a list of all decisions currently in the `stochasticprogram` at `stage`. The decisions are
ordered by creation time.
"""
function all_decision_variables(stochasticprogram::StochasticProgram{N}, stage::Integer) where N
    1 <= stage <= N || error("Stage $stage not in range 1 to $N.")
    decisions::NTuple{stage,Decisions} = stochasticprogram.proxy[stage].ext[:decisions]
    return map(all_decisions(decisions[stage])) do index
        DecisionVariable(stochasticprogram, stage, index)
    end
end
"""
    structure(stochasticprogram::StochasticProgram)

Return the underlying structure of the `stochasticprogram`.
"""
function structure(stochasticprogram::StochasticProgram)
    return stochasticprogram.structure
end
"""
    proxy(stochasticprogram::StochasticProgram, stage::Integer)

Return the proxy model of the `stochasticprogram` at `stage`.
"""
function proxy(stochasticprogram::StochasticProgram{N}, stage::Integer) where N
    1 <= stage <= N || error("Stage $stage not in range 1 to $N.")
    return stochasticprogram.proxy[stage]
end
"""
    Base.getindex(stochasticprogram::StochasticProgram, stage::Integer, name::Symbol)

Returns the decision, or group of decisions, or decision constraint, or group of decision constraints, of the given `name` and `stage` of the `stochasticprogram` which were added to the model. This errors if regular variables or constraints are queried. If so, either annotate the relevant variables with [`@decision`](@ref) or first query the relevant JuMP subproblem and use the regular `[]` syntax.
"""
function Base.getindex(stochasticprogram::StochasticProgram, stage::Integer, name::Symbol)
    obj_dict = object_dictionary(proxy(stochasticprogram, stage))
    if !haskey(obj_dict, name)
        throw(KeyError(name))
    else
        obj = obj_dict[name]
        if obj isa DecisionRef
            return DecisionVariable(stochasticprogram, stage, index(obj))
        elseif obj isa AbstractArray{<:DecisionRef}
            return map(obj) do dvar
                return DecisionVariable(stochasticprogram, stage, index(dvar))
            end
        elseif obj isa ConstraintRef{Model, <:CI{<:DecisionLike}}
            return SPConstraintRef(stochasticprogram, stage, index(obj), obj.shape)
        elseif obj isa AbstractArray{<:ConstraintRef{Model, <:CI{<:DecisionLike}}}
            return map(obj) do cref
                return SPConstraintRef(stochasticprogram, stage, index(cref), cref.shape)
            end
        else
            error("Only decisions and decision constraints can be queried using this syntax. For regular variables and constraints, either annotate the relevant variable with @decision or first query the relevant JuMP subproblem and use the regular `[]` syntax.")
        end
    end
end
"""
    structure_name(stochasticprogram::StochasticProgram)

Return the name of the underlying structure of the `stochasticprogram`.
"""
function structure_name(stochasticprogram::StochasticProgram)
    return structure_name(structure(stochasticprogram))
end
"""
    num_decisions(stochasticprogram::StochasticProgram, stage::Integer = 1)

Return the number of decisions at `stage` in the `stochasticprogram`. Defaults to the first stage.
"""
function num_decisions(stochasticprogram::StochasticProgram{N}, stage::Integer = 1) where N
    1 <= stage <= N || error("Stage $stage not in range 1 to $N.")
    return num_decisions(proxy(stochasticprogram, stage), stage)
end
"""
    num_variables(stochasticprogram::StochasticProgram, stage::Integer = 1)

Return the total number of variables at `stage` in the `stochasticprogram`. Defaults to the first stage.
"""
function JuMP.num_variables(stochasticprogram::StochasticProgram{N}, stage::Integer = 1) where N
    1 <= stage <= N || error("Stage $stage not in range 1 to $N.")
    return num_variables(proxy(stochasticprogram, stage))
end
"""
    first_stage(stochasticprogram::StochasticProgram)

Return the first stage of `stochasticprogram`.
"""
function first_stage(stochasticprogram::StochasticProgram; optimizer = nothing)
    return first_stage(stochasticprogram, structure(stochasticprogram); optimizer = optimizer)
end
function first_stage(stochasticprogram::StochasticProgram, ::AbstractStochasticStructure; optimizer = nothing)
    # Return possibly cached model
    cache = problemcache(stochasticprogram)
    if haskey(cache, :stage_1)
        stage_one = cache[:stage_1]
        optimizer != nothing && set_optimizer(stage_one, optimizer)
        return stage_one
    end
    # Check that the required generators have been defined
    has_generator(stochasticprogram, :stage_1) || error("First-stage problem not defined in stochastic program. Consider @stage 1.")
    # Generate and cache first stage
    generate_stage_one!(stochasticprogram)
    stage_one = cache[:stage_1]
    optimizer != nothing && set_optimizer(stage_one, optimizer)
    return stage_one
end
"""
    scenario(stochasticprogram::StochasticProgram, stage::Integer, scenario_index::Integer)

Return the scenario at `scenario_index` of `stochasticprogram` at stage `stage`.
"""
function scenario(stochasticprogram::StochasticProgram{N}, stage::Integer, scenario_index::Integer) where N
    1 <= stage <= N || error("Stage $stage not in range 1 to $N.")
    stage == 1 && error("The first stage does not have scenarios.")
    return scenario(structure(stochasticprogram), stage, scenario_index)
end
"""
    scenario(stochasticprogram::TwoStageStochasticProgram, scenario_index::Integer)

Return the scenario at `scenario_index` of the two-stage `stochasticprogram`
"""
function scenario(stochasticprogram::TwoStageStochasticProgram, scenario_index::Integer) where N
    return scenario(structure(stochasticprogram), 2, scenario_index)
end
"""
    scenarios(stochasticprogram::StochasticProgram, stage::Integer = 2)

Return an array of all scenarios of the `stochasticprogram` at `stage`. Defaults to the second stage.
"""
function scenarios(stochasticprogram::StochasticProgram{N}, stage::Integer = 2) where N
    1 <= stage <= N || error("Stage $stage not in range 1 to $N.")
    stage == 1 && error("The first stage does not have scenarios.")
    return scenarios(structure(stochasticprogram), stage)
end
"""
    expected(stochasticprogram::StochasticProgram, stage::Integer = 2)

Return the exected scenario of all scenarios of the `stochasticprogram` at `stage`. Defaults to the second stage.
"""
function expected(stochasticprogram::StochasticProgram{N}, stage::Integer = 2) where N
    1 <= stage <= N || error("Stage $stage not in range 1 to $N.")
    stage == 1 && error("The first stage does not have scenarios.")
    return expected(structure(stochasticprogram), stage)
end
"""
    scenario_type(stochasticprogram::StochasticProgram, stage::Integer = 2)

Return the type of the scenario structure associated with `stochasticprogram` at `stage`. Defaults to the second stage.
"""
function scenario_type(stochasticprogram::StochasticProgram{N}, stage::Integer = 2) where N
    1 <= stage <= N || error("Stage $stage not in range 1 to $N.")
    stage == 1 && error("The first stage does not have scenarios.")
    return scenario_type(structure(stochasticprogram), stage)
end
"""
    scenario_types(stochasticprogram::StochasticProgram)

Return a stage-wise list of the scenario types of `stochasticprogram`.
"""
function scenario_types(stochasticprogram::StochasticProgram)
    return ntuple(Val(num_stages(stochasticprogram) - 1)) do s
        scenario_type(stochasticprogram, s + 1)
    end
end
"""
    probability(stochasticprogram::StochasticProgram, stage::Integer, scenario_index::Integer)

Return the probability of the scenario at `scenario_index` in the `stochasticprogram` at `stage` occuring.
"""
function probability(stochasticprogram::StochasticProgram{N}, stage::Integer, scenario_index::Integer) where N
    1 <= stage <= N || error("Stage $stage not in range 1 to $N.")
    s == 1 && error("The first stage does not have scenarios.")
    return probability(structure(stochasticprogram), stage, scenario_index)
end
"""
    probability(stochasticprogram::TwoStageStochasticProgram, scenario_index::Integer)

Return the probability of the scenario at `scenario_index` in the two-stage `stochasticprogram` occuring.
"""
function probability(stochasticprogram::TwoStageStochasticProgram, scenario_index::Integer)
    return probability(structure(stochasticprogram), 2, scenario_index)
end
"""
    stage_probability(stochasticprogram::StochasticProgram, stage::Integer = 2)

Return the probability of any scenario in the `stochasticprogram` at `stage` occuring. A well defined model should return 1. Defaults to the second stage.
"""
function stage_probability(stochasticprogram::StochasticProgram{N}, stage::Integer = 2) where N
    1 <= stage <= N || error("Stage $stage not in range 1 to $N.")
    stage == 1 && error("The first stage does not have scenarios.")
    return stage_probability(structure(stochasticprogram), stage)
end
"""
    has_generator(stochasticprogram::StochasticProgram, key::Symbol)

Return true if a problem generator with `key` exists in `stochasticprogram`.
"""
function has_generator(stochasticprogram::StochasticProgram, key::Symbol)
    return haskey(stochasticprogram.generator, key)
end
"""
    generator(stochasticprogram::StochasticProgram, key::Symbol)

Return the problem generator associated with `key` in `stochasticprogram`.
"""
function generator(stochasticprogram::StochasticProgram, key::Symbol)
    return stochasticprogram.generator[key]
end
"""
    subproblem(stochasticprogram::StochasticProgram, stage::Integer, scenario_index::Integer)

Return the subproblem at `scenario_index` of the `stochasticprogram` at `stage`.
"""
function subproblem(stochasticprogram::StochasticProgram{N}, stage::Integer, scenario_index::Integer) where N
    1 <= stage <= N || error("Stage $stage not in range 1 to $N.")
    stage == 1 || error("The first-stage does not have subproblems.")
    return subproblem(structure(stochasticprogram), stage, scenario_index)
end
"""
    subproblem(stochasticprogram::TwoStasgeStochasticProgram, scenario_index::Integer)

Return the subproblem at `scenario_index` of the two-stage `stochasticprogram`.
"""
function subproblem(stochasticprogram::TwoStageStochasticProgram, scenario_index::Integer)
    return subproblem(structure(stochasticprogram), 2, scenario_index)
end
"""
    subproblems(stochasticprogram::StochasticProgram, stage::Integer = 2)

Return an array of all subproblems of the `stochasticprogram` at `stage`. Defaults to the second stage.
"""
function subproblems(stochasticprogram::StochasticProgram{N}, stage::Integer = 2) where N
    1 <= stage <= N || error("Stage $stage not in range 1 to $N.")
    stage == 1 || error("The first-stage does not have subproblems.")
    return subproblems(structure(stochasticprogram), stage)
end
"""
    num_subproblems(stochasticprogram::StochasticProgram, stage::Integer = 2)

Return the number of subproblems in the `stochasticprogram` at `stage`. Defaults to the second stage.
"""
function num_subproblems(stochasticprogram::StochasticProgram, stage::Integer = 2)
    stage == 1 && return 0
    return num_subproblems(structure(stochasticprogram), stage)
end
"""
    num_scenarios(stochasticprogram::StochasticProgram, stage::Integer = 2)

Return the number of scenarios in the `stochasticprogram` at `stage`. Defaults to the second stage.
"""
function num_scenarios(stochasticprogram::StochasticProgram, stage::Integer = 2)
    stage == 1 && return 0
    return num_scenarios(structure(stochasticprogram), stage)
end
"""
    num_stages(stochasticprogram::StochasticProgram)

Return the number of stages in `stochasticprogram`.
"""
num_stages(::StochasticProgram{N}) where N = N
"""
    distributed(stochasticprogram::StochasticProgram, s::Integer = 2)

Return true if the `stochasticprogram` is memory distributed at stage `s`. Defaults to the second stage.
"""
distributed(stochasticprogram::StochasticProgram, s::Integer = 2) = distributed(structure(stochasticprogram), s)
"""
    deferred(stochasticprogram::StochasticProgram)

Return true if `stochasticprogram` is not fully generated.
"""
deferred(stochasticprogram::StochasticProgram) = num_scenarios(stochasticprogram) == 0 || deferred(structure(stochasticprogram))
"""
    optimizer_constructor(stochasticmodel::StochasticModel)

Return any optimizer constructor supplied to the `stochasticmodel`.
"""
function optimizer_constructor(stochasticmodel::StochasticModel)
    return stochasticmodel.optimizer.optimizer_constructor
end
"""
    optimizer_constructor(stochasticprogram::StochasticProgram)

Return any optimizer constructor supplied to the `stochasticprogram`.
"""
function optimizer_constructor(stochasticprogram::StochasticProgram)
    return stochasticprogram.optimizer.optimizer_constructor
end
"""
    optimizer_name(stochasticmodel::StochasticModel)

Return the currently provided optimizer type of `stochasticmodel`.
"""
function optimizer_name(stochasticmodel::StochasticModel)
    return optimizer_name(optimizer(stochasticmodel))
end
"""
    optimizer_name(stochasticprogram::StochasticProgram)

Return the currently provided optimizer type of `stochasticprogram`.
"""
function optimizer_name(stochasticprogram::StochasticProgram)
    deferred(stochasticprogram) && return "Uninitialized"
    return optimizer_name(optimizer(stochasticprogram))
end
function optimizer_name(optimizer::MOI.AbstractOptimizer)
    if JuMP.moi_mode(optimizer) != DIRECT && MOIU.state(optimizer) == MOIU.NO_OPTIMIZER
        return "No optimizer attached."
    end
    return JuMP._try_get_solver_name(optimizer)
end
"""
    master_optimizer(stochasticprogram::StochasticProgram)

Return a MOI optimizer using the currently provided optimizer of `stochasticprogram`.
"""
function master_optimizer(stochasticprogram::StochasticProgram)
    return master_optimizer(stochasticprogram.optimizer)
end
"""
    subproblem_optimizer(stochasticprogram::StochasticProgram)

Return a MOI optimizer for solving subproblems using the currently provided optimizer of `stochasticprogram`.
"""
function subproblem_optimizer(stochasticprogram::StochasticProgram)
    return subproblem_optimizer(stochasticprogram.optimizer)
end
"""
    optimizer(stochasticprogram::StochasticProgram)

Return the optimizer attached to `stochasticprogram`.
"""
function optimizer(stochasticprogram::StochasticProgram)
    return stochasticprogram.optimizer.optimizer
end
"""
    optimizer(stochasticmodel::StochasticModel)

Return the optimizer attached to `stochasticmodel`.
"""
function optimizer(stochasticmodel::StochasticModel)
    return stochasticmodel.optimizer.optimizer
end
"""
    get_optimizer_attribute(stochasticprogram::StochasticProgram, name::String)

Return the value associated with the solver-specific attribute named `name` in `stochasticprogram`.

See also: [`set_optimizer_attribute`](@ref), [`set_optimizer_attributes`](@ref).
"""
function get_optimizer_attribute(stochasticprogram::StochasticProgram, name::String)
    return get_optimizer_attribute(stochasticprogram, MOI.RawParameter(name))
end
"""
    get_optimizer_attribute(stochasticprogram::StochasticProgram, attr::MOI.AbstractOptimizerAttribute)

Return the value of the solver-specific attribute `attr` in `stochasticprogram`.

See also: [`set_optimizer_attribute`](@ref), [`set_optimizer_attributes`](@ref).
"""
function get_optimizer_attribute(stochasticprogram::StochasticProgram, attr::MOI.AbstractOptimizerAttribute)
    return MOI.get(stochasticprogram, attr)
end
# ========================== #

# Setters
# ========================== #
"""
    update_known_decisions!(stochasticprogram::Stochasticprogram)

Update all known decision values in the first-stage of `stochasticprogram`.
"""
function update_known_decisions!(stochasticprogram::StochasticProgram)
    update_known_decisions!(structure(stochasticprogram))
end
"""
    update_decisions!(stochasticprogram::Stochasticprogram, stage::Integer, scenario_index::Integer)

Update all known decision values in the node at stage `stage` and scenario `scenario_index` of `stochasticprogram`.
"""
function update_decisions!(stochasticprogram::StochasticProgram, stage::Integer, scenario_index::Integer)
    update_known_decisions!(structure(stochasticprogram), stage, scenario_index)
end
"""
    set_optimizer(stochasticmodel::StochasticModel, optimizer)

Set the optimizer of the `stochasticmodel`.
"""
function JuMP.set_optimizer(stochasticmodel::StochasticModel, optimizer)
    set_optimizer!(stochasticmodel.optimizer, optimizer)
    if stochasticmodel.optimizer.optimizer isa AbstractSampledOptimizer
        return nothing
    end
    # Default to SAA
    set_optimizer(stochasticmodel, SAA.Optimizer)
    set_optimizer_attribute(stochasticmodel, InstanceOptimizer(), optimizer)
    return nothing
end
"""
    set_optimizer(stochasticprogram::StochasticProgram, optimizer)

Set the optimizer of the `stochasticprogram`.
"""
function JuMP.set_optimizer(stochasticprogram::StochasticProgram, optimizer)
    if !supports_structure(optimizer(), structure(stochasticprogram))
        @warn "The provided optimizer does not support the underlying structure of the stochastic program. Consider reinstantiating (see `instantiate`) the stochastic program with the optimizer instead."
        return nothing
    end
    set_optimizer!(stochasticprogram.optimizer, optimizer)
    master_opt = master_optimizer(stochasticprogram)
    if master_opt != nothing
        set_master_optimizer!(structure(stochasticprogram), master_opt)
    end
    sub_opt = subproblem_optimizer(stochasticprogram)
    if sub_opt != nothing
        set_subproblem_optimizer!(structure(stochasticprogram), sub_opt)
    end
    return nothing
end
"""
    set_optimizer_attribute(stochasticprogram::StochasticProgram, name::Union{Symbol, String}, value)

Sets solver-specific attribute identified by `name` to `value` in the `stochasticprogram`.

See also: [`get_optimizer_attribute`](@ref)
"""
function JuMP.set_optimizer_attribute(stochasticprogram::StochasticProgram, name::Union{Symbol, String}, value)
    return set_optimizer_attribute(stochasticprogram, MOI.RawParameter(String(name)), value)
end
"""
    set_optimizer_attribute(stochasticprogram::StochasticProgram, attr::MOI.AbstractOptimizerAttribute, value)

Set the solver-specific attribute `attr` in `stochasticprogram` to `value`.

See also: [`get_optimizer_attribute`](@ref)
"""
function JuMP.set_optimizer_attribute(stochasticprogram::StochasticProgram, attr::MOI.AbstractOptimizerAttribute, value)
    return MOI.set(stochasticprogram, attr, value)
end
"""
    set_optimizer_attributes(stochasticprogram::StochasticProgram, pairs::Pair...)

Given a list of `attribute => value` pairs or a collection of keyword arguments, calls
`set_optimizer_attribute(stochasticprogram, attribute, value)` for each pair.

See also: [`get_optimizer_attribute`](@ref)
"""
function JuMP.set_optimizer_attributes(stochasticprogram::StochasticProgram, pairs::Pair...)
    for (name, value) in pairs
        set_optimizer_attribute(stochasticprogram, name, value)
    end
end
function JuMP.set_optimizer_attributes(stochasticprogram::StochasticProgram; kw...)
    for (name, value) in kw
        set_optimizer_attribute(stochasticprogram, name, value)
    end
end
function JuMP.set_silent(stochasticprogram::StochasticProgram)
    return MOI.set(stochasticprogram, MOI.Silent(), true)
end
function JuMP.unset_silent(stochasticprogram::StochasticProgram)
    return MOI.set(stochasticprogram, MOI.Silent(), false)
end
function JuMP.set_time_limit_sec(stochasticprogram::StochasticProgram, limit)
    return MOI.set(stochasticprogram, MOI.TimeLimitSec(), limit)
end
function JuMP.unset_time_limit_sec(stochasticprogram::StochasticProgram)
    return MOI.set(stochasticprogram, MOI.TimeLimitSec(), nothing)
end
function JuMP.time_limit_sec(stochasticprogram::StochasticProgram)
    return MOI.get(stochasticprogram, MOI.TimeLimitSec())
end
function JuMP.simplex_iterations(stochasticprogram::StochasticProgram)
    return MOI.get(stochasticprogram, MOI.SimplexIterations())
end
function JuMP.barrier_iterations(stochasticprogram::StochasticProgram)
    return MOI.get(stochasticprogram, MOI.BarrierIterations())
end
function JuMP.node_count(stochasticprogram::StochasticProgram)
    return MOI.get(stochasticprogram, MOI.NodeCount())
end
function JuMP.set_optimizer_attribute(stochasticmodel::StochasticModel, name::Union{Symbol, String}, value)
    return set_optimizer_attribute(stochasticmodel, MOI.RawParameter(String(name)), value)
end
function JuMP.set_optimizer_attribute(stochasticmodel::StochasticModel, attr::MOI.AbstractOptimizerAttribute, value)
    return MOI.set(stochasticmodel, attr, value)
end
function JuMP.set_optimizer_attributes(stochasticmodel::StochasticModel, pairs::Pair...)
    for (name, value) in pairs
        set_optimizer_attribute(stochasticmodel, name, value)
    end
end
function JuMP.set_optimizer_attributes(stochasticmodel::StochasticModel; kw...)
    for (name, value) in kw
        set_optimizer_attribute(stochasticmodel, name, value)
    end
end
function JuMP.set_silent(stochasticmodel::StochasticModel)
    return MOI.set(stochasticmodel, MOI.Silent(), true)
end
function JuMP.unset_silent(stochasticmodel::StochasticModel)
    return MOI.set(stochasticmodel, MOI.Silent(), false)
end
"""
    add_scenario!(stochasticprogram::StochasticProgram, stage::Integer, scenario::AbstractScenario)

Store the second stage `scenario` in the `stochasticprogram` at `stage`.

If the `stochasticprogram` is distributed, the scenario will be defined on the node that currently has the fewest scenarios.
"""
function add_scenario!(stochasticprogram::StochasticProgram, stage::Integer, scenario::AbstractScenario)
    add_scenario!(structure(stochasticprogram), stage, scenario)
    invalidate_cache!(stochasticprogram)
    return stochasticprogram
end
"""
    add_scenario!(stochasticprogram::TwoStageStochasticProgram, scenario::AbstractScenario)

Store the second stage `scenario` in the two-stage `stochasticprogram`.
"""
function add_scenario!(stochasticprogram::TwoStageStochasticProgram, scenario::AbstractScenario)
    add_scenario!(structure(stochasticprogram), 2,scenario)
    invalidate_cache!(stochasticprogram)
    return stochasticprogram
end
"""
    add_worker_scenario!(stochasticprogram::StochasticProgram, stage::Integer, scenario::AbstractScenario, w::Integer)

Store the second stage `scenario` in worker node `w` of the `stochasticprogram` at `stage`.
"""
function add_worker_scenario!(stochasticprogram::StochasticProgram, stage::Integer, scenario::AbstractScenario, w::Integer)
    add_scenario!(structure(stochasticprogram), stage, scenario, w)
    invalidate_cache!(stochasticprogram)
    return stochasticprogram
end
"""
    add_worker_scenario!(stochasticprogram::StochasticProgram, stage::Integer, scenario::AbstractScenario, w::Integer)

Store the second stage `scenario` in worker node `w` of the two-stage `stochasticprogram`.
"""
function add_worker_scenario!(stochasticprogram::TwoStageStochasticProgram, scenario::AbstractScenario, w::Integer)
    add_scenario!(structure(stochasticprogram), 2, scenario, w)
    invalidate_cache!(stochasticprogram)
    return stochasticprogram
end
"""
    add_scenario!(scenariogenerator::Function, stochasticprogram::StochasticProgram, stage::Integer)

Store the scenario returned by `scenariogenerator` in the `stage` of the `stochasticprogram`. If the `stochasticprogram` is distributed, the scenario will be defined on the node that currently has the fewest scenarios.
"""
function add_scenario!(scenariogenerator::Function, stochasticprogram::StochasticProgram, stage::Integer)
    add_scenario!(scenariogenerator, structure(stochasticprogram), stage)
    invalidate_cache!(stochasticprogram)
    return stochasticprogram
end
"""
    add_scenario!(scenariogenerator::Function, stochasticprogram::StochasticProgram)

Store the scenario returned by `scenariogenerator` in the two-stage `stochasticprogram`. If the `stochasticprogram` is distributed, the scenario will be defined on the node that currently has the fewest scenarios.
"""
function add_scenario!(scenariogenerator::Function, stochasticprogram::TwoStageStochasticProgram)
    add_scenario!(scenariogenerator, structure(stochasticprogram), 2)
    invalidate_cache!(stochasticprogram)
    return stochasticprogram
end
"""
    add_worker_scenario!(scenariogenerator::Function, stochasticprogram::StochasticProgram, stage::Integer, w::Integer)

Store the scenario returned by `scenariogenerator` in worker node `w` of the `stochasticprogram` at `stage`.
"""
function add_worker_scenario!(scenariogenerator::Function, stochasticprogram::StochasticProgram, stage::Integer, w::Integer)
    add_scenario!(scenariogenerator, structure(stochasticprogram), stage, w)
    invalidate_cache!(stochasticprogram)
    return stochasticprogram
end
"""
    add_worker_scenario!(scenariogenerator::Function, stochasticprogram::TwoStageStochasticProgram, w::Integer)

Store the scenario returned by `scenariogenerator` in worker node `w` of the two-stage `stochasticprogram`.
"""
function add_worker_scenario!(scenariogenerator::Function, stochasticprogram::TwoStageStochasticProgram, w::Integer)
    add_scenario!(scenariogenerator, structure(stochasticprogram), 2, w)
    invalidate_cache!(stochasticprogram)
    return stochasticprogram
end
"""
    add_scenarios!(stochasticprogram::StochasticProgram, stage::Integer, scenarios::Vector{<:AbstractScenario})

Store the collection of `scenarios` in the `stochasticprogram` at `stage`. If the `stochasticprogram` is distributed, scenarios will be distributed evenly across workers.
"""
function add_scenarios!(stochasticprogram::StochasticProgram, stage::Integer, scenarios::Vector{<:AbstractScenario})
    add_scenarios!(structure(stochasticprogram), stage, scenarios)
    invalidate_cache!(stochasticprogram)
    return stochasticprogram
end
"""
    add_scenarios!(stochasticprogram::TwoStageStochasticProgram, scenarios::Vector{<:AbstractScenario})

Store the collection of `scenarios` in the two-stage `stochasticprogram`. If the `stochasticprogram` is distributed, scenarios will be distributed evenly across workers.
"""
function add_scenarios!(stochasticprogram::TwoStageStochasticProgram, scenarios::Vector{<:AbstractScenario})
    add_scenarios!(structure(stochasticprogram), 2, scenarios)
    invalidate_cache!(stochasticprogram)
    return stochasticprogram
end
"""
    add_worker_scenarios!(stochasticprogram::StochasticProgram, stage::Integer, scenarios::Vector{<:AbstractScenario}, w::Integer)

Store the collection of `scenarios` in in worker node `w` of the `stochasticprogram` at `stage`.
"""
function add_worker_scenarios!(stochasticprogram::StochasticProgram, stage::Integer, scenarios::Vector{<:AbstractScenario}, w::Integer)
    add_scenarios!(structure(stochasticprogram), stage, scenarios, w)
    invalidate_cache!(stochasticprogram)
    return stochasticprogram
end
"""
    add_worker_scenarios!(stochasticprogram::StochasticProgram, stage::Integer, scenarios::Vector{<:AbstractScenario}, w::Integer)

Store the collection of `scenarios` in in worker node `w` of the two-stage `stochasticprogram`.
"""
function add_worker_scenarios!(stochasticprogram::TwoStageStochasticProgram, scenarios::Vector{<:AbstractScenario}, w::Integer)
    add_scenarios!(structure(stochasticprogram), 2, scenarios, w)
    invalidate_cache!(stochasticprogram)
    return stochasticprogram
end
"""
    add_scenarios!(scenariogenerator::Function, stochasticprogram::StochasticProgram, stage::Integer, n::Integer)

Generate `n` scenarios using `scenariogenerator` and store in the `stochasticprogram` at `stage`. If the `stochasticprogram` is distributed, scenarios will be distributed evenly across workers.
"""
function add_scenarios!(scenariogenerator::Function, stochasticprogram::StochasticProgram, stage::Integer, n::Integer)
    add_scenarios!(scenariogenerator, structure(stochasticprogram), stage, n)
    invalidate_cache!(stochasticprogram)
    return stochasticprogram
end
"""
    add_scenarios!(scenariogenerator::Function, stochasticprogram::TwoStageStochasticProgram, n::Integer)

Generate `n` scenarios using `scenariogenerator` and store in the two-stage `stochasticprogram`. If the `stochasticprogram` is distributed, scenarios will be distributed evenly across workers.
"""
function add_scenarios!(scenariogenerator::Function, stochasticprogram::TwoStageStochasticProgram, n::Integer)
    add_scenarios!(scenariogenerator, structure(stochasticprogram), 2, n)
    invalidate_cache!(stochasticprogram)
    return stochasticprogram
end
"""
    add_worker_scenarios!(scenariogenerator::Function, stochasticprogram::StochasticProgram, stage::Integer, n::Integer, w::Integer)

Generate `n` scenarios using `scenariogenerator` and store them in worker node `w` of the `stochasticprogram` at `stage`.
"""
function add_worker_scenarios!(scenariogenerator::Function, stochasticprogram::StochasticProgram, stage::Integer, n::Integer, w::Integer)
    add_scenarios!(scenariogenerator, structure(stochasticprogram), stage, n, w)
    invalidate_cache!(stochasticprogram)
    return stochasticprogram
end
"""
    add_worker_scenarios!(scenariogenerator::Function, stochasticprogram::TwoStageStochasticProgram, n::Integer, w::Integer)

Generate `n` scenarios using `scenariogenerator` and store them in worker node `w` of the two-stage `stochasticprogram`.
"""
function add_worker_scenarios!(scenariogenerator::Function, stochasticprogram::TwoStageStochasticProgram, n::Integer, w::Integer)
    add_scenarios!(scenariogenerator, structure(stochasticprogram), 2, n, w)
    invalidate_cache!(stochasticprogram)
    return stochasticprogram
end
"""
    sample!(stochasticprogram::StochasticProgram, stage::Integer, sampler::AbstractSampler, n::Integer)

Sample `n` scenarios using `sampler` and add to the `stochasticprogram` at `stage`. If the `stochasticprogram` is distributed, scenarios will be distributed evenly across workers.
"""
function sample!(stochasticprogram::StochasticProgram, stage::Integer, sampler::AbstractSampler, n::Integer)
    sample!(structure(stochasticprogram), stage, sampler, n)
    return stochasticprogram
end
"""
    sample!(stochasticprogram::TwoStageStochasticProgram, sampler::AbstractSampler, n::Integer)

Sample `n` scenarios using `sampler` and add to the two-stage `stochasticprogram`. If the `stochasticprogram` is distributed, scenarios will be distributed evenly across workers.
"""
function sample!(stochasticprogram::TwoStageStochasticProgram, sampler::AbstractSampler, n::Integer)
    sample!(structure(stochasticprogram), 2, sampler, n)
    return stochasticprogram
end
# ========================== #


module PointArrays

export PointArray

const PointArray = Array{UInt32, 2}

const ZERO_UINT32 = zero(UInt32)
const ONE_UINT32  = one(UInt32)
const DEFAULT_OFFSET = (ZERO_UINT32, ZERO_UINT32, ZERO_UINT32)

const MAX_BOUNDARY_DISTANCE = 100000

"""
parameter:
    bin_im: binary array. object voxels are false, non-object voxels are true!
"""
function from_binary_image(bin_im::Union{Array{Bool,3}, BitArray{3}})
    N = length(bin_im) - count(bin_im)
    ret = zeros(UInt32, (N,3))
    i = 0
    for z in ONE_UINT32:UInt32(size(bin_im, 3))
        for y in ONE_UINT32:UInt32(size(bin_im, 2))
            for x in ONE_UINT32:UInt32(size(bin_im, 1))
                if !bin_im[x,y,z]
                    i += 1
                    ret[i,:] = [x,y,z]
                end
            end
        end
    end
    ret
end

"""
find points inside an object from a segmentation array.
"""
function from_seg(seg::Array{T,3}; obj_id::T=convert(T,1),
                     offset::NTuple{3,UInt32}=DEFAULT_OFFSET) where T
    # compute the number of object voxels
    nov = 0
    for i in eachindex(seg)
        if seg[i] == obj_id
            nov += 1
        end
    end
    # initialize the points array with offset
    points = Array{UInt32}(undef, nov, 3)
    points[:,1] .= offset[1]
    points[:,2] .= offset[2]
    points[:,3] .= offset[3]

    # position of current row
    pos = 0
    for z in ONE_UINT32 : convert(UInt32, size(seg, 3))
        for y in ONE_UINT32 : convert(UInt32, size(seg, 2))
            for x in ONE_UINT32 : convert(UInt32, size(seg, 1))
                if seg[x,y,z] == obj_id
                    pos += 1
                    points[pos,:] .+= [x,y,z]
                end
            end
        end
    end
    return points
end

"""
find out the boundary voxels and represent them as indexes in the point array
"""
function get_boundary_point_indexes(self::Array{T,2}, seg::Array{TSeg,3},
                                    obj_id::TSeg = TSeg(1)) where {T, TSeg}
    # compute the indexes of boundary voxels
    ret = Vector{T}()
    for i in 1:size(self,1)
        point = self[i,:]
        x,y,z = (point...,)
        if  z==1 || z==size(seg, 3) ||
            y==1 || y==size(seg, 2) ||
            x==1 || x==size(seg, 1)
            push!(ret, i)
        else
            if  seg[x-1,y,z]!=obj_id || seg[x+1,y,z]!=obj_id ||
                seg[x,y-1,z]!=obj_id || seg[x,y+1,z]!=obj_id ||
                seg[x,y,z-1]!=obj_id || seg[x,y,z+1]!=obj_id
                push!(ret, i)
            end
        end
    end
    return ret
end

"""
add offset to points
"""
function add_offset!(self::Array{T,2}, offset::NTuple{3,T}) where T
    self[:, 1] .+= offset[1]
    self[:, 2] .+= offset[2]
    self[:, 3] .+= offset[3]
end

"""
    merge(self::Array{T,3}, other::Array{T,3})
"""
function merge(self::Array{T,2}, other::Array{T,2}) where T
    vcat(self, other)
end

end # module


using WebAssembly, Test
using WebAssembly.Instructions
using WebAssembly: parsewast, interpretwasm, interpret_module, WType, Func, Module, FuncType, Func, Table, Mem, Global, Elem, Data, Import, Export, i32, i64, f32, f64

@testset "WebAssembly" begin

  b = Block([Nop(), Nop()]) |> WebAssembly.nops
  @test isempty(b.body)

  @testset "Interpreter" begin
    include("interpret.jl")
  end

  @testset "IR" begin
    include("ir.jl")
  end

  @testset "Import-export" begin
    m = WebAssembly.Module(FuncType[], Func[], Table[], Mem[], Global[], Elem[], Data[], Ref(0),
               [Import(:env, :mathfun, :func, [i32, f64], f64)],
               [Export(:fun, :fun, :func)])
  end

end


using Gadfly

function plot_matrix( tiles_matrix, path )
     light_blue = color("blue")
     dark_blue  = color("red")
     sc = Scale.ContinuousColorScale(Scale.lab_gradient( dark_blue,
                                                         light_blue ) )
     num_rows = size(tiles_matrix,1)
     Gadfly.draw(PNG( path, 10inch, 10inch), spy2( tiles_matrix,  sc, Scale.y_continuous(maxvalue=num_rows) ) )
     # or call the R code, because the time this takes can be excessive??
end

function plot_matrix2(tiles_matrix,path; colour_scale =Scale.discrete_color_manual("black","red","green") )
    #border1 = color("blue")
    #border2 = color("green")
    num_rows = size(tiles_matrix,1)
    #sc = Scale.ContinuousColorScale(Scale.lab_gradient( border1,
    #                                                    border2 ) )
    #dc =Scale.discrete_color_manual("black","red","green")
    is, js, values = findnz(tiles_matrix)
    df = DataFrame(i=is, j=js, value=values)
    T = Theme(panel_fill=color("black"))
    p = Gadfly.plot(df, x="j", y="i", color="value",
        T,
        Coord.cartesian(yflip=true),
        #Scale.continuous_color,
        colour_scale,
        Scale.x_continuous,
        Scale.y_continuous(maxvalue=num_rows),
        Geom.rectbin,
        Stat.identity)

    Gadfly.draw(PNG( path, 10inch, 10inch), p )
end

# So you want to get the number of genes in each cluster and make a bar plot of those
# But the bar plot should be coloured by the average methylation level
# So we have to get a value for the average number of tiles methylated

# A. Get the number in each cluster
# B. Get the average methylation for each cluster

function plot_cluster_bar_chart(tiles_matrix, cluster_assignments,path,id=""; gadfly_theme = Theme(major_label_color = color("black"), minor_label_color=color("black")) )

    # GET NUMBERS OF REGIONS IN EACH CLUSTER
    # Make dictionary with keys as cluster_id
    # and value how many regions in each cluster
    cluster_counts = Dict{Int64,Int64}()
    for c in cluster_assignments
        if haskey( cluster_counts,c )
            cluster_counts[c] +=1
        else
            cluster_counts[c]=1
        end
     end

    # GET SUM OF TILES VALUES FOR EACH REGION
    i=1
    cluster_value_sum =Dict{Int64,Float64}()
    tiles_matrix_transpose = tiles_matrix
    for c in cluster_assignments
         sum_of_tiles = sum(tiles_matrix_transpose[i,:])
         if haskey(cluster_value_sum,c)
             cluster_value_sum[c] +=sum_of_tiles
         else
             cluster_value_sum[c] = sum_of_tiles
         end
         i +=1
     end

     # CALCULATE AVERAGE VALUE TILES FROM EACH REGION
     ave_cluster_value = Float64[]
     cluster_num = Int64[]
     for k in sort(collect(keys(cluster_value_sum)))
        println(cluster_value_sum[k])
        println(cluster_counts[k])
        ave_value_tiles = cluster_value_sum[k] / cluster_counts[k]
        push!(ave_cluster_value, ave_value_tiles)
        push!(cluster_num,cluster_counts[k])
     end

     # PLOT
     #sort_by_clusterN = sortperm(cluster_num,rev=true)
     df = DataFrame( average_cluster_value = ave_cluster_value, cluster_num = cluster_num)
     sort!(df, cols = [:cluster_num],rev=true)
     nrows = nrow(df)
     if nrows > 15
         nrows = 15
     end
     df = df[1:nrows,:]
     gp1 = plot(df, x=1:nrows, y="cluster_num", color="average_cluster_value",
         gadfly_theme,Geom.bar,
         Guide.title("Genes in Cluster $id"),
         Guide.xlabel("Cluster Number"),
         Guide.ylabel("Num of Regions"),
         Guide.colorkey("Avg Tiles Value"),
         Scale.y_continuous(format=:plain))
     Gadfly.draw(PNG(path,20cm, 15cm), gp1)
     return gp1
end


module NestedEnvironments

using DifferentialEquations
using Transducers

export AbstractEnv  # types.jl
# internalAPIs.jl
export raw, readable, @reg_env, @raw, @readable  # APIs.jl
export BaseEnv, InputAffineQuadraticCostEnv  # zoo.jl

include("types.jl")
include("internalAPIs.jl")
include("APIs.jl")
include("zoo.jl")

end