#=

ESSE wrapper

Ignacio Quintero Mächler

t(-_-t)

September 26 2017

=#

"""
    esse(states_file ::String,
         tree_file   ::String,
         envdata_file::String,
         cov_mod     ::NTuple{M,String},
         out_file    ::String,
         h           ::Int64;
         constraints ::NTuple{N,String}  = (" ",),
         mvpars      ::NTuple{O,String}  = ("lambda = beta",),
         niter       ::Int64             = 10_000,
         nthin       ::Int64             = 10,
         nburn       ::Int64             = 200,
         nchains     ::Int64             = 1,
         ntakew      ::Int64             = 100,
         winit       ::Float64             = 2.0,
         scale_y     ::NTuple{2,Bool}    = (true, false),
         algorithm   ::String            = "pruning",
         λpriors     ::Float64           = .1,
         μpriors     ::Float64           = .1,
         gpriors     ::Float64           = .1,
         lpriors     ::Float64           = .1,
         qpriors     ::Float64           = .1,
         βpriors     ::NTuple{2,Float64} = (0.0, 10.0),
         hpriors     ::Float64           = .1,
         optimal_w   ::Float64           = 0.8,
         screen_print::Int64             = 5,
         Eδt         ::Float64           = 1e-3,
         ti          ::Float64           = 0.0,
         ρ           ::Array{Float64,1}  = [1.0]) where {M,N,O}

Wrapper for running a SSE model from file.
"""
function esse(states_file ::String,
              tree_file   ::String,
              envdata_file::String,
              cov_mod     ::NTuple{M,String},
              out_file    ::String,
              h           ::Int64;
              constraints ::NTuple{N,String}  = (" ",),
              mvpars      ::NTuple{O,String}  = ("lambda = beta",),
              niter       ::Int64             = 10_000,
              nthin       ::Int64             = 10,
              nburn       ::Int64             = 200,
              nchains     ::Int64             = 1,
              ntakew      ::Int64             = 100,
              winit       ::Float64             = 2.0,
              scale_y     ::NTuple{2,Bool}    = (true, false),
              algorithm   ::String            = "pruning",
              λpriors     ::Float64           = .1,
              μpriors     ::Float64           = .1,
              gpriors     ::Float64           = .1,
              lpriors     ::Float64           = .1,
              qpriors     ::Float64           = .1,
              βpriors     ::NTuple{2,Float64} = (0.0, 10.0),
              hpriors     ::Float64           = .1,
              optimal_w   ::Float64           = 0.8,
              screen_print::Int64             = 5,
              Eδt         ::Float64           = 1e-3,
              ti          ::Float64           = 0.0,
              ρ           ::Array{Float64,1}  = [1.0]) where {M,N,O}

  # read data
  tv, ed, el, bts, x, y =
    read_data_esse(states_file, tree_file, envdata_file)

  @info "Data for $(length(tv)) species successfully read"

  # scale y
  if scale_y[1]
    # if scale each function separately or together
    if scale_y[2]
      ymin = minimum(y)
      ymax = maximum(y)
      for j in axes(y,2), i in axes(y,1)
        y[i,j] = (y[i,j] - ymin)/(ymax - ymin)
      end
    else
      ymin = minimum(y, dims = 1)
      ymax = maximum(y, dims = 1)
      for j in axes(y,2), i in axes(y,1)
        y[i,j] = (y[i,j] - ymin[j])/(ymax[j] - ymin[j])
      end
    end
  end

  # prepare data
  X, p, fp, trios, ns, ned, pupd, phid, nnps, nps, mvps, nngps, mvhfs,
  dcp, dcfp, pardic, k, h, ny, model, af!, assign_hidfacs!, abts, bts, E0 =
        prepare_data(cov_mod, tv, x, y, ed, el, ρ, h, constraints, mvpars)

  @info "Data successfully prepared"

  ## make likelihood function
  # flow algorithm
  if occursin(r"^[f|F][A-za-z]*", algorithm)

    # prepare likelihood
    Gt, Et, lbts, nets, λevent!, rootll =
      prepare_ll(p, bts, E0, k, h, ny, ns, ned, model, Eδt, ti, abts, af!)

    # make likelihood function
    llf = make_loglik(Gt, Et, X, trios, lbts, bts, ns, ned, nets,
                      λevent!, rootll)
  # prunning algorithm
  elseif occursin(r"^[p|P][A-za-z]*", algorithm)

    # prepare likelihood
    X, int, λevent!, rootll, abts1, abts2 =
      prepare_ll(X, p, E0, ns, k, h, ny, model, abts ,af!)

    # make likelihood function
    llf = make_loglik(X, abts1, abts2, trios, int,
      λevent!, rootll, k, h, ns, ned)

  else
    @error "No matching likelihood for algorithm: $algorithm"
  end

  # create prior function
  lpf = make_lpf(pupd, phid,
    λpriors, μpriors, gpriors, lpriors, qpriors, βpriors, hpriors,
    k, h, ny, model)

  # create posterior functions
  lhf = make_lhf(llf, lpf, assign_hidfacs!, dcp, dcfp)

  # number of parameters
  npars = length(pardic)

  # number of samples
  nlogs = fld(niter,nthin)

  # if parallel
  if nchains > 1
    # where to write in the Shared Array
    cits = [(1+j):(nlogs+j) for j in 0:nlogs:(nchains-1)*nlogs]

    # run slice-sampling in parallel
    R = SharedArray{Float64,2}(nlogs*nchains, npars+2)

    # run parallel loop
    @sync @distributed for ci in Base.OneTo(nchains)
      R[cits[ci],:] =
        slice_sampler(lhf, p, fp, nnps, nps, phid, mvps, nngps, mvhfs, npars,
          niter, nthin, nburn, ntakew, winit, optimal_w, screen_print)
    end

    # write output
    write_ssr(R, pardic, out_file, cits)
  else

    R = slice_sampler(lhf, p, fp, nnps, nps, phid, mvps, nngps, mvhfs, npars,
          niter, nthin, nburn, ntakew, winit, optimal_w, screen_print)

    # write output
    write_ssr(R, pardic, out_file)
  end

  return R
end

"""
    esse(tv          ::Dict{Int64,Array{Float64,1}},
         ed          ::Array{Int64,2},
         el          ::Array{Float64,1},
         x           ::Array{Float64,1},
         y           ::Array{Float64,L},
         cov_mod     ::NTuple{M,String},
         out_file    ::String,
         h           ::Int64;
         constraints ::NTuple{N,String}  = (" ",),
         mvpars      ::NTuple{O,String}  = ("lambda = beta",),
         niter       ::Int64             = 10_000,
         nthin       ::Int64             = 10,
         nburn       ::Int64             = 200,
         nchains     ::Int64             = 1,
         ntakew      ::Int64             = 100,
         winit       ::Float64             = 2.0,
         scale_y     ::NTuple{2,Bool}    = (true, false),
         algorithm   ::String            = "pruning",
         λpriors     ::Float64           = .1,
         μpriors     ::Float64           = .1,
         gpriors     ::Float64           = .1,
         lpriors     ::Float64           = .1,
         qpriors     ::Float64           = .1,
         βpriors     ::NTuple{2,Float64} = (0.0, 10.0),
         hpriors     ::Float64           = .1,
         optimal_w   ::Float64           = 0.8,
         screen_print::Int64             = 5,
         Eδt         ::Float64           = 1e-3,
         ti          ::Float64           = 0.0,
         ρ           ::Array{Float64,1}  = [1.0]) where {L,M,N,O}

Wrapper for running a SSE model from simulations.
"""
function esse(tv          ::Dict{Int64,Array{Float64,1}},
              ed          ::Array{Int64,2},
              el          ::Array{Float64,1},
              x           ::Array{Float64,1},
              y           ::Array{Float64,L},
              cov_mod     ::NTuple{M,String},
              out_file    ::String,
              h           ::Int64;
              constraints ::NTuple{N,String}  = (" ",),
              mvpars      ::NTuple{O,String}  = ("lambda = beta",),
              niter       ::Int64             = 10_000,
              nthin       ::Int64             = 10,
              nburn       ::Int64             = 200,
              nchains     ::Int64             = 1,
              ntakew      ::Int64             = 100,
              winit       ::Float64             = 2.0,
              scale_y     ::NTuple{2,Bool}    = (true, false),
              algorithm   ::String            = "pruning",
              λpriors     ::Float64           = .1,
              μpriors     ::Float64           = .1,
              gpriors     ::Float64           = .1,
              lpriors     ::Float64           = .1,
              qpriors     ::Float64           = .1,
              βpriors     ::NTuple{2,Float64} = (0.0, 10.0),
              hpriors     ::Float64           = .1,
              optimal_w   ::Float64           = 0.8,
              screen_print::Int64             = 5,
              Eδt         ::Float64           = 1e-3,
              ti          ::Float64           = 0.0,
              ρ           ::Array{Float64,1}  = [1.0]) where {L,M,N,O}

  # prepare data
  X, p, fp, trios, ns, ned, pupd, phid, nnps, nps, mvps, nngps, mvhfs,
  dcp, dcfp, pardic, k, h, ny, model, af!, assign_hidfacs!, abts, bts, E0 =
        prepare_data(cov_mod, tv, x, y, ed, el, ρ, h, constraints, mvpars)

  @info "Data successfully prepared"

  ## make likelihood function
  # flow algorithm
  if occursin(r"^[f|F][A-za-z]*", algorithm)

    # prepare likelihood
    Gt, Et, lbts, nets, λevent!, rootll =
      prepare_ll(p, bts, E0, k, h, ny, ns, ned, model, Eδt, ti, abts, af!)

    # make likelihood function
    llf = make_loglik(Gt, Et, X, trios, lbts, bts, ns, ned, nets,
                      λevent!, rootll)
  # pruning algorithm
  elseif occursin(r"^[p|P][A-za-z]*", algorithm)

    # prepare likelihood
    X, int, λevent!, rootll, abts1, abts2 =
      prepare_ll(X, p, E0, ns, k, h, ny, model, abts ,af!)

    # make likelihood function
    llf = make_loglik(X, abts1, abts2, trios, int,
      λevent!, rootll, k, h, ns, ned)

  else
    @error "No matching likelihood for algorithm: $algorithm"
  end

  # create prior function
  lpf = make_lpf(pupd, phid,
    λpriors, μpriors, gpriors, lpriors, qpriors, βpriors, hpriors,
    k, h, ny, model)

  # create posterior functions
  lhf = make_lhf(llf, lpf, assign_hidfacs!, dcp, dcfp)

  # number of parameters
  npars = length(pardic)

  # number of samples
  nlogs = fld(niter,nthin)

  # if parallel
  if nchains > 1
    # where to write in the Shared Array
    cits = [(1+j):(nlogs+j) for j in 0:nlogs:(nchains-1)*nlogs]

    # run slice-sampling in parallel
    R = SharedArray{Float64,2}(nlogs*nchains, npars+2)

    # run parallel loop
    @sync @distributed for ci in Base.OneTo(nchains)
      R[cits[ci],:] =
        slice_sampler(lhf, p, fp, nnps, nps, phid, mvps, nngps, mvhfs, npars,
          niter, nthin, nburn, ntakew, winit, optimal_w, screen_print)
    end

    # write output
    write_ssr(R, pardic, out_file, cits)
  else

    R = slice_sampler(lhf, p, fp, nnps, nps, phid, mvps, nngps, mvhfs, npars,
          niter, nthin, nburn, ntakew, winit, optimal_w, screen_print)

    # write output
    write_ssr(R, pardic, out_file)
  end

  return R
end

"""
    read_data_esse(states_file ::String,
                   tree_file   ::String,
                   envdata_file::String)

Process tree and state and environmental data file to run ESSE.
"""
function read_data_esse(states_file ::String,
                        tree_file   ::String,
                        envdata_file::String)

  # read tree in postorder and assign to objects
  tree, bts = read_tree(tree_file, order = "postorder", branching_times = true)
  ntip = tree.nnod + 1
  ed   = tree.ed
  el   = tree.el
  tlab = tree.tlab

  # assign tip labels to edge numbers
  tip_labels = Dict{String,Integer}()
  ii = 0
  for i in Base.OneTo(size(ed,1))
    if ed[i,2] <= ntip
      ii += 1
      tip_labels[tlab[ii]] = ed[i,2]
    end
  end

  # read states text file
  data = readdlm(states_file)

  if size(data,1) != ntip
    data = readdlm(states_file, '\t', '\r')
  end

  if size(data,1) != ntip
    data = readdlm(states_file, '\t', '\n')
  end

  if size(data,1) != ntip
    error("Data file cannot be made of the right dimensions.\n Make sure the data file has the same number of rows as tips in the tree")
  end

  data_tlab    = convert(Array{String,1}, data[:,1])
  data_states  = convert(Array{Float64,2},  data[:,2:end])

  # create dictionary
  tip_states = Dict(tip_labels[val] => data_states[i,:]
                   for (i,val) = enumerate(data_tlab))

  # process environmental data file
  envdata = readdlm(envdata_file)

  x = envdata[:,1]
  y = envdata[:,2:end]

  return tip_states, ed, el, bts, x, y
end

"""
  write_ssr(R       ::Array{Float64,2},
            pardic  ::Dict{String,Int64},
            out_file::String)

Write the samples from an MC sampler data frame
given a Dictionary of parameters.
"""
function write_ssr(R       ::Array{Float64,2},
                   pardic  ::Dict{String,Int64},
                   out_file::String)

  # column names
  col_nam = ["Iteration", "Posterior"]

  for (k,v) in sort!(collect(pardic), by = x -> x[2])
    push!(col_nam, k)
  end

  R = vcat(reshape(col_nam, 1, lastindex(col_nam)), R)

  writedlm(out_file*".log", R)
end

"""
  write_ssr(R       ::SharedArray{Float64,2},
            pardic  ::Dict{String,Int64},
            out_file::String,
            cits    ::Array{UnitRange{Int64},1})

Write the samples from multiple chains of MC sampler data frame
given a Dictionary of parameters.
"""
function write_ssr(R       ::SharedArray{Float64,2},
                   pardic  ::Dict{String,Int64},
                   out_file::String,
                   cits    ::Array{UnitRange{Int64},1})

  # column names
  col_nam = ["Iteration", "Posterior"]

  for (k,v) in sort!(collect(pardic), by = x -> x[2])
    push!(col_nam, k)
  end

  for ci in Base.OneTo(length(cits))
    ri = vcat(reshape(col_nam, 1, lastindex(col_nam)), R[cits[ci],:])
    writedlm(out_file*"_chain_$ci.log", ri)
  end

end


module Torch

using PyCall

using ChainRulesCore
using DLPack
using Functors: @functor
using Adapt

using ..PyCallChainRules: PyAdaptor, fmap

const inspect = PyNULL()
const torch = PyNULL()
const functorch = PyNULL()
const dlpack = PyNULL()
const ispysetup = Ref{Bool}(false)

pyto_dlpack(x) = @pycall dlpack.to_dlpack(x)::PyObject
pyfrom_dlpack(x) = @pycall dlpack.from_dlpack(x)::PyObject

struct TorchModuleWrapper
    torch_stateless_module::PyObject
    dtype::PyObject
    params::Tuple
    buffers::Tuple
end

@functor TorchModuleWrapper (params,)

Base.show(io::IO, f::TorchModuleWrapper) = print(io, f.torch_stateless_module, " ", f.dtype, " ", " params size=", size.(f.params))

Base.length(f::TorchModuleWrapper) = length(f.params)
Base.iterate(f::TorchModuleWrapper) = iterate(f.params)
Base.iterate(f::TorchModuleWrapper, state) = iterate(f.params, state)

function TorchModuleWrapper(torch_module)
    pybuiltin("isinstance")(torch_module, torch.nn.Module) || error("Not a torch.nn.Module")
    funmod, params, buffers = functorch.make_functional_with_buffers(torch_module)
    dtype = params[1].dtype
    jlparams = map(params) do x
        DLPack.wrap(x, pyto_dlpack)
    end
    return TorchModuleWrapper(funmod, dtype, jlparams, buffers)
end

function (wrap::TorchModuleWrapper)(args...; kwargs...)
    params = wrap.params
    out = wrap.torch_stateless_module(Tuple(map(x -> DLPack.share(x, PyObject, pyfrom_dlpack).requires_grad_(true), params)),
        wrap.buffers, fmap(x -> DLPack.share(x, PyObject, pyfrom_dlpack), args)...; kwargs...)
    res = fmap(x->DLPack.wrap(x, pyto_dlpack), out)
    return res
end

function ChainRulesCore.rrule(wrap::TorchModuleWrapper, args...; kwargs...)
    T = typeof(first(wrap.params))
    params = wrap.params
    pyparams = Tuple(map(x -> DLPack.share(x, PyObject, pyfrom_dlpack).requires_grad_(true), params))
    pyargs = fmap(x -> DLPack.share(x, PyObject, pyfrom_dlpack).requires_grad_(true), args)

    torch_primal, torch_vjpfun = functorch.vjp(py"buffer_implicit"(wrap.torch_stateless_module, wrap.buffers), pyparams, pyargs...; kwargs...)
    project = ProjectTo(args)
    function TorchModuleWrapper_pullback(Δ)
        cΔ = fmap(x->Adapt.adapt(PyAdaptor{T}(), x), Δ)
        pycΔ = fmap(x->DLPack.share(x, PyObject, pyfrom_dlpack), cΔ)
        torch_tangent_vals = torch_vjpfun(pycΔ)
        jlparams_tangents = map(x -> DLPack.wrap(x, pyto_dlpack), torch_tangent_vals[1])
        args_tangents = project(fmap(x -> DLPack.wrap(x, pyto_dlpack), torch_tangent_vals[2:end]))
        return (Tangent{TorchModuleWrapper}(; torch_stateless_module = NoTangent(), dtype = NoTangent(), params = jlparams_tangents, buffers = NoTangent()), args_tangents...)
    end
    res = fmap(x->DLPack.wrap(x, pyto_dlpack), torch_primal)
    return res, TorchModuleWrapper_pullback
end

function __init__()
    try
        copy!(torch, pyimport("torch"))
        copy!(dlpack, pyimport("torch.utils.dlpack"))
        copy!(functorch, pyimport("functorch"))
        copy!(inspect, pyimport("inspect"))
        ispysetup[] = true
        py"""
        def buffer_implicit(fn, buffers):
            def newfn(params, *inputs, **kwargs):
                return fn(params, buffers, *inputs, **kwargs)

            return newfn
        """
    catch err
        @warn """PyCallChainRules.jl has failed to import torch and functorch from Python.
                 Please make sure these are installed.
        """
        @debug err
        ispysetup[] = false
        #rethrow(err)
    end
end

end


"""
	ALConstraintSet{T}

An [`AbstractConstraintSet`](@ref) that stores the constraint values as well as Lagrange
multiplier and penalty terms for each constraint.

The cost associated with constraint terms in the augmented Lagrangian can be evaluated for

	cost!(J::Vector, ::ALConstraintSet)

which adds the cost at each time step to the vector `J` of length `N`.

The cost expansion for these terms is evaluated along the trajectory `Z` using

	cost_expansion!(E::Objective, conSet::ALConstraintSet, Z)

which also adds the expansion terms to the terms in `E`.

The penalty and multiplier terms can be updated using

	penalty_update!(::ALConstraintSet)
	dual_update!(::ALConstraintSet)

The current set of active constraint (with tolerance `tol`) can be re-calculated using

	update_active_set!(::ALConstraintSet, ::Val{tol})

The maximum penalty can be queried using `max_penalty(::ALConstraintSet)`, and the
penalties and/or multipliers can be reset using

	reset!(::ALConstraintSet)
	reset_penalties!(::ALConstraintSet)
	reset_duals!(::ALConstraintSet)

# Constructor
	ALConstraintSet(::ConstraintList, ::AbstractModel)
	ALConstraintSet(::Problem)
"""
struct ALConstraintSet{T} <: TO.AbstractConstraintSet
    convals::Vector{ALConVal}
    errvals::Vector{ALConVal}          # TODO: is this needed?
	# ∇c_proj::Vector{<:Vector}        # Jacobians of projected constraints
    # λ::Vector{<:Vector}
    # μ::Vector{<:Vector}
	# active::Vector{<:Vector}
    c_max::Vector{T}
    μ_max::Vector{T}
    μ_maxes::Vector{Vector{T}}
	# params::Vector{TO.ConstraintParams{T}}
	p::Vector{Int}
end

function ALConstraintSet(cons::TO.ConstraintList, model::RD.AbstractModel)
    n,m = cons.n, cons.m
    n̄ = RobotDynamics.state_diff_size(model)
    ncon = length(cons)
    useG = model isa RD.LieGroupModel
	errvals = map(1:ncon) do i
        C,c = TO.gen_convals(n̄, m, cons[i], cons.inds[i])
        ALConVal(n̄, m, cons[i], cons.inds[i], C, c, useG)
    end
    convals = map(errvals) do errval
        ALConVal(n, m, errval)
    end
	errvals = convert(Vector{ALConVal}, errvals)
	convals = convert(Vector{ALConVal}, convals)
	# ∇c_proj = map(1:ncon) do i
	# 	[copy(errvals[i].jac[j,1]) for j in eachindex(cons.inds[i])]
	# end
    # λ = map(1:ncon) do i
	# 	p = length(cons[i])
	# 	convals[i].λ
    #     # [@SVector zeros(p) for j in cons.inds[i]]
    # end
    # μ = map(1:ncon) do i
	# 	p = length(cons[i])
	# 	convals[i].μ
    #     # [@SVector ones(p) for j in cons.inds[i]]
    # end
    # a = map(1:ncon) do i
    #     p = length(cons[i])
    #     [@SVector ones(Bool,p) for j in cons.inds[i]]
	# end
	# if ncon == 0
	# 	∇c_proj = Vector{Float64}[]
	# 	λ = Vector{Float64}[]
	# 	μ = Vector{Float64}[]
	# 	a = Vector{Float64}[]
	# end
    c_max = zeros(ncon)
    μ_max = zeros(ncon)
    μ_maxes = [zeros(length(ind)) for ind in cons.inds]
	# params = [TO.ConstraintParams() for con in cons.constraints]
    # ALConstraintSet(convals, errvals, ∇c_proj, λ, μ, a, c_max, μ_max, μ_maxes, params, copy(cons.p))
    ALConstraintSet(convals, errvals, c_max, μ_max, μ_maxes, copy(cons.p))
end

@inline ALConstraintSet(prob::Problem) = ALConstraintSet(prob.constraints, prob.model)

@inline TO.get_convals(conSet::ALConstraintSet) = conSet.convals
@inline TO.get_errvals(conSet::ALConstraintSet) = conSet.errvals

get_duals(conSet::ALConstraintSet) = [cval.λ for cval in conSet.convals]
function set_duals!(conSet::ALConstraintSet, λ)
	for i = 1:length(conSet)
		set_duals!(conSet.convals[i], λ[i])
	end
end
function set_duals!(cval::ALConVal, λ)
	for i = 1:length(cval.λ)
		cval.λ[i] .= λ[i]
	end
end

for method in (:violation!, :∇violation!, :dual_update!, :penalty_update!, :reset_duals!, :reset_penalties!)
	@eval function $method(conSet::ALConstraintSet)
		for i = 1:length(conSet)
			$method(conSet.convals[i])
		end
	end
end

############################################################################################
#                             Augmented Lagrangian Updates
############################################################################################

# function dual_update!(conSet::ALConstraintSet)
#     for i in eachindex(conSet.λ)
#         dual_update!(conSet.convals[i])
# 	end
# end

function dual_update!(conval::ALConVal)
	c = conval.vals
	λ = conval.λ
	μ = conval.μ
	λ_max = conval.params.λ_max
	cone = TO.sense(conval.con)
	# λ_min = TO.sense(conval.con) == Equality() ? -λ_max : zero(λ_max)
	for i in eachindex(conval.inds)
		λ[i] .= dual_update(cone, SVector(λ[i]), SVector(c[i]), SVector(μ[i]), λ_max)
	end
end

function dual_update(::Equality, λ, c, μ, λmax)
	λbar = λ + μ .* c
	return clamp.(λbar, -λmax, λmax)
end

function dual_update(::Inequality, λ, c, μ, λmax)
 	λbar = λ + μ .* c
	return clamp.(λbar, 0, λmax)  # project onto the dual cone via max(0,x)
end

function dual_update(cone::SecondOrderCone, λ, c, μ, λmax)
	 λbar = λ - μ .* c
	 return TO.projection(cone, λbar)  # project onto the dual cone
end

# function penalty_update!(conSet::ALConstraintSet)
# 	for i in eachindex(conSet.μ)
# 		penalty_update!(conSet.convals[i])
# 	end
# end

function penalty_update!(cval::ALConVal)
	μ = cval.μ
	ϕ = cval.params.ϕ
	μ_max = cval.params.μ_max
	for i = 1:length(μ)
		μ[i] .*= ϕ
		clamp!(μ[i], 0, μ_max)
	end
end

# Active Set
function update_active_set!(conSet::ALConstraintSet, val::Val{tol}=Val(0.0)) where tol
	for i in 1:length(conSet)
		# update_active_set!(conSet.active[i], conSet.λ[i], conSet.convals[i], val)
		update_active_set!(conSet.convals[i], val)
	end
end

function update_active_set!(conval::ALConVal, ::Val{tol}) where tol
	a = conval.active
	λ = conval.λ
	if TO.sense(conval.con) == TO.Inequality()
		for i in eachindex(a)
			a[i] = @. (conval.vals[i] >= -tol) | (λ[i] > zero(tol))
		end
	end
end

"""
	max_penalty(conSet::ALConstraintSet)

Calculate the maximum constrained penalty across all constraints.
"""
function max_penalty(conSet::ALConstraintSet)
	max_penalty!(conSet)
	maximum(conSet.μ_max)
end

function max_penalty!(conSet::ALConstraintSet{T}) where T
    conSet.c_max .*= 0
	for i in 1:length(conSet)
		maxes = conSet.μ_maxes[i]::Vector{T}
        max_penalty!(maxes, conSet.convals[i])
        conSet.μ_max[i] = maximum(maxes)
    end
end

function max_penalty!(μ_max::Vector{<:Real}, cval::ALConVal)
    for i in eachindex(cval.μ)
        μ_max[i] = maximum(cval.μ[i])
    end
    return nothing
end

############################################################################################
#                             Constraint Penalties
############################################################################################
# function projection_jacobians!(conSet::ALConstraintSet)
# 	for i in eachindex(conSet.convals)
# 		projection_jacobians!(conSet.convals[i])
# 	end
# end

# """
# 	constraint_penalty!(conSet::ALConstraintSet)

# Evaluate the penalty of the constraints. For equality constraints, this is the constraint
# itself. For generalized inequalities, this is the projection of the constraint onto the
# appropriate cone.
# """
# function constraint_penalty!(conSet::ALConstraintSet)
# 	for i in eachindex(conSet.convals)
# 		constraint_penalty!(conSet.convals[i], conSet.λ[i])
# 	end
# end

# function constraint_penalty!(conval::TO.ConVal, λ)
# 	for i in eachindex(conval.inds)
# 		conval.vals2[i] = penalty(TO.sense(conval.con), conval.vals[i], λ[i])
# 	end
# end

# """
# 	penalty_jacobian!(conSet::ALConstraintSet)

# Evaluate the Jacobian of the constraint penalty. For equality constraints, this is simply
# the Jacobian of the constraint itself. For generalized inequalities, this is `∇proj*∇c`
# where `∇proj` is the Jacobian of the projection of the constraint onto the corresponding
# cone.
# """
# function penalty_jacobian!(conSet::ALConstraintSet)
# 	for i in eachindex(conSet.errvals)
# 		penalty_jacobian!(conSet.∇c_proj[i], conSet.errvals[i], conSet.λ[i])
# 	end
# end

# function penalty_jacobian!(∇c_proj::Vector, conval::TO.ConVal, λ::Vector)
# 	if size(conval.jac, 2) > 1
# 		throw(ErrorException("Constraint projection not supported for CoupledConstraints"))
# 	end
# 	for i in eachindex(conval.inds)
# 		∇penalty!(TO.sense(conval.con), ∇c_proj[i], conval.jac[i], conval.vals[i], λ[i])
# 	end
# end

############################################################################################
#                                        Cost
############################################################################################

# function TO.cost!(J::Vector{<:Real}, conSet::ALConstraintSet)
# 	for i in eachindex(conSet.convals)
# 		TO.cost!(J, conSet.convals[i], conSet.λ[i], conSet.μ[i], conSet.active[i])
# 	end
# end

# function TO.cost!(J::Vector{<:Real}, conval::ALConVal, λ::Vector{<:StaticVector},
# 		μ::Vector{<:StaticVector}, a::Vector{<:StaticVector})
# 	for (i,k) in enumerate(conval.inds)
# 		c = SVector(conval.vals[i])
# 		Iμ = Diagonal(SVector(μ[i] .* a[i]))
# 		J[k] += λ[i]'c .+ 0.5*c'Iμ*c
# 	end
# end

# function TO.cost_expansion!(E::Objective, conSet::ALConstraintSet, Z::AbstractTrajectory,
# 		init::Bool=false, rezero::Bool=false)
# 	for i in eachindex(conSet.errvals)
# 		TO.cost_expansion!(E, conSet.convals[i], conSet.λ[i], conSet.μ[i], conSet.active[i])
# 	end
# end

# @generated function TO.cost_expansion!(E::QuadraticObjective{n,m}, conval::ALConVal{C}, λ, μ, a) where {n,m,C}
# 	if C <: TO.StateConstraint
# 		expansion = quote
# 			cx = ∇c
# 			E[k].Q .+= cx'Iμ*cx
# 			E[k].q .+= cx'g
# 		end
# 	elseif C <: TO.ControlConstraint
# 		expansion = quote
# 			cu = ∇c
# 			E[k].R .+= cu'Iμ*cu
# 			E[k].r .+= cu'g
# 		end
# 	elseif C<: TO.StageConstraint
# 		ix = SVector{n}(1:n)
# 		iu = SVector{m}(n .+ (1:m))
# 		expansion = quote
# 			cx = ∇c[:,$ix]
# 			cu = ∇c[:,$iu]
# 			E[k].Q .+= cx'Iμ*cx
# 			E[k].q .+= cx'g
# 			E[k].H .+= cu'Iμ*cx
# 			E[k].R .+= cu'Iμ*cu
# 			E[k].r .+= cu'g
# 		end
# 	else
# 		throw(ArgumentError("cost expansion not supported for CoupledConstraints"))
# 	end
# 	quote
# 		for (i,k) in enumerate(conval.inds)
# 			∇c = SMatrix(conval.jac[i])
# 			c = conval.vals[i]
# 			Iμ = Diagonal(a[i] .* μ[i])
# 			g = Iμ*c .+ λ[i]

# 			$expansion
# 		end
# 	end
# end

############################################################################################
#                                  RESET
############################################################################################

function reset!(conSet::ALConstraintSet)
    reset_duals!(conSet)
    reset_penalties!(conSet)
end

# function reset_duals!(conSet::ALConstraintSet)
# 	for i = 1:length(conSet)
# 		reset_duals!(conSet.convals[i])
# 	end
# end

# function reset_penalties!(conSet::ALConstraintSet)
# 	for i = 1:length(conSet)
# 		reset_penalties!(conSet.convals[i])
# 	end
# end

"""
	link_constraints!(set1, set2)

Link any common constraints between `set1` and `set2` by setting elements in `set1` to point
to elements in `set2`
"""
function link_constraints!(set1::ALConstraintSet, set2::ALConstraintSet)
	# Find common constraints
	links = Tuple{Int,Int}[]
	for (i,con1) in enumerate(set1)
		for (j,con2) in enumerate(set2)
			if con1 === con2
				push!(links, (i,j))
			end
		end
	end

	# Link values
	for (i,j) in links
		set1.convals[i] = set2.convals[j]
		set1.errvals[i] = set2.errvals[j]
		# set1.active[i] = set2.active[j]
		# set1.λ[i] = set2.λ[j]
		# set1.μ[i] = set2.μ[j]
	end
	return links
end

function shift_fill!(conSet::ALConstraintSet, n=1)
	for i = 1:length(conSet)
		shift_fill!(conSet.convals[i], n)
	end
end

############################################################################################
#                                    Solver Options
############################################################################################
function reset!(conSet::ALConstraintSet{T}, opts::SolverOptions{T}) where T
    # if !isnan(opts.dual_max)
    #     for params in conSet.params
    #         params.λ_max = opts.dual_max
    #     end
    # end
    # if !isnan(opts.penalty_max)
    #     for params in conSet.params
    #         params.μ_max = opts.penalty_max
    #     end
    # end
    # if !isnan(opts.penalty_initial)
    #     for params in conSet.params
    #         params.μ0 = opts.penalty_initial
    #     end
    # end
    # if !isnan(opts.penalty_scaling)
    #     for params in conSet.params
    #         params.ϕ = opts.penalty_scaling
    #     end
	# end
	for i = 1:length(conSet)
		set_params!(conSet.convals[i], opts)
	end
    if opts.reset_duals
        reset_duals!(conSet)
    end
    if opts.reset_penalties
        reset_penalties!(conSet)
    end
end


# TODO(tim): support NA
# TODO(tim): support Nullable
# TODO(tim): support default value for missing keys

immutable CategoricalDiscretizer{N,D} <: AbstractDiscretizer{N,D}
    n2d :: Dict{N,D}    # maps natural to discrete
    d2n :: Dict{D,N}    # maps discrete to natural
end

function CategoricalDiscretizer{N,D}(natural_to_discrete::Dict{N,D})
    d2n = Dict{D,N}()
    for (k,v) in natural_to_discrete
        d2n[v] = k
    end
    CategoricalDiscretizer{N,D}(natural_to_discrete, d2n)
end
function CategoricalDiscretizer{N, D<:Integer}(data::AbstractArray{N}, ::Type{D}=Int)
    # build a label mapping N -> D <: Integer
    i = zero(D)
    n2d = @compat Dict{N, D}()
    for x in data
        if !haskey(n2d,x)
            n2d[x] = (i += 1)
        end
    end
    CategoricalDiscretizer(n2d)
end

supports_encoding{N,D}(cd::CategoricalDiscretizer{N,D}, x::N) = haskey(cd.n2d, x)
supports_decoding{N,D}(cd::CategoricalDiscretizer{N,D}, d::D) = 1 ≤ d ≤ nlabels(cd)

encode{N,D}(cd::CategoricalDiscretizer{N,D}, x::N) = cd.n2d[x]::D
encode{N,D}(cd::CategoricalDiscretizer{N,D}, x) = cd.n2d[convert(N,x)]::D
function encode{N,D}(cd::CategoricalDiscretizer{N,D}, data::AbstractArray)
    arr = Array{D}(length(data))
    for (i,x) in enumerate(data)
        arr[i] = encode(cd, x)
    end
    reshape(arr, size(data))
end

decode{N,D}(cd::CategoricalDiscretizer{N,D}, x::D) = cd.d2n[x]::N
decode{N,D}(cd::CategoricalDiscretizer{N,D}, x) = cd.d2n[convert(D,x)]::N
function decode{N,D}(cd::CategoricalDiscretizer{N,D}, data::AbstractArray{D})
    arr = Array{N}(length(data))
    for (i,d) in enumerate(data)
        arr[i] = decode(cd, d)
    end
    reshape(arr, size(data))
end

nlabels(cd::CategoricalDiscretizer) = length(cd.n2d)


"""
`process_streaks`
"""

function process_streaks(df::DataFrames.AbstractDataFrame; photometry = false)
    dayly_vars_list = [:MouseID, :Gen, :Drug, :Day, :Daily_Session, :Box, :Stim_Day, :Condition, :ExpDay, :Area, :Session];
    booleans=[:Reward,:Stim,:Wall,:Correct,:Stim_Day]#columns to convert to Bool
    for x in booleans
        df[!,x] = eltype(df[!,x]) == Bool ? df[!,x] : occursin.("true",df[!,x],)
    end
    streak_table = by(df, :Streak) do dd
        dt = DataFrame(
        Num_pokes = size(dd,1),
        Num_Rewards = length(findall(dd[!,:Reward].==1)),
        Start_Reward = dd[1,:Reward],
        Last_Reward = findlast(dd[!,:Reward] .== 1).== nothing ? 0 : findlast(dd[!,:Reward] .== 1),
        Prev_Reward = findlast(dd[!,:Reward] .== 1).== nothing ? 0 : findprev(dd[!,:Reward] .==1, findlast(dd[!,:Reward] .==1)-1),
        Trial_duration = (dd[end,:PokeOut]-dd[1,:PokeIn]),
        Start = (dd[1,:PokeIn]),
        Stop = (dd[end,:PokeOut]),
        Pre_Interpoke = size(dd,1) > 1 ? maximum(skipmissing(dd[!,:Pre_Interpoke])) : missing,
        Post_Interpoke = size(dd,1) > 1 ? maximum(skipmissing(dd[!,:Post_Interpoke])) : missing,
        PokeSequence = [SVector{size(dd,1),Bool}(dd[!,:Reward])],
        Stim = dd[1,:Stim],
        StimFreq = dd[1,:StimFreq],
        Wall = dd[1,:Wall],
        Protocol = dd[1,:Protocol],
        Correct_start = dd[1,:Correct],
        Correct_leave = !dd[end,:Correct],
        Block = dd[1,:Block],
        Streak_within_Block = dd[1,:Streak_within_Block],
        Side = dd[1,:Side],
        ReverseStreak = dd[1,:ReverseStreak]
        )
        for s in dayly_vars_list
            if s in names(df)
                dt[!,s] .= df[1, s]
            end
        end
        return dt
    end
    streak_table[!,:Prev_Reward] = [x .== nothing ? 0 : x for x in streak_table[:,:Prev_Reward]]
    streak_table[!,:AfterLast] = streak_table[!,:Num_pokes] .- streak_table[!,:Last_Reward];
    streak_table[!,:BeforeLast] = streak_table[!,:Last_Reward] .- streak_table[!,:Prev_Reward].-1;
    prov = lead(streak_table[!,:Start],default = 0.0) .- streak_table[!,:Stop];
    streak_table[!,:Travel_to] = [x.< 0 ? 0 : x for x in prov]
    if photometry
        frames = by(df, :Streak) do df
            dd = DataFrame(
            In = df[1,:In],
            Out = df[end,:Out],
            LR_In = findlast(df[:Reward])==0 ? NaN : df[findlast(df[:Reward]),:In],
            LR_Out = findlast(df[:Reward])==0 ? NaN : df[findlast(df[:Reward]),:Out]
            )
        end
        streak_table[:In] = frames[:In]
        streak_table[:Out] = frames[:Out]
        streak_table[:LR_In] = frames[:LR_In]
        streak_table[:LR_Out] = frames[:LR_Out]
    end

    return streak_table
end


using Measures

# 1 function to be moved to Measures.jl
Base.convert(::Type{T} , x::Measure) where T<:Real  = T(x.value)

struct DiscreteSizeScale <: Gadfly.ScaleElement
    f::Function
    levels::Union{Nothing, AbstractVector}
    order::Union{Nothing, AbstractVector}
    preserve_order::Bool
end
DiscreteSizeScale(f; levels=nothing, order=nothing, preserve_order=true) =
        DiscreteSizeScale(f, levels, order, preserve_order)

element_aesthetics(scale::DiscreteSizeScale) = [:size]

default_discrete_sizes(n::Int) = range(0.45mm, 1.8mm, length=n)

"""
    size_discrete2(f=default_discrete_sizes;
                    levels=nothing, order=nothing, preserve_order=true)

A discrete size scale that maps the categorical values in the `size`
aesthetic to x-axis units or `Measure` units (from Measures.jl).  `f` is a function that produces a vector of size units.
`levels` are the categorical levels, and level order will be respected.  `order` is
a vector of integers giving a permutation of the levels default order.  If
`preserve_order` is `true` levels are ordered as they appear in the data.
"""
size_discrete2(f::Function=Gadfly.current_theme().discrete_sizemap; levels=nothing, order=nothing, preserve_order=true) =
        DiscreteSizeScale(f, levels, order, preserve_order)

function Scale.apply_scale(scale::DiscreteSizeScale, aess::Vector{Gadfly.Aesthetics}, datas::Gadfly.Data...)

    d = []
    for (aes, data) in zip(aess, datas)
        data.size === nothing && continue
        append!(d, skipmissing(data.size))
    end
    levelset = unique(d)

    if scale.levels == nothing
        scale_levels = [levelset...]
        scale.preserve_order || sort!(scale_levels)
    else
        scale_levels = scale.levels
    end
    scale.order == nothing || permute!(scale_levels, scale.order)

    sizes = scale.f(length(scale_levels))

    size_map = Dict(s=>string(label) for (s, label) in zip(sizes, scale_levels))
    labeler(xs) = [size_map[x] for x in xs]
    key_vals = OrderedDict(s=>i for (i,s) in enumerate(sizes))

    for (aes, data) in zip(aess, datas)
        data.size === nothing && continue
        ds = discretize([d for d in skipmissing(data.size)], scale_levels)
        vals = sizes[ds.index]
        aes.size = discretize_make_ia(vals, sizes)
        aes.size_label = labeler
        aes.size_key_vals = key_vals
    end
end


#####################################################################
# special.jl
#
# includes non-standard code I needed for papers and whatnot
######################################################################

function norm2(x::LocTuple, y::LocTuple)
	d1 = x[1] - y[1]
	d2 = x[2] - y[2]
	return sqrt(d1*d1 + d2*d2)
end

function sim(j::LocTuple, actions::Vector{Pose}, observations::Vector{Float64}, b::Vector{Matrix{Float64}}, m::SearchDomain, x::Vehicle; show_mean=false, show_cov=false)
	# always start the vehicle in the center
	# NO, don't do that
	#x.x = m.length / 2.0
	#x.y = m.length / 2.0
	theta!(m, j[1], j[2])

	# warn user if jammer is not where it should be
	#if abs(dx) > x.x || abs(dy) > x.y
	#	println("WARNING: jammer outside search domain.")
	#end

	# loop through all observations...
	for (oi,o) in enumerate(observations)
		#update!(f, x, o)
		#plot(m, b[oi], x, show_mean=true, show_cov=true)
		plot(m, b[oi], x, show_mean=show_mean, show_cov=show_cov, obs=o)
		savefig("temp_$(oi).png", format="png")
		hold(false)
		act!(m, x, actions[oi])
	end
end

export sim2
function sim2(j::LocTuple, states::Vector{Pose}, observations::Vector{Float64}, b::Vector{Matrix{Float64}}, m::SearchDomain, x::Vehicle; show_mean=false, show_cov=false, show_path::Bool=true)
	# always start the vehicle in the center
	# NO, don't do that
	#x.x = m.length / 2.0
	#x.y = m.length / 2.0
	theta!(m, j[1], j[2])

	path_x = Float64[]
	path_y = Float64[]

	# loop through all observations...
	for (oi,o) in enumerate(observations)
		if show_path
			#plot path
			plot(path_x, path_y, "k")
		end
		hold(true)
		plot(m, b[oi], x, show_mean=show_mean, show_cov=show_cov, obs=oi)
		#savefig("temp_$(oi).pdf", format="pdf", dpi=300)
		savefig("temp_$(oi).png", format="png")
		hold(false)
		push!(path_x, x.x)
		push!(path_y, x.y)
		x.x = states[oi][1]
		x.y = states[oi][2]
		x.heading = states[oi][3]
		#act!(m, x, actions[oi])
	end
end

# This was for AIAA GNC 2018 paper.
# This probably shouldn't be here.
export eval2
function eval2(j::LocTuple, states::Vector{Pose}, observations::Vector{Float64}, b::Vector{Matrix{Float64}}, m::SearchDomain, x::Vehicle; show_mean=false, show_cov=false, show_path::Bool=true)
	# always start the vehicle in the center
	# NO, don't do that
	#x.x = m.length / 2.0
	#x.y = m.length / 2.0
	theta!(m, j[1], j[2])

    front_cone = Float64[]
    rear_cone = Float64[]
    sides = Float64[]

	# loop through all observations...
	for (oi,o) in enumerate(observations)
        # get true bearing between x
        xt = (x.x,x.y,x.heading)
        rel_bearing = fit_180(xt[3] - true_bearing(xt, m.theta))
        if rel_bearing < 0.0
            rel_bearing = -1.0 * rel_bearing
        end

        cla()
		plot(m, b[oi], x, show_mean=show_mean, show_cov=show_cov, obs=oi)
        #title("o = $o")
        #hold(false)
        if rel_bearing < 60.0
            # expect observation 1
            title("expect 1, got $o")
            push!(front_cone, o)
        elseif rel_bearing < 120.0
            title("expect either, got $o")
            push!(sides, o)
        else
            title("expect 0, got $o")
            push!(rear_cone,o)
        end
        pause(0.15)

        # update vehicle position
		x.x = states[oi][1]
		x.y = states[oi][2]
		x.heading = states[oi][3]
	end
    return front_cone, sides, rear_cone
end

function simulate(m::SearchDomain, uav::SimUnit;
                  video::Bool=true,
                  pause_time=0.3
                 )

    # reset the filter, vehicle, and policy
    # TODO: I think I assume the SimUnit comes in clean and ready to go
    #reset!(uav.f)
    #reset!(m, uav.x)
    #reset!(uav.p)

    # What was the cost to getting this first observation?
    temp_cost = get_cost(uav, m)

    # before doing anything else, we observe
    #  and update filter once
    o = observe(m, uav.x)
    update!(uav, o)

    # This was our first step; steps count number of observations
    step_count = 1

    # plot if need be
    if video
        figure("Simulation")
        plot(m, uav.f, uav.x)
        title("i = $(step_count)")
    end

    while !is_complete(uav.f, uav.tc, step_count)
        # act
        a = action(m, uav, o)
        act!(m, uav.x, a)

        # get cost and update step count
        temp_cost += get_cost(uav, m, a)
        step_count += 1

        # observe and update
        o = observe(m, uav.x)
        update!(uav, o)

        # plot if need be
        if video
            pause(pause_time)
            figure("Simulation")
            cla()
            plot(m, uav.f, uav.x)
            max_b = maximum(uav.f.b)
            title("i = $(step_count), max = $(round(max_b,3))")
        end
    end

    return temp_cost
end


#GroupContribution Struct.
#it statically stores the id and values.
struct GroupContribution{N}
    id::NTuple{N,Int}
    value::NTuple{N,Int}
end

#==
Property Struct:
it holds a value and its metadata.
for now, the metadata is the long name and the units.
the property is parametrically typed on the value, to recognize if is a
string, Int, Float, VariableProperty, or GroupContribution.
it can replace IntVal, StringVal,
TODO:in the future, using Unitful to hold the value, something like:
Property{T,UNITS} where UNITS <: Unitful.Unit
==#

struct Property{T}
    name::String
    units::String
    value::T
end

value(prop::Property) = prop.value
name(prop::Property) = prop.name
units(prop::Property) = prop.units

#==
#those structs are no longer used in this branch.

struct ScalarProperty
    units::String
    value::Float64
end

struct RangeBoundary
    units::String
    value::Float64
end

struct VectorProperty
    units::String
    EqNo::Int
    A::Float64
    B::Float64
    C::Float64
    D::Float64
    E::Float64
    Tmin::RangeBoundary
    Tmax::RangeBoundary
end

struct PureComp
    LibraryIndex::Int
    CompoundID::String
    StructureFormula::String
    Family::Int
    CriticalTemperature::ScalarProperty
    CriticalPressure::ScalarProperty
    CriticalVolume::ScalarProperty
    CriticalComperssibility::ScalarProperty
    NormalBoilingPointTemperature::ScalarProperty
    NormalMeltingPointTemperature::ScalarProperty
    TripplePointTemperature::ScalarProperty
    TripplePointPressure::ScalarProperty
    MolecularWeight::ScalarProperty
    LiquidVolumeAtNormalBoilingPoint::ScalarProperty
    AcentricityFactor::ScalarProperty
    SolubilityParameter::ScalarProperty
    DipoleMoment::ScalarProperty
    HeatOfFormation::ScalarProperty
    GibbsEnergyOfFormation::ScalarProperty
    AbsEntropy::ScalarProperty
    HeatOfFusionAtMeltingPoint::ScalarProperty
    HeatOfCombustion::ScalarProperty
    LiquidDensity::VectorProperty
    VaporPressure::VectorProperty
    HeatOfVaporization::VectorProperty
    LiquidHeatCapacityCp::VectorProperty
    IdealHeatCapacityCp::VectorProperty
    SecondVirialCoefficient::VectorProperty
    LiquidViscosity::VectorProperty
    VaporViscosity::VectorProperty
    LiquidThermalConductivity::VectorProperty
    VaporThermalConductivity::VectorProperty
    RPPHeatCapacityCp::VectorProperty
    RelativeStaticPermittivity::VectorProperty
    AntoineVaporPressure::VectorProperty
    LiquidViscosityRPS::VectorProperty
    COSTALDVolume::ScalarProperty
    DiameterLJ::ScalarProperty
    EnergyLJ::ScalarProperty
    RacketParameter::ScalarProperty
    FullerVolume::ScalarProperty
    Parachor::ScalarProperty
    SpecificGravity::ScalarProperty
    Charge::ScalarProperty
    CostaldAcentricFactor::ScalarProperty
    WilsonVolume::ScalarProperty
    ChaoSeaderAcentricFactor::ScalarProperty
    ChaoSeaderSolubilityParameter::ScalarProperty
    ChaoSeaderLiquidVolume::ScalarProperty
    CAS::String
end
==#


"""
    VF2

An empty concrete type used to dispatch to [`vf2`](@ref) isomorphism functions.
"""
struct VF2 <: IsomorphismAlgorithm end

"""
    VF2State{G, T}

Structure that is internally used by vf2
"""
struct VF2State{G, T}
    g1::G
    g2::G
    core_1::Vector{T}
    core_2::Vector{T}
    in_1::Vector{T}
    in_2::Vector{T}
    out_1::Vector{T}
    out_2::Vector{T}

    function VF2State(g1::G, g2::G) where {G <: AbstractSimpleGraph{T}} where {T <: Integer}
        n1 = nv(g1)
        n2 = nv(g2)
        core_1 = zeros(T, n1)
        core_2 = zeros(T, n2)
        in_1 = zeros(T, n1)
        in_2 = zeros(T, n2)
        out_1 = zeros(T, n1)
        out_2 = zeros(T, n2)

        return new{G, T}(g1, g2, core_1, core_2, in_1, in_2, out_1, out_2)
    end
end

"""
    vf2(callback, g1, g2, problemtype; vertex_relation=nothing, edge_relation=nothing)

Iterate over all isomorphism between the graphs `g1` (or subgraphs thereof) and `g2`.
The problem that is solved depends on the value of `problemtype`:
- IsomorphismProblem(): Only isomorphisms between the whole graph `g1` and `g2` are considered.
- SubGraphIsomorphismProblem(): All isomorphism between subgraphs of `g1` and `g2` are considered.
- InducedSubGraphIsomorphismProblem(): All isomorphism between vertex induced subgraphs of `g1` and `g2` are considered.

Upon finding an isomorphism, the function `callback` is called with a vector `vmap` as an argument.
`vmap` is a vector where `vmap[v] == u` means that vertex `v` in `g2` is mapped to vertex `u` in `g1`.
If the algorithm should look for another isomorphism, then this function should return `true`.

### Optional Arguments
- `vertex_relation`: A binary function that takes a vertex from `g1` and one from `g2`. An
    isomorphism only exists if this function returns `true` for all matched vertices.
- `edge_relation`: A binary function that takes an edge from `g1` and one from `g2`. An
    isomorphism only exists if this function returns `true` for all matched edges.

### References
Luigi P. Cordella, Pasquale Foggia, Carlo Sansone, Mario Vento
“A (Sub)Graph Isomorphism Algorithm for Matching Large Graphs”
"""
function vf2(callback::Function, g1::G, g2::G, problemtype::GraphMorphismProblem;
             vertex_relation::Union{Nothing, Function}=nothing,
             edge_relation::Union{Nothing, Function}=nothing) where {G <: AbstractSimpleGraph}
    if nv(g1) < nv(g2) || (problemtype == IsomorphismProblem() && nv(g1) != nv(g2))
        return
    end

    start_state = VF2State(g1, g2)
    start_depth = 1
    vf2match!(start_state, start_depth, callback, problemtype, vertex_relation, edge_relation)
    return
end

"""
    vf2check_feasibility(u, v, state, problemtype, vertex_relation, edge_relation)

Check whether two vertices of G₁ and G₂ can be matched. Used by [`vf2match!`](@ref).
"""
function vf2check_feasibility(u, v, state::VF2State, problemtype,
                              vertex_relation::Union{Nothing, Function},
                              edge_relation::Union{Nothing, Function})
    @inline function vf2rule_pred(u, v, state::VF2State, problemtype)
        if problemtype != SubGraphIsomorphismProblem()
            @inbounds for u2 in inneighbors(state.g1, u)
                if state.core_1[u2] != 0
                    found = false
                    # TODO can probably be replaced with has_edge for better performance
                    for v2 in inneighbors(state.g2, v)
                        if state.core_1[u2] == v2
                            found = true
                            break
                        end
                    end
                    found || return false
                end
            end
        end
        @inbounds for v2 in inneighbors(state.g2, v)
            if state.core_2[v2] != 0
                found = false
                for u2 in inneighbors(state.g1, u)
                    if state.core_2[v2] == u2
                        found = true
                        break
                    end
                end
                found || return false
            end
        end
        return true
    end

    @inline function vf2rule_succ(u, v, state::VF2State, problemtype)
        if problemtype != SubGraphIsomorphismProblem()
            @inbounds for u2 in outneighbors(state.g1, u)
                if state.core_1[u2] != 0
                    found = false
                    for v2 in outneighbors(state.g2, v)
                        if state.core_1[u2] == v2
                            found = true
                            break
                        end
                    end
                    found || return false
                end
            end
        end
        found = false
        @inbounds for v2 in outneighbors(state.g2, v)
            if state.core_2[v2] != 0
                found = false
                for u2 in outneighbors(state.g1, u)
                    if state.core_2[v2] == u2
                        found = true
                        break
                    end
                end
                found || return false
            end
        end
        return true
    end

    @inline function vf2rule_in(u, v, state::VF2State, problemtype)
        count1 = 0
        count2 = 0
        @inbounds for u2 in outneighbors(state.g1, u)
            if state.in_1[u2] != 0 && state.core_1[u2] == 0
                count1 += 1
            end
        end
        @inbounds for v2 in outneighbors(state.g2, v)
            if state.in_2[v2] != 0 && state.core_2[v2] == 0
                count2 += 1
            end
        end
        if problemtype == IsomorphismProblem()
            count1 == count2 || return false
        else
            count1 >= count2 || return false
        end
        count1 = 0
        count2 = 0
        @inbounds for u2 in inneighbors(state.g1, u)
            if state.in_1[u2] != 0 && state.core_1[u2] == 0
                count1 += 1
            end
        end
        @inbounds for v2 in inneighbors(state.g2, v)
            if state.in_2[v2] != 0 && state.core_2[v2] == 0
                count2 += 1
            end
        end
        problemtype == IsomorphismProblem() && return count1 == count2

        return count1 >= count2
    end

    @inline function vf2rule_out(u, v, state::VF2State, problemtype)
        count1 = 0
        count2 = 0
        @inbounds for u2 in outneighbors(state.g1, u)
            if state.out_1[u2] != 0 && state.core_1[u2] == 0
                count1 += 1
            end
        end
        @inbounds for v2 in outneighbors(state.g2, v)
            if state.out_2[v2] != 0 && state.core_2[v2] == 0
                count2 += 1
            end
        end
        if problemtype == IsomorphismProblem()
            count1 == count2 || return false
        else
            count1 >= count2 || return false
        end

        count1 = 0
        count2 = 0
        @inbounds for u2 in inneighbors(state.g1, u)
            if state.out_1[u2] != 0 && state.core_1[u2] == 0
                count1 += 1
            end
        end
        @inbounds for v2 in inneighbors(state.g2, v)
            if state.out_2[v2] != 0 && state.core_2[v2] == 0
                count2 += 1
            end
        end
        problemtype == IsomorphismProblem() && return count1 == count2

        return count1 >= count2
    end

    @inline function vf2rule_new(u, v, state::VF2State, problemtype)
        problemtype == SubGraphIsomorphismProblem() && return true
        count1 = 0
        count2 = 0
        @inbounds for u2 in inneighbors(state.g1, u)
            if state.in_1[u2] == 0 && state.out_1[u2] == 0
                count1 += 1
            end
        end
        @inbounds for v2 in inneighbors(state.g2, v)
            if state.in_2[v2] == 0 && state.out_2[v2] == 0
                count2 += 1
            end
        end
        if problemtype == IsomorphismProblem()
            count1 == count2 || return false
        else
            count1 >= count2 || return false
        end
        count1 = 0
        count2 = 0
        @inbounds for u2 in outneighbors(state.g1, u)
            if state.in_1[u2] == 0 && state.out_1[u2] == 0
                count1 += 1
            end
        end
        @inbounds for v2 in outneighbors(state.g2, v)
            if state.in_2[v2] == 0 && state.out_2[v2] == 0
                count2 += 1
            end
        end
        problemtype == IsomorphismProblem() && return count1 == count2

        return count1 >= count2
    end

    @inline function vf2rule_self_loops(u, v, state, problemtype)
        u_selflooped = has_edge(state.g1, u, u)
        v_selflooped = has_edge(state.g2, v, v)

        if problemtype == SubGraphIsomorphismProblem()
            return u_selflooped || !v_selflooped
        end
        return u_selflooped == v_selflooped
    end

    syntactic_feasability = vf2rule_pred(u, v, state, problemtype) &&
                            vf2rule_succ(u, v, state, problemtype) &&
                            vf2rule_in(u, v, state, problemtype)   &&
                            vf2rule_out(u, v, state, problemtype)  &&
                            vf2rule_new(u, v, state, problemtype)  &&
                            vf2rule_self_loops(u, v, state, problemtype)
    syntactic_feasability || return false

    if vertex_relation != nothing
        vertex_relation(u, v) || return false
    end
    if edge_relation != nothing
        E1 = edgetype(state.g1)
        E2 = edgetype(state.g2)
        for u2 in outneighbors(state.g1, u)
            state.core_1[u2] == 0 && continue
            v2 = state.core_1[u2]
            edge_relation(E1(u, u2), E2(v, v2)) || return false
        end
        for u2 in inneighbors(state.g1, u)
            state.core_1[u2] == 0 && continue
            v2 = state.core_1[u2]
            edge_relation(E1(u2, u), E2(v2, v)) || return false
        end
    end
    return true
end

"""
    vf2update_state!(state, u, v, depth)

Update state before recursing. Helper function for [`vf2match!`](@ref).
"""
function vf2update_state!(state::VF2State, u, v, depth)
@inbounds begin
     state.core_1[u] = v
     state.core_2[v] = u
     for w in outneighbors(state.g1, u)
         if state.out_1[w] == 0
             state.out_1[w] = depth
         end
     end
     for w in inneighbors(state.g1, u)
         if state.in_1[w] == 0
             state.in_1[w] = depth
         end
     end
     for w in outneighbors(state.g2, v)
         if state.out_2[w] == 0
             state.out_2[w] = depth
         end
     end
     for w in inneighbors(state.g2, v)
         if state.in_2[w] == 0
             state.in_2[w] = depth
         end
     end
end
end

"""
    vf2reset_state!(state, u, v, depth)

Reset state after returning from recursion. Helper function for [`vf2match!`](@ref).
"""
function vf2reset_state!(state::VF2State, u, v, depth)
@inbounds begin
    state.core_1[u] = 0
    state.core_2[v] = 0
    for w in outneighbors(state.g1, u)
        if state.out_1[w] == depth
            state.out_1[w] = 0
        end
    end
    for w in inneighbors(state.g1, u)
        if state.in_1[w] == depth
            state.in_1[w] = 0
        end
    end
    for w in outneighbors(state.g2, v)
        if state.out_2[w] == depth
            state.out_2[w] = 0
        end
    end
    for w in inneighbors(state.g2, v)
        if state.in_2[w] == depth
            state.in_2[w] = 0
        end
    end
end
end

"""
    vf2match!(state, depth, callback, problemtype, vertex_relation, edge_relation)

Perform isomorphic subgraph matching. Called by [`vf2`](@ref).
"""
function vf2match!(state, depth, callback::Function, problemtype::GraphMorphismProblem,
                   vertex_relation, edge_relation)
    n1 = Int(nv(state.g1))
    n2 = Int(nv(state.g2))
    # if all vertices of G₂ are matched we call the callback function. If the
    # algorithm should look for another isomorphism then callback has to return true
    if depth > n2
        keepgoing = callback(state.core_2)
        return keepgoing
    end
    # First we try if there is a pair of unmatched vertices u∈G₁ v∈G₂ that are connected
    # by an edge going out of the set M(s) of already matched vertices
    found_pair = false
    v = 0
    @inbounds for j = 1:n2
        if state.out_2[j] != 0 && state.core_2[j] == 0
            v = j
            break
        end
    end
    if v != 0
        @inbounds for u = 1:n1
            if state.out_1[u] != 0 && state.core_1[u] == 0
                found_pair = true
                if vf2check_feasibility(u, v, state, problemtype, vertex_relation, edge_relation)
                    vf2update_state!(state, u, v, depth)
                    keepgoing = vf2match!(state, depth + 1, callback, problemtype,
                                          vertex_relation, edge_relation)
                    keepgoing || return false
                    vf2reset_state!(state, u, v, depth)
                end
            end
        end
    end
    found_pair && return true
    # If that is not the case we try if there is a pair of unmatched vertices u∈G₁ v∈G₂ that
    # are connected  by an edge coming in from the set M(s) of already matched vertices
    v = 0
    @inbounds for j = 1:n2
        if state.in_2[j] != 0 && state.core_2[j] == 0
            v = j
            break
        end
    end
    if v != 0
        @inbounds for u = 1:n1
            if state.in_1[u] != 0 && state.core_1[u] == 0
                found_pair = true
                if vf2check_feasibility(u, v, state, problemtype, vertex_relation, edge_relation)
                    vf2update_state!(state, u, v, depth)
                    keepgoing = vf2match!(state, depth + 1, callback, problemtype,
                                          vertex_relation, edge_relation)
                    keepgoing || return false
                    vf2reset_state!(state, u, v, depth)
                end
            end
        end
    end
    found_pair && return true
    # If this is also not the case, we try all pairs of vertices u∈G₁ v∈G₂ that are not
    # yet matched
    v = 0
    @inbounds for j = 1:n2
        if state.core_2[j] == 0
            v = j
            break
        end
    end
    if v != 0
        @inbounds for u = 1:n1
            if state.core_1[u] == 0
                if vf2check_feasibility(u, v, state, problemtype, vertex_relation, edge_relation)
                    vf2update_state!(state, u, v, depth)
                    keepgoing = vf2match!(state, depth + 1, callback, problemtype,
                                          vertex_relation, edge_relation)
                    keepgoing || return false
                    vf2reset_state!(state, u, v, depth)
                end
            end
        end
    end
    return true
end

function has_induced_subgraphisomorph(g1::AbstractGraph, g2::AbstractGraph, alg::VF2;
                                 vertex_relation::Union{Nothing, Function}=nothing,
                                 edge_relation::Union{Nothing, Function}=nothing)::Bool
        result = false
        callback(vmap) = (result = true; return false)
        vf2(callback, g1, g2, InducedSubGraphIsomorphismProblem();
                       vertex_relation=vertex_relation, edge_relation=edge_relation)
        return result
end

function has_subgraphisomorph(g1::AbstractGraph, g2::AbstractGraph, alg::VF2;
                                 vertex_relation::Union{Nothing, Function}=nothing,
                                 edge_relation::Union{Nothing, Function}=nothing,)::Bool
    result = false
    callback(vmap) = (result = true; return false)
    vf2(callback, g1, g2, SubGraphIsomorphismProblem();
        vertex_relation=vertex_relation, edge_relation=edge_relation)
    return result
end

function has_isomorph(g1::AbstractGraph, g2::AbstractGraph, alg::VF2;
                         vertex_relation::Union{Nothing, Function}=nothing,
                         edge_relation::Union{Nothing, Function}=nothing)::Bool

    !could_have_isomorph(g1, g2) && return false

    result = false
    callback(vmap) = (result = true; return false)
    vf2(callback, g1, g2, IsomorphismProblem(),
        vertex_relation=vertex_relation,
        edge_relation=edge_relation)
    return result
end

function count_induced_subgraphisomorph(g1::AbstractGraph, g2::AbstractGraph, alg::VF2;
                                   vertex_relation::Union{Nothing, Function}=nothing,
                                   edge_relation::Union{Nothing, Function}=nothing)::Int
    result = 0
    callback(vmap) = (result += 1; return true)
    vf2(callback, g1, g2, InducedSubGraphIsomorphismProblem(),
        vertex_relation=vertex_relation,
        edge_relation=edge_relation)
    return result
end

function count_subgraphisomorph(g1::AbstractGraph, g2::AbstractGraph, alg::VF2;
                                vertex_relation::Union{Nothing, Function}=nothing,
                                edge_relation::Union{Nothing, Function}=nothing)::Int
    result = 0
    callback(vmap) = (result += 1; return true)
    vf2(callback, g1, g2, SubGraphIsomorphismProblem(), vertex_relation=vertex_relation, edge_relation=edge_relation)
    return result
end

function count_isomorph(g1::AbstractGraph, g2::AbstractGraph, alg::VF2;
                        vertex_relation::Union{Nothing, Function}=nothing,
                        edge_relation::Union{Nothing, Function}=nothing)::Int
    !could_have_isomorph(g1, g2) && return 0
    result = 0
    callback(vmap) = (result += 1; return true)
    vf2(callback, g1, g2, IsomorphismProblem(), vertex_relation=vertex_relation, edge_relation=edge_relation)
    return result
end

function all_induced_subgraphisomorph(g1::AbstractGraph, g2::AbstractGraph, alg::VF2;
                                 vertex_relation::Union{Nothing, Function}=nothing,
                                 edge_relation::Union{Nothing, Function}=nothing)::Channel{Vector{Tuple{eltype(g1),eltype(g2)}}}
        make_callback(c) = vmap -> (put!(c, collect(zip(vmap,1:length(vmap)))), return true)
    T = Vector{Tuple{eltype(g1), eltype(g2)}}
    ch::Channel{T} = Channel(ctype=T) do c
        vf2(make_callback(c), g1, g2, InducedSubGraphIsomorphismProblem(),
                       vertex_relation=vertex_relation,
                       edge_relation=edge_relation)
    end
end

function all_subgraphisomorph(g1::AbstractGraph, g2::AbstractGraph, alg::VF2;
                              vertex_relation::Union{Nothing, Function}=nothing,
                              edge_relation::Union{Nothing, Function}=nothing)::Channel{Vector{Tuple{eltype(g1), eltype(g2)}}}

    make_callback(c) = vmap -> (put!(c, collect(zip(vmap,1:length(vmap)))), return true)
    T = Vector{Tuple{eltype(g1), eltype(g2)}}
    ch::Channel{T} = Channel(ctype=T) do c
        vf2(make_callback(c), g1, g2, SubGraphIsomorphismProblem(),
            vertex_relation=vertex_relation,
            edge_relation=edge_relation)
    end
    return ch
end

function all_isomorph(g1::AbstractGraph, g2::AbstractGraph, alg::VF2;
                 vertex_relation::Union{Nothing, Function}=nothing,
                 edge_relation::Union{Nothing, Function}=nothing)::Channel{Vector{Tuple{eltype(g1),eltype(g2)}}}
    T = Vector{Tuple{eltype(g1), eltype(g2)}}
    !could_have_isomorph(g1, g2) && return Channel(_ -> return, ctype=T)
    make_callback(c) = vmap -> (put!(c, collect(zip(vmap,1:length(vmap)))), return true)
    ch::Channel{T} = Channel(ctype=T) do c
        vf2(make_callback(c), g1, g2, IsomorphismProblem(),
            vertex_relation=vertex_relation,
            edge_relation=edge_relation)
    end
    return ch
end


function SimpleExponentialAtmosphere(arg0::BodyShape, arg1::jdouble, arg2::jdouble, arg3::jdouble)
    return SimpleExponentialAtmosphere((BodyShape, jdouble, jdouble, jdouble), arg0, arg1, arg2, arg3)
end

function get_density(obj::SimpleExponentialAtmosphere, arg0::AbsoluteDate, arg1::Vector3D, arg2::Frame)
    return jcall(obj, "getDensity", jdouble, (AbsoluteDate, Vector3D, Frame), arg0, arg1, arg2)
end

function get_density(obj::SimpleExponentialAtmosphere, arg0::FieldAbsoluteDate, arg1::FieldVector3D, arg2::Frame)
    return jcall(obj, "getDensity", RealFieldElement, (FieldAbsoluteDate, FieldVector3D, Frame), arg0, arg1, arg2)
end

function get_frame(obj::SimpleExponentialAtmosphere)
    return jcall(obj, "getFrame", Frame, ())
end


#=

ver_rand and its variants

=#

#using Random

#=
using Random: srand, rand, rand!,
randn, randn!, randexp, randexp!,
bitrand,
randstring,
randsubseq,randsubseq!,
shuffle,shuffle!,
randperm, randcycle
=#

#import Random

rand_ver(::Type{V06}, args...) = RandomV06.rand(args...)
seed_ver!(::Type{V06}, args...) = RandomV06.srand(args...)
rand_ver!(::Type{V06}, args...) = RandomV06.rand!(args...)
randn_ver(::Type{V06}, args...) = RandomV06.randn(args...)
randn_ver!(::Type{V06}, args...) = RandomV06.randn!(args...)
randexp_ver(::Type{V06}, args...) = RandomV06.randexp(args...)
randexp_ver!(::Type{V06}, args...) = RandomV06.randexp!(args...)
bitrand_ver(::Type{V06}, args...) = RandomV06.bitrand(args...)
randstring_ver(::Type{V06}, args...) = RandomV06.randstring(args...)
randsubseq_ver(::Type{V06}, args...) = RandomV06.randsubseq(args...)
randsubseq_ver!(::Type{V06}, args...) = RandomV06.randsubseq!(args...)
shuffle_ver(::Type{V06}, args...) = RandomV06.shuffle(args...)
shuffle_ver!(::Type{V06}, args...) = RandomV06.shuffle!(args...)
randperm_ver(::Type{V06}, args...) = RandomV06.randperm(args...)
randcycle_ver(::Type{V06}, args...) = RandomV06.randcycle(args...)

seed_ver!(::Type{Vcur}, args...) = Random.seed!(args...)
rand_ver(::Type{Vcur}, args...) = Random.rand(args...)
rand_ver!(::Type{Vcur}, args...) = Random.rand!(args...)
randn_ver(::Type{Vcur}, args...) = Random.randn(args...)
randn_ver!(::Type{Vcur}, args...) = Random.randn!(args...)
randexp_ver(::Type{Vcur}, args...) = Random.randexp(args...)
randexp_ver!(::Type{Vcur}, args...) = Random.randexp!(args...)
bitrand_ver(::Type{Vcur}, args...) = Random.bitrand(args...)
randstring_ver(::Type{Vcur}, args...) = Random.randstring(args...)
randsubseq_ver(::Type{Vcur}, args...) = Random.randsubseq(args...)
randsubseq_ver!(::Type{Vcur}, args...) = Random.randsubseq!(args...)
shuffle_ver(::Type{Vcur}, args...) = Random.shuffle(args...)
shuffle_ver!(::Type{Vcur}, args...) = Random.shuffle!(args...)
randperm_ver(::Type{Vcur}, args...) = Random.randperm(args...)
randcycle_ver(::Type{Vcur}, args...) = Random.randcycle(args...)

;


module AxisIndicesGraphs

using AxisIndices, LightGraphs, SimpleWeightedGraphs
using Base: to_index, @propagate_inbounds

struct AxisIndicesGraph{T,G<:AbstractGraph{T},AI} <: AbstractGraph{T}
    graph::G
    axes::AI
end

Base.parent(g::AxisIndicesGraph) = getfield(g, :graph)

# Although it seems weird to overload this for graphs it makes sense in the context
# of how an "axis" is defined (e.g., keys + indices)
Base.axes(g::AxisIndicesGraph) = getfield(g, :axes)

# Shouldn't these be performed on the type
#eltype(g::AbstractMetaGraph) = eltype(g.graph)
#edgetype(g::AbstractMetaGraph) = edgetype(g.graph)

###
### Vertices
###
LightGraphs.nv(g::AxisIndicesGraph) = nv(parent(g))
LightGraphs.vertices(g::AxisIndicesGraph) = vertices(parent(g))

LightGraphs.has_vertex(g::AxisIndicesGraph, x...) = has_vertex(parent(g), x...)
@inline LightGraphs.has_edge(g::AxisIndicesGraph, x...) = has_edge(parent(g), x...)

@propagate_inbounds function LightGraphs.add_vertex!(g::AxisIndicesGraph, u, v,)
    return add_vertex!(parent(g), to_index(axes(g), u), to_index(axes(g), v))
end

@propagate_inbounds function LightGraphs.rem_vertex!(g::AxisIndicesGraph, u, v)
    return rem_vertex!(parent(g), to_index(axes(g), u), to_index(axes(g), v))
end

###
### Edges
###
LightGraphs.ne(g::AxisIndicesGraph) = ne(parent(g))
LightGraphs.edges(g::AxisIndicesGraph) = edges(parent(g))

@propagate_inbounds function LightGraphs.add_edge!(g::AxisIndicesGraph, u, v)
    return add_edge!(parent(g), to_index(axes(g), u), to_index(axes(g), v))
end
@propagate_inbounds function LightGraphs.add_edge!(g::AxisIndicesGraph, u, v, val)
    return add_edge!(parent(g), to_index(axes(g), u), to_index(axes(g), v), val)
end

@propagate_inbounds function LightGraphs.rem_edge!(g::AxisIndicesGraph, u, v)
    return rem_edge!(parent(g), to_index(axes(g), u), to_index(axes(g), v))
end

function LightGraphs.inneighbors(g::AxisIndicesGraph, v)
    return inneighbors(parent(g), to_index(axes(g), v))
end
function LightGraphs.outneighbors(g::AxisIndicesGraph, v)
    return outneighbors(parent(g), to_index(axes(g), v))
end

Base.Base.issubset(g::AxisIndicesGraph, h::AbstractGraph) = issubset(parent(g), h)

LightGraphs.is_directed(::Type{<:AxisIndicesGraph{T,G}}) where {T,G} = is_directed(G)

SimpleWeightedGraphs.weighttype(::Type{<:AxisIndicesGraph{T,G}}) where {T,G} = weighttype(G)

# TODO these should be AxisIndicesArrays
SimpleWeightedGraphs.degree_matrix(g::AxisIndicesGraph, args...) = degree_matrix(parent(g), args...)
function LightGraphs.adjacency_matrix(g::AxisIndicesGraph, args...)
    return adjacency_matrix(parent(g), args...)
end

function LightGraphs.laplacian_matrix(g::AxisIndicesGraph, args...)
    return laplacian_matrix(parent(g), args...)
end

function LightGraphs.pagerank(g::AxisIndicesGraph, α=0.85, n::Integer=100, ϵ=1.0e-6)
    return pagerank(parent(g), α, n, ϵ)
end

function SimpleWeightedGraphs.get_weight(g::AxisIndicesGraph, u, v)
    return get_weight(parent(g), to_index(axes(g), u), to_index(axes(g), v))
end

LightGraphs.connected_components(g::AxisIndicesGraph) = connected_components(parent(g))

function LightGraphs.induced_subgraph(g::AxisIndicesGraph, vlist)
    return induced_subgraph(parent(g), to_index(axes(g), vlist))
end

###
### Base
###
Base.zero(g::AxisIndicesGraph) = AxisIndicesGraph(zero(parent(g)), axes(g))

Base.copy(g::AxisIndicesGraph) = AxisIndicesGraph(copy(parent(g)), copy(axes(g)))

end


import Base: show, dump

function show(io::IO, intrng::Array{Intrng,1})
    write(io, """A 1.0 #regions in PPM
# low field   high field  bias        slope\n""")
    entries = sort(intrng, by=(Δ->Δ[1]), rev=true)
    for (i,(δ1, δ2)) in enumerate(entries)
        write(io, " $δ1  $δ2  0.0  0.0  # for region $i\n")
    end
end

function dump(path, s::Spectrum)
    if !isdir(path)
        mkdir(path)
    end

    # Save each child proc
    for (n,p) in s.procs
        dump(joinpath(path, "pdata", string(n)), p)
    end

    # Save fid
    write(open(joinpath(path,"fid"), "w"), Int32.(round.(s.fid)))

    # TODO: Save acqu/acqus
end

function dump(path, s::Spectrum, templatepath)
    cp(templatepath, path)
    dump(path, s)
end

function dump(path, p::ProcessedSpectrum)
    if !isdir(path)
        mkdir(path)
    end

    # Save 1r and 1i
    open(joinpath(path, "1r"), "w") do f
        write(f, Int32.(round.(p.re_ft)))
    end
    open(joinpath(path, "1i"), "w") do f
        write(f, Int32.(round.(p.im_ft)))
    end

    # Save intrng
    open(joinpath(path, "intrng"), "w") do f
        show(f, p.intrng)
    end

    # Save title
    open(joinpath(path, "title"), "w") do f
        write(f, p.title)
    end

    # TODO: Save proc/procs
end


vis = Visualizer()

if haskey(ENV, "CI")
    port = split(split(url(vis), ':')[end], '/')[1]
    @show port
    stream, proc = open(`julia $(joinpath(@__DIR__, "dummy_websocket_client.jl")) $port`)
else
    proc = nothing
    open(vis)
end

wait(vis)
delete!(vis)

@testset "self-contained visualizer" begin
    @testset "shapes" begin
        v = vis[:shapes]
        delete!(v)
        settransform!(v, Translation(1., 0, 0))
        @testset "box" begin
            setobject!(v[:box], HyperRectangle(Vec(0., 0, 0), Vec(0.1, 0.2, 0.3)))
            settransform!(v[:box], Translation(-0.05, -0.1, 0))
        end

        @testset "cylinder" begin
            setobject!(v[:cylinder], Mesh(
               Cylinder(Point(0., 0, 0), Point(0, 0, 0.2), 0.1),
               MeshLambertMaterial(color=colorant"lime")))
            settransform!(v[:cylinder], Translation(0, 0.5, 0.0))
        end

        @testset "sphere" begin
            setobject!(v[:sphere],
                HyperSphere(Point(0., 0, 0), 0.15),
                MeshLambertMaterial(color=colorant"maroon"))
            settransform!(v[:sphere], Translation(0, 1, 0.15))
        end

        @testset "ellipsoid" begin
            setobject!(v[:ellipsoid], HyperEllipsoid(Point(0., 1.5, 0), Vec(0.3, 0.1, 0.1)))
            settransform!(v[:ellipsoid], Translation(0, 0, 0.1))
        end

        @testset "cube" begin
            setobject!(v[:cube], HyperCube(Vec(-0.1, -0.1, 0), 0.2), MeshBasicMaterial())
            settransform!(v[:cube], Translation(0, 2.0, 0))
        end

        @testset "more complicated cylinder" begin
            setobject!(v[:cylinder2], Cylinder(Point(0, 2.5, 0), Point(0.1, 0.1, 0), 0.05))
            settransform!(v[:cylinder2], Translation(0, 0, 0.05))
        end
    end

    @testset "meshes" begin
        v = vis[:meshes]
        @testset "cat" begin
            mesh = load(joinpath(Pkg.dir("GeometryTypes"), "test", "data", "cat.obj"))
            setobject!(v[:cat], mesh)
            settransform!(v[:cat], Translation(0, -1, 0) ∘ LinearMap(RotZ(π)) ∘ LinearMap(RotX(π/2)))
        end

        @testset "textured valkyrie" begin
            head = Mesh(
                load(joinpath(MeshCat.VIEWER_ROOT, "data", "head_multisense.obj"), GLUVMesh),
                MeshLambertMaterial(
                    map=Texture(
                        image=PngImage(joinpath(MeshCat.VIEWER_ROOT, "data", "HeadTextureMultisense.png"))
                    )
                ))
            setobject!(v[:valkyrie, :head], head)
            settransform!(v[:valkyrie, :head], Translation(0, 0.5, 0.5))
        end
    end

    @testset "points" begin
        v = vis[:points]
        settransform!(v, Translation(-1, 0, 0))
        @testset "random points" begin
            verts = rand(Point3f0, 100000);
            colors = reinterpret(RGB{Float32}, verts);
            setobject!(v[:random], PointCloud(verts, colors))
            settransform!(v[:random], Translation(-0.5, -0.5, 0))
        end
    end

    @testset "Polyhedra" begin
        ext1 = SimpleVRepresentation([0 1 2 3;0 2 1 3; 1 0 2 3; 1 2 0 3; 2 0 1 3; 2 1 0 3;
                                    0 1 3 2;0 3 1 2; 1 0 3 2; 1 3 0 2; 3 0 1 2; 3 1 0 2;
                                    0 3 2 1;0 2 3 1; 3 0 2 1; 3 2 0 1; 2 0 3 1; 2 3 0 1;
                                    3 1 2 0;3 2 1 0; 1 3 2 0; 1 2 3 0; 2 3 1 0; 2 1 3 0])
        poly1 = CDDPolyhedron{4,Rational{BigInt}}(ext1)

        poly2 = project(poly1, [1 1 1; -1 1 1; 0 -2 1; 0 0 -3])
        setobject!(vis[:polyhedron], poly2)
        settransform!(vis[:polyhedron],  Translation(-0.5, 1.0, 0) ∘ LinearMap(UniformScaling(0.1)))
    end

end

close(vis)

if proc !== nothing
    kill(proc)
end


mutable struct Shared
    df
    selectvalues::Vector{SpinBoxType}
    selectlist::Vector{Column}
    plotvalues::Vector{ComboBoxType}
    plotkwargs::TextBoxEntry
    smoother::SliderEntry
    splitting_var::ComboBoxType
    plt::Plots.Plot{Plots.GRBackend}
end
shared = Shared([], SpinBoxType[], Column[], ComboBoxType[],
    TextBoxEntry(""), SliderEntry(0.0), ComboBoxType("", [""]), plot())

"""
    build_window(; kwargs...)

Starts the GUI asking for a suitable csv file.
"""
function build_window(; kwargs...)
    @qmlapp joinpath(Pkg.dir("PlugAndPlot","src"), "QML", "choose_file.qml")
    path = TextBoxEntry("")
    @qmlset qmlcontext().choose = path
    exec()
    return build_window(path.value; kwargs...)
end

"""
    build_window(datafile::AbstractString; kwargs...)

Reads a csv file and starts `build_window` on the corresponding DataFrame
"""
build_window(datafile::AbstractString; kwargs...) = build_window(JuliaDB.loadtable(datafile); kwargs...)

"""
    build_window(dataset; nbox = 5)

Creates a GUI to analyze a data table interactively. Data can be selected either
on continuous columns, with SpinBoxes or on discrete columns with checkboxes,
provided there are less than `nbox` entries.
"""
build_window(dataset; kwargs...) = build_window(table(dataset); kwargs...)

function build_window(dataset::IndexedTables.AbstractIndexedTable; nbox = 5)
    shared.df = dataset
    shared.selectlist = [Column(string(name), string.(union(columns(shared.df, name))))
        for name in colnames(shared.df) if (1 < length(union(columns(shared.df, name))) < nbox)]

    shared.selectvalues = [SpinBoxType(string(name), Float64.([extrema(columns(shared.df, name))...]))
        for name in colnames(shared.df) if length(union(columns(shared.df, name))) >= nbox &&
        eltype(columns(shared.df, name)) <: Real]
    # run QML window
    qml_engine = init_qmlapplicationengine()
    @qmlset qmlcontext()._selectlist = ListModel(shared.selectlist)
    @qmlset qmlcontext()._selectvalues = ListModel(shared.selectvalues)
    shared.plotvalues = get_plotvalues(shared.df)
    @qmlset qmlcontext()._plotvalues = ListModel(shared.plotvalues)
    @qmlset qmlcontext().choose = shared.plotkwargs
    @qmlset qmlcontext().smoother = shared.smoother
    splitting_options = ["Don't split"; string.("splitby:",colnames(shared.df))]
    shared.splitting_var = ComboBoxType("spitting_var", splitting_options)
    @qmlset qmlcontext().splitting_var = shared.splitting_var
    qmlfunction("do_plot", PlugAndPlot.do_plot)
    qmlfunction("save_plot", PlugAndPlot.save_plot)
    qmlfunction("check_equality", PlugAndPlot.check_equality)
    qml_file = joinpath(Pkg.dir("PlugAndPlot","src"), "QML", "gui.qml")
    QML.load(qml_engine,qml_file)

    exec()
    return
end

function do_plot(d::JuliaDisplay, width, height, in_place = false)
    gr(grid = false, size=(Int64(round(width)),Int64(round(height))))
    get_plot!(shared, in_place)
    display(d, shared.plt)
end

save_plot(savename::AbstractString) =  savefig(shared.plt, savename)

function check_equality()
    (shared.plotvalues[1].chosen_value == shared.plotvalues[2].chosen_value) &&
    (shared.plotvalues[1].text_info == shared.plotvalues[2].text_info)
end