import Base: size, length, iterate, getindex, setindex!, lastindex
using Base: @propagate_inbounds

import Adapt
using OffsetArrays

import Oceananigans.Architectures: architecture
import Oceananigans.Utils: datatuple

using Oceananigans: AbstractGrid
using Oceananigans.Architectures: @hascuda, CPU, GPU

@hascuda using CuArrays

"""
    AbstractField{A, G}

Abstract supertype for fields stored on an architecture `A` and defined on a grid `G`.
"""
abstract type AbstractField{A, G} end

"""
    AbstractLocatedField{X, Y, Z, A, G}

Abstract supertype for fields located at `(X, Y, Z)`, stored on an architecture `A`,
and defined on a grid `G`.
"""
abstract type AbstractLocatedField{X, Y, Z, A, G} <: AbstractField{A, G} end

"""
    Cell

A type describing the location at the center of a grid cell.
"""
struct Cell end

"""
	Face

A type describing the location at the face of a grid cell.
"""
struct Face end

"""
    Field{X, Y, Z, A, G} <: AbstractLocatedField{X, Y, Z, A, G}

A field defined at the location (`X`, `Y`, `Z`) which can be either `Cell` or `Face`.
"""
struct Field{X, Y, Z, A, G} <: AbstractLocatedField{X, Y, Z, A, G}
    data :: A
    grid :: G
    function Field{X, Y, Z}(data, grid) where {X, Y, Z}
        return new{X, Y, Z, typeof(data), typeof(grid)}(data, grid)
    end
end

Adapt.adapt_structure(to, field::Field{X, Y, Z}) where {X, Y, Z} =
    Field{X, Y, Z}(adapt(to, data), field.grid)

"""
	Field(L::Tuple, data::AbstractArray, grid)

Construct a `Field` on `grid` using the array `data` with location defined by the tuple `L`
of length 3 whose elements are `Cell` or `Face`.
"""
Field(L::Tuple, data::AbstractArray, grid) = Field{L[1], L[2], L[3]}(data, grid)

"""
    Field(L::Tuple, arch, grid)

Construct a `Field` on architecture `arch` and `grid` at location `L`,
where `L` is a tuple of `Cell` or `Face` types.
"""
Field(L::Tuple, arch, grid) = Field{L[1], L[2], L[3]}(zeros(arch, grid), grid)

"""
    Field(X, Y, Z, arch, grid)

Construct a `Field` on architecture `arch` and `grid` at location `X`, `Y`, `Z`,
where each of `X, Y, Z` is `Cell` or `Face`.
"""
Field(X, Y, Z, arch, grid) =  Field((X, Y, Z), arch, grid)

"""
    CellField([T=eltype(grid)], arch, grid)

Return a `Field{Cell, Cell, Cell}` on architecture `arch` and `grid`.
Used for tracers and pressure fields.
"""
CellField(T, arch, grid) = Field{Cell, Cell, Cell}(zeros(T, arch, grid), grid)

"""
    FaceFieldX([T=eltype(grid)], arch, grid)

Return a `Field{Face, Cell, Cell}` on architecture `arch` and `grid`.
Used for the x-velocity field.
"""
FaceFieldX(T, arch, grid) = Field{Face, Cell, Cell}(zeros(T, arch, grid), grid)

"""
    FaceFieldY([T=eltype(grid)], arch, grid)

Return a `Field{Cell, Face, Cell}` on architecture `arch` and `grid`.
Used for the y-velocity field.
"""
FaceFieldY(T, arch, grid) = Field{Cell, Face, Cell}(zeros(T, arch, grid), grid)

"""
    FaceFieldZ([T=eltype(grid)], arch, grid)

Return a `Field{Cell, Cell, Face}` on architecture `arch` and `grid`.
Used for the z-velocity field.
"""
FaceFieldZ(T, arch, grid) = Field{Cell, Cell, Face}(zeros(T, arch, grid), grid)

 CellField(arch, grid) = Field((Cell, Cell, Cell), arch, grid)
FaceFieldX(arch, grid) = Field((Face, Cell, Cell), arch, grid)
FaceFieldY(arch, grid) = Field((Cell, Face, Cell), arch, grid)
FaceFieldZ(arch, grid) = Field((Cell, Cell, Face), arch, grid)

location(a) = nothing
location(::AbstractLocatedField{X, Y, Z}) where {X, Y, Z} = (X, Y, Z)

architecture(f::Field) = architecture(f.data)
architecture(o::OffsetArray) = architecture(o.parent)

@inline size(f::AbstractField) = size(f.grid)
@inline length(f::Field) = length(f.data)

@propagate_inbounds getindex(f::Field, inds...) = @inbounds getindex(f.data, inds...)
@propagate_inbounds setindex!(f::Field, v, inds...) = @inbounds setindex!(f.data, v, inds...)

@inline lastindex(f::Field) = lastindex(f.data)
@inline lastindex(f::Field, dim) = lastindex(f.data, dim)

"Returns `f.data` for `f::Field` or `f` for `f::AbstractArray."
@inline data(a) = a
@inline data(f::Field) = f.data

# Endpoint for recursive `datatuple` function:
@inline datatuple(obj::AbstractField) = data(obj)

"Returns `f.data.parent` for `f::Field`."
@inline Base.parent(f::Field) = f.data.parent

"Returns a view over the interior points of the `field.data`."
@inline interior(f::Field) = view(f.data, 1:f.grid.Nx, 1:f.grid.Ny, 1:f.grid.Nz)

"Returns a reference to the interior points of `field.data.parent.`"
@inline interiorparent(f::Field) = @inbounds f.data.parent[1+f.grid.Hx:f.grid.Nx+f.grid.Hx,
                                                           1+f.grid.Hy:f.grid.Ny+f.grid.Hy,
                                                           1+f.grid.Hz:f.grid.Nz+f.grid.Hz]

iterate(f::Field, state=1) = iterate(f.data, state)

@inline xnode(::Type{Cell}, i, grid) = @inbounds grid.xC[i]
@inline xnode(::Type{Face}, i, grid) = @inbounds grid.xF[i]

@inline ynode(::Type{Cell}, j, grid) = @inbounds grid.yC[j]
@inline ynode(::Type{Face}, j, grid) = @inbounds grid.yF[j]

@inline znode(::Type{Cell}, k, grid) = @inbounds grid.zC[k]
@inline znode(::Type{Face}, k, grid) = @inbounds grid.zF[k]

@inline xnode(i, ϕ::Field{X, Y, Z}) where {X, Y, Z} = xnode(X, i, ϕ.grid)
@inline ynode(j, ϕ::Field{X, Y, Z}) where {X, Y, Z} = ynode(Y, j, ϕ.grid)
@inline znode(k, ϕ::Field{X, Y, Z}) where {X, Y, Z} = znode(Z, k, ϕ.grid)

xnodes(ϕ::AbstractField) = reshape(ϕ.grid.xC, ϕ.grid.Nx, 1, 1)
ynodes(ϕ::AbstractField) = reshape(ϕ.grid.yC, 1, ϕ.grid.Ny, 1)
znodes(ϕ::AbstractField) = reshape(ϕ.grid.zC, 1, 1, ϕ.grid.Nz)

xnodes(ϕ::Field{Face})                    = reshape(ϕ.grid.xF[1:end-1], ϕ.grid.Nx, 1, 1)
ynodes(ϕ::Field{X, Face}) where X         = reshape(ϕ.grid.yF[1:end-1], 1, ϕ.grid.Ny, 1)
znodes(ϕ::Field{X, Y, Face}) where {X, Y} = reshape(ϕ.grid.zF[1:end-1], 1, 1, ϕ.grid.Nz)

nodes(ϕ) = (xnodes(ϕ), ynodes(ϕ), znodes(ϕ))

# Niceties
const AbstractCPUField =
    AbstractField{A, G} where {A<:OffsetArray{T, D, <:Array} where {T, D}, G}

@hascuda const AbstractGPUField =
    AbstractField{A, G} where {A<:OffsetArray{T, D, <:CuArray} where {T, D}, G}

#####
##### Creating fields by dispatching on architecture
#####

function OffsetArray(underlying_data, grid)
    # Starting and ending indices for the offset array.
    i1, i2 = 1 - grid.Hx, grid.Nx + grid.Hx
    j1, j2 = 1 - grid.Hy, grid.Ny + grid.Hy
    k1, k2 = 1 - grid.Hz, grid.Nz + grid.Hz

    return OffsetArray(underlying_data, i1:i2, j1:j2, k1:k2)
end

function Base.zeros(T, ::CPU, grid)
    underlying_data = zeros(T, grid.Tx, grid.Ty, grid.Tz)
    return OffsetArray(underlying_data, grid)
end

function Base.zeros(T, ::GPU, grid)
    underlying_data = CuArray{T}(undef, grid.Tx, grid.Ty, grid.Tz)
    underlying_data .= 0  # Gotta do this otherwise you might end up with a few NaN values!
    return OffsetArray(underlying_data, grid)
end

Base.zeros(T, ::CPU, grid, Nx, Ny, Nz) = zeros(T, Nx, Ny, Nz)
Base.zeros(T, ::GPU, grid, Nx, Ny, Nz) = zeros(T, Nx, Ny, Nz) |> CuArray

# Default to type of Grid
Base.zeros(arch, grid::AbstractGrid{T}) where T = zeros(T, arch, grid)
Base.zeros(arch, grid::AbstractGrid{T}, Nx, Ny, Nz) where T = zeros(T, arch, grid, Nx, Ny, Nz)


import Base: length, getindex
export SampleSet, MetricNN, QuasiMetricNN
export NearNeighborCache, MutableNNC, ImmutableNNC
export ControlCache, EmptyControlCache, MutableControlCache, ImmutableControlCache
export DistanceDataStructure, EmptyDistanceDS, BruteDistanceDS, TreeDistanceDS
export filter_neighborhood, addpoints
# export resetcache!
# export saveNN, loadNN!
export inball, knn, inball!, knn!, mutualknn!         # for metrics only (symmetric)
export inballF, knnF, inballF!, knnF!, mutualknnF!    # forwards
export inballB, knnB, inballB!, knnB!, mutualknnB!    # backwards

### Typedefs
abstract SampleSet
abstract NearNeighborCache{T}; begin
    ## NearNeighborCache implementations
    # TODO: consider EmptyNNC for online algorithms or streaming-batch
    immutable MutableNNC{T} <: NearNeighborCache{T}
        D::Vector{SparseVector{T,Int}}
        r::Vector{T}
    end
    MutableNNC{T}(N::Int, ::Type{T}) = MutableNNC{T}(Array(SparseVector{T,Int}, N), zeros(T, N))
    immutable ImmutableNNC{T} <: NearNeighborCache{T}
        D::SparseMatrixCSC{T,Int}
        r::Vector{T}
    end
    ImmutableNNC(M::MutableNNC) = ImmutableNNC(hcat(M.D...), M.r)
    MutableNNC{T}(I::ImmutableNNC{T}) = MutableNNC(size(I.D,1), T) # clear the cache
end

abstract ControlCache{U<:ControlInfo}; begin
    ## ControlCache implementations
    immutable EmptyControlCache{U} <: ControlCache{U} end
    EmptyControlCache{U}(::Type{U}) = EmptyControlCache{U}()
    EmptyControlCache() = EmptyControlCache(NullControl)
    immutable MutableControlCache{U} <: ControlCache{U}
        US::Vector{SparseVector{U,Int}}
    end
    MutableControlCache{U}(N::Int, ::Type{U}) = MutableControlCache{U}(Array(SparseVector{U,Int}, N))
    immutable ImmutableControlCache{U} <: ControlCache{U}
        US::SparseMatrixCSC{U,Int}
    end
    ImmutableControlCache(M::MutableControlCache) = ImmutableControlCache(hcat(M.U...))
    MutableControlCache{U}(I::ImmutableControlCache{U}) = MutableControlCache(size(I.US,1), U) # clear the cache
end

abstract DistanceDataStructure{T}; begin
    ## DistanceDataStructure implementations
    immutable EmptyDistanceDS{T} <: DistanceDataStructure{T} end
    EmptyDistanceDS{T}(::Type{T}) = EmptyDistanceDS{T}()
    immutable BruteDistanceDS{A<:AbstractMatrix,T} <: DistanceDataStructure{T}
        Dmat::A
    end
    BruteDistanceDS{T}(Dmat::AbstractMatrix{T}) = BruteDistanceDS{typeof(Dmat),T}(Dmat)
    immutable TreeDistanceDS{NNT<:NearestNeighbors.NNTree,T} <: DistanceDataStructure{T}
        tree::NNT
    end
    TreeDistanceDS{V}(tree::NearestNeighbors.NNTree{V}) = TreeDistanceDS{typeof(tree),eltype(V)}(tree)
end

### Construction
type MetricNN{M<:SymmetricDistance,S<:State,N<:NearNeighborCache,D<:DistanceDataStructure,U<:ControlCache} <: SampleSet
    V::Vector{S}
    dist::M
    init::S
    cache::N
    DS::D
    US::U
end
function MetricNN{S<:State}(V::Vector{S}, dist::SymmetricDistance, init::S)
    N = length(V)
    T = eltype(S)
    MetricNN(V, dist, init, MutableNNC(N,T), helper_data_structures(V, dist)...)
end

type QuasiMetricNN{M<:AsymmetricDistance,S<:State,N<:NearNeighborCache,
                   D<:DistanceDataStructure,U<:ControlCache} <: SampleSet
    V::Vector{S}
    dist::M
    init::S
    cacheF::N   # Everything's just pointers, but F and B may be different (e.g. CSC and CSR), trading memory for speed
    cacheB::N
    DSF::D
    USF::U
    DSB::D
    USB::U
end
function QuasiMetricNN{S<:State}(V::Vector{S}, dist::AsymmetricDistance, init::S)
    N = length(V)
    T = eltype(S)
    QuasiMetricNN(V, dist, init, MutableNNC(N,T), MutableNNC(N,T), helper_data_structures(V, dist)...)
end

## Fallbacks (should usually be extended)
helper_data_structures{S}(V::Vector{S}, d::SymmetricDistance) = (EmptyDistanceDS(eltype(S)),
                                                                 EmptyControlCache(controltype(d)))
function helper_data_structures{S}(V::Vector{S}, d::AsymmetricDistance)
    DS, US = EmptyDistanceDS(eltype(S)), EmptyControlCache(controltype(d))
    (DS, US, DS, US)
end

### Utilities

@inline function filter_neighborhood(n::AbstractSparseVector, f::BitVector)
    inds, ds = nonzeroinds(n), nonzeros(n)
    SparseVector(length(n), inds[f[inds]], ds[f[inds]])
end
addpoints(NN::MetricNN, W) = MetricNN([NN.V; W], NN.dist, NN.init)
addpoints(NN::QuasiMetricNN, W) = QuasiMetricNN([NN.V; W], NN.dist, NN.init)
length(NN::SampleSet) = length(NN.V)
getindex(NN::SampleSet, i::Int) = i > 0 ? NN.V[i] : NN.init
getindex(NN::SampleSet, I::AbstractVector) = [NN[i] for i in I]

### Serialization

### Near Neighbor Computation
# TODO: forwards/backwards should be probably moved into the datastructure; perhaps a type parameterization
@inline inball!(NN::MetricNN, v::Int, r, f::BitVector) = filter_neighborhood(inball!(NN, v, r), f)
@inline inballF!(NN::QuasiMetricNN, v::Int, r, f::BitVector) = filter_neighborhood(inballF!(NN, v, r), f)
@inline inballB!(NN::QuasiMetricNN, v::Int, r, f::BitVector) = filter_neighborhood(inballB!(NN, v, r), f)
for direction in ("", "F", "B")
    NNsym = (direction == "" ? :MetricNN : :QuasiMetricNN)
    inballD! = Symbol(:inball, direction, "!")
    cacheD = Symbol(:cache, direction)
    DSD = Symbol(:DS, direction)
    @eval $inballD!{M,S,N<:ImmutableNNC}(NN::$NNsym{M,S,N}, v::Int, r) = viewcol(NN.$cacheD.D, v)
    @eval function $inballD!{M,S,N<:MutableNNC}(NN::$NNsym{M,S,N}, v::Int, r)
        if !isdefined(NN.$cacheD.D, v)
            NN.$cacheD.D[v] = inball(NN.V, NN.dist, NN.$DSD, v, r, $(direction != "B"))
            NN.$cacheD.r[v] = r
        end
        NN.$cacheD.D[v]
    end
end

function inball(V::Vector, dist::PreMetric, DS::DistanceDataStructure, v::Int, r, forwards::Bool = true)
    allds = forwards ? colwise(dist, V[v], V) : colwise(dist, V, V[v])
    N = length(V)
    inds = Vector{Int}(0)
    ds = similar(allds, 0)
    @inbounds for i in 1:N
        if i != v && allds[i] <= r
            push!(inds, i)
            push!(ds, allds[i])
        end
    end
    SparseVector(N, inds, ds)
end

function inball{A<:Matrix}(V::Vector, dist::PreMetric, DS::BruteDistanceDS{A}, v::Int, r, forwards::Bool = true)
    N = length(V)
    inds = Vector{Int}(0)
    ds = similar(DS.Dmat, 0)
    @inbounds for i in 1:N
        if i != v && DS.Dmat[i,v] <= r
            push!(inds, i)
            push!(ds, allds[i])
        end
    end
    SparseVector(N, inds, ds)
end

function inball{A<:SparseMatrixCSC}(V::Vector, dist::PreMetric, DS::BruteDistanceDS{A}, v::Int, r, forwards::Bool = true)
    vcol = viewcol(DS.Dmat, v)
    allinds, allds = nonzeroinds(vcol), nonzeros(vcol)
    inds = Vector{Int}(0)
    ds = Vector{eltype(A)}(0)
    @inbounds for i in 1:length(allinds)
        if allinds[i] != v && allds[i] <= r
            push!(inds, allinds[i])
            push!(ds, allds[i])
        end
    end
    SparseVector(length(V), inds, ds)
end

function inball{S}(V::Vector{S}, dist::PreMetric, DS::TreeDistanceDS, v::Int, r, forwards::Bool = true)
    inds = inrange(DS.tree, V[v], r, true)
    inds = deleteat!(inds, searchsortedfirst(inds, v))
    SparseVector(length(V), inds, forwards ? colwise(dist, V[v], V[inds]) : colwise(dist, V[inds], V[v]))
end

function inball{S}(V::Vector{S}, dist::ChoppedPreMetric, DS::TreeDistanceDS, v::Int, r, forwards::Bool = true)
    allinds = inrange(DS.tree, V[v], r, true)
    allinds = deleteat!(allinds, searchsortedfirst(allinds, v))
    allds = forwards ? colwise(dist, V[v], V[allinds]) : colwise(dist, V[allinds], V[v])
    inds = Vector{Int}(0)
    ds = similar(allds, 0)
    @inbounds for i in 1:length(allinds)
        if allds[i] <= r
            push!(inds, allinds[i])
            push!(ds, allds[i])
        end
    end
    SparseVector(length(V), inds, ds)
end

for ff in [x -> Symbol("inball", x), x -> Symbol("inball", x, "!")]
    @eval $(ff("F"))(NN::MetricNN, args...) = $(ff(""))(NN, args...)
    @eval $(ff("B"))(NN::MetricNN, args...) = $(ff(""))(NN, args...)
end


module MadNLPPardiso

include(joinpath("..","deps","deps.jl"))

const INPUT_MATRIX_TYPE = :csc

import Libdl: dlopen, RTLD_DEEPBIND
import MadNLP:
    MadNLP, @kwdef, Logger, @debug, @warn, @error,
    SubVector, StrideOneVector, SparseMatrixCSC,
    SymbolicException,FactorizationException,SolveException,InertiaException,
    AbstractOptions, AbstractLinearSolver, set_options!,
    introduce, factorize!, solve!, improve!, is_inertia, inertia
import MKL_jll: libmkl_rt

@isdefined(libpardiso) && include("pardiso.jl")
include("pardisomkl.jl")

function __init__()
    check_deps()
    @isdefined(libpardiso) && dlopen(libpardiso,RTLD_DEEPBIND)
end

export MadNLPPardisoMKL

end # module


##############################################################################################################
# Instantiate BAU model to get some parameters and the the nice model to be used in subsequent optimizations
##############################################################################################################

# Create a folder to save results.
output_directory = joinpath((@__DIR__), "..", "..", "results", results_folder)
mkpath(output_directory)

# Load quintile income distribution scenario data.
income_distribution_raw = DataFrame(load(joinpath(@__DIR__, "..", "..", "data", quintile_income_scenario*"_quintile_distributions_consumption.csv")))

# Clean up and organize time-varying income distribution data into required NICE format (time × regions × quintiles).
income_distributions = get_quintile_income_shares(income_distribution_raw)

#---------------------------------------------------------------------------------------------------
# Create a baseline version of the model without CO₂ mitigation.
#---------------------------------------------------------------------------------------------------

# This includes the user-specifications but has no CO₂ mitigation policy (will be used to calculte global CO₂ policy).
bau_model = create_nice_recycle()
n_steps   = length(dim_keys(bau_model, :time))

# Run model and extract some generic values needed to set up optimizations below.
run(bau_model)
rice_backstop_prices  = bau_model[:emissions, :pbacktime] .* 1000

# Now set user-specified parameter settings.
update_param!(bau_model, :MIU, zeros(n_steps, 12))
update_param!(bau_model, :damage_elasticity, damage_elasticity)
update_param!(bau_model, :quintile_income_shares, income_distributions)
update_param!(bau_model, :recycle_share, recycle_share)
update_param!(bau_model, :global_carbon_tax, zeros(n_steps))
update_param!(bau_model, :rho, ρ)
update_param!(bau_model, :eta, η)

# If selected, allow elasticities to be calculated for all GDP values (not just those observed in studies).
if bound_gdp_elasticity == false
    update_param!(bau_model, :min_study_gdp, 1e-10)
    update_param!(bau_model, :max_study_gdp, +Inf)
end

run(bau_model)

# Create an instance of NICE with revenue recycling.
nice = create_nice_recycle()

# Set user-specified parameter settings.
set_param!(nice, :damage_elasticity, damage_elasticity)
set_param!(nice, :quintile_income_shares, income_distributions)
set_param!(nice, :recycle_share, recycle_share)
set_param!(nice, :rho, ρ)
set_param!(nice, :eta, η)

# If selected, allow elasticities to be calculated for all GDP values (not just those observed in studies).
if bound_gdp_elasticity == false
    set_param!(nice, :min_study_gdp, 1e-10)
    set_param!(nice, :max_study_gdp, +Inf)
end


# This file extends / overrides some of the Julia basic capability for linear algebra to add useful things for FE
# Most of the time these are "in-place" operations that Julia doesn't normally have an interface for, but we need.
# Also: these functions help us have a common interface with PETSc (which has more in-place capability than Julia)

"""
    Set all entries to zero
"""
function zero!(v)
    fill!(v, 0.)
end

"""
    v[i] += u
"""
function plusEquals!(v::Array, u::Array, i)
    v[i] += u
end

"""
    mat[i,j] += u
"""
function plusEquals!(mat::Matrix, u::Matrix, i, j)
    mat[i,j] += u
end

"""
    Zero rows in the matrix
"""
function zeroRows!(mat::Matrix, rows::Array)
    mat[rows,:] = 0
end

"""
    No-op (no need to assemble Julia objects)
"""
function assemble!(::Any)
end

"""
    Pass through
"""
function serializeToZero(array::Array)
    return array
end


#############################################
# BBBackpropTraceState (for backprop_trace) #
#############################################

struct SelectedLeafNode
    addr::Symbol
    value_ref::Expr
    gradient_var::Symbol
end

struct InternalNode
    addr::Symbol
    values_var::Symbol
    gradients_var::Symbol
end

struct BBBackpropTraceState
    generator_type::Type
    ir::BasicBlockIR
    trace::Symbol
    stmts::Vector{Expr}
    schema::Union{StaticAddressSchema,EmptyAddressSchema}
    value_refs::Dict{ValueNode,Expr}
    grad_vars::Dict{ValueNode,Symbol}
    leaf_nodes::Set{SelectedLeafNode}
    internal_nodes::Set{InternalNode}
end

function BBBackpropTraceState(generator_type, trace, stmts, schema, value_refs, grad_vars)
    leaf_nodes = Set{SelectedLeafNode}()
    internal_nodes = Set{InternalNode}()
    ir = get_ir(generator_type)
    BBBackpropTraceState(generator_type, ir, trace, stmts, schema, value_refs, grad_vars,
        leaf_nodes, internal_nodes)
end

###############################################
# BBBackpropParamsState (for backprop_params) #
###############################################

struct BBBackpropParamsState
    generator_type::Type
    ir::BasicBlockIR
    trace::Symbol
    stmts::Vector{Expr}
    value_refs::Dict{ValueNode,Expr}
    grad_vars::Dict{ValueNode,Symbol}
end

function BBBackpropParamsState(generator_type, trace, stmts, value_refs, grad_vars)
    ir = get_ir(generator_type)
    BBBackpropParamsState(generator_type, ir, trace, stmts, value_refs, grad_vars)
end

#####################################################
# code shared by backprop_trace and backprop_params #
#####################################################

function process!(state::Union{BBBackpropParamsState,BBBackpropTraceState}, node::JuliaNode)

    # if the output node is not floating point, don't do anything
    if !haskey(state.grad_vars, node.output)
        return # no statements
    end

    # call gradient function
    input_value_refs = [state.value_refs[in_node] for in_node in node.input_nodes]
    output_value_ref = state.value_refs[node.output]
    input_grad_increments = [gensym("incr") for _ in node.input_nodes]
    output_grad_var = state.grad_vars[node.output]
    gradient_fn = get_grad_fn(state.generator_type, node)
    push!(state.stmts, quote
        ($(input_grad_increments...),) = $gradient_fn(
            $output_grad_var, $output_value_ref, $(input_value_refs...))
    end)

    # increment input gradients
    inputs_do_ad = map((in_node) -> haskey(state.grad_vars, in_node), node.input_nodes)
    for (in_node, incr, do_ad) in zip(node.input_nodes, input_grad_increments, inputs_do_ad)
        if do_ad
            grad_var = state.grad_vars[in_node]
            push!(state.stmts, quote
                $grad_var += $incr
            end)
        end
    end
end

function process!(state::Union{BBBackpropParamsState,BBBackpropTraceState}, node::ArgsChangeNode)
    # skip
end

function process!(state::Union{BBBackpropParamsState,BBBackpropTraceState}, node::AddrChangeNode)
    # skip
end

function initialize_backprop!(ir::BasicBlockIR, stmts::Vector{Expr})

    # create a gradient variable for each value node, initialize them to zero.
    value_refs = Dict{ValueNode,Expr}()
    grad_vars = Dict{ValueNode,Symbol}()
    for (name, node) in ir.value_nodes
        value_refs[node] = value_trace_ref(:trace, node)
        grad_var = gensym("grad_$name")
        if node !== ir.output_node || !ir.output_ad
            if is_differentiable(get_type(node))
                grad_vars[node] = grad_var
                push!(stmts, quote
                    $grad_var = zero($(value_refs[node]))
                end)
            end
        end
    end

    # initialize the gradient variable for the output node
    if ir.output_node !== nothing && ir.output_ad
        grad_var = grad_vars[something(ir.output_node)]
        push!(stmts, quote
            $grad_var = retval_grad
        end)
    end

    (value_refs, grad_vars)
end

function input_gradients(ir::BasicBlockIR, grad_vars)
    input_grads_var = gensym("input_grads")
    input_grads = []
    for (node, has_grad) in zip(ir.arg_nodes, ir.args_ad)
        if has_grad
            push!(input_grads, grad_vars[node])
        else
            push!(input_grads, QuoteNode(nothing))
        end
    end
    Expr(:tuple, input_grads...)
end

function increment_input_gradients!(stmts, node, dist_or_gen, grad_vars, increments)
    inputs_do_ad = has_argument_grads(dist_or_gen)
    for (in_node, do_ad, incr) in zip(node.input_nodes, inputs_do_ad, increments)
        if do_ad
            if !haskey(grad_vars, in_node)
                error("$(dist_or_gen) has AD for an input that is not floating point, node: $node")
            end
            grad_var = grad_vars[in_node]
            push!(stmts, quote
                $grad_var += $incr
            end)
        end
    end
end

function increment_output_gradient!(stmts, node::AddrDistNode, grad_vars, increment)
    # NOTE: if the output has a gradient, then it must be a float...
    # but it may be a float and the may not hae a gradient, currently this is
    # silent; could warn?
    if has_output_grad(node.dist)
        if !haskey(grad_vars, node.output)
            error("Distribution $(node.dist) has AD but the return value is not floating point, node: $node")
        end
        output_grad_var = grad_vars[node.output]
        push!(stmts, quote
            $output_grad_var += $increment
        end)
    end
end

###################
# backprop_params #
###################

function process!(state::BBBackpropParamsState, node::AddrDistNode)

    # get gradient of log density with respect to output and inputs
    input_value_refs = [state.value_refs[in_node] for in_node in node.input_nodes]
    output_value_ref = Expr(:(.), :trace, QuoteNode(node.address))# state.value_refs[node.output]
    output_grad_incr = gensym("incr")
    input_grad_incrs = [gensym("incr") for _ in node.input_nodes]
    push!(state.stmts, quote
        ($output_grad_incr, $(input_grad_incrs...),) = Gen.logpdf_grad(
            $(QuoteNode(node.dist)), $output_value_ref, $(input_value_refs...))
    end)

    increment_output_gradient!(state.stmts, node, state.grad_vars, output_grad_incr)
    increment_input_gradients!(state.stmts, node, node.dist, state.grad_vars, input_grad_incrs)
end

function process!(state::BBBackpropParamsState, node::AddrGeneratorNode)

    # get gradients from generator
    output_do_ad = accepts_output_grad(node.gen)
    output_grad = output_do_ad ? state.grad_vars[node.output] : QuoteNode(nothing)
    addr = node.address
    input_grad_incrs = [gensym("incr") for _ in node.input_nodes]
    subtrace = Expr(:(.), :trace, QuoteNode(addr))
    push!(state.stmts, quote
        ($(input_grad_incrs...),) = backprop_params($(QuoteNode(node.gen)), $subtrace, $output_grad)
    end)

    increment_input_gradients!(state.stmts, node, node.gen, state.grad_vars, input_grad_incrs)
end

function codegen_backprop_params(gen::Type{T}, trace, retval_grad) where {T <: BasicGenFunction}
    ir = get_ir(gen)
    stmts = Expr[]

    # create a gradient variable for each value node, initialize them to zero.
    # also get trace references for each value node
    (value_refs, grad_vars) = initialize_backprop!(ir, stmts)

    # visit statements in reverse topological order
    state = BBBackpropParamsState(gen, :trace, stmts, value_refs, grad_vars)
    for node in reverse(ir.expr_nodes_sorted)
        process!(state, node)
    end

    # increment gradient accumulators for parameters
    for param in ir.params
        value_node = ir.value_nodes[param.name]
        grad_var = grad_vars[value_node]
        push!(stmts, quote
            gen.params_grad[$(QuoteNode(param.name))] += $grad_var
        end)
    end

    # return statement
    push!(stmts, quote
        return $(input_gradients(ir, grad_vars))
    end)
    Expr(:block, stmts...)
end

push!(Gen.generated_functions, quote
@generated function Gen.backprop_params(gen::Gen.BasicGenFunction, trace, retval_grad)
    Gen.codegen_backprop_params(gen, trace, retval_grad)
end
end)

##################
# backprop_trace #
##################

function process!(state::BBBackpropTraceState, node::AddrDistNode)

    # get gradient of log density with respect to output and inputs
    input_value_refs = [state.value_refs[in_node] for in_node in node.input_nodes]
    output_value_ref = state.value_refs[node.output]
    output_grad_incr = gensym("incr")
    input_grad_incrs = [gensym("incr") for _ in node.input_nodes]
    push!(state.stmts, quote
        ($output_grad_incr, $(input_grad_incrs...),) = Gen.logpdf_grad(
            $(QuoteNode(node.dist)), $output_value_ref, $(input_value_refs...))
    end)

    increment_output_gradient!(state.stmts, node, state.grad_vars, output_grad_incr)
    increment_input_gradients!(state.stmts, node, node.dist, state.grad_vars, input_grad_incrs)

    # handle selected address
    addr = node.address
    if isa(state.schema, StaticAddressSchema) && addr in leaf_node_keys(state.schema)
        if !haskey(state.grad_vars, node.output)
            error("Selected a random choice that is not floating point for gradient: $node")
        end
        output_grad_var = state.grad_vars[node.output]
        push!(state.leaf_nodes, SelectedLeafNode(addr, Expr(:(.), :trace, QuoteNode(addr)), output_grad_var))
    end
end

function process!(state::BBBackpropTraceState, node::AddrGeneratorNode)

    # get gradients from generator and handle selection
    output_do_ad = accepts_output_grad(node.gen)
    output_grad = output_do_ad ? state.grad_vars[node.output] : QuoteNode(nothing)
    addr = node.address
    input_grad_incrs = [gensym("incr") for _ in node.input_nodes]
    subtrace = Expr(:(.), :trace, QuoteNode(addr))
    has_selection = isa(state.schema, StaticAddressSchema) && addr in internal_node_keys(state.schema)
    if has_selection
        selection = Expr(:call, :static_get_internal_node, :selection, QuoteNode(Val(addr)))
        value_trie = gensym("value_trie_$addr")
        gradient_trie = gensym("gradient_trie_$addr")
        push!(state.internal_nodes, InternalNode(addr, value_trie, gradient_trie))
    else
        selection = quote Gen.EmptyAddressSet() end
        value_trie = :_
        gradient_trie = :_
    end
    push!(state.stmts, quote
        (($(input_grad_incrs...),), $value_trie, $gradient_trie) = backprop_trace(
            $(QuoteNode(node.gen)), $subtrace, $selection, $output_grad)
    end)

    increment_input_gradients!(state.stmts, node, node.gen, state.grad_vars, input_grad_incrs)
end

const backprop_values_trie = gensym("values")
const backprop_gradients_trie = gensym("gradients")

function choice_trie_construction(leaf_nodes_set, internal_nodes_set)
    leaf_nodes = collect(leaf_nodes_set)
    quoted_leaf_keys = map((node) -> QuoteNode(node.addr), leaf_nodes)
    leaf_values = map((node) -> node.value_ref, leaf_nodes)
    leaf_gradients = map((node) -> node.gradient_var, leaf_nodes)
    internal_nodes = collect(internal_nodes_set)
    quoted_internal_keys = map((node) -> QuoteNode(node.addr), internal_nodes)
    internal_values = map((node) -> node.values_var, internal_nodes)
    internal_gradients = map((node) -> node.gradients_var, internal_nodes)
    quote
        $backprop_values_trie = StaticAssignment(
            NamedTuple{($(quoted_leaf_keys...),)}(($(leaf_values...),)),
            NamedTuple{($(quoted_internal_keys...),)}(($(internal_values...),)))
        $backprop_gradients_trie = StaticAssignment(
            NamedTuple{($(quoted_leaf_keys...),)}(($(leaf_gradients...),)),
            NamedTuple{($(quoted_internal_keys...),)}(($(internal_gradients...),)))
    end
end

function codegen_backprop_trace(gen::Type{T}, trace, selection, retval_grad) where {T <: BasicGenFunction}
    schema = get_address_schema(selection)
    ir = get_ir(gen)
    stmts = Expr[]

    # create a gradient variable for each value node, initialize them to zero.
    # also get trace references for each value node
    (value_refs, grad_vars) = initialize_backprop!(ir, stmts)

    # visit statements in reverse topological order
    state = BBBackpropTraceState(gen, :trace, stmts, schema, value_refs, grad_vars)
    for node in reverse(ir.expr_nodes_sorted)
        process!(state, node)
    end

    # construct values and gradients static choice tries
    push!(stmts, choice_trie_construction(state.leaf_nodes, state.internal_nodes))

    # return statement
    push!(stmts, quote
        return ($(input_gradients(ir, grad_vars)), $backprop_values_trie, $backprop_gradients_trie)
    end)
    Expr(:block, stmts...)
end

push!(Gen.generated_functions, quote
@generated function Gen.backprop_trace(gen::Gen.BasicGenFunction{T,U}, trace::U, selection, retval_grad) where {T,U}
    schema = get_address_schema(selection)
    if !(isa(schema, StaticAddressSchema) || isa(schema, EmptyAddressSchema))
        # try to convert it to a static address set
        return quote backprop_trace(gen, trace, StaticAddressSet(selection), retval_grad) end
    end
    Gen.codegen_backprop_trace(gen, trace, selection, retval_grad)
end
end)


"""
    GenericRoutine{name} <: Function

A `GenericRoutine` cannot be directly execute on a quantum
device. It is a Julia `Function` that returns `Operation`,
and `Operation` can be executed on quanutm device.

!!! note
    An instance of `GenericRoutine` should be treated like
    `Function`.
"""
abstract type GenericRoutine{name} <: Function end

"""
    abstract type Routine end

This defines operations that one can execute on a quantum device.
"""
abstract type Routine end

"""
    abstract type IntrinsicRoutine <: Routine end

`IntrinsicRoutine` are the routine that can be executed
by the compile target that is pre-defined in the compiler.
"""
abstract type IntrinsicRoutine <: Routine end

"""
    Operation{P, Args} <: Routine

An `Operation` is a user defined composite routine
that can be called in other `Operation` or execute
on target device.
"""
struct Operation{P,Args} <: Routine
    parent::P
    args::Args
end

"""
    AdjointOperation{P} <: Routine

An `AdjointOperation` is the adjoint of another
`Routine`.
"""
struct AdjointOperation{P} <: Routine
    parent::P
end

Base.adjoint(x::Routine) = AdjointOperation(x)
Base.adjoint(x::AdjointOperation) = x.parent

routine_name(::Type) = nothing
routine_name(x) = routine_name(typeof(x))
routine_name(::Type{<:GenericRoutine{name}}) where {name} = name
routine_name(::Type{T}) where {T<:IntrinsicRoutine} = nameof(T)
routine_name(::Type{<:Operation{P}}) where {P} = routine_name(P)
routine_name(::Type{<:AdjointOperation{P}}) where {P} = Symbol("adjoint_", routine_name(P))

# this is only a place holder
# we are not using struct because
# singleton types will get const prop
struct MeasureResult{T}
    result::T
end

@inline function Base.:(==)(lhs::MeasureResult, rhs)
    measure_cmp(lhs, rhs)
end

@inline function Base.:(==)(lhs, rhs::MeasureResult)
    measure_cmp(lhs, rhs)
end

@inline function Base.:(==)(lhs::MeasureResult, rhs::MeasureResult)
    measure_cmp(lhs, rhs)
end

@noinline function measure_cmp(lhs, rhs)
    throw(IntrinsicError("cannot compare measurement result outside @device"))
end

struct Chain <: Routine
    args::Vector{Any}
    Chain(args...) = new(collect(Any, args))
end

# prop adjoint to leaves
Base.adjoint(c::Chain) = Chain(map(adjoint, c.args))

struct Typed
    obj
    type
end

struct Gate <: Routine
    operation # SSAValue/Routine
    locations # SSAValue/Locations
end

Base.adjoint(x::Gate) = Gate(adjoint(x.operation), x.locations)

struct Ctrl <: Routine
    gate::Gate
    ctrl # SSAValue/CtrlLocation
end

Base.adjoint(x::Ctrl) = Ctrl(adjoint(x.gate), x.ctrl)

struct BlockIR
    parent::IRCode
    nqubits::Int
    circuit::Chain
end

@as_record Chain
@as_record Gate
@as_record Ctrl
@as_record Typed
@as_record Operation
@as_record AdjointOperation

Base.length(c::Chain) = length(c.args)

function leaves(root::Chain)
    nodes = []
    sizehint!(nodes, length(root))
    for each in root.args
        if each isa Chain
            append!(nodes, leaves(each::Chain))
        else
            push!(nodes, each)
        end
    end
    return nodes
end

function Base.:(==)(lhs::Chain, rhs::Chain)
    mapreduce(==, &, lhs.args, rhs.args)
end

function Base.:(==)(lhs::Gate, rhs::Gate)
    lhs.operation == rhs.operation && lhs.locations == rhs.locations
end

function Base.:(==)(lhs::Ctrl, rhs::Ctrl)
    lhs.gate == rhs.gate && lhs.ctrl == rhs.ctrl
end


#=
  Structs for the approxChol solver
=#

"""
  LLp elements are all in the same column.
  row tells us the row, and val is the entry.
  val is set to zero for some edges that we should remove.
  next gives the next in the column.  It points to itself to terminate.
  reverse is the index into lles of the other copy of this edge,
  since every edge is stored twice as we do not know the order of elimination in advance.
"""
mutable struct LLp{Tind,Tval}
    row::Tind
    val::Tval
    next::LLp{Tind,Tval}
    reverse::LLp{Tind,Tval}

    LLp{Tind,Tval}() where {Tind,Tval} = (x = new(zero(Tind), zero(Tval)); x.next = x; x.reverse = x)
    LLp{Tind,Tval}(row, val, next, rev) where {Tind,Tval} = new(row, val, next, rev)
    LLp{Tind,Tval}(row, val) where {Tind,Tval} = (x = new(row, val); x.next = x; x.reverse = x)
    LLp{Tind,Tval}(row, val, next) where {Tind,Tval} = (x = new(row, val, next); x.reverse = x)
end

"""
  LLmatp is the data structure used to maintain the matrix during elimination.
  It stores the elements in each column in a singly linked list (only next ptrs)
  Each element is an LLp (linked list pointer).
  The head of each column is pointed to by cols.

  We probably can get rid of degs - as it is only used to store initial degrees.
"""
mutable struct LLmatp{Tind,Tval}
    n::Int64
    degs::Array{Tind,1}
    cols::Array{LLp{Tind,Tval},1}
    lles::Array{LLp{Tind,Tval},1}
end

# these are the types we use with a fixed ordering
struct LLord{Tind,Tval}
    row::Tind
    next::Tind
    val::Tval
end

mutable struct LLMatOrd{Tind,Tval}
    n::Int64
    cols::Array{Tind,1}
    lles::Array{LLord{Tind,Tval},1}
end

struct LLcol{Tind,Tval}
      row::Tind
      ptr::Tind
      val::Tval
  end

#=============================================================

LDLinv

=============================================================#

"""
  LDLinv contains the information needed to solve the Laplacian systems.
  It does it by applying Linv, then Dinv, then Linv (transpose).
  But, it is specially constructed for this particular solver.
  It does not explicitly make the matrix triangular.
  Rather, col[i] is the name of the ith col to be eliminated
"""
mutable struct LDLinv{Tind,Tval}
    col::Array{Tind,1}
    colptr::Array{Tind,1}
    rowval::Array{Tind,1}
    fval::Array{Tval,1}
    d::Array{Tval,1}
end

#=============================================================

ApproxCholPQ
the data strcture we use to keep track of degrees

=============================================================#

struct ApproxCholPQElem{Tind}
    prev::Tind
    next::Tind
    key::Tind
end

"""
  An approximate priority queue.
  Items are bundled together into doubly-linked lists with all approximately the same key.
  minlist is the min list we know to be non-empty.
  It should always be a lower bound.
  keyMap maps keys to lists
"""
mutable struct ApproxCholPQ{Tind}
    elems::Array{ApproxCholPQElem{Tind},1} # indexed by node name
    lists::Array{Tind,1}
    minlist::Int
    nitems::Int
    n::Int
end


module SEAL

# SEAL_jll provides `libsealc`, which we will use in this package
using SEAL_jll

"""
    SEALObject

Abstract parent type for all types based on SEAL classes.
"""
abstract type SEALObject end

"""
    gethandle(object::SEALObject)

Return the raw C pointer to where `object` resides in memory.
"""
@inline gethandle(object::SEALObject) = object.handle

"""
    sethandle!(object::SEALObject, handle)

Set the underlying raw C pointer to where `object` resides in memory to `handle`.
"""
@inline sethandle!(object::SEALObject, handle) = object.handle = handle

"""
    destroy!(object::SEALObject)

Call the corresponding destruction function on `object` to free up memory and reset object handle to
a null pointer. If `object` is not allocated, `destroy!` will not do anything.
"""
function destroy!(object::SEALObject) end

"""
    isnull(object::SEALObject)

Return true if the object handle is a null pointer and false otherwise.
"""
@inline isnull(object::SEALObject) = gethandle(object) == C_NULL

"""
    isallocated(object::SEALObject)

Return true if the object is allocated, i.e., if it is not null.
"""
@inline isallocated(object::SEALObject) = !isnull(object)

export SEALObject, gethandle, sethandle!, destroy!, isnull, isallocated

Base.unsafe_convert(::Type{Ptr{Cvoid}}, object::SEALObject) = gethandle(object)

include("auxiliary.jl")
# Julia-only auxiliary methods -> no exports

include("version.jl")
export version_major, version_minor, version_patch, version

include("modulus.jl")
export Modulus, SecLevelType, bit_count, value, coeff_modulus_create, coeff_modulus_bfv_default

include("serialization.jl")
export ComprModeType, SEALHeader, load_header!

include("encryptionparams.jl")
export EncryptionParameters, SchemeType, poly_modulus_degree,
       set_poly_modulus_degree!, set_coeff_modulus!, coeff_modulus,
       scheme, plain_modulus, set_plain_modulus!, plain_modulus_batching, parms_id, save!,
       save_size, load!

include("context.jl")
export SEALContext, first_parms_id, last_parms_id, get_context_data, key_context_data,
       first_context_data, parameter_error_message, using_keyswitching
export ContextData, chain_index, parms, parms_id, total_coeff_modulus_bit_count, qualifiers,
       next_context_data
export EncryptionParameterQualifiers, using_batching

include("publickey.jl")
export PublicKey, parms_id

include("secretkey.jl")
export SecretKey, parms_id, save!, load!

include("galoiskeys.jl")
export GaloisKeys, parms_id

include("relinkeys.jl")
export RelinKeys, parms_id, save_size, save!, load!

include("keygenerator.jl")
export KeyGenerator, create_public_key!, create_public_key, secret_key, create_relin_keys!,
       create_relin_keys, create_galois_keys!

include("plaintext.jl")
export Plaintext, scale, scale!, parms_id, to_string, save_size, save!

include("ciphertext.jl")
export Ciphertext, scale, scale!, parms_id, size, length, save_size, save!, load!, reserve!

include("encryptor.jl")
export Encryptor, set_secret_key!, encrypt!, encrypt_symmetric, encrypt_symmetric!

include("evaluator.jl")
export Evaluator, square!, square_inplace!, relinearize!, relinearize_inplace!, rescale_to_next!,
       rescale_to_next_inplace!, multiply_plain!, multiply_plain_inplace!, multiply!,
       multiply_inplace!, mod_switch_to!, mod_switch_to_inplace!, mod_switch_to_next!,
       mod_switch_to_next_inplace!, add!, add_inplace!, add_plain!, add_plain_inplace!,
       rotate_vector!, rotate_vector_inplace!, rotate_rows!, rotate_rows_inplace!,
       rotate_columns!, rotate_columns_inplace!, complex_conjugate!, complex_conjugate_inplace!,
       negate!, negate_inplace!, sub!, sub_inplace!

include("decryptor.jl")
export Decryptor, decrypt!, invariant_noise_budget

include("ckks.jl")
export CKKSEncoder, slot_count, encode!, decode!

include("batchencoder.jl")
export BatchEncoder, slot_count, encode!, decode!

include("memorymanager.jl")
export MemoryPoolHandle, alloc_byte_count, memory_manager_get_pool

end


module ClassicalSpinMonteCarlo

using StaticArra
using QuantumLattices

using LinearAlgebra::norm,cross

end # module


#***************************************************************************
#                     Molecular Dynamics Potentials (MDP)
#                            CESMIX-MIT Project
#
# Contributing authors: Ngoc-Cuong Nguyen (cuongng@mit.edu, exapde@gmail.com)
#***************************************************************************

function readMDP(filename, mode)

# natom, lattice, pbc, energy, stress, id, type, group, move, atomic number, atomic mass, charge, position, velocity, force

if mode==0 # binary
    frame = reinterpret(Float64,read(filename));
else # text
    tmp = Main.DelimitedFiles.readdlm(filename, ',');
    tmp = String.(reduce(vcat, split.(tmp))); # convert to 1D array of strings
    frame = parse.(Float64, tmp);  # convert to 1D array of Float64
end

config = initializeconfig(0);

natomall = 0;
nframe = 0;
n2 = 0;
n = length(frame);
while (n2 < n)
    nframe = nframe + 1;

    # natom, lattice, pbc, energy, stress, id, type, group, move, atomic number, atomic mass, charge, position, velocity, force
    natom = Int64(frame[n2+1]);  # number of atoms
    ncl = Int64(frame[n2+2]);    # lattice
    ncp = Int64(frame[n2+3]);    # periodic boundary conditions
    nce = Int64(frame[n2+4]);    # energy
    ncs = Int64(frame[n2+5]);    # stress
    nci = Int64(frame[n2+6]);    # atom id
    nct = Int64(frame[n2+7]);    # atom type
    ncg = Int64(frame[n2+8]);    # atom groups
    nco = Int64(frame[n2+9]);    # atom moved/fixed flags
    ncz = Int64(frame[n2+10]);   # atomic numbers
    ncm = Int64(frame[n2+11]);   # atomic masses
    ncq = Int64(frame[n2+12]);   # atom charges
    ncx = Int64(frame[n2+13]);   # positions
    ncv = Int64(frame[n2+14]);   # atom velocities
    ncf = Int64(frame[n2+15]);   # atom forces
    n2 = n2+15;

    if nframe==1
        config.ncx = ncx;
        config.ncv = ncv;
        config.ncq = ncq;
        config.nce = nce;
        config.ncf = ncf;
        config.ncs = ncs;
        config.nci = nci;
        config.nct = nct;
        config.ncg = ncg;
        config.ncz = ncz;
        config.ncm = ncm;
        config.nco = nco;
        config.ncl = ncl;
        config.ncp = ncp;
        config.nconfigs = 1;
        config.dim = ncx;
    end

    config.natom = hcat(config.natom, reshape([natom],1,1))

    # lattice, pbc, energy, stress
    n1 = n2 + 1; n2 = n2 + ncl;
    if nframe == 1 & ncl>0
        config.lattice = reshape(frame[n1:n2], (ncl, 1));
    elseif ncl>0
        config.lattice = hcat(config.lattice, reshape(frame[n1:n2], (ncl, 1)));
    end

    n1 = n2 + 1; n2 = n2 + ncp;
    if nframe == 1 & ncp>0
        config.pbc = reshape(frame[n1:n2], (ncp, 1));
    elseif ncp>0
        config.pbc = hcat(config.pbc, reshape(frame[n1:n2], (ncp, 1)));
    end

    n1 = n2 + 1; n2 = n2 + nce;
    if nframe == 1 & nce>0
        config.e = reshape(frame[n1:n2], (nce, 1));
    elseif nce>0
        config.e = hcat(config.e, reshape(frame[n1:n2], (nce, 1)));
    end

    n1 = n2 + 1; n2 = n2 + ncs;
    if nframe == 1 & ncs>0
        config.stress = reshape(frame[n1:n2], (ncs, 1));
    elseif ncs>0
        config.stress = hcat(config.stress, reshape(frame[n1:n2], (ncs, 1)));
    end

    k = nci + nct + ncg + nco + ncz + ncm + ncq + ncx + ncv + ncf;
    n1 = n2 + 1; n2 = n2 + natom*k;
    tmp = reshape(frame[n1:n2], (k, natom));

    if nframe == 1 & nci>0
        config.tags = tmp[1:nci,:];
    elseif nci>0
        config.tags = hcat(config.tags, tmp[1:nci,:]);
    end

    if nframe == 1 & nct>0
        config.t = tmp[(1+nci):(nci+nct),:];
    elseif nct>0
        config.t = hcat(config.t, tmp[(1+nci):(nci+nct),:]);
    end

    if nframe == 1 & ncg>0
        config.group = tmp[(1+nci+nct):(nci+nct+ncg),:];
    elseif ncg>0
        config.group = hcat(config.group, tmp[(1+nci+nct):(nci+nct+ncg),:]);
    end

    if nframe == 1 & nco>0
        config.move = tmp[(1+nci+nct+ncg):(nci+nct+ncg+nco),:];
    elseif nco>0
        config.move = hcat(config.move, tmp[(1+nci+nct+ncg):(nci+nct+ncg+nco),:]);
    end

    if nframe == 1 & ncz>0
        config.Z = tmp[(1+nci+nct+ncg+nco):(nci+nct+ncg+nco+ncz),:];
    elseif ncz>0
        config.Z = hcat(config.Z, tmp[(1+nci+nct+ncg+nco):(nci+nct+ncg+nco+ncz),:]);
    end

    if nframe == 1 & ncm>0
        config.mass = tmp[(1+nci+nct+ncg+nco+ncz):(nci+nct+ncg+nco+ncz+ncm),:];
    elseif ncm>0
        config.mass = hcat(config.mass, tmp[(1+nci+nct+ncg+nco+ncz):(nci+nct+ncg+nco+ncz+ncm),:]);
    end

    if nframe == 1 & ncq>0
        config.q = tmp[(1+nci+nct+ncg+nco+ncz+ncm):(nci+nct+ncg+nco+ncz+ncm+ncq),:];
    elseif ncq>0
        config.q = hcat(config.q, tmp[(1+nci+nct+ncg+nco+ncz+ncm):(nci+nct+ncg+nco+ncz+ncm+ncq),:]);
    end

    if nframe == 1 & ncx>0
        config.x = tmp[(1+nci+nct+ncg+nco+ncz+ncm+ncq):(nci+nct+ncg+nco+ncz+ncm+ncq+ncx),:];
    elseif ncx>0
        config.x = hcat(config.x, tmp[(1+nci+nct+ncg+nco+ncz+ncm+ncq):(nci+nct+ncg+nco+ncz+ncm+ncq+ncx),:]);
    end

    if nframe == 1 & ncv>0
        config.v = tmp[(1+nci+nct+ncg+nco+ncz+ncm+ncq+ncx):(nci+nct+ncg+nco+ncz+ncm+ncq+ncx+ncv),:];
    elseif ncv>0
        config.v = hcat(config.v, tmp[(1+nci+nct+ncg+nco+ncz+ncm+ncq+ncx):(nci+nct+ncg+nco+ncz+ncm+ncq+ncx+ncv),:]);
    end

    if nframe == 1 & ncf>0
        config.f = tmp[(1+nci+nct+ncg+nco+ncz+ncm+ncq+ncx+ncv):(nci+nct+ncg+nco+ncz+ncm+ncq+ncx+ncv+ncf),:];
    elseif ncf>0
        config.f = hcat(config.f, tmp[(1+nci+nct+ncg+nco+ncz+ncm+ncq+ncx+ncv):(nci+nct+ncg+nco+ncz+ncm+ncq+ncx+ncv+ncf),:]);
    end

    natomall = natomall + natom;
end
config.natomall = natomall;

return config

end


# Analyse multisession solution progression

using Caesar
# IIF and KDE should not be necessary, but adding for late night run safety
using IncrementalInference
using KernelDensityEstimate

# get credentials
addrdict = getcredentials(nparticles=true, multisession=true)

# setup cloudgraphs interface
cloudGraph, addrdict = standardcloudgraphsetup(addrdict=addrdict)

# extras
multisessions = addrdict["multisession"]
N = parse(Int,addrdict["num particles"])

# RESET ENTIRE SYSTEM
removeMultisessions!(cloudGraph, session=multisessions[1])
removeMultisessions!(cloudGraph, session=multisessions[2])
removeMultisessions!(cloudGraph, session=multisessions[3])

addrdict["session"] = multisessions[1]
resetconvertdb(addrdict=addrdict)
setDBAllReady!(cloudGraph.neo4j.connection, addrdict["session"])
convertdb(addrdict=addrdict, N=N)
addrdict["session"] = multisessions[2]
resetconvertdb(addrdict=addrdict)
setDBAllReady!(cloudGraph.neo4j.connection, addrdict["session"])
convertdb(addrdict=addrdict, N=N)
addrdict["session"] = multisessions[3]
resetconvertdb(addrdict=addrdict)
setDBAllReady!(cloudGraph.neo4j.connection, addrdict["session"])
convertdb(addrdict=addrdict, N=N)

resultsfolder = "results_$(Dates.now(Dates.UTC))"
run(`mkdir -p $(resultsfolder)`)

for ANAITER in Int[1;2;3;4;5]
  for session in multisessions
    @show addrdict["session"] = session
    slamindb(addrdict=addrdict, iterations=1)
  end
  for session in multisessions
    multisessionsl = Vector{String}(setdiff(multisessions, [session]))
    # These are the base steps in rmInstMultisessionPriors!
    lm2others = getLandmOtherSessNeoIDs(cloudGraph,
                    session=session,multisessions=multisessionsl)
    # grab local subgraph using NeoIDs
    sfg, lms = getLocalSubGraphMultisession(cloudGraph, lm2others,
                    session=session, numneighbors=1)

    fg = initfg(sessionname=session, cloudgraph=cloudGraph)
    fullLocalGraphCopy!(fg)

    jldfile = string(resultsfolder,"/", session, "_$(ANAITER).jld")
    savejld(fg, file=jldfile)
    for sym in lms
      @show sym
      fv = getVert(fg, sym, nt=:fnc)
      ffv = getData(fv).fnc.usrfnc!
      # plotKDE(ffv.belief)

      p = getVertKDE(fg, sym)
      pl = plotKDE([ffv.belief; p], c=["red";"green"], levels=3, legend=["MultiSess"; "Current"], title=string(sym), fill=true)

      filename = string(resultsfolder,"/", session,"_",sym, "_$(ANAITER)")
      saveplot(pl, name=filename, frt=:pdf, nw=true)
      saveplot(pl, name=filename, frt=:png, nw=true)
    end
  end
end


"""
bio_seq : logic associated with dealing with proteins/translation and
other sequence functions unique to their biological interpretation.
In general our order of aa's is ARNDCQEGHILKMFPSTWYV- .
This is the "alphabetical" order and comes from cross-comptability with Gremlin-deived MRFs.
"""

module bio_seq

using Unicode

islower(s) = all(c->islowercase(c) | isspace(c), s);

nucs = "ACGT"
quick_rev = Dict{String, String}("A" => "T", "C" => "G", "G" => "C", "T" => "A")
codon_table = Dict{String, Char}("..." => '-', "TGT"=>'C',"GAC"=>'D',"TTC"=>'F',"TAG"=>'*',"GTG"=>'V',"CCT"=>'P', "GCT"=>'A',"GGC"=>'G',"CGG"=>'R',
									 "ATT"=>'I',"GAT"=>'D',"CAG"=>'Q',"ATG"=>'M',"CTC"=>'L',"TCT"=>'S',"CGT"=>'R',"ACG"=>'T',"AGA"=>'R',
									 "TGG"=>'W',"TCG"=>'S',"TTA"=>'L',"AGT"=>'S',"CGA"=>'R',"TGC"=>'C',"CAA"=>'Q',"TTG"=>'L',"AAT"=>'N',
									 "AAC"=>'N',"TAA"=>'*',"TAC"=>'Y',"CCA"=>'P',"ACT"=>'T',"TAT"=>'Y',"CAC"=>'H',"CCG"=>'P',"GCA"=>'A',
									 "GAA"=>'E',"ACA"=>'T',"GTA"=>'V',"AGG"=>'R',"AGC"=>'S',"TCA"=>'S',"GGT"=>'G',"GCC"=>'A',"TGA"=>'*',
									 "GTC"=>'V',"TTT"=>'F',"CAT"=>'H',"AAG"=>'K',"AAA"=>'K',"CTT"=>'L',"CTA"=>'L',"ATC"=>'I',"ACC"=>'T',
									 "TCC"=>'S',"CCC"=>'P',"GAG"=>'E',"GCG"=>'A',"CTG"=>'L',"CGC"=>'R',"GTT"=>'V',"GGA"=>'G',"ATA"=>'I',
									 "GGG"=>'G')
aa_to_num = Dict{Char, Int64}('A'=>1, 'R'=>2, 'N'=>3, 'D'=>4, 'C'=>5, 'Q'=>6,
									'E'=>7, 'G'=>8, 'H'=>9, 'I'=>10, 'L'=>11, 'K'=>12, 'M'=>13, 'F'=>14, 'P'=>15,
									'S'=>16, 'T'=>17, 'W'=>18, 'Y'=>19, 'V'=>20, '-'=>21, '*'=>21)

function load_fasta(file_name)
	in_file = open(file_name)
	in_read = readlines(in_file)
	close(in_file)

	seqs = Dict{String, String}()
	cur_head = ""
	cur_seq = String[]
	for line in in_read
		if !(isempty(line)) && line[1] == '>'
			if cur_head != ""
				seqs[cur_head] = join(cur_seq)
			end
			cur_head = strip(line[2:end])
			empty!(cur_seq)
		else
			push!(cur_seq, strip(line))
		end
	end
	seqs[cur_head] = join(cur_seq)
	return seqs
end

function convert_protein(prot_seq)
	prot_len = length(prot_seq)
	big_vector = zeros(Float64, prot_len * 21)
	#and fill it up like a very sparse row vector. One-hot.
	#Order is ARNDCQEGHILKMFPSTWYV-
	prot_fill = Dict{Char, Int64}('A'=>1, 'R'=>2, 'N'=>3, 'D'=>4, 'C'=>5, 'Q'=>6, 'E'=>7, 'G'=>8, 'H'=>9, 'I'=>10,
									'L'=>11, 'K'=>12, 'M'=>13, 'F'=>14, 'P'=>15, 'S'=>16, 'T'=>17, 'W'=>18, 'Y'=>19, 'V'=>20, '-'=>21, 'X'=>21, '*'=>21)
	chr_count = 0
	for chr in prot_seq
		big_vector[(21 * chr_count) + prot_fill[chr]] = 1.0
		chr_count += 1
	end
	return big_vector' #more convenient for it to be a row vector later.
end

function translate(seq, skip=1)
	#adding ellipses for dashes we can't change.
	seq = replace(seq, "*"=>"") #just skip the insertions which we won't change here.
	codons = [seq[x:min(x+2, length(seq))] for x in skip:3:length(seq)]
	out_seq = Char[]
	for codon in codons
		if length(codon) < 3
			continue
		end
		if islower(codon)
			push!(out_seq, lowercase(codon_table[uppercase(codon)]))
		else
			push!(out_seq, codon_table[uppercase(codon)])
		end
	end
	return join(out_seq)
end

function rev_comp(seq)
	return(join([quick_rev[string(B)] for B in reverse(seq)], ""))
end

function get_two_base()
	two_base = [join([one, two]) for one in nucs for two in nucs]
end

function get_three_base() #64.
	three_base = [join([one, two, three]) for one in nucs for two in nucs for three in nucs]
	return three_base
end

function get_four_base() #256 combos.
	four_base = [join([one, two, three, four]) for one in nucs for two in nucs for three in nucs for four in nucs]
	return four_base
end

function get_six_base()
	six_base = [join([one, two, three, four, five, six]) for one in nucs for two in nucs for three in nucs for four in nucs for five in nucs for six in nucs]
	return six_base
end

function precompute_dicodons() #64 * 64 codons = 4096 options in dictionary.
	#Let's precompute all the possible codon pairs to translate very quickly.
	all_bases = get_six_base()
	di_codon_dict = Dict{AbstractString, Tuple{Int64, Int64}}()
	for di_codon in all_bases
		di_codon_dict[di_codon] = translate_to_tuple(di_codon)
	end
	return di_codon_dict
end

function get_base_dict()
	mut_dict = Dict{Int64, Array{AbstractString, 1}}()
	mut_dict[2] = get_two_base()
	mut_dict[3] = get_three_base()
	mut_dict[4] = get_four_base()
	return mut_dict
end

function translate_to_tuple(seq)
	#adding ellipses for dashes we can't change.
	#seq = replace(seq, "*"=>"") #just skip the insertions which we won't change here.
	codons = [seq[x:min(x+2, length(seq))] for x in 1:3:length(seq)]
	out_seq = Char[]
	for codon in codons
		codon = uppercase(codon)
		if length(codon) < 3
			continue
		end
		push!(out_seq, codon_table[uppercase(codon)])
	end
	out_tuple::Tuple{Int64, Int64} = (aa_to_num[out_seq[1]], aa_to_num[out_seq[2]])
	return out_tuple
end

function translate_constrained_maybe_map(nuc, trns; do_map = false)
	gap_poses = Dict{Int64, Int64}()
	for (i, j) in nuc.gap_pos
		if i in keys(gap_poses)
			gap_poses[i] += j
		else
			gap_poses[i] = j
		end
	end

	true_nuc = Char[]
	i_count = 1
	explicit_mapping = Dict{Int64, Int64}()
	aa_counter = 0
	adder = 0

	for unmapped_n in 1:length(nuc.nucs)
		n = unmapped_n + (trns - 1)
		aux_n = n + adder
		if !(n in nuc.skip_trn)
			#then we ignore n for our purposes.
			explicit_mapping[aux_n] = div(aa_counter, 3) + 1
			aa_counter += 1
			push!(true_nuc, nuc.nucs[unmapped_n])
		end
		if n in keys(gap_poses)
			for ii in 1:gap_poses[n]
				for iii in 1:3
					aa_counter += 1
					push!(true_nuc, '.')
				end
			end
		end
	end

	ret_seq = translate(join(true_nuc), 1)
	gap_counter = 0
	if do_map
		return ret_seq, explicit_mapping
	else
		return ret_seq
	end
end

end


type Model
    env::Env # Cplex environment
    lp::Ptr{Void} # Cplex problem (lp)
    has_int::Bool # problem has integer variables?
    has_qc::Bool # problem has quadratic constraints?
    callback::Any
    terminator::Vector{Cint}
end

function Model(env::Env, lp::Ptr{Void})
    model = Model(env, lp, false, false, nothing, Cint[0])
    finalizer(model, free_problem)
    set_terminate(model)
    model
end

function Model(env::Env, name::String)
    @assert is_valid(env)
    stat = Array(Cint, 1)
    tmp = @cpx_ccall(createprob, Ptr{Void}, (Ptr{Void}, Ptr{Cint}, Ptr{Cchar}), env.ptr, stat, name)
    if tmp == C_NULL
        throw(CplexError(env, stat))
    end
    return Model(env, tmp)
end

# internal function that wraps finalizer for model and environment together
function _Model(env::Env)
    @assert is_valid(env)
    stat = Array(Cint, 1)
    tmp = @cpx_ccall(createprob, Ptr{Void}, (Ptr{Void}, Ptr{Cint}, Ptr{Cchar}), env.ptr, stat, "CPLEX.jl")
    if tmp == C_NULL
        throw(CplexError(model.env, stat))
    end
    model = Model(env, tmp, false, false, nothing, Cint[0])
    finalizer(model, model::Model -> begin
            tmp = model.env
            free_problem(model)
            close_CPLEX(tmp)
        end)
    set_terminate(model)
    return model
end

function read_model(model::Model, filename::String)
    stat = @cpx_ccall(readcopyprob, Cint, (Ptr{Void}, Ptr{Void}, Ptr{Cchar}, Ptr{Cchar}), model.env.ptr, model.lp, filename, C_NULL)
    if stat != 0
        throw(CplexError(model.env, stat))
    end
end

function write_model(model::Model, filename::String)
    if endswith(filename,".mps")
        filetype = "MPS"
    elseif endswith(filename,".lp")
        filetype = "LP"
    else
        error("Unrecognized file extension: $filename (Only .mps and .lp are supported)")
    end
    stat = @cpx_ccall(writeprob, Cint, (Ptr{Void}, Ptr{Void}, Ptr{Cchar}, Ptr{Cchar}), model.env.ptr, model.lp, filename, filetype)
    if stat != 0
        throw(CplexError(model.env, stat))
    end
end

## TODO: deep copy model, reset model

function get_sense(model::Model)
    sense_int = @cpx_ccall(getobjsen, Cint, (
                           Ptr{Void},
                           Ptr{Void},
                           ),
                           model.env.ptr, model.lp)
    if sense_int == 1
        sense = :Min
    elseif sense_int == -1
        sense = :Max
    else
        error("CPLEX: problem object or environment does not exist")
    end
    return sense
end

function set_sense!(model::Model, sense)
    if sense == :Min
        @cpx_ccall(chgobjsen, Void, (Ptr{Void}, Ptr{Void}, Cint), model.env.ptr, model.lp, 1)
    elseif sense == :Max
        @cpx_ccall(chgobjsen, Void, (Ptr{Void}, Ptr{Void}, Cint), model.env.ptr, model.lp, -1)
    else
        error("Unrecognized objective sense $sense")
    end
end

function get_obj(model::Model)
    nvars = num_var(model)
    obj = Array(Cdouble, nvars)
    stat = @cpx_ccall(getobj, Cint, (
                      Ptr{Void},
                      Ptr{Void},
                      Ptr{Cdouble},
                      Cint,
                      Cint
                      ),
                      model.env.ptr, model.lp, obj, 0, nvars-1)
    if stat != 0
        throw(CplexError(model.env, stat))
    end
    return obj
end

const type_map = Dict(
     0 => :LP,
     1 => :MILP,
     3 => :MILP,
     5 => :QP,
     7 => :MIQP,
     8 => :MIQP,
    10 => :QCP,
    11 => :MIQCP
)

function get_prob_type(model::Model)
  ret = @cpx_ccall(getprobtype, Cint, (
                   Ptr{Void},
                   Ptr{Void}),
                   model.env.ptr, model.lp)
  ret == -1 && error("No problem of environment")
  return type_map[Int(ret)]
end

function set_obj!(model::Model, c::Vector)
    nvars = num_var(model)
    stat = @cpx_ccall(chgobj, Cint, (
                        Ptr{Void},
                        Ptr{Void},
                        Cint,
                        Ptr{Cint},
                        Ptr{Cdouble}
                        ),
                        model.env.ptr, model.lp, nvars, Cint[0:nvars-1;], float(c))
    if stat != 0
        throw(CplexError(model.env, stat))
    end
end

set_warm_start!(model::Model, x::Vector{Float64}) = set_warm_start!(model, Cint[1:length(x);], x)

function set_warm_start!(model::Model, indx::IVec, val::FVec)
    stat = @cpx_ccall(addmipstarts, Cint, (
                      Ptr{Void},
                      Ptr{Void},
                      Cint,
                      Cint,
                      Ptr{Cint},
                      Ptr{Cint},
                      Ptr{Cdouble},
                      Ptr{Cint},
                      Ptr{Ptr{Cchar}}
                      ),
                      model.env.ptr, model.lp, 1, length(indx), Cint[0], indx.-1, val, Cint[CPX_MIPSTART_AUTO], C_NULL)
    if stat != 0
        throw(CplexError(model.env, stat))
    end
end

function free_problem(model::Model)
    tmp = Ptr{Void}[model.lp]
    stat = @cpx_ccall(freeprob, Cint, (Ptr{Void}, Ptr{Void}), model.env.ptr, tmp)
    if stat != 0
        throw(CplexError(model.env, stat))
    end
end

function close_CPLEX(env::Env)
    tmp = Ptr{Void}[env.ptr]
    stat = @cpx_ccall(closeCPLEX, Cint, (Ptr{Void},), tmp)
    if stat != 0
        throw(CplexError(env, stat))
    end
end

function set_terminate(model::Model)
    stat = @cpx_ccall(setterminate, Cint, (Ptr{Void},Ptr{Cint}), model.env.ptr, model.terminator)
    if stat != 0
        throw(CplexError(env, stat))
    end
end

terminate(model::Model) = (model.terminator[1] = 1)


using Test, BenchmarkTools
include("../src/dqmc_framework.jl") # to be replaced by using DQMC or similar

@testset "All Speed Tests" begin

    # Linear algebra
    include("speedtests_linalg.jl")

end

nothing


"""
Basic initialization
"""
function initialize(user_props;grid_dims=(3,3),seed=0)
    # pulses = []
    # for i in 1:args[:q]
    #     push!(pulses,[])
    # end

    # Global model props
    default_props = Dict(
        :ticks => 0,# # time unit
        :pkt_id => 0,
        :amsg_id =>0,
        :ofmsg_id=>0,
        :capacity_factor=>1.2, #default capacity factor of pkts processed per tick (.2 extra is to have always room for management msgs/pkts.)
        :mapping_ctl_ntw => Dict{Int64,Int64}(), # mapping between (Ctl) Agent and SimNE
        :mapping_ntw_sne => Dict{Int64,Int64}(), #mapping btwn the underlying network and the corresponding simNE agent
        :ntw_links_msgs=>Dict{Tuple{Int,Int},Vector{Vector{OFMessage}}}(),
        :ntw_links_delays =>Dict{Tuple{Int,Int},Int}(),
        :ntw_links_capacity=>Dict{Tuple{Int,Int},Int}(),
        :state_trj => Vector{ModelState}(),
        :base_ntw_graph => user_props[:ntw_graph],
        :dropped_nodes => Vector{Tuple{Int,Int}}()
    )
    #For G6:
    #prob_eq_queries_cycle: 0.2
    #clear_cache_graph_freq: 30 or 15
    #query_cycle: 20 or 10
    # requires redundancy of queries + longer time graph cache (> query cycle)
    #For G7/G4:
    #prob_eq_queries_cycle: 0.82
    #clear_cache_graph_freq: 10
    #query_cycle: 10

    Random.seed!(seed)
    props = merge(default_props,user_props)
    #space = GridSpace(grid_dims, moore=true)
    space = GraphSpace(props[:ntw_graph])
    agent_types = Union{SimNE,Agent}
    model = ABM(agent_types, space;
    scheduler = Agents.Schedulers.randomly #random_activation
    , properties = props)
    init_model!(model)
    #create
    create_agents!(model)
    model
end

"""
    Create simulated asset agents
"""
function create_sim_asset_agents!(model)
    a_params = Dict()

    # create SimNE
    for i in 1:nv(model.properties[:ntw_graph])
        #next_fire = rand(0:0.2:model.:Τ)
        id = nextid(model)

        deterioration = id in model.init_sne_params.ids ? model.init_sne_params.deterioration[first(indexin(id,model.init_sne_params.ids))] : model.deterioration

        mnt = @match model.mnt_policy begin
            1 => MaintenanceInfoPreventive(deterioration,model)
            2 => MaintenanceInfoPredictive(deterioration,model)
            _ => MaintenanceInfoCorrective(deterioration,model)
        end

        capacity_factor =  id in model.init_sne_params.ids ? model.init_sne_params.capacity_factor[first(indexin(id,model.init_sne_params.ids))] : model.capacity_factor
        max_q_capacity =  capacity_factor * model.pkt_per_tick
        a_params[:pkt_per_tick] = max_q_capacity
        a = add_agent_pos!(
                SimNE(id,i,a_params,max_q_capacity,mnt),model
            )

        set_prop!(model.ntw_graph, id, :eid, id )
        # create initial mapping btwn network and SimNE
        model.mapping_ntw_sne[i] = i
    end

    set_indexing_prop!(model.ntw_graph, :eid)
end

"""
    Create control agents
"""
function create_control_agents!(model::ABM)

    a_params::Dict{Symbol,Any} = Dict(:pkt_per_tick=>model.pkt_per_tick)
    mnt::MaintenanceInfo = @match model.mnt_policy begin
                1 => MaintenanceInfoPreventive(model)
                2 => MaintenanceInfoPredictive(model)
                _ => MaintenanceInfoCorrective(model)
            end

    if model.ctrl_model == CENTRALISED
        id = nextid(model)
        a = add_agent_pos!(
                Agent(id,1,mnt,Array{Float64}(undef,0,nv(model.properties[:ntw_graph])),a_params),model
            )
        #set_props!(model.ctl_graph, 1, Dict(:eid => 1, :aid => id) )
         ##assign controller to SimNE
        for j in 1:nv(model.properties[:ntw_graph])
            set_control_agent!(j,id,model)
        end
    else
        for i in 1:nv(model.properties[:ctl_graph])
            id = nextid(model)
            a = add_agent_pos!(
                    Agent(id,i,mnt,Array{Int64}(undef,0,1),a_params),model
                )
            set_props!(model.ctl_graph, i, Dict(:eid => i, :aid => id) )
            ##assign controller to SimNE
            #one-to-one mapping
            set_control_agent!(i,id,model)
        end
        set_indexing_prop!(model.ctl_graph, :aid)
    end

end

"""
Simplest create agents
"""
function create_agents!(model)
    create_sim_asset_agents!(model)

    create_control_agents!(model)

    log_info("Total agents created: $(length(allagents(model))), seed: $(model.seed)")

    init_agents!(model)
end

"""
    It advance model one step. Agents are progressed here to ensure that one action is perform in
    all agents before the next action is performed in all agents, instead of having one action
    executed in some agents and then some start with second action.
"""
function model_step!(model)
    init_step_state!(model)
    # log_info(model.ticks,"==========")
    log_info(".")
    for a in allagents(model)
        init_state!(a)
    end
    generate_traffic!(model)
    for e in edges(model.ntw_graph)
        ntw_link_step!((e.src,e.dst),model)
    end
    ctl_links_step!(model)

    # Run controllers first
    for a in allagents(model)
        if typeof(a) == Agent
            # kn = [ a.ntw_graph[i,:eid] for i=1:nv(a.ntw_graph) ]
            do_agent_step!(a,model)
        end
    end
    #Then run SimNEs
    for a in allagents(model)
        if typeof(a) == SimNE
            do_agent_step!(a,model)
        end
    end
    for a in allagents(model)
        pending_pkt_handler(a,model)
        clear_cache!(a,model)
        calculate_metrics_step!(a,model)
    end
    trigger_random_node_drops!(model)
end

"""
    Progress control agents one step.
"""
function agent_step!(a::Agent,model)
    #placeholder
end

"""
    Progress SimNE one step.
"""
function agent_step!(a::SimNE,model)
    #placeholder
end

"""
Simple run model function
"""
function run_model(n,properties; agent_data, model_data)
    seed = properties[:seed]
    model = initialize(properties;seed)

    # agent_shape(a::Agent) = :square
    agent_size(a) = 7
    # agent_size(a::Agent) = 7

    plot_scheduler = model.scheduler#by_type((SimNE,Agent), false)#

    plotkwargs = (
            # ac = agent_color, am = agent_shape, as = agent_size,
            # scheduler = plot_scheduler, aspect_ratio = 1
            # #Graph space properties
            # , method = :circular, node_size = 0.2
            # ,
            size = (1000, 600),
            showaxis = false
    )

    #p = plotabm(model; plotkwargs...)
    df = init_agent_dataframe(model,agent_data)
    df_m = init_model_dataframe(model,model_data)

        if model.animation

        anim = @animate for i in 0:n
                p = plotabm_networks(model; plotkwargs...)
                title!(p, "step $(i)",position=(10,1))
                #annotate!((1,1,Plots.text("step $(i)", 11, :black, :center)))
                step!(model, agent_step!,model_step!)
                collect_agent_data!(df, model, agent_data, i)
                collect_model_data!(df_m, model, model_data, i)
            end
        plot_label = model.run_label * "_anim"
        check_create_dir!(model.plots_dir)
        gif(anim, model.plots_dir * plot_label * ".gif", fps = 5)
    else
        for i in 0:n
            step!(model, agent_step!,model_step!)
            collect_agent_data!(df, model, agent_data, i)
            collect_model_data!(df_m, model, model_data, i)
        end
    end
    df, df_m
end

function init_agents!(model)
    # for a in sort(allagents(model))
    #     init_agent(a,model)
    # end

    #force ordered start

    ids = [a.id for a in allagents(model)]

    for id in sort(ids)
        init_agent!(getindex(model,id),model)
    end
end

"""
Initialise control agents
"""
function init_agent!(a::Agent,model)

    if model.ctrl_model != GraphModel(1)
        #Calculate sub graphs and init msg channels among agents
        nodes = [get_controlled_assets(a.id,model)...]
        sub_g = get_subgraph(model.ntw_graph,nodes,:eid)

        nodes = [a.id]
        ctl_sub_g = get_subgraph(model.ctl_graph,nodes,:aid)
        a.ntw_graph = sub_g
        a.base_ntw_graph = sub_g
        a.ctl_graph = ctl_sub_g
        a.params[:delay_ctl_link] = 2 # 1: no delay
        #Init vector of msgs
        a.msgs_links = init_array_vectors(AGMessage,a.params[:delay_ctl_link],degree(a.ctl_graph,to_local_vertex(a.ctl_graph,a.id,:aid)))
    else
        a.ntw_graph = model.ntw_graph
        a.base_ntw_graph = model.ntw_graph
    end

    ##
    if a.maintenance.policy !=  CorrectiveM
        schedule_event!(a,CTL_Event(4),a.maintenance.predictive_freq,Array{Int64,1}())
        if a.maintenance.policy ==  PredictiveM
            if Symbol("py_integration") in keys(model.properties)
                #conversion to py
                ajm_py =  model.py_integration.np.matrix(adjacency_matrix(a.ntw_graph))
                model.py_integration.opt_init.optimisation_initialisation( ajm_py,
                model.traffic_dist_params
                #[1,0.05]
                , model.mnt_bc_cost, model.mnt_bc_duration, model.mnt_wc_cost, model.mnt_wc_duration)
            end
        end
    end
end

function label_paths(time::Int64,paths::Array{YenState{Float64,Int64},1})
    npaths = Dict()
    if !isempty(paths)
        for path in paths
            if !isempty(path.paths)
                #obtain only the first path
                npaths[(first(first(path.paths)),last(first(path.paths)))] = [(time,first(path.dists),first(path.paths))]
            end
        end
    end
    return npaths
end

function label_path(path)

    return (first(path),last(path),path)
end

"""
Initialise sne agents
"""
function init_agent!(sne::SimNE,model)
    #print("Initialisation of SimNE agent $(a.id)")
    nbs = all_neighbors(model.ntw_graph,get_address(sne.id,model.ntw_graph))

    push_ep_entry!(sne,(0,"h$(sne.id)")) # link to a host of the same id

    for i in 1:size(nbs,1)
        push_ep_entry!(sne,(i,"s$(nbs[i])"))
        #install flows to host of the same id
        # install_flow!(Flow(sne.id,MRule("$i","$(nbs[i])","$(sne.id)"),[0],OFS_Action(1)),sne,model)
    end

    init_condition!(sne,model)
    init_maintenance!(sne,model)
end

"""
    Next agent message's id
"""
function next_amid!(model)
    model.amsg_id += 1
    return model.amsg_id
end

"""
    Next Open Flow message id
"""
function next_ofmid!(model)
    model.ofmsg_id += 1
    return model.ofmsg_id
end

"""
    Initialise state of model for a given step
"""
function init_step_state!(m::ABM)
    m.ticks += 1
    new_state = ModelState(m.ticks)#deepcopy(get_state(m))
    #for e in edges(m.ntw_graph)
    #    new_state.links_load[(e.src,e.dst)] = 0
           # read from links according to delay
   # put msgs in destination nodes for those where delay is over

        #m.ntw_links[k]
    #end
    push!(m.state_trj,new_state)
end

function get_state(m::ABM)
    return last(m.state_trj)
end

function set_active_flows(m::ABM,af::Vector{Tuple{Int,Int,Flow_Type}})
    current = pop!(m.state_trj)
    current.active_flows = af
    push!(m.state_trj,current_logger)
end

function set_state!(m::ABM,new_state::ModelState)
    pop!(m.state_trj)
    push!(m.state_trj,new_state)
end

function init_model!(m::ABM)
    for e in edges(m.ntw_graph)
        #all delays equal initially
        m.ntw_links_delays[(e.src,e.dst)] = 1
        #set capacity for each link, checking if any init param was passed for the links
        if (e.src,e.dst) in m.init_link_params.ids
            idx = first([i for i=1:length(m.init_link_params.ids) if m.init_link_params.ids[i] == (e.src,e.dst) ])
            m.ntw_links_capacity[(e.src,e.dst)] = m.init_link_params.capacities[idx]
        else
            m.ntw_links_capacity[(e.src,e.dst)] = m.link_capacity
        end
    end
end

function init_link_msg!(l::Tuple{Int,Int},m::ABM)
    link_queue = Vector{Vector{OFMessage}}()
    for q=1:m.ntw_links_delays[l]
        push!(link_queue,Vector{OFMessage}())
    end
    m.ntw_links_msgs[l] = link_queue
end

"""
This function progresses one step each link of the sne network.
"""
function ntw_link_step!(l::Tuple{Int,Int},model)
    # Done for each registered link
    if haskey(model.ntw_links_msgs,l)

        msgs = model.ntw_links_msgs[l]
        to_deliver = first(msgs)
        in_pkt_count = 0

        if !isempty(to_deliver)
            if length(msgs) > 1
                msgs = msgs[begin+1:end]
                push!(msgs,Vector{OFMessage}())
            else
                msgs = [ Vector{OFMessage}() ]
            end
            model.ntw_links_msgs[l] = msgs

            for msg in to_deliver
                #Does it need to check address?, I don't think so
                dst = msg.dpid == l[1] ? getindex(model,l[2]) : getindex(model,l[1])

                if length(dst.queue.data) < dst.queue.sz_max - Int(round(dst.queue.sz_max * .2))

                     #TODO make % a parameter. This the gap for OF control msgs.

                    put!(dst.queue,msg)
                    in_pkt_count = get_state(dst).in_pkt + 1
                    set_in_pkt!(dst,in_pkt_count)
                else
                    #if destination queue is full the msg/pkt is dropped
                    drop_packet!(dst)
                end
            end
        end
    end

end

"""
    Initialize an array of dimensions d1 x d2 that contains vectors of type T
"""
function init_array_vectors(T,d1,d2)
    arr = d2 >= 1 ? Array{Vector{T}}(undef,d1,d2) : Array{Vector{T}}(undef,d1)
    [ arr[i,j] = Vector{T}() for i=1:size(arr,1) for j=1:size(arr,2)]
    return arr
end

function ctl_links_step!(model)
    if model.ctrl_model != GraphModel(1)
        ctl_ags = filter(a->typeof(a) == Agent,Set(allagents(model)))
        ctl_link_step!.(ctl_ags)
    end

end

function ctl_link_step!(a::Agent)
    #Merge msgs from all senders to be processed
    a.msgs_in = vcat(a.msgs_links[1,:]...)
    a.msgs_links[1,:] = init_array_vectors(AGMessage,size(a.msgs_links,2),1)
    a.msgs_links = circshift(a.msgs_links,1)
end

function get_state_trj(m::ABM)::Vector{ModelState}
    return m.state_trj
end

"""
Return total of control messages exchanged by agents
"""
function get_ag_msg(model)
    return cumsum(sum.(eachcol([ [ s.in_ag_msg for s in a.state_trj ] for a in allagents(model) if typeof(a) == Agent ]))...)
end

function is_active_flow(f::Tuple{Int,Int},model)
    v = !isempty(filter(af->(af[1],af[2]) == f || (af[2],af[1]) == f  ,get_state(model).active_flows))
    return v
end


import MAT
import Base.parse
using DataFrames
import Stimulus
#import Spiketrains

@compat function loadWaveformsFile(fname::AbstractString;time_conversion::Real=0.001)
	fid = open(fname,"r")
	waveforms = Array(Int16,0)
	timestamps = Array(Uint64,0)
	try
		headersize = Int(read(fid,Uint32))
		num_spikes = Int(read(fid,Uint32))
		num_chs = Int(read(fid,Uint8))
		sampling_rate = Int(read(fid,Uint32))
		timepts = Int(read(fid,Uint32))
		seek(fid,headersize)
		append!(waveforms,read(fid,Int16,num_chs*timepts*num_spikes))
		waveforms = reshape(waveforms,(num_chs,timepts,num_spikes))
		seek(fid,headersize + num_spikes*num_chs*timepts*sizeof(Int16))
		append!(timestamps,read(fid,Uint64,num_spikes))
	finally
		close(fid)
	end
	return float(waveforms),float(timestamps)*time_conversion
end

function assign_area{T,V}(sptrains::Dict{T,V},config;return_mapping=false)
	newsptrains = Dict{T,V}()
	mapping = Dict()
	for (k,v) in sptrains
            kname = assign_area(k,config)
            newsptrains[kname] = v
            mapping[k] = kname
	end
	if return_mapping
		return newsptrains,mapping
	else
		return newsptrains
	end
end

@compat function assign_area(cell::AbstractString, config::Dict)
	rr = r"g([0-9]*)c([0-9])*[s]*"
        m = match(rr,cell)
        group = int(m.captures[1])
        cellnr = int(m.captures[2])
        ad = filter((kk,vv)->group in vv, config)
        aa = collect(keys(ad))[1]
        cc = collect(values(ad))[1]
        cidx = find(cc.==group)[1]
        kname = "$(aa)_g$(cidx)c$(cellnr)"
end

@compat function getspiketrains(;session::AbstractString="",groups::Array{Int64,1}=Array(Int64,0),checkArtifact::Bool=true,verbose::Integer=0,check_trials::Bool=true,config::Dict=Dict())
	if isempty(session)
		#attempt to get session from the pl2 file
		for d in (".","..")
			files = split(chomp(readall(`find $d -name "*.pl2" -maxdepth 1`)))
			if !isempty(files)
				session = splitext(splitdir(files[1])[2])[1]
				break
			end
		end
	end
	verbose > 0 && println("Analyzing session $session ...")
	if isempty(groups)
		#get the groups from the waveforms files in the current directory
		files = split(chomp(readall(`find . -name "*waveforms.bin"`)))
	else
		files = Array(ASCIIString,0)
		for g in groups
			pp = @sprintf "%sg%04dwaveforms.bin" session g
			if isfile(pp)
				push!(files,pp)
			end
		end
	end
    if isempty(config)
        for ss in (".", "..","../..")
            if isfile("$(ss)/channel_config.csv")
                config = channel_config("$(ss)/channel_config.csv")
                break
            end
        end
    end
	SS = pmap(f->begin
					m = match(r"([A-Za-z0-9_]*)g([0-9]*)",f)
					session = m.captures[1]
					group = int(m.captures[2])
					sptrains = getspiketrains(session,group;checkArtifact=checkArtifact,verbose=verbose)
					sptrains
			end,files)
	sptrains = merge(SS...)
	#check if we have a trial structure
    if check_trials
        for ss in (".","..")
            fname = "$(ss)/event_data.mat"
            if isfile(fname)
                verbose > 0 && println("Found trial info. Checking for relevant spikes....")
                trials = Stimulus.loadTrialInfo(fname)
                rtrials = Stimulus.getTrialType(trials,:reward)
                target_time = Stimulus.gettime(rtrials,:target)
                response_time = Stimulus.gettime(rtrials,:response)
                trial_dur = 1000*maximum(response_time .- target_time) #convert from seconds to ms
                cells = Spiketrains.sortCells(sptrains)
                aligned_spikes = Spiketrains.getTrialRaster(sptrains,rtrials,:target,-200.0,trial_dur)
                #remove cells that never spike during the trial
                for cc in sort(setdiff(1:length(cells),unique(aligned_spikes.cellidx)))
                    pop!(sptrains,cells[cc])
                    verbose > 0 && println("\tRemoved cell $(cells[cc])")
                end
                break
            end
        end
    end
	if !isempty(config)
		newsptrains = assign_area(sptrains,config)
	else
		newsptrains = sptrains
	end
	return newsptrains
end

@compat function getspiketrains(session::AbstractString, group::Integer;checkArtifact::Bool=true,verbose::Integer=0)
	verbose > 0 && println("Processing group $group ...")
	wfname = @sprintf "%sg%04dwaveforms.bin" session group
	overlap_fname = @sprintf "%sg%04dwaveforms.overlap" session group
	cut_fname = @sprintf "%sg%04dwaveforms.cut" session group

	waveforms,timestamps = loadWaveformsFile(wfname)
	D = MAT.matread(overlap_fname)
	cids = map(int,readlines(open(cut_fname,"r")))
	clusters = unique(cids)
	clusters = clusters[clusters.>0]
	sptrains = Dict{ASCIIString,Array{Float64,1}}()
	c2 = 1
	for c in clusters
		spidx = (cids.==c) | ((cids.==-1)&(D["overlaps"][:,c].==1))
		cidx = cids.==c
		if checkArtifact
			if size(cids,1) == size(waveforms,3)

				mm = mean(waveforms[:,:,bool(cidx[:])],3)
				ii = indmax(-mm) #get the minimum point
				goodcluster = ii in div(size(mm,2),3):div(2*size(mm,2),3) #minimum should occur in the middle 3rd of the spike
				goodcluster = goodcluster && -minimum(mm) >= maximum(mm) #minimum should exceed maximum in absolute value
				if goodcluster
					cell = "g$(group)c$c2"
					sptrains[cell] = timestamps[spidx[:]]
					c2 += 1
				else
					verbose > 0 && println("\tCluster $c classified as artifact and removed")
				end
			else
				println("\tsize(cids,1) != size(waveforms,3); Skipping cluster $c...")
			end
		else
			cell = "g$(group)c$c"
			sptrains[cell] = timestamps[spidx[:]]
		end
	end
	return sptrains
end

@compat function parse{T<:Integer}(::Type{Range{T}},a::AbstractString)
	m = match(r"([0-9]*)[\-\:]([0-9]*)",a)
     mc = m.captures
    isempty(mc) && return nothing
	ll = map(int,mc)
    return range(ll[1],ll[2]-ll[1]+1)
end

@compat function channel_config(fname::AbstractString)
	dd = DataFrames.readtable(fname)
	#convert to ranges
	q = map(x->parse(Range{Int64},x),dd[:channels])
	dd[:channels] = q
	return Dict(dd[:area],dd[:channels])
end


# sample implementations of the Abstract Types
# used for testing, not to be used for real code

mutable struct ConcreteDGMesh{Tmsh} <: AbstractDGMesh{Tmsh}
  # counts
  numVert::Int
  numEl::Int
  numNodes::Int
  numDof::Int
  numDofPerNode::Int
  numNodesPerElement::Int
  order::Int
  numNodesPerFace::Int

  # parallel counts
  npeers::Int
  numGlobalEl::Int
  numSharedEl::Int
  peer_face_counts::Array{Int, 1}
  local_element_counts::Array{Int, 1}
  remote_element_counts::Array{Int, 1}

  # MPI info
  comm::Int  # this should really be an MPI.Comm
  myrank::Int
  commsize::Int
  peer_parts::Array{Int, 1}

  # Discretization type
  isDG::Bool
  isInterpolated::Bool

  # Mesh data
  coords::Array{Tmsh, 3}
  dxidx::Array{Tmsh, 4}
  jac::Array{Tmsh, 2}

  # interpolated data
  coords_bndry::Array{Tmsh, 3}
  dxidx_bndry::Array{Tmsh, 4}
  jac_bndry::Array{Tmsh, 2}
  dxidx_face::Array{Tmsh, 4}
  jac_face::Array{Tmsh, 2}

  # parallel data
  coords_sharedface::Array{Array{Tmsh, 3}, 1}
  dxidx_sharedface::Array{Array{Tmsh, 4}, 1}
  jac_sharedface::Array{Array{Tmsh, 2}, 1}

  # boundary condition data
  numBC::Int
  numBoundaryEdges::Int
  bndryfaces::Array{Boundary, 1}
  bndry_offsets::Array{Int, 1}
  bndry_funcs::Array{BCType, 1}

  # interior edge data
  numInterfaces::Int
  interfaces::Array{Interface, 1}

  # dof number data
  dofs::Array{Int, 3}
  dof_offset::Int
  sparsity_bnds::Array{Int, 2}
  sparsity_nodebnds::Array{Int, 2}

  # mesh coloring data
  numColors::Int
  maxColors::Int
  color_masks::Array{BitArray{1}, 1}
  shared_element_colormasks::Array{Array{BitArray{1}, 1}, 1}
  pertNeighborEls::Array{Int, 2}

  # parallel bookkeeping
  bndries_local::Array{Array{Boundary, 1}, 1}
  bndries_remote::Array{Array{Boundary, 1}, 1}
  shared_interfaces::Array{Array{Interface, 1}, 1}
  shared_element_offsets::Array{Int, 1}
  local_element_lists::Array{Array{Int, 1}, 1}

  function ConcreteDGMesh{Tmsh}(opts) where {Tmsh}

    mesh = new{Tmsh}()  # incomplete initialization

    # make up values
    mesh.numVert = 5
    mesh.numEl = 7
    mesh.numNodesPerElement = 4
    mesh.numNodes = mesh.numEl*mesh.numNodesPerElement
    mesh.numDofPerNode = 5
    mesh.numDof = mesh.numNodes*mesh.numDofPerNode
    mesh.order = 2
    mesh.numNodesPerFace = 3

    mesh.npeers = 2
    mesh.numGlobalEl = 42
    mesh.numSharedEl = 4
    mesh.peer_face_counts = [2, 2]
    mesh.local_element_counts = [2, 2]
    mesh.remote_element_counts = [1, 2]

    mesh.comm = 666
    mesh.myrank = 0
    mesh.commsize = 4
    mesh.peer_parts = [1, 3]

    mesh.isDG = true
    mesh.isInterpolated = true

    mesh.coords = rand(2, mesh.numNodesPerElement, mesh.numEl)
    mesh.dxidx = rand(2, 2, mesh.numNodesPerElement, mesh.numEl)
    mesh.jac = rand(mesh.numNodesPerElement, mesh.numEl)

    mesh.numBoundaryEdges = 2
    mesh.coords_bndry = rand(2, mesh.numNodesPerFace, mesh.numBoundaryEdges)
    mesh.dxidx_bndry = rand(2, 2, mesh.numNodesPerFace, mesh.numBoundaryEdges)
    mesh.jac_bndry = rand(mesh.numNodesPerFace, mesh.numBoundaryEdges)

    mesh.numInterfaces = 2
    mesh.dxidx_face = rand(2, 2, mesh.numNodesPerFace, mesh.numInterfaces)
    mesh.jac_face = rand(mesh.numNodesPerFace, mesh.numInterfaces)

    mesh.numBC = 2
    mesh.bndryfaces = [Boundary(2, 1), Boundary(3, 2)]
    mesh.bndry_offsets = [1, 2, 3]
    mesh.bndry_funcs = Array{BCType}(mesh.numBC)

    mesh.interfaces = [Interface(2, 3, 2, 3, 0), Interface(4, 7, 1, 3, 2)]

    mesh.dofs = Array{Int}(mesh.numDofPerNode, mesh.numNodesPerElement, mesh.numEl)
    for i=1:mesh.numDof
      mesh.dofs[i] = i
    end

    mesh.dof_offset = 0
    mesh.sparsity_bnds = Array{Int}(2, mesh.numDof)
    # make the matrix dense
    for i=1:mesh.numDof
      mesh.sparsity_bnds[1, i] = 1
      mesh.sparsity_bnds[2, i] = mesh.numDof
    end

    mesh.sparsity_nodebnds = Array{Int}(2, mesh.numNodes)
    for i=1:mesh.numNodes
      mesh.sparsity_nodebnds[1, i] = 1
      mesh.sparsity_nodebnds[2, i] = mesh.numNodes
    end

    mesh.numColors = 4  # like that's really possible...
    mesh.maxColors = 5
    mesh.color_masks = Array{BitArray{1}}(mesh.numColors)
    mesh.shared_element_colormasks = Array{Array{BitArray{1}, 1}}(mesh.npeers)
    mesh.pertNeighborEls = Array{Int}(4, mesh.numEl)

    mesh.bndries_local = Array{Array{Boundary, 1}}(mesh.npeers)
    mesh.bndries_remote = Array{Array{Boundary, 1}}(mesh.npeers)
    mesh.shared_interfaces = Array{Array{Interface, 1}}(mesh.npeers)
    mesh.shared_element_offsets = Array{Int}(2)
    mesh.local_element_lists = Array{Array{Int, 1}}(mesh.npeers)

    return mesh
  end

end

mutable struct ConcreteParamType{Tdim}
  t::Float64
  order::Int
  time  # this really should have a concrete type

  function ConcreteParamType{Tdim}(mesh, opts) where Tdim

    t = 0.0
    order = mesh.order
    time = 0

    return new(t, order, time)
  end
end

mutable struct ConcreteSolutionData{Tsol, Tres} <: AbstractSolutionData{Tsol, Tres}
  q::Array{Tsol, 3}
  q_vec::Array{Tsol, 1}
  # AbstractSharedFaceData should be a concrete type
  shared_data::Array{AbstractSharedFaceData{Tsol}, 1}
  res::Array{Tres, 3}
  res_vec::Array{Tres, 1}
  M::Array{Float64, 1}
  Minv::Array{Float64, 1}
  disassembleSolution::Function
  assembleSolution::Function
  multiplyA0inv::Function
  majorIterationCallback::Function
  params::ConcreteParamType{2}

  function ConcreteSolutionData{Tsol, Tres}(mesh, sbp, opts)  where {Tsol, Tres}

    numDofPerNode = mesh.numDofPerNode
    numNodesPerElement = mesh.numNodesPerElement
    numEl = mesh.numEl
    numNodesPerFace = mesh.numNodesPerFace
    numDof = mesh.numDof

    q = Array{Tsol}(numDofPerNode, numNodesPerElement, numEl)
    q_vec = Array{Tsol}(numDof)
    # the elements of this array should be populated as well
    shared_data = Array{AbstractSharedFaceData{Tsol}}(mesh.npeers)
    res = Array{Tres}(mesh.numDofPerNode, mesh.numNodesPerElement, mesh.numEl)
    res_vec = Array{Tres}(mesh.numDof)
    M = ones(numDof)
    Minv = ones(numDof)

    # create empty functions
    disassembleSolution = (mesh, sbp, eqn, opts) -> begin end
    assembleSolution = (mesh, sbp, eqn, opts) -> begin end
    multiplyA0inv = (mesh, sbp, eqn, opts, res_arr) -> begin end
    majorIterationCallback = (itr, mesh, sp, eqn, opts) -> begin end

    params = ConcreteParamType{2}(mesh, opts)

    return new(q, q_vec, shared_data, res, res_vec, M, Minv,
               disassembleSolution, assembleSolution, multiplyA0inv,
               majorIterationCallback, params)
  end
end

mutable struct ConcreteSBP
  Q::Array{Float64, 2}

  function ConcreteSBP()
    Q = rand(3, 3)
    return new(Q)
  end
end


function psracos_opt!(objective::Objective, parameter::Parameter)
    psracos = PSRacos(parameter.evaluation_server_num)
    rc = psracos.rc
    rc.objective = objective
    rc.parameter = parameter
    ub = isnull(parameter.uncertain_bits)? 1 : parameter.uncertain_bits

    # require calculator server
    ip = parameter.control_server_ip
    port = parameter.control_server_port
    cs_send = connect(ip, port)
    println(cs_send, "client: require servers#")
    readline(cs_send)
    println(cs_send, string(parameter.evaluation_server_num, "#"))
    msg = readline(cs_send)

    servers_msg = readline(cs_send)
    if servers_msg == "No available evaluation server"
        zoolog("Error: no available evaluation server")
        return Solution()
    end
    servers = split(servers_msg, " ")
    println("get $(length(servers)) servers")
    # close the socket
    close(cs_send)
    parameter.ip_port = RemoteChannel(()->Channel(length(servers)))
    for server in servers
        put!(parameter.ip_port, server)
    end
    psracos_init_attribute!(psracos, parameter)
    psracos_init_sample_set!(psracos, ub)
    println("Initialization succeeds")
    history = []
    # addprocs(1)
    i = parameter.train_size
    output_file = parameter.output_file
    f = Nullable()
    if !isnull(output_file)
        f = open(output_file, "w")
    end
    time_log1 = now()
    while i < parameter.budget
        # parallel estimation
        iteration_num = min(parameter.evaluation_server_num, parameter.budget-i)
        for j = 1:iteration_num
            ip_port = take!(parameter.ip_port)
            sol = take!(psracos.sample_set)
            @spawn begin
                try
                    br = compute_fx(sol, ip_port, parameter)
                    put!(parameter.ip_port, ip_port)
                    put!(psracos.result_set, sol)
                catch e
                    println("Exception")
                    println(e)
                    cs_exception = connect(ip, port)
                    println(cs_exception, "client: restart#")
                    readline(cs_exception)
                    println(cs_exception, string(servers_msg, "#"))
                    close(cs_exception)
                    return Solution()
                end
            end
        end
        # sequential update
        sol = Nullable()
        for j=1:iteration_num
            sol = take!(psracos.result_set)
            push!(history, sol.value)
            bad_ele = sracos_replace!(rc.positive_data, sol, "pos")
            sracos_replace!(rc.negative_data, bad_ele, "neg", strategy=parameter.replace_strategy)
            rc.best_solution = rc.positive_data[1]

    		 if rand(rng, Float64) < rc.parameter.probability
    			 classifier = RacosClassification(rc.objective.dim, rc.positive_data,
    			 rc.negative_data, ub=ub)
    			 mixed_classification(classifier)
    			 solution, distinct_flag = distinct_sample_classifier(rc, classifier, data_num=rc.parameter.train_size)
    		 else
    			 solution, distinct_flag = distinct_sample(rc, rc.objective.dim)
    		 end
    		 if distinct_flag == false
    			 zoolog("ERROR: dimension limited")
    			 break
    		 end
    		 if isnull(solution)
    			 zoolog("ERROR: solution null")
    			 break
    		 end
    		 put!(psracos.sample_set, solution)
    		 i += 1
        end
        time_log2 = now()
        time_pass = Dates.value(time_log2 - time_log1) / 1000
        zoolog("Budget $(i): time=$(floor(time_pass))s, value=$(sol.value), best_solution_value=$(rc.best_solution.value)")
        str = "Budget $(i): time=$(floor(time_pass))s, value=$(sol.value), best_solution_value=$(rc.best_solution.value)\nbest_x=$(rc.best_solution.x)\n"
        if !isnull(f)
            write(f, str)
            flush(f)
        end
        if !isnull(parameter.time_limit) && time_pass > parameter.time_limit
            zoolog("Exceed time limit: $(parameter.time_limit)")
            break
        end
    end
    # finish task
    result = rc.best_solution
    objective.history = history
    parameter.positive_data = rc.positive_data
    parameter.negative_data = rc.negative_data
    cs_receive = connect(ip, port)
    println(cs_receive, "client: return servers#")
    readline(cs_receive)
    println(cs_receive, string(servers_msg, "#"))
    close(cs_receive)
    return result
end

# function psracos_sample(rc, ub)
#     if rand(rng, Float64) < rc.parameter.probability
#         classifier = RacosClassification(rc.objective.dim, rc.positive_data,
#         rc.negative_data, ub=ub)
#         zoolog("before classification")
#         mixed_classification(classifier)
#         zoolog("after classification")
#         solution, distinct_flag = distinct_sample_classifier(rc, classifier, data_num=rc.parameter.train_size)
#     else
#         solution, distinct_flag = distinct_sample(rc, rc.objective.dim)
#     end
#     #painc stop
#     if distinct_flag == false
#         zoolog("ERROR: dimension limited")
#     end
#     if isnull(solution)
#         zoolog("ERROR: solution null")
#     end
#     return solution
# end

function psracos_init_attribute!(psracos::PSRacos, parameter::Parameter)
    f = open("init.txt", "w")
    rc = psracos.rc
    # otherwise generate random solutions
    iteration_num = rc.parameter.train_size
    data_temp = rc.parameter.init_sample
    init_num = 0
    remote_data = RemoteChannel(()->Channel(iteration_num))
    remote_result = RemoteChannel(()->Channel(iteration_num))
    if !isnull(data_temp)
        init_num = length(data_temp) < iteration_num? length(data_temp):iteration_num
        for i = 1:init_num
            # str = "initial sample: $(i), value=$(data_temp[i].value)\nx=$(sol.x)\n"
            put!(remote_data, data_temp[i])
        end
    end
    i = 1
    while i <= iteration_num - init_num
        # distinct_flag: True means sample is distinct(can be use),
        # False means sample is distinct, you should sample again.
        sol, distinct_flag = distinct_sample_from_set(rc, rc.objective.dim, rc.data,
            data_num=iteration_num)
        # panic stop
        if isnull(sol)
            break
        end
        if distinct_flag
            put!(remote_data, sol)
            i += 1
        end
    end
    fn = RemoteChannel(()->Channel(1))
    for i = 1:iteration_num
        d = take!(remote_data)
        if d.value != 0
            put!(remote_result, d)
            if i == iteration_num
                put!(fn, 1)
            end
            continue
        end
        ip_port = take!(parameter.ip_port)
        @spawn begin
            compute_fx(d, ip_port, parameter)
            put!(parameter.ip_port, ip_port)
            put!(remote_result, d)
            if i == iteration_num
                put!(fn, 1)
            end
        end

    end
    result = take!(fn)
    f = open("init.txt", "w")
    for i = 1:iteration_num
        d = take!(remote_result)
        push!(rc.data, d)
        str_print = "init sample: $(i), value=$(d.value)"
        str = "init sample: $(i), value=$(d.value)\n x=$(d.x)\n"
        zoolog(str_print)
        write(f, str)
        flush(f)
    end
    # zoolog("after taking result")
    close(f)
    selection!(rc)
    return
end