#=
Copyright (c) 2015, Intel Corporation
All rights reserved.

Redistribution and use in source and binary forms, with or without
modification, are permitted provided that the following conditions are met:
- Redistributions of source code must retain the above copyright notice,
  this list of conditions and the following disclaimer.
- Redistributions in binary form must reproduce the above copyright notice,
  this list of conditions and the following disclaimer in the documentation
  and/or other materials provided with the distribution.

THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF
THE POSSIBILITY OF SUCH DAMAGE.
=#

using CompilerTools
using CompilerTools.LambdaHandling
using Base.Test

## Tests for CompilerTools.AliasAnalysis

function test_alias_2(x::Int, y::Int, z::Int, s::Int)
     A = rand(x, y)
     D_arr = rand(x, y)
     C = rand(x,y)
     B = D_arr

     for i = 1:x
          if( i%2 == 0 )
              for j = 1:y
              D_arr[i,j] /= 2
          end
          end
      end

      E = B
      if( s > 0 )
          E = A
      end

      return E

end

ast = CompilerTools.Helper.LambdaInfo(test_alias_2, (Int,Int,Int,Int), code_typed(test_alias_2, (Int,Int,Int,Int))[1])
linfo, body = lambdaToLambdaVarInfo(ast)
#cfg_2 = CompilerTools.CFGs.from_lambda(ast) :: CompilerTools.CFGs.CFG

#CompilerTools.AliasAnalysis.set_debug_level(3)
lives = CompilerTools.LivenessAnalysis.from_lambda(linfo, body)
handled = CompilerTools.AliasAnalysis.from_lambda(linfo, body, lives)
handled = map(x -> lookupVariableName(x, linfo), handled)
#println(handled)
#@test (in(:C, handled))
@test (in(:E, handled) == false)
@test (in(:B, handled) == false)
#@test (in(:D_arr, handled) == false)
#@test (in(:A, handled) == false)


### JS_SAA main file
module JS

using Distributions, StatsBase, LinearAlgebra

include("genPurpose.jl")
include("babyNewsvendor.jl")

end #ends module


module Ops

using Compat
using LinearAlgebra
using Zeros
using ..LazyApply

using ..LazyApply: ∑, ∏, α, coefficient, Applied

import ..LazyApply: coefficient

import Base: show, *, /, +, -, Tuple, one, exp, adjoint, promote_rule, convert

export Op, sites, params

#####################################################################################
# General functionality
#

# Helper function to split a `Tuple` according to the function `f`.
# For example:
#
# julia> t = (1, "X", 1, 2, "Y", 2, "Z", 4)
# (1, "X", 1, 2, "Y", 2, "Z", 4)
#
# julia> split(x -> x isa AbstractString, t)
# ((1,), ("X", 1, 2), ("Y", 2), ("Z", 4))
#
function split(f, t::Tuple)
  n = findall(f, t)
  ti = t[1:(first(n) - 1)]
  ts = ntuple(i -> t[n[i]:(n[i + 1] - 1)], length(n) - 1)
  tf = t[last(n):end]
  return ti, ts..., tf
end

#
# General functionality
#####################################################################################

#####################################################################################
# LazyApply extensions
# TODO: Move to `LazyApply`
#

# Helper function for determing the cofficient type of an `Op` related type.
coefficient_type(o::Type) = One
coefficient_type(o::Type{<:α{<:Any,T}}) where {T} = T
coefficient_type(o::Type{<:∑{T}}) where {T} = coefficient_type(T)

coefficient_type(o::Applied) = coefficient_type(typeof(o))

#
# LazyApply extensions
#####################################################################################

const WhichOp = Union{String,AbstractMatrix,UniformScaling}

struct Op
  which_op::WhichOp
  sites::Tuple{Vararg{Int}}
  params::NamedTuple
end
which_op(o::Op) = o.which_op
sites(o::Op) = o.sites
params(o::Op) = o.params
op(o::Op) = o
coefficient_type(o::Op) = One
coefficient(o::Op) = one(coefficient_type(o))

params(o::Vector{Op}) = params(only(o))

# exp
exp(o::Op) = Applied(exp, o)

# adjoint
adjoint(o::Op) = Applied(adjoint, o)

Tuple(o::Op) = (which_op(o), sites(o), params(o))

const OpExpr = Union{Op,∑{Op},α{Op},∑{<:α{Op}},∏{Op},∑{∏{Op}},α{∏{Op}},∑{<:α{∏{Op}}}}

# Type promotion and conversion
convert(::Type{α{Op,T}}, o::Op) where {T} = one(T) * o
convert(::Type{∏{Op}}, o::Op) = ∏([o])
convert(::Type{∑{Op}}, o::Op) = ∑([o])
convert(::Type{α{∏{Op},T}}, o::Op) where {T} = one(T) * convert(∏{Op}, o)
function convert(::Type{∑{α{∏{Op},T}}}, o::Op) where {T}
  return ∑([convert(α{∏{Op},T}, o)])
end

# if VERSION < v"1.6"
const ∑α∏Op = ∑{α{∏{Op},T}} where {T}
const ∑α∏{S} = ∑{α{∏{S},T}} where {T}

convert(O::Type{<:Op}, o::Tuple) = O(o)
convert(O::Type{<:α{Op}}, o::Tuple) = convert(O, Op(o))
convert(O::Type{<:∏{Op}}, o::Tuple) = convert(O, Op(o))
convert(O::Type{<:∑{Op}}, o::Tuple) = convert(O, Op(o))
convert(O::Type{<:α{∏{Op}}}, o::Tuple) = convert(O, Op(o))
convert(O::Type{<:∑{α{∏{Op},T}} where {T}}, o::Tuple) = convert(O, Op(o))
## convert(O::Type{∑{α{∏{Op},T}} where {T}}, o::Tuple) = convert(O, Op(o))
convert(O::Type{∑α∏Op}, o::Tuple) = convert(O, Op(o))

convert(::Type{∑{<:α{∏{Op}}}}, o) = convert(∑{α{∏{Op},T}} where {T}, o)
∑{<:α{∏{Op}}}(o) = (∑{α{∏{Op},T}} where {T})(o)
∑{<:α{∏{Op}}}() = (∑{α{∏{Op},T}} where {T})()

# if VERSION > v"1.5"
## function (∑{α{∏{Op},T}} where {T})(o::OpExpr)
##   return convert(∑{α{∏{Op},T}} where {T}, o)
## end
## function (∑{α{∏{Op},T}} where {T})(o::Tuple)
##   return convert(∑{α{∏{Op},T}} where {T}, o)
## end
## function (∑{α{∏{Op},T}} where {T})(o::Vector{<:Union{OpExpr,Tuple}})
##   return reduce(+, o; init=(∑{α{∏{Op},T}} where {T})())
## end
## function (∑{α{∏{Op},T}} where {T})(o::WhichOp, args...)
##   return convert(∑{α{∏{Op},T}} where {T}, Op(o, args...))
## end
## function (∑{α{∏{Op},T}} where {T})(c::Number, o::WhichOp, args...)
##   return convert(∑{α{∏{Op},T}} where {T}, Op(c, o, args...))
## end
# if VERSION < v"1.6"
function ∑α∏Op(o::OpExpr)
  return convert(∑{α{∏{Op},T}} where {T}, o)
end
function ∑α∏Op(o::Tuple)
  return convert(∑{α{∏{Op},T}} where {T}, o)
end
function ∑α∏Op(o::Vector{<:Union{OpExpr,Tuple}})
  return reduce(+, o; init=(∑{α{∏{Op},T}} where {T})())
end
function ∑α∏Op(o::WhichOp, args...)
  return convert(∑{α{∏{Op},T}} where {T}, Op(o, args...))
end
function ∑α∏Op(c::Number, o::WhichOp, args...)
  return convert(∑{α{∏{Op},T}} where {T}, Op(c, o, args...))
end

# Default constructors
# if VERSION > v"1.5"
## (∑{α{∏{S},T}} where {T})() where {S} = ∑{α{∏{S},Zero}}()
## (∑{α{∏{Op},T}} where {T})(o) = convert(∑{α{∏{Op},T}} where {T}, o)
# if VERSION < v"1.6"
∑α∏{S}() where {S} = ∑{α{∏{S},Zero}}()
∑α∏Op(o) = convert(∑{α{∏{Op},T}} where {T}, o)

function convert(O::Type{α{∏{Op},T}}, o::α{Op}) where {T}
  return convert(T, coefficient(o)) * ∏([op(o)])
end
function convert(O::Type{∑{T}}, o::α{Op}) where {T<:Union{α{Op},α{∏{Op}}}}
  return ∑([convert(T, o)])
end

convert(O::Type{∑{∏{Op}}}, o::∏{Op}) = ∑([o])
convert(O::Type{α{∏{Op},T}}, o::∏{Op}) where {T} = one(T) * o
function convert(O::Type{∑{α{∏{Op},T}}}, o::∏{Op}) where {T}
  return ∑([convert(α{∏{Op},T}, o)])
end
convert(O::Type{∑{α{∏{Op},T}}}, o::∑{∏{Op}}) where {T} = one(T) * o

function convert(O::Type{∑{α{∏{Op},T}}}, o::α{∏{Op}}) where {T}
  return ∑([convert(α{∏{Op},T}, o)])
end

# Versions where the type paramater is left out.
function convert(O::Type{∑{α{∏{Op},T}} where {T}}, o)
  return convert(∑{α{∏{Op},coefficient_type(o)}}, o)
end

#
# Promotion rules.
#
# Rules for promoting Op-like objects when they are being added together.
#
# Should cover promotions between these types:
#
# Op
# α{Op,T}
# ∏{Op}
# ∑{Op}
# ∑{α{Op,T}}
# α{∏{Op},T}}
# ∑{∏{Op}}
# ∑{α{∏{Op},T}}
#

# Conversion of `Op`
promote_rule(::Type{Op}, O::Type{<:α{Op}}) = O
promote_rule(::Type{Op}, O::Type{<:∏{Op}}) = O
promote_rule(::Type{Op}, O::Type{<:∑{Op}}) = O
promote_rule(::Type{Op}, O::Type{<:∑{α{Op}}}) = O
promote_rule(::Type{Op}, O::Type{<:α{∏{Op}}}) = O
promote_rule(::Type{Op}, O::Type{<:∑{∏{Op}}}) = O
promote_rule(::Type{Op}, O::Type{<:∑{α{∏{Op}}}}) = O

# Conversion of `α{Op}`
function promote_rule(::Type{α{Op,T}}, ::Type{α{Op,S}}) where {T,S}
  return α{Op,promote_type(T, S)}
end
promote_rule(::Type{α{Op,T}}, ::Type{∏{Op}}) where {T} = α{∏{Op},T}
promote_rule(::Type{α{Op,T}}, ::Type{∑{Op}}) where {T} = ∑{α{Op,T}}
function promote_rule(::Type{α{Op,T}}, ::Type{∑{α{Op,S}}}) where {T,S}
  return ∑{α{Op,promote_type(T, S)}}
end
function promote_rule(::Type{α{Op,T}}, ::Type{α{∏{Op},S}}) where {T,S}
  return α{∏{Op},promote_type(T, S)}
end
function promote_rule(::Type{α{Op,T}}, ::Type{∑{∏{Op}}}) where {T}
  return ∑{α{∏{Op},T}}
end
function promote_rule(::Type{α{Op,T}}, ::Type{∑{α{∏{Op},S}}}) where {T,S}
  return ∑{α{∏{Op},promote_type(T, S)}}
end

# Conversion of `∏{Op}`
promote_rule(::Type{∏{Op}}, ::Type{∑{Op}}) = ∑{∏{Op}}
function promote_rule(::Type{∏{Op}}, ::Type{∑{α{Op,S}}}) where {S}
  return ∑{α{∏{Op},S}}
end
promote_rule(::Type{∏{Op}}, ::Type{α{∏{Op},S}}) where {S} = α{∏{Op},S}
promote_rule(::Type{∏{Op}}, ::Type{∑{∏{Op}}}) = ∑{∏{Op}}
function promote_rule(::Type{∏{Op}}, ::Type{∑{α{∏{Op},S}}}) where {S}
  return ∑{α{∏{Op},S}}
end

# Conversion of `∑{Op}`
promote_rule(::Type{∑{Op}}, ::Type{∑{α{Op,S}}}) where {S} = ∑{α{Op,S}}
function promote_rule(::Type{∑{Op}}, ::Type{α{∏{Op},S}}) where {S}
  return ∑{α{∏{Op},S}}
end
promote_rule(::Type{∑{Op}}, ::Type{∑{∏{Op}}}) = ∑{∏{Op}}
function promote_rule(::Type{∑{Op}}, ::Type{∑{α{∏{Op},S}}}) where {S}
  return ∑{α{∏{Op},S}}
end

# Conversion of `∑{α{Op,T}}`
function promote_rule(::Type{∑{α{Op,T}}}, ::Type{∑{α{Op,S}}}) where {T,S}
  return ∑{α{Op,promote_type(T, S)}}
end
function promote_rule(::Type{∑{α{Op,T}}}, ::Type{α{∏{Op},S}}) where {T,S}
  return ∑{α{∏{Op},promote_type(T, S)}}
end
function promote_rule(::Type{∑{α{Op,T}}}, ::Type{∑{∏{Op}}}) where {T}
  return ∑{α{∏{Op},T}}
end
function promote_rule(::Type{∑{α{Op,T}}}, ::Type{∑{α{∏{Op},S}}}) where {T,S}
  return ∑{α{∏{Op},promote_type(T, S)}}
end

# Conversion of `α{∏{Op},T}`
function promote_rule(::Type{α{∏{Op},T}}, ::Type{α{∏{Op},S}}) where {T,S}
  return α{∏{Op},promote_type(T, S)}
end
function promote_rule(::Type{α{∏{Op},T}}, ::Type{∑{∏{Op}}}) where {T}
  return ∑{α{∏{Op},T}}
end
function promote_rule(::Type{α{∏{Op},T}}, ::Type{∑{α{∏{Op},S}}}) where {T,S}
  return ∑{α{∏{Op},promote_type(T, S)}}
end

# Conversion of `∑{∏{Op}}`
function promote_rule(::Type{∑{∏{Op}}}, ::Type{∑{α{∏{Op},S}}}) where {S}
  return ∑{α{∏{Op},S}}
end

# Conversion of `∑{α{∏{Op},T}}`
function promote_rule(::Type{∑{α{∏{Op},T}}}, ::Type{∑{α{∏{Op},S}}}) where {T,S}
  return ∑{α{∏{Op},promote_type(T, S)}}
end

op(o::α) = o.args[2]
sites(o::α) = sites(op(o))

which_op(o::α{Op}) = which_op(op(o))
params(o::α{Op}) = params(op(o))
one(o::α{Op}) = one(coefficient(o))

sites(o::Union{∑,∏}) = unique(Iterators.flatten(Iterators.map(sites, o)))

# General definition for single-tensor operations like `exp` or `adjoint`.
# F: exp, adjoint, etc.
op(o::Applied{F}) where {F} = o.args[1]
sites(o::Applied{F}) where {F} = sites(op(o))
which_op(o::Applied{F}) where {F} = which_op(op(o))
params(o::Applied{F}) where {F} = params(op(o))

const OpTuple = Union{Tuple{<:WhichOp,Vararg},Tuple{<:Number,<:WhichOp,Vararg}}

# Conversion from Tuple
Op(o::Tuple) = Op(o...)
Op(which_op::WhichOp, sites::Tuple; kwargs...) = Op(which_op, sites, values(kwargs))
Op(which_op::WhichOp, sites::Int...; kwargs...) = Op(which_op, sites; kwargs...)
Op(which_op::WhichOp, sites::Vector{Int}; kwargs...) = Op(which_op, Tuple(sites); kwargs...)
function Op(which_op::WhichOp, sites_params::Union{Int,<:NamedTuple}...)
  return Op(which_op, Base.front(sites_params), last(sites_params))
end
Op(α::Number, which_op::WhichOp, args...; kwargs...) = α * Op(which_op, args...; kwargs...)
function Op(which_op::WhichOp, sites_params::Union{Int,WhichOp,NamedTuple}...)
  ts = split(x -> x isa WhichOp, (which_op, sites_params...))
  args = filter(x -> !(x isa Tuple{}), ts)
  return ∏(collect(Op.(args)))
end

# Conversion to `∑{Op}` (replacement for `OpSum`)
∑{Op}(o::Vector{<:OpExpr}) = ∑(o)
∑{Op}(o::OpExpr) = ∑{Op}() + o
∑{Op}(o::OpTuple) = ∑{Op}(Op(o))
∑{Op}(which_op::WhichOp, args...; kwargs...) = ∑{Op}(Op(which_op, args...; kwargs...))
function ∑{Op}(α::Number, which_op::WhichOp, args...; kwargs...)
  return ∑{Op}(Op(α, which_op, args...; kwargs...))
end

# Lazy operations with Op
(arg1::Number * arg2::Op) = α(arg1, arg2)
(arg1::Op / arg2::Number) = inv(arg2) * arg1
(arg1::Op * arg2::Op) = ∏([arg1, arg2])
(arg1::Op + arg2::Op) = ∑([arg1, arg2])
-(o::Op) = -𝟏 * o

# Rules for adding, subtracting, and multiplying with Tuples
(arg1::OpExpr + arg2::Tuple) = arg1 + Op(arg2)
(arg1::Tuple + arg2::OpExpr) = Op(arg1) + arg2
(arg1::OpExpr - arg2::Tuple) = arg1 - Op(arg2)
(arg1::Tuple - arg2::OpExpr) = Op(arg1) - arg2
(arg1::OpExpr * arg2::Tuple) = arg1 * Op(arg2)
(arg1::Tuple * arg2::OpExpr) = Op(arg1) * arg2

function print_sites(io::IO, sites)
  nsites = length(sites)
  for n in 1:nsites
    print(io, sites[n])
    if n < nsites
      print(io, ", ")
    end
  end
end

function show(io::IO, ::MIME"text/plain", o::Op)
  print(io, which_op(o), "(")
  print_sites(io, sites(o))
  if !isempty(params(o))
    print(io, ", ", params(o))
  end
  return print(io, ")")
end

function show(io::IO, ::MIME"text/plain", o::∏{Op})
  for n in 1:length(o)
    print(io, o[n])
    if n < length(o)
      print(io, " * ")
    end
  end
end

function print_coefficient(io::IO, o)
  return print(io, o)
end

function print_coefficient(io::IO, o::Complex)
  return print(io, "(", o, ")")
end

function show(io::IO, ::MIME"text/plain", o::Union{α{Op},α{∏{Op}}})
  print_coefficient(io, coefficient(o))
  print(io, " ")
  return print(io, op(o))
end

function show(io::IO, ::MIME"text/plain", o::Union{∑{Op},∑{<:α{Op}},∑{∏{Op}},∑{<:α{∏{Op}}}})
  for n in 1:length(o)
    print(io, o[n])
    if n < length(o)
      print(io, " +\n")
    end
  end
end

show(io::IO, o::OpExpr) = show(io, MIME("text/plain"), o)

end


"""
    fillwt(itype)

Fill weight tables for all order and number of flavours.
`itype` is used to select un-polarised pdfs (1), polarised
pdfs (2) or fragmentation functions (3).

# Returns
- `nwds::Integer`: number of words used in memory.
"""
function fillwt(itype::Integer)

    itype = Ref{Int32}(itype)
    nwds = Ref{Int32}()

    # For now-obsolete integer id vars
    idnum1 = Ref{Int32}()
    idnum2 = Ref{Int32}()

    @ccall fillwt_(itype::Ref{Int32}, idnum1::Ref{Int32},
                   idnum2::Ref{Int32}, nwds::Ref{Int32})::Nothing

    nwds[]
end

"""
    dmpwgt(itype, lun, filename)

Dump weight tables of a given `itype` to `filename`.
"""
function dmpwgt(itype::Integer, lun::Integer, filename::String)

    itype = Ref{Int32}(itype)
    lun = Ref{Int32}(lun)

    @ccall dmpwgt_(itype::Ref{Int32}, lun::Ref{Int32}, filename::Cstring)::Nothing

    nothing
end

"""
    readwt(lun, filename)

Read weight tables from filename.
TODO: Fix implementation which currently crashes.
"""
function readwt(lun::Integer, filename::String)

    lun = Ref{Int32}(lun)
    ierr = Ref{Int32}()
    nwds = Ref{Int32}()

    # For now-obsolete integer id vars
    idnum1 = Ref{Int32}()
    idnum2 = Ref{Int32}()

    @ccall readwt_(lun::Ref{Int32}, filename::Cstring, idnum1::Ref{Int32},
                   idnum2::Ref{Int32}, nwds::Ref{Int32}, ierr::Ref{Int32})::Nothing

    nwds[], ierr[]

end

"""
    wtfile(itype, filename)

Maintains an up-to-date weight table in filename.
"""
function wtfile(itype::Integer, filename::String)

    itype = Ref{Int32}(itype)

    @ccall wtfile_(itype::Ref{Int32}, filename::Cstring)::Nothing

    nothing
end


"""
    Gibbs(n_iters, algs...)

Compositional MCMC interface. Gibbs sampling combines one or more
sampling algorithms, each of which samples from a different set of
variables in a model.

Example:
```julia
@model gibbs_example(x) = begin
    v1 ~ Normal(0,1)
    v2 ~ Categorical(5)
        ...
end

# Use PG for a 'v2' variable, and use HMC for the 'v1' variable.
# Note that v2 is discrete, so the PG sampler is more appropriate
# than is HMC.
alg = Gibbs(1000, HMC(1, 0.2, 3, :v1), PG(20, 1, :v2))
```

Tips:
- `HMC` and `NUTS` are fast samplers, and can throw off particle-based
methods like Particle Gibbs. You can increase the effectiveness of particle sampling by including
more particles in the particle sampler.
"""
mutable struct Gibbs{A} <: InferenceAlgorithm
    n_iters   ::  Int     # number of Gibbs iterations
    algs      ::  A   # component sampling algorithms
    thin      ::  Bool    # if thinning to output only after a whole Gibbs sweep
    gid       ::  Int
end
Gibbs(n_iters::Int, algs...; thin=true) = Gibbs(n_iters, algs, thin, 0)
Gibbs(alg::Gibbs, new_gid) = Gibbs(alg.n_iters, alg.algs, alg.thin, new_gid)

const GibbsComponent = Union{Hamiltonian,MH,PG}

function Sampler(alg::Gibbs, model::Model)
    n_samplers = length(alg.algs)
    samplers = Array{Sampler}(undef, n_samplers)

    space = Set{Symbol}()

    for i in 1:n_samplers
        sub_alg = alg.algs[i]
        if isa(sub_alg, GibbsComponent)
            samplers[i] = Sampler(typeof(sub_alg)(sub_alg, i), model)
        else
            @error("[Gibbs] unsupport base sampling algorithm $alg")
        end
        space = union(space, sub_alg.space)
    end

    # Sanity check for space
    @assert issubset(Set(pvars(model)), space) "[Gibbs] symbols specified to samplers ($space) doesn't cover the model parameters ($(Set(pvars(model))))"

    if Set(pvars(model)) != space
        @warn("[Gibbs] extra parameters specified by samplers don't exist in model: $(setdiff(space, Set(pvars(model))))")
    end

    info = Dict{Symbol, Any}()
    info[:samplers] = samplers

    Sampler(alg, info)
end

function sample(
                model::Model,
                alg::Gibbs;
                save_state=false,         # flag for state saving
                resume_from=nothing,      # chain to continue
                reuse_spl_n=0             # flag for spl re-using
                )

    # Init the (master) Gibbs sampler
    spl = reuse_spl_n > 0 ? resume_from.info[:spl] : Sampler(alg, model)

    @assert typeof(spl.alg) == typeof(alg) "[Turing] alg type mismatch; please use resume() to re-use spl"

    # Initialize samples
    sub_sample_n = []
    for sub_alg in alg.algs
        if isa(sub_alg, GibbsComponent)
            push!(sub_sample_n, sub_alg.n_iters)
        else
            @error("[Gibbs] unsupport base sampling algorithm $alg")
        end
    end

    # Compute the number of samples to store
    n = reuse_spl_n > 0 ? reuse_spl_n : alg.n_iters
    sample_n = n * (alg.thin ? 1 : sum(sub_sample_n))

    # Init samples
    time_total = zero(Float64)
    samples = Array{Sample}(undef, sample_n)
    weight = 1 / sample_n
    for i = 1:sample_n
        samples[i] = Sample(weight, Dict{Symbol, Any}())
    end

    # Init parameters
    varInfo = if resume_from == nothing
        vi_ = VarInfo()
        model(vi_, HamiltonianRobustInit())
        vi_
    else
        resume_from.info[:vi]
    end

    n = spl.alg.n_iters; i_thin = 1

    # Gibbs steps
    PROGRESS[] && (spl.info[:progress] = ProgressMeter.Progress(n, 1, "[Gibbs] Sampling...", 0))
    for i = 1:n
        @debug "Gibbs stepping..."

        time_elapsed = zero(Float64)
        lp = nothing; epsilon = nothing; lf_num = nothing; eval_num = nothing

        for local_spl in spl.info[:samplers]
            last_spl = local_spl
      # PROGRESS[] && haskey(spl.info, :progress) && (local_spl.info[:progress] = spl.info[:progress])

            @debug "$(typeof(local_spl)) stepping..."

            if isa(local_spl.alg, GibbsComponent)
                for _ = 1:local_spl.alg.n_iters
                    @debug "recording old θ..."
                    time_elapsed_thin = @elapsed varInfo, is_accept = step(model, local_spl, varInfo, Val(i==1))

                    if ~spl.alg.thin
                        samples[i_thin].value = Sample(varInfo).value
                        samples[i_thin].value[:elapsed] = time_elapsed_thin
                        if ~isa(local_spl.alg, Hamiltonian)
                            # If statement below is true if there is a HMC component which provides lp and epsilon
                            if lp != nothing samples[i_thin].value[:lp] = lp end
                            if epsilon != nothing samples[i_thin].value[:epsilon] = epsilon end
                            if lf_num != nothing samples[i_thin].value[:lf_num] = lf_num end
                            if eval_num != nothing samples[i_thin].value[:eval_num] = eval_num end
                        end
                        i_thin += 1
                    end
                    time_elapsed += time_elapsed_thin
                end

                if isa(local_spl.alg, Hamiltonian)
                    lp = getlogp(varInfo)
                    epsilon = getss(local_spl.info[:wum])
                    lf_num = local_spl.info[:lf_num]
                    eval_num = local_spl.info[:eval_num]
                end
            else
                @error("[Gibbs] unsupport base sampler $local_spl")
            end
        end

        time_total += time_elapsed

        if spl.alg.thin
            samples[i].value = Sample(varInfo).value
            samples[i].value[:elapsed] = time_elapsed
            # If statement below is true if there is a HMC component which provides lp and epsilon
            if lp != nothing samples[i].value[:lp] = lp end
            if epsilon != nothing samples[i].value[:epsilon] = epsilon end
            if lf_num != nothing samples[i].value[:lf_num] = lf_num end
            if eval_num != nothing samples[i].value[:eval_num] = eval_num end
        end

        if PROGRESS[]
            if haskey(spl.info, :progress)
                ProgressMeter.update!(spl.info[:progress], spl.info[:progress].counter + 1)
            end
        end
    end

    @info("[Gibbs] Finished with")
    @info("  Running time    = $time_total;")

    if resume_from != nothing   # concat samples
        pushfirst!(samples, resume_from.value2...)
    end
    c = Chain(0, samples)       # wrap the result by Chain

    if save_state               # save state
        save!(c, spl, model, varInfo)
    end

    return c
end


module CustomVertexGraphs
  using AbstractBijections
  using Dictionaries
  using Graphs

  export set_vertices

  import Graphs: src, dst, nv, vertices, has_vertex, ne, edges, has_edge, neighbors, outneighbors, inneighbors, all_neighbors, is_directed, add_edge!, add_vertex!, add_vertices!, induced_subgraph, adjacency_matrix, blockdiag, edgetype

  import Base: show, eltype, copy

  # TODO: restrict to `AbstractSimpleGraph` (graph with contiguous integer vertices)
  # and `InvVectorBijection`, a Bijection where the inverse is a `Vector` (maps a contiguous
  # space of integers to another space).
  # TODO: rename `parent_graph` to something more descriptive like `simple_graph`, `int_vertex_graph`, etc.
  struct CustomVertexGraph{V,G<:AbstractGraph,B<:AbstractBijection} <: AbstractGraph{V}
    parent_graph::G
    vertex_to_parent_vertex::B # Invertible map from the vertices to the parent vertices
    function CustomVertexGraph(parent_graph::G, vertex_to_parent_vertex::B) where {G<:AbstractGraph,B<:AbstractBijection}
      @assert issetequal(vertices(parent_graph), image(vertex_to_parent_vertex))
      V = domain_eltype(vertex_to_parent_vertex)
      return new{V,G,B}(parent_graph, vertex_to_parent_vertex)
    end
  end
  parent_graph(graph::CustomVertexGraph) = graph.parent_graph
  vertex_to_parent_vertex(graph::CustomVertexGraph) = graph.vertex_to_parent_vertex

  copy(graph::CustomVertexGraph) = CustomVertexGraph(copy(parent_graph(graph)), copy(vertex_to_parent_vertex(graph)))

  eltype(g::CustomVertexGraph{V}) where {V} = V

  # Convenient constructor
  set_vertices(graph::AbstractGraph, vertices) = CustomVertexGraph(graph, vertices)

  vertices(graph::CustomVertexGraph) = domain(vertex_to_parent_vertex(graph))

  parent_graph_type(::Type{<:CustomVertexGraph{<:Any,G}}) where {G} = G

  parent_vertices(graph::CustomVertexGraph) = vertices(parent_graph(graph))
  parent_edges(graph::CustomVertexGraph) = edges(parent_graph(graph))
  parent_edgetype(graph::CustomVertexGraph) = edgetype(parent_graph(graph))

  parent_vertex(graph::CustomVertexGraph, vertex) = vertex_to_parent_vertex(graph)[vertex]
  parent_edge(graph::CustomVertexGraph, edge) = parent_edgetype(graph)(parent_vertex(graph, src(edge)), parent_vertex(graph, dst(edge)))
  parent_vertices(graph::CustomVertexGraph, vertices) = [parent_vertex(graph, vertex) for vertex in vertices]
  parent_vertex_to_vertex(graph::CustomVertexGraph, parent_vertex) = vertices(graph)[parent_vertex]

  CustomVertexGraph(vertices::Vector{T}) where T = CustomVertexGraph{Graph{Int}}(vertices)
  CustomVertexDiGraph(vertices::Vector{T}) where T = CustomVertexGraph{DiGraph{Int}}(vertices)

  import Base: Pair, Tuple, show, ==, hash, eltype
  import Graphs: AbstractEdge, src, dst, reverse

  abstract type AbstractCustomVertexEdge{V} <: AbstractEdge{V} end
  struct CustomVertexEdge{V} <: AbstractCustomVertexEdge{V}
    src::V
    dst::V
  end

  CustomVertexEdge{T}(e::CustomVertexEdge{T}) where {T} = e

  CustomVertexEdge(t::Tuple) = CustomVertexEdge(t[1], t[2])
  CustomVertexEdge(p::Pair) = CustomVertexEdge(p.first, p.second)
  CustomVertexEdge{T}(p::Pair) where {T} = CustomVertexEdge(T(p.first), T(p.second))
  CustomVertexEdge{T}(t::Tuple) where {T} = CustomVertexEdge(T(t[1]), T(t[2]))

  eltype(::Type{<:ET}) where ET<:AbstractCustomVertexEdge{T} where T = T

  src(e::AbstractCustomVertexEdge) = e.src
  dst(e::AbstractCustomVertexEdge) = e.dst

  function show(io::IO, mime::MIME"text/plain", e::AbstractCustomVertexEdge)
    show(io, src(e))
    print(io, " => ")
    show(io, dst(e))
    return nothing
  end

  show(io::IO, edge::AbstractCustomVertexEdge) = show(io, MIME"text/plain"(), edge)

  # Conversions
  Pair(e::AbstractCustomVertexEdge) = Pair(src(e), dst(e))
  Tuple(e::AbstractCustomVertexEdge) = (src(e), dst(e))

  CustomVertexEdge{T}(e::AbstractCustomVertexEdge) where {T} = CustomVertexEdge{T}(T(e.src), T(e.dst))

  # Convenience functions
  reverse(e::T) where T<:AbstractCustomVertexEdge = T(dst(e), src(e))
  ==(e1::AbstractCustomVertexEdge, e2::AbstractCustomVertexEdge) = (src(e1) == src(e2) && dst(e1) == dst(e2))
  hash(e::AbstractCustomVertexEdge, h::UInt) = hash(src(e), hash(dst(e), h))

  edgetype(graph::CustomVertexGraph{V}) where {V} = CustomVertexEdge{V}

  default_vertices(graph::AbstractGraph) = Vector(vertices(graph))

  function CustomVertexGraph(graph::AbstractGraph, vertices=default_vertices(graph))
    if length(vertices) != nv(graph)
      throw(ArgumentError("Vertices and parent graph's vertices must have equal length."))
    end
    if !allunique(vertices)
      throw(ArgumentError("Vertices have to be unique."))
    end
    return CustomVertexGraph(graph, inv(Bijection(vertices)))
  end

  function CustomVertexGraph(graph::AbstractGraph, dims::Tuple{Vararg{Integer}})
    return CustomVertexGraph(graph, vec(Tuple.(CartesianIndices(dims))))
  end

  function CustomVertexGraph(dims::Tuple{Vararg{Integer}})
    return CustomVertexGraph(Graph(prod(dims)), vec(Tuple.(CartesianIndices(dims))))
  end

  function CustomVertexGraph{S}(vertices::Vector) where {S<:AbstractGraph}
    return CustomVertexGraph(S(length(vertices)), vertices)
  end

  has_vertex(g::CustomVertexGraph, v) = v in vertices(g)

  function edges(graph::CustomVertexGraph)
    vertex(parent_vertex) = inv(vertex_to_parent_vertex(graph))[parent_vertex]
    edge(parent_edge) = CustomVertexEdge(vertex(src(parent_edge)), vertex(dst(parent_edge)))
    return map(edge, parent_edges(graph))
  end

  for f in [:outneighbors, :inneighbors, :all_neighbors, :neighbors]
    @eval begin
      function $f(graph::CustomVertexGraph, v)
        parent_vertices = $f(parent_graph(graph), parent_vertex(graph, v))
        return [parent_vertex_to_vertex(graph, u) for u ∈ parent_vertices]
      end
    end
  end

  # Ambiguity errors with Graphs.jl
  for f in [
    :neighbors, :inneighbors, :outneighbors, :all_neighbors
  ]
    @eval begin
      $f(tn::CustomVertexGraph, vertex::Integer) = $f(parent_graph(tn), vertex)
    end
  end

  function add_edge!(graph::CustomVertexGraph, edge::CustomVertexEdge)
    add_edge!(parent_graph(graph), parent_edge(graph, edge))
    return graph
  end

  function has_edge(graph::CustomVertexGraph, edge::CustomVertexEdge)
    return has_edge(parent_graph(graph), parent_edge(graph, edge))
  end

  # handles single-argument edge constructors such as pairs and tuples
  has_edge(g::CustomVertexGraph, x) = has_edge(g, edgetype(g)(x))
  add_edge!(g::CustomVertexGraph, x) = add_edge!(g, edgetype(g)(x))

  # handles two-argument edge constructors like src,dst
  has_edge(g::CustomVertexGraph, x, y) = has_edge(g, edgetype(g)(x, y))
  add_edge!(g::CustomVertexGraph, x, y) = add_edge!(g, edgetype(g)(x, y))

  function add_vertex!(graph::CustomVertexGraph, v)
    if v ∈ vertices(graph)
      throw(ArgumentError("Duplicate vertices are not allowed"))
    end
    add_vertex!(parent_graph(graph))
    insert!(vertex_to_parent_vertex(graph), v, last(parent_vertices(graph)))
    return graph
  end

  function add_vertices!(graph::CustomVertexGraph, vertices::Vector)
    if any(v ∈ vertices(graph) for v ∈ vertices)
      throw(ArgumentError("Duplicate vertices are not allowed"))
    end
    for vertex in vertices
      add_vertex!(graph, vertex)
    end
    return graph
  end

  function induced_subgraph(graph::CustomVertexGraph, sub_vertices::Vector)
    return _induced_subgraph(graph, _get_vertices(graph, sub_vertices))
  end

  function _get_vertices(graph::CustomVertexGraph, vertex)
    return vertex
  end

  function _get_vertices(graph::CustomVertexGraph, vertices::Vector)
    return mapreduce(vertex -> _get_vertices(graph, vertex), vcat, vertices)
  end

  function _induced_subgraph(graph::CustomVertexGraph, vertices::Vector)
    sub_graph, _ = induced_subgraph(parent_graph(graph), parent_vertices(graph, vertices))
    return CustomVertexGraph(sub_graph, vertices), vertices
  end

  is_directed(LG::Type{<:CustomVertexGraph}) = is_directed(parent_graph_type(LG))

  function blockdiag(graph1::CustomVertexGraph, graph2::CustomVertexGraph)
    new_parent_graph = blockdiag(parent_graph(graph1), parent_graph(graph2))
    new_vertices = vcat(vertices(graph1), vertices(graph2))
    return CustomVertexGraph(new_parent_graph, new_vertices)
  end

  for f in [:nv, :ne, :adjacency_matrix]
    @eval begin
      $f(graph::CustomVertexGraph, args...) = $f(parent_graph(graph), args...)
    end
  end

  function show(io::IO, mime::MIME"text/plain", graph::CustomVertexGraph)
    println(io, "CustomVertexGraph with $(nv(graph)) vertices:")
    show(io, mime, vertices(graph))
    println(io, "\n")
    println(io, "and $(ne(graph)) edge(s):")
    for e in edges(graph)
      show(io, mime, e)
      println(io)
    end
    return nothing
  end

  show(io::IO, graph::CustomVertexGraph) = show(io, MIME"text/plain"(), graph)
end # module CustomVertexGraphs


using Exam
using Test
using CSV
using Random
using DataFrames

function generate_data()
    open("exam_results.csv", "w") do io
        write(io, "name,id,bonus,1,2,3\n")
        for i = 1:20
           write(io, randstring('a':'z', 6), ",", string(rand(1:Int(1e12))), ",",string(rand((false, true))), ",",
                     string(rand(0:5)), ",", string(rand(0:2)), ",", string(rand(0:10)), "\n")
        end
    end
end

@testset  begin
    correct_results = CSV.read("simple_test_output.csv", DataFrame)
    results = evaluate("simple_results.csv", "simple_exam.yml", histogram=false)
    @test results == correct_results

    correct_results = CSV.read("simple_test_output_anonymous.csv", DataFrame)
    results = evaluate("simple_results.csv", "simple_exam.yml", anonymize=true, histogram=false)
    @test results == correct_results

end


using DataDeps

register(DataDep(
  "indian-liver-patient",
  "http://archive.ics.uci.edu/ml/datasets/ILPD+%28Indian+Liver+Patient+Dataset%29",
  "http://archive.ics.uci.edu/ml/machine-learning-databases/00225/Indian%20Liver%20Patient%20Dataset%20(ILPD).csv",
  "84feac16488de5cf89bd22bd802c77f25841fe93e9ddd32355683e94d46c3425",
  post_fetch_method=(path -> begin
    UCIData.process_dataset(path,
        target_index=11,
        feature_indices=1:10,
        categoric_indices=2,
    )
  end),
))


################################################################################
## UNIT TESTING JUST ARITHMETIC operations

include("./unit/unit-addition.jl")
include("./unit/unit-subtraction.jl")
include("./unit/unit-multiplication.jl")
include("./unit/unit-division.jl")

include("./unit/unit-boundstesting.jl")


Base.@kwdef struct ScreeningParams
  start_time::Float64 = 0.0
  precision::Float64 = 0.8
  period::Float64 = 7.0
  lower_bound_age::Int64 = 8
  upper_bound_age::Int64 = 16
end

function screening!(state::AbstractSimState, params::AbstractSimParams, event::Event)
  if params.screening_params != nothing
    for id in 1:numindividuals(state)
      health = MocosSim.health(state, id)
      if health == Healthy || health == Recovered || health == Incubating
        continue
      end

      age = params.ages[id]
      if age < params.screening_params.lower_bound_age || age > params.screening_params.upper_bound_age
        continue
      end

      if rand(state.rng) >= params.screening_params.precision
        continue
      end
      push!(
        state.queue,
        Event(
          Val(DetectionEvent),
          time(event),
          id,
          OutsideQuarantineDetection),
          immediate=true)
    end
  end
end

function add_screening!(state::AbstractSimState, params::AbstractSimParams, time_limit::TimePoint=typemax(TimePoint))
  for screening_time in params.screening_params.start_time:params.screening_params.period:time_limit
    event = Event(Val(ScreeningEvent), screening_time)
    push!(state.queue, event)
  end
end


# Refine elements in the DG solver based on a list of cell_ids that should be refined
function refine!(u_ode::AbstractVector, adaptor, mesh::TreeMesh{1},
                 equations, dg::DGSEM, cache, cells_to_refine)
  # Return early if there is nothing to do
  if isempty(cells_to_refine)
    return
  end

  # Determine for each existing element whether it needs to be refined
  needs_refinement = falses(nelements(dg, cache))

  # The "Ref(...)" is such that we can vectorize the search but not the array that is searched
  elements_to_refine = searchsortedfirst.(Ref(cache.elements.cell_ids[1:nelements(dg, cache)]),
                                          cells_to_refine)
  needs_refinement[elements_to_refine] .= true

  # Retain current solution data
  old_n_elements = nelements(dg, cache)
  old_u_ode = copy(u_ode)
  GC.@preserve old_u_ode begin # OBS! If we don't GC.@preserve old_u_ode, it might be GC'ed
    old_u = wrap_array(old_u_ode, mesh, equations, dg, cache)

    # Get new list of leaf cells
    leaf_cell_ids = local_leaf_cells(mesh.tree)

    # re-initialize elements container
    @unpack elements = cache
    resize!(elements, length(leaf_cell_ids))
    init_elements!(elements, leaf_cell_ids, mesh, dg.basis.nodes)
    @assert nelements(dg, cache) > old_n_elements

    resize!(u_ode, nvariables(equations) * nnodes(dg)^ndims(mesh) * nelements(dg, cache))
    u = wrap_array(u_ode, mesh, equations, dg, cache)

    # Loop over all elements in old container and either copy them or refine them
    element_id = 1
    for old_element_id in 1:old_n_elements
      if needs_refinement[old_element_id]
        # Refine element and store solution directly in new data structure
        refine_element!(u, element_id, old_u, old_element_id,
                        adaptor, equations, dg)
        element_id += 2^ndims(mesh)
      else
        # Copy old element data to new element container
        @views u[:, .., element_id] .= old_u[:, .., old_element_id]
        element_id += 1
      end
    end
    # If everything is correct, we should have processed all elements.
    # Depending on whether the last element processed above had to be refined or not,
    # the counter `element_id` can have two different values at the end.
    @assert element_id == nelements(dg, cache) + 1 || element_id == nelements(dg, cache) + 2^ndims(mesh) "element_id = $element_id, nelements(dg, cache) = $(nelements(dg, cache))"
  end # GC.@preserve old_u_ode

  # re-initialize interfaces container
  @unpack interfaces = cache
  resize!(interfaces, count_required_interfaces(mesh, leaf_cell_ids))
  init_interfaces!(interfaces, elements, mesh)

  # re-initialize boundaries container
  @unpack boundaries = cache
  resize!(boundaries, count_required_boundaries(mesh, leaf_cell_ids))
  init_boundaries!(boundaries, elements, mesh)

  # Sanity check
  if isperiodic(mesh.tree)
    @assert ninterfaces(interfaces) == 1 * nelements(dg, cache) ("For 1D and periodic domains, the number of interfaces must be the same as the number of elements")
  end

  return nothing
end

# TODO: Taal compare performance of different implementations
# Refine solution data u for an element, using L2 projection (interpolation)
function refine_element!(u::AbstractArray{<:Any,3}, element_id,
                         old_u, old_element_id,
                         adaptor::LobattoLegendreAdaptorL2, equations, dg)
  @unpack forward_upper, forward_lower = adaptor

  # Store new element ids
  left_id  = element_id
  right_id = element_id + 1

  @boundscheck begin
    @assert old_element_id >= 1
    @assert size(old_u, 1) == nvariables(equations)
    @assert size(old_u, 2) == nnodes(dg)
    @assert size(old_u, 3) >= old_element_id
    @assert     element_id >= 1
    @assert size(    u, 1) == nvariables(equations)
    @assert size(    u, 2) == nnodes(dg)
    @assert size(    u, 3) >= element_id + 1
  end

  # Interpolate to left element
  for i in eachnode(dg)
    acc = zero(get_node_vars(u, equations, dg, i, element_id))
    for k in eachnode(dg)
      acc += get_node_vars(old_u, equations, dg, k, old_element_id) * forward_lower[i, k]
    end
    set_node_vars!(u, acc, equations, dg, i, left_id)
  end

  # Interpolate to right element
  for i in eachnode(dg)
    acc = zero(get_node_vars(u, equations, dg, i, element_id))
    for k in eachnode(dg)
      acc += get_node_vars(old_u, equations, dg, k, old_element_id) * forward_upper[i, k]
    end
    set_node_vars!(u, acc, equations, dg, i, right_id)
  end

  return nothing
end

# Coarsen elements in the DG solver based on a list of cell_ids that should be removed
function coarsen!(u_ode::AbstractVector, adaptor, mesh::TreeMesh{1},
                  equations, dg::DGSEM, cache, child_cells_to_coarsen)
  # Return early if there is nothing to do
  if isempty(child_cells_to_coarsen)
    return
  end

  # Determine for each old element whether it needs to be removed
  to_be_removed = falses(nelements(dg, cache))
  # The "Ref(...)" is such that we can vectorize the search but not the array that is searched
  elements_to_remove = searchsortedfirst.(Ref(cache.elements.cell_ids[1:nelements(dg, cache)]),
                                          child_cells_to_coarsen)
  to_be_removed[elements_to_remove] .= true

  # Retain current solution data
  old_n_elements = nelements(dg, cache)
  old_u_ode = copy(u_ode)
  GC.@preserve old_u_ode begin # OBS! If we don't GC.@preserve old_u_ode, it might be GC'ed
    old_u = wrap_array(old_u_ode, mesh, equations, dg, cache)

    # Get new list of leaf cells
    leaf_cell_ids = local_leaf_cells(mesh.tree)

    # re-initialize elements container
    @unpack elements = cache
    resize!(elements, length(leaf_cell_ids))
    init_elements!(elements, leaf_cell_ids, mesh, dg.basis.nodes)
    @assert nelements(dg, cache) < old_n_elements

    resize!(u_ode, nvariables(equations) * nnodes(dg)^ndims(mesh) * nelements(dg, cache))
    u = wrap_array(u_ode, mesh, equations, dg, cache)

    # Loop over all elements in old container and either copy them or coarsen them
    skip = 0
    element_id = 1
    for old_element_id in 1:old_n_elements
      # If skip is non-zero, we just coarsened 2^ndims elements and need to omit the following elements
      if skip > 0
        skip -= 1
        continue
      end

      if to_be_removed[old_element_id]
        # If an element is to be removed, sanity check if the following elements
        # are also marked - otherwise there would be an error in the way the
        # cells/elements are sorted
        @assert all(to_be_removed[old_element_id:(old_element_id+2^ndims(mesh)-1)]) "bad cell/element order"

        # Coarsen elements and store solution directly in new data structure
        coarsen_elements!(u, element_id, old_u, old_element_id,
                          adaptor, equations, dg)
        element_id += 1
        skip = 2^ndims(mesh) - 1
      else
        # Copy old element data to new element container
        @views u[:, .., element_id] .= old_u[:, .., old_element_id]
        element_id += 1
      end
    end
    # If everything is correct, we should have processed all elements.
    @assert element_id == nelements(dg, cache) + 1 "element_id = $element_id, nelements(dg, cache) = $(nelements(dg, cache))"
  end # GC.@preserve old_u_ode

  # re-initialize interfaces container
  @unpack interfaces = cache
  resize!(interfaces, count_required_interfaces(mesh, leaf_cell_ids))
  init_interfaces!(interfaces, elements, mesh)

  # re-initialize boundaries container
  @unpack boundaries = cache
  resize!(boundaries, count_required_boundaries(mesh, leaf_cell_ids))
  init_boundaries!(boundaries, elements, mesh)

  # Sanity check
  if isperiodic(mesh.tree)
    @assert ninterfaces(interfaces) == 1 * nelements(dg, cache) ("For 1D and periodic domains, the number of interfaces must be the same as the number of elements")
  end

  return nothing
end

# TODO: Taal compare performance of different implementations
# Coarsen solution data u for two elements, using L2 projection
function coarsen_elements!(u::AbstractArray{<:Any,3}, element_id,
                           old_u, old_element_id,
                           adaptor::LobattoLegendreAdaptorL2, equations, dg)
  @unpack reverse_upper, reverse_lower = adaptor

  # Store old element ids
  left_id  = old_element_id
  right_id = old_element_id + 1

  @boundscheck begin
    @assert old_element_id >= 1
    @assert size(old_u, 1) == nvariables(equations)
    @assert size(old_u, 2) == nnodes(dg)
    @assert size(old_u, 3) >= old_element_id + 1
    @assert     element_id >= 1
    @assert size(    u, 1) == nvariables(equations)
    @assert size(    u, 2) == nnodes(dg)
    @assert size(    u, 3) >= element_id
  end

  for i in eachnode(dg)
    acc = zero(get_node_vars(u, equations, dg, i, element_id))

    # Project from lower left element
    for k in eachnode(dg)
      acc += get_node_vars(old_u, equations, dg, k, left_id) * reverse_lower[i, k]
    end

    # Project from lower right element
    for k in eachnode(dg)
      acc += get_node_vars(old_u, equations, dg, k, right_id) * reverse_upper[i, k]
    end

    # Update value
    set_node_vars!(u, acc, equations, dg, i, element_id)
  end
end

# this method is called when an `ControllerThreeLevel` is constructed
function create_cache(::Type{ControllerThreeLevel}, mesh::TreeMesh{1}, equations, dg::DG, cache)

  controller_value = Vector{Int}(undef, nelements(dg, cache))
  return (; controller_value)
end

function create_cache(::Type{ControllerThreeLevelCombined}, mesh::TreeMesh{1}, equations, dg::DG, cache)

  controller_value = Vector{Int}(undef, nelements(dg, cache))
  return (; controller_value)
end


using HorizonSideRobots
include("BackPath.jl")
include("Structs.jl")

function count_borders!(robot::Robot)
    to_sud_west_corner = BackPath(r)
    robot = CountbordersRobot(robot)
    snake(robot)
    return get_num_borders(robot) - 1
end


module NARemovers

using Random
using DataFrames: DataFrame, nrow

using ..AbsTypes
using ..BaseFilters
using ..Utils

import ..AbsTypes: fit, fit!, transform, transform!
export fit, fit!, transform, transform!
export NARemover

"""
    NARemover(
      Dict(
        :name => "nadetect",
        :acceptance => 0.10 # tolerable NAs percentage
      )
    )

Removes columns with NAs greater than acceptance rate.
This assumes that it processes columns of features.
The output column should not be part of input to avoid
it being excluded if it fails the acceptance critera.

Implements `fit!` and `transform!`.
"""
mutable struct NARemover <: Transformer
   name::String
   model::Dict{Symbol,Any}

   function NARemover(args::Dict = Dict())
      default_args = Dict{Symbol,Any}(
         :name => "nadetect",
         :acceptance => 0.10
      )
      cargs=nested_dict_merge(default_args,args)
      cargs[:name] = cargs[:name]*"_"*randstring(3)
      new(cargs[:name],cargs)
   end
end

"""
NARemover(acceptance::Float64)

Helper function for NARemover.
"""
NARemover(acceptance::Float64) = NARemover(Dict(:acceptance => acceptance))

"""
    fit!(nad::NARemover,features::DataFrame,labels::Vector=[])

Checks and exit of df is empty

# Arguments
- `nad::NARemover`: custom type
- `features::DataFrame`: input
- `labels::Vector=[]`:
"""
function fit!(nad::NARemover, features::DataFrame, labels::Vector=[])::Nothing
   return nothing
end

function fit(nad::NARemover, features::DataFrame, labels::Vector=[])::NARemover
   fit!(nad,features,labels)
   return deepcopy(nad)
end

"""
    transform!(nad::NARemover,nfeatures::DataFrame)

Removes columns with NAs greater than acceptance rate.

# Arguments
- `nad::NARemover`: custom type
- `nfeatures::DataFrame`: input
"""
function transform!(nad::NARemover, nfeatures::DataFrame)::DataFrame
   isempty(nfeatures) && return DataFrame()
   features = deepcopy(nfeatures)
   sz = nrow(features)
   tol = nad.model[:acceptance]
   colnames = []
   for (colname,dat) in collect(pairs(eachcol(features)))
      if sum(ismissing.(dat)) < tol*sz
         push!(colnames,colname)
      end
   end
   xtr =  features[:,colnames]
   return xtr
end

function transform(nad::NARemover, nfeatures::DataFrame)::DataFrame
   return transform!(nad,nfeatures)
end

end


import Trebuchet as Trebuchets

println(names(Trebuchets))

@doc Trebuchets.shoot

Trebuchets.shoot(5, 0.25pi, 500)


# policy file better sanity fewer contacts
adaptivepolicythresholdoff = 0
adaptivepolicythresholdon = 9999999999999

socialmaxyyt = socialmaxyy_o # max no. people met outside firm young_young
socialmaxoot = socialmaxoo_o # max no. people met outside firm old_old
socialmaxoyt = socialmaxoy_o
socialmaxyot = socialmaxyo_o

genhomeoffice = false

ptravel = ptravel_o

pshopt = [[1,1,1],[1,1,1]]

tauadj = 0.05

shorttimeprogram = false


using Test
using Revise
using Deeplearning

@testset "basic.jl" begin
    Deeplearning.df1()
    Deeplearning.df_arg_matrix()
    Deeplearning.df_m()
    Deeplearning.simple_models()
    Deeplearning.build_layers()
    Deeplearning.build_layers_o1()
    Deeplearning.build_layers_dense()
end


using TermWin

TermWin.initsession()
#s = "a test"
v = newTwEntry( rootTwScreen, Rational{Int}; width=25, posy=:random, posx=:random, showHelp=true, box=true )
v.title = "Input: "
v.data.tickSize = 1
activateTwObj( rootTwScreen )
ret = v.value
#unregisterTwObj( rootTwScreen, v )
TermWin.endsession()
println( "You entered ", string( ret ) )


# Consensus and fine splitting functions

"""
   fine_clustering(kmer_vecs; rough_radius::Float64 = 0.01, fine_radius::Float64 = 1.0, clust_multiplier = 1.0,
                        min_devs = 6, initial_devs = 20, user_min_clust = 10, distfunc = euclidean, center = mean,
                        verbose::Int64 = 1, cycle_lim = 20, triangle::Bool = false, p_value::Float64 = 0.01)

Preforms fine clustering on the already rough clustered sequences.

"""

function fine_clustering(kmer_vecs; rough_radius::Float64=0.01, fine_radius::Float64=1.0, clust_multiplier=1.0,
                        min_devs=6, initial_devs=20, user_min_clust=10, distfunc=euclidean, center=mean, verbose::Int64=1,
                        cycle_lim=20, triangle::Bool=false, p_value::Float64=0.01)

    if verbose > 0
        println("Your inputs have begun rough clustering...")
    end
    if verbose == 1 ||verbose == 2
        v = 2
    elseif verbose == 0
        v = 0
    end
    #"rough" clustering step to get the massively different sequences into different clusters
    μs, sizes, indices, centroid_inds = cluster(kmer_vecs, rough_radius, distfunc=corrected_kmer_dist,
                                        center=center,verbose=v, cycle_lim=cycle_lim, triangle=triangle)
    #only sequences greater than indicated size may form a cluster
    keep_inds = find(user_min_clust.<sizes)
    indices = indices[keep_inds]

    fine_indices = Array{Array{Int64, 1}, 1}()
    #get array which converts kmer indices into kmer sequences
    nuc_inds = gen_arrangements(Int(log(4, length(kmer_vecs[1]))))

    if verbose > 0
        println("Your inputs have begun fine clustering...")
    end
    if verbose == 1 || verbose == 0
        in_verbose = 0
    elseif verbose == 2
        in_verbose = 1
    end

    #call cluster_split on every cluster created by rough clustering
    for ind in 1:length(indices)
        if verbose > 0
            println("Splitting cluster number $(ind)")
        end
        temp =cluster_split(kmer_vecs[indices[ind]], 1, nuc_inds, rough_radius=rough_radius, fine_radius=fine_radius,
            clust_multiplier=clust_multiplier, min_devs=min_devs, initial_devs=initial_devs, user_min_clust=user_min_clust,
            distfunc=euclidean,center=center, verbose=in_verbose, cycle_lim=cycle_lim, triangle=triangle, p_value=p_value)
        for arr in temp
            push!(fine_indices, indices[ind][arr])
        end
    end
    if verbose > 0
        println("Returning $(length(fine_indices)) templates")
    end
    return fine_indices
end

"""
     gen_arrangements(k)

Returns an array containing all possible kmer sequences of length k

"""
function gen_arrangements(k)
    alph = [["A"],["C"],["G"],["T"]]
    current = copy(alph)
    for i in 1:k-1
        current = vcat([[vcat(j,letter) for letter in alph] for j in current]...)
    end
    return [join(i) for i in current]
end
"""
	cluster_split(kmer_vecs, count, nuc_inds; rough_radius::Float64=0.01, fine_radius::Float64=1.0, clust_multiplier=1.0,
                        min_devs=6, initial_devs=20, user_min_clust=10, distfunc=euclidean,
                        center=mean, verbose::Int64=1, cycle_lim=20, triangle::Bool=false, p_value::Float64=0.01)

Determines if passed cluster is a 'real' cluster and recursively splits non-error clusters
"""
function cluster_split(kmer_vecs, count, nuc_inds; rough_radius::Float64=0.01, fine_radius::Float64=1.0, clust_multiplier=1.0,
                        min_devs=6, initial_devs=20, user_min_clust=10, distfunc=euclidean,
                        center=mean, verbose::Int64=1, cycle_lim=20, triangle::Bool=false, p_value::Float64=0.01)
    if verbose == 2
        println("______________________________________________")
        println("Current level of cluster_split recursion is $(count)")
        count+=1
    end

    k = Int(log(4, length(kmer_vecs[1])))
    num_vecs = length(kmer_vecs)

    #arr of each kmer's stdev throughout the kmer vectors
    stdevs = [std([kmer_vecs[seq_ind][kmer_ind] for seq_ind in 1:length(kmer_vecs)]) for kmer_ind in 1:4^k]
    #arr of indices of high->low stdev kmer indices
    dev_inds = reverse(sortperm(stdevs))[1:initial_devs]#setting the number of kmers high results in a lack of splitting

    #eliminate kmer indices that represent homopolymer differences
    dev_inds = get_non_homopolymers(dev_inds, nuc_inds, min_devs=min_devs)

    if length(dev_inds) == 0
        return [[i for i in 1:length(kmer_vecs)]]
    end
    #kmer counts of high stdev kmer indices
    top_kmers =[kmer[dev_inds] for kmer in kmer_vecs]
    #cluster kmers by high stdev indices
    us,sizes,indices,C_inds =cluster(top_kmers, fine_radius,distfunc=distfunc,center=center,
                                            verbose=verbose,cycle_lim=cycle_lim, triangle=triangle)
    #replace 2nd largest with each size to determine whether or not to split on it
    #p_val_arr = [(1-cdf(Poisson(maximum(sizes)*(error_rate/4)), sizes[i])) *sum(kmer_vecs[1]) for i in 1:length(sizes)]

    #minimum non-error cluster size. Estimate of the sum of poisson distributions
    #min_clust = clust_multiplier*.25*2.5*sqrt(rough_radius*sum(kmer_vecs[1])*2*log(maximum(sizes)))
    #minimum cluster size is the largest number between user input and the minimum non-error size
    keep_inds = find(user_min_clust.<sizes)
    #keep_inds = Array{Int64, 1}()
    #for ind in first_keep_inds

    #keep_inds = [ind for ind in first_keep_inds if (1-cdf(Poisson(maximum(sizes)*(rough_radius/4)), sizes[ind])) *sum(kmer_vecs[1]) < p_value]

    #if there are no clusters to keep, add everything to the largest one
    if length(keep_inds) == 0 || length(keep_inds) == 1
        return [[i for i in 1:length(kmer_vecs)]]
    else
        final_indices = Array{Array{Int64, 1}, 1}()
        keep_centers = [center(kmer_vecs[indices[keep]]) for keep in keep_inds]

        for keep in keep_inds
            push!(final_indices, indices[keep])
        end

        for i in 1:length(indices) #go through every cluster index
            if !(i in keep_inds) #if that cluster is not one of the kept clusters
                C_arr = indices[i] #get the array of seq indices corresponding to the cluster
                for index in C_arr #go through every seq indix in the non-kept cluster
                    push!(final_indices[indmin([distfunc(kmer_vecs[index], center) for center in keep_centers])],
                            index)
                end
            end
        end
    end

    #if the input has gotten this far, it must be split. Call cluster_split on each current cluster
    for clust_ind in 1:length(final_indices)
        outer_clust = copy(final_indices[clust_ind])
        new_clusts = cluster_split(kmer_vecs[outer_clust], count, nuc_inds, rough_radius=rough_radius, fine_radius=fine_radius,
                                    min_devs=min_devs, distfunc=distfunc,center=center,verbose=verbose,
                                    cycle_lim=cycle_lim, triangle=triangle, p_value=p_value)
        for new_clust in new_clusts
            push!(final_indices, outer_clust[new_clust])
        end
        final_indices = deleteat!(final_indices, clust_ind)

    end
    return final_indices
end

"""
    get_non_homopolymers(dev_inds::Array{Int64, 1}, nuc_inds; min_devs = 6)

Takes in array of kmer indices. Returns the highest standard deviation non-homopolymer indices

"""
function get_non_homopolymers(dev_inds::Array{Int64, 1}, nuc_inds; min_devs=6)
    k=log(4, length(nuc_inds))
    len = length(dev_inds)
    dev_seqs = nuc_inds[dev_inds]
    to_remove = Array{Int64, 1}()
    #go through every high deviation kmer index
    for outer_ind in 1:len
        #make sure to git every combination only once
        for inner_ind in outer_ind + 1:len
            #align corresponding sequences and determine if they represent homopolymer differences
            if are_homopolymers(nw_align(dev_seqs[outer_ind], dev_seqs[inner_ind], mismatch_cost=-100.0), k=k)
                push!(to_remove, outer_ind, inner_ind)
            end
        end
    end
    #difference between all indices and the homopolymer indices is the non homopolymer indices
    new_dev_inds = setdiff(dev_inds, dev_inds[unique(to_remove)])

    if length(new_dev_inds) < min_devs #if there aren't enough indices, empty array says "don't split!"
        return []
    else
        return new_dev_inds[1:min_devs] #pick the top non-homopolymer kmer indices
    end
end

"""
    are_homopolymers(alignment::Tuple{String,String}; k = 6, polylen = 3)

Determines if the aligned kmers represent a homopolymer difference. If the
kmers represent a homopolymer difference, the function returns true.

"""
function are_homopolymers(alignment::Tuple{String,String};k=6,polylen::Int64=3)

    al1 = alignment[1]
    al2 = alignment[2]
    len = length(alignment[1])
    if len != k + 1 #homopolymers only differ by one
        return false
    end
    if polylen > k - 1
        return false
    end

    ind1 = findfirst(alignment[1], '-')#findfirst(isequal('-'), alignment[1])
    ind2 = findfirst(alignment[2], '-')#findfirst(isequal('-'), alignment[2])
    diff = ind1-ind2
    #determine if inds are only on edges
    #        XXXXX- (ind1)
    # (ind2) -XXXXX
    if abs(diff) == k #both gaps are on edges
        str = al1[1:(1+polylen)] * al2[1:(1+polylen)]
        if (length(unique(al1[1:(1+polylen)] * al2[1:(1+polylen)])) == 2) || (length(unique(al1[(len-polylen):len] * al2[(len-polylen):len])) == 2)
            return true
        else
            return false
        end
    elseif ind1 == 1 || ind2 == 1 #one gap on left edge, one on non-edge
        if ind1 == 1
            temp = ind2
        else
            temp = ind1
        end
        if length(unique(al1[1:(1+polylen)] * al2[1:(1+polylen)])) == 2 #edge box is true. THIS MEANS NO HOMOPOLY!!!
            return false

        else #edge box is false. May be a homopoly.
            if temp + polylen <= len && length(unique(al1[temp:(temp+polylen)]*al2[temp:(temp+polylen)])) == 2
                return true
            elseif temp - polylen >= 1 && length(unique(al1[(temp-polylen):temp]*al2[(temp-polylen):temp])) == 2
                return true
            end
        end
        if ind1 == 1
            al = al2
        else
            al=al1
        end
        #if you're here, it means that the non-edge gap is inside, not bordering, the homopoly diff
        if polylen + 2 < len
            if abs(diff) == 1 && length(unique(al[1:polylen+2])) == 2
                return true
            elseif abs(diff) == 2 && length(unique(al[2:polylen+3])) == 2 && length(unique(al1[[ind1, ind2]]*al2[[ind1, ind2]])) == 2
                return true
            end
        end
    else #one of the indices is k+1
        if ind1 == len
            temp = ind2
        else
            temp = ind1
        end
        if length(unique(al1[(len-polylen):len] * al2[(len-polylen):len])) == 2 #edge box is true. THIS MEANS NO HOMOPOLY!!!
            return false
        else #edge box is false. May be a homopoly

            if temp + polylen <= len && length(unique(al1[temp:(temp+polylen)]*al2[temp:(temp+polylen)])) == 2

                return true
            elseif temp - polylen >= 1 && length(unique(al1[(temp-polylen):temp]*al2[(temp-polylen):temp])) == 2
                return true
            end
        end
        if ind1 == len
            al=al2
        else
            al=al1
        end
        if polylen + 2 < len
            if abs(diff) == 1 && length(unique(al[len-polylen-1:len])) == 2
                return true
            elseif abs(diff) == 2 && length(unique(al[len-polylen-2:len-1])) == 2 && length(unique(al1[[ind1, ind2]]*al2[[ind1, ind2]])) == 2
                return true
            end
        end
    end
    #if the alignment did not meet any of these conditions, it must be a nonhomopolymer diff
    return false
    #=
    arr = [alignment[1]... , alignment[2]... ]
    #max homopolymer length difference is 2
    if length(findin(arr, '-')) > 2
        return false
    end

    len = length(alignment[1])

    #check left->right
    for i in 1:len
        if i+polylen > len #no chance of sufficiently large polymer, break
            break;
        end
        if (alignment[1][i] == '-') #test if this is start of homopolymer
            @inbounds temp = alignment[2][(i+1):(i+polylen)]
            temp2 = alignment[1][i:i+polylen]

            if length(unique(temp)) == 1 && length(unique(temp2))== 2
                return true
            end
        elseif (alignment[2][i] == '-')
            @inbounds temp = alignment[1][(i+1):(i+polylen)]
            temp2 = alignment[2][i:i+polylen]
            if length(unique(temp)) == 1 && length(unique(temp2))== 2
                return true
            end
        end
    end
    #check right->left
    for i in len:-1:1
        if i - polylen < 1 #no chance of sufficiently large polymer, return
            return false
        end
        if (alignment[1][i] == '-')
            @inbounds temp = alignment[2][(i-polylen):(i-1)]
            temp2 = alignment[1][i-polylen:i]
            if length(unique(temp)) == 1 && length(unique(temp2)) == 2
                return true
            end
        elseif (alignment[2][i] == '-')
            @inbounds temp = alignment[1][(i-polylen):(i-1)]
            temp2 = alignment[2][i-polylen:i]
            if length(unique(temp)) == 1 && length(unique(temp2)) == 2
                return true
            end
        end
    end
    return false
    =#
end

#==
      OLD VERSION
==#

"""
    alignment_consensus(seqs)

Computes a consensus sequence for array of given sequences that have already been aligned.
Consensus is computed by site-wise mode among given sequences.
"""
function alignment_consensus(seqs)
    string_lengths = [length(i) for i in seqs]
    ncols = length(seqs[1])
    if !all([x == ncols for x in string_lengths])
        error("Aligned sequences must all be the same length.")
    end
    result = string([mode([seq[i] for seq in seqs]) for i in 1:ncols]...)
    return degap(result)
end

"""
    coords(ref, read)

Returns an array of degapped coordinates, such that
coords(ref, read)[i] gives you the position the aligned read/ref
that matches the i'th ungapped position in ref.
"""
function coords(ref, read)
    if length(ref) != length(read)
        error("Aligned strings are meant to be the same length.")
    end
    degappedRef = degap(ref)
    coordMap = Array{Int64}(length(degappedRef))
    count = 1
    for i in 1:length(degappedRef)
        while ref[count] == '-'
            count += 1
        end
        coordMap[i] = count
        count += 1
    end
    return coordMap
end

"""
    get_centroid(reads, k, distfunc = corrected_kmer_dist)

Computes a rough cluster centroid by returning the sequence in `reads`
that has the nearest distance in `k`-mer space to the mean `k`-mer vector.
"""
function get_centroid(reads, k, distfunc = corrected_kmer_dist)
    kmerVecs = [kmer_count(read, k) for read in reads]
    meanVec = mean(kmerVecs)
    dists = [distfunc(kmerVec, meanVec, k=k) for kmerVec in kmerVecs]
    return reads[indmin(dists)]
end

"""
    consensus_seq(reads; thresh = 0.7, shift = 1, k = 6,
                       distfunc = corrected_kmer_dist)

Computes a cluster centroid by taking the nearest read to the mean `k`-mer
vector (see `get_centroid(reads, k, distfunc)`) and refining it (see `refine_ref(ref, reads)`).
`shift` determines the window size of comparison between sequences when refining rough centroid
locally (actual window size is `shift+1`).
"""
function consensus_seq(reads; degap_param = true, thresh = 0.7, shift = 1, k = 6,
                       distfunc = corrected_kmer_dist)
    ref = get_centroid(reads, k, distfunc)
    return refine_ref(ref, reads, degap_param = degap_param, thresh = thresh, shift = shift)
end

#=function get_matches(candidate_ref, reads, shift; kmer_align = true)
    alignments = []
    if kmer_align
        alignments = [kmer_seeded_align(candidate_ref, i) for i in reads]
    else
        alignments = [nw_align(candidate_ref, i) for i in reads]
    end
    maps = [coords(i...) for i in alignments]
    matchContent = [[degap(alignments[i][2][maps[i][k]:maps[i][k+shift]]) for i in 1:length(maps)] for k in 1:length(candidate_ref)-shift]
    matches = [freq(matchContent[k], degap(candidate_ref[k:k+shift])) for k in 1:length(matchContent)]
    return alignments, maps, matches, matchContent
end=#
#warn("Multithreaded version active as sole get_matches")

"""
    get_matches(candidate_ref, reads, shift; kmer_align = true)

Determine alignments between each sequence of `reads` and `candidate_ref`, along with coordinate mappings
and local matches. Used internally in `refine_ref` and `peak_split`.
"""
function get_matches(candidate_ref, reads, shift; degap_param = true, kmer_align = true)
    alignments = []
    if kmer_align
        alignments = map(i -> kmer_seeded_align(candidate_ref, i), reads)
    else
        alignments = map(i -> nw_align(candidate_ref, i), reads)
    end
    maps = [coords(i...) for i in alignments]

    if (degap_param)
        matchContent = [[degap(alignments[i][2][maps[i][k]:maps[i][k+shift]]) for i in 1:length(maps)] for k in 1:length(candidate_ref)-shift]

        matches = [freq(matchContent[k], degap(candidate_ref[k:k+shift])) for k in 1:length(matchContent)]

    else
       matchContent = [[(alignments[i][2][maps[i][k]:maps[i][k+shift]]) for i in 1:length(maps)] for k in 1:length(candidate_ref)-shift]

       matches = [freq(matchContent[k], candidate_ref[k:k+shift]) for k in 1:length(matchContent)]

    end

    return alignments, maps, matches, matchContent
end

"""
    refine_ref(candidate_ref, reads;  thresh = 0.7, shift = 1)

Takes a candidate consensus sequence, `candidate_ref`, and corresponding array of `reads` and refines the candidate via majority
votes from reads at local windows of size `shift+1`. If after alignment the frequency of a local region of the candidate
is less than `thresh`, this part of the candidate is refined.

"""
function refine_ref(candidate_ref, reads; degap_param = true, thresh = 0.7, shift = 1)
    # need to add in code to handle deletions at the start and end of
    # the candidate!

    alignments, maps, matches, _ = get_matches(candidate_ref, reads, shift, degap_param = degap_param)

    boolVec = BitArray{1}(1+zeros(length(candidate_ref)))
    for i in 1:length(matches)
        if matches[i] < thresh
            boolVec[i:i+shift] = false
        end
    end

    # need to explicitly handle the case where the reference is shorter than the reads.
    # refine front
    frontStrCol = [degap((alignments[j][2])[1:maps[j][1]-1]) for j in 1:length(maps)]
    frontNonEmpties = [length(s) != 0 for s in frontStrCol]
    frontConsensus = ""
    majority = length(alignments) / 2
    while countnz(frontNonEmpties) > majority
        frontStrCol = frontStrCol[frontNonEmpties]
        frontConsensus = mode(frontStrCol)
        frontAli, frontMaps, _, _ = get_matches(frontConsensus, frontStrCol, 1, degap_param = degap_param, kmer_align=false)
        frontStrCol = [degap(frontAli[j][2][1:frontMaps[j][1]-1]) for j in 1:length(frontMaps)]
        frontNonEmpties = [length(s) != 0 for s in frontStrCol]
        majority = length(frontAli) / 2
    end
    frontStr = frontConsensus

    # refine end
    endStrCol = [degap((alignments[j][2])[maps[j][end]+1:end]) for j in 1:length(maps)]
    endNonEmpties = [length(s) != 0 for s in endStrCol]
    endConsensus = ""
    majority = length(alignments) / 2
    while countnz(endNonEmpties) > majority
        endStrCol = endStrCol[endNonEmpties]
        endConsensus = mode(endStrCol)
        endAli, endMaps, _, _ = get_matches(endConsensus, endStrCol, 1, degap_param = degap_param, kmer_align=false)
        endStrCol = [degap(endAli[j][2][endMaps[j][end]+1:end]) for j in 1:length(endMaps)]
        endNonEmpties = [length(s) != 0 for s in endStrCol]
        majority = length(endAli) / 2
    end
    endStr = endConsensus

    # # if something breaks, CHECK THIS FIRST!
    if sum(1-boolVec) < 0.5
        return frontStr * candidate_ref * endStr
    end
    ranges = []
    i=1
    while i < length(boolVec)
        if !boolVec[i]
            start = i
            i += 1
            while !boolVec[i] && (i < length(boolVec))
                i += 1
            end
            push!(ranges,[start, i-1])
        end
        i += 1
    end

    newRef = candidate_ref[1:ranges[1][1]-1]
    for i in 1:length(ranges)-1
        low, high = ranges[i][1], ranges[i][2]
        insStr = mode([degap((alignments[j][2])[maps[j][low]:maps[j][high]])
                       for j in 1:length(maps)])
        newRef = newRef*insStr
        newRef = newRef*candidate_ref[ranges[i][2]+1:ranges[i+1][1]-1]
    end

    low, high = ranges[length(ranges)][1], ranges[length(ranges)][2]
    insStr = mode([degap((alignments[j][2])[maps[j][low]:maps[j][high]]) for j in 1:length(maps)])
    newRef = newRef*insStr
    newRef = newRef*candidate_ref[ranges[length(ranges)][2]+1:length(candidate_ref)]
    return frontStr*newRef*endStr
end

"""
    recursive_split(cluster, glom; centroid = "", shift = 7, minreads = 5, polylen = 3, k = 6, reads = nothing, phreds = nothing, names = nothing)

Recursively splits input `cluster` using `peak_split` and populates `glom` with the separated sub-clusters.
Splits clusters according to `centroid` (unless not given, in which case it is inferred).
`polylen` is lower bound for run-on subsequence to by considered a homopolymer region and not be split on.
Does not split on clusters of size less than `minReads`.
If `reads` is given as a string array, this will populate `reads` with tuples of `(cluster, phreds, names)` for
determining the actual clusters that are split, rather than just corresponding consensus sequences, where `phreds`
and `names` may be arrays of corresponding values, or `nothing`.
"""
function recursive_split(cluster::Array{String}, glom; centroid = "", shift = 7, minreads = 5, polylen = 3, k = 6, reads = nothing, phreds = nothing, names = nothing)
    consensus = ""
    if length(centroid) == 0
        consensus = consensus_seq(cluster, k=k)
    else
        consensus = refine_ref(centroid, cluster)
    end
    splitClus = peak_split(consensus, cluster, shift=shift, minReads=minreads, polylen=polylen)

    if length(splitClus) > 1
        for clust in splitClus
            split_phreds = phreds==nothing ? nothing : phreds[clust]
            split_names = names==nothing  ? nothing : names[clust]
            recursive_split(cluster[clust], glom, shift=shift, minreads=minreads, polylen=polylen, k=k, reads=reads, phreds=split_phreds, names=split_names)
        end
    else
        push!(glom, [consensus, length(cluster)])
        if reads != nothing
            push!(reads, (cluster, phreds, names))
        end
    end
end

"""
    peak_split(candidate_ref, reads; thresh = 0.7, shift = 3, minReads = 5, polylen = 3)

Creates and returns sub-clusters of `reads` split by variation from `candidateRef` if, locally, the most common read is less frequent than given `thresh`, default 0.7.
If the size of the input cluster `reads` is large, then `thresh` may be changed: thresh = max(1-10/length(reads), thresh).
Size of local window/site for comparison = `shift+1`.
Return values are arrays of boolean arrays, where each boolean array can be used to index into the input reads data to give the split clusters.
Determines  treshold for a cluster split, where if the candidate has local region with frequency less than `thresh` the cluster is split,
 to be `thresh = max(1-10/length(reads), thresh)` to account for larger cluster sizes.
`polylen` is lower bound for run-on subsequence to by considered a homopolymer region and not be split on.
Does not split on clusters of size less than `minReads`.
"""
function peak_split(candidate_ref, reads; thresh = 0.7, shift = 3, minReads = 5, polylen = 3)
    candidateRef = candidate_ref
    threshProp = thresh
    if length(reads) < minReads
        return [BitArray([true for r in reads])]
    end
    # Adjusting thresh is meant to account for the fact that 30% of a very large cluster is still very large,
    #  so we assume that, for large clusters, if we have 10 seqs, that is enough to split into a new cluster.
    # However, this will fail for very very large clusters that have 10 sequences that have the same mismatch base error.
    # Need to build in the expected number of mismatched base errors at any given base to control the false positive rate here.
    thresh = max(1-10/length(reads), thresh)
    numSteps = length(candidateRef) - shift
    #Need to add in code to handle deletions at the start and end of the candidate!
    alignments, maps, matches, matchContent = get_matches(candidateRef, reads, shift)
    boolVec = BitArray{1}(1 + zeros(length(candidateRef)))
    for i in 1:length(matches)
        if (matches[i] < thresh)
            boolVec[i:i+shift] = false
        end
    end
    # TODO: make ranges statically sized to avoid push!()
    ranges = []
    i = 1
    while i <= numSteps
        if !boolVec[i]
            start = i
            while !boolVec[i] && i < numSteps
                i += 1
            end
            # Julia subranges are inclusive, => window size = (shift+1)
            if shift+1 < polylen
                start = max(1, start - (polylen - (shift+1)))
                i = min(numSteps, i + (polylen - (shift+1)))
            end
            push!(ranges, [start, i-1])
        end
        i += 1
    end

    # add 2 to lengths to account for front and end as well
    scoreList = 1 + zeros(length(ranges) + 2)
    scoreInds = zeros(UInt32, length(ranges) + 2)
    for i in 1:length(ranges)
        low, high = ranges[i][1], ranges[i][2]
        sortedFreqs = sorted_freqs([degap(alignments[j][2][maps[j][low]:maps[j][high]]) for j in 1:length(maps)])
        # The following might happen with discrepancies at ends of reads; this is taken care of later in this function
        if length(sortedFreqs) == 1
            continue
        end
        # This will not handle the case where the third most frequent variant is a true mutation and the second is noise, but this case only occurs or is important when:
        #  1) the second most frequent, noisy read is more common than reads of a mutation (ie in a very long homopolymer)
        #  2) the third variant, the true mutation, doesn't cause a split ANYWHERE else, so that the mutation is basically a single base change at a homopolymer region (as opposed to a noisy indel)
        # So it is very uncommon for this case to happen but it might need to be handled.
        ali = nw_align(sortedFreqs[1][2], sortedFreqs[2][2])
        if !differ_by_just_one_gap(ali, polylen=polylen)
            scoreList[i] = minimum(matches[ranges[i][1]:ranges[i][2]])
            scoreInds[i] = indmin(matches[ranges[i][1]:ranges[i][2]]) + ranges[i][1] - 1
        end
    end

    # take care of splits at overhanging beginnings of reads, where consensus is empty ("")
    # if any overhang is more common than threshold, potentially split on most common overhang
    frontStrCol = [degap(alignments[j][2][1:maps[j][1]-1]) for j in 1:length(maps)]
    frontOverhangs = [length(s) != 0 for s in frontStrCol]
    if freq(frontStrCol, "") < thresh
        sortedFreqs = sorted_freqs([degap(alignments[j][2][1:maps[j][polylen]])
                    for j in 1:length(maps) if frontOverhangs[j]])
        if length(sortedFreqs) == 1
            scoreList[end-1] = freq(frontStrCol, mode(frontStrCol[frontOverhangs]))
            scoreInds[end-1] = 1
        else
            ali = nw_align(sortedFreqs[1][2][1:min(polylen, end)], sortedFreqs[2][2][1:min(polylen, end)])
            if !differ_by_just_one_gap(ali, polylen=polylen)
                scoreList[end-1] = freq(frontStrCol, mode(frontStrCol[frontOverhangs]))
                scoreInds[end-1] = 1
            end
        end
    end
    # take care of splits at overhanging ends
    endStrCol = [degap(alignments[j][2][maps[j][end]+1:end]) for j in 1:length(maps)]
    endOverhangs = [length(s) != 0 for s in endStrCol]
    if freq(endStrCol, "") < thresh
        sortedFreqs = sorted_freqs([degap(alignments[j][2][maps[j][end-polylen]:end])
                    for j in 1:length(maps) if endOverhangs[j]])
        if length(sortedFreqs) == 1
            scoreList[end] = freq(endStrCol, mode(endStrCol[endOverhangs]))
            scoreInds[end] = 1
        else
            ali = nw_align(sortedFreqs[1][2][max(end-polylen, 0):end], sortedFreqs[2][2][max(end-polylen, 0):end])
            if !differ_by_just_one_gap(ali, polylen=polylen)
                scoreList[end] = freq(endStrCol, mode(endStrCol[endOverhangs]))
                scoreInds[end] = 1
            end
        end
    end
    # TODO: maybe add singleton clusters back to cluster with closest distance
    if minimum(scoreList) < thresh
        ind = indmin(scoreList)
        splitpos = scoreInds[ind]
        # if in middle of sequence
        if ind < length(scoreInds) - 1
            unio = union(matchContent[splitpos])
            splitCluster = [matchContent[splitpos] .== s for s in unio]
            return splitCluster
        # else if on front
        elseif ind == length(scoreInds) - 1
            unio = union(frontStrCol)
            splitCluster = [frontStrCol .== s for s in unio]
            return splitCluster
        # else if on end
        else
            unio = union(endStrCol)
            splitCluster = [endStrCol .== s for s in unio]
            return splitCluster
        end
    else
        return [BitArray([true for r in reads])]
    end
end

"""
    differ_by_just_one_gap(alignment; polylen = 3)

`alignment` is an array containing two already aligned strings.
Returns true if first and second sequences in given array differ by
just one gap where that gap is in a homopolymer region.
A homopolymer region is determined to be a region of a single repeated
nucleotide of length at least `polylen`.
"""
function differ_by_just_one_gap(alignment; polylen = 3)
    gapDiffCount = 0
    otherDiffCount = 0
    gapInd = 0
    noGapSeq = 0
    for i in 1:length(alignment[1])
        if (alignment[1][i] == '-')
            gapDiffCount += 1
            gapInd = i
            noGapSeq = 2
        elseif (alignment[2][i] == '-')
            gapDiffCount += 1
            gapInd = i
            noGapSeq = 1
        elseif (alignment[1][i] != alignment[2][i])
            otherDiffCount += 1
        end
    end
    if otherDiffCount == 0 && gapDiffCount == 1
        # check if homopolymer region by matching regex to local string
        lo, hi = max(gapInd - (polylen-1), 1), min(gapInd + (polylen-1), length(alignment[1]))
        regex = Regex(string(alignment[noGapSeq][gapInd]) * "{" * string(polylen) * "}")
        return ismatch(regex, alignment[noGapSeq][lo:hi])
    else
        return false
    end
end

"""
    consensus_viz(candidate_ref, reads; thresh = 0.7, shift = 3,
                  intitle = "Consensus agreement.")

Creates a plot to visualize the agreement of a consensus sequence, `candidate_ref`,
with its cluster of `reads`.
Size of local window/site for comparison = `shift+1`.
I'm pretty sure `thresh` does nothing here.
"""
function consensus_viz(candidate_ref, reads;
                       thresh = 0.7, shift = 3, intitle = "Consensus agreement.")
    #Need to add in code to handle deletions at the start and end of the candidate!
    alignments, maps, matches, _ = get_matches(candidate_ref, reads, shift)
    frontStrCol = freq([degap((alignments[j][2])[1:maps[j][1]-1]) for j in 1:length(maps)], "")
    endStrCol = freq([degap((alignments[j][2])[maps[j][end]+1:end]) for j in 1:length(maps)], "")
    w, h = plt[:figaspect](0.1)
    figure(figsize=(w, h))
    plot(1:length(matches)+2, vcat([frontStrCol], matches,[endStrCol]))
    ax = gca()
    ax[:set_xlim]([-10, length(matches)+12])
    ax[:set_ylim]([0, 1])
    title(intitle)
    xlabel("Consensus Position")
    ylabel("Agreement")
end

"""
    disagreements(candidate_ref, reads; thresh = 0.7, shift = 3)

Prints local disagreements between `candidate_ref` and each sequence of `reads` after aligning.
Size of local window/site for comparison = `shift+1`.
A disagreement is a region where the candidate has a local region with frequency less than `thresh`.
"""
function disagreements(candidate_ref, reads; thresh = 0.7, shift = 3)
    # need to add in code to handle deletions at the start and end of
    # the candidate!
    alignments = [kmer_seeded_align(candidate_ref, i) for i in reads]
    maps = [coords(i...) for i in alignments]

    matches = [freq([(alignments[i][2])[maps[i][k]:maps[i][k+shift]]
                     for i in 1:length(maps)], candidate_ref[k:k+shift])
               for k in 1:length(candidate_ref)-shift]

    boolVec = BitArray{1}(1+zeros(length(candidate_ref)))
    for i in 1:length(matches)
        if matches[i] < thresh
            boolVec[i:i+shift] = boolVec[i:i+shift] & [false for k in 1:1+shift]
        end
    end

    ranges = []
    i = 1
    while i < length(boolVec)
        if !boolVec[i]
            start = i
            i += 1
            while !boolVec[i] && i < length(boolVec)
                i += 1
            end
            push!(ranges,[start, i-1])
        end
        i += 1
    end

    frontStrCol = [degap((alignments[j][2])[1:maps[j][1]-1]) for j in 1:length(maps)]
    if freq(frontStrCol, "") < thresh
        println("Before ref starts:")
        freq_dict_print(countmap(frontStrCol), thresh=1.1)
    end

    for i in 1:length(ranges)
        low, high = ranges[i][1], ranges[i][2]
        insStrCol = [degap((alignments[j][2])[maps[j][low]:maps[j][high]]) for j in 1:length(maps)]
        if freq(insStrCol, candidate_ref[low:high]) < thresh
            println("Ref pos ", low, ":", candidate_ref[low:high])
            freq_dict_print(countmap(insStrCol), thresh=1.1)
        end
    end

    endStrCol = [degap((alignments[j][2])[maps[j][end]+1:end]) for j in 1:length(maps)]
    if freq(endStrCol, "") < thresh
        println("After ref ends:")
        freq_dict_print(countmap(endStrCol), thresh=1.1)
    end
end

"""
    diff_in_homopolymer_region(alignment::Array{String, 1}; polylen=3)

Returns true if two aligned sequences differ only by single gaps in homopolyer regions
(ie one gap per region). `alignment` is an array of two strings that have already been aligned.
A homopolymer region is determined to be a region of a single repeated
nucleotide of length at least `polylen`.
"""
function diff_in_homopolymer_region(alignment::Array{String, 1}; polylen=3)
    for i in 1:length(alignment[1])
        if alignment[1][i] == '-' || alignment[2][i] == '-'
            lo = max(1, i - polylen + 1)
            hi = min(length(alignment[1]), i + polylen - 1)
            if !differ_by_just_one_gap([alignment[1][lo:hi], alignment[2][lo:hi]], polylen=polylen)
                return false
            end
        end
    end
    return true
end

"""
    merge_consensuses(seqs)

Takes in an array of 2-array [sequence, size] pairs and returns a list with
reach sequence unique with summed sizes among all previous copies.
"""
function merge_consensuses(seqs)
    tmp = sort(seqs, by=s->s[1]);
    i = 1
    while i < length(tmp)
        if tmp[i][1] == tmp[i+1][1]
            # consolidate cluster sizes
            tmp[i][2] += tmp[i+1][2]
            deleteat!(tmp, i+1)
        else
            i += 1
        end
    end
    return tmp
end

"""
    merge_diff_homopolymers(seqs; thresh = 0.002, k = 6, polylen = 3)

Takes in an array of 2-array [sequence, size] pairs and combines sequences that are identical
or differ by single gaps in homopolymer regions, summing respective frequencies. Slower than
naive `merge_consensuses(seqs)`.
A homopolymer region is determined to be a region of a single repeated
nucleotide of length at least `polylen`.
`thresh` gives an upper bound on distance in `k`-mer space for sequences to be compared for
homopolymer differences.
"""
function merge_diff_homopolymers(seqs; thresh = 0.002, k = 6, polylen = 3)
    # sort by decreasing size
    tmp = reverse(sort(seqs, by=s->s[2]))
    kmer_vecs = [kmer_count(s[1], k) for s in tmp]
    i = 1
    while i < length(tmp)
        j = i + 1
        while j < length(tmp)
            # if identical
            if tmp[i][1] == tmp[j][1]
                # consolidate cluster sizes
                tmp[i][2] += tmp[j][2]
                deleteat!(tmp, j)
                deleteat!(kmer_vecs, j)
            # if differ in homopolyer region
            elseif corrected_kmer_dist(kmer_vecs[i], kmer_vecs[j], k=k) < thresh &&
                    diff_in_homopolymer_region(kmer_seeded_align(tmp[i][1], tmp[j][1]), polylen=polylen)
                # consolidate cluster sizes
                tmp[i][2] += tmp[j][2]
                deleteat!(tmp, j)
                deleteat!(kmer_vecs, j)
            else
                j += 1
            end
        end
        i += 1
    end
    return tmp
end

"""
    get_coarse_centroid(seqs::Array{String, 1}; subsample = 1000, k = 4)

Returns a 'master consensus' representing largest coarse cluster of sequences,
computed among a subsample of `reads`. The consensus is the closest sequence of
`reads` to the mean `k`-mer vector in the largest computed cluster.
"""
function get_coarse_centroid(seqs::Array{String, 1}; subsample = 1000, k = 4)
    if subsample >= length(seqs)
        sampled_seqs = seqs
    else
        inds = shuffle([1:length(seqs)...])[1:subsample]
        sampled_seqs = seqs[inds]
    end
    kmer_vecs = [kmer_count(seq, k) for seq in sampled_seqs]
    @time _, sizes, _, centroid_inds = dp_centers(kmer_vecs, 0.15)
    return sampled_seqs[centroid_inds[indmax(sizes)]]
end


"""
    summarize(source::Data.Table, f::Formula)
    summarize!(sch::Data.Schema, source::Data.Table, f::Formula)

Compute summaries of the data from source as necessary to evaluate a formula.
Summaries are stored in the metadata of the schema for source.  This is mediated
by dispatching on the types of the terms and the types of the data via
`make_summarizer`.
"""
summarize(source::Data.Table, fs::Formula...) = summarize!(Data.schema(source), source, fs...)

function summarize!(sch::Data.Schema, source::Data.Table, fs::Formula...)
    terms = mapreduce(f->f.term.rhs, vcat, [], fs)
    types = Dict(Symbol(k) => Data.types(sch)[sch[k]] for k in Data.header(sch))
    summarizers = unique(reduce(vcat, [], make_summarizer(t, types) for t in terms))
    isempty(summarizers) && return sch
    for nt in Data.rows(source)
        for s in summarizers
            update!(s, nt)
        end
    end
    summaries = Dict()
    for s in summarizers
        push!(get!(Dict, summaries, name(s)), get(s))
    end
    sch.metadata[:summaries] = summaries
    return sch
end


# Functions for representing graphs in GraphViz's dot format
# http://www.graphviz.org/
# http://www.graphviz.org/Documentation/dotguide.pdf
# http://www.graphviz.org/pub/scm/graphviz2/doc/info/lang.html

# Write the dot representation of a graph to a file by name.
function to_dot(graph::AbstractGraph, filename::AbstractString, attrs::AttributeDict=AttributeDict())
    open(filename,"w") do f
        to_dot(graph, f, attrs)
    end
end

# Get the dot representation of a graph as a string.
function to_dot(graph::AbstractGraph, attrs::AttributeDict=AttributeDict())
    str = IOBuffer()
    to_dot(graph, str, attrs)
    takebuf_string(str)
end

# Write the dot representation of a graph to a stream.
function to_dot{G<:AbstractGraph}(graph::G, stream::IO,attrs::AttributeDict=AttributeDict())
    has_vertex_attrs = method_exists(attributes, (vertex_type(graph), G))
    has_edge_attrs = method_exists(attributes, (edge_type(graph), G))

    write(stream, "$(graph_type_string(graph)) graphname {\n")
    write(stream, "$(to_dot_graph(attrs))")
    if implements_edge_list(graph) && implements_vertex_map(graph)
        for vtx in  vertices(graph)
            attrs = has_vertex_attrs ?  "\t$(to_dot(attributes(vtx,graph)))" : ""
            write(stream,"$(vertex_index(vtx,graph))$attrs\n")
        end
        for edge in edges(graph)
            write(stream,"$(vertex_index(source(edge), graph)) $(edge_op(graph)) $(vertex_index(target(edge), graph))\n")
        end
    elseif implements_vertex_list(graph) && (implements_incidence_list(graph) || implements_adjacency_list(graph))
        for vertex in vertices(graph)
            if has_vertex_attrs && !isempty(attributes(vertex, graph))
                write(stream, "$(vertex_index(vertex, graph)) $(to_dot(attributes(vertex, graph)))\n")
            end
            if implements_incidence_list(graph)
                for e in out_edges(vertex, graph)
                    n = target(e, graph)
                    if is_directed(graph) || vertex_index(n, graph) > vertex_index(vertex, graph)
                        write(stream,"$(vertex_index(vertex, graph)) $(edge_op(graph)) $(vertex_index(n, graph))$(has_edge_attrs ? string(" ", to_dot(attributes(e, graph))) : "")\n")
                    end
                end
            else # implements_adjacency_list
                for n in out_neighbors(vertex, graph)
                    if is_directed(graph) || vertex_index(n, graph) > vertex_index(vertex, graph)
                        write(stream,"$(vertex_index(vertex, graph)) $(edge_op(graph)) $(vertex_index(n,graph))\n")
                    end
                end
            end
        end
    else
        throw(ArgumentError("More graph Concepts needed: dot serialization requires iteration over edges or iteration over vertices and neighbors."))
    end
    write(stream, "}\n")
    stream
end
#write node attributes example: [shape=box, style=filled]
function to_dot(attrs::AttributeDict)
    if isempty(attrs)
        ""
    else
        string("[",join(map(a -> to_dot(a[1], a[2]), collect(attrs)),","),"]")
    end
end
# write a graph wide attributes example: size = "4,4";
function to_dot_graph(attrs::AttributeDict)
    if isempty(attrs)
        ""
    else
        string(join(map(a -> to_dot(a[1], a[2]), collect(attrs)),";\n"),";\n")
    end
end

to_dot(attr::AbstractString, value) = "\"$attr\"=\"$value\""

to_dot(attr_tuple::@compat Tuple{String, Any}) = "\"$(attr_tuple[1])\"=\"$(attr_tuple[2])\""

function graph_type_string(graph::AbstractGraph)
    is_directed(graph) ? "digraph" : "graph"
end

function edge_op(graph::AbstractGraph)
    is_directed(graph) ? "->" : "--"
end

function plot(g::AbstractGraph)
    stdin, proc = open(`neato -Tx11`, "w")
    to_dot(g, stdin)
    close(stdin)
end


using Random
traj_folder = joinpath(dirname(pathof(TrajectoryOptimization)),"..")
urdf_folder = joinpath(traj_folder, "dynamics","urdf")
urdf_kuka_orig = joinpath(urdf_folder, "kuka_iiwa.urdf")
urdf_kuka = joinpath(urdf_folder, "temp","kuka.urdf")

function write_kuka_urdf()
    kuka_mesh_dir = joinpath(TrajectoryOptimization.root_dir(),"dynamics","urdf","kuka_iiwa_mesh")
    temp_dir = joinpath(TrajectoryOptimization.root_dir(),"dynamics","urdf","temp")
    if !isdir(temp_dir)
        mkdir(temp_dir)
    end
    open(urdf_kuka_orig,"r") do f
        open(urdf_kuka, "w") do fnew
            for ln in eachline(f)
                pre = findfirst("<mesh filename=",ln)
                post = findlast("/>",ln)
                if !(pre isa Nothing) && !(post isa Nothing)
                    inds = pre[end]+2:post[1]-2
                    pathstr = ln[inds]
                    file = splitdir(pathstr)[2]
                    ln = ln[1:pre[end]+1] * joinpath(kuka_mesh_dir,file) * ln[post[1]-1:end]
                end
                println(fnew,ln)
            end
        end
    end
end

function get_kuka_ee(kuka)
    ee_body = findbody(kuka, "iiwa_link_ee")
    ee_point = Point3D(default_frame(ee_body),0.,0.,0.)
    return ee_body, ee_point
end

function get_kuka_ee_postition_fun(kuka::Mechanism,statecache=StateCache(kuka)) where {O}
    ee_body, ee_point = Dynamics.get_kuka_ee(kuka)
    world = root_frame(kuka)
    nn = num_positions(kuka)

    function ee_position(x::AbstractVector{T}) where T
        state = statecache[T]
        set_configuration!(state, x[1:nn])
        RigidBodyDynamics.transform(state, ee_point, world).v
    end
end

function calc_ee_position(kuka::Mechanism,X::Trajectory)
    ee = zero.(X)
    N = length(X)
    state = MechanismState(kuka)
    world = root_frame(kuka)
    ee_point = get_kuka_ee(kuka)[2]
    nn = num_positions(kuka)
    for k = 1:N
        set_configuration!(state, X[k][1:nn])
        ee[k] = RigidBodyDynamics.transform(state, ee_point, world).v
    end
    return ee
end

function kuka_ee_ik(kuka::Mechanism,point::Vector,ik_iterations=1000,attempts=20,tol=1e-2)
    state = MechanismState(kuka)
    world = root_frame(kuka)

    # Get end-effector
    ee_body, ee_point = get_kuka_ee(kuka)

    # Run IK
    err = Inf
    iter = 1
    while err > tol
        rand!(state)
        goal = Point3D(world,point)
        ik_res = jacobian_transpose_ik!(state,ee_body,ee_point,goal,iterations=ik_iterations)
        point_res = RigidBodyDynamics.transform(ik_res,ee_point,world).v
        err = norm(point-point_res)
        if iter > attempts
            error("IK cannot get sufficiently close to the goal")
        end
        return ik_res
    end
end

function jacobian_transpose_ik!(state::MechanismState,
                               body::RigidBody,
                               point::Point3D,
                               desired::Point3D;
                               α=0.1,
                               iterations=100)
    mechanism = state.mechanism
    world = root_frame(mechanism)

    # Compute the joint path from world to our target body
    p = path(mechanism, root_body(mechanism), body)
    # Allocate the point jacobian (we'll update this in-place later)
    Jp = point_jacobian(state, p, RigidBodyDynamics.transform(state, point, world))

    q = copy(configuration(state))

    for i in 1:iterations
        # Update the position of the point
        point_in_world = RigidBodyDynamics.transform(state, point, world)
        # Update the point's jacobian
        point_jacobian!(Jp, state, p, point_in_world)
        # Compute an update in joint coordinates using the jacobian transpose
        Δq = α * Array(Jp)' * (RigidBodyDynamics.transform(state, desired, world) - point_in_world).v
        # Apply the update
        q .= configuration(state) .+ Δq
        set_configuration!(state, q)
    end
    state
end

function hold_trajectory(n,m,N, mech::Mechanism, q)
    state = MechanismState(mech)
    nn = num_positions(state)
    set_configuration!(state, q[1:nn])
    vd = zero(state.q)
    u0 = dynamics_bias(state)

    if length(q) > m
        throw(ArgumentError("system must be fully actuated to hold an arbitrary position ($(length(q)) should be > $m)"))
    end
    U0 = zeros(m,N)
    for k = 1:N
        U0[:,k] = u0
    end
    return U0
end

# Write new urdf file with correct absolute paths
write_kuka_urdf()

kuka = Model(urdf_kuka)
end_effector_function = Dynamics.get_kuka_ee_postition_fun(parse_urdf(urdf_kuka,remove_fixed_tree_joints=false))


using ArrayFire
using AutoGrad
importall Base
using AutoGrad: matmul2arg,broadcast2arg,unbroadcast,id2,math2arg,float2arg,broadcast2cmp
import AutoGrad: tofloat,_dbg

tofloat{T<:AbstractFloat}(a::AFArray{T})=a
_dbg{T}(x::AFArray{T})=Symbol("AF$(join([id2(x),size(x)...,8*sizeof(T)],'_'))")
typealias AForN Union{AFArray,Number}
for (f,d) in matmul2arg
    @eval @primitive $f(x1::AFArray,x2::AFArray)::y $(d[1]) $(d[2])
end
for (f,g) in broadcast2arg
    @eval @primitive $f(x1::AFArray,x2::AFArray)::y  unbroadcast(y,x1,dy->dy.*$(g[1]))  unbroadcast(y,x2,dy->dy.*$(g[2]))
    @eval @primitive $f(x1::Number,x2::AFArray)::y  unbroadcast(y,x1,dy->dy.*$(g[1]))  unbroadcast(y,x2,dy->dy.*$(g[2]))
    @eval @primitive $f(x1::AFArray,x2::Number)::y  unbroadcast(y,x1,dy->dy.*$(g[1]))  unbroadcast(y,x2,dy->dy.*$(g[2]))
end
for (f,g) in math2arg
    @eval @primitive $f(x1::AFArray,x2::AFArray)::y  unbroadcast(y,x1,dy->dy.*$(g[1]))  unbroadcast(y,x2,dy->dy.*$(g[2]))
    @eval @primitive $f(x1::Number,x2::AFArray)::y  unbroadcast(y,x1,dy->dy.*$(g[1]))  unbroadcast(y,x2,dy->dy.*$(g[2]))
    @eval @primitive $f(x1::AFArray,x2::Number)::y  unbroadcast(y,x1,dy->dy.*$(g[1]))  unbroadcast(y,x2,dy->dy.*$(g[2]))
end
for (f,g) in float2arg
    @eval @primitive $f(x1::AFArray,x2::AFArray)::y unbroadcast(y,x1,$(g[1])) unbroadcast(y,x2,$(g[2]))
    @eval @primitive $f(x1::Number,x2::AFArray)::y unbroadcast(y,x1,$(g[1])) unbroadcast(y,x2,$(g[2]))
    @eval @primitive $f(x1::AFArray,x2::Number)::y unbroadcast(y,x1,$(g[1])) unbroadcast(y,x2,$(g[2]))
end
for f in broadcast2cmp
    @eval begin
        @zerograd $f(x1::AFArray,x2::AFArray)
        @zerograd $f(x1::AFArray,x2::Number)
        @zerograd $f(x1::Number,x2::AFArray)
    end
end

Base.eltype{T}(x::AFArray{T})=T
@primitive  sum(x::AFArray)  (dy->convert(eltype(x),dy).+zeros(x))
@primitive  sum(x::AFArray,i...)  (dy->dy.+zeros(x))
Base.size(x::AFArray,i)=(if i>ndims(x); 1; else; size(x)[i]; end)

include("profile.jl")
d0af = map(i->map(AFArray,i),d0)
# w1af = map(AFArray,w1)
w2af = map(AFArray,w2)
# timeall(w2af,d0af,10)
# loop(fun[9],w2af,d0af,10)
# (x1af,y1af)=first(d0af)
# fun[10](w2af,x1af,y1af)

function timeall_af(w=w2af,d=d0af,t=10)
    for i=1:length(fun)
        print(i); printfun(fun[i])
        for j=1:2
            sleep(2)
            gc_enable(false)
            @time (loop_af(fun[i],w,d,t); sync())
            gc_enable(true)
            @show deviceMemInfo()
        end
    end
end

function loop_af(f,w,d,t)
    for i in 1:t
        for (x,y) in d
            # ArrayFire.af_device_gc()
            f(w,x,y)
        end
    end
end

gc_enable(false); @time loop_af(fun[9],w2af,d0af,10); gc_enable(true)
gc_enable(false); @time loop_af(fun[9],w2af,d0af,10); gc_enable(true)
gc_enable(false); @time loop_af(fun[9],w2af,d0af,10); gc_enable(true)
gc_enable(false); @time loop_af(fun[9],w2af,d0af,10); gc_enable(true)
gc_enable(false); @time loop_af(fun[9],w2af,d0af,10); gc_enable(true)
gc_enable(false); @profile loop_af(fun[9],w2af,d0af,10); gc_enable(true)
Profile.print()
Profile.print(format=:flat)


while true
       start = 0
       space = 110
       x = fill('.', space)
       i = 0

       d = Dict{String,Char}()
       open("day12.txt") do f
       for (k,v) in enumerate(split(readline(f), ""))
       x[1, k + 5] = v[1]
       end
       for line in eachline(f)
       d[line[1:5]] = line[7]
       end
       end
       y = deepcopy(x)

       while i < 50000000000
       x = deepcopy(y)
       y = fill('.', space)
       for j in 1:146
       b = string(x[j:j + 4]...)
       y[j + 2] = (b in keys(d) ? d[b] : '.')
       end

       if any(==('#'), y[1:4])
       space += 2; start -= 2
       else
       space -= 2; start += 2
       any(==('#'), y[end - 4:end]) && (space += 2)

       end; s = 0
       for i in 1:150
       x[21,i] == '#' && (s += i - 26)
       end; s
       i += 1; i > 50000000000 && break