export zoo_cnf,
    zoo_dnf,
    zoo_sdd,
    load_logic_circuit,
    load_smooth_logic_circuit,
    load_struct_smooth_logic_circuit,
    load_cnf,
    load_dnf,
    zoo_cnf_file,
    zoo_dnf_file,
    zoo_lc_file,
    zoo_psdd_file,
    zoo_sdd_file

#####################
# loaders from model zoo
#####################

const zoo_version = "/Circuit-Model-Zoo-0.1.2"

zoo_cnf(name) =
    load_cnf(zoo_cnf_file(name))

zoo_dnf(name) =
    load_dnf(zoo_dnf_file(name))

# zoo file identifiers

zoo_cnf_file(name) =
    artifact"circuit_model_zoo" * zoo_version * "/cnfs/$name"

zoo_dnf_file(name) =
    artifact"circuit_model_zoo" * zoo_version * "/dnfs/$name"

zoo_lc_file(name) =
    artifact"circuit_model_zoo" * zoo_version * "/lcs/$name"

zoo_psdd_file(name) =
    artifact"circuit_model_zoo" * zoo_version * "/psdds/$name"

zoo_sdd_file(name) =
    artifact"circuit_model_zoo" * zoo_version * "/sdds/$name"

zoo_sdd(name) =
    load_logic_circuit(zoo_sdd_file(name))

#####################
# general parser infrastructure for circuits
#####################

# The `ParserCombinator` library works correctly but is orders of magnitude too slow.
# Instead here we hardcode some simpler parsers to speed things up

"""
    load_logic_circuit(file::String)::PlainLogicCircuit

Load a logical circuit from file.
Support file formats:
 * ".sdd" for SDD files
 * ".psdd" for PSDD files
 * ".circuit" for Logistic Circuit files
"""
function load_logic_circuit(file::String)::PlainLogicCircuit
    compile_logical(parse_circuit_file(file))
end

"""
    load_smooth_logic_circuit(file::String)::PlainLogicCircuit

Load a smooth logic circuit from file.
Support file formats:
 * ".psdd" for PSDD files
 * ".circuit" for Logistic Circuit files
"""
function load_smooth_logic_circuit(file::String)::PlainLogicCircuit
    compile_smooth_logical(parse_circuit_file(file))
end

"""
    load_struct_smooth_logic_circuit(circuit_file::String, vtree_file::String)::Tuple{StructLogicCircuit,PlainVtree}

Load a smooth structured logic circuit and its vtree from file.
Supported circuit file formats:
 * ".psdd" for PSDD files
 * ".circuit" for Logistic Circuit files
Supported vtree file formats:
 * ".vtree" for Vtree files
"""
function load_struct_smooth_logic_circuit(circuit_file::String, vtree_file::String)::Tuple{StructLogicCircuit,PlainVtree}
    circuit_lines = parse_circuit_file(circuit_file)
    vtree_lines = parse_vtree_file(vtree_file)
    compile_smooth_struct_logical(circuit_lines, vtree_lines)
end

#####################
# parse based on file extension
#####################

function parse_circuit_file(file::String)::CircuitFormatLines
    if endswith(file,".circuit")
        parse_lc_file(file)
    elseif endswith(file,".psdd")
        parse_psdd_file(file)
    elseif endswith(file,".sdd")
        parse_sdd_file(file)
    else
        throw("Cannot parse this file type as a circuit: $file")
    end
end

#####################
# parser of logistic circuit file format
#####################

const parens = r"\(([^\)]+)\)"

function parse_lc_decision_line(ln::String)::DecisionLine{LCElement}
    @assert startswith(ln, "D")
    head::SubString, tail::SubString = split(ln,'(',limit=2)
    head_tokens = split(head)
    head_ints::Vector{UInt32} = map(x->parse(UInt32,x),head_tokens[2:4])
    elems_str::String = "("*tail
    elems = Vector{LCElement}()
    for x in eachmatch(parens::Regex, elems_str)
        tokens = split(x[1], limit=3)
        weights::Vector{Float64} = map(x->parse(Float64,x), split(tokens[3]))
        elem = LCElement(parse(UInt32,tokens[1]), parse(UInt32,tokens[2]), weights)
        push!(elems,elem)
    end
    DecisionLine(head_ints[1],head_ints[2],head_ints[3],elems)
end

function parse_lc_literal_line(ln::String)::WeightedLiteralLine
    @assert startswith(ln, "T") || startswith(ln, "F")
    tokens = split(ln)
    head_ints = map(x->parse(UInt32,x),tokens[2:4])
    weights = map(x->parse(Float64,x), tokens[5:end])
    lit = var2lit(head_ints[3])
    if startswith(ln, "F")
       lit = -lit # negative literal
    end
    WeightedLiteralLine(head_ints[1],head_ints[2],lit,true,weights)
end

function parse_comment_line(ln::String)
    @assert startswith(ln, "c")
    CircuitCommentLine(lstrip(chop(ln, head = 1, tail = 0)))
end

function parse_lc_header_line(ln::String)
    @assert (ln == "Logistic Circuit") || (ln == "Logisitic Circuit")
    LcHeaderLine()
end

function parse_bias_line(ln::String)::BiasLine
    @assert startswith(ln, "B")
    tokens = split(ln)
    weights = map(x->parse(Float64,x), tokens[2:end])
    BiasLine(weights)
end

parse_lc_file(file::String)::CircuitFormatLines = open(parse_lc_file, file)

function parse_lc_file(file::Core.IO)::CircuitFormatLines
    q = Vector{CircuitFormatLine}()
    for ln in eachline(file)
        @assert !isempty(ln)
        if ln[1] == 'D'
            push!(q, parse_lc_decision_line(ln))
        elseif ln[1] == 'T' || ln[1] == 'F'
            push!(q, parse_lc_literal_line(ln))
        elseif ln[1] == 'c'
            push!(q, parse_comment_line(ln))
        elseif ln[1] == 'L'
            push!(q, parse_lc_header_line(ln))
        elseif ln[1] == 'B'
            push!(q, parse_bias_line(ln))
        else
            error("Cannot parse logistic circuit file format line '$ln'")
        end
    end
    q
end

#####################
# parser for PSDD circuit file format
#####################

function parse_psdd_decision_line(ln::String)::DecisionLine{PSDDElement}
    @assert startswith(ln, "D")
    tokens = split(ln)
    head_ints::Vector{UInt32} = map(x->parse(UInt32,x),tokens[2:4])
    elems = Vector{PSDDElement}()
    for (p,s,w) in Iterators.partition(tokens[5:end],3)
        prime = parse(UInt32,p)
        sub = parse(UInt32,s)
        weight = parse(Float64,w)
        elem = PSDDElement(prime, sub, weight)
        push!(elems,elem)
    end
    DecisionLine(head_ints[1],head_ints[2],head_ints[3],elems)
end

function parse_psdd_true_leaf_line(ln::String)::WeightedNamedConstantLine
    @assert startswith(ln, "T")
    tokens = split(ln)
    @assert length(tokens)==5
    head_ints = map(x->parse(UInt32,x),tokens[2:4])
    weight = parse(Float64,tokens[5])
    WeightedNamedConstantLine(head_ints[1],head_ints[2],head_ints[3],weight)
end

function parse_literal_line(ln::String, normalized::Bool)::UnweightedLiteralLine
    @assert startswith(ln, "L")
    tokens = split(ln)
    @assert length(tokens)==4 "line has too many tokens: $ln"
    head_ints = map(x->parse(UInt32,x),tokens[2:3])
    lit = parse(Int32,tokens[4])
    UnweightedLiteralLine(head_ints[1],head_ints[2],lit,normalized)
end

function parse_psdd_header_line(ln::String)
    @assert startswith(ln, "psdd")
    PsddHeaderLine(parse(Int,split(ln)[2]))
end

parse_psdd_file(file::String)::CircuitFormatLines = open(parse_psdd_file, file)

function parse_psdd_file(file::Core.IO)::CircuitFormatLines
    q = Vector{CircuitFormatLine}()
    for ln in eachline(file)
        @assert !isempty(ln)
        if ln[1] == 'D'
            push!(q, parse_psdd_decision_line(ln))
        elseif ln[1] == 'T'
            push!(q, parse_psdd_true_leaf_line(ln))
        elseif ln[1] == 'L'
            push!(q, parse_literal_line(ln, true))
        elseif ln[1] == 'c'
            push!(q, parse_comment_line(ln))
        elseif startswith(ln,"psdd")
            push!(q, parse_psdd_header_line(ln))
        else
            error("Cannot parse PSDD file format line '$ln'")
        end
    end
    q
end

#####################
# parser for SDD circuit file format
#####################

function parse_sdd_decision_line(ln::String)::DecisionLine{SDDElement}
    @assert startswith(ln, "D")
    tokens = split(ln)
    head_ints::Vector{UInt32} = map(x->parse(UInt32,x),tokens[2:4])
    elems = Vector{SDDElement}()
    for (p,s) in Iterators.partition(tokens[5:end],2)
        prime = parse(UInt32,p)
        sub = parse(UInt32,s)
        elem = SDDElement(prime, sub)
        push!(elems,elem)
    end
    DecisionLine(head_ints[1],head_ints[2],head_ints[3],elems)
end

function parse_sdd_constant_leaf_line(ln::String)::AnonymousConstantLine
    @assert startswith(ln, "T") || startswith(ln, "F")
    tokens = split(ln)
    @assert length(tokens)==2
    AnonymousConstantLine(parse(UInt32,tokens[2]), startswith(ln, "T"), false)
end

function parse_sdd_header_line(ln::String)
    @assert startswith(ln, "sdd")
    SddHeaderLine(parse(Int,split(ln)[2]))
end

parse_sdd_file(file::String)::CircuitFormatLines = open(parse_sdd_file, file)

function parse_sdd_file(file::Core.IO)::CircuitFormatLines
    q = Vector{CircuitFormatLine}()
    for ln in eachline(file)
        @assert !isempty(ln)
        if ln[1] == 'D'
            push!(q, parse_sdd_decision_line(ln))
        elseif ln[1] == 'T' || ln[1] == 'F'
            push!(q, parse_sdd_constant_leaf_line(ln))
        elseif ln[1] == 'L'
            push!(q, parse_literal_line(ln, false))
        elseif ln[1] == 'c'
            push!(q, parse_comment_line(ln))
        elseif startswith(ln,"sdd")
            push!(q, parse_sdd_header_line(ln))
        else
            error("Cannot parse SDD file format line '$ln'")
        end
    end
    q
end

#####################
# loader for CNF/DNF file format
#####################

"""
    load_cnf(file::String; dual=false)::PlainLogicCircuit

Load a CNF/DNF as a logic circuit from file.
Suppported file formats:
* ".cnf" for CNF files  (`dual=False`)
* ".dnf" for DNF files  (`dual=True`)
"""
function load_cnf(file::String; dual=false)::PlainLogicCircuit
    @assert !dual && endswith(file, ".cnf") || dual && endswith(file, ".dnf")

    Clause = dual ? Plain⋀Node : Plain⋁Node
    # record the current clause
    clause = Clause([])
    # linearized clauses (disjunctions)
    clauses = Vector{Clause}()

    # literal cache is responsible for making leaf literals nodes unique and adding them to `circuit`
    lit_cache = Dict{Lit,PlainLiteralNode}()

    open(file) do file

        for ln in eachline(file)
            @assert !isempty(ln)
            if ln[1] == 'c' ||
                !dual && startswith(ln, "p cnf") ||
                dual && startswith(ln, "p dnf")
                # skip comment and header lines
                continue
            elseif ln[1] == '%'
                # some CNF files appear to have trailing lines after '%' which should not be parsed
                break
            else
                tokens = split(ln)
                for token in tokens
                    occursin(r"^\s*[-]?[0-9]+\s*$", token) || error("Cannot parse CNF file format line '$ln'")
                    literal = parse(Lit, token)
                    if literal == 0
                        push!(clauses, clause)
                        clause = Clause([])
                    else
                        push!(clause.children, get!(lit_cache, literal) do
                            PlainLiteralNode(literal)
                        end)
                    end
                end
            end
        end

    end

    @assert isempty(clause.children)

    dual ? disjoin(clauses) : conjoin(clauses)
end

"""
    load_dnf(file::String)::PlainLogicCircuit

Load a DNF as a logical circuit from file.
"""
load_dnf(file::String)::PlainLogicCircuit =
    load_cnf(file; dual=true)


module GeoPhyInv

# include modules (note: due to dependencies, order is important!)
include("Interpolation/Interpolation.jl")
include("Utils/Utils.jl")
include("Operators.jl")
include("Smooth.jl")
include("IO.jl")
include("Models/Models.jl")
include("Acquisition/Acquisition.jl")
include("Coupling.jl")
include("Data/Data.jl")
include("Born/Born.jl")
include("Fdtd/Fdtd.jl")
include("FWI/FWI.jl")
include("Poisson/Poisson.jl")
include("Gallery/Gallery.jl")
include("Plots.jl")

# export the Expt for Seismic Forward Modelling
const SeisForwExpt=GeoPhyInv.Fdtd.Param
export SeisForwExpt
mod!(a::SeisForwExpt)=GeoPhyInv.Fdtd.mod!(a)

# export the Expt for Seismic Inversion
const SeisInvExpt=GeoPhyInv.FWI.Param
export SeisInvExpt

# export the Expt for Poisson
const PoissonExpt=GeoPhyInv.Poisson.ParamExpt
export PoissonExpt
mod!(a::PoissonExpt,b,c)=GeoPhyInv.Poisson.mod!(a,b,c)
mod!(a::PoissonExpt,b)=GeoPhyInv.Poisson.mod!(a,b)
mod!(a::PoissonExpt)=GeoPhyInv.Poisson.mod!(a)
operator_Born(a::PoissonExpt,b)=GeoPhyInv.Poisson.operator_Born(a,b)

export mod!, operator_Born

const JuMIT=GeoPhyInv
export JuMIT

const GIPh=GeoPhyInv
export GIPh

const J=GeoPhyInv
export J

const GInterp=GeoPhyInv.Interpolation
export GInterp

const JP=GeoPhyInv.Plots
export JP

const JF=GeoPhyInv.FWI
export JF

const JD=GeoPhyInv.Data
export JD

const JA=GeoPhyInv.Acquisition
export JA

const JG=GeoPhyInv.Gallery
export JG

end # module


# function to generate a Kitaev colorcode
function generateBondColorcodeKitaev(
        lattice :: L
    ) :: Dict{LB,RGB} where {
        LB, N, U,
        S <: AbstractSite,
        B <: AbstractBond{LB,N},
        L <: AbstractLattice{S,B,U}
    }

    # put the default kitaev labels with default colors into a dictonary
    return Dict{LB,RGB}(
        getDefaultLabelX(LB) => RGB(1.0, 0.0, 0.0),
        getDefaultLabelY(LB) => RGB(0.0, 1.0, 0.0),
        getDefaultLabelZ(LB) => RGB(0.0, 0.0, 1.0)
    )
end

# function to generate a Kitaev colorcode
function generateBondColorcodeRandom(lattice::AbstractLattice)
    # TODO
    # This would be way better:
    # labels = unique(label.(bonds(lattice)))
    # colors = distinguishable_colors(length(labels)) # or a colormap
    # Dict([lb => color for (lb, color) in zip(labels, colors)])

    return Dict([
        lb => rand(RGB) for lb in unique(label.(bonds(lattice)))
    ])
end


export iauCp
"""
Copy a p-vector.

This function is part of the International Astronomical Union's
SOFA (Standards Of Fundamental Astronomy) software collection.

Status:  vector/matrix support function.

Given:
   p        double[3]     p-vector to be copied

Returned:
   c        double[3]     copy

This revision:  2013 June 18

SOFA release 2018-01-30

Copyright (C) 2018 IAU SOFA Board.  See notes at end.
"""

# void iauCp(double p[3], double c[3])

function iauCp(p::AbstractVector{<:Real})

   # Allocate return value
   c = zeros(Float64, 3)

   ccall((:iauCp, libsofa_c), Cvoid,
         (Ptr{Cdouble}, Ptr{Cdouble}),
         convert(Array{Float64, 1}, p), c)

   return SVector{3}(c)
end


# Code taken from FluxML/Optimisers.jl

const Param{T<:Number} = Union{AbstractArray{T},T}

_apply(opt, x, x̄, state) = apply(opt, x, x̄, state)
_apply(opt, x, x̄, ::Nothing) = apply(opt, x, x̄)

# Immutable updates
update!(opt, x, xb, state=state_init(opt, x)) = _update!(opt, x, xb, state)
update(opt, x, xb, state=state_init(opt, x))  = _update(opt, x, xb, state)

function _update(opt, x::Param, x̄::Param, state = nothing)
  Δ, state = _apply(opt, x, x̄, state)
  return x .- Δ, state
end

# Mutable updates

# Figure out if we can do in-place
inplace(x, y) = false
inplace(x, y::Nothing) = true
inplace(x::AbstractArray, x̄::AbstractArray) = true
inplace(x, x̄::NamedTuple) = all(inplace(getfield(x, f), getfield(x̄, f)) for f in fieldnames(typeof(x̄)))

function _update!(opt, x::AbstractArray{<:Number}, x̄::AbstractArray, state = nothing)
  Δ, state = _apply(opt, x, x̄, state)
  #println("fr = ", x)
  #println("to = ", Δ)
  x .-= Δ
  #println("after = ", x)
  return state
end

#function _update!(opt, x::Tuple, x̄::Tuple, state=nothing)
function _update!(opt, x, x̄::Tuple, state=nothing)
  for f in propertynames(x̄)
    f̄ = getindex(x̄, f)

    f̄ === nothing || _update!(opt, getindex(x, f), f̄, state)
  end
  return state
end

function _update!(opt, x, x̄::NamedTuple, state=nothing)
  for f in propertynames(x̄)
    f == :tuple_all_weights && continue
    f̄ = getproperty(x̄, f)

    f̄ === nothing || _update!(opt, getfield(x, f), f̄, state)
  end
  return state
end

function _update!(opt, x, x̄, state=nothing)
  for f in propertynames(x̄)
    if f isa Int
      f̄ = getindex(x̄, f)
      f̄ === nothing || _update!(opt, getindex(x, f), f̄, state)
    else
      f̄ = getproperty(x̄, f)
      f̄ === nothing || _update!(opt, getfield(x, f), f̄, state)
    end
  end
  return state
end


"""
    predict(method, model, input[; device, context])

Predict a `target` from `input` using `model`. Optionally apply function `device`
to `x` before passing to `model` and use `context` instead of the default
context [`Inference`](#).
"""
function predict(method, model, input; device = identity, undevice = identity, context = Inference())
    if shouldbatch(method)
        return predictbatch(method, model, [input,]; device = device, undevice = undevice, context = context) |> only
    else
        return decodeŷ(method, context, undevice(model(device(encodeinput(method, context, input)))))
    end
end

"""
    predictbatch(method, model, inputs[; device, context])

Predict `targets` from a vector of `inputs` using `model` by batching them.
Optionally apply function `device` to batch before passing to `model` and
use `context` instead of the default [`Inference`](#).
"""
function predictbatch(method, model, inputs; device = identity, undevice = identity, context = Inference())
    xs = device(DataLoaders.collate([copy(encodeinput(method, context, input)) for input in inputs]))
    ŷs = undevice(model(xs))
    targets = [decodeŷ(method, context, ŷ) for ŷ in DataLoaders.obsslices(ŷs)]
    return targets
end

function _predictx(method, model, x, device = identity)
    if shouldbatch(method)
        x = DataLoaders.collate([x])
    end
    ŷs = device(model)(device(x))
    if shouldbatch(method)
        ŷ = ŷs[((:) for _ in 1:ndims(ŷs)-1)..., 1]
    end
    return ŷ
end


#
# This file is part of the DiscreteEvents.jl Julia package, MIT license
#
# Paul Bayer, 2019
#
# This is a Julia package for discrete event simulation
#
using Printf, Statistics, DiscreteEvents, BenchmarkTools
import Dates.now

struct RES
    min::Float64
    mean::Float64
    nevents::Int64
    nsample::Int64
    nactions::Int64
    minact::Float64
    meanact::Float64
end
Base.Array(r::RES) = [r.min, r.mean, r.nevents, r.nsample, r.nactions, r.minact, r.meanact]
function benchresult(trial::BenchmarkTools.Trial, clk::Clock)
    _min = minimum(t.times)*1e-6
    _mean = mean(t.times)*1e-6
    nevents = clk.evcount
    nsample = clk.scount
    nactions = clk.evcount+clk.scount
    minact = _min/nactions*1e3
    meanact = _mean/nactions*1e3
    return RES(_min,_mean,nevents,nsample,nactions,minact,meanact)
end
res = Dict()

btime = now()

println("... channel benchmark ...")
include("bench_channel.jl")
res["channel"] = benchresult(t, 𝐶)

println("... dice benchmark ...")
include("bench_dice.jl")
res["dice"] = benchresult(t, clk)

println("... heating benchmark ...")
include("bench_heating.jl")
res["heating"]  = benchresult(t, 𝐶)

arr = sum(hcat(Array(res["channel"]), Array(res["dice"]), Array(res["heating"])), dims=2)
arr[6] = arr[1]/arr[5]*1e3
arr[7] = arr[2]/arr[5]*1e3
res["sum"] = RES(arr...)

println()
println("Benchmark results:")
println("==================")
@printf "time    \tdatetime               \tchannel [ms] \tdice [ms] \theating [ms] \t    sum [ms]\n"
@printf "minimum \t%s\t%12.3f \t%9.3f \t%12.3f \t%12.3f\n" btime res["channel"].min res["dice"].min res["heating"].min res["sum"].min
@printf "mean    \t%s\t%12.3f \t%9.3f \t%12.3f \t%12.3f\n" btime res["channel"].mean res["dice"].mean res["heating"].mean res["sum"].mean
println()
println("action times:")
println("-------------")
@printf "time    \tdatetime               \tchannel [μs] \tdice [μs] \theating [μs] \toverall [μs]\n"
@printf "minimum \t%s\t%12.3f \t%9.3f \t%12.3f \t%12.3f\n" btime res["channel"].minact res["dice"].minact res["heating"].minact res["sum"].minact
@printf "mean    \t%s\t%12.3f \t%9.3f \t%12.3f \t%12.3f\n" btime res["channel"].meanact res["dice"].meanact res["heating"].meanact res["sum"].meanact


Base.run(`clear`)
println("Starting")
println("----------------------------")
using Leapfrog
println("JIT compilation")
mp, cost = init_all(1);
stags, ess, number_of_equilibria = solve(cost, mp);
println("----------------------------")
println("Checking # of equilibria")
println(number_of_equilibria)
println("----------------------------")
println("Timing code")
times = [@elapsed solve(cost,mp) for i = 1:10];
times_mean = mean(times)
times_min = minimum(times)
times_max = maximum(times)
times_var = std(times)
println("----------------------------")
@printf("Mean    Min    Max    Std\n")
@printf("%4.4f  %4.4f %4.4f %4.4f\n", times_mean, times_min, times_max, times_var)
println("----------------------------")


export add_scalar_inverse

function add_scalar_inverse(mgr::Ptr{Manager}, f::Ptr{Node}, epsilon::Ptr{Node})
    result = ccall((:Cudd_addScalarInverse, libcudd),
        Ptr{Node}, (Ptr{Manager}, Ptr{Node}, Ptr{Node}), mgr, f, epsilon)
    if result == C_NULL
        throw(OutOfMemoryError())
    end
    return result
end


# ---
# jupyter:
#   jupytext:
#     formats: ipynb,jl:hydrogen
#     text_representation:
#       extension: .jl
#       format_name: hydrogen
#       format_version: '1.3'
#       jupytext_version: 1.11.2
#   kernelspec:
#     display_name: Julia 1.8.0-DEV
#     language: julia
#     name: julia-1.8
# ---

# %% [markdown]
# https://discourse.julialang.org/t/handling-of-package-inter-dependency/64568/24

# %%
module SCI
export fun1,fun2
fun1(x) = x
fun2(x) = x^2
end

module COM
using ..SCI
"""function to initialize a container"""
init_container() = []
"""function to read a message from a pipe, assume it reads from the argument for simplicity"""
read_message(x) = x
# function to check if instruction is authorized (omitted for simplicity)

"""function doing something based on the message received"""
function compute!(container, msg)
    if msg == :fun1
        push!(container, fun1(rand()))
    elseif msg ==:fun2
        push!(container, fun2(rand()))
    end
    return container
end
end

# Load COM+SCI utilities
using .COM

@show container = COM.init_container()
@show msg = COM.read_message(:fun1)
@show container = COM.compute!(container, msg)
@show msg = COM.read_message(:fun2)
@show container = COM.compute!(container, msg);

# %%

0012/BambOoxX's MWE with minor corrections.jl


module ColorVectorSpaceTests

using FactCheck, Base.Test, ColorVectorSpace, ColorTypes, FixedPointNumbers, Compat

macro test_colortype_approx_eq(a, b)
    :(test_colortype_approx_eq($(esc(a)), $(esc(b)), $(string(a)), $(string(b))))
end

u8sum(x,y) = Float64(UFixed8(x)) + Float64(UFixed8(y))

FactCheck.roughly(x::Gray) = (y::Gray) -> isapprox(y, x)
FactCheck.roughly(x::RGB) = (y::RGB) -> isapprox(y, x)

facts("Colortypes") do
    function test_colortype_approx_eq(a::Colorant, b::Colorant, astr, bstr)
        @fact typeof(a) --> typeof(b)
        n = length(fieldnames(typeof(a)))
        for i = 1:n
            @fact getfield(a, i) --> roughly(getfield(b,i))
        end
    end

    context("nan") do
        function make_checked_nan{T}(::Type{T})
            x = nan(T)
            isa(x, T) && isnan(x)
        end
        for S in (Float32, Float64)
            @fact make_checked_nan(S) --> true
            @fact make_checked_nan(Gray{S}) --> true
            @fact make_checked_nan(AGray{S}) --> true
            @fact make_checked_nan(GrayA{S}) --> true
            @fact make_checked_nan(RGB{S}) --> true
            @fact make_checked_nan(ARGB{S}) --> true
            @fact make_checked_nan(ARGB{S}) --> true
        end
    end

    context("Arithmetic with Gray") do
        cf = Gray{Float32}(0.1)
        ccmp = Gray{Float32}(0.2)
        @fact 2*cf --> ccmp
        @fact cf*2 --> ccmp
        @fact ccmp/2 --> cf
        @fact 2.0f0*cf --> ccmp
        @test_colortype_approx_eq cf*cf Gray{Float32}(0.01)
        @test_colortype_approx_eq cf^2 Gray{Float32}(0.01)
        @test_colortype_approx_eq cf^3.0f0 Gray{Float32}(0.001)
        @fact eltype(2.0*cf) --> Float64
        @fact abs2(ccmp) --> 0.2f0^2
        @fact norm(cf) --> 0.1f0
        @fact sumabs2(ccmp) --> 0.2f0^2
        cu = Gray{U8}(0.1)
        @fact 2*cu --> Gray(2*cu.val)
        @fact 2.0f0*cu --> Gray(2.0f0*cu.val)
        f = U8(0.5)
        @fact (f*cu).val --> roughly(f*cu.val)
        @fact 2.*cf --> ccmp
        @fact cf.*2 --> ccmp
        @fact cf/2.0f0 --> Gray{Float32}(0.05)
        @fact cu/2 --> Gray(cu.val/2)
        @fact cu/0.5f0 --> Gray(cu.val/0.5f0)
        @fact cf+cf --> ccmp
        @fact cf --> isfinite
        @fact cf --> not(isinf)
        @fact cf --> not(isnan)
        @fact Gray(NaN) --> not(isfinite)
        @fact Gray(NaN) --> not(isinf)
        @fact Gray(NaN) --> isnan
        @fact Gray(Inf) --> not(isfinite)
        @fact Gray(Inf) --> isinf
        @fact Gray(Inf) --> not(isnan)
        @fact abs(Gray(0.1)) --> roughly(0.1)
        @fact eps(Gray{U8}) --> Gray(eps(U8))  # #282

        acu = Gray{U8}[cu]
        acf = Gray{Float32}[cf]
        @fact typeof(acu+acf) --> Vector{Gray{Float32}}
        @fact typeof(acu-acf) --> Vector{Gray{Float32}}
        @fact typeof(acu.+acf) --> Vector{Gray{Float32}}
        @fact typeof(acu.-acf) --> Vector{Gray{Float32}}
        @fact typeof(acu+cf) --> Vector{Gray{U8}}
        @fact typeof(acu-cf) --> Vector{Gray{U8}}
        @fact typeof(acu.+cf) --> Vector{Gray{U8}}
        @fact typeof(acu.-cf) --> Vector{Gray{U8}}
        @fact typeof(2*acf) --> Vector{Gray{Float32}}
        @fact typeof(2.*acf) --> Vector{Gray{Float32}}
        @fact typeof(0x02*acu) --> Vector{Gray{Float32}}
        @fact typeof(acu/2) --> Vector{Gray{typeof(U8(0.5)/2)}}
        @fact typeof(acf.^2) --> Vector{Gray{Float32}}
        @fact (acu/Gray{U8}(0.5))[1] --> gray(acu[1])/U8(0.5)
        @fact (acf/Gray{Float32}(2))[1] --> roughly(0.05f0)
        @fact (acu/2)[1] --> Gray(gray(acu[1])/2)
        @fact (acf/2)[1] --> roughly(Gray{Float32}(0.05f0))
        @fact sumabs2([cf,ccmp]) --> roughly(0.05f0)

        @fact gray(0.8) --> 0.8

        a = Gray{U8}[0.8,0.7]
        @fact sum(a) --> Gray(u8sum(0.8,0.7))
        @fact abs( var(a) - (a[1]-a[2])^2 / 2 ) --> less_than(0.001)
        @fact isapprox(a, a) --> true
        @fact real(Gray{Float32}) <: Real --> true
        @fact zero(ColorTypes.Gray)-->0
        @fact one(ColorTypes.Gray)-->1
        a = Gray{U8}[0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0]
        @fact histrange(a,10)-->0.1:0.1:1

    end

    context("Comparisons with Gray") do
        g1 = Gray{U8}(0.2)
        g2 = Gray{U8}(0.3)
        @fact isless(g1, g2) --> true
        @fact isless(g2, g1) --> false
        @fact g1 < g2 --> true
        @fact g2 < g1 --> false
        @fact isless(g1, 0.5) --> true
        @fact isless(0.5, g1) --> false
        @fact g1 < 0.5 --> true
        @fact 0.5 < g1 --> false
        @fact @inferred(max(g1, g2)) --> g2
        @fact @inferred(max(g1, 0.1)) --> 0.2
        @fact @inferred(min(g1, g2)) --> g1
        @fact @inferred(min(g1, 0.1)) --> 0.1
        a = Gray{Float64}(0.9999999999999999)
        b = Gray{Float64}(1.0)

        @fact isapprox(a, b) --> true
        a = Gray{Float64}(0.99)
        @fact isapprox(a, b, rtol = 0.01) --> false
        @fact isapprox(a, b, rtol = 0.1) --> true
    end

    context("Unary operations with Gray") do
        for g in (Gray(0.4), Gray{U8}(0.4))
            for op in ColorVectorSpace.unaryOps
                try
                    v = @eval $op(gray(g))  # if this fails, don't bother
                    @fact $op(g) --> v
                end
            end
        end
        u = U8(0.4)
        @fact ~Gray(u) --> Gray(~u)
        @fact -Gray(u) --> Gray(-u)
    end

    context("Arithmetic with GrayA") do
        p1 = GrayA{Float32}(Gray(0.8), 0.2)
        p2 = GrayA{Float32}(Gray(0.6), 0.3)
        @test_colortype_approx_eq p1+p2 GrayA{Float32}(Gray(1.4),0.5)
        @test_colortype_approx_eq (p1+p2)/2 GrayA{Float32}(Gray(0.7),0.25)
        @test_colortype_approx_eq 0.4f0*p1+0.6f0*p2 GrayA{Float32}(Gray(0.68),0.26)
        @test_colortype_approx_eq ([p1]+[p2])[1] GrayA{Float32}(Gray(1.4),0.5)
        @test_colortype_approx_eq ([p1].+[p2])[1] GrayA{Float32}(Gray(1.4),0.5)
        @test_colortype_approx_eq ([p1]+p2)[1] GrayA{Float32}(Gray(1.4),0.5)
        @test_colortype_approx_eq ([p1].+p2)[1] GrayA{Float32}(Gray(1.4),0.5)
        @test_colortype_approx_eq ([p1]-[p2])[1] GrayA{Float32}(Gray(0.2),-0.1)
        @test_colortype_approx_eq ([p1].-[p2])[1] GrayA{Float32}(Gray(0.2),-0.1)
        @test_colortype_approx_eq ([p1]-p2)[1] GrayA{Float32}(Gray(0.2),-0.1)
        @test_colortype_approx_eq ([p1].-p2)[1] GrayA{Float32}(Gray(0.2),-0.1)
        @test_colortype_approx_eq ([p1]/2)[1] GrayA{Float32}(Gray(0.4),0.1)
        @test_colortype_approx_eq (0.4f0*[p1]+0.6f0*[p2])[1] GrayA{Float32}(Gray(0.68),0.26)

        a = GrayA{U8}[GrayA(0.8,0.7), GrayA(0.5,0.2)]
        @fact sum(a) --> GrayA(u8sum(0.8,0.5), u8sum(0.7,0.2))
        @fact isapprox(a, a) --> true
        a = AGray{Float64}(1.0, 0.9999999999999999)
        b = AGray{Float64}(1.0, 1.0)

        @fact isapprox(a, b) --> true
        a = AGray{Float64}(1.0, 0.99)
        @fact isapprox(a, b, rtol = 0.01) --> false
        @fact isapprox(a, b, rtol = 0.1) --> true

    end

    context("Arithemtic with RGB") do
        cf = RGB{Float32}(0.1,0.2,0.3)
        ccmp = RGB{Float32}(0.2,0.4,0.6)
        @fact 2*cf --> ccmp
        @fact cf*2 --> ccmp
        @fact ccmp/2 --> cf
        @fact 2.0f0*cf --> ccmp
        @fact eltype(2.0*cf) --> Float64
        cu = RGB{U8}(0.1,0.2,0.3)
        @test_colortype_approx_eq 2*cu RGB(2*cu.r, 2*cu.g, 2*cu.b)
        @test_colortype_approx_eq 2.0f0*cu RGB(2.0f0*cu.r, 2.0f0*cu.g, 2.0f0*cu.b)
        f = U8(0.5)
        @fact (f*cu).r --> roughly(f*cu.r)
        @fact 2.*cf --> ccmp
        @fact cf.*2 --> ccmp
        @fact cf/2.0f0 --> RGB{Float32}(0.05,0.1,0.15)
        @fact cu/2 --> RGB(cu.r/2,cu.g/2,cu.b/2)
        @fact cu/0.5f0 --> RGB(cu.r/0.5f0, cu.g/0.5f0, cu.b/0.5f0)
        @fact cf+cf --> ccmp
        @fact cu * 1//2 --> roughly(RGB{Float64}(U8(0.1)/2, U8(0.2)/2, U8(0.3)/2))
        @test_colortype_approx_eq (cf*[0.8f0])[1] RGB{Float32}(0.8*0.1,0.8*0.2,0.8*0.3)
        @test_colortype_approx_eq ([0.8f0]*cf)[1] RGB{Float32}(0.8*0.1,0.8*0.2,0.8*0.3)
        @test_colortype_approx_eq (cf.*[0.8f0])[1] RGB{Float32}(0.8*0.1,0.8*0.2,0.8*0.3)
        @test_colortype_approx_eq ([0.8f0].*cf)[1] RGB{Float32}(0.8*0.1,0.8*0.2,0.8*0.3)
        @fact cf --> isfinite
        @fact cf --> not(isinf)
        @fact cf --> not(isnan)
        @fact RGB(NaN, 1, 0.5) --> not(isfinite)
        @fact RGB(NaN, 1, 0.5) --> not(isinf)
        @fact RGB(NaN, 1, 0.5) --> isnan
        @fact RGB(1, Inf, 0.5) --> not(isfinite)
        @fact RGB(1, Inf, 0.5) --> isinf
        @fact RGB(1, Inf, 0.5) --> not(isnan)
        @fact abs(RGB(0.1,0.2,0.3)) --> roughly(0.6)
        @fact sumabs2(RGB(0.1,0.2,0.3)) --> roughly(0.14)
        @fact norm(RGB(0.1,0.2,0.3)) --> roughly(sqrt(0.14))

        acu = RGB{U8}[cu]
        acf = RGB{Float32}[cf]
        @fact typeof(acu+acf) --> Vector{RGB{Float32}}
        @fact typeof(acu-acf) --> Vector{RGB{Float32}}
        @fact typeof(acu.+acf) --> Vector{RGB{Float32}}
        @fact typeof(acu.-acf) --> Vector{RGB{Float32}}
        @fact typeof(acu+cf) --> Vector{RGB{U8}}
        @fact typeof(acu-cf) --> Vector{RGB{U8}}
        @fact typeof(acu.+cf) --> Vector{RGB{U8}}
        @fact typeof(acu.-cf) --> Vector{RGB{U8}}
        @fact typeof(2*acf) --> Vector{RGB{Float32}}
        @fact typeof(convert(UInt8, 2)*acu) --> Vector{RGB{Float32}}
        @fact typeof(acu/2) --> Vector{RGB{typeof(U8(0.5)/2)}}

        a = RGB{U8}[RGB(1,0,0), RGB(1,0.8,0)]
        @fact sum(a) --> RGB(2.0,0.8,0)
        @fact isapprox(a, a) --> true
        a = RGB{Float64}(1.0, 1.0, 0.9999999999999999)
        b = RGB{Float64}(1.0, 1.0, 1.0)

        @fact isapprox(a, b) --> true
        a = RGB{Float64}(1.0, 1.0, 0.99)
        @fact isapprox(a, b, rtol = 0.01) --> false
        @fact isapprox(a, b, rtol = 0.1) --> true
    end

    context("Arithemtic with RGBA") do
        cf = RGBA{Float32}(0.1,0.2,0.3,0.4)
        ccmp = RGBA{Float32}(0.2,0.4,0.6,0.8)
        @fact 2*cf --> ccmp
        @fact cf*2 --> ccmp
        @fact ccmp/2 --> cf
        @fact 2.0f0*cf --> ccmp
        @fact eltype(2.0*cf) --> Float64
        cu = RGBA{U8}(0.1,0.2,0.3,0.4)
        @test_colortype_approx_eq 2*cu RGBA(2*cu.r, 2*cu.g, 2*cu.b, 2*cu.alpha)
        @test_colortype_approx_eq 2.0f0*cu RGBA(2.0f0*cu.r, 2.0f0*cu.g, 2.0f0*cu.b, 2.0f0*cu.alpha)
        f = U8(0.5)
        @fact (f*cu).r --> roughly(f*cu.r)
        @fact 2.*cf --> ccmp
        @fact cf.*2 --> ccmp
        @fact cf/2.0f0 --> RGBA{Float32}(0.05,0.1,0.15,0.2)
        @fact cu/2 --> RGBA(cu.r/2,cu.g/2,cu.b/2,cu.alpha/2)
        @fact cu/0.5f0 --> RGBA(cu.r/0.5f0, cu.g/0.5f0, cu.b/0.5f0, cu.alpha/0.5f0)
        @fact cf+cf --> ccmp
        @test_colortype_approx_eq (cf*[0.8f0])[1] RGBA{Float32}(0.8*0.1,0.8*0.2,0.8*0.3,0.8*0.4)
        @test_colortype_approx_eq ([0.8f0]*cf)[1] RGBA{Float32}(0.8*0.1,0.8*0.2,0.8*0.3,0.8*0.4)
        @test_colortype_approx_eq (cf.*[0.8f0])[1] RGBA{Float32}(0.8*0.1,0.8*0.2,0.8*0.3,0.8*0.4)
        @test_colortype_approx_eq ([0.8f0].*cf)[1] RGBA{Float32}(0.8*0.1,0.8*0.2,0.8*0.3,0.8*0.4)
        @fact cf --> isfinite
        @fact cf --> not(isinf)
        @fact cf --> not(isnan)
        @fact RGBA(NaN, 1, 0.5, 0.8) --> not(isfinite)
        @fact RGBA(NaN, 1, 0.5) --> not(isinf)
        @fact RGBA(NaN, 1, 0.5) --> isnan
        @fact RGBA(1, Inf, 0.5) --> not(isfinite)
        @fact RGBA(1, Inf, 0.5) --> isinf
        @fact RGBA(1, Inf, 0.5) --> not(isnan)
        @fact RGBA(0.2, 1, 0.5, NaN) --> not(isfinite)
        @fact RGBA(0.2, 1, 0.5, NaN) --> not(isinf)
        @fact RGBA(0.2, 1, 0.5, NaN) --> isnan
        @fact RGBA(0.2, 1, 0.5, Inf) --> not(isfinite)
        @fact RGBA(0.2, 1, 0.5, Inf) --> isinf
        @fact RGBA(0.2, 1, 0.5, Inf) --> not(isnan)
        @fact abs(RGBA(0.1,0.2,0.3,0.2)) --> roughly(0.8)

        acu = RGBA{U8}[cu]
        acf = RGBA{Float32}[cf]
        @fact typeof(acu+acf) --> Vector{RGBA{Float32}}
        @fact typeof(acu-acf) --> Vector{RGBA{Float32}}
        @fact typeof(acu.+acf) --> Vector{RGBA{Float32}}
        @fact typeof(acu.-acf) --> Vector{RGBA{Float32}}
        @fact typeof(acu+cf) --> Vector{RGBA{U8}}
        @fact typeof(acu-cf) --> Vector{RGBA{U8}}
        @fact typeof(acu.+cf) --> Vector{RGBA{U8}}
        @fact typeof(acu.-cf) --> Vector{RGBA{U8}}
        @fact typeof(2*acf) --> Vector{RGBA{Float32}}
        @fact typeof(convert(UInt8, 2)*acu) --> Vector{RGBA{Float32}}
        @fact typeof(acu/2) --> Vector{RGBA{typeof(U8(0.5)/2)}}

        a = RGBA{U8}[RGBA(1,0,0,0.8), RGBA(0.7,0.8,0,0.9)]
        @fact sum(a) --> RGBA(u8sum(1,0.7),0.8,0,u8sum(0.8,0.9))
        @fact isapprox(a, a) --> true
        a = ARGB{Float64}(1.0, 1.0, 1.0, 0.9999999999999999)
        b = ARGB{Float64}(1.0, 1.0, 1.0, 1.0)

        @fact isapprox(a, b) --> true
        a = ARGB{Float64}(1.0, 1.0, 1.0, 0.99)
        @fact isapprox(a, b, rtol = 0.01) --> false
        @fact isapprox(a, b, rtol = 0.1) --> true
    end
end

isinteractive() || FactCheck.exitstatus()

end


function contact_forces(u, p, t)
    @unpack footang, g, iff, ihat, irf, ish, ith, k1, k2, k3, k4, k5, k6, k7, k8, lff, lffo, lhat, lmtpxi, lrf, lrff, lrffo, lrfo, lsh, lsho, lth, ltho, ltoexi, mff, mhat, mrf, msh, mth, mtpb, mtpk, rmtpxi, rtoexi, u4, u5, u6, u7, u8, u9, z = p
    @inbounds q1, q2, q3, u1, u2, u3 = u

    # specified variables
    rmtp = _rmtp(t)
    ra = _ra(t)
    rk = _rk(t)
    rh = _rh(t)
    la = _la(t)
    lmtp = _lmtp(t)
    lh = _lh(t)
    lk = _lk(t)
    lkp = _lkp(t)
    lap = _lap(t)
    lhp = _lhp(t)
    lmtpp = _lmtpp(t)
    rhp = _rhp(t)
    rkp = _rkp(t)
    rap = _rap(t)
    rmtpp = _rmtpp(t)

    # grf
    pop2y = q2 - lff * sin(la + lh + lmtp - lk - q3)
    pop10y = q2 + lth * sin(lh - q3) + lsh * sin(rh - rk - q3) + lrf * sin(la + lh - lk - q3) - lth * sin(rh - q3) - lsh * sin(lh - lk - q3) - lrf * sin(ra + rh - rk - q3) - lff * sin(la + lh + lmtp - lk - q3)
    pop11y = q2 + lth * sin(lh - q3) + lsh * sin(rh - rk - q3) + lrf * sin(la + lh - lk - q3) + lff * sin(ra + rh + rmtp - rk - q3) - lth * sin(rh - q3) - lsh * sin(lh - lk - q3) - lrf * sin(ra + rh - rk - q3) - lff * sin(la + lh + lmtp - lk - q3)

    # left foot
    if q2 < 0.0
        dp1x = q1 - ltoexi
        lry1 = -k3 * q2 - k4 * abs(q2) * u2
        lrx1 = lry1 * (-k1 * dp1x - k2 * u1) # (-k1*(q1-pop1xi) - k2*u1)*ry1
    else
        lrx1 = 0.0
        lry1 = 0.0
    end

    if pop2y < 0.0
        pop2x = q1 + lff * cos(la + lh + lmtp - lk - q3)
        vop2x = u1 - lff * sin(la + lh + lmtp - lk - q3) * (lap + lhp + lmtpp - lkp - u3 - u5 - u7 - u9)
        vop2y = u2 - lff * cos(la + lh + lmtp - lk - q3) * (lap + lhp + lmtpp - lkp - u3 - u5 - u7 - u9)
        dp2x = pop2x - lmtpxi
        lry2 = -k7 * pop2y - k8 * abs(pop2y) * vop2y
        lrx2 = lry2 * (-k5 * dp2x - k6 * vop2x)
    else
        lrx2 = 0.0
        lry2 = 0.0
    end

    # right foot
    if pop11y < 0.0
        pop11x = q1 + lth * cos(rh - q3) + lsh * cos(lh - lk - q3) + lrf * cos(ra + rh - rk - q3) + lff * cos(la + lh + lmtp - lk - q3) - lth * cos(lh - q3) - lsh * cos(rh - rk - q3) - lrf * cos(la + lh - lk - q3) - lff * cos(ra + rh + rmtp - rk - q3)
        vop11x = u1 + lth * sin(lh - q3) * (lhp - u3 - u5) + lsh * sin(rh - rk - q3) * (rhp - rkp - u3 - u4 - u6) + lrf * sin(la + lh - lk - q3) * (lap + lhp + lmtpp - lkp - u3 - u5 - u7 - u9) + lff * sin(ra + rh + rmtp - rk - q3) * (rap + rhp + rmtpp - rkp - u3 - u4 - u6 - u8) - lth * sin(rh - q3) * (rhp - u3 - u4) - lsh * sin(lh - lk - q3) * (lhp - lkp - u3 - u5 - u7) - lrf * sin(ra + rh - rk - q3) * (rap + rhp - rkp - u3 - u4 - u6 - u8) - lff * sin(la + lh + lmtp - lk - q3) * (lap + lhp + lmtpp - lkp - u3 - u5 - u7 - u9)
        vop11y = u2 + lth * cos(lh - q3) * (lhp - u3 - u5) + lsh * cos(rh - rk - q3) * (rhp - rkp - u3 - u4 - u6) + lrf * cos(la + lh - lk - q3) * (lap + lhp + lmtpp - lkp - u3 - u5 - u7 - u9) + lff * cos(ra + rh + rmtp - rk - q3) * (rap + rhp + rmtpp - rkp - u3 - u4 - u6 - u8) - lth * cos(rh - q3) * (rhp - u3 - u4) - lsh * cos(lh - lk - q3) * (lhp - lkp - u3 - u5 - u7) - lrf * cos(ra + rh - rk - q3) * (rap + rhp - rkp - u3 - u4 - u6 - u8) - lff * cos(la + lh + lmtp - lk - q3) * (lap + lhp + lmtpp - lkp - u3 - u5 - u7 - u9)
        dp11x = pop11x - rtoexi
        rry1 = -k3 * pop11y - k4 * abs(pop11y) * vop11y
        rrx1 = rry1 * (-k1 * dp11x - k2 * vop11x) # (-k1*(q1-pop1xi) - k2*u1)*ry1
    else
        rrx1 = 0.0
        rry1 = 0.0
    end

    if pop10y < 0.0
        pop10x = q1 + lth * cos(rh - q3) + lsh * cos(lh - lk - q3) + lrf * cos(ra + rh - rk - q3) + lff * cos(la + lh + lmtp - lk - q3) - lth * cos(lh - q3) - lsh * cos(rh - rk - q3) - lrf * cos(la + lh - lk - q3)
        vop10x = u1 + lth * sin(lh - q3) * (lhp - u3 - u5) + lsh * sin(rh - rk - q3) * (rhp - rkp - u3 - u4 - u6) + lrf * sin(la + lh - lk - q3) * (lap + lhp + lmtpp - lkp - u3 - u5 - u7 - u9) - lth * sin(rh - q3) * (rhp - u3 - u4) - lsh * sin(lh - lk - q3) * (lhp - lkp - u3 - u5 - u7) - lrf * sin(ra + rh - rk - q3) * (rap + rhp - rkp - u3 - u4 - u6 - u8) - lff * sin(la + lh + lmtp - lk - q3) * (lap + lhp + lmtpp - lkp - u3 - u5 - u7 - u9)
        vop10y = u2 + lth * cos(lh - q3) * (lhp - u3 - u5) + lsh * cos(rh - rk - q3) * (rhp - rkp - u3 - u4 - u6) + lrf * cos(la + lh - lk - q3) * (lap + lhp + lmtpp - lkp - u3 - u5 - u7 - u9) - lth * cos(rh - q3) * (rhp - u3 - u4) - lsh * cos(lh - lk - q3) * (lhp - lkp - u3 - u5 - u7) - lrf * cos(ra + rh - rk - q3) * (rap + rhp - rkp - u3 - u4 - u6 - u8) - lff * cos(la + lh + lmtp - lk - q3) * (lap + lhp + lmtpp - lkp - u3 - u5 - u7 - u9)
        dp10x = pop10x - rmtpxi
        rry2 = -k7 * pop10y - k8 * abs(pop10y) * vop10y
        rrx2 = rry2 * (-k5 * dp10x - k6 * vop10x)
    else
        rrx2 = 0.0
        rry2 = 0.0
    end

    return lrx1, lry1, lrx2, lry2, rrx1, rry1, rrx2, rry2
end


mutable struct MazeDirected{T, R} <: AbstractGridWorld
    world::GridWorldBase{Tuple{Agent, Wall, Goal}}
    agent_pos::CartesianIndex{2}
    agent_dir::AbstractDirection
    reward::T
    rng::R
    terminal_reward::T
    goal_pos::CartesianIndex{2}
    done::Bool
end

@generate_getters(MazeDirected)
@generate_setters(MazeDirected)

mutable struct MazeUndirected{T, R} <: AbstractGridWorld
    world::GridWorldBase{Tuple{Agent, Wall, Goal}}
    agent_pos::CartesianIndex{2}
    reward::T
    rng::R
    terminal_reward::T
    goal_pos::CartesianIndex{2}
    done::Bool
end

@generate_getters(MazeUndirected)
@generate_setters(MazeUndirected)

#####
# Directed
#####

"""
Maze generation uses the iterative implementation of randomized depth-first search from [Wikipedia](https://en.wikipedia.org/wiki/Maze_generation_algorithm#Iterative_implementation).
"""
function MazeDirected(; T = Float32, height = 9, width = 9, rng = Random.GLOBAL_RNG)
    @assert isodd(height) && isodd(width) "height and width must be odd numbers"
    vertical_range = 2:2:height-1
    horizontal_range = 2:2:width-1

    objects = (AGENT, WALL, GOAL)
    world = GridWorldBase(objects, height, width)

    world[WALL, :, :] .= true
    for i in vertical_range, j in horizontal_range
        world[WALL, i, j] = false
    end

    goal_pos = CartesianIndex(height - 1, width - 1)
    world[GOAL, goal_pos] = true

    agent_pos = CartesianIndex(2, 2)
    world[AGENT, agent_pos] = true

    agent_dir = RIGHT
    reward = zero(T)
    terminal_reward = one(T)
    done = false

    env = MazeDirected(world, agent_pos, agent_dir, reward, rng, terminal_reward, goal_pos, done)

    RLBase.reset!(env)

    return env
end

RLBase.state_space(env::MazeDirected, ::RLBase.Observation, ::RLBase.DefaultPlayer) = nothing
const MAZE_DIRECTED_LAYERS = SA.SVector(2, 3)
RLBase.state(env::MazeDirected, ::RLBase.Observation, ::RLBase.DefaultPlayer) = get_grid(get_world(env), get_agent_pos(env), get_agent_dir(env), get_half_size(env), MAZE_DIRECTED_LAYERS)

RLBase.state_space(env::MazeDirected, ::RLBase.InternalState, ::RLBase.DefaultPlayer) = nothing
RLBase.state(env::MazeDirected, ::RLBase.InternalState, ::RLBase.DefaultPlayer) = (copy(get_grid(env)), get_agent_dir(env))

RLBase.action_space(env::MazeDirected, ::RLBase.DefaultPlayer) = DIRECTED_NAVIGATION_ACTIONS
RLBase.reward(env::MazeDirected, ::RLBase.DefaultPlayer) = get_reward(env)
RLBase.is_terminated(env::MazeDirected) = get_done(env)

function RLBase.reset!(env::MazeDirected{T}) where {T}
    world = get_world(env)
    rng = get_rng(env)
    height = get_height(env)
    width = get_width(env)

    vertical_range = 2:2:height-1
    horizontal_range = 2:2:width-1

    world[AGENT, get_agent_pos(env)] = false
    world[GOAL, get_goal_pos(env)] = false

    world[WALL, :, :] .= true
    for i in vertical_range, j in horizontal_range
        world[WALL, i, j] = false
    end

    generate_maze!(env)

    new_goal_pos = rand(rng, pos -> !any(@view world[:, pos]), env)
    set_goal_pos!(env, new_goal_pos)
    world[GOAL, new_goal_pos] = true

    agent_start_pos = rand(rng, pos -> !any(@view world[:, pos]), env)
    set_agent_pos!(env, agent_start_pos)
    world[AGENT, agent_start_pos] = true

    set_agent_dir!(env, rand(rng, DIRECTIONS))

    set_reward!(env, zero(T))
    set_done!(env, false)

    return nothing
end

function (env::MazeDirected{T})(action::AbstractTurnAction) where {T}
    new_dir = turn(action, get_agent_dir(env))
    set_agent_dir!(env, new_dir)
    world = get_world(env)

    if world[GOAL, get_agent_pos(env)]
        set_done!(env, true)
        set_reward!(env, get_terminal_reward(env))
    else
        set_done!(env, false)
        set_reward!(env, zero(T))
    end

    return nothing
end

function (env::MazeDirected{T})(action::AbstractMoveAction) where {T}
    world = get_world(env)

    agent_pos = get_agent_pos(env)
    dest = move(action, get_agent_dir(env), agent_pos)

    if !world[WALL, dest]
        world[AGENT, agent_pos] = false
        set_agent_pos!(env, dest)
        world[AGENT, dest] = true
    end

    if world[GOAL, get_agent_pos(env)]
        set_done!(env, true)
        set_reward!(env, get_terminal_reward(env))
    else
        set_done!(env, false)
        set_reward!(env, zero(T))
    end

    return nothing
end

#####
# Undirected
#####

"""
Maze generation uses the iterative implementation of randomized depth-first search from [Wikipedia](https://en.wikipedia.org/wiki/Maze_generation_algorithm#Iterative_implementation).
"""
function MazeUndirected(; T = Float32, height = 9, width = 9, rng = Random.GLOBAL_RNG)
    @assert isodd(height) && isodd(width) "height and width must be odd numbers"
    vertical_range = 2:2:height-1
    horizontal_range = 2:2:width-1

    objects = (AGENT, WALL, GOAL)
    world = GridWorldBase(objects, height, width)

    world[WALL, :, :] .= true
    for i in vertical_range, j in horizontal_range
        world[WALL, i, j] = false
    end

    goal_pos = CartesianIndex(height - 1, width - 1)
    world[GOAL, goal_pos] = true

    agent_pos = CartesianIndex(2, 2)
    world[AGENT, agent_pos] = true

    reward = zero(T)
    terminal_reward = one(T)
    done = false

    env = MazeUndirected(world, agent_pos, reward, rng, terminal_reward, goal_pos, done)

    RLBase.reset!(env)

    return env
end

RLBase.state_space(env::MazeUndirected, ::RLBase.Observation, ::RLBase.DefaultPlayer) = nothing
const MAZE_UNDIRECTED_LAYERS = SA.SVector(2, 3)
RLBase.state(env::MazeUndirected, ::RLBase.Observation, ::RLBase.DefaultPlayer) = get_grid(get_world(env), get_agent_pos(env), get_half_size(env), MAZE_UNDIRECTED_LAYERS)

RLBase.state_space(env::MazeUndirected, ::RLBase.InternalState, ::RLBase.DefaultPlayer) = nothing
RLBase.state(env::MazeUndirected, ::RLBase.InternalState, ::RLBase.DefaultPlayer) = copy(get_grid(env))

RLBase.action_space(env::MazeUndirected, player::RLBase.DefaultPlayer) = UNDIRECTED_NAVIGATION_ACTIONS
RLBase.reward(env::MazeUndirected, ::RLBase.DefaultPlayer) = get_reward(env)
RLBase.is_terminated(env::MazeUndirected) = get_done(env)

function RLBase.reset!(env::MazeUndirected{T}) where {T}
    world = get_world(env)
    rng = get_rng(env)
    height = get_height(env)
    width = get_width(env)

    vertical_range = 2:2:height-1
    horizontal_range = 2:2:width-1

    world[AGENT, get_agent_pos(env)] = false
    world[GOAL, get_goal_pos(env)] = false

    world[WALL, :, :] .= true
    for i in vertical_range, j in horizontal_range
        world[WALL, i, j] = false
    end

    generate_maze!(env)

    new_goal_pos = rand(rng, pos -> !any(@view world[:, pos]), env)
    set_goal_pos!(env, new_goal_pos)
    world[GOAL, new_goal_pos] = true

    agent_start_pos = rand(rng, pos -> !any(@view world[:, pos]), env)
    set_agent_pos!(env, agent_start_pos)
    world[AGENT, agent_start_pos] = true

    set_reward!(env, zero(T))
    set_done!(env, false)

    return nothing
end

function (env::MazeUndirected{T})(action::AbstractMoveAction) where {T}
    world = get_world(env)

    agent_pos = get_agent_pos(env)
    dest = move(action, agent_pos)

    if !world[WALL, dest]
        world[AGENT, agent_pos] = false
        set_agent_pos!(env, dest)
        world[AGENT, dest] = true
    end

    if world[GOAL, get_agent_pos(env)]
        set_done!(env, true)
        set_reward!(env, get_terminal_reward(env))
    else
        set_done!(env, false)
        set_reward!(env, zero(T))
    end

    return nothing
end

#####
# Common
#####

function get_candidate_neighbors(pos::CartesianIndex{2})
    shifts = ((-2, 0), (0, -2), (0, 2), (2, 0))
    return map(shift -> CartesianIndex(pos.I .+ shift), shifts)
end

function generate_maze!(env::Union{MazeDirected, MazeUndirected})
    world = get_world(env)
    rng = get_rng(env)
    height = get_height(env)
    width = get_width(env)

    vertical_range = 2:2:height-1
    horizontal_range = 2:2:width-1

    visited = falses(height, width)
    stack = DS.Stack{CartesianIndex{2}}()

    first_cell = CartesianIndex(rand(rng, vertical_range), rand(rng, horizontal_range))
    visited[first_cell] = true
    push!(stack, first_cell)

    while !isempty(stack)
        current_cell = pop!(stack)
        candidate_neighbors = get_candidate_neighbors(current_cell)
        neighbors = filter(pos -> (pos.I[1] in vertical_range) && (pos.I[2] in horizontal_range), candidate_neighbors)
        unvisited_neighbors = filter(pos -> !visited[pos], neighbors)
        if length(unvisited_neighbors) > 0
            push!(stack, current_cell)

            next_cell = rand(rng, unvisited_neighbors)

            mid_pos = CartesianIndex((current_cell.I .+ next_cell.I) .÷ 2)
            world[WALL, mid_pos] = false

            visited[next_cell] = true
            push!(stack, next_cell)
        end
    end

    return nothing
end


### A Pluto.jl notebook ###
# v0.12.10

using Markdown
using InteractiveUtils

# ╔═╡ cbdf4e42-2985-11eb-2041-dbd79e434686
using Plots, OTSIN, StatsBase, NamedArrays

# ╔═╡ d94c4784-29bb-11eb-1050-517f3dd80fa0
md"""
Illustration of the optimal transportation on a seed dispersal dataset `M_SD_004`.

```

Spindalis portoricensis
Euphonia musica
Nesospingus sp1 M_SD_004
Loxigilla portoricensis
Vireo altiloquous
Melanerpes portoricensis
Columba squamosa
Margarops fuscatus
Tyrannus dominicensis
Zenaida asiatica
Coereba flaveola
Dendroica tigrina
Vireo flavirostris
Dendroica caerulescens
Icterus dominicensis
Myarchus antillarum
Tiaris bicolor
Todus mexicanus
Tyrannus caudifasciatus
Vireo latimeri
Phoradendron sp1 M_SD_004	10	55	3	9	1	0	0	0	0	0	1	0	1	0	0	0	0	0	0	0
Schefflera morototoni	56	1	4	1	3	0	5	1	4	0	1	1	0	0	0	0	0	0	1	0
Cecropia schreberiana	29	0	6	6	4	3	0	1	0	0	2	3	0	1	0	0	0	0	0	0
Ficus sp1 M_SD_004	34	0	4	0	1	0	0	1	0	0	0	0	0	0	0	0	0	0	0	0
Guarea guidonia	0	0	5	17	14	0	0	1	0	0	0	0	1	0	0	0	0	0	0	0
Cordia sulcata	28	0	0	0	0	0	0	0	0	0	0	0	0	0	0	0	0	0	0	0
Alchornea latifolia	3	0	0	14	3	0	0	0	0	5	0	0	0	0	0	0	0	1	0	1
Miconia serrulata	6	2	12	0	0	0	0	0	0	0	0	0	0	0	0	0	0	0	0	0
Clusia rosea	0	0	1	0	11	3	0	0	1	0	0	0	0	0	1	1	0	0	0	0
Anthurium scandens	1	16	0	0	0	0	0	0	0	0	0	0	0	0	0	0	0	0	0	0
Miconia affinis	2	0	10	0	1	0	0	0	0	0	0	0	0	0	0	0	0	0	0	0
Inga vera	0	0	8	0	1	0	0	0	0	0	0	0	0	0	0	0	0	0	0	0
Musa acuminata	1	0	3	4	0	0	0	0	0	0	0	0	0	0	0	0	0	0	0	0
Inga laurina	4	0	0	0	1	0	0	0	0	0	0	0	0	0	0	0	0	0	0	0
Dendropemon bicolor	0	0	3	1	0	0	0	0	0	0	0	0	0	0	0	0	0	0	0	0
Myrcine coriacea	0	0	2	1	1	0	0	0	0	0	0	0	0	0	0	0	0	0	0	0
Syzigium jambos	0	0	1	3	0	0	0	0	0	0	0	0	0	0	0	0	0	0	0	0
Andira inermis	1	0	0	0	2	0	0	0	0	0	0	0	0	0	0	0	0	0	0	0
Coffea arabica	0	0	0	1	2	0	0	0	0	0	0	0	0	0	0	0	0	0	0	0
Miconia racemosa	0	0	2	1	0	0	0	0	0	0	0	0	0	0	0	0	0	0	0	0
Palicourea guianensis	0	0	1	2	0	0	0	0	0	0	0	0	0	0	0	0	0	0	0	0
Tilandsia sp1 M_SD_004	0	0	3	0	0	0	0	0	0	0	0	0	0	0	0	0	0	0	0	0
Buchenavia capitata	0	0	0	1	1	0	0	0	0	0	0	0	0	0	0	0	0	0	0	0
Casearia arborea	1	0	1	0	0	0	0	0	0	0	0	0	0	0	0	0	0	0	0	0
Citrus sinensis	1	0	1	0	0	0	0	0	0	0	0	0	0	0	0	0	0	0	0	0
Dendropanax arboreus	1	0	0	0	0	0	0	0	0	0	0	0	0	0	0	0	0	0	0	0
Henriettea fascicularis	0	1	0	0	0	0	0	0	0	0	0	0	0	0	0	0	0	0	0	0
Momordica charantia	0	1	0	0	0	0	0	0	0	0	0	0	0	0	0	0	0	0	0	0
Palicourea crocea	0	0	0	1	0	0	0	0	0	0	0	0	0	0	0	0	0	0	0	0
Philodendron angustatum	0	0	1	0	0	0	0	0	0	0	0	0	0	0	0	0	0	0	0	0
Piper sp1 M_SD_004	0	0	0	0	0	0	0	0	0	0	0	0	0	0	0	0	1	0	0	0
Solanum rugosum	0	0	0	0	0	0	0	1	0	0	0	0	0	0	0	0	0	0	0	0
Unidentified sp1 M_SD_004	0	0	1	0	0	0	0	0	0	0	0	0	0	0	0	0	0	0	0	0
Zanthoxylum martinicensis	0	0	0	1	0	0	0	0	0	0	0	0	0	0	0	0	0	0	0	0
```

"""

# ╔═╡ 8dae9448-29bc-11eb-18b0-fb55e9fb1be3
plant_species = ["Phoradendron sp1 M_SD_004",
	"Schefflera morototoni",
	"Cecropia schreberiana",
	"Ficus sp1 M_SD_004",
	"Guarea guidonia",
	"Cordia sulcata",
	"Alchornea latifolia",
	"Miconia serrulata",
	"Clusia rosea",
	"Anthurium scandens",
	"Miconia affinis",
"Inga vera",
"Musa acuminata",
"Inga laurina",
"Dendropemon bicolor",
"Myrcine coriacea",
"Syzigium jambos",
"Andira inermis",
"Coffea arabica",
"Miconia racemosa",
"Palicourea guianensis",
"Tilandsia sp1 M_SD_004",
"Buchenavia capitata",
"Casearia arborea",
"Citrus sinensis",
"Dendropanax arboreus",
"Henriettea fascicularis",
"Momordica charantia",
"Palicourea crocea",
"Philodendron angustatum",
"Piper sp1 M_SD_004",
"Solanum rugosum",
"Unidentified sp1 M_SD_004",
"Zanthoxylum martinicensis"]

# ╔═╡ 67d66742-2b32-11eb-1a58-370b8ca72b53
length(plant_species)

# ╔═╡ 37e622b0-29bc-11eb-2ede-a9430c312543
bird_species = ["Spindalis portoricensis",
	"Euphonia musica",
	"Nesospingus sp1 M_SD_004",
	"Loxigilla portoricensis",
	"Vireo altiloquous",
	"Melanerpes portoricensis",
	"Columba squamosa",
	"Margarops fuscatus",
	"Tyrannus dominicensis",
	"Zenaida asiatica",
	"Coereba flaveola",
	"Dendroica tigrina",
	"Vireo flavirostris",
	"Dendroica caerulescens",
	"Icterus dominicensis",
	"Myarchus antillarum",
	"Tiaris bicolor",
	"Todus mexicanus",
	"Tyrannus caudifasciatus",
	"Vireo latimeri"]

# ╔═╡ 60219008-2b32-11eb-3279-37c1b9489162
length(bird_species)

# ╔═╡ ed1c921c-2986-11eb-28ad-e1f7d3489193
Y = [10 55 3 9 1 0 0 0 0 0 1 0 1 0 0 0 0 0 0 0;
56 1 4 1 3 0 5 1 4 0 1 1 0 0 0 0 0 0 1 0;
29 0 6 6 4 3 0 1 0 0 2 3 0 1 0 0 0 0 0 0;
34 0 4 0 1 0 0 1 0 0 0 0 0 0 0 0 0 0 0 0;
0 0 5 17 14 0 0 1 0 0 0 0 1 0 0 0 0 0 0 0;
28 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0;
3 0 0 14 3 0 0 0 0 5 0 0 0 0 0 0 0 1 0 1;
6 2 12 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0;
0 0 1 0 11 3 0 0 1 0 0 0 0 0 1 1 0 0 0 0;
1 16 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0;
2 0 10 0 1 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0;
0 0 8 0 1 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0;
1 0 3 4 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0;
4 0 0 0 1 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0;
0 0 3 1 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0;
0 0 2 1 1 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0;
0 0 1 3 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0;
1 0 0 0 2 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0;
0 0 0 1 2 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0;
0 0 2 1 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0;
0 0 1 2 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0;
0 0 3 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0;
0 0 0 1 1 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0;
1 0 1 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0;
1 0 1 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0;
1 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0;
0 1 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0;
0 1 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0;
0 0 0 1 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0;
0 0 1 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0;
0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 1 0 0 0;
0 0 0 0 0 0 0 1 0 0 0 0 0 0 0 0 0 0 0 0;
0 0 1 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0;
0 0 0 1 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0];

# ╔═╡ 06c42ff4-29be-11eb-23d3-bffa5ba9a554
n_interactions = sum(Y)

# ╔═╡ 08688ab4-298a-11eb-0ee0-0f00e8ee6248
n, m = size(Y)

# ╔═╡ 3ec1637c-2987-11eb-27c3-7d291ca7306e
a, b = marginals(Y)

# ╔═╡ 4859100e-2987-11eb-3f77-b5825e36a555
P = normalize(Y)

# ╔═╡ 398ed8c4-29bb-11eb-3089-afafbc8345af
plot_Y = marginalsheatmap(Y, cmap=cgrad([:white, :blue]), showaxis=true, xticks=(5:5:m), yticks=5:5:n, xlabel="bird species", ylabel="plant species")

# ╔═╡ 50b53a72-2987-11eb-27b5-a53dc0b26e5d
M = fitM(P, fix_a=true, fix_b=true, γ=1e-4)

# ╔═╡ 88a8f162-2987-11eb-2ed8-d5689f6e2895
plot_M = heatmap(M, color=cgrad([:blue,:white,:red]), xlabel="bird species", ylabel="plant species", title="fitted M")

# ╔═╡ 66caee3c-2988-11eb-188a-ef54b7e65c39
Q = optimaltransport(M, a, b)

# ╔═╡ c72c0ccc-29bb-11eb-1819-f361cf9caa6f
Qpu = optimaltransport(M, ones(n)/n, b)

# ╔═╡ 8de8b87c-29be-11eb-3022-83ee7e3f562b
Qbo = optimaltransport(M, a)

# ╔═╡ 368d73e2-2989-11eb-3dc5-2f89432bff64
plot_Qpu = marginalsheatmap(Qpu, cmap=cgrad([:white, :purple]), showaxis=true, xticks=(5:5:m), yticks=5:5:n, xlabel="bird species", ylabel="plant species")

# ╔═╡ 98e014b6-29be-11eb-0713-1de8527dc3d2
plot_Qbo = marginalsheatmap(Qbo, cmap=cgrad([:white, :purple]), showaxis=true, xticks=(5:5:m), yticks=5:5:n, xlabel="bird species", ylabel="plant species")

# ╔═╡ 265d6474-29bf-11eb-2e20-a1180d19e9e8
p_total = plot(plot_Y, plot_M, plot_Qpu, plot_Qbo, size=(1200, 1000))

# ╔═╡ 7597e448-29c3-11eb-2262-951017f11586
savefig(p_total, "../figures/SD_sim.svg")

# ╔═╡ 7762f2c6-2988-11eb-33b1-a3817c1041b7
mean(abs.(Y .- Q))

# ╔═╡ 26ca4e9e-2989-11eb-06ae-fb1ec5ab97a8
kldivergence(P, normalize(Q))

# ╔═╡ Cell order:
# ╟─d94c4784-29bb-11eb-1050-517f3dd80fa0
# ╠═8dae9448-29bc-11eb-18b0-fb55e9fb1be3
# ╠═67d66742-2b32-11eb-1a58-370b8ca72b53
# ╠═37e622b0-29bc-11eb-2ede-a9430c312543
# ╠═60219008-2b32-11eb-3279-37c1b9489162
# ╠═cbdf4e42-2985-11eb-2041-dbd79e434686
# ╠═ed1c921c-2986-11eb-28ad-e1f7d3489193
# ╠═06c42ff4-29be-11eb-23d3-bffa5ba9a554
# ╠═08688ab4-298a-11eb-0ee0-0f00e8ee6248
# ╠═3ec1637c-2987-11eb-27c3-7d291ca7306e
# ╠═4859100e-2987-11eb-3f77-b5825e36a555
# ╠═398ed8c4-29bb-11eb-3089-afafbc8345af
# ╠═50b53a72-2987-11eb-27b5-a53dc0b26e5d
# ╠═88a8f162-2987-11eb-2ed8-d5689f6e2895
# ╠═66caee3c-2988-11eb-188a-ef54b7e65c39
# ╠═c72c0ccc-29bb-11eb-1819-f361cf9caa6f
# ╠═8de8b87c-29be-11eb-3022-83ee7e3f562b
# ╠═368d73e2-2989-11eb-3dc5-2f89432bff64
# ╠═98e014b6-29be-11eb-0713-1de8527dc3d2
# ╠═265d6474-29bf-11eb-2e20-a1180d19e9e8
# ╠═7597e448-29c3-11eb-2262-951017f11586
# ╠═7762f2c6-2988-11eb-33b1-a3817c1041b7
# ╠═26ca4e9e-2989-11eb-06ae-fb1ec5ab97a8


module GraphJuMP

using JuMP
using LightGraphs
using MetaGraphs

export setgraph!, removegraph!, addpath!

const extKey = :GraphJuMP

"""
Several formulations of the flows are possible:

* `EdgeFlow`: each edge of the graph is explicitly modelled. This model tends to scale poorly to large number of nodes
  and edges in the graph
* `Path`: only some paths in the graph are modelled. This formulation is typical for column-generation schemes
  (set `delayed=true` for this)

To create a new formulation, subtype `GraphFormulation` and `GraphFormulationExt`; implement the required functions:

* `supportspath(::GraphFormulation)`: whether this formulation uses the notion of path. This implies that functions like
  `addpath!` must be supported.
* `buildformulation(m::Model, gf::GraphFormulation)`: builds the formulation, i.e. creates the required JuMP variables and
  constraints for the formulation.
"""
abstract type GraphFormulation end
struct EdgeFlow <: GraphFormulation end
struct Path <: GraphFormulation end

supportspath(::GraphFormulation) = false
supportspath(::Path) = true

"""
Each formulation should have its subtype of `GraphFormulationExt` to store its internal parameters.
"""
abstract type GraphFormulationExt end

mutable struct EdgeFlowFormulationExt <: GraphFormulationExt
  flows::Matrix{Variable} # First index: edge; second index: commodity.
  id_to_edge::Dict{Int, Edge} # Mapping an integer index (used in the model) to a graph edge.
  edge_to_id::Dict{Edge, Int} # Reverse mapping: from a graph edge to an integer index.
end

"""
Stores the internal state of the graph model.
"""
mutable struct GraphModelExt{T <: Int, U <: Real}
  graph::MetaDiGraph{T, U} # The graph below this model. Must be a "meta" graph, as many properties might be useful for each edge or vertex (such as max capacity).
  formulation::GraphFormulation # A graph formulation, see GraphFormulation enumeration.
  delayed::Bool # For a path formulation, controls when paths are computed: false, before solving; true, after solving.
  flowcategory::Symbol # The JuMP category for flow variables. For now, these are: :Cont, :Int, :Bin, :SemiCont, :SemiInt, :SDP
  commodities::Vector{String} # Whether several commodities are present.
  formulationext::Union{GraphFormulationExt, Void} # The formulation object, once it is built (after calling `buildformulation`).

  function GraphModelExt{T, U}(graph::MetaDiGraph{T, U}, formulation::GraphFormulation=EdgeFlow(),
                               delayed::Bool=false, flowcategory::Symbol=:Cont, commodities::Vector{String}=String["default"]
                              ) where{T, U}
    return new(graph, formulation, delayed, flowcategory, commodities, nothing)
  end
end

# Link between GraphModelExt and JuMP's model.
hasgraph(m::Model) = haskey(m.ext, extKey)
checkgraph(m::Model) = hasgraph(m) || error("The model has no associated graph model.")
getgraphext(m::Model) = hasgraph(m) && m.ext[extKey]
getgraph(m::Model) = getgraphext(m).graph

hasformulation(m::Model) = getgraphext(m).formulationext !== nothing
checkformulation(m::Model) = hasformulation(m) || error("The model has no associated graph formulation. " *
  "Did you forget to call `buildformulation` or to link the graph model to a JuMP model with `setgraph!`?")

function setgraph!(m::Model, graph::GraphModelExt)
  if hasgraph(m)
    warn("Model already has an associated graph model! It will be replaced.")
  end
  m.ext[extKey] = graph
  m.ext[extKey].formulationext = buildformulation(m)
end

setgraph!{T, U}(m::Model, graph::MetaDiGraph{T, U}, args...) = setgraph!(m, GraphModelExt{T, U}(graph, args...))
setgraph!{T}(m::Model, graph::AbstractMetaGraph{T}, args...) = error("GraphJuMP only uses graphs with metadata from the MetaGraphs.jl package. " *
  "You can convert your `graph` to a MetaGraph using `MetaGraph(graph)`. ")

function removegraph!(m::Model)
  if hasgraph(m)
    delete!(m.ext, extKey)
  end
  # No error if there is no graph.
end

# Build the formulations.
function buildformulation(m::Model) # Makes a dispatch on the formulation.
  checkgraph(m)
  return buildformulation(m, getgraphext(m).formulation)
end

function buildformulation(m::Model, gf::EdgeFlow)
  # Create the flow variables, one per edge and per commodity.
  n_edges = ne(getgraph(m))
  n_commodities = length(getgraphext(m).commodities)
  flows = @variable(m, [e=1:n_edges, c=1:n_commodities], category=getgraphext(m).flowcategory)

  # Map edge and their IDs in the model.
  id_to_edge = Dict([idx => edge for (idx, edge) in enumerate(edges(getgraph(m)))])
  edge_to_id = map(reverse, id_to_edge)

  # Give names to the flow variables.
  for e in 1:n_edges
    for c in 1:n_commodities
      edge = id_to_edge[e]
      setname(flows[e, c], "flow_" * string(edge.src) * "_to_" * string(edge.dst) * "_commodity_" * string(getgraphext(m).commodities[c]))
    end
  end

  # Done!
  return EdgeFlowFormulationExt(flows, id_to_edge, edge_to_id)
end

# For path formulations.
function addpath!(m::Model, path)
  checkgraph(m)

  if ! supportspath(getgraphext(m).formulation)
    error("addpath! can only be called on graph models using a path-based formulation, such as `Path`.")
  end

  # TODO!
end

end


module Mortchartgen

using CSV, DataFrames, DataStructures, JSON, Loess, Mustache, ODBC, Statistics, PyCall, Missings

const bo = PyNULL()
const bp = PyNULL()
const bml  = PyNULL()
const bpal = PyNULL()

function __init__()
	copy!(bo, pyimport("bokeh"))
	copy!(bp, pyimport("bokeh.plotting"))
	copy!(bml, pyimport("bokeh.models"))
	copy!(bpal, pyimport("bokeh.palettes"))
end

mainpath = normpath(@__DIR__, "..")
datapath = normpath(mainpath, "data")
chartpath = normpath(mainpath, "charts")
tmpoutpath = normpath(tempdir(), "mout.html")
mkpath(chartpath)
conf = JSON.parsefile(normpath(datapath, "chartgen.json"),
	dicttype=DataStructures.OrderedDict)
tables = Dict(:deaths => "Deaths", :pop => "Pop")
wpp = CSV.read(normpath(datapath, "wppconverted.csv"))
dfarrmatch(col, arr) = map((x) -> in(x, arr), Vector(col))
ctrycodes = map((x)->parse(Int, x), collect(keys(conf["countries"])))
perc_round(value) = replace("$(round(value, 4))", ".", ",")
dthalias(language) = uppercasefirst(conf["deaths"]["alias"][language])
obend = "svg"

function caalias(cause, language)
	if cause=="pop"
		return conf["pop"]["alias"][language]
	else
		return conf["causes"][cause]["alias"][language]
	end
end

function cgen_frames(causes = keys(conf["causes"]))
	conn_config = conf["settings"]["conn_config"]
	#conn = MySQL.connect(conn_config["host"], conn_config["user"],
	#	conn_config["password"], db = conn_config["database"],
	#	unix_socket = conn_config["unix_socket"])
	conn = ODBC.DSN(conn_config["dsn"], conn_config["user"], conn_config["password"])
	qstr = *("select Sex,Year,Country,Admin1,",
		join(map((x)->"$(tables[:pop])$x,", 1:25)),
		"$(tables[:pop])26 from $(tables[:pop])")
	popframe = ODBC.query(conn, qstr)
	rename!(popframe,
		Dict(zip(map((x)->Symbol("$(tables[:pop])$(x)"), 1:26),
			 map((x)->Symbol("Age$(x)"), 1:26))))
	dthframes = DataStructures.OrderedDict()
	for cause in causes
		causeexpr = conf["causes"][cause]["causeexpr"]
		selstat = *("Sex,Year,List,Country,Admin1,",
			join(map((x)->"$(tables[:deaths])$x,", 1:25)),
			"$(tables[:deaths])26")
		#Prepared statements give incorrect results.
		qstr = *("select ", selstat,
			" from $(tables[:deaths]) where ",
			"(case when List='07A' then Cause REGEXP '$(causeexpr["07A"])' ",
			"when List='08A' then Cause REGEXP '$(causeexpr["08A"])' ",
			"when List REGEXP '09(B|N)' then Cause REGEXP '$(causeexpr["09B"])' ",
			"when List REGEXP '10(M|[3-4])' then Cause REGEXP '$(causeexpr["10"])' ",
			"when List='101' then Cause REGEXP '$(causeexpr["101"])' end) ",
			"order by Year")
		dthframe = aggregate(ODBC.query(conn, qstr),
				[:Sex, :Year, :List, :Country, :Admin1], sum)
		rename!(dthframe,
			Dict(zip(map((x)->Symbol("$(tables[:deaths])$(x)_sum"), 1:26),
				 map((x)->Symbol("Age$(x)"), 1:26))))
		dthframes[cause] = hcat(DataFrame(Cause = fill(cause, size(dthframe)[1])),
			dthframe)
	end
	ODBC.disconnect!(conn)
	frames = Dict(:deaths => vcat(values(dthframes)...), :pop => popframe)
	for tblkey in keys(tables)
		frame = frames[tblkey]
		for r in eachrow(frame)
			if r[:Country]==4100 && r[:Year]<1990
				r[:Country] = 4085
			elseif r[:Country]==3150 && r[:Year]>1974 && !ismissing(r[:Admin1])
				r[:Country] = 0
			end
		end
		delete!(frame, :Admin1)
		ncols = size(frame)[2]
		frame_trimmed = frame[dfarrmatch(frame[:Country], ctrycodes), :]
		frame_trimmed_long = stack(frame_trimmed, ncols-25:ncols)
		frames[tblkey] = frame_trimmed_long
	end
	for ctrycode in ctrycodes
		ctry = conf["countries"]["$ctrycode"]
		if haskey(ctry, "wppstart")
			wppctry = wpp[((wpp[:Country].==ctrycode).&
				(wpp[:Year].>=ctry["wppstart"])), :]
			frames[:pop] = vcat(frames[:pop], wppctry)
		end
	end
	frames
end

function save_frames(framedict)
	for tblkey in keys(tables)
		writetable(normpath(datapath, *(string(tblkey), ".csv")),
			framedict[tblkey])
	end
end

function load_frames()
	framedict = Dict()
	for tblkey in keys(tables)
		framedict[tblkey] = CSV.read(normpath(datapath, *(string(tblkey), ".csv")))
		framedict[tblkey][:variable] = map(Symbol, framedict[tblkey][:variable])
	end
	framedict
end

subframe_sray(df, sex, countries, agelist, years) = df[((df[:Sex].==sex)
	.& (dfarrmatch(df[:Country], countries)) .& (dfarrmatch(df[:variable], agelist))
	.& (dfarrmatch(df[:Year], years))), :]
dfgrp_agemean(df, grpcol, f = mean) = by(df, grpcol, x -> DataFrame(value = f(x[:value])))
dfgrp_sum(df, grpcol, f = sum) = by(df, grpcol,
x -> DataFrame(value = f(x[:value]), value_1 = f(x[:value_1])))

function listchanges(country, years, framedict)
	countryframe = subframe_sray(framedict[:deaths], 2, country, [:Age1], years)
	listframe = countryframe[countryframe[:Cause].=="all", [:Year; :List]]
	nrows = size(listframe)[1]
	lcomp = DataFrame(Year = listframe[:Year][2:nrows],
		List = listframe[:List][2:nrows], Listprev = listframe[:List][1:nrows-1])
	lcomp[lcomp[:List].!=lcomp[:Listprev], :]
end

function grpprop(numframe_sub, denomframe_sub, grpcol, agemean)
	numdenomframe_sub =
		join(numframe_sub, denomframe_sub, on = grpcol, makeunique = true)
	if agemean
		propfr_agesp = DataFrame()
		propfr_agesp[grpcol] = numdenomframe_sub[grpcol]
		propfr_agesp[:value] = numdenomframe_sub[:value]./numdenomframe_sub[:value_1]
		return dfgrp_agemean(propfr_agesp, grpcol)
	else
		numdenomgrp = dfgrp_sum(numdenomframe_sub, grpcol)
		propfr_agegr = DataFrame()
		propfr_agegr[grpcol] = numdenomgrp[grpcol]
		propfr_agegr[:value] = numdenomgrp[:value]./numdenomgrp[:value_1]
		return propfr_agegr
	end
end

function propgrp(numframe, denomframe, sex, countries, agelist, years, agemean, grpcol)
	numframe_sub = subframe_sray(numframe, sex, countries, agelist, years)
	denomframe_sub = subframe_sray(denomframe, sex, countries, agelist, years)
	grpprop(numframe_sub, denomframe_sub, grpcol, agemean)
end

function ageslice(sage, eage, agemean, language)
	infsymb = "\u03c9"
	ages = map((x)->Symbol("Age$x"), 1:25)
	agest = [0;0:4;5:5:90;95]
	ageend = [infsymb;0:4;9:5:94;infsymb]
	if agemean
		agemeanstr = " $(conf["agemean"]["alias"][language])"
	else
		agemeanstr = ""
	end
	Dict(:agelist => ages[sage:eage],
		:alias => "$(agest[sage])\u2013$(ageend[eage])$agemeanstr",
		:agest => agest[sage], :ageend => ageend[eage],
		:color => bpal.Category20[20][mod1(sage, 20)])
end

function propplotframes(ca1, ca2, framedict, language)
	deaths = framedict[:deaths]
	ca1frame = deaths[deaths[:Cause].==ca1, :]
	ca1alias = caalias(ca1, language)
	ca2alias = caalias(ca2, language)
	if ca2=="pop"
		ca2frame = framedict[:pop]
	else
		ca2frame = deaths[deaths[:Cause].==ca2, :]
	end
	Dict(:ca1frame => ca1frame, :ca1alias => ca1alias,
		:ca2frame => ca2frame, :ca2alias => ca2alias)
end

function listlabels(country, years, minvals, framedict)
	lch = listchanges(country, years, framedict)
	lchdata = bml[:ColumnDataSource](data = Dict("year" => lch[:Year],
		"list" => lch[:List]))
	listlabels = bml[:LabelSet](x = "year", y = minimum(minvals),
		text = "list", text_color = "red", angle = pi/2, render_mode = "canvas",
		source = lchdata)
end

function preplot(outfile)
	bp[:reset_output]()
	bp[:output_file](outfile)
end

function postplot(figure, showplot)
	if showplot
		bp[:show](figure)
	else
		bp[:save](figure)
	end
end

function propplot_sexesyrs(ca1, ca2, sexes, country, sage, eage, years, agemean,
	framedict, language, outfile, showplot, apctype = "pa", acbcoh = 1900)
	ages = ageslice(sage, eage, agemean, language)
	if (apctype == "pa" || apctype == "pc")
		xdict = (sage, eage)
		xcol = :Year
		ylab = ages[:alias]
		ages = [(sage, eage)]
	elseif (apctype == "ap" || apctype == "ac")
		if apctype == "ap"
			xdict = years
			ylab = string(years)
		else
			xdict = acbcoh
			ylab = string(acbcoh)
		end
		xcol = :Age
		ages = collect(map((a)->(a,a), sage:eage))
	end
	meta = Dict(:ca1 => ca1, :ca2 => ca2, :country => country,
		:sage => sage, :eage => eage, :agemean => agemean,
		:apctype => apctype, :years => years, :acbcoh => acbcoh,
		:xcol => xcol, :ylab => ylab, :plclass => "sexesyrs")
	propframes = Dict()
	for sex in sexes
		propframe = propplot_agesyrs(ca1, ca2, sex, country, ages,
			years, agemean, framedict, language, "linear", "",
			false, apctype)[xdict][:frame]
		yrfloatarr = convert(Array{Float64}, propframe[xcol])
		valarr = convert(Array, propframe[:value])
		propsm = Loess.predict(loess(yrfloatarr, valarr), yrfloatarr)
		propframes[sex] = Dict(:propframe => propframe, :propsm => propsm)
	end
	propframes[:meta] = meta
	if outfile != ""
		propframes_render(propframes, language, sexes, outfile, framedict, showplot)
	end
	propframes
end

function propframes_render(propframes, language, propframeds, outfile, framedict, showplot,
			  x_axis_type = "linear", y_axis_type = "linear")
	preplot(outfile)
	meta = propframes[:meta]
	plclass = meta[:plclass]
	ca1 = meta[:ca1]
	ca2 = meta[:ca2]
	years = meta[:years]
	apctype = meta[:apctype]
	ca1alias = caalias(ca1, language)
	ca2alias = caalias(ca2, language)
	legends = []
	minvals = []
	if plclass == "sexesyrs"
		country = meta[:country]
		xcol = meta[:xcol]
		ylab = meta[:ylab]
		ctryalias = conf["countries"][string(country)]["alias"][language]
		figtitle = "$(dthalias(language)) $ca1alias/$ca2alias $ctryalias"
		p = bp[:figure](output_backend = obend, title = figtitle, y_axis_label = ylab,
			toolbar_location = "below", toolbar_sticky = false,
			plot_width = 600, plot_height = 600)
		for sex in propframeds
			sexalias = conf["sexes"][string(sex)]["alias"][language]
			col = conf["sexes"][string(sex)]["color"]
			propframed = propframes[sex]
			propframe = propframed[:propframe]
			propsm = propframed[:propsm]
			sexcirc = p[:circle](propframe[xcol], propframe[:value], color = col)
			sexline = p[:line](propframe[xcol], propsm, color = col)
			legends = vcat(legends,
				[(sexalias, [sexcirc]); ("$sexalias loess", [sexline])]...)
			minvals = vcat(minvals, minimum(propframe[:value]))
		end
		legend = bml[:Legend](items = legends, location = (0, -30))
		p[:add_layout](legend, "right")
		if apctype == "pa"
			p[:add_layout](listlabels(country, years, minvals, framedict))
		end
	elseif plclass == "agesyrs"
		country = meta[:country]
		sex = meta[:sex]
		ctryalias = conf["countries"][string(country)]["alias"][language]
		sexalias = conf["sexes"][string(sex)]["alias"][language]
		figtitle = "$(dthalias(language)) $ca1alias/$ca2alias $sexalias $ctryalias"
		p = bp[:figure](output_backend = obend, title = figtitle, y_axis_type = y_axis_type,
			toolbar_location = "below", toolbar_sticky = false,
			plot_width = 600, plot_height = 600)
		for agetuple in propframeds
			propframed = propframes[agetuple]
			propframe = propframed[:frame]
			ages = propframed[:ages]
			agealias = ages[:alias]
			if (apctype == "pa" || apctype == "pc")
				ageline = p[:line](propframe[:Year], propframe[:value], color = ages[:color])
				legends = vcat(legends, (ages[:alias], [ageline]))
			end
		end
		legend = bml[:Legend](items = legends, location = (0, -30))
		p[:add_layout](legend, "right")
	end
	p[:add_tools](bml[:CrosshairTool]())
	postplot(p, showplot)
end

function sexesrat(sexesyrsplot, numsex, denomsex, xcol)
	xs = sexesyrsplot[numsex][:propframe][xcol]
	values = sexesyrsplot[numsex][:propframe][:value]./sexesyrsplot[denomsex][:propframe][:value]
	df = DataFrame()
	df[xcol] = xs
	df[:value] = values
	df
end

function srataliases(sexesyrsplots, language)
	ca1 = map((s)->caalias(s[:meta][:ca1], language), sexesyrsplots)
	ca2 = map((s)->caalias(s[:meta][:ca2], language), sexesyrsplots)
	country = map((s)->conf["countries"][string(s[:meta][:country])]["alias"][language],
		sexesyrsplots)
	ages = map((s)->ageslice(s[:meta][:sage], s[:meta][:eage],
		s[:meta][:agemean], language)[:alias], sexesyrsplots)
	years = map((s)->string(s[:meta][:years]), sexesyrsplots)
	acbcohs = map((s)->string(s[:meta][:acbcoh]), sexesyrsplots)
	aldict = Dict(:ca1 => ca1, :ca2 => ca2, :country => country, :ages => ages,
		:years => years, :acbcohs => acbcohs)
end

unicount(i) = size(unique(i))[1]

# Takes an array of sexesyrsplots with the same apctype.
function plot_sexesrats(sexesyrsplots, language, outfile, showplot,
	sexesratio = true, numsex = 1, denomsex = 2, y_axis_type = "linear")
	preplot(outfile)
	numsexalias = conf["sexes"][string(numsex)]["alias"][language]
	denomsexalias = conf["sexes"][string(denomsex)]["alias"][language]
	srals = srataliases(sexesyrsplots, language)
	figtitle = "$(dthalias(language))"
	apctype = sexesyrsplots[1][:meta][:apctype]
	if sexesratio
		figtitle = "$figtitle $numsexalias/$denomsexalias"
	else
		figtitle = "$figtitle $numsexalias"
	end
	ylab = ""
	if unicount(srals[:ca1]) == 1
		figtitle = "$figtitle $(srals[:ca1][1])"
	end
	if unicount(srals[:ca2]) == 1
		figtitle = "$figtitle/$(srals[:ca2][1])"
	end
	if unicount(srals[:country]) == 1
		figtitle = "$figtitle, $(srals[:country][1])"
	end
	if (apctype == "pa" || apctype == "pc")
		xcol = :Year
		if unicount(srals[:ages]) == 1
			ylab = srals[:ages][1]
		end
	elseif apctype == "ap"
		xcol = :Age
		if unicount(srals[:years]) == 1
			ylab = srals[:years][1]
		end
	elseif apctype == "ac"
		xcol = :Age
		if unicount(srals[:acbcohs]) == 1
			ylab = srals[:acbcohs][1]
		end
	end
	p = bp[:figure](output_backend = obend, title = figtitle, y_axis_label = ylab,
		toolbar_location = "below", toolbar_sticky = false,
		plot_width = 800, plot_height = 600, y_axis_type = y_axis_type)
	sratlegends = []
	for (i, syplot) in enumerate(sexesyrsplots)
		meta = syplot[:meta]
		col = bpal[:Category20][20][mod1(i, 20)]
		if sexesratio
			sratframe = sexesrat(syplot, numsex, denomsex, xcol)
			yrfloatarr = convert(Array{Float64}, sratframe[xcol])
			valarr = convert(Array, sratframe[:value])
			sratsm = Loess.predict(loess(yrfloatarr, valarr), yrfloatarr)
		else
			sratframe = syplot[numsex][:propframe]
			sratsm = syplot[numsex][:propsm]
		end
		sratcirc = p[:circle](sratframe[xcol], sratframe[:value], color = col)
		sratline = p[:line](sratframe[xcol], sratsm, color = col)
		sratleg = ""
		if unicount(srals[:ca1]) > 1
			sratleg = "$(srals[:ca1][i])"
		end
		if unicount(srals[:ca2]) > 1
			sratleg = "$sratleg/$(srals[:ca2][i])"
		end
		if unicount(srals[:country]) > 1
			sratleg = "$sratleg, $(srals[:country][i])"
		end
		if (apctype == "pa" || apctype == "pc")
			if unicount(srals[:ages]) > 1
				sratleg = "$sratleg, $(srals[:ages][i])"
			end
		elseif apctype == "ap"
			if unicount(srals[:years]) > 1
				sratleg = "$sratleg, $(srals[:years][i])"
			end
		elseif apctype == "ac"
			if unicount(srals[:acbcohs]) > 1
				sratleg = "$sratleg, $(srals[:acbcohs][i])"
			end
		end
		sratlegends = vcat(sratlegends, (sratleg, [sratcirc]))
	end
	legend = bml[:Legend](items = sratlegends, location = (0, -30))
	p[:add_layout](legend, "right")
	p[:add_tools](bml[:CrosshairTool]())
	postplot(p, showplot)
end

function propplot_agesyrs(ca1, ca2, sex, country, agetuples, years, agemean,
	framedict, language, y_axis_type, outfile, showplot, apctype = "pa")
	pframes = propplotframes(ca1, ca2, framedict, language)
	minvals = []
	meta = Dict(:ca1 => ca1, :ca2 => ca2, :sex => sex, :country => country,
		    :apctype => apctype, :years => years,
		    :agemean => agemean, :plclass => "agesyrs")
	propframes = Dict()
	for agetuple in agetuples
		ages = ageslice(agetuple[1], agetuple[2], agemean, language)
		agelist = ages[:agelist]
		propframe = propgrp(pframes[:ca1frame], pframes[:ca2frame],
			sex, country, agelist, years, agemean, :Year)
		if (apctype == "pc" || apctype == "ac")
			propframe[:Year] = propframe[:Year].-ages[:agest]
		end
		minvals = vcat(minvals, minimum(propframe[:value]))
		propframes[agetuple] = Dict(:frame => propframe, :ages => ages)
	end
	if (apctype == "ap" || apctype == "ac")
		framedicts = values(propframes)
		ageframes = map((agdict)->DataFrame(Year = agdict[:frame][:Year],
			value = agdict[:frame][:value],
			Age = agdict[:ages][:agest]),
			framedicts)
		ageframe = vcat(ageframes...)
		propframes = Dict()
		uniyears = sort!(unique(ageframe[:Year]))
		for year in uniyears
			propframe = ageframe[ageframe[:Year].==year, :]
			sort!(propframe, cols=[:Age])
			propframes[year] = Dict(:frame => propframe)
		end
	end
	propframes[:meta] = meta
	if outfile != ""
		propframes_render(propframes, language, agetuples, outfile, framedict, showplot)
	end
	propframes
end

function propscat_yrsctry(ca1, ca2, sex, countries, sage, eage, year1, year2, agemean,
	framedict, language, outfile, showplot)
	preplot(outfile)
	pframes = propplotframes(ca1, ca2, framedict, language)
	ages = ageslice(sage, eage, agemean, language)
	agealias = ages[:alias]
	agelist = ages[:agelist]
	sexalias =  conf["sexes"][string(sex)]["alias"][language]
	figtitle = "$(dthalias(language)) $(pframes[:ca1alias])/$(pframes[:ca2alias])\n$agealias $sexalias"
	meta = Dict(:ca1 => ca1, :ca2 => ca2, :sex => sex, :countries => countries,
		:sage => sage, :eage => eage, :agemean => agemean)
	yr1propframe = propgrp(pframes[:ca1frame], pframes[:ca2frame], sex,
		countries, agelist, year1, agemean, :Country)
	yr2propframe = propgrp(pframes[:ca1frame], pframes[:ca2frame], sex,
		countries, agelist, year2, agemean, :Country)
	propframe = join(yr1propframe, yr2propframe, on=:Country)
	isos = map((c)->conf["countries"][string(c)]["iso3166"], propframe[:Country])
	ctrynames = map((c)->conf["countries"][string(c)]["alias"][language], propframe[:Country])
	scatdata = bml[:ColumnDataSource](data = Dict("year1prop" => propframe[:value],
		"year2prop" => propframe[:value_1], "isos" => isos, "ctrynames" => ctrynames))
	hover = bml[:HoverTool](tooltips =
			[("befolkning", "@ctrynames"),
			("$year1", "@year1prop"),
			("$year2", "@year2prop")])
	p = bp[:figure](output_backend = obend, title = figtitle, x_axis_label = "$year1",
		y_axis_label = "$year2", plot_width = 600, plot_height = 600)
	p[:add_tools](hover)
	p[:circle](x = "year1prop", y = "year2prop", size = 12, source = scatdata)
	isolabels = bml[:LabelSet](x = "year1prop", y = "year2prop", text = "isos",
		level = "glyph", x_offset = 5, y_offset = 5, source = scatdata)
	p[:add_layout](isolabels)
	postplot(p, showplot)
	Dict(:meta => meta, year1 => yr1propframe, year2 => yr2propframe)
end

function propscat_sexesctry(ca1, ca2, countries, sage, eage, year, agemean,
	framedict, language, outfile, showplot)
	pframes = propplotframes(ca1, ca2, framedict, language)
	preplot(outfile)
	ages = ageslice(sage, eage, agemean, language)
	agealias = ages[:alias]
	agelist = ages[:agelist]
	figtitle = "$(dthalias(language)) $(pframes[:ca1alias])/$(pframes[:ca2alias])\n$agealias $year"
	meta = Dict(:ca1 => ca1, :ca2 => ca2, :countries => countries,
		:sage => sage, :eage => eage, :agemean => agemean)
	fempropframe = propgrp(pframes[:ca1frame], pframes[:ca2frame], 2,
		countries, agelist, year, agemean, :Country)
	malepropframe = propgrp(pframes[:ca1frame], pframes[:ca2frame], 1,
		countries, agelist, year, agemean, :Country)
	propframe = join(fempropframe, malepropframe, on=:Country)
	femalias =  conf["sexes"]["2"]["alias"][language]
	malealias = conf["sexes"]["1"]["alias"][language]
	isos = map((c)->conf["countries"][string(c)]["iso3166"], propframe[:Country])
	ctrynames = map((c)->conf["countries"][string(c)]["alias"][language], propframe[:Country])
	scatdata = bml[:ColumnDataSource](data = Dict("femprop" => propframe[:value],
		"maleprop" => propframe[:value_1], "isos" => isos, "ctrynames" => ctrynames))
	hover = bml[:HoverTool](tooltips =
			[("befolkning", "@ctrynames"),
			("$femalias", "@femprop"),
			("$malealias", "@maleprop")])
	p = bp[:figure](output_backend = obend, title = figtitle, x_axis_label = femalias,
		y_axis_label = malealias, plot_width = 600, plot_height = 600)
	p[:add_tools](hover)
	p[:circle](x = "femprop", y = "maleprop", size = 12, source = scatdata)
	isolabels = bml[:LabelSet](x = "femprop", y = "maleprop", text = "isos",
		level = "glyph", x_offset = 5, y_offset = 5, source = scatdata)
	p[:add_layout](isolabels)
	postplot(p, showplot)
	Dict(:meta => meta, 2 => fempropframe, 1 => malepropframe)
end

batchages_caflt(cause) = filter((age)->
	age["ca2"]!=cause &&
	((age["startage"]==1 && age["endage"]==1) ||
	(!(haskey(conf["causes"][cause], "lowerage")) ||
	age["endage"]>=conf["causes"][cause]["lowerage"]) &&
	(!(haskey(conf["causes"][cause], "upperage")) ||
	age["startage"]<=conf["causes"][cause]["upperage"])),
	conf["batchages"])

fname_sexesyrs(ca1, ca2, country, sage, eage, agemean) =
*(ca1, ca2, country, "s", string(sage), "e", string(eage),
"mean", string(agemean), ".html")

agedict_bothsexes(sage, eage, agemean, language, ca2, fname) =
Dict("alias" => *(ageslice(sage, eage, agemean, language)[:alias],
	"/$(caalias(ca2, language))"),
	"fname" => fname)
agedict_sex(sage, eage, sex, agemean, language, ca2, fname) =
Dict("alias" => *(ageslice(sage, eage, agemean, language)[:alias],
	"/$(caalias(ca2, language)) ",
	conf["sexes"][string(sex)]["alias"][language]),
	"fname" => fname)

cadict(cause, children, language) = Dict("alias" => uppercasefirst(caalias(cause, language)),
"children" => children, "name" => cause)

function ctryints_flt(countries, year1, year2)
	countries_flt = filter((ctry)->conf["countries"][ctry]["startyear"]<=year1
				&& conf["countries"][ctry]["endyear"]>=year2, countries)
	map((c)->parse(Int, c), countries_flt)
end

fname_yrsctry(ca1, ca2, sex, sage, eage, agemean, year1, year2) =
 *(ca1, ca2, string(sex), "s", string(sage), "e", string(eage),
"mean", string(agemean), "ctries", string(year1), "vs", string(year2), ".html")

causes_flt_skipyrs(causes, year1, year2) =
filter((ca)-> !(haskey(conf["causes"][ca], "skipyrs")) ||
(!(year1 in conf["causes"][ca]["skipyrs"]) &&
!(year2 in conf["causes"][ca]["skipyrs"])),
causes)

fname_sexesctry(ca1, ca2, sage, eage, agemean, year) =
*(ca1, ca2, "s", string(sage), "e", string(eage), "mean", string(agemean),
"sexesctries", string(year), ".html")

function agebatchplot(framedict, age, plottype, language, ca1, child, countries, yr1, yr2, sexes)
	ca2 = age["ca2"]
	sage = age["startage"]
	eage = age["endage"]
	agemean = age["agemean"]
	if plottype == "sexesyrs"
		fname = fname_sexesyrs(ca1, ca2, child, sage, eage, agemean)
		outfile = normpath(chartpath, fname)
		propplot_sexesyrs(ca1, ca2, sexes, parse(Int, child),
			sage, eage, yr1:yr2,
			agemean, framedict, language, outfile, false)
		return agedict_bothsexes(sage, eage, agemean, language, ca2, fname)
	elseif plottype == "yrsctry"
		if ca1 in causes_flt_skipyrs(keys(conf["causes"]), yr1, yr2)
			agesexdicts = []
			ctryints = ctryints_flt(countries, yr1, yr2)
			for sex in sexes
				fname = fname_yrsctry(ca1, ca2, sex, sage, eage,
					agemean, yr1, yr2)
				outfile = normpath(chartpath, fname)
				propscat_yrsctry(ca1, ca2, sex, ctryints,
					sage, eage, yr1, yr2,
					agemean, framedict, language, outfile, false)
				agesexdict = agedict_sex(sage, eage, sex,
					agemean, language, ca2, fname)
				agesexdicts = vcat(agesexdicts, agesexdict)
			end
			return agesexdicts
		end
	elseif plottype == "sexesctry"
		if (ca1 in causes_flt_skipyrs(keys(conf["causes"]), yr1, yr2) &&
			conf["causes"][ca1]["sex"]==[2;1])
			ctryints = ctryints_flt(countries, yr1, yr2)
			year = yr1
			fname = fname_sexesctry(ca1, ca2, sage, eage, agemean, year)
			outfile = normpath(chartpath, fname)
			propscat_sexesctry(ca1, ca2, ctryints,
				sage, eage, year, agemean,
				framedict, language, outfile, false)
			return agedict_bothsexes(sage, eage, agemean, language, ca2, fname)
		end
	end
end

function batchplotalias(child, plottype, language)
	if plottype == "sexesyrs"
		return conf["countries"][child]["alias"][language]
	elseif plottype == "yrsctry"
		return "$(child[1]) vs $(child[2])"
	elseif plottype == "sexesctry"
		return "$child"
	end
end

function batchplotyrs(child, plottype)
	if plottype == "sexesyrs"
		return [conf["countries"][child]["startyear"];
			conf["countries"][child]["endyear"]]
	elseif plottype == "yrsctry"
		return child
	elseif plottype == "sexesctry"
		return [child; child]
	end
end

function childplot(framedict, language, plottype, ca1, countries, years, yrtups)
	ages = batchages_caflt(ca1)
	sexes = conf["causes"][ca1]["sex"]
	childdicts = []
	if plottype == "sexesyrs"
		children = countries
	elseif plottype == "yrsctry"
		children = yrtups
	elseif plottype == "sexesctry"
		children = years
	end
	for child in children
		print("$child\n")
		agedicts = []
		ages = batchages_caflt(ca1)
		yrs = batchplotyrs(child, plottype)
		for age in ages
			print("$age\n")
			agedict = agebatchplot(framedict, age, plottype, language, ca1, child,
				countries, yrs[1], yrs[2], sexes)
			agedicts = vcat(agedicts, agedict)
		end
		if !(agedicts[1]==nothing)
			childdict = Dict(
				"alias" => batchplotalias(child, plottype, language),
				"ages" => agedicts)
			childdicts = vcat(childdicts, childdict)
		end
	end
	childdicts
end

function batchplot(framedict, language, plottype, causes = collect(keys(conf["causes"])),
	countries = collect(keys(conf["countries"])), years = conf["batchsexesyrs"],
	yrtups = conf["batchyrtups"])
	sort!(causes, by=((c)->
		(conf["causes"][c]["causeclass"],
		!(conf["causes"][c]["classtot"]),
		conf["causes"][c]["alias"][language])))
	sort!(countries, by=((c)->conf["countries"][c]["alias"][language]))
	cadicts = []
	if plottype == "sexesctry"
		causes = filter((ca)->conf["causes"][ca]["sex"]==[2;1], causes)
	end
	for ca1 in causes
		print("$plottype, $ca1\n")
		children = childplot(framedict, language, plottype, ca1,
			countries, years, yrtups)
		cadicts = vcat(cadicts, cadict(ca1, children, language))
	end
	Dict(
		"plottype" => plottype,
		"cadicts" => cadicts
	)
end

function writeplotlist(batchplotdict, outfile)
	tpl = read(normpath(datapath, "plotlist.mustache"), String)
	write(outfile, render(tpl, batchplotdict))
end

function ccflt(causeclass, language)
	cavals = filter((c)->c["causeclass"]==causeclass, collect(values(conf["causes"])))
	cadicts = map((c)->Dict("alias" => uppercasefirst(c["alias"][language]),
		"classtot" => c["classtot"], "codedesc" => c["codedesc"][language],
		"note" => c["note"][language]), cavals)
	sort!(cadicts, by=((c)->(!(c["classtot"]), c["alias"])))
end

function writetempl(language, fname, sitepath)
	haktemplpath = normpath(sitepath, "templates")
	mkpath(haktemplpath)
	tpl = read(normpath(datapath, "$fname.mustache"), String)
	maintempl = Dict()
	if (fname == "default" || fname == "site")
		maintempl["maintempldicts"] = map((p)->
			Dict("plottype" => p, "alias" => conf["plottypes"][p]["alias"][language]),
			keys(conf["plottypes"]))
		if fname == "default"
			outpath = haktemplpath
			ext = "html"
			maintempl["language"] = language
			for siteelem in ["dochead"; "homehead"; "sitehead";
				"sitenavhead"; "blognavhead"; "othlanghead"]
				maintempl[siteelem] = conf[siteelem][language]
			end
			maintempl["othlangs"] = map((od)->Dict(od),
				conf["othlangs"][language])
			maintempl["footer"] =
				read(normpath(datapath, "footer-$language.$ext"), String)
		elseif fname == "site"
			outpath = sitepath
			maintempl["indexpagetitle"] = conf["indexpagetitle"][language]
			ext = "hs"
		end
	elseif (fname == "index" || fname == "mortchartdoc")
		if fname == "index"
			ext = "html"
		elseif fname == "mortchartdoc"
			ext = "md"
			maintempl["maintempldicts"] = map((cc)->
				Dict("class" => cc,
				"alias" => conf["causeclasses"][cc]["alias"][language],
				"causes" => ccflt(parse(Int, cc), language)),
				keys(conf["causeclasses"]))
			maintempl["refhead"] = conf["refhead"][language]
		end
		outpath = sitepath
		maintempl["body"] = read(normpath(datapath, "$fname-$language.$ext"), String)
	end
	write(normpath(outpath, "$fname.$ext"), render(tpl, maintempl = maintempl))
end

function writeplotsite(framedict, language,
	sitepath = normpath(mainpath, "mortchart-site"))
	sitesubpaths = DataStructures.OrderedDict(
		"siteroot" => Dict("path" => "", "files" =>
			["default.csl"; "mortchartdoc_biber.bib"]),
		"css" => Dict("path" => "css", "files" => ["default.css"]),
		"images" => Dict("path" => "images", "files" => ["mortchartico.png"]),
		"charts" => Dict("path" => "charts", "files" => []))
	sitefullsubpaths = DataStructures.OrderedDict()
	for subpath in keys(sitesubpaths)
		sitefullsubpath = (normpath(sitepath, sitesubpaths[subpath]["path"]))
		sitefullsubpaths[subpath] = sitefullsubpath
		mkpath(sitefullsubpath)
		for file in sitesubpaths[subpath]["files"]
			cp(normpath(datapath, file), normpath(sitefullsubpath, file),
				remove_destination = true)
		end
	end
	for plottype in keys(conf["plottypes"])
		batchplotdict = batchplot(framedict, language, plottype)
		writeplotlist(batchplotdict, normpath(sitepath, "$plottype.html"))
	end
	chartdestpath = sitefullsubpaths["charts"]
	for chartfile in readdir(chartpath)
		mv(normpath(chartpath, chartfile), normpath(chartdestpath, chartfile),
		remove_destination = true)
	end
	for fname in ["default"; "index"; "site"; "mortchartdoc"]
		writetempl(language, fname, sitepath)
	end
end

end # module


using OmegaCore, OmegaGrad, Distributions
# using OmegaTestModels

# import ForwardDiff
# using Flux

# const ΩTs = [LinearΩ{Vector{Int}, UnitRange{Int64}, Vector{Float64}},
#              Omega.SimpleΩ{Vector{Int}, Float64},
#              LinearΩ{Vector{Int}, UnitRange{Int64}, TrackedArray{Float64, 1, Array{Float64,1}}}]

# const gradalgs = [Omega.TrackerGrad,
#                   Omega.ForwardDiffGrad,
#                   # Omega.ZygoteGrad
#                   ]

# function testgrad(OT = Omega.SimpleΩ{Vector{Int}, Float64})
#   μ = uniform(0.0, 1.0)
#   x = normal(μ, 1.0) + normal(μ, 1.0)
#   y = logerr(x ==ₛ 1.0)
#   ω = OT()
#   y(ω)
#   Omega.lineargradient(y, ω, Omega.ForwardDiffGrad)
# end

# testgrad(LinearΩ{Vector{Int}, UnitRange{Int64}, Vector{Float64}})

# # Test them all
# istracked(::Type{LinearΩ{I, SEG, T}}) where {I, SEG, T <: TrackedArray} = true
# istracked(ΩT) = false

# function testgrad2(; modelcond, ΩT, gradalg)
#   @show ΩT
#   println()
#   Omega.lineargradient(logerr(modelcond), ΩT(), gradalg)
# end

# # Filter to differentiable models
# models = filter(hascond ∧ isdiff, allmodels)

# # Dont use TrackerGrad on wrong Omegas, etc
# f1(m, ΩT, gradalg) = gradalg == Omega.TrackerGrad ?  istracked(ΩT) : true
# f2(m, ΩT, gradalg) = gradalg == Omega.ForwardDiffGrad ? !istracked(ΩT) : true

# function runtests()
#   for model in models, gradalg in gradalgs, ΩT in ΩTs
#     if f1(model, ΩT, gradalg) && f2(model, ΩT, gradalg)
#       @show model.name
#       @show ΩT
#       @show gradalg
#       testgrad2(; modelcond = model.cond, ΩT = ΩT, gradalg = gradalg)
#     end
#   end
# end

# runtests()

include("zygotetests.jl")


module test_h3_api_inspection

using Test
using H3.API # h3GetResolution h3GetBaseCell stringToH3 h3ToString h3IsValid h3IsResClassIII h3IsPentagon

@test h3GetResolution(0x085283473fffffff) == 5

@test h3GetBaseCell(0x085283473fffffff) == 20

@test stringToH3("85283473fffffff") == 0x085283473fffffff

@test h3ToString(0x085283473fffffff) == "85283473fffffff"

@test h3IsValid(0x85283473fffffff)
@test !h3IsValid(0x5004295803a8800)

@test h3IsResClassIII(0x085283473fffffff)

@test !h3IsPentagon(0x085283473fffffff)

end # module test_h3_api_inspection


for (func, elty) in [(:CLBlastSgbmv, Float32), (:CLBlastDgbmv, Float64),
                     (:CLBlastCgbmv, ComplexF32), (:CLBlastZgbmv, ComplexF64)]
    #TODO: (:CLBlastHgbmv, Float16)

    @eval function $func(layout::CLBlastLayout, a_transpose::CLBlastTranspose,
                         m::Integer, n::Integer, kl::Integer, ku::Integer,
                         alpha::$elty,
                         a_buffer::cl.CL_mem, a_offset::Integer, a_ld::Integer,
                         x_buffer::cl.CL_mem, x_offset::Integer, x_inc::Integer,
                         beta::$elty,
                         y_buffer::cl.CL_mem, y_offset::Integer, y_inc::Integer,
                         queue::cl.CmdQueue, event::cl.Event)
        err = ccall(
            ($(string(func)), libCLBlast),
            cl.CL_int,
            (Cint, Cint, Csize_t, Csize_t, Csize_t, Csize_t, $elty, Ptr{Cvoid}, Csize_t, Csize_t,
              Ptr{Cvoid}, Csize_t, Csize_t, $elty, Ptr{Cvoid}, Csize_t, Csize_t, Ptr{Cvoid}, Ptr{Cvoid}),
            Cint(layout), Cint(a_transpose), m, n, kl, ku, alpha, a_buffer, a_offset, a_ld,
              x_buffer, x_offset, x_inc, beta, y_buffer, y_offset, y_inc, Ref(queue), Ref(event)
        )
        if err != cl.CL_SUCCESS
            println(stderr, "Calling function $(string($func)) failed!")
            throw(CLBlastError(err))
        end
        return err
    end

    @eval function gbmv!(tA::Char, m::Integer, kl::Integer, ku::Integer, α::Number,
                         A::cl.CLArray{$elty,2}, x::cl.CLArray{$elty},
                         β::Number, y::cl.CLArray{$elty};
                         queue::cl.CmdQueue=cl.queue(y))
        # check and convert arguments
        if tA == 'N'
            a_transpose = CLBlastTransposeNo
        elseif tA == 'T' || (tA == 'C' && $elty <: Real)
            a_transpose = CLBlastTransposeYes
        elseif tA == 'C' && $elty <: Complex
            a_transpose = CLBlastTransposeConjugate
        else
            throw(ArgumentError("Transpose marker `tA` is $(tA) but only 'N', 'T', and 'C' are allowed."))
        end
        alpha = convert($elty, α)
        beta  = convert($elty, β)
        layout = CLBlastLayoutColMajor

        # output event
        event::cl.Event = cl.Event(C_NULL)

        $func(layout, a_transpose,
              m, size(A,2), kl, ku,
              alpha,
              pointer(A), 0, size(A,1),
              pointer(x), 0, 1,
              beta,
              pointer(y), 0, 1,
              queue, event)

        # wait for kernel
        cl.wait(event)

        y
    end

end


mutable struct ExactSimulator{M} <: AbstractSimulator
  algorithm::M
  next_jump_index::Int
  next_jump_time::Float64
  t::Float64
end

##### constructors ######
function ExactSimulator(algorithm::M) where M
  return ExactSimulator{M}(algorithm, 0, 0.0, 0.0)
end

##### accessors #####

"""
```
cumulative_intensity(simulator::ExactSimulator)
```

Retrieve the cumulative intensity of the process.
"""
@inline cumulative_intensity(simulator::ExactSimulator) = first(simulator.algorithm.total_rate)

"""
```
jump_rates(simulator::ExactSimulator)
```

Retrieve the current rates for each jump.

If the cache type is `HasRates()`, then `rate[j]` is the rate for reaction channel `j`.
Otherwise if the cache type is `HasSums()`, then `rate[j]` is the cumulative `sum(rate[1:j])`.
"""
@inline jump_rates(simulator::ExactSimulator) = simulator.algorithm.rates
# TODO: Fix this for FirstReactionMethod

"""
```
next_jump_index(simulator::ExactSimulator)
```

Retrieve the index of the next jump.
"""
@inline next_jump_index(simulator::ExactSimulator) = simulator.next_jump_index

"""
```
next_jump_time(simulator::ExactSimulator)
```

Retrieve the time to the next jump.
"""
@inline next_jump_time(simulator::ExactSimulator) = simulator.next_jump_time

# start a new simulation and initialize the stepper with the next event
@inline function initialize!(simulator::ExactSimulator, state, model, tfinal)
  # reset the current system time
  simulator.t = 0.0

  # initialize the simulation algorithm
  initialize!(simulator.algorithm, state, model)

  # seed the stepper with the first event
  generate_next_jump!(simulator)

  return nothing
end

# carry out the next jump event
@inline function step!(simulator::ExactSimulator, state, model)
  # unpack information
  j = simulator.next_jump_index
  s = simulator.next_jump_time

  # update the system time
  simulator.t = get_new_time(simulator)

  # update the system state according to jump j
  @inbounds @fastmath execute_jump!(state, model, j)

  # update statistics
  # TODO

  # update the simulation algorithm given that jump j occured
  update!(simulator.algorithm, state, model, j)

  # generate the next jump
  generate_next_jump!(simulator)

  return nothing
end

# generate the next (j,s) pair by invoking the stepper's event sampler
@inline function generate_next_jump!(simulator::ExactSimulator)
  # ask the algorithm for the next random jump
  j, s = generate_jump(simulator.algorithm)

  # update the stepper's fields
  simulator.next_jump_index = j
  simulator.next_jump_time = s

  return nothing
end

@inline get_new_time(simulator::ExactSimulator) = get_new_time(simulator.algorithm, simulator.t, simulator.next_jump_time)

##### helper functions #####

function no_update_step!(simulator::ExactSimulator, state, model)
  # unpack information
  j = simulator.next_jump_index
  s = simulator.next_jump_time

  # update the system time
  simulator.t = get_new_time(simulator)

  # update the system state according to jump j
  @inbounds @fastmath execute_jump!(state, model, j)

  return nothing
end