# include("Move.jl")
# include("Wavefunction.jl")
# include("SwapWavefunction.jl")

abstract Walk

is_nontrivial_transition(walk::Walk) = walk.transition_is_nontrivial

function get_probability_ratio(walk::Walk)
    if !walk.transition_is_nontrivial
        warn("Asking for the probability ratio ( = 1) of a trivial transition.")
        return 1.
    else
        return walk.probability_ratio
    end
end

# Lots of code copying below since Julia does't really have abstract base classes.  There has to be a better way..
# I could have one class with finish_move! and cancel_move! functions passed to the constructor and held as members,
# but that seems less safe and less explicit.

# standard walk for ordinary wavefunctions / measurements
type StandardWalk <: Walk
    wf::Wavefunction
    transition_function::Function
    probability_ratio::Float64
    transition_is_nontrivial::Bool
    # transition_in_progress::Bool
    StandardWalk(wf::Wavefunction, transition_function::Function) = new(wf, transition_function)
end

function propose_and_assess_random_transition!(walk::StandardWalk)
    move = walk.transition_function(walk.wf)

    if !is_nontrivial(move)
        walk.transition_is_nontrivial = false
    else
        probrat = assess_move!(walk.wf, move)
        walk.transition_is_nontrivial = true
        walk.probability_ratio = probrat
    end
end

accept_transition!(walk::StandardWalk) = finish_move!(walk.wf)
reject_transition!(walk::StandardWalk) = cancel_move!(walk.wf)

check_for_numerical_error!(walk::StandardWalk) = check_for_numerical_error!(walk.wf)

# walk for Renyi mod measurement
type RenyiModWalk <: Walk
    wf::Union{SwapWavefunction, SwapWavefunctionNS}
    transition_function::Function
    probability_ratio::Float64
    transition_is_nontrivial::Bool
    # transition_in_progress::Bool
    RenyiModWalk(wf::Union{SwapWavefunction, SwapWavefunctionNS}, transition_function::Function) = new(wf, transition_function)
end

function propose_and_assess_random_transition!(walk::RenyiModWalk)
    move = walk.transition_function(walk.wf)
    # move here is either a 2-tuple of moves for the different copies (SwapWavefunction) or a Tuple{ParticlesMove{LatticeParticleMove},Int64} (SwapWavefunctionNS),
    # where the Int here is the copy_to_move
    @assert length(move) == 2

    if isa(move[2], Int) # for SwapWavefunctionNS
        trivial_condition = !is_nontrivial(move[1])
    else # for SwapWavefunction
        trivial_condition = all(map(!, map(is_nontrivial, move)))
    end

    if trivial_condition
        walk.transition_is_nontrivial = false
    else
        probrat = assess_move_for_mod!(walk.wf, move...)
        walk.transition_is_nontrivial = true
        walk.probability_ratio = probrat
    end
end

accept_transition!(walk::RenyiModWalk) = finish_move_for_mod!(walk.wf)
reject_transition!(walk::RenyiModWalk) = cancel_move_for_mod!(walk.wf)

check_for_numerical_error!(walk::RenyiModWalk) = check_for_numerical_error_for_mod!(walk.wf)

# walk for Renyi sign measurement
type RenyiSignWalk <: Walk
    wf::Union{SwapWavefunction, SwapWavefunctionNS}
    transition_function::Function
    probability_ratio::Float64
    transition_is_nontrivial::Bool
    # transition_in_progress::Bool
    RenyiSignWalk(wf::Union{SwapWavefunction, SwapWavefunctionNS}, transition_function::Function) = new(wf, transition_function)
end

function propose_and_assess_random_transition!(walk::RenyiSignWalk)
    move = walk.transition_function(walk.wf)
    @assert length(move) == 2 # a 2-tuple of moves for the different copies

    if isa(move[2], Int) # for SwapWavefunctionNS
        trivial_condition = !is_nontrivial(move[1])
    else # for SwapWavefunction
        trivial_condition = all(map(!, map(is_nontrivial, move)))
    end

    if trivial_condition
        walk.transition_is_nontrivial = false
    else
        probrat = assess_move_for_sign!(walk.wf, move...)
        walk.transition_is_nontrivial = true
        walk.probability_ratio = probrat
    end
end

accept_transition!(walk::RenyiSignWalk) = finish_move_for_sign!(walk.wf)
reject_transition!(walk::RenyiSignWalk) = cancel_move_for_sign!(walk.wf)

check_for_numerical_error!(walk::RenyiSignWalk) = check_for_numerical_error_for_sign!(walk.wf)

# type RenyiWalk <: Walk
#     wf::SwapWavefunction
#     transition_function::Function
#     probability_ratio::Float64
#     transition_is_nontrivial::Bool
#     # transition_in_progress::Bool
#     RenyiWalk(wf::SwapWavefunction, transition_function::Function) = new(wf, transition_function)
# end

# type RenyiModWalk
#     walk::RenyiWalk
#     RenyiModWalk(wf::SwapWavefunction, transition_function::Function) = new(RenyiWalk(wf, transition_function))
# end

# How to then make, e.g., RenyiModWalk.wf invoke RenyiModWalk.walk.wf?


module CountMin

export CountMinSketch,
       push!,
       pop!,
       add!,
       haskey,
       getindex,
       collisionrate,
       readcms,
       writecms,
       eltype,
       size,
       show

import Base: read!,
             write,
             push!,
             pop!,
             getindex,
             eltype,
             size,
             show

import HDF5: h5open

"""
A Count-min Sketch

### Implementation Notes

Sketches are stored such that each table is contiguous, which in Julia means
that tables need to be columns. To ease transition between those using this
datastructure in Python or C, we preseve the (number of tables, table size)
way of referring to the sketch's size, and just reverse this internally.
"""
type CountMinSketch{T<:Unsigned}
    # tables and tablesize are purely for convenience, and are just
    # shape(sketch)

    # Num tables (to save getting this from shape(sketch)
    tables::UInt

    # Number of bins per table
    tablesize::UInt

    # We store the sketch as a Matrix
    sketch::Matrix{T}

    # Type min and max precomuted to avoid allocs
    tpmax::Int
    tpmin::Int

    """
    The constructor for `CountMinSketch`es

    Arguments
    ---------

    * `tables::Int`: The number of tables to create
    * `tablesize::Int`: The size of each table
    """
    function CountMinSketch(tables::Int, tablesize::Int)
        if !(1 <= tables <= 20)
            error("Must have between 1 and 20 tables")
        end
        tablesize > 1 || error("Table size must be greater than 1")
        sketch = zeros(T, tablesize, tables)
        return new(tables, tablesize, sketch, typemax(T), typemin(T))
    end
end

"""
A trivial constructor for `CountMinSketch`es.

Creates a sketch with 4 tables of 1000 UInt8-typed bins.

WARNING: This trivial constructor instantiates a tiny `CountMinSketch`, which
         will be useless for anything but a simple test. Storing more than 10
         items  is not recommended.
"""
CountMinSketch() = CountMinSketch{UInt8}(4, 1000)

function show{T}(io::IO, cms::CountMinSketch{T})
    nt = cms.tables
    ts = cms.tablesize
    print("CountMinSketch with $nt tables of $ts `$T`s")
end

"""
Adds `count` `item`s to a CountMinSketch

This function will never allow the counts for an item to overflow or underflow.

Arguments
---------
* `cms`: A count-min sketch
* `item`: any hashable item
* `count::Integer`: a (potentially negative) number of `item`s to add.
"""
function add!(cms::CountMinSketch, item, count::Int)
    i::UInt = 0
    while (i+=1) <= cms.tables
        offset = (hash(item, hash(i)) % cms.tablesize) + 1
        try
            @inbounds cms.sketch[offset, i] += count
        catch
            # clamp overflow to typemin/typemax(eltype(sketch))
            @inbounds cms.sketch[offset, i] = count > 0 ? cms.tpmax : cms.tpmin
        end
    end
end

"""
Determine count of an item in a CountMinSketch.

NB: This function returns an estimate of the true count. This estimate is never
    lower than the true count, but has a very low probablity of being higher.

Arguments
---------
* `cms`: A count-min sketch
* `item`: any hashable item
"""
function getindex(cms::CountMinSketch, item)
    minval = cms.tpmax
    i::UInt = 0
    while (i+=1) <= cms.tables
        offset = (hash(item, hash(i)) % cms.tablesize) + 1
        @inbounds val = cms.sketch[offset, i]
        if val < minval
            minval = val
        end
    end
    return minval
end

"""
Add one `item` from a CountMinSketch

This function will never allow the counts for an item to overflow.

Arguments
---------
* `cms`: A count-min sketch
* `item`: any hashable item
"""
function push!(cms::CountMinSketch, item)
    add!(cms, item, 1)
end

"""
Remove one `item` from a CountMinSketch

This function will never allow the counts for an item to underflow.

Arguments
---------
* `cms`: A count-min sketch
* `item`: any hashable item
"""
function pop!(cms::CountMinSketch, item)
    add!(cms, item, -1)
end

"""
Test for presence of an item in a CountMinSketch.

NB: This function never returns a false negative, but has a very low probablity
    of returning a false positive due to collision.
"""
function haskey(cms::CountMinSketch, item)
    return cms[item] > 0
end

"""
Element type of a CountMinSketch
"""
function eltype(cms::CountMinSketch)
    return eltype(cms.sketch)
end

"""
Shape of a CountMinSketch, i.e. (number of tables, tablesize).
"""
function size(cms::CountMinSketch)
    # It's the reverse as we store the sketch in table-major order (which in
    # Julia means transposed (tables are columns), as Julia is Column major).
    return reverse(size(cms.sketch))
end

function collisionrate(cms::CountMinSketch)
    # Expectected collision rate, modified from Khmer
    occupancy = sum(cms.sketch .> 0, 1)
    # The total FPR is the product of all rates, as we assume each is truly
    # independent
    rate = prod(float(occupancy) / float(cms.tablesize))
    return rate
end

function writecms(filename::AbstractString, cms::CountMinSketch)
    h5open(filename, "w") do h5f
        # Chunks of up to 1MB by one table
        chunksize = (min(2^20, cms.tablesize), 1)
        h5f["sketch", "blosc", 9, "chunk", chunksize] = cms.sketch
    end
end

function readcms(filename::AbstractString)
    cms = 0
    h5open(filename, "r") do h5f
        sketch = h5f["sketch"]
        # Sketch is stored column major
        ts, nt = size(sketch)
        cms = CountMinSketch{eltype(sketch)}(nt, ts)
        cms.sketch = read(sketch)
        cms.tables = nt
        cms.tablesize = ts
    end
    return cms
end

end # module CountMin


# GurobiMathProgInterface
# Standardized MILP interface

export GurobiSolver

import MathProgBase
const MPB = MathProgBase

mutable struct GurobiMathProgModel <: MPB.AbstractLinearQuadraticModel
    inner::Model
    last_op_type::Symbol  # To support arbitrary order of addVar/addCon
                          # Two possibilities :Var :Con
    changed_constr_bounds::Bool # have we updated the bounds below?
    obj::Vector{Float64} # objective vector
    lb::Vector{Float64} # persistent bounds on constraints to maintain
    ub::Vector{Float64} # abstraction for lb ≤ Ax ≤ ub
    lazycb
    cutcb
    heuristiccb
    infocb
    options
end

function GurobiMathProgModel(env=nothing;options...)
   finalize_env = (env == nothing)
   if env == nothing
       env = Env()
   end
   m = GurobiMathProgModel(Model(env,""; finalize_env=finalize_env), :Con, false, Float64[], Float64[], Float64[], nothing, nothing, nothing, nothing, options)
   setparams!(m)
   return m
end

function Base.copy(m::GurobiMathProgModel)

    m.lazycb == nothing || @warn("Callbacks can't be copied, lazy callback ignored")
    m.cutcb == nothing || @warn("Callbacks can't be copied, cut callback ignored")
    m.heuristiccb == nothing || @warn("Callbacks can't be copied, heuristic callback ignored")
    m.infocb == nothing || @warn("Callbacks can't be copied, info callback ignored")

    return GurobiMathProgModel(copy(m.inner),
                               m.last_op_type,
                               m.changed_constr_bounds,
                               copy(m.obj),
                               copy(m.lb),
                               copy(m.ub),
                               nothing,
                               nothing,
                               nothing,
                               nothing,
                               deepcopy(m.options))
end

function setparams!(m::GurobiMathProgModel)
    # Helper to set the parameters on the model's copy of env rather than
    # modifying the global env (ref: http://www.gurobi.com/support/faqs#P)

    # Set `InfUnbdInfo` to 1 by default so infeasibility rays available
    setparam!(m.inner, "InfUnbdInfo", 1)
    for (name,value) in m.options
        setparam!(m.inner, string(name), value)
    end
end

mutable struct GurobiSolver <: MPB.AbstractMathProgSolver
    env
    options
end
GurobiSolver(env=nothing; kwargs...) = GurobiSolver(env, kwargs)
MPB.LinearQuadraticModel(s::GurobiSolver) = GurobiMathProgModel(s.env; s.options...)
MPB.ConicModel(s::GurobiSolver) = MPB.LPQPtoConicBridge(MPB.LinearQuadraticModel(s))

MPB.supportedcones(::GurobiSolver) = [:Free,:Zero,:NonNeg,:NonPos,:SOC,:SOCRotated]

function MPB.setparameters!(s::GurobiSolver; mpboptions...)
    opts = collect(Any,s.options)
    for (optname, optval) in mpboptions
        if optname == :TimeLimit
            push!(opts, (:TimeLimit,optval))
        elseif optname == :Silent
            if optval == true
                push!(opts, (:OutputFlag,0))
            end
        else
            error("Unrecognized parameter $optname")
        end
    end
    s.options = opts
    return
end

function MPB.setparameters!(m::GurobiMathProgModel; mpboptions...)
    for (optname, optval) in mpboptions
        if optname == :TimeLimit
            setparam!(m.inner, "TimeLimit", optval)
        elseif optname == :Silent
            if optval == true
                setparam!(m.inner,"OutputFlag",0)
            end
        else
            error("Unrecognized parameter $optname")
        end
    end
end

function MPB.loadproblem!(m::GurobiMathProgModel, filename::AbstractString)
    read_model(m.inner, filename)
    m.obj = MPB.getobj(m)
    if checkvalue(m.obj, GRB_INFINITY)
        _objwarning(m.obj)
    end
    _truncateobj!(m.obj)
    m.lb = MPB.getconstrLB(m)
    m.ub = MPB.getconstrUB(m)
end

function MPB.loadproblem!(m::GurobiMathProgModel, A::LinearAlgebra.Adjoint{T, Array{T, 2}},
                        collb, colub, obj, rowlb, rowub, sense) where T
    MPB.loadproblem!(m, collect(A), collb, colub, obj, rowlb, rowub, sense)
end

function MPB.loadproblem!(m::GurobiMathProgModel, A, collb, colub, obj, rowlb, rowub, sense)
  # throw away old model but keep env and finalize_env
  env = m.inner.env
  finalize_env = m.inner.finalize_env
  m.inner.finalize_env = false
  free_model(m.inner)
  m.inner = Model(env, "", finalize_env=finalize_env)
  # Re-set options on new model's copy of env
  setparams!(m)

  add_cvars!(m.inner, float(copy(obj)), float(collb), float(colub))
  update_model!(m.inner)

  neginf = typemin(eltype(rowlb))
  posinf = typemax(eltype(rowub))

  # check if we have any range constraints
  # to properly support these, we will need to keep track of the
  # slack variables automatically added by gurobi.
  rangeconstrs = sum((rowlb .!= rowub) .& (rowlb .> neginf) .& (rowub .< posinf))
  if rangeconstrs > 0
      @warn("Julia Gurobi interface doesn't properly support range " *
            "(two-sided) constraints. See Gurobi.jl issue #14")
      add_rangeconstrs!(m.inner, float(A), float(rowlb), float(rowub))
      # A work around for the additional slack variables introduced and the
      # reformulation bug warning
      append!(obj, zeros(rangeconstrs))
  else
      b = Array{Float64}(undef, length(rowlb))
      senses = Array{Cchar}(undef, length(rowlb))
      for i in 1:length(rowlb)
          if rowlb[i] == rowub[i]
              senses[i] = '='
              b[i] = rowlb[i]
          elseif rowlb[i] > neginf
              senses[i] = '>'
              b[i] = rowlb[i]
          else
              @assert rowub[i] < posinf
              senses[i] = '<'
              b[i] = rowub[i]
          end
      end
      add_constrs!(m.inner, float(A), senses, b)
  end

  m.lb, m.ub = rowlb, rowub
  m.obj = copy(obj)
  _truncateobj!(m.obj)
  update_model!(m.inner)
  MPB.setsense!(m,sense)
end

function _truncateobj(v::Real)
    # Gurobi truncates objective coefficients to +/-GRB_INFINITY
    # We need to do this to the m.obj vector for consistency
    #   see #94
    if v > GRB_INFINITY
        return GRB_INFINITY
    elseif v < -GRB_INFINITY
        return -GRB_INFINITY
    end
    return v
end
_truncateobj!(obj::Vector) = map!(_truncateobj, obj, obj)

MPB.writeproblem(m::GurobiMathProgModel, filename::AbstractString) = write_model(m.inner, filename)

MPB.getvarLB(m::GurobiMathProgModel)     = get_dblattrarray( m.inner, "LB", 1, num_vars(m.inner))
function MPB.setvarLB!(m::GurobiMathProgModel, l)
    if checkvalue(l, GRB_BOUNDMAX)
        _boundwarning(l, getvarUB(m))
    end
    set_dblattrarray!(m.inner, "LB", 1, num_vars(m.inner), l)
end

MPB.getvarUB(m::GurobiMathProgModel)     = get_dblattrarray( m.inner, "UB", 1, num_vars(m.inner))
function MPB.setvarUB!(m::GurobiMathProgModel, u)
    if checkvalue(u, GRB_BOUNDMAX)
        _boundwarning(getvarLB(m), u)
    end
    set_dblattrarray!(m.inner, "UB", 1, num_vars(m.inner), u)
end

function MPB.getconstrLB(m::GurobiMathProgModel)
    sense = get_charattrarray(m.inner, "Sense", 1, num_constrs(m.inner))
    ret   = get_dblattrarray(m.inner, "RHS", 1, num_constrs(m.inner))
    for i = 1:num_constrs(m.inner)
        if sense[i] == '>' || sense[i] == '='
            # Do nothing
        else
            # LEQ constraint, so LB is -Inf
            ret[i] = -Inf
        end
     end
     return ret
end

function MPB.getconstrUB(m::GurobiMathProgModel)
    sense = get_charattrarray(m.inner, "Sense", 1, num_constrs(m.inner))
    ret   = get_dblattrarray(m.inner, "RHS", 1, num_constrs(m.inner))
    for i = 1:num_constrs(m.inner)
        if sense[i] == '<' || sense[i] == '='
            # Do nothing
        else
            # GEQ constraint, so UB is +Inf
            ret[i] = +Inf
        end
    end
    return ret
end

# setconstrLB!(m::GurobiMathProgModel, lb) = (m.changed_constr_bounds = true; m.last_op_type = :Con; m.lb = copy(lb))
# setconstrUB!(m::GurobiMathProgModel, ub) = (m.changed_constr_bounds = true; m.last_op_type = :Con; m.ub = copy(ub))
MPB.setconstrLB!(m::GurobiMathProgModel, lb) = (m.changed_constr_bounds = true; m.lb = copy(lb))
MPB.setconstrUB!(m::GurobiMathProgModel, ub) = (m.changed_constr_bounds = true; m.ub = copy(ub))

MPB.getobj(m::GurobiMathProgModel, i::Int=1) = get_dblattrarray( m.inner, "Obj", i, num_vars(m.inner))
function MPB.setobj!(m::GurobiMathProgModel, c, i::Int=1)
    if checkvalue(c, GRB_INFINITY)
        _objwarning(c)
    end
    m.obj = copy(c)
    _truncateobj!(m.obj)
    set_dblattrarray!(m.inner, "Obj", i, num_vars(m.inner), c)
end

set_multiobj_n!(m::Gurobi.Model, n::Int) = Gurobi.set_intattr!(m, "NumObj", n)
get_multiobj_n(m::Gurobi.Model)          = Gurobi.get_intattr(m, "NumObj")

function set_multiobj_c!(m::Model, i::Int, c::AbstractVector{Float64})
    _chklen(c, num_vars(m))
    i0 = get_int_param(m, "ObjNumber")
    set_int_param!(m, "ObjNumber", i-1)
    set_dblattrarray!(m, "ObjN", 1, num_vars(m), c)
    set_int_param!(m, "ObjNumber", i0)
end

function get_multiobj_c(m::Model, i::Int)
    i0 = get_int_param(m, "ObjNumber")
    set_int_param!(m, "ObjNumber", i-1)
    c = get_dblattrarray(m, "ObjN", 1, num_vars(m))
    set_int_param!(m, "ObjNumber", i0)
    return c
end

function set_multiobj_priority!(m::Model, i::Int, priority::Int)
    i0 = get_int_param(m, "ObjNumber")
    set_int_param!(m, "ObjNumber", i-1)
    set_intattr!(m, "ObjNPriority", priority)
    set_int_param!(m, "ObjNumber", i0)
end

function get_multiobj_priority(m::Model, i::Int)::Int
    i0 = get_int_param(m, "ObjNumber")
    set_int_param!(m, "ObjNumber", i-1)
    p = get_intattr(m, "ObjNPriority")
    set_int_param!(m, "ObjNumber", i0)
    return p
end

function set_multiobj_weight!(m::Model, i::Int, w::Float64)
    i0 = get_int_param(m, "ObjNumber")
    set_int_param!(m, "ObjNumber", i-1)
    set_dblattr!(m, "ObjNWeight", w)
    set_int_param!(m, "ObjNumber", i0)
end

function get_multiobj_weight(m::Model, i::Int)::Float64
    i0 = get_int_param(m, "ObjNumber")
    set_int_param!(m, "ObjNumber", i-1)
    w::Float64 = get_dblattr(m, "ObjNWeight")
    set_int_param!(m, "ObjNumber", i0)
    return w
end

function set_multiobj!(m::Model, i::Int, c::AbstractVector{Float64}, p::Int, w::Float64)
    nobj = get_multiobj_n(m)
    if !(1 ≤ i ≤ nobj)
        error("Tried to set objective $i of $nobj. You must use the set_multiobj_n to specify the correct number of objective functions for your problem.")
    end
    set_multiobj_c!(m, i, c)
    set_multiobj_priority!(m, i, p)
    set_multiobj_weight!(m, i, w)
end

function MPB.addvar!(m::GurobiMathProgModel, constridx, constrcoef, l, u, objcoef)
    if m.last_op_type == :Con
        updatemodel!(m)
        m.last_op_type = :Var
    end
    push!(m.obj, _truncateobj(objcoef))
    add_var!(m.inner, length(constridx), constridx, float(constrcoef), float(objcoef), float(l), float(u), GRB_CONTINUOUS)
end
MPB.addvar!(m::GurobiMathProgModel, l, u, objcoef) = MPB.addvar!(m, Int[], Float64[], l, u, objcoef)

function MPB.delvars!(m::GurobiMathProgModel, idx)
    deleteat!(m.obj, idx)
    del_vars!(m.inner, idx)
end

function MPB.addconstr!(m::GurobiMathProgModel, varidx, coef, lb, ub)
    if m.last_op_type == :Var
        updatemodel!(m)
        m.last_op_type = :Con
    end
    if lb == -Inf
        # <= constraint
        add_constr!(m.inner, varidx, coef, '<', ub)
    elseif ub == +Inf
        # >= constraint
        add_constr!(m.inner, varidx, coef, '>', lb)
    elseif lb == ub
        # == constraint
        add_constr!(m.inner, varidx, coef, '=', lb)
    else
        # Range constraint
        error("Adding range constraints not supported yet.")
    end
end
MPB.delconstrs!(m::GurobiMathProgModel, idx) = del_constrs!(m.inner, idx)

MPB.changecoeffs!(m::GurobiMathProgModel, cidx, vidx, val) = chg_coeffs!(m.inner, cidx, vidx, val)

function updatemodel!(m::GurobiMathProgModel)
    update_model!(m.inner)
    if Gurobi.version < v"7.0" && m.obj != MPB.getobj(m)
        @warn("""
            You have encountered a known bug in Gurobi. Any information you query from the model may be incorrect.
            This bug has existed since the first version of Gurobi but is fixed in Gurobi v7.0.

            For more information go to https://github.com/JuliaOpt/Gurobi.jl/issues/60
            Please leave a comment stating that you encountered this bug! We would like to know how prevalent it is.
        """)
    end
    if m.changed_constr_bounds
    # update lower/upper bounds on linear constraints (if they're consistent...)
        sense = get_charattrarray(m.inner, "Sense", 1, num_constrs(m.inner))
        rhs   = get_dblattrarray( m.inner, "RHS",   1, num_constrs(m.inner))
        lb, ub = m.lb, m.ub
        @assert (n_rows = num_constrs(m.inner)) == length(lb) == length(ub) ==
                                                   length(sense) == length(rhs)
        for i = 1:n_rows
            if -Inf < lb[i] == ub[i] < Inf
                sense[i] = '='
                rhs[i]   = lb[i]
            elseif (-Inf < lb[i] && ub[i] < Inf)
                error("Gurobi.jl does not currently support range constraints")
            elseif (lb[i] == -Inf && ub[i] < Inf)
                sense[i] = '<'
                rhs[i]   = ub[i]
            elseif (lb[i] > -Inf && ub[i] == Inf)
                sense[i] = '>'
                rhs[i]   = lb[i]
            else
                error("Internal error.")
            end
        end
        set_charattrarray!(m.inner, "Sense", 1, num_constrs(m.inner), sense)
        set_dblattrarray!(m.inner, "RHS", 1, num_constrs(m.inner), rhs)
        m.changed_constr_bounds = false
        update_model!(m.inner)
    end
end

MPB.getconstrmatrix(m::GurobiMathProgModel) = get_constrmatrix(m.inner)

function MPB.setsense!(m::GurobiMathProgModel, sense)
  if sense == :Min
    set_sense!(m.inner, :minimize)
  elseif sense == :Max
    set_sense!(m.inner, :maximize)
  else
    error("Unrecognized objective sense $sense")
  end
end
function MPB.getsense(m::GurobiMathProgModel)
  v = get_intattr(m.inner, "ModelSense")
  if v == -1
    return :Max
  else
    return :Min
  end
end

MPB.numvar(m::GurobiMathProgModel)    = (updatemodel!(m); num_vars(m.inner))
MPB.numconstr(m::GurobiMathProgModel) = num_constrs(m.inner) + num_qconstrs(m.inner)
MPB.numlinconstr(m::GurobiMathProgModel) = num_constrs(m.inner)
MPB.numquadconstr(m::GurobiMathProgModel) = num_qconstrs(m.inner)

function MPB.optimize!(m::GurobiMathProgModel)
    # set callbacks if present
    if m.lazycb != nothing || m.cutcb != nothing || m.heuristiccb != nothing || m.infocb != nothing
        updatemodel!(m)
        if !is_mip(m.inner)
            @warn("Gurobi ignores branch-and-bound callbacks when no discrete elements are present in the model.")
        end
        setmathprogcallback!(m)
    end
    if m.lazycb != nothing
      setparam!(m.inner, "LazyConstraints", 1)
    end
    updatemodel!(m)
    optimize(m.inner)
end

function MPB.status(m::GurobiMathProgModel)
  s = get_status(m.inner)
  if s == :optimal
    return :Optimal
  elseif s == :infeasible
    return :Infeasible
  elseif s == :unbounded
    return :Unbounded
  elseif s == :inf_or_unbd
    @warn("Gurobi reported infeasible or unbounded. Set InfUnbdInfo=1 for more specific status.")
    return :InfeasibleOrUnbounded
  elseif s == :iteration_limit || s == :node_limit || s == :time_limit || s == :solution_limit
    return :UserLimit
  elseif s == :numeric
    return :Numeric
  elseif s == :suboptimal
    return :Suboptimal # not very useful status
  elseif s == :interrupted # ended by user?
    return :UserLimit
  elseif s == :inprogress
    return :InProgress
  elseif s == :user_obj_limit
    return :UserObjLimit
  elseif s == :cutoff
    return :Cutoff
  else
    error("Unrecognized solution status: $s")
  end
end

MPB.getobjval(m::GurobiMathProgModel)   = get_objval(m.inner)
MPB.getobjbound(m::GurobiMathProgModel) = get_objbound(m.inner)
MPB.getsolution(m::GurobiMathProgModel) = get_solution(m.inner)

function MPB.getconstrsolution(m::GurobiMathProgModel)
    sense = get_charattrarray(m.inner, "Sense", 1, num_constrs(m.inner))
    rhs   = get_dblattrarray( m.inner, "RHS",   1, num_constrs(m.inner))
    slack = get_dblattrarray( m.inner, "Slack", 1, num_constrs(m.inner))
    ret   = zeros(num_constrs(m.inner))
    for i = 1:num_constrs(m.inner)
        if sense[i] == '='
            # Must be equal to RHS if feasible
            ret[i] = rhs[i]
        else
            # Slack variable is non-negative for <= constraints, and non-positive for >= constraints
            ret[i] = rhs[i] - slack[i]
        end
    end
    return ret
end

function MPB.getreducedcosts(m::GurobiMathProgModel)
    if is_qcp(m.inner) && get_int_param(m.inner, "QCPDual") == 0
        return fill(NaN, num_vars(m.inner))
    else
        return get_dblattrarray(m.inner, "RC", 1, num_vars(m.inner))
    end
end

function MPB.getconstrduals(m::GurobiMathProgModel)
    if is_qcp(m.inner) && get_int_param(m.inner, "QCPDual") == 0
        return fill(NaN, num_constrs(m.inner))
    else
        return get_dblattrarray(m.inner, "Pi", 1, num_constrs(m.inner))
    end
end

function MPB.getquadconstrduals(m::GurobiMathProgModel)
    if is_qcp(m.inner) && get_int_param(m.inner, "QCPDual") == 0
        return fill(NaN, num_qconstrs(m.inner))
    else
        return get_dblattrarray(m.inner, "QCPi", 1, num_qconstrs(m.inner))
    end
end

MPB.getinfeasibilityray(m::GurobiMathProgModel) = -get_dblattrarray(m.inner, "FarkasDual", 1, num_constrs(m.inner)) # note sign is flipped
MPB.getunboundedray(m::GurobiMathProgModel) = get_dblattrarray(m.inner, "UnbdRay", 1, num_vars(m.inner))

MPB.getbasis(m::GurobiMathProgModel) = get_basis(m.inner)

MPB.getrawsolver(m::GurobiMathProgModel) = m.inner

const var_type_map = Dict(
  'C' => :Cont,
  'B' => :Bin,
  'I' => :Int,
  'S' => :SemiCont,
  'N' => :SemiInt
)

const rev_var_type_map = Dict(
  :Cont => 'C',
  :Bin => 'B',
  :Int => 'I',
  :SemiCont => 'S',
  :SemiInt => 'N'
)

function MPB.setvartype!(m::GurobiMathProgModel, vartype::Vector{Symbol})
    # do this to make sure we deal with new columns
    updatemodel!(m)
    nvartype = map(x->rev_var_type_map[x], vartype)
    set_charattrarray!(m.inner, "VType", 1, length(nvartype), nvartype)
    updatemodel!(m) # otherwise getvartype! will return old values
end

function MPB.getvartype(m::GurobiMathProgModel)
    ret = get_charattrarray(m.inner, "VType", 1, num_vars(m.inner))
    map(x->var_type_map[x], ret)
end

function MPB.setwarmstart!(m::GurobiMathProgModel, v)
    for j = 1:length(v)
        if isnan(v[j])
            v[j] = 1e101  # GRB_UNDEFINED
        end
    end
    set_dblattrarray!(m.inner, "Start", 1, num_vars(m.inner), v)
end

MPB.addsos1!(m::GurobiMathProgModel, idx, weight) = add_sos!(m.inner, :SOS1, idx, weight)
MPB.addsos2!(m::GurobiMathProgModel, idx, weight) = add_sos!(m.inner, :SOS2, idx, weight)

# Callbacks

MPB.setlazycallback!(m::GurobiMathProgModel,f) = (m.lazycb = f)
MPB.setcutcallback!(m::GurobiMathProgModel,f) = (m.cutcb = f)
MPB.setheuristiccallback!(m::GurobiMathProgModel,f) = (m.heuristiccb = f)
MPB.setinfocallback!(m::GurobiMathProgModel,f) = (m.infocb = f)

mutable struct GurobiCallbackData <: MPB.MathProgCallbackData
    cbdata::CallbackData
    state::Symbol
    where::Cint
    sol::Vector{Float64}  # Used for heuristic callbacks
#    model::GurobiMathProgModel # not needed?
end

function MPB.cbgetmipsolution(d::GurobiCallbackData)
    @assert d.state == :MIPSol
    return cbget_mipsol_sol(d.cbdata, d.where)
end

function MPB.cbgetmipsolution(d::GurobiCallbackData,output)
    @assert d.state == :MIPSol
    return cbget_mipsol_sol(d.cbdata, d.where, output)
end

function MPB.cbgetlpsolution(d::GurobiCallbackData)
    @assert d.state == :MIPNode
    return cbget_mipnode_rel(d.cbdata, d.where)
end

function MPB.cbgetlpsolution(d::GurobiCallbackData, output)
    @assert d.state == :MIPNode
    return cbget_mipnode_rel(d.cbdata, d.where, output)
end

# TODO: macro for these getters?
function MPB.cbgetobj(d::GurobiCallbackData)
    if d.state == :MIPNode
        return cbget_mipnode_objbst(d.cbdata, d.where)
    elseif d.state == :Intermediate
        return cbget_mip_objbst(d.cbdata, d.where)
    elseif d.state == :MIPSol
        error("Gurobi does not implement cbgetobj when state == MIPSol")
        # https://groups.google.com/forum/#!topic/gurobi/Az_X6Ag-y6k
        # https://github.com/JuliaOpt/MathProgBase.jl/issues/23
        # A complex workaround is possible but hasn't been implemented.
        # return cbget_mipsol_objbst(d.cbdata, d.where)
    else
        error("Unrecognized callback state $(d.state)")
    end
end

function MPB.cbgetbestbound(d::GurobiCallbackData)
    if d.state == :MIPNode
        return cbget_mipnode_objbnd(d.cbdata, d.where)
    elseif d.state == :MIPSol
        return cbget_mipsol_objbnd(d.cbdata, d.where)
    elseif d.state == :Intermediate
        return cbget_mip_objbnd(d.cbdata, d.where)
    else
        error("Unrecognized callback state $(d.state)")
    end
end

function MPB.cbgetexplorednodes(d::GurobiCallbackData)
    if d.state == :MIPNode
        return cbget_mipnode_nodcnt(d.cbdata, d.where)
    elseif d.state == :MIPSol
        return cbget_mipsol_nodcnt(d.cbdata, d.where)
    elseif d.state == :Intermediate
        return cbget_mip_nodcnt(d.cbdata, d.where)
    else
        error("Unrecognized callback state $(d.state)")
    end
end

# returns :MIPNode :MIPSol :Intermediate
MPB.cbgetstate(d::GurobiCallbackData) = d.state

function MPB.cbaddcut!(d::GurobiCallbackData,varidx,varcoef,sense,rhs)
    @assert d.state == :MIPNode
    cbcut(d.cbdata, convert(Vector{Cint}, varidx), float(varcoef), sense, float(rhs))
end

function MPB.cbaddlazy!(d::GurobiCallbackData,varidx,varcoef,sense,rhs)
    @assert d.state == :MIPNode || d.state == :MIPSol
    cblazy(d.cbdata, convert(Vector{Cint}, varidx), float(varcoef), sense, float(rhs))
end

function MPB.cbaddsolution!(d::GurobiCallbackData)
    # Gurobi doesn't support adding solutions on MIPSol.
    # TODO: support this anyway
    @assert d.state == :MIPNode
    cbsolution(d.cbdata, d.sol)
    # "Wipe" solution back to GRB_UNDEFINIED
    for i in 1:length(d.sol)
        d.sol[i] = 1e101  # GRB_UNDEFINED
    end
end

function MPB.cbsetsolutionvalue!(d::GurobiCallbackData,varidx,value)
    d.sol[varidx] = value
end

# breaking abstraction, define our low-level callback to eliminatate
# a level of indirection

function mastercallback(ptr_model::Ptr{Cvoid}, cbdata::Ptr{Cvoid}, where::Cint, userdata::Ptr{Cvoid})

    model = unsafe_pointer_to_objref(userdata)::GurobiMathProgModel
    grbrawcb = CallbackData(cbdata,model.inner)
    if where == CB_MIPSOL
        state = :MIPSol
    elseif where == CB_MIPNODE
        state = :MIPNode
        # skip callback if node is reported to be cut off or infeasible --
        # nothing to do.
        # TODO: users may want this information
        status = cbget_mipnode_status(grbrawcb, where)
        if status != 2
            return convert(Cint,0)
        end
    elseif where == CB_MIP
        state = :Intermediate
    else
        # State with no relevant callbacks
        return convert(Cint,0)
    end

    grbcb = GurobiCallbackData(grbrawcb, state, where, [0.0])
    if model.infocb != nothing
        ret = model.infocb(grbcb)
        if ret == :Exit
            terminate(model.inner)
        end
    end
    if model.cutcb != nothing && state == :MIPNode
        ret = model.cutcb(grbcb)
        if ret == :Exit
            terminate(model.inner)
        end
    end
    if model.heuristiccb != nothing && state == :MIPNode
        grbcb.sol = fill(1e101, MPB.numvar(model))  # GRB_UNDEFINED
        ret = model.heuristiccb(grbcb)
        if ret == :Exit
            terminate(model.inner)
        end
    end
    if model.lazycb != nothing && (state == :MIPSol || state == :MIPNode)
        ret = model.lazycb(grbcb)
        if ret == :Exit
            terminate(model.inner)
        end
    end

    return convert(Cint,0)
end

# User callback function should be of the form:
# callback(cbdata::MPB.MathProgCallbackData)
# return :Exit to indicate an error

function setmathprogcallback!(model::GurobiMathProgModel)
    if Sys.iswindows() && Sys.WORD_SIZE != 64
        error("Callbacks not currently supported on Win32. Use 64-bit Julia with 64-bit Gurobi.")
    end
    grbcallback = @cfunction(mastercallback, Cint, (Ptr{Cvoid}, Ptr{Cvoid}, Cint, Ptr{Cvoid}))
    ret = @grb_ccall(setcallbackfunc, Cint, (Ptr{Cvoid}, Ptr{Cvoid}, Any), model.inner.ptr_model, grbcallback, model)
    if ret != 0
        throw(GurobiError(model.env, ret))
    end
    nothing
end

# QCQP

function MPB.setquadobj!(m::GurobiMathProgModel, rowidx, colidx, quadval)
    delq!(m.inner)
    update_model!(m.inner)
    scaledvals = similar(quadval)
    for i in 1:length(rowidx)
      if rowidx[i] == colidx[i]
        # rescale from matrix format to "terms" format
        scaledvals[i] = quadval[i] / 2
      else
        scaledvals[i] = quadval[i]
      end
    end
    add_qpterms!(m.inner, rowidx, colidx, scaledvals)
    update_model!(m.inner)
end

function MPB.addquadconstr!(m::GurobiMathProgModel, linearidx, linearval, quadrowidx, quadcolidx, quadval, sense, rhs)
    if m.last_op_type == :Var
        updatemodel!(m)
        m.last_op_type = :Con
    end
    add_qconstr!(m.inner, linearidx, linearval, quadrowidx, quadcolidx, quadval, sense, rhs)
end

MPB.getsolvetime(m::GurobiMathProgModel) = get_runtime(m.inner)
MPB.getnodecount(m::GurobiMathProgModel) = get_node_count(m.inner)
MPB.getsimplexiter(m::GurobiMathProgModel) = get_dblattr(m.inner, "IterCount")


export
    LineNode,
    set!

import ..Nodes.draw

mutable struct LineNode <: Ranger.AbstractNode
    base::Nodes.NodeData
    transform::Nodes.TransformProperties{Float64}

    mesh::Geometry.Mesh

    color::Rendering.Palette

    function LineNode(world::Ranger.World, name::String, parent::Ranger.AbstractNode)
        o = new()

        o.base = Nodes.NodeData(name, parent, world)
        o.transform = Nodes.TransformProperties{Float64}()
        o.mesh = Geometry.Mesh()
        o.color = Rendering.White()

        Geometry.add_vertex!(o.mesh, 0.0, 0.0)
        Geometry.add_vertex!(o.mesh, 0.0, 0.0)

        Geometry.build!(o.mesh)

        o
    end
end

function Nodes.draw(node::LineNode, context::Rendering.RenderContext)
    if Nodes.is_dirty(node)
        Rendering.transform!(context, node.mesh)
        Nodes.set_dirty!(node, false)
    end

    Rendering.set_draw_color(context, node.color)

    v1 = node.mesh.bucket[1]
    v2 = node.mesh.bucket[2]
    Rendering.draw_line(context, v1.x, v1.y, v2.x, v2.y);

end

function set!(node::LineNode, x1::Float64, y1::Float64, x2::Float64, y2::Float64)
    Geometry.set!(node.mesh.vertices[1], x1, y1)
    Geometry.set!(node.mesh.vertices[2], x2, y2)

    Nodes.set_dirty!(node, true);
end


using Base.Test

using DataFrames

import Celeste: GalsimBenchmark

GALSIM_CASES_EXERCISED = [
    "simple_star",
    "star_with_noise",
    "angle_and_axis_ratio_1",
    "galaxy_with_all",
    "galaxy_with_noise",
]

function assert_estimates_are_close(benchmark_results)
    for row in eachrow(benchmark_results)
        if isna(row[:truth])
            continue
        end
        if row[:variable] == :is_star
            maximum_error = 0.1
        elseif row[:variable] == :de_vaucouleurs_mixture_weight
            maximum_error = 0.2
        elseif row[:variable] == :angle_deg
            maximum_error = 5
        #elseif !isna(row[1, :error_sds])
        #    maximum_error = 2.5 * row[1, :error_sds]
        else
            maximum_error = 0.2 * abs(row[:truth])
        end
        if !isapprox(row[:truth], row[:estimate], atol=maximum_error)
            @show row
            @show maximum_error
            @test false
        else
            @test true # just so test framework will count test cases
        end
    end
end

@testset "GalSim benchmark tests, single-source inference" begin
    truth, results = GalsimBenchmark.run_benchmarks(
        test_case_names=GALSIM_CASES_EXERCISED,
        joint_inference=false
    )
    assert_estimates_are_close(GalsimBenchmark.truth_comparison_df(truth, results))
end

@testset "GalSim benchmark tests, joint inference" begin
    truth, results = GalsimBenchmark.run_benchmarks(
        test_case_names=GALSIM_CASES_EXERCISED,
        joint_inference=true
    )
    assert_estimates_are_close(GalsimBenchmark.truth_comparison_df(truth, results))
end


# this is an example of how to augment Wavs
using WAV
using SpeechAugment

##------------------------------------##
##---------- exmaple one -------------##
##------------------------------------##

# 1. read a wav file as a speech example
batchsize = 8;
data,fs = wavread("/home/xxx/ASpeechExample.wav");

# 2. init all the augmentation functions you want
echo  = initAddEcho(fs, (0.05,0.4), (3.0,3.2,2.5,3.5,2.0,3.0));
noise = initAddNoise("XXPathFullOfNoiseWAVs", 2, (5,15));
clip  = initClipWav((0.5,2.0));
drop  = initDropWav(fs, (0.09,0.15));
far   = initFarfieldWav(fs, (0.5,0.9));
speed = initSpeedWav((0.8,1.2));

# 3. make a function list
fnlist = [echo noise clip drop far speed];

# 4. augment one wav into #batchSize audios
wavlist = augmentWav(fnlist, data, batchsize)
for i = 1:8
    wavwrite(wavlist[i], "A$i.wav",Fs=16000,nbits=32)
end

##------------------------------------##
##---------- exmaple two -------------##
##------------------------------------##

# 1. prepare #batchSize audios
wavs = Vector(undef, batchsize);
for i = 1:batchsize
    wavs[i] = copy(data)
end
# 2. augment #batchSize audios
wavs = augmentWavs(fnlist, wavs)
for i = 1:batchsize
    wavwrite(wavs[i], "B$i.wav",Fs=16000,nbits=32)
end


B_data = [ -2.5           0.5           0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   3.0          -2.0           1.0           0.0   0.0   0.0   0.0   0.0   0.0   0.0           0.0          0.0           0.0           0.0          0.0           0.0        0.0        0.0       0.0        0.0
1.0          -3.0           0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   3.0          -2.0           1.0           0.0   0.0   0.0   0.0   0.0   0.0   0.0           0.0          0.0           0.0           0.0          0.0           0.0        0.0        0.0       0.0        0.0
0.0           0.0          -2.5   0.0  -0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   2.0           0.0           0.0           0.0   0.0   0.0   0.0   0.0   0.0   0.5           0.0          0.0           0.0           0.0          0.0           0.0        0.0        0.0       0.0        0.0
0.0           0.0          -0.0  -2.5   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0           2.0           0.0           0.0   0.0   0.0   0.0   0.0   0.0   0.0           0.5          0.0           0.0           0.0          0.0           0.0        0.0        0.0       0.0        0.0
0.0           0.0           0.0  -0.0  -2.5   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0           0.0           2.0           0.0   0.0   0.0   0.0   0.0   0.0   0.0           0.0          0.5           0.0           0.0          0.0           0.0        0.0        0.0       0.0        0.0
0.0           0.0           0.0   0.0   0.0  -2.5   0.0  -0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0           0.0           0.0           2.0   0.0   0.0   0.0   0.0   0.0   0.0           0.0          0.0           0.5           0.0          0.0           0.0        0.0        0.0       0.0        0.0
0.0           0.0           0.0   0.0   0.0  -0.0  -2.5   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0           0.0           0.0           0.0   2.0   0.0   0.0   0.0   0.0   0.0           0.0          0.0           0.0           0.5          0.0           0.0        0.0        0.0       0.0        0.0
0.0           0.0           0.0   0.0   0.0   0.0  -0.0  -2.5   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0           0.0           0.0           0.0   0.0   2.0   0.0   0.0   0.0   0.0           0.0          0.0           0.0           0.0          0.5           0.0        0.0        0.0       0.0        0.0
0.0           0.0           0.0   0.0   0.0   0.0   0.0   0.0  -2.5   0.0  -0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0           0.0           0.0           0.0   0.0   0.0   2.0   0.0   0.0   0.0           0.0          0.0           0.0           0.0          0.0           0.5        0.0        0.0       0.0        0.0
0.0           0.0           0.0   0.0   0.0   0.0   0.0   0.0  -0.0  -2.5   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0           0.0           0.0           0.0   0.0   0.0   0.0   2.0   0.0   0.0           0.0          0.0           0.0           0.0          0.0           0.0        0.5        0.0       0.0        0.0
0.0           0.0           0.0   0.0   0.0   0.0   0.0   0.0   0.0  -0.0  -2.5   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0           0.0           0.0           0.0   0.0   0.0   0.0   0.0   2.0   0.0           0.0          0.0           0.0           0.0          0.0           0.0        0.0        0.5       0.0        0.0
0.0           0.0           0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0  -2.5   0.0  -0.0   0.0   0.0   0.0   0.0   0.0   0.0   2.0           0.0           0.0           0.0   0.0   0.0   0.0   0.0   0.0   0.5           0.0          0.0           0.0           0.0          0.0           0.0        0.0        0.0       0.0        0.0
0.0           0.0           0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0  -0.0  -2.5   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0           2.0           0.0           0.0   0.0   0.0   0.0   0.0   0.0   0.0           0.5          0.0           0.0           0.0          0.0           0.0        0.0        0.0       0.0        0.0
0.0           0.0           0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0  -0.0  -2.5   0.0   0.0   0.0   0.0   0.0   0.0   0.0           0.0           2.0           0.0   0.0   0.0   0.0   0.0   0.0   0.0           0.0          0.5           0.0           0.0          0.0           0.0        0.0        0.0       0.0        0.0
0.0           0.0           0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0  -2.5   0.0  -0.0   0.0   0.0   0.0   0.0           0.0           0.0           2.0   0.0   0.0   0.0   0.0   0.0   0.0           0.0          0.0           0.5           0.0          0.0           0.0        0.0        0.0       0.0        0.0
0.0           0.0           0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0  -0.0  -2.5   0.0   0.0   0.0   0.0   0.0           0.0           0.0           0.0   2.0   0.0   0.0   0.0   0.0   0.0           0.0          0.0           0.0           0.5          0.0           0.0        0.0        0.0       0.0        0.0
0.0           0.0           0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0  -0.0  -2.5   0.0   0.0   0.0   0.0           0.0           0.0           0.0   0.0   2.0   0.0   0.0   0.0   0.0           0.0          0.0           0.0           0.0          0.5           0.0        0.0        0.0       0.0        0.0
0.0           0.0           0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0  -2.5   0.0  -0.0   0.0           0.0           0.0           0.0   0.0   0.0   2.0   0.0   0.0   0.0           0.0          0.0           0.0           0.0          0.0           0.5        0.0        0.0       0.0        0.0
0.0           0.0           0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0  -0.0  -2.5   0.0   0.0           0.0           0.0           0.0   0.0   0.0   0.0   2.0   0.0   0.0           0.0          0.0           0.0           0.0          0.0           0.0        0.5        0.0       0.0        0.0
0.0           0.0           0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0  -0.0  -2.5   0.0           0.0           0.0           0.0   0.0   0.0   0.0   0.0   2.0   0.0           0.0          0.0           0.0           0.0          0.0           0.0        0.0        0.5       0.0        0.0
0.0           0.0           1.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0  -5.0           5.0          -2.0           0.0   0.0   0.0   0.0   0.0   0.0   1.0           0.0          0.0           0.0           0.0          0.0           0.0        0.0        0.0       0.0        0.0
0.0           0.0           0.0   1.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0  -0.5          -2.0           0.5           0.0   0.0   0.0   0.0   0.0   0.0   0.0           1.0          0.0           0.0           0.0          0.0           0.0        0.0        0.0       0.0        0.0
0.0           0.0           0.0   0.0   1.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.5          -2.0          -3.5           4.5  -3.0   1.5   0.0   0.0   0.0   0.0           0.0          1.0           0.0           0.0          0.0           0.0        0.0        0.0       0.0        0.0
0.0           0.0           0.0   0.0   0.0   1.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0           0.0           0.0          -5.0   5.0  -2.0   0.0   0.0   0.0   0.0           0.0          0.0           1.0           0.0          0.0           0.0        0.0        0.0       0.0        0.0
0.0           0.0           0.0   0.0   0.0   0.0   1.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0           0.0           0.0          -0.5  -2.0   0.5   0.0   0.0   0.0   0.0           0.0          0.0           0.0           1.0          0.0           0.0        0.0        0.0       0.0        0.0
0.0           0.0           0.0   0.0   0.0   0.0   0.0   1.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0           0.0           0.0           0.5  -2.0  -3.5   4.5  -3.0   1.5   0.0           0.0          0.0           0.0           0.0          1.0           0.0        0.0        0.0       0.0        0.0
0.0           0.0           0.0   0.0   0.0   0.0   0.0   0.0   1.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0           0.0           0.0           0.0   0.0   0.0  -5.0   5.0  -2.0   0.0           0.0          0.0           0.0           0.0          0.0           1.0        0.0        0.0       0.0        0.0
0.0           0.0           0.0   0.0   0.0   0.0   0.0   0.0   0.0   1.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0           0.0           0.0           0.0   0.0   0.0  -0.5  -2.0   0.5   0.0           0.0          0.0           0.0           0.0          0.0           0.0        1.0        0.0       0.0        0.0
0.0           0.0           0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   1.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0           0.0           0.0           0.0   0.0   0.0   0.5  -2.0  -3.5   0.0           0.0          0.0           0.0           0.0          0.0           0.211839  -0.423678   1.63552   0.436525   2.1398
3.26466       0.683008      0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   1.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   2.8285       -0.552334      0.276167      0.0   0.0   0.0   0.0   0.0   0.0  -5.25         -3.0          0.75          0.0           0.0          0.0           0.0        0.0        0.0       0.0        0.0
1.35919e-16   2.84359e-17   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   1.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   3.44933e-17   2.0           1.14978e-17   0.0   0.0   0.0   0.0   0.0   0.0   0.75         -3.0         -0.75          0.0           0.0          0.0           0.0        0.0        0.0       0.0        0.0
-9.06125e-16  -1.89573e-16   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   1.0   0.0   0.0   0.0   0.0   0.0   0.0  -2.29955e-16   1.53303e-16   2.0           0.0   0.0   0.0   0.0   0.0   0.0  -3.0           7.5         -7.5           0.0           0.0          0.0           0.0        0.0        0.0       0.0        0.0
0.0           0.0           0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   1.0   0.0   0.0   0.0   0.0   0.0   0.0           0.0           0.0           2.0   0.0   0.0   0.0   0.0   0.0   2.25         -4.5          6.75         -5.25         -3.0          0.75          0.0        0.0        0.0       0.0        0.0
0.0           0.0           0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   1.0   0.0   0.0   0.0   0.0   0.0           0.0           0.0           0.0   2.0   0.0   0.0   0.0   0.0   9.36751e-17  -1.8735e-16   2.81025e-16   0.75         -3.0         -0.75          0.0        0.0        0.0       0.0        0.0
0.0           0.0           0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   1.0   0.0   0.0   0.0   0.0           0.0           0.0           0.0   0.0   2.0   0.0   0.0   0.0  -6.245e-16     1.249e-15   -1.8735e-15   -3.0           7.5         -7.5           0.0        0.0        0.0       0.0        0.0
0.0           0.0           0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   1.0   0.0   0.0   0.0           0.0           0.0           0.0   0.0   0.0   2.0   0.0   0.0   0.0           0.0          0.0           2.25         -4.5          6.75         -5.25      -3.0        0.75      0.0        0.0
0.0           0.0           0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   1.0   0.0   0.0           0.0           0.0           0.0   0.0   0.0   0.0   2.0   0.0   0.0           0.0          0.0           9.36751e-17  -1.8735e-16   2.81025e-16   0.75      -3.0       -0.75      0.0        0.0
0.0           0.0           0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   1.0   0.0           0.0           0.0           0.0   0.0   0.0   0.0   0.0   2.0   0.0           0.0          0.0          -6.245e-16     1.249e-15   -1.8735e-15   -3.0        7.5       -7.5       0.0        0.0
0.0           0.0           0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0           0.0           0.0           0.0   0.0   0.0   0.0   0.0   0.0   0.0           0.0          0.0           0.0           0.0          0.0           0.25      -0.5        0.75     -2.5        2.0
0.0           0.0           0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0           0.0           0.0           0.0   0.0   0.0   0.0   0.0   0.0   0.0           0.0          0.0           0.0           0.0          0.0           0.5       -1.0        1.5       1.0       -2.0]


export SupervisedModel
export iterations, cost, fitness, predict, predictProb, accuracy
export stop!, state, remember!, history, trainingCurve, name
export @remember!, @history

import StatsBase.predict

function defaultCallback()
end

# ==========================================================================
# Interface for all models

cost(model::Any) = @_not_implemented
cost(model::Any, data::DataSource) = @_not_implemented
fitness(model::Any) = -1 * cost(model)
fitness(model::Any, data::DataSource) = -1 * cost(model, data)
#predict(model::Any, data::DataSource) = @_not_implemented # Statsbase takes care of this
predictProb(model::Any, data::DataSource) = @_not_implemented
accuracy(model::Any, data::DataSource) = @_not_implemented
train!(callback::Function, model::Any, data::DataSource, solver::AbstractSolver, args...; nargs...) = @_not_implemented

# ==========================================================================
# Use default callback if no other is specified

train!(model::Any, data::DataSource, solver::AbstractSolver, args...; nargs...) =
  train!(defaultCallback, model, data, solver, args...; nargs...)

# ==========================================================================
# Embedding supervised learning model for convenience

type SupervisedModel{T<:Any}
  _state::State
  _iterations::Int
  _native::T
  _history::TrainingHistory{Int}
end

function SupervisedModel{T<:Any}(model::T)
  SupervisedModel{T}(:uninitialized, zero(Int), model, TrainingHistory(Int))
end

# ==========================================================================
# Macros for simpler function defintion

macro _static_always(f, command, args...)
  esc(:(($f{T<:Any})(model::SupervisedModel{T},$(args...)) = $command))
end

macro _static_training_only(f, command, args...)
  msg = string(f, " only defined during model training")
  esc(:(($f{T<:Any})(model::SupervisedModel{T},$(args...)) = :training <= model._state <= :trystop ? $command : throw(StateError($msg))))
end

macro _static_initialized_only(f, command, args...)
  msg = string(f, " not defined for uninitialized or erroneous models")
  esc(:(($f{T<:Any})(model::SupervisedModel{T},$(args...)) = model._state < :initialized ? throw(StateError($msg)) : $command))
end

macro _delegate_initialized_only(f, args...)
  vars = map(x->x.args[1], args)
  msg = string(f, " not defined for uninitialized or erroneous models")
  esc(:(($f{T<:Any})(model::SupervisedModel{T},$(args...)) = model._state < :initialized ? throw(StateError($msg)) : $f(model._native,$(vars...))))
end

# ==========================================================================
# Define the local and delegating functions

@_static_always(name, string(typeof(model._native)))
@_static_always(state, model._state) # :untrained, :training, :trained
@_static_always(iterations, model._iterations)
@_static_training_only(stop!, model._state = :trystop)
@_static_training_only(remember!, push!(model._history, model._iterations, key, value), key::Symbol, value)
@_static_initialized_only(history, get(model._history, key), key::Symbol)
@_static_initialized_only(trainingCurve, get(model._history, :trainingCurve))
@_delegate_initialized_only(cost)
@_delegate_initialized_only(cost, data::DataSource)
@_delegate_initialized_only(fitness)
@_delegate_initialized_only(fitness, data::DataSource)
@_delegate_initialized_only(predict, data::DataSource)
@_delegate_initialized_only(predictProb, data::DataSource)
@_delegate_initialized_only(accuracy, data::DataSource)

macro remember!(model, func)
  key = string(func)
  esc(:(@remember!($model, symbol($key), $func)))
end

macro history(model, func)
  key = string(func)
  esc(:(@get($model, symbol($key))))
end

# ==========================================================================
# Train method

function train!{T<:Any, D<:LabeledDataSource}(
    userCallback::Function,
    model::SupervisedModel{T},
    data::D,
    solver::AbstractSolver,
    args...;
    max_iter::Int = 1000,
    break_every::Int = -1, # -1 = automatic, 0 = off
    nargs...)
  @assert max_iter > 0
  break_every = break_every < 0 ? safeFloor(max_iter / 10) : break_every
  enable_callback = break_every > 0

  # Define callback for the native model
  function localCallback()
    model._iterations += 1
    if model._iterations % break_every == 0
      remember!(model, :trainingCurve, cost(model)) # This should always be trivial cheap
      userCallback()
      model._state._idx == @stateidx(:trystop) ? :stop : nothing
    end
  end

  # If callback is disabled only pass the empty function
  cb = enable_callback ? localCallback : defaultCallback

  # Train the native model
  model._state = :training
  try
    itersOld = model._iterations
    itersNew = train!(cb, model._native, data, solver, args...; max_iter=max_iter, nargs...)
    model._iterations = itersOld + itersNew
    model._state = :trained
  catch ex
    print("An exception was thrown while training: ", ex)
    model._state = :error
  end
  model
end


module SeqEvol

export evolMSA, extract_params

using DelimitedFiles
using FastaIO
using LinearAlgebra
using StatsBase
using GZip

include("read_write.jl")
include("evol.jl")

end # module


# Equations of state

using WaterData
import JLD, Dierckx, VoronoiDelaunay, GeometricalPredicates
import WaterData: extracteos
using IterTools: chain

# The bulk of the EOSes and heat capacities are provided in the WaterData
# package. We define one more EOS which is mass-piecewise (for modelling layered
# planets)
""" Equation of state which is piecewise in the mass coordinate

    * `eoses`: Vector of EOSes for each piece.
    * `edges`: Vector defining the edge of each piece. Evaluating
      past the ends of this vector will just evaluate the nearest EOS. """
struct MassPiecewiseEOS{E<:EOS, M<:Mass} <: WaterData.PiecewiseEOS
    eoses::Vector{E}
    edges::Vector{M}
end
function MassPiecewiseEOS(eoses, M::Mass, mass_fractions)
    layer_edges = vcat([-Inf], cumsum(mass_fractions)) .* M
    layer_edges[end] = Inf
    MassPiecewiseEOS(eoses, layer_edges)
end

# EOS evaluation
# (we loop through a bunch of EOSes to avoid ambiguity clashes)

# Split piecewise EOSes appropriately
function extracteos(eos::MassPiecewiseEOS, vs::ValueSet)
    extracteos(eos, mass(vs))
end
function extracteos(eos::WaterData.PressurePiecewiseEOS, vs::ValueSet)
    extracteos(eos, pressure(vs))
end
for T in (MassRadiusPressure, PhysicalValues)
    (eos::MassPiecewiseEOS)(vs::T) = extracteos(eos, vs)(vs)
end

# call definitions have gone to atmosphere.jl as a workaround


include("deflump.jl")

filename = "methanecomplete.kpp"

f = open(filename)
lines = [replace(ln,r"[;\h]+","") for ln in split(replace(readstring(f),r":\h*\n",":"),"\n")]
eqns = [split(eq,r"[\}=:]") for eq in filter(x->ismatch(r"\{[\d\.]+\}", x), lines)]

specs = Set(vcat(pmap(getspecs,eqns)...))
nspecs = length(specs)

#rate number dictionary
rated = Set([r[4] for r in eqns])
r_rated = Dict(zip(UInt32(1):UInt32(length(rated)),rated))
rated = Dict(zip(rated,UInt32(1):UInt32(length(rated))))

#spec dicionary
specd = Dict(zip(specs,UInt32(1):UInt32(length(specs))))
rspec = Dict(zip(UInt32(1):UInt32(length(specs)),specs))

@eval @everywhere rated=$rated
@eval @everywhere specd=$specd

#each eq numbered
eqnum = pmap(numerate,eqns)
@eval @everywhere eqnum=$eqnum

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

lumplist = [["CH3O2","CH3O2NO2"],["HCHO","CH3O"]]
#lumplist = [["CH3O2","CH3O2NO2"]]
#lumplist = [["CL"]]#,["HCHO","CH3O"]]

pspecs = [specd[x] for x in collect(Base.flatten(lumplist))]
parr = pmap(production_array,pspecs)
parr = Dict(zip(pspecs,parr))

lumpdict = []
lumpspec = []
  ppair = []

counter = 1
for pairing in enumerate(lumplist)
  for j in pairing[2]
    push!(lumpdict, [j,pairing[1]] )
    push!(lumpspec, [j,counter] )
    counter +=1
  end
end

  @eval @everywhere lumpspec =  Dict(lumpspec)
lumpdict = Dict(lumpdict)

for element in enumerate(lumplist)

  @eval @everywhere pnum = $element[1]
  pairing = element[2]
  p(pairing)

  pair = [specd[i] for i in pairing]

  #a-> b :r  if l->l ignore . Therefore ignore and cancel in normal eqnum
  p(pair)

  push!(ppair, "\n! LMP$(pnum): "*join(["$(lumpspec[i]) -> $i" for i in pairing], ", ")*"\n")

  for n in 1:length(pair)
    i = pair[n]
    temp = "0"
    for j in parr[i]
      s = ""
      self = false
      for k in j[1][1]
        if k != "1"
          s = s * k *"*"
        end
      end
      for k in j[2][1]
        #if k in pair
          #self = true
          #break
        #end
        try
          s=s * "C(ind_LMP$(lumpdict[rspec[k]]))*LC_$(lumpspec[rspec[k]])_*"
        catch
          s=s * "C(ind_"* rspec[k] * ")*"
        end
      end
      s=s * r_rated[j[3]]
      #if self == false
        temp = temp*"+"*s*""
      #end
    end

    thisspec = rspec[i]
    thisno = lumpspec[thisspec]

    #temp = replace(tostr(sympy.N(sympy.expand(replace(temp,r"(\d)D([\+\-\d])",s"\1E\2")),3)),"+","&\n      +")
    temp = replace(temp,"+","&\n      +")

    temp = replace(temp,r"EXP\(([-\d\.]+)/TEMP\)",s"ET(\1)")
    temp = replace(temp,r"LC_(\d)_",s"LC(\1)")
    push!(ppair, "PA($(thisno))="*temp)

    push!(ppair, "LC($(thisno)) =  PA($(thisno))/("*join(["PA($(lumpspec[rspec[e]]))" for e in pair],"+")*")")

  end

  @eval @everywhere pgroup=$pairing

  eqns = pmap(replace_match,eqns)

end

include("writelump.jl")

x= 1


#!/usr/bin/env julia

using Cooler

coolName = "MB286.mcool"
startCoord=1000
endCoord = 10000000
resolution = "resolutions/1000"

matrix = Cooler.matrix(coolName,startCoord,endCoord,resolution)
println(length(matrix))


function f_min(iter, factor)
    return  (iter * factor) -(factor-1)
end

function f_max(iter, factor)
    return (iter * factor)
end

function projection()

    println("Predefined vars for projections:")
    println("------------------------------------------------")
    println("=====================[gas]:=====================")
    println("       -all the non derived hydro vars-")
    println(":cpu, :level, :rho, :cx, :cy, :cz, :vx, :vy, :vz, :p, var6,...")
    println("further possibilities: :rho, :density, :ρ")
    println("              -derived hydro vars-")
    println(":x, :y, :z")
    println(":sd or :Σ or :surfacedensity")
    println(":mass, :cellsize, :freefall_time")
    println(":cs, :mach, :jeanslength, :jeansnumber")
    println()
    println("==================[particles]:==================")
    println("        all the non derived  vars:")
    println(":cpu, :level, :id, :family, :tag ")
    println(":x, :y, :z, :vx, :vy, :vz, :mass, :birth, :metal....")
    println()
    println("              -derived particle vars-")
    println(":age")
    println()
    println("==============[gas or particles]:===============")
    println(":v, :ekin")
    println("squared => :vx2, :vy2, :vz2")
    println("velocity dispersion => σx, σy, σz, σ")
    println()
    println("related to a given center:")
    println("---------------------------")
    println(":vr_cylinder, vr_sphere (radial components)")
    println(":vϕ_cylinder, :vθ")
    println("squared => :vr_cylinder2, :vϕ_cylinder2")
    println("velocity dispersion => σr_cylinder, σϕ_cylinder ")
    #println(":l, :lx, :ly, :lz :lr, :lϕ, :lθ")
    println()
    println("2d maps (not projected):")
    println(":r_cylinder") #, :r_sphere")
    println(":ϕ") # :θ
    println("------------------------------------------------")
    println()
    return
end

"""
#### Remap projected data of DataMapsType onto a coarser grid:
- select the level of the coarser grid
- select the weighting
- toggle verbose mode

```julia
remap(dataobject::DataMapsType, lmax::Real; weighting::Symbol=:volume, verbose::Bool=verbose_mode)

return DataMapsType
```
#### Arguments
##### Required:
- **`dataobject`:** needs to be of type: "DataMapsType", e.g. created by the function *projection*
- **`lmax`:** the level of the coarser grid to be created
##### Predefined/Optional Keywords:
- **`weighting`:** choose between: :volume, :mass
- **`verbose`:** print the dimensions/pixel size of the provided and created maps on screen; default: set by the variable `verbose_mode`
"""
function remap(dataobject::DataMapsType, lmax::Real; weighting::Symbol=:volume, verbose::Bool=verbose_mode)

    # checks to use maps instead of mass weighting
    rcheck = [:r_cylinder, :r_sphere]
    anglecheck = [:ϕ]
    ranglecheck = [rcheck..., anglecheck...]

    mcheck = [:sd, :Σ, :surfacedensity,:density, :rho, :ρ]
    mvar = :no
    if weighting == :mass
        notonly_ranglecheck_vars = checkformaps(dataobject, ranglecheck)
        if notonly_ranglecheck_vars
            no_mass_map = false
            for ivar in keys(dataobject.maps)
                if in(ivar, mcheck)
                    mvar = ivar
                    no_mass_map = false
                    break
                else
                    no_mass_map = true
                end
            end
            if no_mass_map
                error("""[Mera]: For mass weighting a map with "rho" or "sd" is necessary.""")
            end
        end
    end

    simlmax = dataobject.lmax
    lmax_projected = lmax

    maps = dataobject.maps
    selected_vars= Symbol[]
    for k in keys(maps)
        push!(selected_vars, k)
    end

    maps_unit = dataobject.maps_unit
    units= Symbol[]
    for k in selected_vars
        kunit = maps_unit[ Symbol(string(k) ) ]
        push!(units, kunit)
    end
    σcheck = [:σx, :σy, :σz, :σ, :σr_cylinder, :σϕ_cylinder]
    σ_to_v = SortedDict(  :σx => [:vx, :vx2],
                          :σy => [:vy, :vy2],
                          :σz => [:vz, :vz2],
                          :σ  => [:v,  :v2],
                          :σr_cylinder => [:vr_cylinder, :vr_cylinder2],
                          :σϕ_cylinder => [:vϕ_cylinder, :vϕ_cylinder2] )

    wcheck = [:r_cylinder, :r_sphere, :ϕ, :sd, :Σ, :surfacedensity,:density, :rho, :ρ] # 2d maps (no projections)

    maps_lmax = SortedDict()
    # remap onto lmax grid
    if simlmax > lmax
        if verbose
            println()
            println("remap from:")
            println("level ", simlmax, " => ", lmax)

            min_cellsize, min_unit  = humanize(dataobject.info.boxlen / 2^lmax, dataobject.info.scale, 2, "length")
            max_cellsize, max_unit  = humanize(dataobject.info.boxlen / 2^simlmax, dataobject.info.scale, 2, "length")
            println("cellsize ", max_cellsize," [$max_unit] => ", min_cellsize ," [$min_unit]")
        end

        first_time  =1
        for ivar in selected_vars
            if !in(ivar, σcheck)
                if mvar == :no || in(ivar, wcheck)
                    maps_buffer = maps[Symbol(ivar)]
                else # use mass-weighting
                    maps_buffer = maps[Symbol(ivar)] .* maps[Symbol(mvar)]
                end

                s = size(maps_buffer)
                lmax_ratio = 2^simlmax / 2^lmax
                scaled_length1 = round(Int, s[1] / lmax_ratio)
                scaled_length2 = round(Int, s[2] / lmax_ratio)
                if first_time == 1
                    if verbose println("pixels ", s, " => ($scaled_length1, $scaled_length2)" ) end
                    first_time = 0
                end
                maps_lmax[Symbol(ivar)] = [maps_buffer[floor(Int,x),floor(Int,y)]  for x in range(1, stop=s[1], length=scaled_length1), y in range(1, stop=s[2], length=scaled_length2)]
            end

        end
    else
        error("lmax of simulation is =< given lmax.")
    end

    if mvar != :no # reverse mass-weighting
        for ivar in selected_vars
            if  !in(ivar, wcheck) && !in(ivar, σcheck) # exclude 2d maps (no projections) and velocity dispersion maps
                maps_lmax[Symbol(ivar)] = maps_lmax[Symbol(ivar)] ./ maps_lmax[Symbol(mvar)]
            end
        end
    end

    # create velocity dispersion maps from rebinned maps
    if simlmax > lmax
        for ivar in selected_vars

            if in(ivar, σcheck)
                    selected_unit, unit_name= getunit(dataobject, ivar, selected_vars, units, uname=true)

                        selected_v = σ_to_v[ivar]
                        iv  = maps_lmax[selected_v[1]]
                        iv_unit = maps_unit[Symbol( string(selected_v[1])  )]
                        iv2 = maps_lmax[selected_v[2]]
                        iv2_unit = maps_unit[Symbol( string(selected_v[2])  )]
                        if iv_unit == iv2_unit
                            diff_iv = iv2 .- iv .^2
                            if typeof(dataobject) == PartMapsType
                                diff_iv[isnan.(diff_iv)] .= 0
                                diff_iv[diff_iv .< 0.] .= 0
                            end

                            if iv_unit == unit_name
                                maps_lmax[Symbol(ivar)] = sqrt.( diff_iv )
                            elseif iv_unit == :standard
                                maps_lmax[Symbol(ivar)] = sqrt.( diff_iv )  .* selected_unit
                            elseif iv_unit == :km_s
                                maps_lmax[Symbol(ivar)] = sqrt.( diff_iv )  ./ dataobject.info.scale.km_s
                            end
                        elseif iv_unit != iv2_unit
                            if iv_unit == :km_s && unit_name == :standard
                                iv = iv ./ dataobject.info.scale.km_s
                            elseif iv_unit == :standard && unit_name == :km_s
                                iv = iv .* dataobject.info.scale.km_s
                            end
                            if iv2_unit == :km_s && unit_name == :standard
                                iv2 = iv2 ./ dataobject.info.scale.km_s.^2
                            elseif iv2_unit == :standard && unit_name == :km_s
                                iv2 = iv2 .* dataobject.info.scale.km_s.^2
                            end

                            diff_iv = iv2 .- iv .^2
                            if typeof(dataobject) == PartMapsType
                                diff_iv[isnan.(diff_iv)] .= 0
                                diff_iv[diff_iv .< 0.] .= 0
                            end

                            maps_lmax[Symbol(ivar)]  = sqrt.( diff_iv )
                        end

            end
        end
    end

    if typeof(dataobject) == HydroMapsType
        return HydroMapsType(maps, maps_unit, maps_lmax, SortedDict(),
                        lmax_projected,
                        dataobject.lmin,
                        simlmax,
                        dataobject.ranges,
                        dataobject.extent,
                        dataobject.cextent,
                        dataobject.ratio,
                        dataobject.boxlen,
                        dataobject.smallr, dataobject.smallc, dataobject.scale,
                        dataobject.info)
    elseif typeof(dataobject) == PartMapsType
        return PartMapsType(maps, maps_unit, maps_lmax, SortedDict(),
                            lmax_projected,
                            dataobject.lmin,
                            simlmax,
                            dataobject.ref_time,
                            dataobject.ranges,
                            dataobject.extent,
                            dataobject.cextent,
                            dataobject.ratio,
                            dataobject.boxlen,
                            dataobject.scale,
                            dataobject.info)
    end

end

# check if only variables from ranglecheck are selected
function checkformaps(selected_vars::Array{Symbol,1}, reference_vars::Array{Symbol,1})
    Nvars = length(selected_vars)
    cw = 0
    for iw in selected_vars
        if in(iw,reference_vars)
            cw +=1
        end
    end
    Ndiff = Nvars-cw
    return Ndiff != 0
end

function checkformaps(dataobject::DataMapsType, reference_vars::Array{Symbol,1})
    Nvars =0
    cw = 0
    for iw in keys(dataobject.maps)
        Nvars +=1
        if in(iw,reference_vars)
            cw +=1
        end
    end
    Ndiff = Nvars-cw
    return Ndiff != 0
end


include("TwoPartCost.jl")
include("ThreePartCost.jl")
include("MarketBidCost.jl")
include("MultiStartCost.jl")
include("Area.jl")
include("LoadZone.jl")
include("Bus.jl")
include("Arc.jl")
include("Line.jl")
include("MonitoredLine.jl")
include("PhaseShiftingTransformer.jl")
include("TapTransformer.jl")
include("Transformer2W.jl")
include("HVDCLine.jl")
include("VSCDCLine.jl")
include("InterruptibleLoad.jl")
include("FixedAdmittance.jl")
include("PowerLoad.jl")
include("HydroEnergyReservoir.jl")
include("HydroDispatch.jl")
include("HydroPumpedStorage.jl")
include("RenewableDispatch.jl")
include("RenewableFix.jl")
include("ThermalStandard.jl")
include("ThermalMultiStart.jl")
include("GenericBattery.jl")
include("BatteryEMS.jl")
include("StaticReserve.jl")
include("StaticReserveNonSpinning.jl")
include("StaticReserveGroup.jl")
include("ReserveDemandCurve.jl")
include("VariableReserve.jl")
include("VariableReserveNonSpinning.jl")
include("AGC.jl")
include("Transfer.jl")
include("AVRFixed.jl")
include("AVRSimple.jl")
include("ESDC1A.jl")
include("ESDC2A.jl")
include("IEEET1.jl")
include("AVRTypeI.jl")
include("AVRTypeII.jl")
include("SCRX.jl")
include("ESAC1A.jl")
include("EXAC1A.jl")
include("EXAC1.jl")
include("EXAC2.jl")
include("ESAC6A.jl")
include("ESST1A.jl")
include("EXPIC1.jl")
include("ESST4B.jl")
include("BaseMachine.jl")
include("RoundRotorMachine.jl")
include("SalientPoleMachine.jl")
include("AndersonFouadMachine.jl")
include("FullMachine.jl")
include("MarconatoMachine.jl")
include("OneDOneQMachine.jl")
include("SimpleAFMachine.jl")
include("SimpleFullMachine.jl")
include("SimpleMarconatoMachine.jl")
include("PSSFixed.jl")
include("PSSSimple.jl")
include("IEEEST.jl")
include("SingleMass.jl")
include("FiveMassShaft.jl")
include("TGFixed.jl")
include("GasTG.jl")
include("GeneralGovModel.jl")
include("SteamTurbineGov1.jl")
include("HydroTurbineGov.jl")
include("IEEETurbineGov1.jl")
include("TGTypeI.jl")
include("TGTypeII.jl")
include("AverageConverter.jl")
include("FixedDCSource.jl")
include("ZeroOrderBESS.jl")
include("LCLFilter.jl")
include("LCFilter.jl")
include("KauraPLL.jl")
include("VirtualInertia.jl")
include("ReactivePowerDroop.jl")
include("CurrentControl.jl")
include("Source.jl")

export get_A1
export get_A2
export get_A3
export get_A4
export get_A5
export get_A6
export get_AT
export get_A_set
export get_Ae
export get_At
export get_Be
export get_D
export get_DB_h
export get_DB_l
export get_D_12
export get_D_23
export get_D_34
export get_D_45
export get_D_T
export get_D_ex
export get_D_hp
export get_D_ip
export get_D_lp
export get_D_turb
export get_Dm
export get_E_sat
export get_Efd_lim
export get_H
export get_H_ex
export get_H_hp
export get_H_ip
export get_H_lp
export get_I_lr
export get_K
export get_K0
export get_K1
export get_K2
export get_K3
export get_K4
export get_K5
export get_K6
export get_K7
export get_K8
export get_K_d
export get_K_ex
export get_K_hp
export get_K_i
export get_K_im
export get_K_ip
export get_K_ir
export get_K_lp
export get_K_lr
export get_K_p
export get_K_pm
export get_K_pr
export get_K_turb
export get_K_ω
export get_Ka
export get_Kb
export get_Kc
export get_Kd
export get_Kd_gov
export get_Ke
export get_Kf
export get_Kg
export get_Kh
export get_Ki
export get_Ki_gov
export get_Ki_load
export get_Ki_mw
export get_Kl
export get_Kp
export get_Kp_gov
export get_Kp_load
export get_Ks
export get_Kt
export get_Kv
export get_L_1d
export get_L_1q
export get_L_ad
export get_L_aq
export get_L_d
export get_L_f1d
export get_L_ff
export get_L_q
export get_Ld_ref
export get_Ls_lim
export get_PSS_flags
export get_P_ref
export get_R
export get_R_1d
export get_R_1q
export get_R_close
export get_R_f
export get_R_lim
export get_R_open
export get_R_th
export get_Rselect
export get_Se
export get_T1
export get_T2
export get_T3
export get_T4
export get_T5
export get_T6
export get_T7
export get_T_AA
export get_T_act
export get_T_eng
export get_T_rate
export get_Ta
export get_Ta_2
export get_Ta_3
export get_Ta_4
export get_Ta_Tb
export get_Tb
export get_Tb1
export get_Tc
export get_Tc1
export get_Td0_p
export get_Td0_pp
export get_Td_gov
export get_Te
export get_Tf
export get_Tf_1
export get_Tf_2
export get_Tf_load
export get_Tg
export get_Th
export get_Tj
export get_Tk
export get_Tpelec
export get_Tq0_p
export get_Tq0_pp
export get_Tr
export get_Ts
export get_Tsa
export get_Tsb
export get_Tw
export get_U0
export get_UEL_flags
export get_U_c
export get_VB_max
export get_VELM
export get_VFE_lim
export get_VH_max
export get_V_lim
export get_V_lr
export get_V_pss
export get_V_ref
export get_Va_lim
export get_Vcl
export get_Vcu
export get_Vdc_ref
export get_Vf
export get_Vi_lim
export get_Vm_lim
export get_Vr_lim
export get_Wf_nl
export get_X_th
export get_Xd
export get_Xd_p
export get_Xd_pp
export get_Xl
export get_Xq
export get_Xq_p
export get_Xq_pp
export get_Y
export get_active_power
export get_active_power_flow
export get_active_power_limits
export get_active_power_limits_from
export get_active_power_limits_pump
export get_active_power_limits_to
export get_ancillary_services
export get_angle
export get_angle_limits
export get_arc
export get_area
export get_available
export get_b
export get_base_power
export get_base_voltage
export get_battery_resistance
export get_battery_voltage
export get_bias
export get_bus
export get_bustype
export get_cf
export get_contributing_services
export get_conversion_factor
export get_db
export get_dc_dc_inductor
export get_dc_link_capacitance
export get_delta_t
export get_dynamic_injector
export get_efficiency
export get_energy_value
export get_eq_p
export get_ext
export get_fixed
export get_flow_limits
export get_from
export get_fs
export get_fuel
export get_fuel_flag
export get_gate_position_limits
export get_inflow
export get_initial_ace
export get_initial_energy
export get_initial_storage
export get_input_active_power_limits
export get_input_code
export get_internal_angle
export get_internal_voltage
export get_inv_d_fluxlink
export get_inv_q_fluxlink
export get_inverter_firing_angle
export get_inverter_tap_limits
export get_inverter_xrc
export get_kad
export get_kd
export get_kffi
export get_kffv
export get_ki_pll
export get_kic
export get_kii
export get_kiv
export get_kp_pll
export get_kpc
export get_kpi
export get_kpv
export get_kq
export get_kω
export get_lf
export get_lg
export get_load_response
export get_load_zone
export get_loss
export get_lv
export get_magnitude
export get_max_active_power
export get_max_reactive_power
export get_model
export get_must_run
export get_n_states
export get_no_load
export get_number
export get_operation_cost
export get_outflow
export get_output_active_power_limits
export get_peak_active_power
export get_peak_reactive_power
export get_penalty_cost
export get_power_factor
export get_power_trajectory
export get_primary_shunt
export get_prime_mover
export get_pump_efficiency
export get_q_nl
export get_r
export get_ramp_limits
export get_ramp_limits_pump
export get_rate
export get_rated_current
export get_rated_voltage
export get_rating
export get_rating_pump
export get_rc_rfd
export get_reactive_power
export get_reactive_power_flow
export get_reactive_power_limits
export get_reactive_power_limits_from
export get_reactive_power_limits_pump
export get_reactive_power_limits_to
export get_rectifier_firing_angle
export get_rectifier_tap_limits
export get_rectifier_xrc
export get_remote_bus_control
export get_requirement
export get_rf
export get_rg
export get_rv
export get_saturation_coeffs
export get_services
export get_shut_down
export get_speed_error_signal
export get_start_time_limits
export get_start_types
export get_start_up
export get_state_of_charge_limits
export get_states
export get_states_types
export get_status
export get_storage_capacity
export get_storage_target
export get_switch
export get_tap
export get_time_at_status
export get_time_frame
export get_time_limits
export get_time_limits_pump
export get_to
export get_valve_position_limits
export get_variable
export get_voltage
export get_voltage_limits
export get_x
export get_α
export get_γ_d1
export get_γ_d2
export get_γ_q1
export get_γ_q2
export get_γ_qd
export get_γd
export get_γq
export get_θp
export get_θp_rad
export get_τ_limits
export get_ω_lp
export get_ωad
export get_ωf
export set_A1!
export set_A2!
export set_A3!
export set_A4!
export set_A5!
export set_A6!
export set_AT!
export set_A_set!
export set_Ae!
export set_At!
export set_Be!
export set_D!
export set_DB_h!
export set_DB_l!
export set_D_12!
export set_D_23!
export set_D_34!
export set_D_45!
export set_D_T!
export set_D_ex!
export set_D_hp!
export set_D_ip!
export set_D_lp!
export set_D_turb!
export set_Dm!
export set_E_sat!
export set_Efd_lim!
export set_H!
export set_H_ex!
export set_H_hp!
export set_H_ip!
export set_H_lp!
export set_I_lr!
export set_K!
export set_K0!
export set_K1!
export set_K2!
export set_K3!
export set_K4!
export set_K5!
export set_K6!
export set_K7!
export set_K8!
export set_K_d!
export set_K_ex!
export set_K_hp!
export set_K_i!
export set_K_im!
export set_K_ip!
export set_K_ir!
export set_K_lp!
export set_K_lr!
export set_K_p!
export set_K_pm!
export set_K_pr!
export set_K_turb!
export set_K_ω!
export set_Ka!
export set_Kb!
export set_Kc!
export set_Kd!
export set_Kd_gov!
export set_Ke!
export set_Kf!
export set_Kg!
export set_Kh!
export set_Ki!
export set_Ki_gov!
export set_Ki_load!
export set_Ki_mw!
export set_Kl!
export set_Kp!
export set_Kp_gov!
export set_Kp_load!
export set_Ks!
export set_Kt!
export set_Kv!
export set_L_1d!
export set_L_1q!
export set_L_ad!
export set_L_aq!
export set_L_d!
export set_L_f1d!
export set_L_ff!
export set_L_q!
export set_Ld_ref!
export set_Ls_lim!
export set_PSS_flags!
export set_P_ref!
export set_R!
export set_R_1d!
export set_R_1q!
export set_R_close!
export set_R_f!
export set_R_lim!
export set_R_open!
export set_R_th!
export set_Rselect!
export set_Se!
export set_T1!
export set_T2!
export set_T3!
export set_T4!
export set_T5!
export set_T6!
export set_T7!
export set_T_AA!
export set_T_act!
export set_T_eng!
export set_T_rate!
export set_Ta!
export set_Ta_2!
export set_Ta_3!
export set_Ta_4!
export set_Ta_Tb!
export set_Tb!
export set_Tb1!
export set_Tc!
export set_Tc1!
export set_Td0_p!
export set_Td0_pp!
export set_Td_gov!
export set_Te!
export set_Tf!
export set_Tf_1!
export set_Tf_2!
export set_Tf_load!
export set_Tg!
export set_Th!
export set_Tj!
export set_Tk!
export set_Tpelec!
export set_Tq0_p!
export set_Tq0_pp!
export set_Tr!
export set_Ts!
export set_Tsa!
export set_Tsb!
export set_Tw!
export set_U0!
export set_UEL_flags!
export set_U_c!
export set_VB_max!
export set_VELM!
export set_VFE_lim!
export set_VH_max!
export set_V_lim!
export set_V_lr!
export set_V_pss!
export set_V_ref!
export set_Va_lim!
export set_Vcl!
export set_Vcu!
export set_Vdc_ref!
export set_Vf!
export set_Vi_lim!
export set_Vm_lim!
export set_Vr_lim!
export set_Wf_nl!
export set_X_th!
export set_Xd!
export set_Xd_p!
export set_Xd_pp!
export set_Xl!
export set_Xq!
export set_Xq_p!
export set_Xq_pp!
export set_Y!
export set_active_power!
export set_active_power_flow!
export set_active_power_limits!
export set_active_power_limits_from!
export set_active_power_limits_pump!
export set_active_power_limits_to!
export set_ancillary_services!
export set_angle!
export set_angle_limits!
export set_arc!
export set_area!
export set_available!
export set_b!
export set_base_power!
export set_base_voltage!
export set_battery_resistance!
export set_battery_voltage!
export set_bias!
export set_bus!
export set_bustype!
export set_cf!
export set_contributing_services!
export set_conversion_factor!
export set_db!
export set_dc_dc_inductor!
export set_dc_link_capacitance!
export set_delta_t!
export set_dynamic_injector!
export set_efficiency!
export set_energy_value!
export set_eq_p!
export set_ext!
export set_fixed!
export set_flow_limits!
export set_from!
export set_fs!
export set_fuel!
export set_fuel_flag!
export set_gate_position_limits!
export set_inflow!
export set_initial_ace!
export set_initial_energy!
export set_initial_storage!
export set_input_active_power_limits!
export set_input_code!
export set_internal_angle!
export set_internal_voltage!
export set_inv_d_fluxlink!
export set_inv_q_fluxlink!
export set_inverter_firing_angle!
export set_inverter_tap_limits!
export set_inverter_xrc!
export set_kad!
export set_kd!
export set_kffi!
export set_kffv!
export set_ki_pll!
export set_kic!
export set_kii!
export set_kiv!
export set_kp_pll!
export set_kpc!
export set_kpi!
export set_kpv!
export set_kq!
export set_kω!
export set_lf!
export set_lg!
export set_load_response!
export set_load_zone!
export set_loss!
export set_lv!
export set_magnitude!
export set_max_active_power!
export set_max_reactive_power!
export set_model!
export set_must_run!
export set_n_states!
export set_no_load!
export set_number!
export set_operation_cost!
export set_outflow!
export set_output_active_power_limits!
export set_peak_active_power!
export set_peak_reactive_power!
export set_penalty_cost!
export set_power_factor!
export set_power_trajectory!
export set_primary_shunt!
export set_prime_mover!
export set_pump_efficiency!
export set_q_nl!
export set_r!
export set_ramp_limits!
export set_ramp_limits_pump!
export set_rate!
export set_rated_current!
export set_rated_voltage!
export set_rating!
export set_rating_pump!
export set_rc_rfd!
export set_reactive_power!
export set_reactive_power_flow!
export set_reactive_power_limits!
export set_reactive_power_limits_from!
export set_reactive_power_limits_pump!
export set_reactive_power_limits_to!
export set_rectifier_firing_angle!
export set_rectifier_tap_limits!
export set_rectifier_xrc!
export set_remote_bus_control!
export set_requirement!
export set_rf!
export set_rg!
export set_rv!
export set_saturation_coeffs!
export set_services!
export set_shut_down!
export set_speed_error_signal!
export set_start_time_limits!
export set_start_types!
export set_start_up!
export set_state_of_charge_limits!
export set_states!
export set_states_types!
export set_status!
export set_storage_capacity!
export set_storage_target!
export set_switch!
export set_tap!
export set_time_at_status!
export set_time_frame!
export set_time_limits!
export set_time_limits_pump!
export set_to!
export set_valve_position_limits!
export set_variable!
export set_voltage!
export set_voltage_limits!
export set_x!
export set_α!
export set_γ_d1!
export set_γ_d2!
export set_γ_q1!
export set_γ_q2!
export set_γ_qd!
export set_γd!
export set_γq!
export set_θp!
export set_θp_rad!
export set_τ_limits!
export set_ω_lp!
export set_ωad!
export set_ωf!


module TimingDiagrams

using ..Plots: plot, plot!, annotate!, Shape, text
using LinearAlgebra

export timingdiagram, node!, tx!

struct Node
  name::String
  pos::Vector{Float64}
  ypos::Float64
end

const nodes = Node[]
const speed = Ref(1.0)
const duration = Ref(0.0)
const timings = Float64[]

const theme = Dict{Symbol,Any}(
  :txheight => 10.0,
  :fontsize => 8,
  :ypad => 5.0,
  :tpad => 0.01,
  :txcolor => :blue,
  :rxcolor => :green,
  :linecolor => :lightgray,
  :nodecolor => :black
)

speed!(c) = (speed[] = c)
duration!(d) = (duration[] = d)

function timingdiagram(f; kwargs...)
  plt = plot(; border=:none, legend=false, yflip=true, kwargs...)
  empty!(nodes)
  empty!(timings)
  f()
  for n ∈ nodes
    plot!(plt, [minimum(timings) - theme[:tpad], maximum(timings) + theme[:tpad]], [n.ypos, n.ypos]; color=theme[:nodecolor])
    annotate!(plt, -theme[:tpad], n.ypos, text(n.name * " ", theme[:fontsize], :right))
  end
  empty!(nodes)
  empty!(timings)
  plt
end

function node!(name, pos; ypos=norm(pos))
  push!(nodes, Node(name, pos isa Real ? [Float64(pos)] : pos, ypos))
  length(nodes)
end

function tx!(t, from, to=setdiff(eachindex(nodes),from); duration=duration[], txlabel=nothing, txtiminglabel=nothing, rxlabels=nothing, rxtiminglabels=nothing)
  rxtimes = [t + norm(nodes[i].pos - nodes[from].pos) / speed[] for i ∈ to]
  duration > 0.0 && rect!(t, nodes[from].ypos, duration, theme[:txheight]; color=theme[:txcolor])
  txlabel === nothing || annotate!(t + duration/2, nodes[from].ypos - theme[:txheight] - theme[:ypad], text(txlabel, theme[:fontsize], :bottom))
  txtiminglabel === nothing || annotate!(t, nodes[from].ypos + theme[:ypad], text(txtiminglabel, theme[:fontsize], :top))
  for (i, rx) ∈ enumerate(to)
    plot!([t, rxtimes[i]], [nodes[from].ypos, nodes[rx].ypos]; color=theme[:linecolor])
    duration > 0.0 && rect!(rxtimes[i], nodes[rx].ypos, duration, theme[:txheight]; color=theme[:rxcolor])
    rxlabels === nothing || annotate!(rxtimes[i] + duration/2, nodes[rx].ypos - theme[:txheight] - theme[:ypad], text(rxlabels[i], theme[:fontsize], :bottom))
    rxtiminglabels === nothing || annotate!(rxtimes[i], nodes[rx].ypos + theme[:ypad], text(rxtiminglabels[i], theme[:fontsize], :top))
    end
  push!(timings, t)
  push!(timings, minimum(rxtimes))
  push!(timings, maximum(rxtimes) + duration)
  rxtimes
end

function rect!(t, y, dt, dy; color)
  plot!(Shape([t, t+dt, t+dt, t], [y, y, y-dy, y-dy]); color)
end

end # module


"""
Subtype of `LDScoreRegression`.
"""

mutable struct Hsq <: LDScoreRegression
    y; x; w; N; M;
    n_blocks; intercept; slow; step1_ii; old_weights;
    __null_intercept__

    constrain_intercept; n_annot; intercept_se; twostep_filtered;

    function Hsq(
        y, x, w, N, M; n_blocks = 200, slow = false, old_weights = false,
        intercept = nothing, step1_ii = nothing,
    )
        hsq = new(
            y, x, w, N, M, n_blocks, intercept, slow, step1_ii, old_weights, 1,
        )
        return ld_score_regression(
            hsq,
            y, x, w, N, M, n_blocks, intercept, slow, step1_ii, old_weights,
        )
    end
end

function hsq_weights(ld, w_ld, N, M, hsq; intercept=nothing)
    if intercept == nothing intercept = 1.0 end
    M = M[1][1]
    hsq = max(maximum(hsq), 0.0)
    hsq = min(hsq, 1.0)
    ld = fmax(ld, 1.0)
    w_ld = fmax(w_ld, 1.0)
    c = hsq .* N ./ M

    het_w = 1.0 ./ (2 .* (c .* ld .+ intercept).^2)
    oc_w = 1.0 ./ w_ld
    w = het_w .* oc_w
    return w
end

function update_weights(reg::Hsq, ld, w_ld, N, M, hsq, intercept)
    if intercept == nothing intercept = reg.__null_intercept__ end

    return hsq_weights(ld, w_ld, N, M, hsq; intercept=intercept)
end

function overlap_output(
    hsq::Hsq,
    category_names, overlap_matrix, M_annot, M_tot, print_coefficients,
)
    overlap_matrix_prop = zeros([hsq.n_annot, hsq.n_annot])
    for i in range(hsq.n_annot)
        overlap_matrix_prop[i, :] = overlap_matrix[i, :] / M_annot
    end

    # TODO missing code here

    #=
    df = DataFrames.DataFrame({
        "Category": category_names,
        "Prop._SNPs": one_d_convert(prop_M_overlap),
        "Prop._h2": one_d_convert(prop_hsq_overlap),
        "Prop._h2_std_error": one_d_convert(prop_hsq_overlap_se),
        "Enrichment": one_d_convert(enrichment),
        "Enrichment_std_error": one_d_convert(enrichment_se),
        "Enrichment_p":diff_p,
        "Coefficient": one_d_convert(hsq.coef),
        "Coefficient_std_error": hsq.coef_se,
        "Coefficient_z-score": one_d_convert(hsq.coef) / one_d_convert(hsq.coef_se)
    })
    if print_coefficients
        return df[
            !,
            [
                "Category", "Prop._SNPs", "Prop._h2", "Prop._h2_std_error",
                "Enrichment","Enrichment_std_error", "Enrichment_p",
                "Coefficient", "Coefficient_std_error","Coefficient_z-score",
            ],
        ]
    else
        return df[
            !,
            [
                "Category", "Prop._SNPs", "Prop._h2", "Prop._h2_std_error",
                "Enrichment","Enrichment_std_error", "Enrichment_p",
            ],
        ]
    end
    =#
end


#### thermodynamics

export new_thermo_state_anelastic, recover_thermo_state_anelastic

"""
    new_thermo_state_anelastic(atmos::AtmosModel, state::Vars, aux::Vars)

Create a new thermodynamic state, based on the `state`, and _not_
the `aux` state.

!!! note
    This method calls the iterative saturation adjustment
    procedure for EquilMoist models.
"""
new_thermo_state_anelastic(atmos::AtmosModel, state::Vars, aux::Vars) =
    new_thermo_state_anelastic(
        atmos,
        energy_model(atmos),
        moisture_model(atmos),
        state,
        aux,
    )

"""
    recover_thermo_state_anelastic(atmos::AtmosModel, state::Vars, aux::Vars)

An atmospheric thermodynamic state.

!!! warn
    For now, we are directly calling new_thermo_state_anelastic to avoid
    inconsistent aux states in kernels where the aux states are
    out of sync with the boundary state.

# TODO: Define/call `recover_thermo_state_anelastic` when it's safely implemented
  (see https://github.com/CliMA/ClimateMachine.jl/issues/1648)
"""
recover_thermo_state_anelastic(atmos::AtmosModel, state::Vars, aux::Vars) =
    new_thermo_state_anelastic(
        atmos,
        energy_model(atmos),
        moisture_model(atmos),
        state,
        aux,
    )

function new_thermo_state_anelastic(
    atmos::AtmosModel,
    energy::TotalEnergyModel,
    moist::DryModel,
    state::Vars,
    aux::Vars,
)
    param_set = parameter_set(atmos)
    e_int = internal_energy(atmos, state, aux)
    p = aux.ref_state.p
    return PhaseDry_pe(param_set, p, e_int)
end

function new_thermo_state_anelastic(
    atmos::AtmosModel,
    energy::TotalEnergyModel,
    moist::EquilMoist,
    state::Vars,
    aux::Vars,
)
    param_set = parameter_set(atmos)
    e_int = internal_energy(atmos, state, aux)
    p = aux.ref_state.p
    ρ = density(atmos, state, aux)
    return PhaseEquil_peq(
        param_set,
        p,
        e_int,
        state.moisture.ρq_tot / ρ,
        moist.maxiter,
        moist.tolerance,
    )
end

function new_thermo_state_anelastic(
    atmos::AtmosModel,
    energy::TotalEnergyModel,
    moist::NonEquilMoist,
    state::Vars,
    aux::Vars,
)
    param_set = parameter_set(atmos)
    e_int = internal_energy(atmos, state, aux)
    p = aux.ref_state.p
    ρ = density(atmos, state, aux)
    q = PhasePartition(
        state.moisture.ρq_tot / ρ,
        state.moisture.ρq_liq / ρ,
        state.moisture.ρq_ice / ρ,
    )

    return PhaseNonEquil_peq(param_set, p, e_int, q)
end


export rand_controller

function rand_controller(cfg::Dict)
    Random.seed!(cfg["seed"])
    nouts = output_lengths(cfg)
    cell_division(x::Array{Float64}) = rand(Bool)
    new_cell_params(x::Array{Float64}) = rand(nouts[2])
    cell_death(x::Array{Float64}) = rand(Bool)
    cell_state_update(x::Array{Float64}) = rand(nouts[4])
    cell_param_update(x::Array{Float64}) = rand(nouts[5])
    connect(x::Array{Float64}) = rand(Bool)
    new_conn_params(x::Array{Float64}) = rand(nouts[7])
    disconnect(x::Array{Float64}) = rand(Bool)
    conn_state_update(x::Array{Float64}) = rand(nouts[9])
    conn_param_update(x::Array{Float64}) = rand(nouts[10])

    Controller(cell_division, new_cell_params, cell_death,
               cell_state_update, cell_param_update,
               connect, new_conn_params, disconnect,
               conn_state_update, conn_param_update)
end


"""
Struct for Active Brownian particle system
"""
mutable struct ActiveBrownianSystem <: AbstractActiveSystem
  dim::Int64
  v0::Float64
  t::Float64
  x::Vector{Float64}
  f::Vector{Float64}
  thermostat::ActiveLangevin
  function ActiveBrownianSystem(dim::Int64,v0::Float64,t::Float64,x::Vector{Float64},f::Vector{Float64},thermostat::ActiveLangevin)
    new(dim,v0,t,x,f,thermostat)
  end
end

function ActiveBrownianSystem(v0::Float64,model::AbstractPotential,thermostat::ActiveLangevin)
  dim::Int64 = getdimensionality(model)
  x::Vector{Float64} = zeros(dim+1)
  f::Vector{Float64} = zeros(dim+1)
  ActiveBrownianSystem(dim+1,v0,0.0,x,f,thermostat)
end

# TODO idk yet, but this needs to be fixed to allow the different gammas
function addforce!(dt::Float64,sys::ActiveBrownianSystem,model::AbstractSinglePotential)
  @. @views sys.x[1:sys.dim-1] .+= dt .* model.f[1:sys.dim-1] ./ sys.thermostat.gammat
  if model.dim == sys.dim
    sys.x[end] += dt .* model.f[end] ./ sys.thermostat.gammar
  end
end


module DRAW

using Knet, Sloth
import Knet: minibatch, params, train!
using Statistics, Random, Dates
import Base: push!, empty!

using ArgParse
using Images, ImageMagick
using JLD2
using MosaicViews

include(Knet.dir("data","mnist.jl"))
_etype = Float32
_atype = gpu() >= 0 ? KnetArray{_etype} : Array{_etype}
SEED = -1 # FIXME: find a better way

include("utils.jl")
include("layers.jl")
include("network.jl")
include("visualize.jl")
include("train.jl")

end # module