"""
branch_mostinfeasible(m, node, disc2var_idx)

Get the index of an integer variable which is currently continuous which is most unintegral.
(nearest to *.5)
"""
function branch_mostinfeasible(m, node, disc2var_idx)
    x = node.solution
    idx = 0
    max_diff = 0
    for i in disc2var_idx
        diff = abs(x[i]-round(x[i]))
        if diff > max_diff
            idx = i
            max_diff = diff
        end
    end
    return idx
end

"""
init_strong_restart!(node, opts, var_idx, disc_var_idx, l_nd, r_nd, reasonable_disc_vars, infeasible_disc_vars,
 left_node, right_node, strong_restart, max_gain_var, atol, var_type, disc2var_idx)

Tighten the bounds for the node and check if there are variables that need to be checked for a restart.
"""
function init_strong_restart!(node, opts, var_idx, disc_var_idx, l_nd, r_nd,
                                reasonable_disc_vars, infeasible_disc_vars,
                                left_node, right_node, strong_restart, max_gain_var,
                                atol, var_type, disc2var_idx)
    restart = false
    set_to_last_var = false
    need_to_resolve = false

    # set the bounds directly for the node
    # also update the best bound and the solution
    if !state_is_optimal(l_nd.relaxation_state; allow_almost=opts.allow_almost_solved)
        # the solution shouldn't be updated when the current max_gain_var is then type correct
        if max_gain_var == 0 || !is_type_correct(r_nd.solution[max_gain_var],var_type[max_gain_var],atol)
            # if all variables reasonable_disc_vars are type correct => just continue with the tree
            # no restarts and no change in bounds
            if !all_reasonable_type_correct(r_nd.solution, disc2var_idx, reasonable_disc_vars, atol)
                node.l_var[var_idx] = ceil(node.solution[var_idx])
                node.best_bound = r_nd.best_bound
                node.solution = r_nd.solution
                need_to_resolve = true
                push!(infeasible_disc_vars, disc_var_idx)
                strong_restart && (restart = true)
            end
        end
    else
        # the solution shouldn't be updated when the current max_gain_var is then type correct
        if max_gain_var == 0 || !is_type_correct(l_nd.solution[max_gain_var],var_type[max_gain_var],atol)
            if !all_reasonable_type_correct(l_nd.solution, disc2var_idx, reasonable_disc_vars, atol)
                node.u_var[var_idx] = floor(node.solution[var_idx])
                node.best_bound = l_nd.best_bound
                node.solution = l_nd.solution
                need_to_resolve = true
                push!(infeasible_disc_vars, disc_var_idx)
                strong_restart && (restart = true)
            end
        end
    end

    if length(reasonable_disc_vars) == length(infeasible_disc_vars)
        # basically branching on the last infeasible variable
        set_to_last_var = true
        restart = false
    end
    return need_to_resolve, restart, infeasible_disc_vars, set_to_last_var
end

"""
    branch_strong_on!(m,opts,step_obj,reasonable_disc_vars, disc2var_idx, strong_restart, counter)

Try to branch on a few different variables and choose the one with highest obj_gain.
Update obj_gain for the variables tried and average the other ones.
"""
function branch_strong_on!(m,opts,step_obj,
    reasonable_disc_vars, disc2var_idx, strong_restart, counter)

    function set_temp_gains!(gains, gain_l, gain_r, disc_var_idx)
        if !isinf(gain_l)
            gains.minus[disc_var_idx] = gain_l
            gains.minus_counter[disc_var_idx] = 1
        else
            gains.inf_counter[disc_var_idx] = 1
        end
        if !isinf(gain_r)
            gains.plus[disc_var_idx] = gain_r
            gains.plus_counter[disc_var_idx] = 1
        else
            gains.inf_counter[disc_var_idx] = 1
        end

        if !isinf(gain_l) && !isinf(gain_r)
            gains.inf_counter[disc_var_idx] -= 1
        end
    end

    function get_current_gains(opts, node, l_nd, r_nd)
        gain_l = sigma_minus(opts, node, l_nd, node.solution[node.var_idx])
        gain_r = sigma_plus( opts, node,  r_nd, node.solution[node.var_idx])

        if isinf(gain_l)
            gain = gain_r
        elseif isinf(gain_r)
            gain = gain_l
        else
            gain = (gain_l+gain_r)/2
        end
        if isnan(gain)
            gain = Inf
        end
        return gain_l, gain_r, gain
    end

    function add_strong_step(strong_step,l_nd,r_nd; restart=false)
        strong_step.l_relaxation_state = l_nd.relaxation_state
        strong_step.r_relaxation_state = r_nd.relaxation_state
        strong_step.init_restart = restart
        push!(step_obj.strong_branching, strong_step)
    end

    strong_time = time()

    node = step_obj.node
    strong_restarts = -1
    restart = true
    status = :Normal
    node = step_obj.node
    infeasible_disc_vars = zeros(Int64,0)
    gains = init_gains(m.num_disc_var)

    left_node, right_node = nothing, nothing
    atol = opts.atol

    need_to_resolve = false
    opts.debug && (step_obj.strong_branching = [])

    max_gain = -Inf # then one is definitely better
    max_gain_var = 0
    max_gain_disc_var = 0

    allow_almost = opts.allow_almost_solved
    while restart
        strong_restarts += 1 # is init with -1
        restart = false
        for disc_var_idx in reasonable_disc_vars
            # don't rerun if the variable has already one infeasible node
            if disc_var_idx in infeasible_disc_vars
                continue
            end
            var_idx = disc2var_idx[disc_var_idx]
            step_obj.var_idx = var_idx
            u_b, l_b = node.u_var[var_idx], node.l_var[var_idx]
            # don't rerun if bounds are exact or is type correct
            if isapprox(u_b,l_b; atol=atol) || is_type_correct(node.solution[var_idx],m.var_type[var_idx], atol)
                continue
            end

            # branch on the current variable and get the corresponding children
            l_nd,r_nd = branch!(m, opts, step_obj, counter, disc2var_idx; temp=true)
            opts.debug && (strong_step = new_default_strong_branch_step_obj(var_idx))

            # no current restart => we can set max_gain and variable
            gain_l, gain_r, gain = get_current_gains(opts, node, l_nd, r_nd)

            if !state_is_optimal(l_nd.relaxation_state; allow_almost=allow_almost) &&
                !state_is_optimal(r_nd.relaxation_state; allow_almost=allow_almost) && counter == 1
                # TODO: Might be Error/UserLimit instead of infeasible
                status = :Infeasible
                left_node = l_nd
                right_node = r_nd
                opts.debug && add_strong_step(strong_step,l_nd,r_nd)
                break
            end

            # check if one part is infeasible => update bounds & restart if strong restart is true
            if !state_is_optimal(l_nd.relaxation_state; allow_almost=allow_almost) ||
                !state_is_optimal(r_nd.relaxation_state; allow_almost=allow_almost)
                if !state_is_optimal(l_nd.relaxation_state; allow_almost=allow_almost) && !state_is_optimal(r_nd.relaxation_state; allow_almost=allow_almost)
                    # TODO: Might be Error/UserLimit instead of infeasible
                    status = :PartlyInfeasible
                    left_node = l_nd
                    right_node = r_nd
                    opts.debug && add_strong_step(strong_step,l_nd,r_nd)
                    break
                end
                need_to_resolve, restart, new_infeasible_disc_vars, set_to_last_var = init_strong_restart!(node, opts, var_idx, disc_var_idx, l_nd, r_nd, reasonable_disc_vars, infeasible_disc_vars, left_node, right_node, strong_restart, max_gain_var, opts.atol, m.var_type, disc2var_idx)
                infeasible_disc_vars = new_infeasible_disc_vars

                # only variables where one branch is infeasible => no restart and break
                if set_to_last_var
                    max_gain_var = var_idx
                    max_gain_disc_var = disc_var_idx
                    left_node = l_nd
                    right_node = r_nd
                    restart = false
                    need_to_resolve = false
                    gain_l, gain_r, gain = get_current_gains(opts, node, l_nd, r_nd)
                    max_gain = gain
                    set_temp_gains!(gains, gain_l, gain_r, disc_var_idx)
                    opts.debug && add_strong_step(strong_step,l_nd,r_nd)
                    break
                end

                if restart && time()-strong_time > opts.strong_branching_time_limit
                    restart = false
                end
            end

            # don't update maximum if restart
            if gain > max_gain && !restart
                max_gain = gain
                max_gain_var = var_idx
                max_gain_disc_var = disc_var_idx
                left_node = l_nd
                right_node = r_nd
                # we are changing the bounds if one branch is infeasible then the current
                # selection might not work anymore => we have to resolve it later
                need_to_resolve = false
            end
            set_temp_gains!(gains, gain_l, gain_r, disc_var_idx)
            opts.debug && add_strong_step(strong_step,l_nd,r_nd;restart=restart)
            if restart
                break
            end

            if time()-m.start_time >= opts.time_limit
                break
            end

            if time()-strong_time >= opts.strong_branching_time_limit && !isinf(max_gain)
                @info "Breaking out of strong branching as the time limit of $(opts.strong_branching_time_limit) seconds got reached."
                break
            end
        end
    end

    # check if need to resolve (not if Local/Global Infeasible)
    if need_to_resolve && status == :Normal
        status = :Resolve
    end

    return status, max_gain_var, left_node, right_node, gains, strong_restarts
end

"""
branch_strong!(m,opts,disc2var_idx,step_obj,counter)

Try to branch on a few different variables and choose the one with highest obj_gain.
Update obj_gain for the variables tried and average the other ones.
"""
function branch_strong!(m,opts,disc2var_idx,step_obj,counter)
    node = step_obj.node

    # generate an of variables to branch on
    num_strong_var = Int(round((opts.strong_branching_perc/100)*m.num_disc_var))
    # if smaller than 2 it doesn't make sense
    num_strong_var = num_strong_var < 2 ? 2 : num_strong_var

    # get reasonable candidates (not type correct and not already perfectly bounded)
    disc_vars = m.num_disc_var
    atol = opts.atol
    reasonable_disc_vars = get_reasonable_disc_vars(node, m.var_type, disc_vars,  disc2var_idx, atol)
    if num_strong_var < length(reasonable_disc_vars)
        shuffle!(reasonable_disc_vars)
        reasonable_disc_vars = reasonable_disc_vars[1:num_strong_var]
    end

    # compute the gain for each reasonable candidate and choose the highest
    status, max_gain_var, left_node, right_node, gains, strong_restarts = branch_strong_on!(m,opts,step_obj,
        reasonable_disc_vars, disc2var_idx, opts.strong_restart, counter)

    step_obj.obj_gain += gains

    if status != :Resolve
        step_obj.l_nd = left_node
        step_obj.r_nd = right_node

        # set the variable to branch (best gain)
        node.state = :Done
        node.var_idx = max_gain_var
    end

    @assert max_gain_var != 0 || status == :PartlyInfeasible || status == :Infeasible || node.state == :Infeasible
    return status, max_gain_var, strong_restarts
end

function branch_reliable!(m,opts,step_obj,disc2var_idx,gains,counter)
    idx = 0
    node = step_obj.node
    mu = opts.gain_mu
    reliability_param = opts.reliability_branching_threshold
    reliability_perc = opts.reliability_branching_perc
    num_strong_var = Int(round((reliability_perc/100)*m.num_disc_var))
    # if smaller than 2 it doesn't make sense
    num_strong_var = num_strong_var < 2 ? 2 : num_strong_var

    gmc_r = gains.minus_counter .< reliability_param
    gpc_r = gains.plus_counter  .< reliability_param

    strong_restarts = 0
    reasonable_disc_vars = []
    atol = opts.atol
    for i=1:length(gmc_r)
        if gmc_r[i] || gpc_r[i]
            idx = disc2var_idx[i]
            u_b = node.u_var[idx]
            l_b = node.l_var[idx]
            if isapprox(u_b,l_b; atol=atol) || is_type_correct(node.solution[idx],m.var_type[idx],atol)
                continue
            end
            push!(reasonable_disc_vars,i)
        end
    end
    if length(reasonable_disc_vars) > 0
        unrealiable_idx = sortperm(gains.minus_counter[reasonable_disc_vars])
        reasonable_disc_vars = reasonable_disc_vars[unrealiable_idx]
        num_reasonable = num_strong_var < length(reasonable_disc_vars) ? num_strong_var : length(reasonable_disc_vars)
        reasonable_disc_vars = reasonable_disc_vars[1:num_reasonable]

        status, max_gain_var, left_node, right_node, gains, strong_restarts = branch_strong_on!(m,opts,step_obj,
            reasonable_disc_vars, disc2var_idx, opts.strong_restart, counter)

        step_obj.obj_gain += gains

        step_obj.upd_gains = :GainsToTree
        new_gains = copy(step_obj.obj_gain)

        if status == :Infeasible
            return :Infeasible, 0, strong_restarts
        end
    else
        step_obj.upd_gains = :GuessAndUpdate
        new_gains = copy(gains)
    end
    idx = branch_pseudo(m, node, disc2var_idx, new_gains, mu, atol)
    return :Normal, idx, strong_restarts
end

function branch_pseudo(m, node, disc2var_idx, obj_gain, mu, atol)
    # use the one with highest obj_gain which is currently continuous
    idx = 0
    scores, sort_idx = sorted_score_idx(node.solution, obj_gain, disc2var_idx, mu)
    for l_idx in sort_idx
        var_idx = disc2var_idx[l_idx]
        if !is_type_correct(node.solution[var_idx], m.var_type[var_idx], atol)
            u_b = node.u_var[var_idx]
            l_b = node.l_var[var_idx]
            # if the upper bound is the lower bound => should be type correct
            @assert !isapprox(u_b, l_b; atol=atol)
            idx = var_idx
            break
        end
    end
    return idx
end

function sorted_score_idx(x, gains, i2v, mu)
    g_minus, g_minus_c = gains.minus, gains.minus_counter
    g_plus, g_plus_c = gains.plus, gains.plus_counter
    g_minus_c += map(i -> (i == 0) && (i = 1), g_minus_c)
    g_plus_c += map(i -> (i == 0) && (i = 1), g_plus_c)
    scores = [score(f_minus(x[i2v[i]])*g_minus[i]/g_minus_c[i],f_plus(x[i2v[i]])*g_plus[i]/g_plus_c[i],mu) for i=1:length(g_minus)]
    sortedidx = sortperm(scores; rev=true)
    infsortedidx = sortperm(gains.inf_counter[sortedidx]; rev=true)
    return scores,sortedidx[infsortedidx]
end

"""
    Score function from
    http://citeseerx.ist.psu.edu/viewdoc/download?doi=10.1.1.92.7117&rep=rep1&type=pdf
"""
function score(q_m, q_p, mu)
    minq = q_m < q_p ? q_m : q_p
    maxq = q_m < q_p ? q_p : q_m
    return (1-mu)*minq+mu*maxq
end

function diff_obj(opts, node, cnode)
    if state_is_optimal(cnode.relaxation_state; allow_almost=opts.allow_almost_solved)
        return abs(node.best_bound - cnode.best_bound)
    else
        return Inf
    end
end

f_plus(x) = ceil(x)-x
f_minus(x) = x-floor(x)
sigma_plus(opts, node,r_nd,x) = diff_obj(opts, node,r_nd)/f_plus(x)
sigma_minus(opts, node,l_nd,x) = diff_obj(opts, node,l_nd)/f_minus(x)


""" Graphviz support for Catlab's graph types.
"""
module GraphvizGraphs
export parse_graphviz, to_graphviz

using StaticArrays: StaticVector, SVector

using ...Present, ...Theories
using ...Graphs, ...CategoricalAlgebra.Subobjects
import ..Graphviz

# Property graphs
#################

""" Parse Graphviz output in JSON format.

Returns a property graph with graph layout and other metadata. Each node has a
position and size.

All units are in points. Note that Graphviz has 72 points per inch.
"""
function parse_graphviz(doc::AbstractDict)::AbstractPropertyGraph
  graph = doc["directed"] ? PropertyGraph{Any}() : SymmetricPropertyGraph{Any}()
  nsubgraphs = doc["_subgraph_cnt"] # Subgraphs are ignored.

  # Graph-level layout: bounds and padding.
  # It seems, but is not documented, that the first two numbers in the Graphviz
  # bounding box are always zero.
  set_gprops!(graph,
    bounds = SVector{2}(parse_vector(doc["bb"])[3:4]),
    pad = 72*parse_point(get(doc, "pad", "0,0")),
    rankdir = get(doc, "rankdir", "TB"),
  )

  # Add vertex for each Graphviz node.
  node_keys = ("id", "name", "comment", "label", "shape", "style")
  for node in doc["objects"][nsubgraphs+1:end]
    props = Dict{Symbol,Any}(
      Symbol(k) => node[k] for k in node_keys if haskey(node, k))
    props[:position] = parse_point(node["pos"])
    props[:size] = 72*SVector(parse(Float64, node["width"]),
                              parse(Float64, node["height"]))
    add_vertex!(graph, props)
  end

  # Add edge for each Graphviz edge.
  edge_keys = ("id", "comment", "label", "xlabel", "headlabel", "taillabel",
               "headport", "tailport")
  for edge in doc["edges"]
    if get(edge, "style", nothing) == "invis"
      # Omit invisible edges, which are used to tweak the layout in Graphviz.
      continue
    end
    props = Dict{Symbol,Any}(
      Symbol(k) => edge[k] for k in edge_keys if haskey(edge, k))
    props[:spline] = parse_spline(edge["pos"])
    src = edge["tail"] - nsubgraphs + 1
    tgt = edge["head"] - nsubgraphs + 1
    add_edge!(graph, src, tgt, props)
  end

  graph
end

""" Parse an array of floats in Graphviz's comma-separated format.
"""
parse_vector(s::AbstractString) = [ parse(Float64, x) for x in split(s, ",") ]

""" Parse Graphviz point.

http://www.graphviz.org/doc/info/attrs.html#k:point
"""
parse_point(s::AbstractString) = SVector{2}(parse_vector(s))

""" Parse Graphviz spline.

In Graphviz, a "spline" is a cubic B-spline of overlapping cubic Bezier curves.
It consists of 3n+1 points, where n is the number of Bezier curves.

http://www.graphviz.org/doc/info/attrs.html#k:splineType
https://web.archive.org/web/20170418034924/http://www.graphviz.org/content/how-convert-b-spline-bezier
"""
function parse_spline(spline::AbstractString)
  points = StaticVector{2,Float64}[]
  start, stop = nothing, nothing
  for s in split(spline, " ")
    if startswith(s, "s,"); start = parse_point(s[3:end])
    elseif startswith(s, "e,"); stop = parse_point(s[3:end])
    else push!(points, parse_point(s)) end
  end
  # Prefer explicit start or end points to the spline start and end points.
  # Thus, endpoints may pass into the node shape but should not fall short.
  if !isnothing(start); points[0] = start end
  if !isnothing(stop); points[end] = stop end
  points
end

""" Convert a property graph to a Graphviz graph.

This method is usually more convenient than direct AST manipulation for creating
simple Graphviz graphs. For more advanced features, like nested subgraphs, you
must use the Graphviz AST.
"""
function to_graphviz(g::AbstractPropertyGraph)::Graphviz.Graph
  gv_name(v::Int) = "n$v"
  gv_path(e::Int) = [gv_name(src(g,e)), gv_name(tgt(g,e))]

  # Add Graphviz node for each vertex in property graph.
  stmts = Graphviz.Statement[]
  for v in vertices(g)
    push!(stmts, Graphviz.Node(gv_name(v), vprops(g, v)))
  end

  # Add Graphviz edge for each edge in property graph.
  is_directed = !(g isa SymmetricPropertyGraph)
  for e in edges(g)
    # In undirected case, only include one edge from each pair.
    if is_directed || (e <= inv(g,e))
      push!(stmts, Graphviz.Edge(gv_path(e), eprops(g, e)))
    end
  end

  attrs = gprops(g)
  Graphviz.Graph(
    name = get(attrs, :name, "G"),
    directed = is_directed,
    prog = get(attrs, :prog, is_directed ? "dot" : "neato"),
    stmts = stmts,
    graph_attrs = Graphviz.as_attributes(get(attrs, :graph, Dict())),
    node_attrs = Graphviz.as_attributes(get(attrs, :node, Dict())),
    edge_attrs = Graphviz.as_attributes(get(attrs, :edge, Dict())),
  )
end

# Graphs
########

# FIXME: Should be an abstract type after StructACSets refactor.
const IsGraph = Union{AbstractGraph, AbstractSymmetricGraph,
                      AbstractReflexiveGraph, AbstractSymmetricReflexiveGraph}

""" Convert a graph to a Graphviz graph.

A simple default style is applied. For more control over the visual appearance,
first convert the graph to a property graph, define the Graphviz attributes as
needed, and then convert to a Graphviz graph.
"""
function to_graphviz(g::IsGraph; kw...)
  to_graphviz(to_graphviz_property_graph(g; kw...))
end

function to_graphviz_property_graph(g::AbstractGraph;
    prog::AbstractString="dot", graph_attrs::AbstractDict=Dict(),
    node_attrs::AbstractDict=Dict(), edge_attrs::AbstractDict=Dict(),
    node_labels::Bool=false, edge_labels::Bool=false)
  node_labeler(v) = Dict(:label => node_labels ? string(v) : "")
  edge_labeler(e) = if edge_labels
    Dict(:label => string(e))
  else
    Dict{Symbol,String}()
  end
  PropertyGraph{Any}(g, node_labeler, edge_labeler;
    prog = prog,
    graph = merge!(Dict(:rankdir => "LR"), graph_attrs),
    node = merge!(default_node_attrs(node_labels), node_attrs),
    edge = merge!(Dict(:arrowsize => "0.5"), edge_attrs),
  )
end

default_node_attrs(show_labels::Bool) = Dict(
  :shape => show_labels ? "circle" : "point",
  :width => "0.05",
  :height => "0.05",
  :margin => "0",
)

# Symmetric graphs
##################

function to_graphviz_property_graph(g::AbstractSymmetricGraph;
    prog::AbstractString="neato", graph_attrs::AbstractDict=Dict(),
    node_attrs::AbstractDict=Dict(), edge_attrs::AbstractDict=Dict(),
    node_labels::Bool=false, edge_labels::Bool=false)
  node_labeler(v) = Dict(:label => node_labels ? string(v) : "")
  edge_labeler(e) = if edge_labels
    e′ = inv(g,e)
    Dict(:label => "($(min(e,e′)),$(max(e,e′)))")
  else
    Dict{Symbol,String}()
  end
  SymmetricPropertyGraph{Any}(g, node_labeler, edge_labeler;
    prog = prog,
    graph = graph_attrs,
    node = merge!(default_node_attrs(node_labels), node_attrs),
    edge = merge!(Dict(:len => "0.5"), edge_attrs),
  )
end

# Reflexive graphs
##################

function to_graphviz_property_graph(g::AbstractReflexiveGraph;
    prog::AbstractString="dot", graph_attrs::AbstractDict=Dict(),
    node_attrs::AbstractDict=Dict(), edge_attrs::AbstractDict=Dict(),
    node_labels::Bool=false, edge_labels::Bool=false,
    show_reflexive::Bool=false)
  pg = PropertyGraph{Any}(; prog = prog,
    graph = graph_attrs,
    node = merge!(default_node_attrs(node_labels), node_attrs),
    edge = merge!(Dict(:arrowsize => "0.5"), edge_attrs),
  )
  for v in vertices(g)
    add_vertex!(pg, label=node_labels ? string(v) : "")
  end
  reflexive_edges = Set(refl(g))
  for e in edges(g)
    is_reflexive = e ∈ reflexive_edges
    if show_reflexive || !is_reflexive
      e′ = add_edge!(pg, src(g,e), tgt(g,e))
      if is_reflexive; set_eprop!(pg, e′, :style, "dashed") end
      if edge_labels; set_eprop!(pg, e′, :label, string(e)) end
    end
  end
  pg
end

# Symmetric reflexive graphs
############################

function to_graphviz_property_graph(g::AbstractSymmetricReflexiveGraph;
    prog::AbstractString="neato", graph_attrs::AbstractDict=Dict(),
    node_attrs::AbstractDict=Dict(), edge_attrs::AbstractDict=Dict(),
    node_labels::Bool=false, edge_labels::Bool=false,
    show_reflexive::Bool=false)
  pg = SymmetricPropertyGraph{Any}(; prog = prog,
    graph = graph_attrs,
    node = merge!(default_node_attrs(node_labels), node_attrs),
    edge = merge!(Dict(:len => "0.5"), edge_attrs),
  )
  for v in vertices(g)
    add_vertex!(pg, label=node_labels ? string(v) : "")
  end
  reflexive_edges = Set(refl(g))
  for e in edges(g)
    is_reflexive = e ∈ reflexive_edges
    if is_reflexive && show_reflexive
      e′ = add_edge!(pg, src(g,e), tgt(g,e), style="dashed")
      if edge_labels; set_eprop!(pg, e′, :label, string(e)) end
    elseif !is_reflexive && e <= inv(g,e)
      e′ = add_edge!(pg, src(g,e), tgt(g,e))
      if edge_labels; set_eprop!(pg, e′, :label, "($e,$(inv(g,e)))") end
    end
  end
  pg
end

# Half-edge graphs
##################

to_graphviz(g::AbstractHalfEdgeGraph; kw...) =
  to_graphviz(to_graphviz_property_graph(g; kw...))

function to_graphviz_property_graph(g::AbstractHalfEdgeGraph;
    prog::AbstractString="neato", graph_attrs::AbstractDict=Dict(),
    node_attrs::AbstractDict=Dict(), edge_attrs::AbstractDict=Dict(),
    node_labels::Bool=false, edge_labels::Bool=false)
  pg = SymmetricPropertyGraph{Any}(; prog = prog,
    graph = graph_attrs,
    node = merge!(default_node_attrs(node_labels), node_attrs),
    edge = merge!(Dict(:len => "0.5"), edge_attrs),
  )
  for v in vertices(g)
    add_vertex!(pg, label=node_labels ? string(v) : "")
  end
  for e in half_edges(g)
    if e == inv(g,e)
      # Dangling half-edge: add an invisible vertex.
      v = add_vertex!(pg, style="invis", shape="none", label="")
      e′ = add_edge!(pg, vertex(g,e), v)
      if edge_labels; set_eprop!(pg, e′, :label, string(e)) end
    elseif e < inv(g,e)
      # Pair of distict half-edges: add a standard edge.
      e′ = add_edge!(pg, vertex(g,e), vertex(g,inv(g,e)))
      if edge_labels; set_eprop!(pg, e′, :label, "($e,$(inv(g,e)))") end
    end
  end
  pg
end

# Subgraphs
###########

to_graphviz(subgraph::Subobject{<:IsGraph}; kw...) =
  to_graphviz(to_graphviz_property_graph(subgraph; kw...))

function to_graphviz_property_graph(
    subgraph::Subobject{<:Union{AbstractGraph,AbstractSymmetricGraph}};
    subgraph_node_attrs::AbstractDict=Dict(:color => "cornflowerblue"),
    subgraph_edge_attrs::AbstractDict=Dict(:color => "cornflowerblue"), kw...)
  pg = to_graphviz_property_graph(ob(subgraph); kw...)
  f = hom(subgraph)
  for v in vertices(dom(f))
    set_vprops!(pg, f[:V](v), subgraph_node_attrs)
  end
  for e in edges(dom(f))
    set_eprops!(pg, f[:E](e), subgraph_edge_attrs)
  end
  pg
end

# Schemas
#########

to_graphviz(pres::Presentation{Schema}) =
  to_graphviz(to_graphviz_property_graph(pres))

function to_graphviz_property_graph(pres::Presentation{Schema})
  obs,homs,attr_types,attrs = generators.(Ref(pres), [:Ob,:Hom,:AttrType,:Attr])
  g = PropertyGraph{Any}()

  ob_vertices = add_vertices!(g,length(obs))
  for (v, ob) in zip(ob_vertices, obs)
    set_vprop!(g,v,:label,string(nameof(ob)))
    set_vprop!(g,v,:shape,"plain")
    set_vprop!(g,v,:margin,"2")
  end

  attr_vertices = add_vertices!(g,length(attr_types))
  for (v, attr_type) in zip(attr_vertices, attr_types)
    set_vprop!(g,v,:label,string(nameof(attr_type)))
  end

  hom_edges = add_edges!(g,
    ob_vertices[generator_index.(Ref(pres), nameof.(dom.(homs)))],
    ob_vertices[generator_index.(Ref(pres), nameof.(codom.(homs)))])
  for (e, hom) in zip(hom_edges, homs)
    set_eprop!(g,e,:label,string(nameof(hom)))
    set_eprop!(g,e,:len,"2")
  end

  attr_edges = add_edges!(g,
    ob_vertices[generator_index.(Ref(pres), nameof.(dom.(attrs)))],
    attr_vertices[generator_index.(Ref(pres), nameof.(codom.(attrs)))])
  for (e, attr) in zip(attr_edges, attrs)
    set_eprop!(g,e,:label,string(nameof(attr)))
    set_eprop!(g,e,:len,"2")
  end

  set_gprop!(g,:graph,Dict(:rankdir => "LR"))
  set_gprop!(g,:prog,"neato")
  g
end

end


"""
FaultTree
"""

export AbstractFaultTreeNode, AbstractFaultTreeOperation, FaultTreeOperation, FaultTreeEvent, FaultTreeKoutofN
export ftevent

abstract type AbstractFaultTreeNode end
abstract type AbstractFaultTreeOperation <: AbstractFaultTreeNode end

struct FaultTreeOperation <: AbstractFaultTreeOperation
    params::Set{Symbol}
    op::Symbol
    args::Vector{AbstractFaultTreeNode}
end

struct FaultTreeKoutofN <: AbstractFaultTreeOperation
    params::Set{Symbol}
    op::Symbol
    k::Int
    args::Vector{AbstractFaultTreeNode}
end

struct FaultTreeEvent <: AbstractFaultTreeNode
    params::Set{Symbol}
    var::Symbol
end

"""
ftree
"""

function ftevent(var::Symbol)
    FaultTreeEvent(Set([var]), var)
end

function ftevent(var::Vararg)
    s = Symbol(var...)
    FaultTreeEvent(Set([s]), s)
end


@testset "477.total-hamming-distance.jl" begin
    @test total_hamming_distance([4, 14, 2]) == 6
    @test total_hamming_distance([4, 14, 4]) == 4
end


using SparseMatrixAssemblers
using SparseArrays
using Test

@testset "Constructors" begin
    # Mainly testing that expressions compile
    @test typeof(SparseMatrixAssembler()) == SparseMatrixAssembler{Float64,Int}
    @test typeof(SparseMatrixAssembler{Int32}()) == SparseMatrixAssembler{Int32,Int}
    @test typeof(SparseMatrixAssembler{Int32,Int32}()) == SparseMatrixAssembler{Int32,Int32}
end

@testset "Conversion" begin
    # Mainly testing that expressions compile
    @test typeof(convert(SparseMatrixAssembler, SparseMatrixAssembler())) == SparseMatrixAssembler{Float64,Int}
    @test typeof(convert(SparseMatrixAssembler{Int32}, SparseMatrixAssembler())) == SparseMatrixAssembler{Int32,Int}
    @test typeof(convert(SparseMatrixAssembler{Int32,Int32}, SparseMatrixAssembler())) == SparseMatrixAssembler{Int32,Int32}
end

@testset "Assignment" begin
    @testset "Scalar" begin
        A = SparseMatrixAssembler{Int}()
        A[1,2] = 1
        @test sparse(A) == [0 1]
    end

    @testset "Col vector" begin
        A = SparseMatrixAssembler{Int}()
        A[1:2,2] = [1,2]
        @test sparse(A) == [0 1; 0 2]
    end

    @testset "Row vector" begin
        A = SparseMatrixAssembler{Int}()
        A[2,1:2] = [1,2]
        @test sparse(A) == [0 0; 1 2]
    end

    @testset "Matrix" begin
        A = SparseMatrixAssembler{Int}()
        A[2:3,1:2] = [1 2; 3 4]
        @test sparse(A) == [0 0; 1 2; 3 4]
    end
end

@testset "Conversion to other types" begin
    # Mainly testing that expressions compile
    i = [1,1]
    j = [1,1]
    v = [1,2]
    @test sparse(SparseMatrixAssembler(i,j,v)) == sparse(i,j,v)
    @test sparse(SparseMatrixAssembler(i,j,v),2,2) == sparse(i,j,v,2,2)
    @test sparse(SparseMatrixAssembler(i,j,v),2,2,*) == sparse(i,j,v,2,2, *)

    @test SparseMatrixCSC(SparseMatrixAssembler(i,j,v)) == sparse(i,j,v)
    @test SparseMatrixCSC(SparseMatrixAssembler(i,j,v),2,2) == sparse(i,j,v,2,2)
    @test SparseMatrixCSC(SparseMatrixAssembler(i,j,v),2,2,*) == sparse(i,j,v,2,2, *)

    @test Matrix(SparseMatrixAssembler(i,j,v)) == sparse(i,j,v)
    @test Matrix(SparseMatrixAssembler(i,j,v),2,2) == sparse(i,j,v,2,2)
    @test Matrix(SparseMatrixAssembler(i,j,v),2,2,*) == sparse(i,j,v,2,2, *)
end


############################################################################################
"""
$(TYPEDEF)

Contains temporary buffer values to avoid allocations during particle propagation.

# Fields
$(TYPEDFIELDS)
"""
struct ParticleBuffer{A,I,P}
    "Contains buffer values for particles and ancestor for one iteration."
    parameter::BaytesCore.ParameterBuffer{A, I}
    "Proposal trajectory and predicted latent variable."
    proposal::Vector{A}
    "Predicted latent and oberved data"
    prediction::Vector{P}
    "Stores boolean if resampled at each iteration."
    resampled::Vector{Bool}
    function ParticleBuffer(
        prediction::Vector{P}, reference::AbstractVector{A}, Nparticles::Integer, Ndata::Integer, F::Type{I}
    ) where {P,A,I<:Integer}
        parameter = BaytesCore.ParameterBuffer(reference, Nparticles, F)
        proposal = typeof(reference)(undef, Ndata)
        resampled = zeros(Bool, Ndata)
        return new{A,I,P}(parameter, proposal, prediction, resampled)
    end
end
function update!(buffer::ParticleBuffer, Nparticles::Integer, Ndata::Integer)
    resize!(buffer.parameter, Nparticles)
    resize!(buffer.proposal, Ndata)
    resize!(buffer.resampled, Ndata)
    return nothing
end

############################################################################################
#export
export ParticleBuffer, update!


module GroupedArrays
using Missings
using DataAPI
using Base.Threads
include("spawn.jl")
include("utils.jl")

"""
GroupedArray{T,N} <: AbstractArray{T,N}
N-dimensional dense array with elements of type T, where T <: Union{Int, Missing}
"""
mutable struct GroupedArray{T <: Union{Int, Missing}, N} <: AbstractArray{T, N}
	groups::Array{Int, N}
	ngroups::Int
end
const GroupedVector{T} = GroupedArray{T, 1}
const GroupedMatrix{T} = GroupedArray{T, 2}

##############################################################################
##
## Arrays API
##
##############################################################################

Base.size(g::GroupedArray) = size(g.groups)
Base.axes(g::GroupedArray) = axes(g.groups)
Base.IndexStyle(g::GroupedArray) = Base.IndexLinear()
Base.LinearIndices(g::GroupedArray) = axes(g.groups, 1)

Base.@propagate_inbounds function Base.getindex(g::GroupedArray{Int}, i::Number)
	@boundscheck checkbounds(g, i)
	@inbounds g.groups[i]
end

Base.@propagate_inbounds function Base.getindex(g::GroupedArray, i::Number)
	@boundscheck checkbounds(g, i)
	@inbounds x = g.groups[i]
	x == 0 ? missing : x
end

Base.@propagate_inbounds function Base.setindex!(g::GroupedArray, x::Number,  i::Number)
	@boundscheck checkbounds(g, i)
	x > 0 || throw(ArgumentError("The number x must be positive"))
	x > g.ngroups && (g.ngroups = x)
	@inbounds g.groups[i] = x
end

Base.@propagate_inbounds function Base.setindex!(g::GroupedArray{T}, ::Missing,  i::Number) where {T >: Missing}
	@boundscheck checkbounds(g, i)
	@inbounds g.groups[i] = 0
end

##############################################################################
##
## Constructor
##
##############################################################################

"""

	GroupedArray(args... [; coalesce = false, sort = nothing])

Construct a `GroupedArray` taking on distinct values for the groups formed by elements of `args`

### Arguments
* `args...`: `AbstractArrays` of same sizes.

### Keyword arguments
* `coalesce::Bool`: should missing values considered as distinct grotups indicators?
* `sort::Union{Bool, Nothing}`: should the order of the groups be the sort order?

### Examples
```julia
using GroupedArrays
p1 = ["a", "a", "b", "b", missing, missing]
GroupedArray(p1)
GroupedArray(p1; coalesce = true)
p2 = [1, 1, 1, 2, 2, 2]
GroupedArray(p1, p2)
```
"""
function GroupedArray(args...; coalesce = false, sort = nothing)
	all(x isa AbstractArray for x in args) || throw(DimensionMismatch("arguments are not AbstractArrays"))
	s = size(first(args))
	all(size(x) == s for x in args) || throw(DimensionMismatch("cannot match array  sizes"))
	groups = Vector{Int}(undef, prod(s))
	ngroups, rhashes, gslots, sorted = row_group_slots(vec.(args), Val(false), groups, !coalesce, sort)
	# sort groups if row_group_slots hasn't already done that
	if sort === true && !sorted
		idx = find_index(GroupedVector{Int}(groups, ngroups))
		group_invperm = invperm(sortperm(collect(zip(map(x -> view(x, idx), args)...))))
		@inbounds for (i, gix) in enumerate(groups)
			groups[i] = gix > 0 ? group_invperm[gix] : 0
		end
	end
	T = !coalesce && any(eltype(x) >: Missing for x in args) ? Union{Int, Missing} : Int
	GroupedArray{T, length(s)}(reshape(groups, s), ngroups)
end

# Find index of representative row for each group
function find_index(g::GroupedArray)
	groups, ngroups = g.groups, g.ngroups
	idx = Vector{Int}(undef, ngroups)
	filled = fill(false, ngroups)
	nfilled = 0
	@inbounds for (i, gix) in enumerate(groups)
		if gix > 0 && !filled[gix]
			filled[gix] = true
			idx[gix] = i
			nfilled += 1
			nfilled == ngroups && break
		end
	end
	return idx
end

##############################################################################
##
## Conversions
##
##############################################################################
Base.convert(::Type{GroupedArray{T, N}}, g::GroupedArray{T, N}) where {T, N} = g
function Base.convert(::Type{GroupedArray{Union{Int, Missing},N}}, g::GroupedArray{Int, N}) where {N}
    return GroupedArray{Union{Int, Missing},N}(g.groups, g.ngroups)
end
function Base.convert(::Type{GroupedArray{Int, N}}, g::GroupedArray{Union{Int, Missing}, N}) where {N}
	@assert all(x > 0 for x in g.groups)
    return GroupedArray{Int,N}(g.groups, g.ngroups)
end

Base.convert(::Type{GroupedArray}, g::GroupedArray) = g

function Base.convert(::Type{GroupedArray{T,N}}, a::AbstractArray) where {T, N}
    convert(GroupedArray{T, N}, GroupedArray(a))
end
Base.convert(::Type{GroupedArray}, a::AbstractArray) = GroupedArray(a)

##############################################################################
##
## Data API
##
##############################################################################

DataAPI.refarray(g::GroupedArray) = g.groups
DataAPI.levels(g::GroupedArray) = 1:g.ngroups
DataAPI.refvalue(g::GroupedArray, ref::Integer) = ref > 0 ? ref : missing

# refpool is such that refpool[refarray[i]] = x
struct GroupedRefPool{T <: Union{Int, Missing}} <: AbstractVector{T}
	ngroups::Int
end
Base.size(x::GroupedRefPool{T}) where T = (x.ngroups + (T >: Missing),)
Base.axes(x::GroupedRefPool{T}) where T = ((1-(T >: Missing)):x.ngroups,)
Base.IndexStyle(::Type{<: GroupedRefPool}) = Base.IndexLinear()
@inline function Base.getindex(x::GroupedRefPool{T}, i::Integer) where T
    @boundscheck checkbounds(x, i)
    T >: Missing && i == 0 ? missing : i
end
Base.allunique(x::GroupedRefPool) = true

DataAPI.refpool(g::GroupedArray{T}) where {T} = GroupedRefPool{T}(g.ngroups)
# invrefpool is such that invrefpool[refpool[x]] = x.
# In words, for each element of refpool, it associates the corresponding index in the pool
# here, this gives
#  missing -> 0
#  i -> i for 1 ≤ i ≤ ngroups
struct GroupedInvRefPool{T}
	ngroups::Int
end
@inline Base.haskey(x::GroupedInvRefPool{T}, ::Missing) where {T} = T >: Missing
@inline Base.haskey(x::GroupedInvRefPool, v::Integer) = 1 <= v <= x.ngroups
@inline function Base.getindex(x::GroupedInvRefPool{T}, ::Missing) where {T}
	@boundscheck T >: Missing
	0
end
@inline function Base.getindex(x::GroupedInvRefPool, i::Integer)
	@boundscheck 1 <= i <= x.ngroups
	i
end
@inline Base.get(x::GroupedInvRefPool{T}, ::Missing, default) where {T} = T >: Missing ? 0 : default
@inline Base.get(x::GroupedInvRefPool, i::Integer, default) = 1 <= v <= x.ngroups ? i : default
DataAPI.invrefpool(g::GroupedArray{T}) where {T} = GroupedInvRefPool{T}(g.ngroups)

export GroupedArray, GroupedVector, GroupedMatrix
end # module


"""
    module NewtonsCradle

Model of [Newtons' Cradle](https://en.wikipedia.org/wiki/Newton%27s_cradle).

The balls of the Newton Cradle are initially placed in such a way, that there is a small
distance between the balls, so the balls are **not** touching each other. This distance was
determined (by trial-and-error) in such a way, that collisions between the balls
occur in sequence (so if ball1 hits ball2, then ball3 is hit by ball2 only after the
contact between ball1 and ball2 has already ended; this can be easily determined by
setting `simulate!(..., log=true)`). Only in this case, the simulation matches the
real experiment. If the balls are directly touching each other, multiple contacts
occur at the same time instant and the simulation result is then no longer according
to reality.

See also:

M. Payr (2008): **[An Experimental and Theoretical Study of Perfect Multiple
                Contact Collisions in Linear Chains of Balls.](http://www.zfm.ethz.ch/e/dynamics/payr_collisions.htm)**
                Dissertation, ETH Zürich.
"""
module NewtonsCradle

using ModiaLang
import Modia3D
using  Modia3D.ModiaInterface

vmatVisual = VisualMaterial(color="LightBlue", transparency=0.1)
vmatSolids   = VisualMaterial(color="Red"      , transparency=0.0)

diameter   = 0.02
lengthWire = 0.15
rsmall     = diameter/100

Lx = diameter/2
Ly = 5*diameter
Lz = Lx

Pendulum = Model(
    wire = Object3D(feature=Visual(shape=Cylinder(axis=3, diameter=diameter/5, length=lengthWire),
                                   visualMaterial=vmatVisual)),
    frame1 = Object3D(parent=:wire,
                      translation=[0.0, 0.0, lengthWire/2]),
    frame2 = Object3D(parent=:wire,
                      translation=[0.0, 0.0, -lengthWire/2]),
    cylinder = Object3D(parent=:frame1,
                        feature=Visual(shape=Cylinder(axis=1, diameter=0.6*Lx, length=1.4*Lx),
                                       visualMaterial=vmatSolids)),
    sphere = Object3D(parent=:frame2,
                      feature=Solid(shape=Sphere(diameter=diameter),
                                    solidMaterial="Steel",
                                    visualMaterial=vmatSolids,
                                    contactMaterial="BilliardBall", collision=true)),
)

Cradle = Model(
    gravField = UniformGravityField(g=9.81, n=[0, 0,-1]),
    world = Object3D(feature=Scene(gravityField=:gravField,
                                   nominalLength=Ly,
                                   enableContactDetection=true,
                                   elasticContactReductionFactor=1.0,
                                   animationFile="NewtonsCradle.json",
                                   visualizeBoundingBox=true)),
    box = Object3D(parent=:world,
                   feature=Visual(shape=Box(lengthX=Lx, lengthY=Ly, lengthZ=Lz),
                                  visualMaterial=vmatVisual)),

    frame1 = Object3D(parent=:box, translation=[0.0, -2*(diameter+rsmall), 0.0]),
    frame2 = Object3D(parent=:box, translation=[0.0,   -(diameter+rsmall), 0.0]),
    frame3 = Object3D(parent=:box, translation=[0.0,                  0.0, 0.0]),
    frame4 = Object3D(parent=:box, translation=[0.0,    (diameter+rsmall), 0.0]),
    frame5 = Object3D(parent=:box, translation=[0.0,  2*(diameter+rsmall), 0.0]),

    pendulum1 = Pendulum,
    pendulum2 = Pendulum,
    pendulum3 = Pendulum,
    pendulum4 = Pendulum,
    pendulum5 = Pendulum,

    rev1 = Revolute(obj1=:frame1, obj2=:(pendulum1.frame1), axis=1, canCollide=true, phi=Var(init=-pi/3)),
    rev2 = Revolute(obj1=:frame2, obj2=:(pendulum2.frame1), axis=1, canCollide=true),
    rev3 = Revolute(obj1=:frame3, obj2=:(pendulum3.frame1), axis=1, canCollide=true),
    rev4 = Revolute(obj1=:frame4, obj2=:(pendulum4.frame1), axis=1, canCollide=true),
    rev5 = Revolute(obj1=:frame5, obj2=:(pendulum5.frame1), axis=1, canCollide=true)
)

newtonsCradle = @instantiateModel(buildModia3D(Cradle), unitless=true, log=false, logStateSelection=false, logCode=false)

stopTime = 5.0
tolerance = 1e-8
requiredFinalStates = [-1.180791527568564, -1.033934458288708, 0.04982982739055273, 0.018466029517559813, 0.049844032881367094, 0.018471344490359712, 0.04986127142907837, 0.018476465719447804, 0.04987561416137295, 0.018481494292598495]
simulate!(newtonsCradle, stopTime=stopTime, tolerance=tolerance, log=true, logStates=true, logEvents=true, requiredFinalStates=requiredFinalStates)

@usingModiaPlot
plot(newtonsCradle, ["rev1.phi" "rev5.phi"; "rev1.w" "rev5.w"], figure=1)

end


import FixedPointNumbers: N0f8

#defined in Plotting/palettes.jl
colormapdict = Dict{String,Vector{RGB{Float32}}}(

	 "Spectral"  =>  Spectral,
	 "RdYlGn"    =>  RdYlGn,
	 "GnYlRd"    =>  reverse(RdYlGn),
	 "RdBu"      =>  RdBu,
	 "PiYG"      =>  PiYG,
	 "PRGn"      =>  PRGn,
	 "RdYlBu"    =>  RdYlBu,
	 "BuYlRd"    =>  BuYlRd,#reverse(RdYlBu),
	 "BrBG"      =>  BrBG,
	 "RdGy"      =>  RdGy,
	 "GyRd"      =>  reverse(RdGy),
	 "PuOr"      =>  PuOr,
	 "Arnold"    =>  reverse(PiYG),
	 "BuRd"      =>  reverse(RdBu),
	 "GnYlRd"    =>  reverse(RdYlGn),
	 "Paired"    =>  Paired,
	 "Purples"   =>  Purples,
	 "BlPu"		 =>  BlPu,
	 "BkOr"   =>  BkOr,
	 "GrBl" => GrBl,
	 "distgclrs" => distgclrs,
	 "bone"		=>	bone,
	 "bone_rev"		=>	reverse(bone),
	 "YG" => YG)

abstract type AbstractColorMap
end

mutable struct ColorMap <: AbstractColorMap
	 map::Vector{RGB{Float32}}
	 function ColorMap(str::String)
		  clrmap = new()
		  clrmap.map = generate(colormapdict[str],1:256)
		  return clrmap
	 end
end

using Nullables

mutable struct BoundedColorMap
	 clrmap::ColorMap
	 mindta::Nullable{Float32}
	 maxdta::Nullable{Float32}
	 function BoundedColorMap(clrmap::ColorMap)
		  return new(clrmap,Nullable{Float32}(),Nullable{Float32}())
	 end
end

import Statistics.mean
function mean(clrs::Vector{RGB{Float32}})
	 r,g,b,a = zeros(Float32,4)
	 n = length(clrs)
	 for clr in clrs
		  r+=clr.r
		  g+=clr.g
		  b+=clr.b
	 end
	 r/=n;g/=n;b/=n
	 return RGB{Float32}(r,g,b)
end

function generate(colormap::Vector{RGB{Float32}},data::AbstractVector)
	 inds = map(x->Int(floor(x)),linspace(1,length(data),length(colormap)))
	 clrs = map(mean,map(x->colormap[find(y->y==x,inds)],unique(inds)))
	 inds[1] += sum(map(i->inds[i+1]-inds[i],1:(length(inds)-1)))-length(data)
	 return length(data) == length(clrs) ? clrs : vcat(map(i->(inds[i+1]-inds[i]) > 1 ? linspace(clrs[i],clrs[i+1],inds[i+1]-inds[i]) : clrs[i] ,1:(length(inds)-1))...)
end

mutable struct JetColorMap <: AbstractColorMap
	 map::Vector{RGBA{Float32}}
	 function JetColorMap()
		  cmap = Array(RGBA{Float32}, 256)
		  for i = 0:255
				n=4*i/256
				cmap[i+1] = RGBA(min(max(min(n-1.5,-n+4.5),0.),1.),
									 min(max(min(n-0.5,-n+3.5),0.),1.),
									 min(max(min(n+0.5,-n+2.5),0.),1.),1.)
		  end
		  return new(cmap)
	 end
end

mutable struct HSVColorMap <: AbstractColorMap
	 map::Vector{RGBA{Float32}}
	 function HSVColorMap()
		  return new(linspace(Colors.HSV(0,1,1),Colors.HSV(330,1,1),256))
	 end
end

function apply!(cm::BoundedColorMap, data::Union{Array{Float32,1},Array{Float64,1},AbstractVector{Int}})
	 cm.mindta = minimum(data)
	 cm.maxdta = maximum(data)
	 return apply(cm.clrmap,data)
end

function apply(cm::AbstractColorMap,data::Union{Array{Float64,1},Array{Float32,1}})
	 color = Array{RGBA{Float32}}(length(data))
	 mindta = minimum(data)
	 maxdta = maximum(data)
	 interv = maxdta-mindta
	 if isapprox(abs(interv),0.,atol=1e-10)
		  fill!(color,cm.map[floor(Int,length(cm.map)/2)])
	 else
		  pos = 0.0
		  for i = 1:length(data)
				pos = (data[i]-mindta)/interv * 255.0 + 1.0
				clr_1 = cm.map[floor(Int64,pos)]
				clr_2 = cm.map[ceil(Int64,pos)]
				weight = pos%1
				clr = weighted_color_mean(weight,clr_2,clr_1)
				#clr = (1-weight)*clr_1+weight*clr_2
				color[i] = clr
		  end
	 end
	 return color
end

function apply(cm::AbstractColorMap,data::AbstractVector{Int})
	 color = Array{RGBA{Float32}}(length(data))
	 mindta = minimum(data)
	 maxdta = maximum(data)
	 interv = maxdta-mindta
	 if interv < 1
		  fill!(color,cm.map[floor(Int,length(cm.map)/2)])
	 else
		  pos = 0.0
		  for i = 1:length(data)
				pos = (data[i]-mindta)/interv * 255.0 + 1.0
				clr_1 = cm.map[floor(Int64,pos)]
				clr_2 = cm.map[ceil(Int64,pos)]
				weight = pos%1
				clr = weighted_color_mean(weight,clr_2,clr_1)
				#clr = (1-weight)*clr_1+weight*clr_2
				color[i] = clr
		  end
	 end
	 return color
end


module StateVectorGates

    export SV_Gate, SV_OneQubitGate, SV_TwoQubitGate
    abstract type SV_Gate end
    abstract type SV_OneQubitGate <: SV_Gate end
    abstract type SV_TwoQubitGate <: SV_Gate end

    struct X <: SV_OneQubitGate
        _target::Int64
    end

    struct Y <: SV_OneQubitGate
        _target::Int64
    end

    struct Z <: SV_OneQubitGate
        _target::Int64
    end

    struct H <: SV_OneQubitGate
        _target::Int64
    end

    struct T <: SV_OneQubitGate
        _target::Int64
    end

    struct S <: SV_OneQubitGate
        _target::Int64
    end

    struct RX <: SV_OneQubitGate
        _target::Int64
        _theta::Float64
    end

    struct RY <: SV_OneQubitGate
        _target::Int64
        _theta::Float64
    end

    struct RZ <: SV_OneQubitGate
        _target::Int64
        _theta::Float64
    end

    struct U3 <: SV_OneQubitGate
        _target::Int64
        _theta::Float64
        _phi::Float64
        _lambd::Float64
    end

    struct CZ <: SV_TwoQubitGate
        _control::Int64
        _target::Int64
    end

    struct CX <: SV_TwoQubitGate
        _control::Int64
        _target::Int64
    end

    struct CP <: SV_TwoQubitGate
        _control::Int64
        _target::Int64
    end

    struct CRX <: SV_TwoQubitGate
        _control::Int64
        _target::Int64
        _theta::Float64
    end

    struct CRY <: SV_TwoQubitGate
        _control::Int64
        _target::Int64
        _theta::Float64
    end

    struct CRZ <: SV_TwoQubitGate
        _control::Int64
        _target::Int64
        _theta::Float64
    end

    mutable struct M <: SV_OneQubitGate
        _target::Int64
        _cregidx::Int64
        _result::Int64
    end

    # export all SV_Gate and TwoQubitGate
    for n in names(@__MODULE__; all=true)
        if Base.isidentifier(n) && n ∉ (Symbol(@__MODULE__), :eval, :include)
            if @eval typeof($n) <: DataType
                if @eval $n <: SV_Gate
                    @eval export $n
                end
            end
        end
    end
end


import Base: merge, hcat, vcat, map

###### merge ####################

function _merge_outer(::Type{IndexType}, ta1::TimeArray{T,N,D}, ta2::TimeArray{T,M,D}, padvalue, meta) where {IndexType,T,N,M,D}
    timestamps, new_idx1, new_idx2 = sorted_unique_merge(IndexType, timestamp(ta1), timestamp(ta2))
    vals = fill(convert(T, padvalue), (length(timestamps), length(colnames(ta1)) + length(colnames(ta2))))
    insertbyidx!(vals, values(ta1), new_idx1)
    insertbyidx!(vals, values(ta2), new_idx2, size(values(ta1), 2))
    TimeArray(timestamps, vals, [colnames(ta1); colnames(ta2)], meta; unchecked = true)
end

"""
    merge(ta1::TimeArray{T}, ta2::TimeArray{T}, [tas::TimeArray{T}...];
          method = :inner, colnames = [...], padvalue = NaN)

Merge several `TimeArray`s along with the time index.

## Argument

- `method::Symbol`: `:inner`, `:outer`, `:left` or `:right`.
"""
function merge(ta1::TimeArray{T,N,D}, ta2::TimeArray{T,M,D};
               method::Symbol = :inner, colnames::Vector = Symbol[], meta = nothing,
               padvalue=NaN) where {T,N,M,D}

    if colnames isa Vector{<:AbstractString}
        @warn "`merge(...; colname::Vector{<:AbstractString})` is deprecated, " *
              "use `merge(...; colnames=Symbol.(colnames))` instead."
        colnames = Symbol.(colnames)
    end

    meta = if _meta(ta1) == _meta(ta2) && meta ≡ nothing
        _meta(ta1)
    elseif typeof(_meta(ta1)) <: AbstractString && typeof(_meta(ta2)) <: AbstractString && meta ≡ nothing
        string(_meta(ta1), "_", _meta(ta2))
    else
        meta
    end

    if method == :inner

        idx1, idx2 = overlap(timestamp(ta1), timestamp(ta2))
        vals = [values(ta1[idx1]) values(ta2[idx2])]
        ta = TimeArray(timestamp(ta1[idx1]), vals, [_colnames(ta1); _colnames(ta2)], meta; unchecked = true)

    elseif method == :left

        new_idx2, old_idx2 = overlap(timestamp(ta1), timestamp(ta2))
        right_vals = fill(convert(T, padvalue), (length(ta1), length(_colnames(ta2))))
        insertbyidx!(right_vals, values(ta2), new_idx2, old_idx2)
        ta = TimeArray(timestamp(ta1), [values(ta1) right_vals], [_colnames(ta1); _colnames(ta2)], meta; unchecked = true)

    elseif method == :right

        ta = merge(ta2, ta1, method = :left; padvalue = padvalue)
        ncol2 = length(_colnames(ta2))
        vals = [values(ta)[:, (ncol2+1):end] values(ta)[:, 1:ncol2]]
        ta = TimeArray(timestamp(ta), vals, [_colnames(ta1); _colnames(ta2)], meta; unchecked = true)

    elseif method == :outer

        ta = if (length(timestamp(ta1)) + length(timestamp(ta2))) > typemax(Int32)
            _merge_outer(Int64, ta1, ta2, padvalue, meta)
        else
            _merge_outer(Int32, ta1, ta2, padvalue, meta)
        end

    else
        throw(ArgumentError(
            "merge method must be one of :inner, :left, :right, :outer"))
    end

    return rename!(ta, colnames)

end

merge(x::TimeArray{T}, y::TimeArray{T}, z::TimeArray{T}, a::TimeArray{T}...;
      colnames::Vector = Symbol[], kw...) where {T} =
    rename!(merge(merge(x, y; kw...), z, a...; kw...), colnames)

# hcat ##########################

function hcat(x::TimeArray, y::TimeArray)
    tsx = timestamp(x)
    tsy = timestamp(y)

    if length(tsx) != length(tsx) || tsx != tsy
        throw(DimensionMismatch(
            "timestamps not consistent, please checkout `merge`."))
    end

    m = ifelse(meta(x) == meta(y), meta(x), nothing)

    TimeArray(tsx, [values(x) values(y)], [colnames(x); colnames(y)], m)
end

hcat(x::TimeArray, y::TimeArray, zs::Vararg{TimeArray}) =
    hcat(hcat(x, y), zs...)

# collapse ######################

# accessors functions
# https://github.com/JuliaLang/julia/blob/b617e8d3a77e49cd5625ca663af88e40f7796f15/stdlib/Dates/src/accessors.jl#L50

# `year` doesn't need this wrapper
for (F, T, P) ∈ ((:quarter,     :TimeType,               :Quarter),
                 (:month,       :TimeType,               :Month),
                 (:week,        :TimeType,               :Week),
                 (:day,         :TimeType,               :Day),
                 (:hour,        :(Union{DateTime,Time}), :Hour),
                 (:minute,      :(Union{DateTime,Time}), :Minute),
                 (:second,      :(Union{DateTime,Time}), :Second),
                 (:millisecond, :(Union{DateTime,Time}), :Millisecond),
                 (:microsecond, :Time,                   :Microsecond),
                 (:nanosecond,  :Time,                   :Nanosecond))
    if F === :quarter
        (VERSION < v"1.6") && continue
        F = :(Dates.quarter)
    end
    @eval let
        global collapse
        wrapper(x::$T) = floor(x, $P(1))
        collapse(ta::TimeArray, ::typeof($F), f::Function, g::Function = f; kw...) =
            collapse(ta, wrapper, f, g; kw...)
    end
end

function collapse(ta::TimeArray, period::Function, timestamp::Function,
                  value::Function = timestamp)

    isempty(ta) && return ta

    m, n = length(ta), length(colnames(ta))
    ts   = _timestamp(ta)
    val  = _values(ta)
    idx  = UnitRange{Int}[]
    sizehint!(idx, m)

    t₀ = period(ts[1])
    j  = 1
    for i in 1:m-1
        t₁ = period(ts[i+1])
        t₀ == t₁ && continue
        push!(idx, j:i)
        j = i + 1
        t₀ = t₁
    end
    push!(idx, j:m)

    ts′  = [timestamp(@view(ts[i])) for i ∈ idx]
    val′ = if n == 1
        [value(@view(val[i])) for i ∈ idx]
    else
        [value(@view(val[i, k])) for i ∈ idx, k ∈ 1:n]
    end

    TimeArray(ts′, val′, colnames(ta), meta(ta))
end

collapse(ta::TimeArray, p::Period, t::Function, v::Function = t; kw...) =
    collapse(ta, x -> floor(x, p), t, v; kw...)

# vcat ######################

"""
    $(SIGNATURES)

Concatenate two ``TimeArray`` into single object.

If there are duplicated timestamps, we will keep order as the function input.

```julia-repl
julia> a = TimeArray([Date(2015, 10, 1), Date(2015, 10, 2), Date(2015, 10, 3)], [1, 2, 3]);

julia> b = TimeArray([Date(2015, 10, 2), Date(2015, 10, 3)], [4, 5]);

julia> [a; b]
5×1 TimeArray{Int64,1,Date,Array{Int64,1}} 2015-10-01 to 2015-10-03
│            │ A     │
├────────────┼───────┤
│ 2015-10-01 │ 1     │
│ 2015-10-02 │ 2     │
│ 2015-10-02 │ 4     │
│ 2015-10-03 │ 3     │
│ 2015-10-03 │ 5     │
```
"""
function vcat(tas::TimeArray...)
    # Check all meta fields are identical.
    prev_meta = meta(tas[1])
    for ta in tas
        if meta(ta) != prev_meta
            throw(ArgumentError("metadata doesn't match"))
        end
    end

    # Check column names are identical.
    prev_colnames = colnames(tas[1])
    for ta in tas
        if colnames(ta) != prev_colnames
            throw(ArgumentError("column names don't match"))
        end
    end

    # Concatenate the contents.
    ts  = vcat([timestamp(ta) for ta in tas]...)
    val = vcat([values(ta) for ta in tas]...)

    order = sortperm(ts)
    if ndims(tas[1]) == 1 # Check for 1D to ensure values remains a 1D vector.
        return TimeArray(ts[order], val[order], prev_colnames, prev_meta)
    else
        return TimeArray(ts[order], val[order, :], prev_colnames, prev_meta)
    end
end

# map ######################

@generated function map(f, ta::TimeArray{T,N}) where {T,N}
    input_val  = (N == 1) ? :(values(ta)[i]) : :(vec(values(ta)[i, :]))
    output_val = (N == 1) ? :(vals[i]) : :(vals[i, :])

    output_vals = (N == 1) ? :(vals[order]) : :(vals[order, :])

    quote
        ts   = similar(timestamp(ta))
        vals = similar(values(ta))

        for i in eachindex(ta)
            @inbounds ts[i], $output_val = f(timestamp(ta)[i], $input_val)
        end

        order = sortperm(ts)
        TimeArray(ts[order], $output_vals, colnames(ta), meta(ta))
    end
end


#
# Copyright (c) 2021 Tobias Thummerer, Lars Mikelsons, Josef Kircher
# Licensed under the MIT license. See LICENSE file in the project root for details.
#

using FMI.FMIImport.FMICore: fmi2Real

myFMU = fmiLoad("SpringPendulum1D", ENV["EXPORTINGTOOL"], ENV["EXPORTINGVERSION"])
comp = fmiInstantiate!(myFMU; loggingOn=false)
@test comp != 0

# choose FMU or FMUComponent
fmuStruct = nothing
envFMUSTRUCT = ENV["FMUSTRUCT"]
if envFMUSTRUCT == "FMU"
    fmuStruct = myFMU
elseif envFMUSTRUCT == "FMUCOMPONENT"
    fmuStruct = comp
end
@assert fmuStruct != nothing "Unknwon fmuStruct, environment variable `FMUSTRUCT` = `$envFMUSTRUCT`"

@test fmiSetupExperiment(fmuStruct) == 0
@test fmiEnterInitializationMode(fmuStruct) == 0
@test fmiExitInitializationMode(fmuStruct) == 0

targetValues = [[0.0, -10.0], [1.0, 0.0]]
dir_ders_buffer = zeros(fmi2Real, 2)
sample_ders_buffer = zeros(fmi2Real, 2, 1)
for i in 1:fmiGetNumberOfStates(myFMU)

    if fmiProvidesDirectionalDerivative(myFMU)
        # multi derivatives calls
        sample_ders = fmiSampleDirectionalDerivative(fmuStruct, myFMU.modelDescription.derivativeValueReferences, [myFMU.modelDescription.stateValueReferences[i]])
        fmiSampleDirectionalDerivative!(fmuStruct, myFMU.modelDescription.derivativeValueReferences, [myFMU.modelDescription.stateValueReferences[i]], sample_ders_buffer)

        @test sum(abs.(sample_ders[:,1] - targetValues[i])) < 1e-3
        @test sum(abs.(sample_ders_buffer[:,1] - targetValues[i])) < 1e-3

        dir_ders = fmiGetDirectionalDerivative(fmuStruct, myFMU.modelDescription.derivativeValueReferences, [myFMU.modelDescription.stateValueReferences[i]])
        @test fmiGetDirectionalDerivative!(fmuStruct, myFMU.modelDescription.derivativeValueReferences, [myFMU.modelDescription.stateValueReferences[i]], dir_ders_buffer) == 0

        @test sum(abs.(dir_ders - targetValues[i])) < 1e-3
        @test sum(abs.(dir_ders_buffer - targetValues[i])) < 1e-3

        # single derivative call
        dir_der = fmiGetDirectionalDerivative(fmuStruct, myFMU.modelDescription.derivativeValueReferences[1], myFMU.modelDescription.stateValueReferences[1])
        @test dir_der == targetValues[1][1]

    else
        @warn "Skipping directional derivative testing, FMU from $(ENV["EXPORTINGTOOL"]) doesn't support directional derivatives."
    end

end

fmiUnload(myFMU)


function iterative_greedy_construction(tuning_run::Run,
                                       channel::RemoteChannel;
                                       cutoff::Integer = 10_000)
    key_set = collect(keys(tuning_run.starting_point.parameters))

    function iterate(tuning_run::Run)
        for key in key_set
            result = greedy_construction(tuning_run,
                                         key,
                                         cutoff = cutoff)
        end
        shuffle!(key_set)

        return result
    end

    technique(tuning_run,
              channel,
              iterate,
              name = "Iterative Greedy Construction")
end


#=

Macro interface is inspired by InfiniteOpt.jl. Many of the code structures are
reused. Transitioning this over to a fully InfiniteOpt.jl in the future is desirable
but most of our use cases involve nonlinear functions so that InfiniteOpt.jl
nonlinear support is gating.

The general idea here is pretty simple. When we add a variable to the model
the user specifies whether or not it is a decision variable or an uncertain
variable we then classify constraints. See https://github.com/PSORLab/EAGO.jl/issues/61
for a discussion of the roadmap here.
=#

#=
@enum(SIPCons, DECISION, UNCERTAIN, SEMIINFINITE, SIPNOTSET)
const MOIU_AUTO_CACHE = MOIU.CachingOptimizerMode.AUTOMATIC

Base.@kwdef mutable struct SIPData
    p::Dict{JuMP.VariableRef,Bool}                  = Dict{JuMP.VariableRef,Bool}()
    constr_type::Dict{JuMP.ConstraintRef, SIPCons}  = Dict{JuMP.ConstraintRef, SIPCons}()
    nl_constraint_type::Dict{Int, SIPCons}          = Dict{Int, SIPCons}()
    nl_expression_type::Dict{Int, SIPCons}          = Dict{Int, SIPCons}()
    model_decision::Model                           = Model(caching_mode = MOIU_AUTO_CACHE)
    model_uncertain::Model                          = Model((aching_mode = MOIU_AUTO_CACHE)
    model_semiinfinite::Model                       = Model(caching_mode = MOIU_AUTO_CACHE)
end
uncertain_variable_num(d::SIPData) = any(x -> x, d.p)

function _copy_nlp_objective!(out::JuMP.Model, m::JuMP.Model, nlp_data::JuMP._NLPData)
    if nlp_data.nlobj !== nothing
        # TODO CHECK THAT THE FORMULATION IS PEACHY
        obj_expr = MOI.objective_expr(nlp_new.evaluator)
        # obj_expr_jref = substitute moi variables to jump refs...
        @NLobjective(m.ext[:sip].model_decision, Min, obj_expr_jref)
    end
    return nothing
end

function _copy_nlp_constraints!(m::JuMP.Model, nlp_data::JuMP._NLPData)
    for i = 1:length(nlp_data.nlconstr)
        con_expr = MOI.constraint_expr(nlp_data.evaluator, i)
    end
    return nothing
end

"""

Populates variables and constraints into one of three models based on classification:
1) store decision variables and constraints containing only decision variables in
`m.model_decision`, 2) store uncertain variables and constraints containing only
uncertain variables in `m.model_uncertain`, 3) store all variables and all
constraints containing both decision variables to `m.model_semiinfinite`.
"""
function initialize_pure_models!(m::JuMP.Model)
    if haskey(m.ext, :sip)

        # extract sip data and storage models
        sip_data = _get_sip_data(m)::SIPData
        m_dec = sip_data.model_decision
        m_unc = sip_data.model_uncertain
        m_sip = sip_data.model_semiinfinite

        # get index map
        indx_map_dec = MOI.copy_to(backend(m_dec), backend(m), copy_names = true)
        indx_map_unc = MOI.copy_to(backend(m_unc), backend(m), copy_names = true)
        indx_map_sip = MOI.copy_to(backend(m_sip), backend(m), copy_names = true)

        # copy extension data
        for (key, data) in m.ext
            if key != :sip
                m_dec.ext[key] = copy_extension_data(data, m_dec, m)
                m_unc.ext[key] = copy_extension_data(data, m_unc, m)
                m_sip.ext[key] = copy_extension_data(data, m_sip, m)
            end
        end

        # copy objects (TODO need to check on SIP, DEC, UNC)
        reference_map = ReferenceMap(new_model, index_map)
        for (name, value) in object_dictionary(model)
            new_model[name] = getindex.(reference_map, value)
        end

        if m.nlp_data !== nothing
            nlp_data = deepcopy(m.nlp_data)
            nlp_data.evaluator = NLPEvaluator(m)
            MOI.initialize(nlp_data.evaluator, Symbol[:ExprGraph])
            copy_nlp_objective!(m, nlp_data)
            copy_nlp_expressions!(m, nlp_data)
            copy_nlp_constraints!(m, nlp_data)
        end

        for (vi, is_uncertain) in m.ext.p
            model_semiinfinite
            model_decision_
    end
    return nothing
end

function initialize_sip_data(m::JuMP.Model)
    m.ext[:sip] = SIPData()
end

function sip_optimizehook(m::JuMP.Model; kwargs...)
    data = _getsipdata(m)::SIPData
    if uncertain_variable_num(data) == 0.0
        ret = JuMP.optimize!(m::JuMP.Model, ignore_optimize_hook = true, kwargs...)
    else
        initialize_pure_models!(m)
        ret = sip_solve(data)
    end
    return ret
end

function enable_semiinfinite(m::JuMP.Model)
    haskey(m.ext, :sip) && error("Model already has semiinfinite programs enabled")
    initialize_sip_data(m)
    JuMP.set_optimize_hook(m, sip_optimizehook)
    return nothing
end

function ModelWithSIP(args...; kwargs...)
    m = JuMP.Model(args...; kwargs...)
    set_optimizer(SIPOptimizer)
    enable_semiinfinite(m)
    return m
end

"""
    @uncertain_variable

Add an *anonymous* variable to `m::SIPModel` described by the keyword
arguments `kw_args` and returns the variable.

    @uncertain_variable(m, expr, args..., kw_args...)

This macro simply denotes the variable as belonging to the set of uncertain
variables and then performs `@variable(m, expr, args..., kw_args...)` to
add the variable to the basic JuMP model for storage. All JuMP syntax is
supported.
"""
macro uncertain_variable(args...)
    esc(quote
        vi = @variable($(args...))
        $(args[1]).ext[:sip].p[vi] = true
        vi
        end)
end

"""
    @decision_variable

Add an *anonymous* variable to `m::SIPModel` described by the keyword
arguments `kw_args` and returns the variable.

    @decision_variable(m, expr, args..., kw_args...)

This macro simply denotes the variable as belonging to the set of decision
variables and then performs `@variable(m, expr, args..., kw_args...)` to
add the variable to the basic JuMP model for storage. All JuMP syntax is
supported.
"""
macro decision_variable(args...)
    esc(quote
        vi = @variable($(args...))
        $(args[1]).ext[:sip].p[vi] = false
        vi
        end)
end
=#


using Ghost
using Ghost: Tape, Input, Constant, mkcall, Variable, V

struct ONNXCtx
    name2var::Dict{String, Variable}
    backends::Vector{Symbol}
    exec::Bool
end

ONNXCtx(backends; exec=true) = ONNXCtx(Dict(), backends, exec)
ONNXCtx(;exec=true) = ONNXCtx(Dict(), [:ONNX], exec)

# TODO: implement rebind_context!()

"""
    getindex(tape::Tape{ONNXCtx}, onnx_name::String)

Get operation on the tape using the name in ONNX graph
"""
Base.getindex(tape::Tape{ONNXCtx}, onnx_name::String) =
    tape[tape.c.name2var[onnx_name]]

###############################################################################
#                               Operations                                    #
###############################################################################

"""
    push_call!(tape::Tape{ONNXCtx}, fn, args...; kwargs)

Shortcut for `push!(tape, mkcall(fn, args..))` also handling
keyword arguments and respecting `ONNXCtx.exec` setting.
"""
function push_call!(tape::Tape{ONNXCtx}, fn, args...; kwargs...)
    kwargs = NamedTuple(kwargs)
    if !isempty(kwargs)
        args = (kwargs, fn, args...)
        fn = Core.kwfunc(fn)
    end
    op = tape.c.exec ? mkcall(fn, args...) : mkcall(fn, args...; val=nothing)
    return push!(tape, op)
end

# A few constants to keep function signatures concise
struct OpConfig{BE, Op} end
const VarVec = Vector{Ghost.Variable}
const AttrDict = Dict{Symbol, Any}

function load_node!(tape::Tape, nd::NodeProto, backend::Symbol)
    args = [tape.c.name2var[name] for name in nd.input]
    attrs = convert(Dict{Symbol, Any}, Dict(nd.attribute))
    conf = OpConfig{backend, Symbol(nd.op_type)}()
    try
        out = load_node!(tape, conf, args, attrs)
        ismissing(out) && return out
        if out isa Tuple
            for i=1:length(nd.output)
                tape.c.name2var[nd.output[i]] = out[i]
            end
        else
            tape.c.name2var[nd.output[1]] = out
        end
    catch
        @error "Error while loading node $nd"
        rethrow()
    end
end

function load_node!(tape::Tape, ::OpConfig{:ONNX, :Gemm}, args::VarVec, attrs::AttrDict)
    if (length(args) == 2 && get(attrs, :alpha, 1) == 1 &&
        get(attrs, :transA, 0) == 0 && get(attrs, :transB, 0) == 0)
        # simplified version: just matrix multiplication
        # note: arguments are swapped to account for row-major arrays
        return push_call!(tape, *, args[2], args[1])
    else
        # complete GEMM version
        kw = rename_keys(attrs, Dict(
            :transA => :tA,
            :transB => :tB,
            :alpha => :α,
            :beta => :β
        ))
        return push_call!(tape, onnx_gemm, args...; kw...)
    end
end

function load_node!(tape::Tape, ::OpConfig{:ONNX, :Conv}, args::VarVec, attrs::AttrDict)
    kw = from_onnx_conv(attrs) |> NamedTuple
    return push_call!(tape, conv, args...; kw...)
end

function load_node!(tape::Tape, ::OpConfig{:ONNX, :MaxPool}, args::VarVec, attrs::AttrDict)
    kw = from_onnx_conv(attrs; pooling=true) |> NamedTuple
    return push_call!(tape, maxpool, args[1]; kw...)
end

function load_node!(tape::Tape, ::OpConfig{:ONNX, :GlobalAveragePool}, args::VarVec, attrs::AttrDict)
    return push_call!(tape, global_average_pool, args...)
end

function load_node!(tape::Tape, ::OpConfig{:ONNX, :Flatten}, args::VarVec, attrs::AttrDict)
    return push_call!(tape, onnx_flatten, args...; attrs...)
end

function load_node!(tape::Tape, ::OpConfig{:ONNX, :Add}, args::VarVec, attrs::AttrDict)
    return push_call!(tape, add, args...)
end

function load_node!(tape::Tape, ::OpConfig{:ONNX, :Mul}, args::VarVec, attrs::AttrDict)
    return push_call!(tape, mul, args...)
end

function load_node!(tape::Tape, ::OpConfig{:ONNX, :Relu}, args::VarVec, attrs::AttrDict)
    return push_call!(tape, relu, args[1])
end

function load_node!(tape::Tape, ::OpConfig{:ONNX, :MatMul}, args::VarVec, attrs::AttrDict)
    A_ndims = ndims(args[1]._op.val)
    B_ndims = ndims(args[2]._op.val)
    if A_ndims == 2 && B_ndims == 2
        return push_call!(tape, *, args[2], args[1])
    elseif A_ndims in (2, 3) && B_ndims in (2, 3)
        return push_call!(tape, NNlib.batched_mul, args[2], args[1])
    else
        error("MatMul with arrays of $A_ndims and $B_ndims is not implemented yet")
    end
end

function load_node!(tape::Tape, ::OpConfig{:ONNX, :Sigmoid}, args::VarVec, attrs::AttrDict)
    return push_call!(tape, Broadcast.broadcast, NNlib.sigmoid, args...)
end

function load_node!(tape::Tape, ::OpConfig{:ONNX, :BatchNormalization},
        args::VarVec, attrs::AttrDict)
    kw = from_onnx_norm(attrs)
    bn = push_call!(tape, batch_norm, args...; kw...)
    if bn._op.val isa Tuple
        # usual in training mode
        # unpack tuples into calls to getfield
        y = push_call!(tape, getfield, bn, 1)
        μnext = push_call!(tape, getfield, bn, 2)
        σ²next = push_call!(tape, getfield, bn, 3)
        return y, μnext, σ²next
    else
        return bn
    end
end

function load_node!(tape::Tape, ::OpConfig{:ONNX, :Shape}, args::VarVec, attrs::AttrDict)
    # TODO: handle start and end attributes
    return push_call!(tape, size_vector, args[1])
end

function load_node!(tape::Tape, ::OpConfig{:ONNX, :Constant}, args::VarVec, attrs::AttrDict)
    val_attr = first(keys(attrs))
    val = if val_attr == :value
        array(attrs[val_attr])
    else
        error("Don't know how to load constant value from attribute $val_attr")
    end
    return push!(tape, Constant(val))
end

function load_node!(tape::Tape, ::OpConfig{:ONNX, :Gather}, args::VarVec, attrs::AttrDict)
    axis = get(attrs, :axis, 0)
    data = tape[args[1]].val
    dim = ndims(data) - axis
    return push_call!(tape, onnx_gather, args...; dim=dim)
end

function load_node!(tape::Tape, ::OpConfig{:ONNX, :Unsqueeze}, args::VarVec, attrs::AttrDict)
    if length(args) == 2
        # ONNX >= v13
        return push_call!(tape, onnx_unsqueeze, args...)
    elseif length(args) == 1
        # ONNX < v13
        axes = attrs[:axes]
        v_axes = push!(tape, Constant(axes))
        return push_call!(tape, onnx_unsqueeze, args[1], v_axes)
    else
        throw(ArgumentError("Cannot load node from Unsqueeze with $(length(args)) arguments"))
    end
end

function load_node!(tape::Tape, ::OpConfig{:ONNX, :Slice}, args::VarVec, attrs::AttrDict)
    return push_call!(tape, onnx_slice, args...)
end

###############################################################################
#                                    API                                      #
###############################################################################

"""
    load(io::IO, model_args...; backends=[:ONNX], exec::Bool=true)
    load(filename::String, model_args...; backends=[:ONNX], exec::Bool=true)

Load an ONNX model as a Ghost.Tape. The way a particular ONNX node is deserialized is
controlled by methods of [load_node!](@ref) dispatched by backend and node's op_type.

`backends` parameter can be used to customize the loading process.

`exec` parameter instructs the loader to execute every added operation just after
the addition, making the debugging easier. Default is `true`.

See also: [`save!`](@ref)
"""
function load(io::IO, args...; backends=[:ONNX], exec::Bool=true)
    onnx_model = readproto(io, ModelProto());
    g = onnx_model.graph;
    tape = Tape(ONNXCtx(backends; exec=exec))
    # create map of initializers
    init_vals = Dict{String, Any}(init.name => array(init)
        for init in g.initializer)
    # load inputs; if input has init value, take it
    # otherwise take the next available argument value
    arg_idx = 1
    used_init_names = Set([])
    for inp in g.input
        val = get(init_vals, inp.name, missing)
        v = V(0)   # will be overwritten
        if val === missing && exec == true
            @assert(
                arg_idx <= length(args),
                "Neither initializer, nor argument is provided for input $(inp.name)"
            )
            val = args[arg_idx]
            arg_idx += 1
            v = push!(tape, Input(val))
        else
            # convert inputs that also have initializers to constants
            # these are usually model parameters, but may
            v = push!(tape, Constant(val))
        end
        tape.c.name2var[inp.name] = v
        push!(used_init_names, inp.name)
    end
    # load the rest of initilizers as constants
    for init in g.initializer
        name = init.name
        if !in(name, used_init_names)
            val = init_vals[name]
            v = push!(tape, Ghost.Constant(val))
            tape.c.name2var[name] = v
        end
    end
    # load nodes
    for nd in g.node
        success = false
        for backend in tape.c.backends
          if !ismissing(load_node!(tape, nd, backend))
            success = true
            @debug "Loaded $(nd.op_type) using backend $(backend)"
            break
          end
        end
        success || error("Couldn't load node for $(nd.op_type), " *
                         "tried the following backends: $(tape.c.backends)")
    end
    if length(g.output) == 1
        tape.result = Ghost.bound(tape, V(length(tape)))
    else
        # tuple output: we expect tape to contain these outputs as vars  destructured
        # from a multi-ouput op using a sequence of `getfield()` calls
        vars = [tape.c.name2var[o.name] for o in g.output]
        @assert(all(tape[v] isa Call && tape[v].fn == getfield for v in vars),
            "Don't understand this multi-output result of the graph")
        tape.result = tape[vars[1]].args[1]
    end
    return tape
end

function load(filename::String, args...; backends=[:ONNX], exec::Bool=true)
    return open(filename) do io
        load(io, args...; backends=backends, exec=exec)
    end
end


################################################################################
#  Copyright 2020, Marta Vanin, Tom Van Acker                                  #
################################################################################
# PowerModelsDistributionStateEstimation.jl                                    #
# An extention package of PowerModels(Distribution).jl for Static Power System #
# State Estimation.                                                            #
################################################################################
"""
    assign_start_to_variables!(data)
This function gives the measurement value in the data dictionary
as a starting value to its associated variable.
"""
function assign_start_to_variables!(data::Dict{String, Any})
    for (_,meas) in data["meas"]
        if all(i -> i ∈ [_DST.Normal{Float64}, Float64], typeof.(meas["dst"]))
            data[string(meas["cmp"])][string(meas["cmp_id"])]["$(string(meas["var"]))_start"] = _DST.mean.(meas["dst"])
        end
    end
end
"""
    assign_start_to_variables!(data, start_values_source)
This function assigns start values to the problem variables based on a dictionary where they are collected: start_values_source.
This dictionary must have the form of a powerflow solution dictionary.
"""
function assign_start_to_variables!(data::Dict{String, Any}, start_values_source::Dict{String, Any})
    for (_,meas) in data["meas"]
        msr_cmp = string(meas["cmp"])
        cmp_id = string(meas["cmp_id"])
        msr_var = string(meas["var"])
        if haskey(start_values_source["solution"][msr_cmp][cmp_id], msr_var)
            if msr_var != "w"
                data[msr_cmp][cmp_id]["$(msr_var)_start"] = start_values_source["solution"][msr_cmp][cmp_id][msr_var]
            else
                data[msr_cmp][cmp_id]["$(msr_var)_start"] = start_values_source["solution"][msr_cmp][cmp_id][msr_var][1]
            end
        else
            @warn "$(msr_var) is not in $(start_values_source["solution"][msr_cmp][cmp_id]), possible formulation mismatch"
        end
    end
end