using ResumableFunctions
using DataStructures
import Base.deepcopy

VertexType = String
KeyType = Int64
EdgeType = Tuple{VertexType, VertexType}
MultiEdgeType = Tuple{VertexType, VertexType, KeyType}
AdjencyView = Dict{VertexType, Dict{VertexType, Set{KeyType}}}
WeightType = Float64
FlowType = Float64
FlowDictType = Dict{MultiEdgeType, FlowType}

mutable struct MultiDiGraph
    succ::AdjencyView
    pred::AdjencyView
end

function deepcopy(mdg::MultiDiGraph) :: MultiDiGraph
    return MultiDiGraph(Base.deepcopy(mdg.succ), Base.deepcopy(mdg.pred))
end

function _init_multidigraph(edges) :: MultiDiGraph
    mdg = MultiDiGraph(AdjencyView(), AdjencyView())
    for e in edges
        add_edge!(mdg, e)
    end
    return mdg
end

function fill_adjency_view(adj::AdjencyView, e::MultiEdgeType)
    if e[1] in keys(adj)
        if e[2] in keys(adj[e[1]])
            push!(adj[e[1]][e[2]], e[3])
        else
            adj[e[1]][e[2]] = Set{KeyType}([e[3]])
        end
        if ! (e[2] in keys(adj))
            adj[e[2]] = Dict{VertexType, Set{KeyType}}()
        end
    else
        adj[e[1]] = Dict{VertexType, Set{KeyType}}(e[2] => Set{KeyType}([e[3]]))
    end
end

MultiDiGraph(edges) = _init_multidigraph(edges)

function get_sub_graph(mdg::MultiDiGraph, sub_nodes::Set{VertexType}, from_new=true) :: MultiDiGraph
    if from_new
        sub_graph = MultiDiGraph(AdjencyView(), AdjencyView())
        for e in get_edges(mdg)
            if (e[1] in sub_nodes) && (e[2] in sub_nodes)
                add_edge!(sub_graph, e)
            end
        end
        return sub_graph
    else
        sub_graph = deepcopy(mdg)
        for u in setdiff(get_nodes(mdg), sub_nodes)
            remove_node!(sub_graph, u)
        end
        return sub_graph
    end
end

function get_edge_sub_graph(mdg::MultiDiGraph, edges, from_new=true) :: MultiDiGraph
    if from_new
        sub_graph = MultiDiGraph(AdjencyView(), AdjencyView())
        for e in edges
            add_edge!(sub_graph, e)
        end
        return sub_graph
    else
        sub_graph = deepcopy(mdg)
        for e in get_edges(mdg)
            if !(e in edges)
                remove_edge!(sub_graph, e)
            end
        end
        remove_isolates!(sub_graph)
        return sub_graph
    end
end

"""
Return set of nodes from an adjency dictionary.
"""
function get_nodes(mdg::MultiDiGraph)
    return keys(mdg.succ)
end

#=
Yields all edges in the given graph.
=#
@resumable function get_edges(mdg::MultiDiGraph)
    for (u, ns) in mdg.succ
        for (v, ks) in ns
            for k in ks
                @yield (u, v, k)
            end
        end
    end
end

#=
Yields an outgoing edge to the given node with the given adjency dictionary.
=#
@resumable function out_edges(mdg::MultiDiGraph, u::VertexType) :: Union{MultiEdgeType, Nothing}
    if u in keys(mdg.succ)
        for (v, ks) in mdg.succ[u]
            for k in ks
                @yield (u, v, k)
            end
        end
    end
end

#=
Yields an incoming edge to the given node with the given adjency dictionary.
=#
@resumable function in_edges(mdg::MultiDiGraph, v::VertexType) :: Union{MultiEdgeType, Nothing}
    if v in keys(mdg.pred)
        for (u, ks) in mdg.pred[v]
            for k in ks
                @yield (u, v, k)
            end
        end
    end
end

#=
Yields reachable edges in a DFS way from a start node.
=#
@resumable function edges_dfs(mdg::MultiDiGraph, start_node::VertexType) :: Union{MultiEdgeType, Nothing}
    visited_nodes = Set{VertexType}()
    visited_edges = Set{MultiEdgeType}()
    edges = Dict{VertexType, Any}()
    stack = Array{VertexType, 1}([start_node])
    while ! isempty(stack)
        current_node = last(stack)
        if ! (current_node in visited_nodes)
            edges[current_node] = out_edges(mdg, current_node)
            push!(visited_nodes, current_node)
        end
        edge = edges[current_node]()
        if edge === nothing
            pop!(stack)
        else
            if ! (edge in visited_edges)
                push!(visited_edges, edge)
                push!(stack, edge[2])
                @yield edge
            end
        end
    end
end

"""
Find a cycle with the given adjency dictionary. Return `nothing` if no cycle is
found.
"""
function find_cycle(mdg::MultiDiGraph, start_nodes=nothing::Union{Nothing, Set{VertexType}}) :: Union{Array{MultiEdgeType, 1}, Nothing}
    explored = Set{VertexType}()
    cycle = Array{MultiEdgeType, 1}()
    final_node = nothing
    if start_nodes === nothing
        nodes = get_nodes(mdg)
    else
        nodes = start_nodes
    end
    for start_node in nodes
        if start_node in explored
            # No loop is possible.
            continue
        end
        edges = Array{MultiEdgeType, 1}()
        # All nodes seen in this iteration of edge_dfs
        seen = Set{VertexType}()
        push!(seen, start_node)
        # Nodes in active path.
        active_nodes = Set{VertexType}()
        push!(active_nodes, start_node)
        previous_head = nothing
        for edge in edges_dfs(mdg, start_node)
            # Determine if this edge is a continuation of the active path.
            tail, head = edge[1], edge[2]
            if head in explored
                # Then we've already explored it. No loop is possible.
                continue
            end
            if (previous_head !== nothing) && (tail != previous_head)
                # This edge results from backtracking.
                # Pop until we get a node whose head equals the current tail.
                # So for example, we might have:
                #  (0, 1), (1, 2), (2, 3), (1, 4)
                # which must become:
                #  (0, 1), (1, 4)
                while true
                    if ! isempty(edges)
                        popped_edge = pop!(edges)
                        popped_head = popped_edge[2]
                        filter!(x -> x != popped_head, active_nodes)
                    else
                        edges = Array{MultiEdgeType, 1}()
                        active_nodes = Set{VertexType}([tail])
                        break
                    end
                    if ! isempty(edges)
                        last_head = last(edges)[2]
                        if tail == last_head
                            break
                        end
                    end
                end
            end
            push!(edges, edge)
            if head in active_nodes
                # We have a loop!
                cycle = vcat(cycle, edges)
                final_node = head
                break
            else
                push!(seen, head)
                push!(active_nodes, head)
                previous_head = head
            end
        end
        if ! isempty(cycle)
            break
        else
            for n in seen
                push!(explored, n)
            end
        end
    end
    if isempty(cycle)
        return nothing
    end
    # We now have a list of edges which ends on a cycle.
    # So we need to remove from the beginning edges that are not relevant.
    trimming_i = 1
    for (i, edge) in enumerate(cycle)
        tail, head = edge[1], edge[2]
        if tail == final_node
            trimming_i = i
            break
        end
    end
    return cycle[trimming_i:length(cycle)]
end

"""
Return all nodes reachable from `source` in the given adjency dictionary.
"""
function descendants(mdg::MultiDiGraph, source::VertexType, strict=false) :: Set{VertexType}
    des = Set{VertexType}()
    node_stack = [source]
    while ! isempty(node_stack)
        current_node = pop!(node_stack)
        if ! (current_node in des)
            push!(des, current_node)
            for e in out_edges(mdg, current_node)
                push!(node_stack, e[2])
            end
        end
    end
    if strict
        filter!(x -> x != source, des)
    end
    return des
end

"""
Compute the shortest path from a set of sources to a target node.
"""
function dijkstra_multisource(mdg::MultiDiGraph, sources::Set{VertexType}, target::VertexType, weight::Dict)
    dists = DefaultDict{VertexType, WeightType}(typemax(WeightType))
    visited = DefaultDict{VertexType, Bool}(false)
    parents = Dict{VertexType, VertexType}()
    fringe = PriorityQueue{VertexType, WeightType}()

    for s=sources
        dists[s] = zero(WeightType)
        visited[s] = true
        fringe[s] = zero(WeightType)
    end

    while !isempty(fringe)
        u = dequeue!(fringe)
        d = dists[u]
        if u == target
            break
        end
        if u in keys(mdg.succ)
            for v in keys(mdg.succ[u])
                alt = d + weight[(u, v)]
                if (visited[v]) && alt < dists[u]
                    throw(ErrorException("Contradictory paths found:
                                            negative weights?: alt='$(alt)', dists[v]='$(dists[v])"))
                end
                if (!visited[v]) || (alt < dists[v])
                    visited[v] = true
                    dists[v] = alt
                    parents[v] = u
                    fringe[v] = alt
                end
            end
        end
    end
    if ! visited[target]
        return nothing
    end
    path = [target]
    node = target
    while ! (node in sources)
        node = parents[node]
        push!(path, node)
    end
    return (path[end:-1:1], dists[target])
end

"""
Return the lowest common ancestor of the two given nodes in the given DAG.
Return `nothing` if none is found.
"""
function lowest_common_ancestor(mdg::MultiDiGraph, node1::VertexType, node2::VertexType) :: Union{VertexType, Nothing}
    # First node backtracking
    nodes_stack = [node1]
    ancestors1 = Set{VertexType}([node1])
    while !isempty(nodes_stack)
        current_node = pop!(nodes_stack)
        for e in in_edges(mdg, current_node)
            if e[1] in ancestors1
                continue
            end
            push!(ancestors1, e[1])
            push!(nodes_stack, e[1])
        end
    end
    # Second node backtracking
    ancestors2 = Set{VertexType}([node2])
    nodes_stack = [node2]
    while !isempty(nodes_stack)
        current_node = pop!(nodes_stack)
        if current_node in ancestors1
            return current_node  # lca found
        end
        for e in in_edges(mdg, current_node)
            if e[1] in ancestors2
                continue
            end
            push!(ancestors2, e[1])
            push!(nodes_stack, e[1])
        end
    end
    # no common ancestor found
    return nothing
end

function add_node!(graph::MultiDiGraph, n::VertexType)
    if ! (n in keys(graph.succ))
        graph.succ[n] = Dict{VertexType, Set{KeyType}}()
    end
    if ! (n in keys(graph.pred))
        graph.pred[n] = Dict{VertexType, Set{KeyType}}()
    end
end

function add_edge!(graph::MultiDiGraph, e::MultiEdgeType)
    u, v, k = e
    add_node!(graph, u)
    add_node!(graph, v)
    if v in keys(graph.succ[u])
        push!(graph.succ[u][v], k)
    else
        graph.succ[u][v] = Set{KeyType}([k])
    end
    if u in keys(graph.pred[v])
        push!(graph.pred[v][u], k)
    else
        graph.pred[v][u] = Set{KeyType}([k])
    end
end

function remove_node!(graph::MultiDiGraph, n::VertexType)
    for u in keys(graph.pred[n])
        delete!(graph.succ[u], n)
    end
    for v in keys(graph.succ[n])
        delete!(graph.pred[v], n)
    end
    delete!(graph.succ, n)
    delete!(graph.pred, n)
end

function remove_edge!(graph::MultiDiGraph, e::MultiEdgeType)
    u, v, k = e
    delete!(graph.succ[u][v], k)
    if isempty(graph.succ[u][v])
        delete!(graph.succ[u], v)
    end
    delete!(graph.pred[v][u], k)
    if isempty(graph.pred[v][u])
        delete!(graph.pred[v], u)
    end
end

function has_one_parent(graph::MultiDiGraph, n::VertexType) :: Bool
    parent_found = false
    for (u, ks) in graph.pred[n]
        if parent_found
            return false
        end
        parent_found = true
    end
    return parent_found
end

function has_no_parent(graph::MultiDiGraph, n::VertexType) :: Bool
    return isempty(graph.pred[n])
end

function remove_isolates!(graph::MultiDiGraph)
    nodes_to_remove = Set{VertexType}()
    for (u, ns) in graph.succ
        if isempty(ns) && isempty(graph.pred[u])
            push!(nodes_to_remove, u)
        end
    end
    for n in nodes_to_remove
        remove_node!(graph, n)
    end
end


using Test
using Dates
using ExtensibleScheduler
using TimeFrames

@testset "CustomTrigger" begin
    f = (dt_previous_fire, dt_now) -> dt_now + Dates.Minute(5)
    trigger = Trigger(f)

    n_executed = 0

    function sample_action()
        n_executed += 1
    end

    dt_start = DateTime(2010, 1, 4, 0, 0)
    dt = dt_start
    clock = SimClock(dt)
    sched = BlockingScheduler(clock=clock, delayfunc=NoSleep)
    action = Action(sample_action)

    add(sched, action, trigger)
    run_pending(sched)
    @test n_executed == 0

    for i in 1:13
        dt += Dates.Minute(1)
        #println("Clock set to $dt")
        set(clock, dt)
        run_pending(sched)
    end

    @test n_executed == 2

end


module VTK

using WriteVTK
using WriteVTK.VTKCellTypes: VTK_VERTEX
using WriteVTK.VTKCellTypes: VTK_LINE
using WriteVTK.VTKCellTypes: VTK_TRIANGLE
using WriteVTK.VTKCellTypes: VTK_QUAD
using WriteVTK.VTKCellTypes: VTK_TETRA
using WriteVTK.VTKCellTypes: VTK_HEXAHEDRON

using UnstructuredGrids.Core

export writevtk

function writevtk(grid::UGrid,filebase;celldata=Dict(),pointdata=Dict())
  points = coordinates(grid)
  cells = _vtkcells(grid)
  vtkfile = vtk_grid(filebase, points, cells, compress=false)
  for (k,v) in celldata
    vtk_cell_data(vtkfile,_prepare_data(v),k)
  end
  for (k,v) in pointdata
    vtk_point_data(vtkfile,_prepare_data(v),k)
  end
  outfiles = vtk_save(vtkfile)
end

function writevtk(r::RefCell,filebase)
  for d in 0:(ndims(r)-1)
    f = "$(filebase)_$d"
    grid = UGrid(r,dim=d)
    writevtk(grid,f)
  end
end

_prepare_data(v) = v

function _prepare_data(v::AbstractVector{Bool})
  b = similar(v,Int)
  b .= 0
  b[v] .= 1
  b
end

function _vtkcells(grid::UGrid)
  vtkid_to_vtktype = _vtkcelltypedict()
  cell_to_ctype = celltypes(grid)
  cell_to_points_data = list(connections(grid))
  cell_to_points_ptrs = ptrs(connections(grid))
  ctype_to_refcell = refcells(grid)
  ctype_to_vtkid = [ vtkid(rc) for rc in ctype_to_refcell ]
  ctype_to_vtknodes = [ vtknodes(rc) for rc in ctype_to_refcell ]
  ncells = length(cell_to_ctype)
  vtkcells = Array{MeshCell}(undef,(ncells))
  for cell in 1:ncells
    a = cell_to_points_ptrs[cell]
    b = cell_to_points_ptrs[cell+1]-1
    cellpoints = cell_to_points_data[a:b]
    ctype = cell_to_ctype[cell]
    vtkid = ctype_to_vtkid[ctype]
    vtknodes = ctype_to_vtknodes[ctype]
    vtktype = vtkid_to_vtktype[vtkid]
    vtkcells[cell] = MeshCell(vtktype,cellpoints[vtknodes])
  end
  vtkcells
end

function _vtkcelltypedict()
  d = Dict{Int,WriteVTK.VTKCellTypes.VTKCellType}()
  d[VTK_VERTEX.vtk_id] = VTK_VERTEX
  d[VTK_LINE.vtk_id] = VTK_LINE
  d[VTK_LINE.vtk_id] = VTK_LINE
  d[VTK_TRIANGLE.vtk_id] = VTK_TRIANGLE
  d[VTK_QUAD.vtk_id] = VTK_QUAD
  d[VTK_TETRA.vtk_id] = VTK_TETRA
  d[VTK_HEXAHEDRON.vtk_id] = VTK_HEXAHEDRON
  d
end

end #module VTK


module CIAOAlgorithms #how we use this module? always paths are saved in a same format/ any other formats? how to switch to dev mode? in the dev mode, what happens to the other packages? how to add a package to the dev mode?

const RealOrComplex{R} = Union{R,Complex{R}}
const Maybe{T} = Union{T, Nothing}

# utilities
include("utilities/indexingUtilities.jl")
include("utilities/IndexIterator.jl")
include("utilities/lbfgs.jl")
# include("utilities/indNonnegativeBallL2.jl")
# algorithms
include("algorithms/Finito/Finito.jl")
include("algorithms/ProShI/ProShI.jl")
include("algorithms/SVRG/SVRG.jl")
include("algorithms/SGD/SGD.jl")
include("algorithms/SARAH/SARAH.jl")
include("algorithms/SAGA_SAG/SAGA.jl")

end # module


"""
    solve_decomposition_by_id(manager::ExperimentManager, instance_name, added_edges; resolve=false)

Retrieve the data from the database and solve one decomposition.
"""
function solve_decomposition(manager::ExperimentManager,
                             instance_name::String,
                             oid::Union{UUID,AbstractString},
                             set_cholesky_time_limit::Bool=true,
                             mpi_infos::AbstractDict=Dict();
                             time_limit::Float64=100000.)
    paths = db.getOPFpath_one(manager.instances, instance_name)
    path_opf_ctr = paths["ctr"]
    path_opf_mat = paths["mat"]
    solve_decomposition(manager,
                        instance_name,
                        oid,
                        path_opf_ctr,
                        path_opf_mat,
                        set_cholesky_time_limit,
                        mpi_infos;
                        time_limit=time_limit)
end

"""
    solve_decomposition(manager::ExperimentManager, instance_name, added_edges, path_opf_ctr, path_opf_mat; resolve=false)

Solve one decomposition.
"""
function solve_decomposition(manager::ExperimentManager,
                             instance_name::String,
                             oid::Union{UUID,AbstractString},
                             path_opf_ctr::String,
                             path_opf_mat::String,
                             set_cholesky_time_limit::Bool=true,
                             mpi_infos::AbstractDict=Dict();
                             path_mosek::AbstractString=joinpath("data", "mosek_logs", instance_name, "$instance_name.log"),
                             time_limit::Float64=100000.)
    decomposition = db.getdecomposition(manager.decompositions, instance_name, oid)
    instance_name = decomposition["_id"]["instance_name"]
    if decomposition["nb_added_edges"] != 0 && set_cholesky_time_limit
        time_cholesky = db.getcholeskytime(manager.decompositions, instance_name)
        if !isnothing(time_cholesky)
            time_limit = time_cholesky * 2.
            @info "Computing time limit: $instance_name: $time_limit"
        end
    end
    cliques = decomposition["cliques"]
    cliquetree = decomposition["cliquetree"]
    time, nb_iter, path_mosek, objective, m, nlc = slv.solve_sdp(instance_name,
                                                                 cliques,
                                                                 cliquetree,
                                                                 path_opf_ctr,
                                                                 path_opf_mat;
                                                                 time_limit=time_limit)
    #time = parse(Float64, time)
    features = Mongoc.BSON("solver" => Mongoc.BSON("solving_time" => time,
                                                   "nb_iter" => nb_iter,
                                                   "nonzerosfactor" => readf.get_nonzeros_factor(path_mosek),
                                                   "objective" => objective,
                                                   "m" => m,
                                                   "nlc" => nlc,
                                                   "molzahn_heuristic" => compute_molzahn_heuristic(cliques, m, nlc)))
    fields = Mongoc.BSON("paths.mosek_log" => path_mosek)
    db.setfields_decomposition!(manager.decompositions, instance_name, oid, fields)
    db.setfeatures_decomposition!(manager.decompositions, instance_name, oid, features)
    !isempty(mpi_infos) && db.setMPIinfos(manager.decompositions, instance_name, oid, mpi_infos)
end


"""
`FlipY <: ImageTransformation`

Description
============

Reverses the pixel order along the y-axis. In other words it
mirrors the Image along the x-axis (i.e. vertically).

If created using the parameter `chance`, the operation will be
lifted into a `Either`. See the documentation of `Either` for more
information.

Usage
======

    FlipY()

    FlipY(chance)

Arguments
==========

- **`chance`** : The probability of the operation to be applied to
the image if the operation is used in `transform`. Must be within
the set [0, 1].

Methods
========

- **`transform`** : Applies the transformation to the given Image
or set of images

Author(s)
==========

- Christof Stocker (Github: https://github.com/Evizero)

Examples
========

    using Images
    using Augmentor

    # Create a test image with 4 pixel.
    img = grayim(UInt8[200 150; 50 1])

    # Apply the image operation to the image
    img_new = transform(FlipY(), img)

    # Check manually if the result is as expected
    @assert img_new == grayim(UInt8[150 200; 1 50])

see also
=========

`ImageTransformation`, `Either`, `transform`
"""
immutable FlipY <: ImageTransformation
end

FlipY(chance) = ProbableOperation(FlipY(); chance = chance)

Base.showcompact(io::IO, op::FlipY) = print(io, "Flip y-axis.")
multiplier(::FlipY) = 1

function transform{T<:AbstractImage}(op::FlipY, img::T)
    result = copyproperties(img, flipy(img))::T
    _log_operation!(result, op)
end


#=

EGeoHiSSE equations

Ignacio Quintero Mächler

t(-_-t)

Created 18 03 2019
=#

"""
    exp_expr(k    ::Int64,
             h    ::Int64,
             ny   ::Int64,
             model::NTuple{3,Bool})

Return exponential expression for one time evaluation of covariates.
"""
function exp_expr(k    ::Int64,
                  h    ::Int64,
                  ny   ::Int64,
                  model::NTuple{3,Bool})

  # last non β parameters
  bbase = h*(k+1) + 2*k*h + k*(k-1)*h + h*(h-1)

  # ncov
  ncov = model[1]*k + model[2]*k + model[3]*k*(k-1)

  # y per parameter
  yppar = isone(ny) ? 1 : div(ny,ncov)

  # starting indices for models 2 and 3
  m2s = model[1]*k*h
  m3s = m2s + model[2]*k*h

  y2s = model[1]*yppar*k
  y3s = y2s + model[2]*yppar*k

  # start expression
  ex = quote end
  pop!(ex.args)

  isone(ny) && push!(ex.args, :(r1 = r[1]))

  # speciation
  if model[1]
    for j = Base.OneTo(h), i = Base.OneTo(k)
      coex = Expr(:call, :+)
      for yi in Base.OneTo(yppar)
        rex = isone(ny) ? :(r1) : :(r[$(yi+yppar*(i-1))])
        push!(coex.args, :(p[$(bbase + yi + yppar*((i-1) + k*(j-1)))] * $rex))
      end
      push!(ex.args, :(eaft[$(i + k*(j-1))] = p[$(i + (k+1)*(j-1))]*exp($coex)))
    end
  end
  # extinction
  if model[2]
    for j = Base.OneTo(h), i = Base.OneTo(k)
      coex = Expr(:call, :+)
      for yi in Base.OneTo(yppar)
        rex = isone(ny) ? :(r1) : :(r[$(y2s + yi+yppar*(i-1))])
        push!(coex.args,
          :(p[$(bbase + yi + m2s*yppar + yppar*((i-1) + k*(k-1)*(j-1)))] * $rex))
      end
      push!(ex.args, :(eaft[$(i + k*(j-1) + m2s)] =
        p[$((k+1)*h + i + k*(j-1))]*exp($coex)))
    end
  end
  # transition
  if model[3]
    for j = Base.OneTo(h), i = Base.OneTo(k*(k-1))
      coex = Expr(:call, :+)
      for yi in Base.OneTo(yppar)
        rex = isone(ny) ? :(r1) : :(r[$(y3s + yi+yppar*(i-1))])
        push!(coex.args,
          :(p[$(bbase + yi + m3s*yppar + yppar*((i-1) + k*(k-1)*(j-1)))] * $rex))
      end
      push!(ex.args, :(eaft[$(i + k*(j-1) + m3s)] =
        p[$(h*(2k+1) + i + k*(k-1)*(j-1))]*exp($coex)))
    end
  end

  return ex
end

"""
    noevents_expr(si   ::Int64,
                  s    ::Sgh,
                  ls   ::Int64,
                  oa   ::Array{Int64,1},
                  k    ::Int64,
                  h    ::Int64,
                  ns   ::Int64,
                  model::NTuple{3,Bool},
                  ext  ::Bool)

Return expression for no events in Vector form.
"""
function noevents_expr(si   ::Int64,
                       s    ::Sgh,
                       ls   ::Int64,
                       oa   ::Array{Int64,1},
                       k    ::Int64,
                       h    ::Int64,
                       ns   ::Int64,
                       model::NTuple{3,Bool},
                       ext  ::Bool)

  # starting indices for models 2 and 3
  m2s = model[1]*k*h
  m3s = m2s + model[2]*k*h

  # if extinction ODE
  wu = ext ? ns : 0

  ts = isone(ls) ? 0 : k*h*k

  # between-region speciation
  ex = Expr(:call, :+)
  ls > 1 && push!(ex.args, :($(2.0^(ls-1) - 1.0) * p[$(k+1+s.h*(k+1))]))

  for v in s.g

    #speciation
    if model[1]
      push!(ex.args, :(eaft[$(v + s.h*k)]))
    else
      push!(ex.args, :(p[$(v + s.h*(k+1))]))
    end

    # extinction
    if model[2]
      push!(ex.args, :(eaft[$(m2s + v + s.h*k)]))
    else
      push!(ex.args, :(p[$(v + (k+1)*h + s.h*k + ts)]))
    end

    # dispersal
    if model[3]
      for j = oa
        j -= v <= j ? 1 : 0
        push!(ex.args,
          :(eaft[$(m3s + s.h*k*(k-1) + (k-1)*(v-1) + j)]))
      end
    else
      for j = oa
        j -= v <= j ? 1 : 0
        push!(ex.args,
          :(p[$(2h*k + h + s.h*k*(k-1) + (k-1)*(v-1) + j)]))
      end
    end
  end

  # add hidden state shifts
  for hi in setdiff(0:(h-1), s.h)
    hi -= s.h <= hi ? 0 : -1
    push!(ex.args, :(p[$(h*(3k + 1 + k*(k-1)) + s.h*(h-1) + hi)]))
  end

  # multiply by u
  ex = :(-1.0 * $ex * u[$(si + wu)])

  return ex
end

"""
    noevents_expr(si     ::Int64,
                  s      ::Sgh,
                  ls     ::Int64,
                  oa     ::Array{Int64,1},
                  k      ::Int64,
                  h      ::Int64,
                  ns     ::Int64,
                  i      ::Int64,
                  model  ::NTuple{3,Bool},
                  ext    ::Bool)

Return expression for no events in Matrix form.
"""
function noevents_expr(si     ::Int64,
                       s      ::Sgh,
                       ls     ::Int64,
                       oa     ::Array{Int64,1},
                       k      ::Int64,
                       h      ::Int64,
                       ns     ::Int64,
                       i      ::Int64,
                       model  ::NTuple{3,Bool})

  # starting indices for models 2 and 3
  m2s = model[1]*k*h
  m3s = m2s + model[2]*k*h

  ts = isone(ls) ? 0 : k*h*k

  # between-region speciation
  ex = Expr(:call, :+)
  ls > 1 && push!(ex.args, :($(2.0^(ls-1) - 1.0) * p[$(k+1+s.h*(k+1))]))

  for v in s.g

    #speciation
    if model[1]
      push!(ex.args, :(eaft[$(v + s.h*k)]))
    else
      push!(ex.args, :(p[$(v + s.h*(k+1))]))
    end

    # extinction
    if model[2]
      push!(ex.args, :(eaft[$(m2s + v + s.h*k)]))
    else
      push!(ex.args, :(p[$(v + (k+1)*h + s.h*k + ts)]))
    end

    # dispersal
    if model[3]
      for j = oa
        j -= v <= j ? 1 : 0
        push!(ex.args,
          :(eaft[$(m3s + s.h*k*(k-1) + (k-1)*(v-1) + j)]))
      end
    else
      for j = oa
        j -= v <= j ? 1 : 0
        push!(ex.args,
          :(p[$(2h*k + h + s.h*k*(k-1) + (k-1)*(v-1) + j)]))
      end
    end
  end

  # add hidden state shifts
  for hi in setdiff(0:(h-1), s.h)
    hi -= s.h <= hi ? 0 : -1
    push!(ex.args, :(p[$(h*(3k + 1 + k*(k-1)) + s.h*(h-1) + hi)]))
  end

  # multiply by u
  ex = :(-1.0 * $ex * u[$si, $i])

  return ex
end

"""
    localext_expr(s ::Sgh,
                  S ::Array{Sgh,1},
                  k ::Int64,
                  h ::Int64,
                  model::NTuple{3,Bool})

Return expression for local extinction in Vector form.
"""
function localext_expr(s ::Sgh,
                       S ::Array{Sgh,1},
                       k ::Int64,
                       h ::Int64,
                       model::NTuple{3,Bool})

  ex = Expr(:call, :+)

  if model[2]
    # starting indices for models 2 and 3
    m2s = model[1]*k*h

    for i = s.g
      push!(ex.args, :(eaft[$(m2s + i + s.h*k)] *
                       u[$(findfirst(x -> isSghequal(x,
                                          Sgh(setdiff(s.g, i),s.h)), S))]))
    end
  else
    for i = s.g
      push!(ex.args, :(p[$(i + (k+1)*h + s.h*k + k*h*k)] *
                       u[$(findfirst(x -> isSghequal(x,
                                          Sgh(setdiff(s.g, i),s.h)), S))]))
    end
  end

  return ex
end

"""
    localext_expr(s     ::Sgh,
                  S     ::Array{Sgh,1},
                  k     ::Int64,
                  h     ::Int64,
                  i     ::Int64
                  model ::NTuple{3,Bool})

Return expression for local extinction in Matrix form.
"""
function localext_expr(s     ::Sgh,
                       S     ::Array{Sgh,1},
                       k     ::Int64,
                       h     ::Int64,
                       i     ::Int64,
                       model ::NTuple{3,Bool})

  ex = Expr(:call, :+)

  if model[2]
    # starting indices for models 2 and 3
    m2s = model[1]*k*h

    for j = s.g
      push!(ex.args, :(eaft[$(m2s + j + s.h*k)] *
                       u[$(findfirst(x -> isSghequal(x,
                                          Sgh(setdiff(s.g, j),s.h)), S)),$i]))
    end
  else
    for j = s.g
      push!(ex.args, :(p[$(j + (k+1)*h + s.h*k + k*h*k)] *
                       u[$(findfirst(x -> isSghequal(x,
                                          Sgh(setdiff(s.g, j),s.h)), S)),$i]))
    end
  end

  return ex
end

"""
    dispersal_expr(s    ::Sgh,
                   oa   ::Array{Int64,1},
                   S    ::Array{Sgh,1},
                   ns   ::Int64,
                   k    ::Int64,
                   h    ::Int64,
                   model::NTuple{3,Bool},
                   ext  ::Bool)

Return expression for dispersal in Vector form.
"""
function dispersal_expr(s    ::Sgh,
                        oa   ::Array{Int64,1},
                        S    ::Array{Sgh,1},
                        ns   ::Int64,
                        k    ::Int64,
                        h    ::Int64,
                        model::NTuple{3,Bool},
                        ext  ::Bool)

  m3s = model[1]*k*h + model[2]*k*h

  wu = ext ? ns : 0

  ida = findall(x -> length(union(s.g, x.g))     == length(s.g)+1 &&
                     length(intersect(s.g, x.g)) == length(s.g)   &&
                     s.h == x.h, S)

  ex = Expr(:call, :+)

  # if dispersal covariate
  if model[3]
    for a = s.g, (i, j) = enumerate(oa)
      j -= a <= j ? 1 : 0
      push!(ex.args, :(eaft[$(m3s + s.h*k*(k-1) + (k-1)*(a-1) + j)] *
                       u[$(ida[i] + wu)]))
    end
  else
    for a = s.g, (i, j) = enumerate(oa)
      j -= a <= j ? 1 : 0
      push!(ex.args, :(p[$(2h*k + h + s.h*k*(k-1) + (k-1)*(a-1) + j)] *
                       u[$(ida[i] + wu)]))
    end
  end

  return ex
end

"""
    dispersal_expr(s    ::Sgh,
                   oa   ::Array{Int64,1},
                   S    ::Array{Sgh,1},
                   ns   ::Int64,
                   k    ::Int64,
                   h    ::Int64,
                   i    ::Int64,
                   model::NTuple{3,Bool})

Return expression for dispersal in Matrix form.
"""
function dispersal_expr(s    ::Sgh,
                        oa   ::Array{Int64,1},
                        S    ::Array{Sgh,1},
                        ns   ::Int64,
                        k    ::Int64,
                        h    ::Int64,
                        i    ::Int64,
                        model::NTuple{3,Bool})

  m3s = model[1]*k*h + model[2]*k*h

  ida = findall(x -> length(union(s.g, x.g))     == length(s.g)+1 &&
                     length(intersect(s.g, x.g)) == length(s.g)   &&
                     s.h == x.h, S)

  ex = Expr(:call, :+)

  # if dispersal covariate
  if model[3]
    for a = s.g, (l, j) = enumerate(oa)
      j -= a <= j ? 1 : 0
      push!(ex.args, :(eaft[$(m3s + s.h*k*(k-1) + (k-1)*(a-1) + j)] *
                       u[$(ida[l]), $i]))
    end
  else
    for a = s.g, (l, j) = enumerate(oa)
      j -= a <= j ? 1 : 0
      push!(ex.args, :(p[$(2h*k + h + s.h*k*(k-1) + (k-1)*(a-1) + j)] *
                       u[$(ida[l]), $i]))
    end
  end

  return ex
end

"""
    hidtran_expr(s  ::Sgh,
                 S  ::Array{Sgh,1},
                 ns ::Int64,
                 k  ::Int64,
                 h  ::Int64,
                 i  ::Int64)

Return expression for hidden states transitions in Matrix form.
"""
function hidtran_expr(s  ::Sgh,
                      S  ::Array{Sgh,1},
                      ns ::Int64,
                      k  ::Int64,
                      h  ::Int64,
                      i  ::Int64)

  hs = findall(x -> isequal(s.g, x.g) &&
                   !isequal(s.h, x.h), S)

  ex = Expr(:call, :+)
  for j in hs
    hi = S[j].h
    hi -= s.h <= hi ? 0 : -1
    push!(ex.args, :(p[$(h*(3k + 1 + k*(k-1)) + s.h*(h-1) + hi)] *
                     u[$j, $i]))
  end
  return ex
end

"""
    hidtran_expr(s  ::Sgh,
                 S  ::Array{Sgh,1},
                 ns ::Int64,
                 k  ::Int64,
                 h  ::Int64,
                 ext::Bool)

Return expression for hidden states transitions in Vector form.
"""
function hidtran_expr(s  ::Sgh,
                      S  ::Array{Sgh,1},
                      ns ::Int64,
                      k  ::Int64,
                      h  ::Int64,
                      ext::Bool)

  wu = ext ? ns : 0

  hs = findall(x -> isequal(s.g, x.g) &&
                   !isequal(s.h, x.h), S)

  ex = Expr(:call, :+)
  for i in hs
    hi = S[i].h
    hi -= s.h <= hi ? 0 : -1
    push!(ex.args, :(p[$(h*(3k + 1 + k*(k-1)) + s.h*(h-1) + hi)] *
                     u[$(i + wu)]))
  end
  return ex
end

"""
    wrspec_expr(si   ::Int64,
                s    ::Sgh,
                ns   ::Int64,
                k    ::Int64,
                i    ::Int64,
                model::NTuple{3,Bool})

Return expression for within-region speciation in Matrix form.
"""
function wrspec_expr(si   ::Int64,
                     s    ::Sgh,
                     ns   ::Int64,
                     k    ::Int64,
                     i    ::Int64,
                     model::NTuple{3,Bool})

  # if model speciation
  if model[1]
    if isone(length(s.g))
      ex = Expr(:call, :*, 2.0)
      for j = s.g
        push!(ex.args, :(eaft[$(j + s.h*k)] *
               rE[$(j + s.h*(2^k-1))] * u[$(j + s.h*(2^k-1)), $i]))
      end
    else
      ex = Expr(:call, :+)
      for j = s.g
        push!(ex.args, :(eaft[$(j + s.h*k)] *
          (rE[$(j + s.h*(2^k-1))] * u[$(si), $i] +
           rE[$(si)] * u[$(j + s.h*(2^k-1)), $i])))
      end
    end

  # if model *NOT* speciation
  else
    if isone(length(s.g))
      ex = :(2.0 * p[$si] * rE[$(si)] * u[$(si)])
    else
      ex = Expr(:call, :+)
      for j = s.g
        push!(ex.args, :(p[$(j + s.h*(k+1))] *
          (rE[$(j + s.h*(2^k-1))] * u[$(si), $i] +
           rE[$(si)] * u[$(j + s.h*(2^k-1)), $i])))
      end
    end
  end

  return ex
end

"""
    wrspec_expr(si   ::Int64,
                s    ::Sgh,
                ns   ::Int64,
                k    ::Int64,
                model::NTuple{3,Bool})

Return expression for within-region speciation in Vector form.
"""
function wrspec_expr(si   ::Int64,
                     s    ::Sgh,
                     ns   ::Int64,
                     k    ::Int64,
                     model::NTuple{3,Bool})

  # if model speciation
  if model[1]
    if isone(length(s.g))
      ex = Expr(:call, :*, 2.0)
      for i = s.g
        push!(ex.args, :(eaft[$(i + s.h*k)] *
               u[$(i + s.h*(2^k-1) + ns)] * u[$(i + s.h*(2^k-1))]))
      end
    else
      ex = Expr(:call, :+)
      for i = s.g
        push!(ex.args, :(eaft[$(i + s.h*k)] *
          (u[$(i + s.h*(2^k-1) + ns)] * u[$(si)] +
           u[$(si + ns)] * u[$(i + s.h*(2^k-1))])))
      end
    end

  # if model *NOT* speciation
  else
    if isone(length(s.g))
      ex = :(2.0 * p[$si] * u[$(si + ns)] * u[$(si)])
    else
      ex = Expr(:call, :+)
      for i = s.g
        push!(ex.args, :(p[$(i + s.h*(k+1))] *
          (u[$(i + s.h*(2^k-1) + ns)] * u[$(si)] +
           u[$(si + ns)] * u[$(i + s.h*(2^k-1))])))
      end
    end
  end

  return ex
end

"""
    brspec_expr(s   ::Sgh,
                S   ::Array{Sgh,1},
                ns  ::Int64,
                k   ::Int64,
                i   ::Int64,
                ext ::Bool,
                flow::Bool)

Return expression for between-region speciation in Vector form.
"""
function brspec_expr(s   ::Sgh,
                     S   ::Array{Sgh,1},
                     ns  ::Int64,
                     k   ::Int64,
                     ext ::Bool,
                     flow::Bool)

  if ext
    va = vicsubsets(s.g)[1:div(end,2)]
    wu1 = flow ? 0 : ns
    wu2 = flow ? 0 : ns
  else
    va = vicsubsets(s.g)
    wu1 = flow ? 0 : ns
    wu2 = 0
  end

  ex = Expr(:call, :+)
  for (la, ra) = va
  push!(ex.args,
    :(u[$(findfirst(x -> isequal(ra,x.g), S) + (2^k-1)*s.h + wu1)] *
      u[$(findfirst(x -> isequal(la,x.g), S) + (2^k-1)*s.h + wu2)]))
  end

  ex = :(p[$(k+1+s.h*(k+1))] * $ex)

  return ex
end

"""
    brspec_expr(s   ::Sgh,
                S   ::Array{Sgh,1},
                ns  ::Int64,
                k   ::Int64,
                i   ::Int64)

Return expression for between-region speciation in Matrix form.
"""
function brspec_expr(s   ::Sgh,
                     S   ::Array{Sgh,1},
                     ns  ::Int64,
                     k   ::Int64,
                     i   ::Int64)

  va = vicsubsets(s.g)

  ex = Expr(:call, :+)

  for (la, ra) = va
  push!(ex.args,
    :(rE[$(findfirst(x -> isequal(ra,x.g), S) + (2^k-1)*s.h)] *
       u[$(findfirst(x -> isequal(la,x.g), S) + (2^k-1)*s.h), $i]))
  end

  ex = :(p[$(k+1+s.h*(k+1))] * $ex)

  return ex
end

"""
    ext_expr(s    ::Sgh,
             S    ::Array{Sgh,1},
             ns   ::Int64,
             k    ::Int64,
             h    ::Int64,
             model::NTuple{3,Bool})

Return expression for extinction.
"""
function ext_expr(s    ::Sgh,
                  S    ::Array{Sgh,1},
                  ns   ::Int64,
                  k    ::Int64,
                  h    ::Int64,
                  model::NTuple{3,Bool},
                  flow ::Bool)

  wu = flow ? 0 : ns

  m2s = model[1]*k*h

  if isone(length(s.g))
    if model[2]
      ex = Expr(:call, :+)
      for i in s.g push!(ex.args, :(eaft[$(m2s + k*s.h + i)])) end
    else
      ex = Expr(:call, :+)
      for i in s.g push!(ex.args, :(p[$((k+1)*h + k*s.h + i)])) end
    end
  else
    if model[2]
      ex = Expr(:call, :+)
      for i = s.g
        push!(ex.args, :(eaft[$(m2s + k*s.h + i)] *
                         u[$(findfirst(x -> isSghequal(x,
                                            Sgh(setdiff(s.g, i),s.h)), S) + wu)]))
      end
    else
      ex = Expr(:call, :+)
      for i = s.g
        push!(ex.args, :(p[$(i + (k+1)*h + s.h*k + k*h*k)] *
                         u[$(findfirst(x -> isSghequal(x,
                                            Sgh(setdiff(s.g, i),s.h)), S) + wu)]))
      end
    end
  end

  return ex
end

"""
    wrsext_expr(si   ::Int64,
                s    ::Sgh,
                ns   ::Int64,
                k    ::Int64,
                model::NTuple{3,Bool},
                flow ::Bool)

Return expression of extinction for within-region speciation.
"""
function wrsext_expr(si   ::Int64,
                     s    ::Sgh,
                     ns   ::Int64,
                     k    ::Int64,
                     model::NTuple{3,Bool},
                     flow ::Bool)

  wu = flow ? 0 : ns

  # if speciation model
  if model[1]
    if isone(length(s.g))
      for i = s.g
        ex = :(eaft[$(i + s.h*k)] * u[$(i + s.h*(2^k-1) + wu)]^2)
      end
    else
      ex = Expr(:call, :+)
      for i = s.g
        push!(ex.args, :(eaft[$(i + s.h*k)] *
                         u[$(i + s.h*(2^k-1) + wu)] * u[$(si + wu)]))
      end
    end
  # if *NOT* speciation model
  else
    if isone(length(s.g))
      ex = :(p[$si] * u[$(si + wu)]^2)
    else
      ex = Expr(:call, :+)
      for i = s.g
        push!(ex.args, :(p[$(i + s.h*(k+1))] *
                         u[$(i + s.h*(2^k-1) + wu)] * u[$(si + wu)]))
      end
    end
  end

  return ex
end

"""
    geohisse_full(du  ::Array{Float64,1},
                  u   ::Array{Float64,1},
                  p   ::Array{Float64,1},
                  t   ::Float64,
                  r   ::Array{Float64,1},
                  eaft::Array{Float64,1},
                  af! ::Function,
                  ::Val{k},
                  ::Val{h},
                  ::Val{ny},
                  ::Val{model}) where {k, h, ny, model}

Creates Covariate GeoHiSSE ODE equation function for `k` areas,
`h` hidden states and `ny` covariates.
"""
@generated function geohisse_full(du  ::Array{Float64,1},
                                  u   ::Array{Float64,1},
                                  p   ::Array{Float64,1},
                                  t   ::Float64,
                                  r   ::Array{Float64,1},
                                  eaft::Array{Float64,1},
                                  af! ::Function,
                                  ::Val{k},
                                  ::Val{h},
                                  ::Val{ny},
                                  ::Val{model}) where {k, h, ny, model}

  # n states
  ns = (2^k - 1)*h

  # create individual areas subsets
  S = create_states(k, h)

  # start Expression for ODE equations
  eqs = quote end
  popfirst!(eqs.args)

  # add environmental function
  push!(eqs.args, :(af!(t, r)))

  # compute exponential for each β*covariable
  expex = exp_expr(k, h, ny, model)

  for ex in expex.args
    push!(eqs.args, ex)
  end

  for si = Base.OneTo(ns)

    # state range
    s = S[si]

    # length of geographic range
    ls = length(s.g)

    # which single areas do not occur in r
    oa = setdiff(1:k, s.g)

    #=
    likelihoods
    =#

    # no events
    nev = noevents_expr(si, s, ls, oa, k, h, ns, model, false)

    # remove if !isone(lr)
    if ls > 1
      # local extinction
      lex = localext_expr(s, S, k, h, model)

      # between-region speciation
      brs = brspec_expr(s, S, ns, k, false, false)
    end

    # dispersal
    # remove if lr == k
    if ls != k
      dis = dispersal_expr(s, oa, S, ns, k, h, model, false)
    end

    # hidden states transitions
    hid = h > 1 ? hidtran_expr(s, S, ns,k, h, false) : :0.0

    # within-region speciation
    wrs = wrspec_expr(si, s, ns, k, model)

    # push `D` equation to to eqs
    if isone(ls)
      push!(eqs.args,
        :(du[$si] = $nev + $dis + $hid + $wrs))
    elseif ls == k
      push!(eqs.args,
        :(du[$si] = $nev + $lex + $hid + $wrs + $brs))
    else
      push!(eqs.args,
        :(du[$si] = $nev + $lex + $dis + $hid + $wrs + $brs))
    end

    #=
    extinctions
    =#

    # no events
    nev = noevents_expr(si, s, ls, oa, k, h, ns, model, true)

    # extinction
    ext = ext_expr(s, S, ns, k, h, model, false)

    # dispersal and extinction
    if ls != k
      dis = dispersal_expr(s, oa, S, ns, k, h, model, true)
    end

    # hidden states transitions
    hid = h > 1 ? hidtran_expr(s, S, ns,k, h, true) : :0.0

    # within-region extinction
    wrs = wrsext_expr(si, s, ns, k, model, false)

    # between-region extinction
    if ls > 1
      brs = brspec_expr(s, S, ns, k, true, false)
    end

    # push `E` equation to to eqs
    if isone(ls)
      push!(eqs.args,
        :(du[$(si + ns)] = $nev + $ext + $dis + $hid + $wrs))
    elseif ls == k
      push!(eqs.args,
        :(du[$(si + ns)] = $nev + $ext + $hid + $wrs + $brs))
    else
      push!(eqs.args,
        :(du[$(si + ns)] = $nev + $ext + $dis + $hid + $wrs + $brs))
    end

  end

  ## aesthetic touches
  # sort `du`s
  eqs.args[[(end-(ns*2)+1):end...]] =
    eqs.args[append!([(end-(ns*2)+1):2:end...],
            [(end-(ns*2)+2):2:end...])]

  #Core.println(eqs)

  return quote
    @inbounds begin
      $eqs
    end
    return nothing
  end

end

"""
    make_egeohisse(::Val{k},
                   ::Val{h},
                   ::Val{ny},
                   ::Val{model},
                   af!::Function) where {k, h, ny, model}

Make closure for EGeoHiSSE.
"""
function make_egeohisse(::Val{k},
                        ::Val{h},
                        ::Val{ny},
                        ::Val{model},
                        af!::Function) where {k, h, ny, model}

  r    = Array{Float64,1}(undef, ny)
  eaft = Array{Float64,1}(undef,
    model[1]*k*h + model[2]*k*h + model[3]*k*(k-1)*h)

  # make ode function with closure
  ode_fun = (du::Array{Float64,1},
             u::Array{Float64,1},
             p::Array{Float64,1},
             t::Float64) ->
    begin
      geohisse_full(du::Array{Float64,1},
                    u ::Array{Float64,1},
                    p ::Array{Float64,1},
                    t ::Float64,
                    r,eaft,af!,Val(k),Val(h),Val(ny),Val(model))
      return nothing
    end

  return ode_fun
end

"""
    geohisse_E(du  ::Array{Float64,1},
               u   ::Array{Float64,1},
               p   ::Array{Float64,1},
               t   ::Float64,
               r   ::Array{Float64,1},
               eaft::Array{Float64,1},
               af! ::Function,
               ::Val{k},
               ::Val{h},
               ::Val{ny},
               ::Val{model}) where {k, h, ny, model}

Creates Covariate GeoHiSSE Extinction ODE equation for `k` areas,
`h` hidden states and `ny` covariates for the specific `model`.
"""
@generated function geohisse_E(du  ::Array{Float64,1},
                               u   ::Array{Float64,1},
                               p   ::Array{Float64,1},
                               t   ::Float64,
                               r   ::Array{Float64,1},
                               eaft::Array{Float64,1},
                               af! ::Function,
                               ::Val{k},
                               ::Val{h},
                               ::Val{ny},
                               ::Val{model}) where {k, h, ny, model}

  # n states
  ns = (2^k - 1)*h

  # create individual areas subsets
  S = create_states(k, h)

  # start Expression for ODE equations
  eqs = quote end
  popfirst!(eqs.args)

  # add environmental function
  push!(eqs.args, :(af!(t, r)))

  # compute exponential for each β*covariable
  expex = exp_expr(k, h, ny, model)

  for ex in expex.args
    push!(eqs.args, ex)
  end

  for si = Base.OneTo(ns)

    # state range
    s = S[si]

    # length of geographic range
    ls = length(s.g)

    # which single areas do not occur in r
    oa = setdiff(1:k, s.g)

    # no events
    nev = noevents_expr(si, s, ls, oa, k, h, ns, model, false)

    # extinction
    ext = ext_expr(s, S, ns, k, h, model, true)

    # dispersal and extinction
    if ls != k
      dis = dispersal_expr(s, oa, S, ns, k, h, model, false)
    end

    # hidden states transitions
    hid = h > 1 ? hidtran_expr(s, S, ns, k, h, false) : :0.0

    # within-region extinction
    wrs = wrsext_expr(si, s, ns, k, model, true)

    # between-region extinction
    if ls > 1
      brs = brspec_expr(s, S, ns, k, true, true)
    end

    # push `E` equation to to eqs
    if isone(ls)
      push!(eqs.args,
        :(du[$(si)] = $nev + $ext + $dis + $hid + $wrs))
    elseif ls == k
      push!(eqs.args,
        :(du[$(si)] = $nev + $ext + $hid + $wrs + $brs))
    else
      push!(eqs.args,
        :(du[$(si)] = $nev + $ext + $dis + $hid + $wrs + $brs))
    end

  end

  # print for checking
  #Core.println(eqs)

  return quote
    @inbounds begin
      $eqs
    end
    return nothing
  end

end

"""
    make_egeohisse_E(::Val{k},
                     ::Val{h},
                     ::Val{ny},
                     ::Val{model},
                     af!::Function) where {k, h, ny, model}

Make closure for Extinction function of EGeoHiSSE.
"""
function make_egeohisse_E(::Val{k},
                        ::Val{h},
                        ::Val{ny},
                        ::Val{model},
                        af!::Function) where {k, h, ny, model}

  r    = Array{Float64,1}(undef, ny)
  eaft = Array{Float64,1}(undef,
    model[1]*k*h + model[2]*k*h + model[3]*k*(k-1)*h)

  # make ode function with closure
  ode_fun = (du::Array{Float64,1},
             u::Array{Float64,1},
             p::Array{Float64,1},
             t::Float64) ->
    begin
      geohisse_E(du::Array{Float64,1},
                 u ::Array{Float64,1},
                 p ::Array{Float64,1},
                 t ::Float64,
                 r, eaft, af!, Val(k), Val(h), Val(ny), Val(model))
      return nothing
    end

  return ode_fun
end

"""
    geohisse_M(du  ::Array{Float64,1},
               u   ::Array{Float64,1},
               p   ::Array{Float64,1},
               t   ::Float64,
               r   ::Array{Float64,1},
               eaft::Array{Float64,1},
               af! ::Function,
               ::Val{k},
               ::Val{h},
               ::Val{ny},
               ::Val{model}) where {k, h, ny, model}

Creates Covariate GeoHiSSE Likelihood ODE equation function for `k` areas,
`h` hidden states and `ny` covariates in Matrix form.
"""
@generated function geohisse_M(du  ::Array{Float64,2},
                               u   ::Array{Float64,2},
                               pp  ::Array{Array{Float64,1},1},
                               t   ::Float64,
                               r   ::Array{Float64,1},
                               eaft::Array{Float64,1},
                               rE  ::Array{Float64,1},
                               af! ::Function,
                               afE!::Function,
                               idxl::Int64,
                               ::Val{k},
                               ::Val{h},
                               ::Val{ny},
                               ::Val{model}) where {k, h, ny, model}

  # n states
  ns = (2^k - 1)*h

  # create individual areas subsets
  S = create_states(k, h)

  # start Expression for ODE equations
  eqs = quote end
  popfirst!(eqs.args)

  # add parameter subset
  push!(eqs.args, :(p = pp[idxl]))

  # add extinction function
  push!(eqs.args, :(afE!(t, rE, pp)))

  # add environmental function
  push!(eqs.args, :(af!(t, r)))

  # compute exponential for each β*covariable
  expex = exp_expr(k, h, ny, model)

  for ex in expex.args
    push!(eqs.args, ex)
  end

  for i in Base.OneTo(ns), si = Base.OneTo(ns)

    # state range
    s = S[si]

    # length of geographic range
    ls = length(s.g)

    # which single areas do not occur in r
    oa = setdiff(1:k, s.g)

    # no events
    nev = noevents_expr(si, s, ls, oa, k, h, ns, i, model)

    # remove if !isone(lr)
    if ls > 1
      # local extinction
      lex = localext_expr(s, S, k, h, i, model)

      # between-region speciation
      brs = brspec_expr(s, S, ns, k, i)
    end

    # dispersal
    # remove if lr == k
    if ls != k
      dis = dispersal_expr(s, oa, S, ns, k, h, i, model)
    end

    # hidden states transitions
    hid = h > 1 ? hidtran_expr(s, S, ns,k, h, i) : :0.0

    # within-region speciation
    wrs = wrspec_expr(si, s, ns, k, i, model)

    # push `D` equation to to eqs
    if isone(ls)
      push!(eqs.args,
        :(du[$si, $i] = $nev + $dis + $hid + $wrs))
    elseif ls == k
      push!(eqs.args,
        :(du[$si, $i] = $nev + $lex + $hid + $wrs + $brs))
    else
      push!(eqs.args,
        :(du[$si, $i] = $nev + $lex + $dis + $hid + $wrs + $brs))
    end

  end

  ## aesthetic touches
  # sort `du`s

  #Core.println(eqs)

  return quote
    @inbounds begin
      $eqs
    end
    return nothing
  end

end

"""
    make_egeohisse_M(::Val{k},
                     ::Val{h},
                     ::Val{ny},
                     ::Val{model},
                     af!::Function) where {k, h, ny, model}

Make closure for EGeoHiSSE given `E(t)` and in matrix form.
"""
function make_egeohisse_M(::Val{k},
                          ::Val{h},
                          ::Val{ny},
                          ::Val{model},
                          af! ::Function,
                          afE!::Function,
                          idxl::Int64) where {k, h, ny, model}

  ns = (2^k - 1)*h

  r    = Array{Float64,1}(undef, ny)
  eaft = Array{Float64,1}(undef,
    model[1]*k*h + model[2]*k*h + model[3]*k*(k-1)*h)
  rE  = Array{Float64,1}(undef,ns)

  # make ode function with closure
  ode_fun = (du::Array{Float64,2},
             u ::Array{Float64,2},
             pp::Array{Array{Float64,1},1},
             t ::Float64) ->
    begin
      geohisse_M(du::Array{Float64,2},
                 u ::Array{Float64,2},
                 pp::Array{Array{Float64,1},1},
                 t ::Float64,
                 r, eaft, rE, af!, afE!, idxl,
                 Val(k), Val(h), Val(ny), Val(model))
      return nothing
    end

  return ode_fun
end


# singleton for indicating if no default arguments are present
struct NoDefault end
const NO_DEFAULT = NoDefault()

"""
    @addlogprob!(ex)

Add the result of the evaluation of `ex` to the joint log probability.
"""
macro addlogprob!(ex)
    return quote
        acclogp!($(esc(:(__varinfo__))), $(esc(ex)))
    end
end

"""
    getargs_dottilde(x)

Return the arguments `L` and `R`, if `x` is an expression of the form `L .~ R` or
`(~).(L, R)`, or `nothing` otherwise.
"""
getargs_dottilde(x) = nothing
function getargs_dottilde(expr::Expr)
    return MacroTools.@match expr begin
        (.~)(L_, R_) => (L, R)
        (~).(L_, R_) => (L, R)
        # Julia 1.6: see https://github.com/TuringLang/Turing.jl/issues/1525
        (L_ .~ R_) => (L, R)
        x_ => nothing
    end
end

"""
    getargs_tilde(x)

Return the arguments `L` and `R`, if `x` is an expression of the form `L ~ R`, or `nothing`
otherwise.
"""
getargs_tilde(x) = nothing
function getargs_tilde(expr::Expr)
    return MacroTools.@match expr begin
        (~)(L_, R_) => (L, R)
        x_ => nothing
    end
end

function to_namedtuple_expr(syms, vals=syms)
    length(syms) == 0 && return :(NamedTuple())

    names_expr = Expr(:tuple, QuoteNode.(syms)...)
    vals_expr = Expr(:tuple, vals...)
    return :(NamedTuple{$names_expr}($vals_expr))
end

#####################################################
# Helper functions for vectorize/reconstruct values #
#####################################################

vectorize(d::UnivariateDistribution, r::Real) = [r]
vectorize(d::MultivariateDistribution, r::AbstractVector{<:Real}) = copy(r)
vectorize(d::MatrixDistribution, r::AbstractMatrix{<:Real}) = copy(vec(r))

# NOTE:
# We cannot use reconstruct{T} because val is always Vector{Real} then T will be Real.
# However here we would like the result to be specifric type, e.g. Array{Dual{4,Float64}, 2},
# otherwise we will have error for MatrixDistribution.
# Note this is not the case for MultivariateDistribution so I guess this might be lack of
# support for some types related to matrices (like PDMat).
reconstruct(d::UnivariateDistribution, val::AbstractVector) = val[1]
reconstruct(d::MultivariateDistribution, val::AbstractVector) = copy(val)
function reconstruct(d::MatrixDistribution, val::AbstractVector)
    return reshape(copy(val), size(d))
end
function reconstruct!(r, d::Distribution, val::AbstractVector)
    return reconstruct!(r, d, val)
end
function reconstruct!(r, d::MultivariateDistribution, val::AbstractVector)
    r .= val
    return r
end
function reconstruct(d::Distribution, val::AbstractVector, n::Int)
    return reconstruct(d, val, n)
end
function reconstruct(d::UnivariateDistribution, val::AbstractVector, n::Int)
    return copy(val)
end
function reconstruct(d::MultivariateDistribution, val::AbstractVector, n::Int)
    return copy(reshape(val, size(d)[1], n))
end
function reconstruct(d::MatrixDistribution, val::AbstractVector, n::Int)
    tmp = reshape(val, size(d)[1], size(d)[2], n)
    orig = [tmp[:, :, i] for i in 1:size(tmp, 3)]
    return orig
end
function reconstruct!(r, d::Distribution, val::AbstractVector, n::Int)
    return reconstruct!(r, d, val, n)
end
function reconstruct!(r, d::MultivariateDistribution, val::AbstractVector, n::Int)
    r .= val
    return r
end

# Uniform random numbers with range 4 for robust initializations
# Reference: https://mc-stan.org/docs/2_19/reference-manual/initialization.html
randrealuni(rng::Random.AbstractRNG) = 4 * rand(rng) - 2
randrealuni(rng::Random.AbstractRNG, args...) = 4 .* rand(rng, args...) .- 2

const Transformable = Union{
    PositiveDistribution,
    UnitDistribution,
    TransformDistribution,
    SimplexDistribution,
    PDMatDistribution,
}
istransformable(dist) = false
istransformable(::Transformable) = true

#################################
# Single-sample initialisations #
#################################

inittrans(rng, dist::UnivariateDistribution) = invlink(dist, randrealuni(rng))
function inittrans(rng, dist::MultivariateDistribution)
    return invlink(dist, randrealuni(rng, size(dist)[1]))
end
inittrans(rng, dist::MatrixDistribution) = invlink(dist, randrealuni(rng, size(dist)...))

################################
# Multi-sample initialisations #
################################

inittrans(rng, dist::UnivariateDistribution, n::Int) = invlink(dist, randrealuni(rng, n))
function inittrans(rng, dist::MultivariateDistribution, n::Int)
    return invlink(dist, randrealuni(rng, size(dist)[1], n))
end
function inittrans(rng, dist::MatrixDistribution, n::Int)
    return invlink(dist, [randrealuni(rng, size(dist)...) for _ in 1:n])
end

#######################
# Convenience methods #
#######################
collectmaybe(x) = x
collectmaybe(x::Base.AbstractSet) = collect(x)


using FunctionalCollections
using FluxionArrays

function rearrange_dims(arr, arr_dims, target_dims)
    """
    Contracts:
    - elements of `arr_dims` and `target_dims` support the `==` operator
    - all elements of `arr_dims` are different
    - each element of `arr_dims` is present in `target_dims` exactly once
    - `arr` needs to support `reshape()`, `permutedims()` and `size()` methods.
    """
    if length(target_dims) == 0 || length(arr_dims) == 0
        return arr
    end

    old_shape = [size(arr, dim_num) for dim_num in 1:length(arr_dims)]

    # transpose the array so that the order of its dimensions
    # is the same as in `target_dims`
    dims_correspondence = [
        (findfirst(target_dims .== dim), old_dim_num)
        for (old_dim_num, dim) in enumerate(arr_dims)]
    new_dims_order = [old_dim_num for (new_dim_num, old_dim_num) in sort(dims_correspondence)]

    # Don't permute if the required order is already achieved
    if new_dims_order != collect(1:length(new_dims_order))
        arr = permutedims(arr, new_dims_order)
        new_shape = old_shape[new_dims_order]
        new_dims = arr_dims[new_dims_order]
    else
        new_shape = old_shape
        new_dims = arr_dims
    end

    # We may need to reshape the array adding new array dimensions (of size 1)
    # in place of unused array dimensions.
    # Currently just checking for the difference in shape lengths;
    # we probably don't need to add dummy dimensions if broadcasting works for us,
    # but that's harder to check
    if length(new_shape) != length(target_dims)

        new_dim_num = 1
        target_shape = Array{Any, 1}(length(target_dims))
        for target_dim_num = 1:length(target_dims)
            if (new_dim_num <= length(new_dims)
                    && target_dims[target_dim_num] == new_dims[new_dim_num])
                target_shape[target_dim_num] = new_shape[new_dim_num]
                new_dim_num += 1
            else
                target_shape[target_dim_num] = 1
            end
        end

        arr = reshape(arr, target_shape...)
    end

    return arr
end

function ensure_shape(arr, arr_shape, target_shape)

    if arr_shape == target_shape
        return arr
    end

    if length(arr_shape) < length(target_shape)
        extend_dims = length(target_shape) - length(arr_shape)
        extend_shape = [1 for i in 1:extend_dims]
        arr_shape = [arr_shape..., extend_shape...]
        arr = reshape(arr, arr_shape...)
    end

    outer = [
        (arr_dim == 1 && target_dim > 1) ? target_dim : 1
        for (arr_dim, target_dim) in zip(arr_shape, target_shape)
    ]

    repeat(arr, outer=outer)
end

function map_accum{T}(func, state, vals::Array{T, 1}, args...)
    new_vals = Array{Any, 1}(length(vals))
    for (i, val) in enumerate(vals)
        state, val = func(state, val, args...)
        new_vals[i] = val
    end
    state, new_vals
end

immutable ArrayExprBuildingState
    unknowns :: PersistentHashMap{FieldExpr, Unknown}
    random_states :: Nullable{Pair{Unknown, ArrayExpr}}
    random_variables :: PersistentHashMap{RandomVariable, ArrayExpr}
end

ArrayExprBuildingState() = ArrayExprBuildingState(
    PersistentHashMap{FieldExpr, Unknown}(),
    nothing,
    PersistentHashMap{RandomVariable, ArrayExpr}())

function add_if_missing(gen_val, state::ArrayExprBuildingState, key::FieldExpr)
    if haskey(state.unknowns, key)
        state, state.unknowns[key]
    else
        val = gen_val(key)
        new_unknowns = assoc(state.unknowns, key, val)
        new_state = ArrayExprBuildingState(
            new_unknowns, state.random_states, state.random_variables)
        new_state, val
    end
end

function set_random_state(state::ArrayExprBuildingState, rs::Unknown)
    ArrayExprBuildingState(
        state.unknowns, Pair{Unknown, ArrayExpr}(rs, rs), state.random_variables)
end

function update_random_state(state::ArrayExprBuildingState, rs::ArrayExpr)
    input_rs, output_rs = get(state.random_states)
    ArrayExprBuildingState(
        state.unknowns, Pair{Unknown, ArrayExpr}(input_rs, rs), state.random_variables)
end

function add_random_variable(state::ArrayExprBuildingState, fexpr::RandomVariable, aexpr::ArrayExpr)
    ArrayExprBuildingState(
        state.unknowns, state.random_states,
        assoc(state.random_variables, fexpr, aexpr))
end

function clear_random_variables(state::ArrayExprBuildingState)
    ArrayExprBuildingState(
        state.unknowns, state.random_states,
        PersistentHashMap{RandomVariable, ArrayExpr}())
end

function set_random_variables(state::ArrayExprBuildingState, random_variables)
    ArrayExprBuildingState(
        state.unknowns, state.random_states,
        random_variables)
end

random_states_set(state::ArrayExprBuildingState) = !isnull(state.random_states)

random_states(state::ArrayExprBuildingState) = get(state.random_states)
output_random_state(state::ArrayExprBuildingState) = random_states(state)[2]

function unknown_array_like(x::UnknownDimension)
    return unknown_array(shape=[], eltype=Float64)
end

function unknown_array_like(x::UnknownField)
    return unknown_array(
        shape=[Nullable{Int64}(length(vertices(grid(dim)))) for dim in known_dimensions(x)],
        eltype=Float64)
end

_as_array_expr(state, x::Number, target_dimensions) = state, lazy_array(x)
_as_array_expr(state, x::KnownDimension, target_dimensions) =
    state, rearrange_dims(lazy_array(vertices(grid(x))), [x], target_dimensions)
_as_array_expr(state, x::TransformedDimension, target_dimensions) =
    state, rearrange_dims(lazy_array(vertices(grid(x))), [x], target_dimensions)
_as_array_expr(state, x::KnownField, target_dimensions) =
    state, rearrange_dims(x.data, known_dimensions(x), target_dimensions)
_as_array_expr(state, x::UnknownDimension, target_dimensions) =
    add_if_missing(unknown_array_like, state, x)

function _as_array_expr(state, x::UnknownField, target_dimensions)
    state, arr_var = add_if_missing(unknown_array_like, state, x)
    state, rearrange_dims(arr_var, known_dimensions(x), target_dimensions)
end

function _as_array_expr(state, x::Fourier, target_dimensions)

    dims_dict = map(reverse, x.dimensions)

    fft_axes = []
    target_dimensions_inner = []
    for (i, dim) in enumerate(target_dimensions)
        if haskey(dims_dict, dim)
            push!(fft_axes, i)
            push!(target_dimensions_inner, dims_dict[dim])
        else
            push!(target_dimensions_inner, dim)
        end
    end

    state, expr = _as_array_expr(state, x.fexpr, target_dimensions_inner)

    if x.inverse
        result = ifft(expr, fft_axes)
    else
        result = fft(expr, fft_axes)
    end
    state, result
end

function _as_array_expr(state, fexpr::Sample, target_dimensions)
    target_dimensions = [
        x === STOCHASTIC_DIMENSION ? KnownStochasticDimension(fexpr.trajectories) : x
        for x in target_dimensions]

    # Random variables inside a sample() call are independent from the rest
    outer_rvs = state.random_variables
    inner_state = clear_random_variables(state)
    state, aexpr = _as_array_expr(inner_state, fexpr.fexpr, target_dimensions)
    state = set_random_variables(state, outer_rvs)
    state, aexpr
end

function _as_array_expr(state, fexpr::StochasticMean, target_dimensions)
    trs = trajectories_number(fexpr.fexpr)
    target_dimensions = [target_dimensions..., KnownStochasticDimension(trs)]
    state, aexpr = _as_array_expr(state, fexpr.fexpr, target_dimensions)
    state, sum(aexpr, length(target_dimensions)) ./ trs
end

immutable KnownStochasticDimension <: KnownDimension
    id :: Symbol
    trajectories :: Int64
    KnownStochasticDimension(trajectories) = new(gensym(), trajectories)
end

function _as_array_expr(state, fexpr::RandomVariable, target_dimensions)

    if haskey(state.random_variables, fexpr)
        return state, state.random_variables[fexpr]
    end

    if fexpr.distribution == :random_normal
        sdims = find(dim -> typeof(dim) == KnownStochasticDimension, target_dimensions)
        @assert length(sdims) == 1
        random_dimensions = [fexpr.dimensions..., target_dimensions[sdims[1]]]

        shape = [
            typeof(x) == KnownStochasticDimension ? x.trajectories : length(vertices(grid(x)))
            for x in random_dimensions]

        if random_states_set(state)
            rs = output_random_state(state)
        else
            rs = unknown_random_state()
            state = set_random_state(state, rs)
        end

        # TODO: we need to either use the known mean/std here,
        # or expose the basic random_normal(rs, shape...) from Arrays.
        res, new_rs = FluxionArrays.random_normal(rs, 0, 1, shape...)

        res = rearrange_dims(res, random_dimensions, target_dimensions)

        state = update_random_state(state, new_rs)
        state = add_random_variable(state, fexpr, res)

        return state, res
    end

end

function _as_array_expr(state, fexpr::FieldExprNode, target_dimensions)
    const broadcasted_operators = Set([:+, :-, :*, :/, :^])

    op = fexpr.head
    state, args = map_accum(_as_array_expr, state, fexpr.args, target_dimensions)
    if in(op, broadcasted_operators)
        res = broadcast(getfield(Base, op), args...)
    else
        res = getfield(Base, op)(args...)
    end
    state, res
end

function as_array_expr(state, fexpr::FieldExpr, target_dimensions)

    state, aexpr = _as_array_expr(state, fexpr, target_dimensions)

    if is_stochastic(fexpr)
        tr_num = trajectories_number(fexpr)
    else
        tr_num = nothing
    end
    target_shape = [
        dim === STOCHASTIC_DIMENSION ? tr_num : length(vertices(grid(dim)))
        for dim in target_dimensions]

    shape = size(aexpr)
    @assert !isa(shape, ArrayExpr)
    shape = collect(shape)

    aexpr = ensure_shape(aexpr, shape, target_shape)

    state, aexpr
end

immutable WrappedArrayFunction
    array_function
    uses_random_state :: Bool
end

function wrapped_array_function(returns, parameters)

    array_returns = []

    state = ArrayExprBuildingState()

    for (fexpr, result_dimensions) in returns
        state, aexpr = as_array_expr(state, fexpr, result_dimensions)
        push!(array_returns, aexpr)
    end

    array_parameters = [
        get(state.unknowns, parameter, unknown_array())
        for parameter in parameters]

    uses_random_state = !isnull(state.random_states)

    if uses_random_state
        input_rs, output_rs = get(state.random_states)
        array_parameters = [input_rs, array_parameters...]
        array_returns = [output_rs, array_returns...]
    end

    afunc = array_function(array_returns, array_parameters)

    WrappedArrayFunction(afunc, uses_random_state)
end

immutable WrappedBackendArrayFunction
    backend_array_function
    uses_random_state :: Bool
end

function compile_wrapped_array_function(backend::ArrayBackend, afunc::WrappedArrayFunction)
    WrappedBackendArrayFunction(
        compile_array_function(backend, afunc.array_function), afunc.uses_random_state)
end

import FluxionArrays: array_backend

array_backend(wbafunc::WrappedBackendArrayFunction) = array_backend(wbafunc.backend_array_function)

function (wcafunc::WrappedBackendArrayFunction)(args...; seed=nothing)
    if wcafunc.uses_random_state
        backend = array_backend(wcafunc.backend_array_function)
        rs = random_state(backend, seed=seed)
        rs, result = wcafunc.backend_array_function(rs, args...)
        return result
    else
        wcafunc.backend_array_function(args...)
    end
end

immutable FieldFunction
    wrapped_array_function
    output_dimensions
end

function field_function(returns, parameters)
    wafunc = wrapped_array_function(returns, parameters)
    output_dimensions = [rd for (_, rd) in returns]
    FieldFunction(wafunc, output_dimensions)
end

immutable BackendFieldFunction
    wrapped_backend_array_function :: WrappedBackendArrayFunction
    output_dimensions
end

function compile_field_function(backend::ArrayBackend, ffunc::FieldFunction)
    BackendFieldFunction(
        compile_wrapped_array_function(backend, ffunc.wrapped_array_function),
        ffunc.output_dimensions)
end

unwrap_field(x::Number) = x
unwrap_field(x::KnownField) = x.data

function (bffunc::BackendFieldFunction)(args...; seed=nothing)

    args = map(unwrap_field, args)

    if bffunc.wrapped_backend_array_function.uses_random_state
        backend = array_backend(bffunc.wrapped_backend_array_function)
        rs = random_state(backend, seed=seed)
        results = bcfunc.wrapped_backend_array_function(rs, args...)
        results = results[2:end]
    else
        results = bffunc.wrapped_backend_array_function(args...)
        if length(bffunc.output_dimensions) == 1
            results = [results]
        end
    end

    wrapped_results = []
    for (result, odims) in zip(results, bffunc.output_dimensions)
        push!(wrapped_results, known_field(result, odims))
    end

    if length(wrapped_results) == 1
        return wrapped_results[1]
    else
        return tuple(wrapped_results)
    end
end


function ImageSegmentation(
		blocks::Tuple{<:Image{N},<:Mask{N}},
        data=nothing;
		size=ntuple(i -> 128, N),
		aug_projections=DataAugmentation.Identity(),
		aug_image=DataAugmentation.Identity(),
        C=RGB{N0f8},
        computestats=false,
	) where N
	return SupervisedTask(
		blocks,
		(
			ProjectiveTransforms(size; augmentations=aug_projections),
			getimagepreprocessing(data, computestats; C=C, augmentations=aug_image),
        	OneHot()
		)
	)
end

"""
    ImageSegmentation(size, classes; kwargs...)

Learning task for image segmentation. Images are
resized to `size` and a class is predicted for every pixel.

## Keyword arguments

- `computestats = false`: Whether to compute image statistics on dataset `data` or use
    default ImageNet stats.
- `aug_projections = `[`DataAugmentation.Identity`](#): augmentation to apply during
  [`ProjectiveTransforms`](#) (resizing and cropping)
- `aug_image = `[`DataAugmentation.Identity`](#): pixel-level augmentation to apply during
  [`ImagePreprocessing`](#)
- `C = RGB{N0f8}`: Color type images are converted to before further processing. Use `Gray{N0f8}`
    for grayscale images.
"""
function ImageSegmentation(size::NTuple{N,Int}, classes::AbstractVector; kwargs...) where N
    blocks = (Image{N}(), Mask{N}(classes))
    return ImageSegmentation(blocks; size=size, kwargs...)
end

registerlearningtask!(FASTAI_METHOD_REGISTRY, ImageSegmentation, (Image, Mask))

# ## Tests

@testset "ImageSegmentation [task]" begin
    @testset "2D" begin
        task = ImageSegmentation((16, 16), 1:4)
        testencoding(getencodings(task), getblocks(task).sample)
        FastAI.checktask_core(task)
        @test_nowarn tasklossfn(task)
        @test_nowarn taskmodel(task, Models.xresnet18())
        @testset "Show backends" begin
            @testset "ShowText" begin
                FastAI.test_task_show(task, ShowText(Base.DevNull()))
            end
        end
    end
    @testset "3D" begin
        task = SupervisedTask(
            (Image{3}(), Mask{3}(1:4)),
            (
                ProjectiveTransforms((16, 16, 16), inferencefactor=8),
                ImagePreprocessing(),
                FastAI.OneHot()
            )
        )
        testencoding(getencodings(task), getblocks(task).sample)
        FastAI.checktask_core(task)
        @test_nowarn tasklossfn(task)
    end

end


# API definition file for data layer abstraction.
# The idea is for a one line call to change from internal Graphs.jl to CloudGraphs.jl
# switch for as data layer.

type DataLayerAPI
  addvertex!::Function
  getvertex::Function
  # setupvertgraph!::Function
  # setupfncvertgraph!::Function
  makeaddedge!::Function
  getedge::Function
  outneighbors::Function
  updatevertex!::Function
  updateedge!::Function
  deletevertex!::Function
  deleteedge!::Function
  cgEnabled::Bool
end

# global dlapi

dlapi = DataLayerAPI(addGraphsVert!,          # addvertex
                     getVertNode,             # getvertex
                     makeAddEdge!,            # makeaddedge
                     graphsGetEdge,           # getedge
                     graphsOutNeighbors,      # outneighbors
                     updateFullVertData!,     # updatevertex!
                     +,                       # updateedge!
                     graphsDeleteVertex!,     # deletevertex!
                     +,                       # deleteedge!
                     false )

# Remember 3rd party users interact with
# addNode!
# addFactor!
# prepBatchTree!
# inferOverTree!
#
## Visualization functions -- Visualization may well be separated or abstracted out
# writeGraphPdf
# drawHorBeliefsList


# all of these must be added as a test dependency!!
using BenchmarkTools
using Random
using Mycelia

const SUITE = BenchmarkGroup()

SUITE["utf8"] = BenchmarkGroup(["string", "unicode"])
teststr = String(join(rand(MersenneTwister(1), 'a':'d', 10^4)))
SUITE["utf8"]["replace"] = @benchmarkable replace($teststr, "a" => "b")
SUITE["utf8"]["join"] = @benchmarkable join($teststr, $teststr)
SUITE["utf8"]["plots"] = BenchmarkGroup()

SUITE["trigonometry"] = BenchmarkGroup(["math", "triangles"])
SUITE["trigonometry"]["circular"] = BenchmarkGroup()
for f in (sin, cos, tan)
    for x in (0.0, pi)
        SUITE["trigonometry"]["circular"][string(f), x] = @benchmarkable ($f)($x)
    end
end

SUITE["trigonometry"]["hyperbolic"] = BenchmarkGroup()
for f in (sin, cos, tan)
    for x in (0.0, pi)
        SUITE["trigonometry"]["hyperbolic"][string(f), x] = @benchmarkable ($f)($x)
    end
end

# any time that I update the benchmarks or want to modify the tunings,
# just delete the tune.json file and then re run this script
# https://github.com/JuliaCI/BenchmarkTools.jl/blob/master/doc/manual.md#caching-parameters
tuning_file = "$(@__DIR__)/tune.json"
if !isfile(tuning_file)

    # tune the suite to configure benchmark parameters
    tune!(SUITE);

    # save the suite's parameters using a thin wrapper
    # over JSON (this wrapper maintains compatibility
    # across BenchmarkTools versions)
    BenchmarkTools.save("tune.json", params(SUITE));
end


## Overlay for animation
@with_kw struct BeliefOverlay <: SceneOverlay
    # belief::SingleOCFBelief = SingleOCFBelief()
     b_dict::Dict{Int64, SingleOCFBelief} = Dict{Int64, SingleOCFBelief}()
     ego_vehicle::Vehicle = Vehicle(VehicleState(VecSE2(0., 0., 0.), 0.), VehicleDef(), 1)
     color::Colorant = RGBA(0.0, 0.0, 1.0, 0.2)
 end

 function AutoViz.render!(rendermodel::RenderModel, overlay::BeliefOverlay, scene::Scene, roadway::R) where R

     b_dict = overlay.b_dict
     ego = overlay.ego_vehicle

     for b_id in keys(b_dict)
         for (s, prob) in weighted_iterator( b_dict[b_id])

             ped = Vehicle(VehicleState(VecSE2(s.ped_s+ego.def.length/2+AutomotivePOMDPs.PEDESTRIAN_DEF.width/2, s.ped_T, s.ped_theta), 0.), AutomotivePOMDPs.PEDESTRIAN_DEF, 1)
             if prob > 1e-6
                 prob_color = RGBA(0.0, 0.0, 1.0, prob*10)
                 add_instruction!(rendermodel, render_vehicle, (ego.state.posG.x+ped.state.posG.x, ego.state.posG.y+ped.state.posG.y, ego.state.posG.θ + ped.state.posG.θ, ped.def.length, ped.def.width, prob_color))
             end
         end
     end

     return rendermodel
 end

function AutomotivePOMDPs.animate_record(scenes::Vector{Scene},dt::Float64, env::CrosswalkEnv, sensor::GaussianSensor, sensor_o::Vector{Vector{Vehicle}}, risk::Vector{Float64}, belief_dict::Vector{Dict{Int64, SingleOCFBelief}}, ego_vehicle::Vector{Vehicle}, action_pomdp::Vector{SingleOCFAction}, prediction::Vector{Array{Float64}}, cam=FitToContentCamera(0.0))

    duration =length(scenes)*dt::Float64
    fps = Int(1/dt)
    function render_rec(t, dt)

        frame_index = Int(floor(t/dt)) + 1
        text_to_visualize = [string("v-ego: ", round(ego_vehicle[frame_index].state.v*3.6,digits=1), " km/h" ,
                                " y: ", ego_vehicle[frame_index].state.posG.y, " m",
                                " ax: ", action_pomdp[frame_index].acc, " m/s²")]
        sensor_overlay = GaussianSensorOverlay(sensor=sensor, o=sensor_o[frame_index])
        occlusion_overlay = OcclusionOverlay(obstacles=env.obstacles)

        text_overlay = TextOverlay(text=text_to_visualize,pos=VecE2(80.,10.),incameraframe=true,color=colorant"white",font_size=20)
        belief_overlay = BeliefOverlay(b_dict=belief_dict[frame_index], ego_vehicle=ego_vehicle[frame_index])
      #  max_speed = 14.0
      #  histogram_overlay = HistogramOverlay(pos = VecE2(65.0, 10.0), val=ego_vehicle[frame_index].state.v/max_speed, label="v speed")

        if length(prediction) > 0
            prediction_overlay = EmergencyBrakingSystem.PredictionOverlay(prediction=prediction[frame_index])
            return AutoViz.render(scenes[frame_index], env, [prediction_overlay, belief_overlay, sensor_overlay, occlusion_overlay, text_overlay], cam=cam)
        else
            return AutoViz.render(scenes[frame_index], env, [belief_overlay, sensor_overlay, occlusion_overlay, text_overlay], cam=cam)
        end
    end
    return duration, fps, render_rec
end


using BinaryBuilder

# Collection of sources required to build FastTransforms
name = "FastTransforms"
version = v"0.3.2"
sources = [
    ArchiveSource("https://github.com/MikaelSlevinsky/FastTransforms/archive/v$(version).tar.gz",
                  "ef401079751e570d3dcfd4ac3c40e7d7d5858a44ea66c62808cf319a68c223c8"),
]

# Bash recipe for building across all platforms
script = raw"""
cd $WORKSPACE/srcdir/FastTransforms-*
if [[ ${target} == x86_64-w64-mingw32 ]]; then
    export CFLAGS="-O3 -mavx -fno-asynchronous-unwind-tables "
else
    export CFLAGS="-O3 -mavx "
fi
if [[ ${nbits} == 64 ]]; then
    SYMBOL_DEFS=()
    SYMBOLS=(dgemm dtrmm dtrmv dtrsm sgemm strmm strsm ztrmm)
    for sym in ${SYMBOLS[@]}; do
        SYMBOL_DEFS+=("-Dcblas_${sym}=cblas_${sym}64_")
    done
    CFLAGS+=${SYMBOL_DEFS[@]}
    BLAS=openblas64_
else
    BLAS=openblas
fi
make assembly CC=gcc
make lib CC=gcc FT_PREFIX=${prefix} FT_BLAS=${BLAS} FT_FFTW_WITH_COMBINED_THREADS=1
mv -f libfasttransforms.${dlext} ${libdir}
"""

platforms = expand_gfortran_versions(supported_platforms())
platforms = [p for p in platforms if BinaryBuilder.proc_family(p) == :intel]

# The products that we will ensure are always built
products = [
    LibraryProduct("libfasttransforms", :libfasttransforms),
]

# Dependencies that must be installed before this package can be built
dependencies = [
    Dependency("CompilerSupportLibraries_jll"),
    Dependency("FFTW_jll"),
    Dependency("MPFR_jll"),
    Dependency("OpenBLAS_jll"),
]

# Build the tarballs, and possibly a `build.jl` as well.
build_tarballs(ARGS, name, version, sources, script, platforms, products, dependencies)


module Level
    import ..Log4jl

    bitstype 32 EventLevel
    const levels = Dict{Int32,Symbol}(
        Int32(typemax(Int16)) => :ALL,
        Int32(600) => :TRACE,
        Int32(500) => :DEBUG,
        Int32(400) => :INFO,
        Int32(300) => :WARN,
        Int32(200) => :ERROR,
        Int32(100) => :FATAL,
        Int32(0)   => :OFF
    )
    Base.isless(x::EventLevel, y::EventLevel) = isless(Int32(x), Int32(y))
    Base.convert{T<:Integer}(::Type{T}, x::EventLevel) = convert(T, Intrinsics.box(Int32, x))

    for (v, l) in levels
        eval(:(const $l = Intrinsics.box(EventLevel, $v)))
    end

    function Base.print(io::IO, level::EventLevel)
        vlvl = Int32(level)
        if haskey(levels, vlvl)
            print(io, levels[vlvl])
        end
    end
    Base.show(io::IO,x::EventLevel) = print(io, x, "::EventLevel")

    "Create a custom level and generate convenience functions"
    function add(lvl::Symbol, lvlval::Int32)
        lvlup = symbol(uppercase(string(lvl)))
        lvllw = symbol(lowercase(string(lvl)))
        @assert !isdefined(Level, lvlup) "Level $lvlup is already defined"
        eval(Level, :(const $lvlup = Intrinsics.box(EventLevel, $lvlval)))
        levels[lvlval] = lvlup
        isdefined(Base, lvllw) && eval(Log4jl, :(import Base.$lvllw))
        eval(Log4jl, parse("makemethods(:$lvllw, Level.$lvlup)"))
    end
end
typealias LEVEL Nullable{Level.EventLevel}


module Colors

using FixedPointNumbers
using Reexport
using Printf

@reexport using ColorTypes

# TODO: why these types are defined here? Can they move to ColorTypes.jl?
AbstractAGray{C<:AbstractGray,T} = AlphaColor{C,T,2}
AbstractGrayA{C<:AbstractGray,T} = ColorAlpha{C,T,2}

import Base: ==, +, -, *, /
import Base: convert, eltype, hex, isless, range, show, typemin, typemax

# Additional exports, not exported by ColorTypes
export weighted_color_mean,
       hex, @colorant_str,
       protanopic, deuteranopic, tritanopic,
       distinguishable_colors, whitebalance,
       colordiff, DE_2000, DE_94, DE_JPC79, DE_CMC, DE_BFD, DE_AB, DE_DIN99, DE_DIN99d, DE_DIN99o,
       MSC, sequential_palette, diverging_palette, colormap,
       colormatch, CIE1931_CMF, CIE1964_CMF, CIE1931J_CMF, CIE1931JV_CMF

# Early utilities
include("utilities.jl")

# Include other module components
include("conversions.jl")
include("promotions.jl")
include("algorithms.jl")
include("parse.jl")
include("differences.jl")
include("colormaps.jl")
include("display.jl")
include("colormatch.jl")

end # module


type recdata
  component :: Array{String} # provide information of output data component
  recnumber :: Int64 # Number of Receivers
  location :: Array{Float64} # Receiver coordinates
  nlocation :: Array{Int64}
  BDnlocation :: Array{Int64}
  timesamp :: Int64 # Total time sampling Tn
  data :: Array{Float64}
end

function readrec(DataPath::String)
    fid = open(DataPath,"r")
    CpntDict = Dict(1=>"P",2=>"Vx",3=>"Vz",
                    4=>"Txx",5=>"Tzz",6=>"Txz")

        rn = read(fid,Int64,1)
        nCpnt = read(fid,Int64,1)
        NumCpnt = read(fid,Int64,nCpnt)
        component = []
        for i in NumCpnt
            if i == 23
                push!(component,CpntDict[2])
                push!(component,CpntDict[3])
            elseif i == 45
                push!(component,CpntDict[4])
                push!(component,CpntDict[5])
            else push!(component,CpntDict[i])
            end
        end
        nT = read(fid,Int64,1) # Number of Receivers and Total time sampling Tn
        RowNum = length(component)*rn
        ColNum = nT
        loc = read(fid, Float64, 2*rn)
        loc = reshape(loc, rn,2)
        nloc = read(fid,Int64, 2*rn)
        nloc = reshape(nloc,rn,2)
        BDnloc = read(fid,Int64, 2*rn)
        BDnloc = reshape(BDnloc,rn,2)

        data = read(fid,Float64,RowNum*ColNum)
        data = reshape(data,RowNum,ColNum)

        if eof(fid)
            println("Data reading successful!")
            return recdata(Cpnt,rn,loc,nloc,BDnloc,nT,data)
            close(fid)
        elseif !(eof(fid))
            close(fid)
            error("Check data length, not finished!")
        end

    end


"""
    initGigaSOM(data, args...)

Initializes a SOM by random selection from the training data. A generic
overload that works for matrices and DataFrames that can be coerced to
`Matrix{Float64}`. Other arguments are passed to the data-independent
`initGigaSOM`.

Arguments:
- `data`: matrix of data for running the initialization
"""
function initGigaSOM(data::Union{Matrix,DataFrame}, args...; kwargs...)

    d = Matrix{Float64}(data)

    (n, ncol) = size(d)
    means = [sum(d[:, i]) / n for i = 1:ncol]
    sdevs = [sqrt(sum((d[:, i] .- means[i]) .^ 2.0) / n) for i = 1:ncol]

    return initGigaSOM(ncol, means, sdevs, args...; kwargs...)
end

"""
    function initGigaSOM(data::Dinfo,
                         xdim::Int64, ydim::Int64 = xdim;
                         seed=rand(Int), rng=StableRNG(seed))

`initGigaSOM` overload for working with distributed-style `Dinfo`
data. The rest of the arguments is passed to the data-independent
`initGigaSOM`.

Arguments:
- `data`: a `Dinfo` object with the distributed dataset matrix
"""
function initGigaSOM(data::Dinfo, args...; kwargs...)
    ncol = get_val_from(data.workers[1], :(size($(data.val))[2]))
    (means, sdevs) = dstat(data, Vector(1:ncol))

    initGigaSOM(ncol, means, sdevs, args...; kwargs...)
end

"""
    function initGigaSOM(ncol::Int64,
                         means::Vector{Float64}, sdevs::Vector{Float64},
                         xdim::Int64, ydim::Int64 = xdim;
                         seed = rand(Int), rng = StableRNG(seed))

Generate a stable random initial SOM with the random distribution that matches the parameters.

Arguments:
- `ncol`: number of desired data columns
- `means`, `sdevs`: vectors that describe the data distribution, both of size `ncol`
- `xdim`, `ydim`: Size of the SOM
- `seed`: a seed (defaults to random seed from the current default random generator
- `rng`: a random number generator to be used (defaults to a `StableRNG` initialized with the `seed`)

Returns: a new `Som` structure
"""
function initGigaSOM(
    ncol::Int64,
    means::Vector{Float64},
    sdevs::Vector{Float64},
    xdim::Int64,
    ydim::Int64 = xdim;
    seed = rand(UInt),
    rng = StableRNG(seed),
)

    numCodes = xdim * ydim
    grid = gridRectangular(xdim, ydim)

    # Initialize with an unbiased random gaussian with same mean/sd as the data
    # in each dimension
    codes = randn(rng, (numCodes, ncol))
    for col = 1:ncol
        codes[:, col] .*= sdevs[col]
        codes[:, col] .+= means[col]
    end

    return Som(codes = codes, xdim = xdim, ydim = ydim, grid = grid)
end

"""
    trainGigaSOM(
        som::Som,
        dInfo::Dinfo;
        kernelFun::Function = gaussianKernel,
        metric = Euclidean(),
        somDistFun = distMatrix(Chebyshev()),
        knnTreeFun = BruteTree,
        rStart = 0.0,
        rFinal = 0.1,
        radiusFun = expRadius(-5.0),
        epochs = 20,
        eachEpoch = (e, r, som) -> nothing,
    )

# Arguments:
- `som`: object of type Som with an initialised som
- `dInfo`: `Dinfo` object that describes a loaded dataset
- `kernelFun::function`: optional distance kernel; one of (`bubbleKernel, gaussianKernel`)
            default is `gaussianKernel`
- `metric`: Passed as metric argument to the KNN-tree constructor
- `somDistFun`: Function for computing the distances in the SOM map
- `knnTreeFun`: Constructor of the KNN-tree (e.g. from NearestNeighbors package)
- `rStart`: optional training radius. If zero (default), it is computed from the SOM grid size.
- `rFinal`: target radius at the last epoch, defaults to 0.1
- `radiusFun`: Function that generates radius decay, e.g. `linearRadius` or `expRadius(10.0)`
- `epochs`: number of SOM training iterations (default 10)
- `eachEpoch`: a function to call back after each epoch, accepting arguments
  `(epochNumber, radius, som)`. For simplicity, this gets additionally called
  once before the first epoch, with `epochNumber` set to zero.
"""
function trainGigaSOM(
    som::Som,
    dInfo::Dinfo;
    kernelFun::Function = gaussianKernel,
    metric = Euclidean(),
    somDistFun = distMatrix(Chebyshev()),
    knnTreeFun = BruteTree,
    rStart = 0.0,
    rFinal = 0.1,
    radiusFun = expRadius(-5.0),
    epochs = 20,
    eachEpoch = (e, r, som) -> nothing,
)

    # set the default radius
    if rStart == 0.0
        rStart = (som.xdim + som.ydim) / 2
        @debug "The radius has been determined automatically." rStart rFinal
    end

    # get the SOM neighborhood distances
    dm = somDistFun(som.grid)

    result_som = copy(som)
    result_som.codes = copy(som.codes) # prevent rewriting by reference

    eachEpoch(0, rStart, result_som)

    for epoch = 1:epochs
        @debug "Epoch $epoch..."

        numerator, denominator = distributedEpoch(
            dInfo,
            result_som.codes,
            knnTreeFun(Array{Float64,2}(transpose(result_som.codes)), metric),
        )

        r = radiusFun(rStart, rFinal, epoch, epochs)
        @debug "radius: $r"
        if r <= 0
            @error "Sanity check failed: radius must be positive"
            error("Radius check")
        end

        wEpoch = kernelFun(dm, r)
        result_som.codes = (wEpoch * numerator) ./ (wEpoch * denominator)

        eachEpoch(epoch, r, result_som)
    end

    return result_som
end

"""
    trainGigaSOM(som::Som, train;
                 kwargs...)

Overload of `trainGigaSOM` for simple DataFrames and matrices. This slices the
data, distributes them to the workers, and runs normal `trainGigaSOM`. Data is
`unscatter`d after the computation.
"""
function trainGigaSOM(som::Som, train; kwargs...)

    train = Matrix{Float64}(train)

    #this slices the data into parts and and sends them to workers
    dInfo = scatter_array(:GigaSOMtrainDataVar, train, workers())
    som_res = trainGigaSOM(som, dInfo; kwargs...)
    unscatter(dInfo)
    return som_res
end

"""
    doEpoch(x::Array{Float64, 2}, codes::Array{Float64, 2}, tree)

vectors and the adjustment in radius after each epoch.

# Arguments:
- `x`: training Data
- `codes`: Codebook
- `tree`: knn-compatible tree built upon the codes
"""
function doEpoch(x::Array{Float64,2}, codes::Array{Float64,2}, tree)

    # initialise numerator and denominator with 0's
    sumNumerator = zeros(Float64, size(codes))
    sumDenominator = zeros(Float64, size(codes)[1])

    # for each sample in dataset / trainingsset
    for s = 1:size(x, 1)
        (bmuIdx, bmuDist) = knn(tree, x[s, :], 1)

        target = bmuIdx[1]

        sumNumerator[target, :] .+= x[s, :]
        sumDenominator[target] += 1
    end

    return sumNumerator, sumDenominator
end

"""
    distributedEpoch(dInfo::Dinfo, codes::Matrix{Float64}, tree)

Execute the `doEpoch` in parallel on workers described by `dInfo` and collect
the results. Returns pair of numerator and denominator matrices.
"""
function distributedEpoch(dInfo::Dinfo, codes::Matrix{Float64}, tree)
    return dmapreduce(
        dInfo,
        (data) -> doEpoch(data, codes, tree),
        ((n1, d1), (n2, d2)) -> (n1 + n2, d1 + d2),
    )
end

"""
    mapToGigaSOM(som::Som, dInfo::Dinfo;
        knnTreeFun = BruteTree, metric = Euclidean(),
        output::Symbol=tmp_symbol(dInfo)::Dinfo

Compute the index of the BMU for each row of the input data.

# Arguments
- `som`: a trained SOM
- `dInfo`: `Dinfo` that describes the loaded and distributed data
- `knnTreeFun`: Constructor of the KNN-tree (e.g. from NearestNeighbors package)
- `metric`: Passed as metric argument to the KNN-tree constructor
- `output`: Symbol to save the result, defaults to `tmp_symbol(dInfo)`

Data must have the same number of dimensions as the training dataset
and will be normalised with the same parameters.
"""
function mapToGigaSOM(
    som::Som,
    dInfo::Dinfo;
    knnTreeFun = BruteTree,
    metric = Euclidean(),
    output::Symbol = tmp_symbol(dInfo),
)::Dinfo

    tree = knnTreeFun(Array{Float64,2}(transpose(som.codes)), metric)

    return dtransform(dInfo, (d) -> (vcat(knn(tree, transpose(d), 1)[1]...)), output)
end

"""
    mapToGigaSOM(som::Som, data;
                 knnTreeFun = BruteTree,
                 metric = Euclidean())
Overload of `mapToGigaSOM` for simple DataFrames and matrices. This slices the
data using `DistributedArrays`, sends them the workers, and runs normal
`mapToGigaSOM`. Data is `unscatter`d after the computation.
"""
function mapToGigaSOM(som::Som, data; knnTreeFun = BruteTree, metric = Euclidean())

    data = Matrix{Float64}(data)

    if size(data, 2) != size(som.codes, 2)
        @error "Data dimension ($(size(data,2))) does not match codebook dimension ($(size(som.codes,2)))."
        error("Data dimensions do not match")
    end

    dInfo = scatter_array(:GigaSOMmappingDataVar, data, workers())
    rInfo = mapToGigaSOM(som, dInfo, knnTreeFun = knnTreeFun, metric = metric)
    res = gather_array(rInfo)
    unscatter(dInfo)
    unscatter(rInfo)
    return DataFrame(index = res)
end

"""
    scaleEpochTime(iteration::Int64, epochs::Int64)

Convert iteration ID and epoch number to relative time in training.
"""
function scaleEpochTime(iteration::Int64, epochs::Int64)
    # prevent division by zero on 1-epoch training
    if epochs > 1
        epochs -= 1
    end

    return Float64(iteration - 1) / Float64(epochs)
end