text:
struct Production{V,U}
    value::V
    children::Vector{U}
end

struct ProductionDiff <: Diff
    value_diff::Diff
    children_diff::VectorDiff
end

# the production kernel takes a value (of type U)
# and returns a value of type Tuple{V,Vector{U}}

# the reduction kernel takes an input value of tyoe Tuple{V,Vector{W}} and
# returns a value of type W. The Diff type for the Vector{W} is a VectorDiff.

using FunctionalCollections: assoc, dissoc, PersistentHashMap
using DataStructures: PriorityQueue, dequeue!, enqueue!

#################
# recurse trace #
#################

struct RecurseTrace{S,T,U,V,W,X,Y} <: Trace
    gen_fn::GenerativeFunction
    production_traces::PersistentHashMap{Int,S}
    aggregation_traces::PersistentHashMap{Int,T}
    max_branch::Int
    score::Float64
    root_idx::Int
    num_has_choices::Int
end

get_gen_fn(trace::RecurseTrace) = trace.gen_fn

function get_args(trace::RecurseTrace)
    root_arg = get_args(trace.production_traces[trace.root_idx])[1]
    (root_arg, trace.root_idx)
end

function get_retval(trace::RecurseTrace)
    get_retval(trace.aggregation_traces[trace.root_idx])
end

get_score(trace::RecurseTrace) = trace.score

# TODO assumes there is no untraced randomness
project(trace::RecurseTrace, ::EmptySelection) = 0.

##############################
# recurse assignment wrapper #
##############################

struct RecurseTraceChoiceMap <: ChoiceMap
    trace::RecurseTrace
end

get_choices(trace::RecurseTrace) = RecurseTraceChoiceMap(trace)

function Base.isempty(choices::RecurseTraceChoiceMap)
    choices.trace.num_has_choices == 0
end

get_address_schema(::Type{RecurseTraceChoiceMap}) = DynamicAddressSchema()

function get_submap(choices::RecurseTraceChoiceMap,
                      addr::Tuple{Int,Val{:production}})
    idx = addr[1]
    if !haskey(choices.trace.aggregation_traces, idx)
        EmptyChoiceMap()
    else
        get_choices(choices.trace.production_traces[idx])
    end
end

function get_submap(choices::RecurseTraceChoiceMap,
                      addr::Tuple{Int,Val{:aggregation}})
    idx = addr[1]
    if !haskey(choices.trace.aggregation_traces, idx)
        EmptyChoiceMap()
    else
        get_choices(choices.trace.aggregation_traces[idx])
    end
end

function get_submap(choices::RecurseTraceChoiceMap, addr::Pair)
    _get_submap(choices, addr)
end

function has_value(choices::RecurseTraceChoiceMap, addr::Pair)
    _has_value(choices, addr)
end

function get_value(choices::RecurseTraceChoiceMap, addr::Pair)
    _get_value(choices, addr)
end

get_values_shallow(choices::RecurseTraceChoiceMap) = ()

function get_submaps_shallow(choices::RecurseTraceChoiceMap)
    production_iter = (((idx, Val(:production)), get_choices(subtrace))
        for (idx, subtrace) in choices.trace.production_traces)
    aggregation_iter = (((idx, Val(:aggregation)), get_choices(subtrace))
        for (idx, subtrace) in choices.trace.aggregation_traces)
    Iterators.flatten((production_iter, aggregation_iter))
end

# TODO when lightweight Gen functions properly declare their argument and return types, use:
# production_kern::GenerativeFunction{Tuple{V,Vector{U}},S}
# aggregation_kern::GenerativeFunction{W,T}

""""
    Recurse(production_kernel, aggregation_kernel, max_branch,
         ::Type{U}, ::Type{V}, ::Type{W})

Constructor for recurse production and aggregation function.
"""
struct Recurse{S,T,U,V,W,X,Y} <: GenerativeFunction{W,RecurseTrace{S,T,U,V,W,X,Y}}
    production_kern::GenerativeFunction{X,S}
    aggregation_kern::GenerativeFunction{Y,T}
    max_branch::Int
end

function Recurse(production_kernel::GenerativeFunction{X,S},
                 aggregation_kernel::GenerativeFunction{Y,T},
                 max_branch::Int, ::Type{U}, ::Type{V}, ::Type{W}) where {S,T,U,V,W,X,Y}
    Recurse{S,T,U,V,W,X,Y}(production_kernel, aggregation_kernel, max_branch)
end

# TODO
accepts_output_grad(::Recurse) = false

function (gen_fn::Recurse)(args...)
    (_, _, retval) = propose(gen_fn, args)
    retval
end

function get_child(parent::Int, child_num::Int, max_branch::Int)
    @assert child_num >= 1 && child_num <= max_branch
    (parent - 1) * max_branch + child_num + 1
end

function get_child_num(child::Int, max_branch::Int)
    @assert max_branch > 0
    (child - 2) % max_branch + 1
end

function get_parent(child::Int, max_branch::Int)
    @assert max_branch > 0
    child_num = get_child_num(child, max_branch)
    div(child - 1 - child_num, max_branch) + 1
end

function get_num_children(production_trace)
    length(get_retval(production_trace).children)
end

function get_production_input(gen_fn::Recurse{S,T,U,V,W}, cur::Int,
                              production_traces::AbstractDict{Int,S},
                              root_idx::Int, root_production_input::U) where {S,T,U,V,W}
    if cur == root_idx
        return root_production_input
    else
        parent = get_parent(cur, gen_fn.max_branch)
        child_num = get_child_num(cur, gen_fn.max_branch)
        # return type of parent is Tuple{V,Vector{U}}
        @assert haskey(production_traces, parent)
        return get_retval(production_traces[parent]).children[child_num]::U
    end
end

function get_aggregation_input(gen_fn::Recurse{S,T,U,V,W}, cur::Int,
                               production_traces::AbstractDict{Int,S},
                               aggregation_traces::AbstractDict{Int,T}) where {S,T,U,V,W}
    # requires that the corresponding production trace exists (for the v)
    # also requires that the child aggregation traces exist (for the w's)
    # does not require that this aggregation trace already exists
    vinput::V = get_retval(production_traces[cur]).value
    num_children = get_num_children(production_traces[cur])
    winputs::Vector{W} = [
        get_retval(aggregation_traces[get_child(cur, i, gen_fn.max_branch)])
        for i=1:num_children]
    return (vinput, winputs)
end

function get_production_constraints(constraints::ChoiceMap, cur::Int)
    get_submap(constraints, (cur, Val(:production)))
end

function get_aggregation_constraints(constraints::ChoiceMap, cur::Int)
    get_submap(constraints, (cur, Val(:aggregation)))
end

############
# simulate #
############

function simulate(gen_fn::Recurse{S,T,U,V,W,X,Y}, args::Tuple{U,Int}) where {S,T,U,V,W,X,Y}
    (root_production_input::U, root_idx::Int) = args
    production_traces = PersistentHashMap{Int,S}()
    aggregation_traces = PersistentHashMap{Int,T}()
    score = 0.
    num_has_choices = 0

    # production phase
    # does not matter in which order we visit (since children are inserted after parents)
    prod_to_visit = Set{Int}([root_idx])
    while !isempty(prod_to_visit)
        local subtrace::S
        local input::U
        cur = first(prod_to_visit)
        delete!(prod_to_visit, cur)
        input = get_production_input(gen_fn, cur, production_traces, root_idx, root_production_input)
        subtrace = simulate(gen_fn.production_kern, (input,))
        score += get_score(subtrace)
        production_traces = assoc(production_traces, cur, subtrace)
        children_inputs::Vector{U} = get_retval(subtrace).children
        for child_num in 1:length(children_inputs)
            push!(prod_to_visit, get_child(cur, child_num, gen_fn.max_branch))
        end
        if !isempty(get_choices(subtrace))
            num_has_choices += 1
        end
    end

    # aggregation phase
    # visit children first
    agg_to_visit = sort(collect(keys(production_traces)), rev=true)
    for cur in agg_to_visit
        local subtrace::T
        local input::Tuple{V,Vector{W}}
        input = get_aggregation_input(gen_fn, cur, production_traces, aggregation_traces)
        subtrace = simulate(gen_fn.aggregation_kern, input)
        score += get_score(subtrace)
        aggregation_traces = assoc(aggregation_traces, cur, subtrace)
        if !isempty(get_choices(subtrace))
            num_has_choices += 1
        end
    end

    RecurseTrace{S,T,U,V,W,X,Y}(gen_fn,
                production_traces, aggregation_traces, gen_fn.max_branch,
                score, root_idx, num_has_choices)
end

############
# generate #
############

function generate(gen_fn::Recurse{S,T,U,V,W,X,Y}, args::Tuple{U,Int},
                    constraints::ChoiceMap) where {S,T,U,V,W,X,Y}
    (root_production_input::U, root_idx::Int) = args
    production_traces = PersistentHashMap{Int,S}()
    aggregation_traces = PersistentHashMap{Int,T}()
    weight = 0.
    score = 0.
    num_has_choices = 0

    # production phase
    # does not matter in which order we visit (since children are inserted after parents)
    prod_to_visit = Set{Int}([root_idx])
    while !isempty(prod_to_visit)
        local subtrace::S
        local input::U
        cur = first(prod_to_visit)
        delete!(prod_to_visit, cur)
        input = get_production_input(gen_fn, cur, production_traces, root_idx, root_production_input)
        subconstraints = get_production_constraints(constraints, cur)
        (subtrace, subweight) = generate(gen_fn.production_kern, (input,), subconstraints)
        score += get_score(subtrace)
        production_traces = assoc(production_traces, cur, subtrace)
        weight += subweight
        children_inputs::Vector{U} = get_retval(subtrace).children
        for child_num in 1:length(children_inputs)
            push!(prod_to_visit, get_child(cur, child_num, gen_fn.max_branch))
        end
        if !isempty(get_choices(subtrace))
            num_has_choices += 1
        end
    end

    # aggregation phase
    # visit children first
    agg_to_visit = sort(collect(keys(production_traces)), rev=true)
    for cur in agg_to_visit
        local subtrace::T
        local input::Tuple{V,Vector{W}}
        input = get_aggregation_input(gen_fn, cur, production_traces, aggregation_traces)
        subconstraints = get_aggregation_constraints(constraints, cur)
        (subtrace, subweight) = generate(gen_fn.aggregation_kern, input, subconstraints)
        score += get_score(subtrace)
        aggregation_traces = assoc(aggregation_traces, cur, subtrace)
        weight += subweight
        if !isempty(get_choices(subtrace))
            num_has_choices += 1
        end
    end

    trace = RecurseTrace{S,T,U,V,W,X,Y}(gen_fn,
                production_traces, aggregation_traces, gen_fn.max_branch,
                score, root_idx, num_has_choices)
    return (trace, weight)
end

##########
# update #
##########

struct NodeQueue{T}
    ordering::T
    pq::PriorityQueue{Int,Int}
end

function NodeQueue(ordering)
    pq = PriorityQueue{Int,Int}(ordering)
    NodeQueue(ordering, pq)
end

function Base.push!(queue::NodeQueue, node::Int)
    if !haskey(queue.pq, node)
        enqueue!(queue.pq, node, node)
    end
end
Base.pop!(queue::NodeQueue) = dequeue!(queue.pq)
Base.in(node::Int, queue::NodeQueue) = haskey(queue.pq, node)
Base.isempty(queue::NodeQueue) = isempty(queue.pq)

function recurse_unpack_constraints(constraints::ChoiceMap)
    production_constraints = Dict{Int, Any}()
    aggregation_constraints = Dict{Int, Any}()
    for (addr, node) in get_submaps_shallow(constraints)
        idx::Int = addr[1]
        if addr[2] == Val(:production)
            production_constraints[idx] = node
        elseif addr[2] == Val(:aggregation)
            aggregation_constraints[idx] = node
        else
            error("Unknown address: $addr")
        end
    end
    if length(get_values_shallow(constraints)) > 0
        error("Unknown address: $(first(get_values_shallow(constraints))[1])")
    end
    return (production_constraints, aggregation_constraints)
end

function dissoc_subtree!(discard::DynamicChoiceMap,
						 production_traces::PersistentHashMap{Int,S},
                         aggregation_traces::PersistentHashMap{Int,T},
                         root::Int, max_branch::Int) where {S,T}
    num_has_choices = 0
    production_subtrace = production_traces[root]
    aggregation_subtrace = aggregation_traces[root]
	set_submap!(discard, (root, Val(:production)), get_choices(production_subtrace))
	set_submap!(discard, (root, Val(:aggregation)), get_choices(aggregation_subtrace))
    if !isempty(get_choices(production_subtrace))
        num_has_choices += 1
    end
    if !isempty(get_choices(aggregation_subtrace))
        num_has_choices += 1
    end
    num_children = get_num_children(production_subtrace)
    removed_score = get_score(production_subtrace) + get_score(aggregation_subtrace)
    for child_num=1:num_children
        child = get_child(root, child_num, max_branch)
        (production_traces, aggregation_traces, child_removed_score, child_num_has_choices) = dissoc_subtree!(
            discard, production_traces, aggregation_traces, child, max_branch)
        removed_score += child_removed_score
        num_has_choices += child_num_has_choices
    end
    production_traces = dissoc(production_traces, root)
    aggregation_traces = dissoc(aggregation_traces, root)
    return (production_traces, aggregation_traces, removed_score, num_has_choices)
end

function get_production_argdiff(production_retdiffs::Dict{Int,Diff},
                                root_idx::Int, root_argdiff, cur::Int,
                                max_branch::Int)
    if cur == root_idx
        return root_argdiff
    else
        parent = get_parent(cur, max_branch)
        if !haskey(production_retdiffs, parent)
            return NoChange()
        else
            retdiff = production_retdiffs[parent]
            @assert retdiff != NoChange()
            child_num = get_child_num(cur, max_branch)
            if isa(retdiff, ProductionDiff)
                # TODO: test this code path
                children_diff = production_retdiffs[parent].children_diff
                if child_num > children_diff.prev_len
                    return UnknownChange()
                else
                    if haskey(children_diff.updated, child_num)
                        return children_diff.updated[child_num]
                    else
                        return NoChange()
                    end
                end
            else
                return UnknownChange()
            end
        end
    end
end

function get_aggregation_argdiffs(production_retdiffs::Dict{Int,Diff},
                                  aggregation_retdiffs::Dict{Int,Diff},
                                  idx_to_prev_num_children::Dict{Int,Int},
                                  production_traces, cur::Int)
    if haskey(production_retdiffs, cur)
        production_retdiff = production_retdiffs[cur]
        if isa(production_retdiff, ProductionDiff)
            dv = production_retdiff.value_diff
        else
            dv = UnknownChange()
        end
    else
        dv = NoChange()
    end
    new_num_children = get_num_children(production_traces[cur])
    if haskey(idx_to_prev_num_children, cur)
        prev_num_children = idx_to_prev_num_children[cur]
    else
        prev_num_children = new_num_children
    end
    dws = Dict{Int,Diff}()
    for child_num=1:min(prev_num_children,new_num_children)
        if haskey(aggregation_retdiffs, child_num)
            dws[child_num] = aggregation_retdiffs[child_num]
        end
    end
    (dv, VectorDiff(new_num_children, prev_num_children, dws))
end

function update(trace::RecurseTrace{S,T,U,V,W,X,Y},
                new_args::Tuple{U,Int},
                argdiffs::Tuple,
                constraints::ChoiceMap) where {S,T,U,V,W,X,Y}
    gen_fn = get_gen_fn(trace)
    (root_production_input::U, root_idx::Int) = new_args
    if root_idx != get_args(trace)[2]
        error("Cannot change root_idx argument")
    end

    (root_argdiff::Diff,) = argdiffs

    production_traces = trace.production_traces
    aggregation_traces = trace.aggregation_traces
    (production_constraints, aggregation_constraints) = recurse_unpack_constraints(constraints)
    discard = choicemap()

    # initial score from previous trace
    score = trace.score
    weight = 0.
    num_has_choices = trace.num_has_choices

    # initialize set of production nodes to visit
    # visit nodes in forward order (starting with root)
    prod_to_visit = NodeQueue(Base.Order.Forward)
    for node in keys(production_constraints)
        push!(prod_to_visit, node)
    end

    # if the root argdiff is not noargdiff then we visit the root production node
    if root_argdiff != NoChange() && !(root_idx in prod_to_visit)
        push!(prod_to_visit, root_idx)
    end

    # initialize set of aggregation nodes to visit
    # visit nodes in reverse order (starting with leaves)
    agg_to_visit = NodeQueue(Base.Order.Reverse)
    for node in keys(aggregation_constraints)
        push!(agg_to_visit, node)
    end

    # only store for nodes that are retained
    # if a value is not present for a retained node, then the number of children of that node did not change
    idx_to_prev_num_children = Dict{Int,Int}()

    # production phase
    production_retdiffs = Dict{Int,Diff}() # elements with no difference are ommitted
    while !isempty(prod_to_visit)
        local subtrace::S

        cur = pop!(prod_to_visit)
        subconstraints = get_production_constraints(constraints, cur)
        input = (get_production_input(gen_fn, cur, production_traces, root_idx, root_production_input)::U,)

        if haskey(production_traces, cur)
            # the node exists already
            local subargdiff
            local subretdiff::Diff

            # get argdiff for this production node
            subargdiff = get_production_argdiff(production_retdiffs, root_idx,
                                                root_argdiff, cur, gen_fn.max_branch)

            # call update on production kernel
            prev_subtrace = production_traces[cur]
            (subtrace, subweight, subretdiff, subdiscard) = update(
                prev_subtrace, input, (subargdiff,), subconstraints)
            prev_num_children = get_num_children(production_traces[cur])
            new_num_children = length(get_retval(subtrace).children)
            idx_to_prev_num_children[cur] = prev_num_children
            set_submap!(discard, (cur, Val(:production)), subdiscard)
            production_retdiffs[cur] = subretdiff

            # update trace, weight, and score
            production_traces = assoc(production_traces, cur, subtrace)
            weight += subweight
            score += subweight

            # update num_has_choices
            if !isempty(get_choices(prev_subtrace)) && isempty(get_choices(subtrace))
                num_has_choices -= 1
            elseif isempty(get_choices(prev_subtrace)) && !isempty(get_choices(subtrace))
                num_has_choices += 1
            end

            # delete children (and their descendants), both production and aggregation nodes
            for child_num=new_num_children+1:prev_num_children
                child = get_child(cur, child_num, gen_fn.max_branch)
                set_submap!(discard, (child, Val(:production)), get_choices(production_traces[child]))
                set_submap!(discard, (child, Val(:aggregation)), get_choices(aggregation_traces[child]))
                (production_traces, aggregation_traces, removed_score, removed_num_has_choices) = dissoc_subtree!(
                    discard, production_traces, aggregation_traces, child, gen_fn.max_branch)
                score -= removed_score
                num_has_choices -= removed_num_has_choices
                weight -= removed_score
            end

            # mark new children for processing
            for child_num=prev_num_children+1:new_num_children
                child = get_child(cur, child_num, gen_fn.max_branch)
                push!(prod_to_visit, child)
            end

            if isa(subretdiff, ProductionDiff)

                # maybe mark existing children for processing, if they have a custom retdiff
                # (otherwise they do not change)
                for child_num in keys(subretdiff.children_diff.updated)
                    @assert child_num <= min(prev_num_children,new_num_children)
                    child = get_child(cur, child_num, gen_fn.max_branch)
                    push!(prod_to_visit, child)
                end

                # mark corresponding aggregation node for processing if v has
                # changed, or if the number of children has changed
                if subretdiff.value_diff != NoChange() || prev_num_children != new_num_children
                    push!(agg_to_visit, cur)
                end
            else
                # treat it like UnknownDiff
                for child_num=1:min(prev_num_children,new_num_children)
                    child = get_child(cur, child_num, gen_fn.max_branch)
                    push!(prod_to_visit, child)
                end
                push!(agg_to_visit, cur)
            end

        else
            # the node does not exist already (and none of its children exist either)
            (subtrace, ) = generate(gen_fn.production_kern, input, subconstraints)

            # update trace, weight, and score
            production_traces = assoc(production_traces, cur, subtrace)
            weight += get_score(subtrace)
            score += get_score(subtrace)

            # update num_has_choices
            if !isempty(get_choices(subtrace))
                num_has_choices += 1
            end

            # mark corresponding aggregation node for processsing
            push!(agg_to_visit, cur)

            # mark children (which are all new) for processing
            for child_num=1:get_num_children(subtrace)
                child = get_child(cur, child_num, gen_fn.max_branch)
                push!(prod_to_visit, child)
            end
        end
    end

    # aggregation phase

    # reasons to visit an aggregation node:
    # - it is new
    # - it is directly constrained
    # - its corresponding production node has !isnodiff(dv)
    # - one or more of of its children has a !isnodiff(dw)
    # note: we already deleted aggregation nodes that are to be deleted above

    aggregation_retdiffs = Dict{Int,Diff}() # NoChange()s are ommitted
    local subtrace::T
    local subargdiffs::Tuple # (diff for V, diff for vector of Ws)
    local subretdiff # Diff
    local retdiff
    retdiff = NoChange()
    while !isempty(agg_to_visit)

        cur = pop!(agg_to_visit)
        subconstraints = get_aggregation_constraints(constraints, cur)
        input::Tuple{V,Vector{W}} = get_aggregation_input(
            gen_fn, cur, production_traces, aggregation_traces)

        # if the node exists already
        if haskey(aggregation_traces, cur)

            # get argdiff for this aggregation node
            subargdiffs = get_aggregation_argdiffs(production_retdiffs, aggregation_retdiffs,
                                                   idx_to_prev_num_children, production_traces, cur)

            # call update on aggregation kernel
            prev_subtrace = aggregation_traces[cur]
            (subtrace, subweight, subretdiff, subdiscard) = update(
                prev_subtrace, input, subargdiffs, subconstraints)

            # update trace, weight, and score, and discard
            aggregation_traces = assoc(aggregation_traces, cur, subtrace)
            weight += subweight
            score += subweight
            set_submap!(discard, (cur, Val(:aggregation)), subdiscard)

            # update num_has_choices
            if !isempty(get_choices(prev_subtrace)) && isempty(get_choices(subtrace))
                num_has_choices -= 1
            elseif isempty(get_choices(prev_subtrace)) && !isempty(get_choices(subtrace))
                num_has_choices += 1
            end

            # take action based on our subretdiff
            if cur == root_idx
                retdiff = subretdiff
            else
                if subretdiff != NoChange()
                    aggregation_retdiffs[cur] = subretdiff
                    push!(agg_to_visit, get_parent(cur, gen_fn.max_branch))
                end
            end

        # if the node does not exist (but its children do, since we created them already)
        else
            (subtrace, _) = generate(gen_fn.aggregation_kern, input, subconstraints)

            # update trace, weight, and score
            aggregation_traces = assoc(aggregation_traces, cur, subtrace)
            weight += get_score(subtrace)
            score += get_score(subtrace)

            # update num_has_choices
            if !isempty(get_choices(subtrace))
                num_has_choices += 1
            end

            # the parent should have been marked during production phase
            @assert cur != 1
            @assert get_parent(cur, gen_fn.max_branch) in agg_to_visit
        end

    end

    new_trace = RecurseTrace{S,T,U,V,W,X,Y}(gen_fn,
        production_traces, aggregation_traces, gen_fn.max_branch,
        score, root_idx, num_has_choices)

    return (new_trace, weight, retdiff, discard)
end

export Recurse
export Production, ProductionDiff
export get_child

max_stars_repo_name: dendisuhubdy/Gen

max_issues_repo_name: kumsh/Gen

max_forks_repo_name: kumsh/Gen


text:
# Copyright (c) 2017: Miles Lubin and contributors
# Copyright (c) 2017: Google Inc.
#
# Use of this source code is governed by an MIT-style license that can be found
# in the LICENSE.md file or at https://opensource.org/licenses/MIT.

"""
    parse_expression(data::Model, input)::Expression

Parse `input` into a [`Expression`](@ref).
"""
function parse_expression(data::Model, input)
    expr = Expression()
    parse_expression(data, expr, input, -1)
    return expr
end

"""
    parse_expression(
        data::Model,
        expr::Expression,
        input::Any,
        parent_index::Int,
    )::Expression

Parse `input` into a [`Expression`](@ref), and add it to `expr` as a
child of `expr.nodes[parent_index]`. Existing subexpressions and parameters are
stored in `data`.

You can extend parsing support to new types of objects by overloading this
method with a different type on `input::Any`.
"""
function parse_expression(::Model, ::Expression, x::Any, ::Int)
    return error(
        "Unexpected object $x of type $(typeof(x)) in nonlinear expression.",
    )
end

function parse_expression(
    data::Model,
    expr::Expression,
    x::Expr,
    parent_index::Int,
)
    if isexpr(x, :call)
        if length(x.args) == 2 && !isexpr(x.args[2], :...)
            _parse_univariate_expression(data, expr, x, parent_index)
        else
            _parse_multivariate_expression(data, expr, x, parent_index)
        end
    elseif isexpr(x, :comparison)
        _parse_comparison_expression(data, expr, x, parent_index)
    elseif isexpr(x, :...)
        _parse_splat_expression(data, expr, x, parent_index)
    elseif isexpr(x, :&&) || isexpr(x, :||)
        _parse_logic_expression(data, expr, x, parent_index)
    else
        error("Unsupported expression: $x")
    end
    return
end

function _parse_splat_expression(data, expr, x, parent_index)
    @assert isexpr(x, :...) && length(x.args) == 1
    if parent_index == -1
        error(
            "Unsupported use of the splatting operator. This is only " *
            "supported in the arguments of a function call.",
        )
    elseif x.args[1] isa Expr
        error(
            "Unsupported use of the splatting operator. JuMP supports " *
            "splatting only symbols. For example, `x...` is ok, but " *
            "`(x + 1)...`, `[x; y]...` and `g(f(y)...)` are not.",
        )
    end
    for xi in x.args[1]
        parse_expression(data, expr, xi, parent_index)
    end
    return
end

function _parse_univariate_expression(
    data::Model,
    expr::Expression,
    x::Expr,
    parent_index::Int,
)
    @assert isexpr(x, :call, 2)
    id = get(data.operators.univariate_operator_to_id, x.args[1], nothing)
    if id === nothing
        # It may also be a multivariate operator like * with one argument.
        if haskey(data.operators.multivariate_operator_to_id, x.args[1])
            _parse_multivariate_expression(data, expr, x, parent_index)
            return
        end
        error("Unable to parse: $x")
    end
    push!(expr.nodes, Node(NODE_CALL_UNIVARIATE, id, parent_index))
    parse_expression(data, expr, x.args[2], length(expr.nodes))
    return
end

function _parse_multivariate_expression(
    data::Model,
    expr::Expression,
    x::Expr,
    parent_index::Int,
)
    @assert isexpr(x, :call)
    id = get(data.operators.multivariate_operator_to_id, x.args[1], nothing)
    if id === nothing
        @assert x.args[1] in data.operators.comparison_operators
        _parse_inequality_expression(data, expr, x, parent_index)
        return
    end
    push!(expr.nodes, Node(NODE_CALL_MULTIVARIATE, id, parent_index))
    parent_var = length(expr.nodes)
    for i in 2:length(x.args)
        parse_expression(data, expr, x.args[i], parent_var)
    end
    return
end

# This function parses single inequalities like `a <= b`. It's not to be
# confused with `_parse_comparison_expression`, which handles things like
# `a <= b <= c`.
function _parse_inequality_expression(
    data::Model,
    expr::Expression,
    x::Expr,
    parent_index::Int,
)
    operator_id = data.operators.comparison_operator_to_id[x.args[1]]
    push!(expr.nodes, Node(NODE_COMPARISON, operator_id, parent_index))
    parent_var = length(expr.nodes)
    for i in 2:length(x.args)
        parse_expression(data, expr, x.args[i], parent_var)
    end
    return
end

# This function parses double inequalities like `a <= b <= c`. It's not to be
# confused with `_parse_inequality_expression`, which handles things like
# `a <= b`.
function _parse_comparison_expression(
    data::Model,
    expr::Expression,
    x::Expr,
    parent_index::Int,
)
    for k in 2:2:length(x.args)-1
        @assert x.args[k] == x.args[2] # don't handle a <= b >= c
    end
    operator_id = data.operators.comparison_operator_to_id[x.args[2]]
    push!(expr.nodes, Node(NODE_COMPARISON, operator_id, parent_index))
    parent_var = length(expr.nodes)
    for i in 1:2:length(x.args)
        parse_expression(data, expr, x.args[i], parent_var)
    end
    return
end

function _parse_logic_expression(
    data::Model,
    expr::Expression,
    x::Expr,
    parent_index::Int,
)
    id = data.operators.logic_operator_to_id[x.head]
    push!(expr.nodes, Node(NODE_LOGIC, id, parent_index))
    parent_var = length(expr.nodes)
    parse_expression(data, expr, x.args[1], parent_var)
    parse_expression(data, expr, x.args[2], parent_var)
    return
end

function parse_expression(
    ::Model,
    expr::Expression,
    x::MOI.VariableIndex,
    parent_index::Int,
)
    push!(expr.nodes, Node(NODE_MOI_VARIABLE, x.value, parent_index))
    return
end

function parse_expression(::Model, expr::Expression, x::Real, parent_index::Int)
    push!(expr.values, convert(Float64, x)::Float64)
    push!(expr.nodes, Node(NODE_VALUE, length(expr.values), parent_index))
    return
end

function parse_expression(
    ::Model,
    expr::Expression,
    x::ParameterIndex,
    parent_index::Int,
)
    push!(expr.nodes, Node(NODE_PARAMETER, x.value, parent_index))
    return
end

function parse_expression(
    ::Model,
    expr::Expression,
    x::ExpressionIndex,
    parent_index::Int,
)
    push!(expr.nodes, Node(NODE_SUBEXPRESSION, x.value, parent_index))
    return
end

function parse_expression(
    ::Model,
    expr::Expression,
    x::AbstractArray,
    parent_index::Int,
)
    return error(
        "Unexpected array $x in nonlinear expression. Nonlinear expressions " *
        "may contain only scalar expressions.",
    )
end

"""
    convert_to_expr(data::Model, expr::Expression)

Convert the [`Expression`](@ref) `expr` into a Julia `Expr`.

 * subexpressions are represented by a [`ExpressionIndex`](@ref) object.
 * parameters are represented by a [`ParameterIndex`](@ref) object.
 * variables are represennted by an [`MOI.VariableIndex`](@ref) object.
"""
function convert_to_expr(model::Model, expr::Expression)
    tree = Any[]
    for node in expr.nodes
        node_expr = if node.type == NODE_CALL_UNIVARIATE
            Expr(:call, model.operators.univariate_operators[node.index])
        elseif node.type == NODE_CALL_MULTIVARIATE
            Expr(:call, model.operators.multivariate_operators[node.index])
        elseif node.type == NODE_COMPARISON
            Expr(:call, model.operators.comparison_operators[node.index])
        elseif node.type == NODE_LOGIC
            Expr(model.operators.logic_operators[node.index])
        elseif node.type == NODE_MOI_VARIABLE
            MOI.VariableIndex(node.index)
        elseif node.type == NODE_PARAMETER
            ParameterIndex(node.index)
        elseif node.type == NODE_SUBEXPRESSION
            ExpressionIndex(node.index)
        else
            @assert node.type == NODE_VALUE
            expr.values[node.index]
        end
        if 1 <= node.parent <= length(tree)
            push!(tree[node.parent].args, node_expr)
        end
        push!(tree, node_expr)
    end
    return _replace_comparison(tree[1])
end

# TODO(odow): NODE_COMPARISON is a bit of a pain to deal with...
_replace_comparison(expr) = expr

function _replace_comparison(expr::Expr)
    if isexpr(expr, :call, 4) && expr.args[1] in (:<=, :>=, :(==), :<, :>)
        return Expr(
            :comparison,
            expr.args[2],
            expr.args[1],
            expr.args[3],
            expr.args[1],
            expr.args[4],
        )
    end
    for i in 1:length(expr.args)
        expr.args[i] = _replace_comparison(expr.args[i])
    end
    return expr
end

ext: jl

max_stars_repo_name: JuliaOpt/MathOptInterface

max_issues_repo_name: JuliaOpt/MathOptInterface

max_forks_repo_name: JuliaOpt/MathOptInterface


text:
"""
Class of fields defined by the abstract concept of Field.
"""
abstract type Field end


text:
# ---
# title: 390. Elimination Game
# id: problem390
# author: Indigo
# date: 2021-06-20
# difficulty: Medium
# categories:
# link: <https://leetcode.com/problems/elimination-game/description/>
# hidden: true
# ---
#
# There is a list of sorted integers from 1 to _n_. Starting from left to right,
# remove the first number and every other number afterward until you reach the
# end of the list.
#
# Repeat the previous step again, but this time from right to left, remove the
# right most number and every other number from the remaining numbers.
#
# We keep repeating the steps again, alternating left to right and right to
# left, until a single number remains.
#
# Find the last number that remains starting with a list of length _n_.
#
# **Example:**
#
#
#
#     Input:
#     n = 9,
#     _1_ 2 _3_ 4 _5_ 6 _7_ 8 _9_
#     2 _4_ 6 _8_
#     _2_ 6
#     6
#
#     Output:
#     6
#
#
#
## @lc code=start
using LeetCode

last_remaining(n::Int)::Int = n == 1 ? 1 : 2 * (n ÷ 2 + 1 - last_remaining(n ÷ 2))

## @lc code=end


text:
using FactCheck
using RickTracy

num_item_fields = length(fieldnames(TraceItem))

facts("first things first") do
    fred = "flintstone"
    barney = 10

    @snap fred barney

    @fact (@tracevals fred) --> ["flintstone"]
    @fact (@tracevals barney) --> [10]
end

facts("snap the Nth") do
    for person in ["wilma", "fred", "betty", "barney"]
        @snap (N=2) person
    end
    @fact (@tracevals person) --> ["wilma", "betty"]
end

facts("snapat all on the floow") do
    @resetallsnaps
    fred = "Savage"
    winnie = 8.5
    @snap loc="McKinley" fred winnie
    traceitems = @traceitems loc="McKinley"
    @fact length(traceitems) --> 2
    @fact traceitems[1].exprstr --> "fred"
    @fact traceitems[1].val --> "Savage"
    @fact traceitems[2].exprstr --> "winnie"
    @fact traceitems[2].val --> 8.5
end

facts("snapifat") do
    @resetallsnaps
    for i in 1:10
        @snap loc="looptown" iff= i%3 == 0 && i != 6 i
        i%3 == 0 || @snap loc="loopcity" i
    end
    traceitems = @traceitems loc="looptown"
    @fact (@tracevals loc="looptown" i) --> [3,9]
    @fact (@tracevals loc="loopcity" i) --> [1,2,4,5,7,8,10]
end

facts("snapif") do
    @resetallsnaps
    for i in 1:10
        i%3 == 0 && @snap i
    end
    @fact (@tracevals i) --> [3,6,9]
end

facts("snapall") do
    @resetallsnaps
    @watch fred barney bambam

    fred = "1"
    barney = "2"
    bambam = "3"
    @snapall loc="first" N=1

    for i in 1:10
        fred = 2*i
        barney = 10*i
        bambam = 100*i
        @snapall loc="second" N=2
    end

    @fact (@tracevals bambam) --> Any["3",100,300,500,700,900]
end

facts("undefined is defined (as :undefined)") do
    @resetallsnaps
    @snap barney
    @fact (@tracevals barney) --> [:undefined]
    @fact_throws DomainError (@snap throw(DomainError()))
end

facts("tracevalsdic") do
    @resetallsnaps
    @watch fred barney bambam

    fred = "1"
    barney = "2"
    bambam = "3"
    @snapall loc="first" N=1

    for i in 1:10
        fred = 2*i
        barney = 10*i
        bambam = 100*i
        @snapall loc="second" N=5
    end
    @fact (@tracevalsdic) --> Dict("fred"=>Any["1",2,12],"barney"=>Any["2",10,60],"bambam"=>Any["3",100,600])
    @fact (@tracevalsdic loc=second) --> Dict("fred"=>[2,12], "barney"=>[10,60],
                                            "bambam"=>[100,600])
    @fact (@tracevalsdic loc=second bambam) --> Dict("bambam"=>[100,600])
    @fact (@tracevalsdic barney) --> Dict("barney"=>["2",10,60])
end

facts("tracesdf & tracesdic") do
    @resetallsnaps
    for i in 1:10
        @snap loc=ok1 i
        @snap loc=ok2 "dingo$i"
    end
    df1 = @tracesdf loc=ok1
    @fact sort!(names(df1)) --> sort!(fieldnames(TraceItem))
    @fact size(df1) --> (10, num_item_fields)

    df2 = @tracesdf
    @fact sort!(names(df1)) --> sort!(fieldnames(TraceItem))
    @fact size(df2) --> (20, num_item_fields)

    dicky = @tracesdic i
    @fact dicky[:val] --> collect(1:10)
end

facts("filter multiple expressions") do
    @resetallsnaps
    for i in 1:10
        @snap loc=ok1 i i^2
        @snap loc="I love you bad momo" "$i personal space"
    end
    dfi = @tracesdf i i^2
    df1 = @tracesdf loc=ok1
    df1i = @tracesdf loc=ok1 i i^2
    @fact sort!(names(dfi)) --> sort!(fieldnames(TraceItem))
    @fact size(dfi) --> (20, num_item_fields)
    @fact dfi --> df1
    @fact dfi --> df1i

    df2 = @tracesdf
    @fact sort!(names(df1)) --> sort!(fieldnames(TraceItem))
    @fact size(df2) --> (30, num_item_fields)

    dickfor = @tracevalsdic i i^2
    @fact dickfor["i"] --> collect(1:10)
    @fact dickfor[string(:(i^2))] --> collect(1:10) .^2

    oohyeah = @tracevals loc="I love you bad momo"
    @fact oohyeah[1] --> "1 personal space"
end

FactCheck.exitstatus()
include("iter_extrema.jl")
include("traces4symvals.jl")

max_stars_repo_name: JobJob/RickTracy


text:
if VERSION < v"1.7.0-DEV.849"
    include("generated_emoji_symbols_below_170.jl")
else
    include("generated_emoji_symbols_170.jl")
end


text: gaussdata = [-27.020, 3.570, 8.191, 9.898, 9.603, 9.945, 10.056]


text:
# https://i.stack.imgur.com/egDbs.png
"""
Does a math thing?
"""
function extendedeuclidean(a::Int64, b::Int64)::Tuple{Int64,String}
    d0 = (a, 1, 0)
    d1 = (b, 0, 1)

    while d1[1] != 0
        d0, d1 = d1, d0 .- (d0[1] ÷ d1[1]) .* d1
    end

    (d0, "$(d0[1]) = $(a) * $(d0[2]) + $(b) * $(d0[3])")
end

max_stars_repo_name: stellartux/advent-of-code

max_issues_repo_name: stellartux/advent-of-code

max_forks_repo_name: stellartux/advent-of-code


text:
# init_ranges = Dict(
#   "cm" => (1, 1000),
#   "Gleak" => (0.001, 1.0),
#   "Eleak" => (-0.2, 0.2),
#   "G" => (0.001, 1.0),
#   "μ" => (0.3, 0.8),
#   "σ" => (3, 8),
#   "E" => (1, 1.00001), # pol_mask
# )

struct LTCCellE{W,SENS,SYN,SENSP,SENSRE,SYNP,SYNRE,MI,MO,V,S,SOLVER,SENSEALG}

  wiring::W

  mapin::MI
  mapout::MO

  sens_f::SENS
  sens_p::SENSP
  sens_pl::Int
  sens_re::SENSRE
  syn_f::SYN
  syn_p::SYNP
  syn_pl::Int
  syn_re::SYNRE

  cm::V
  Gleak::V
  Eleak::V

  state0::S

  solver::SOLVER
  sensealg::SENSEALG

  function LTCCellE(wiring, mapin, mapout, sens_f, sens_p, sens_pl, sens_re, syn_f, syn_p, syn_pl, syn_re, cm, Gleak, Eleak, state0, solver, sensealg)
    new{typeof(wiring),typeof(sens_f),typeof(syn_f),typeof(sens_p),typeof(sens_re),typeof(syn_p),typeof(syn_re),typeof(mapin),typeof(mapout),typeof(cm),typeof(state0),typeof(solver),typeof(sensealg)}(
      wiring,mapin,mapout,sens_f,sens_p,sens_pl,sens_re,syn_f,syn_p,syn_pl,syn_re,cm,Gleak,Eleak,state0,solver,sensealg)
  end
end

function LTCCellE(wiring, solver, sensealg)
  n_in = wiring.n_in
  out = wiring.out
  n_total = wiring.n_total
  # cm    = rand_uniform(init_ranges["cm"][1], init_ranges["cm"][2], n_total)
  # Gleak = rand_uniform(init_ranges["Gleak"][1], init_ranges["Gleak"][2], n_total)
  # Eleak = rand_uniform(init_ranges["Eleak"][1], init_ranges["Eleak"][2], n_total)
  cm    = rand_uniform(1.6, 2.5, n_total)
  Gleak = rand_uniform(0.001, 1, n_total)
  Eleak = rand_uniform(-0.2, 0.2, n_total)
  state0 = fill(Float32(0.01), n_total, 1)

  mapin = Mapper(n_in)
  mapout = Mapper(n_total)

  sens_f = gNN(n_in, n_total)
  syn_f = gNN(n_total, n_total)
  sens_p, sens_re = Flux.destructure(sens_f)
  syn_p, syn_re = Flux.destructure(syn_f)

  LTCCellE(wiring,mapin,mapout,sens_f,sens_p,length(sens_p),sens_re,syn_f,syn_p,length(syn_p),syn_re, cm,Gleak,Eleak,state0,solver,sensealg)
end

Flux.@functor LTCCellE (mapin, mapout, sens_p, syn_p, cm, Gleak, Eleak, state0,)
#Flux.@functor LTCCellE (mapin, mapout, sens_p, syn_p, cm, Gleak, Eleak)
Flux.trainable(m::LTCCellE)=(Flux.trainable(m.mapin), Flux.trainable(m.mapout), m.sens_p, m.syn_p, m.cm, m.Gleak, m.Eleak, m.state0)

NCP(args...) = Flux.Recur(LTCCellE(args...))
Flux.Recur(m::LTCCellE) = Flux.Recur(m, m.state0)

#function (m::LTCCellE{W,SENS,SYN,SENSP,SENSRE,SYNP,SYNRE,MI,MO,V,<:AbstractArray{T}})(h::AbstractVecOrMat{T}, x::AbstractVecOrMat{T}) where {W,SENS,SYN,SENSP,SENSRE,SYNP,SYNRE,MI,MO,V,T}
function (m::LTCCellE)(h, x)

  _p = vcat(m.cm, m.Gleak, m.Eleak, m.sens_p, m.syn_p)
  #_p = vcat(m.cm, m.Gleak, m.Eleak)

  function prob_func(prob,i,repeat)
    #f = (du,u,p,t) -> dudt!(du,u,p,t,i)
    Ii = @view I[:,i]
    hi = @view h[:,i]
    #f = ODEFunction((x,p,t)->oop(x,p,t, m))
    #f = ODEFunction((x,p,t)->oop(x,p,t, m, Ii, sens_f,syn_f))
    f = ODEFunction((dx, x,p,t)->densltcdt!(dx, x,p,t, m, m.sens_f, m.syn_f))
    remake(prob, f=f, p=[_p;Ii[:]], u0=hi)
  end
  # function output_func(sol,i)
  #
  #   #i == 1 && @show sol.t
  #
  #   #@show size(sol)
  #   #sol_b[:,1,:] .= sol[:,end]
  #   #sol[:,end], false
  #   out = [x[:,1] for x in sol]
  #   return (DiffEqArray(out,sol.t), false)
  # end

  x = m.mapin(x)
  h = repeat(h, 1, size(x,2)-size(h,2)+1)
  # @show size(h)

  # h_new = Flux.Zygote.Buffer(h, size(h))
  # for i in eachindex(h_new)
  #   h_new[i] = h[i]
  # end

  # h_new = Flux.Zygote.Buffer(h, size(h,1),size(h,2))
  # for i in 1:size(h,1)
  #   h_new[i,:] .= h[i,:]
  # end

  #sol_b = Flux.Zygote.Buffer(h, size(h,1), 1, size(h,2))

  I = x
  sens_f = m.sens_re(m.sens_p)
  syn_f  = m.syn_re(m.syn_p)
  prob = ODEProblem((dx,x,p,t)->densltcdt!(dx,x,p,t, m, m.sens_f, m.syn_f), h[:,1], Float32.((0, 1)), _p)
  #prob = ODEProblem(ODEFunction{false}((x,p,t)->oop(x,p,t, m)), h[:,1], Float32.((0, 1)), [_p;I[:,1][:]])
  ensemble_prob = EnsembleProblem(prob, prob_func=prob_func)
  #sol = Array(solve(ensemble_prob,m.solver,EnsembleThreads(),trajectories=size(x,2), sensealg=m.sensealg, save_everystep=false, save_start=false))[:,end,:]
  sol = solve(ensemble_prob,m.solver,EnsembleThreads(),trajectories=size(x,2), dt=0.1, adaptive=false, sensealg=m.sensealg, save_everystep=false, save_start=false)[:,end,:]
  #sol = copy(sol_b)

  h = sol#[:,:,end]
  #@show size(h)

  # @show size(sol)
  # @show size(sol[1])
  # @show size(Array(sol))
  # @show size(Array(sol)[1])
  # for b in size(sol,1)
  #   h_new[:,b] .= sol[b]
  # end
  #h = Array(sol)[:,1,:]
  # h = copy(h_new)
  #h = reshape(sol, size(h,1),:)

  #sol = solve(prob, solver; sensealg, save_everystep=false, save_start=false, reltol=1e-3, abstol=1e-3)
  #sol = solve(prob, solver; sensealg, save_everystep=false, save_start=false, reltol=1e-3, abstol=1e-3)
  #sol = solve(prob, m.solver; sensealg=m.sensealg, dtmax=1.0)
  #@show length(sol.t)
  #sol = Array(solve(prob, solver; sensealg, save_everystep=false, save_start=false, reltol=1e-2, abstol=1e-2))
  #@show size(sol)
  #@show size(sol,3)

  out = m.mapout(h)
  return h, @view out[end-m.wiring.n_motor+1:end, :]
end

function densltcdt!(dx,x,p,t, m, sens_f, syn_f)

  cm_pl = size(x,1)
  Gleak_pl = size(x,1)
  Eleak_pl = size(x,1)
  sens_pl = m.sens_pl
  syn_pl = m.syn_pl

  @views cm    = p[1 : cm_pl]
  @views Gleak = p[cm_pl + 1 : cm_pl + Gleak_pl]
  @views Eleak = p[cm_pl + Gleak_pl + 1 : cm_pl + Gleak_pl + Eleak_pl]

  @views sens_p = p[cm_pl + Gleak_pl + Eleak_pl + 1 : cm_pl + Gleak_pl + Eleak_pl + sens_pl]
  @views syn_p = p[cm_pl + Gleak_pl + Eleak_pl + sens_pl + 1 : cm_pl + Gleak_pl + Eleak_pl + sens_pl + syn_pl]

  @views I = p[cm_pl + Gleak_pl + Eleak_pl + sens_pl + syn_pl + 1 : end]
  I = reshape(I, m.wiring.n_in, :)

  #@unpack cm, Gleak, Eleak, sens_p, syn_p, I = p

  # sens_f = m.sens_re(sens_p)
  # syn_f  = m.syn_re(syn_p)

  I_sens = sens_f(sens_p, x, I, m.wiring.sens_mask, m.wiring.sens_pol)
  I_syn = syn_f(sens_p, x, x, m.wiring.syn_mask, m.wiring.syn_pol)

  #tout(cm, Gleak, x, Eleak, I_syn, I_sens) = @tullio out[j,b] := cm[j] * (-(Gleak[j] * (x[j,b] - Eleak[j])) + I_syn[j,b] + I_sens[j,b])
  #dx .= tout(cm, Gleak, x, Eleak, I_syn, I_sens)
  @. dx = cm * (-(Gleak * (x - Eleak)) + I_syn + I_sens)
  nothing
end

function oop(x,p,t, m)

  cm_pl = size(x,1)
  Gleak_pl = size(x,1)
  Eleak_pl = size(x,1)
  sens_pl = m.sens_pl
  syn_pl = m.syn_pl

  @views cm    = p[1 : cm_pl]
  @views Gleak = p[cm_pl + 1 : cm_pl + Gleak_pl]
  @views Eleak = p[cm_pl + Gleak_pl + 1 : cm_pl + Gleak_pl + Eleak_pl]

  @views sens_p = p[cm_pl + Gleak_pl + Eleak_pl + 1 : cm_pl + Gleak_pl + Eleak_pl + sens_pl]
  @views syn_p = p[cm_pl + Gleak_pl + Eleak_pl + sens_pl + 1 : cm_pl + Gleak_pl + Eleak_pl + sens_pl + syn_pl]

  I = p[cm_pl + Gleak_pl + Eleak_pl + sens_pl + syn_pl + 1 : end]
  I = reshape(I, m.wiring.n_in, :)

  #@unpack cm, Gleak, Eleak, sens_p, syn_p, I = p

  sens_f = m.sens_re(sens_p)
  syn_f  = m.syn_re(syn_p)

  I_sens = sens_f(x, I, m.wiring.sens_mask, m.wiring.sens_pol)
  I_syn = syn_f(x, x, m.wiring.syn_mask, m.wiring.syn_pol)

  #tout(cm, Gleak, x, Eleak, I_syn, I_sens) = @tullio out[j,b] := cm[j] * (-(Gleak[j] * (x[j,b] - Eleak[j])) + I_syn[j,b] + I_sens[j,b])
  #dx .= tout(cm, Gleak, x, Eleak, I_syn, I_sens)
  @. cm * (-(Gleak * (x - Eleak)) + I_syn + I_sens)
end

struct Mapper{V<:AbstractArray}
  W::V
  b::V
  Mapper(W::V,b::V) where {V<:AbstractArray} = new{V}(W,b)
end
Mapper(in::Integer) = Mapper(ones(Float32,in), zeros(Float32,in))

(m::Mapper)(x::AbstractVecOrMat) = m.W .* x .+ m.b

Flux.@functor Mapper

struct gNN{A<:AbstractArray}
  G::A
  μ::A
  σ::A
  E::A
  #gNN(G::A, μ::A, σ::A, E::A)  where {A} = new{typeof(G)}(G, μ, σ, E)
  gNN(G::A, μ::A, σ::A, E::A) where {A<:AbstractArray} = new{A}(G, μ, σ, E)
end
function gNN(in::Integer, out::Integer)
  # G = rand_uniform(init_ranges["G"][1], init_ranges["G"][2], in, out)
  # μ = rand_uniform(init_ranges["μ"][1], init_ranges["μ"][2], in, out)
  # σ = rand_uniform(init_ranges["σ"][1], init_ranges["σ"][2], in, out)
  # E = rand_uniform(init_ranges["E"][1], init_ranges["E"][2], in, out)
  G = rand_uniform(0.001, 1, in, out)
  μ = rand_uniform(0.3, 0.8, in, out)
  σ = rand_uniform(3, 8, in, out)
  E = rand_uniform(-0.3, 0.3, in, out)
  gNN(G,μ,σ,E)
end

function (m::gNN)(h::AbstractVector, x, bitmask::Matrix{Float32}, polmask::Matrix{Float32})
  G, μ, σ, E = m.G, m.μ, m.σ, m.E

  function fun1e(bitmask, G, E, x, polmask, h, μ, σ)
    #sigmoid((x .- μ) .* σ)
    reshape(sum(bitmask .* G .* sigmoid.((x .- μ) .* σ) .* (polmask .* E .- reshape(h, 1,:)),dims=1),:)
  end

  function fun2e(bitmask, G, E, x, polmask, h, μ, σ)
    @tullio out[j] := bitmask[i,j] * G[i,j] * 1 / (1 + exp(-((x[i] - μ[i,j]) * σ[i,j]))) * (h[j] - polmask[i,j] * E[i,j]) #verbose=true
  end

  #fun4e(bitmask, G, E, x, polmask, h, μ, σ) = @tullio out[j] := bitmask[i,j] * G[i,j] * sigmoid((x[j] - μ[i,j]) * σ[i,j]) * (polmask[i,j] * E[i,j] - h[j])

  return fun1e(bitmask, G, E, x, polmask, h, μ, σ)
end

function (m::gNN)(p::AbstractVector, h::AbstractVector, x, bitmask::Matrix{Float32}, polmask::Matrix{Float32})
  Gl = length(m.G)
  μl = length(m.μ)
  σl = length(m.σ)
  El = length(m.E)
  @views G = p[1 : Gl]
  @views μ = p[Gl + 1 : Gl + μl]
  @views σ = p[Gl + μl + 1 : Gl + μl + σl]
  @views E = p[Gl + μl + σl + 1 : Gl + μl + σl + El]

  G = reshape(G, size(m.G))
  μ = reshape(μ, size(m.μ))
  σ = reshape(σ, size(m.σ))
  E = reshape(E, size(m.E))

  function fun1e(bitmask, G, E, x, polmask, h, μ, σ)
    #sigmoid((x .- μ) .* σ)
    reshape(sum(bitmask .* G .* sigmoid.((x .- μ) .* σ) .* (polmask .* E .- reshape(h, 1,:)),dims=1),:)
  end

  function fun2e(bitmask, G, E, x, polmask, h, μ, σ)
    @tullio out[j] := bitmask[i,j] * G[i,j] * 1 / (1 + exp(-((x[i] - μ[i,j]) * σ[i,j]))) * (h[j] - polmask[i,j] * E[i,j]) #verbose=true
  end

  #fun4e(bitmask, G, E, x, polmask, h, μ, σ) = @tullio out[j] := bitmask[i,j] * G[i,j] * sigmoid((x[j] - μ[i,j]) * σ[i,j]) * (polmask[i,j] * E[i,j] - h[j])

  return fun1e(bitmask, G, E, x, polmask, h, μ, σ)
end

Flux.@functor gNN


text:
#############################################
# Get the braking points
#############################################
function get_edges(X::AbstractMatrix{T}, nbins=250) where {T}
    edges = Vector{Vector{T}}(undef, size(X,2))
    @threads for i in 1:size(X, 2)
        edges[i] = quantile(view(X, :,i), (1:nbins)/nbins)
        if length(edges[i]) == 0
            edges[i] = [minimum(view(X, :,i))]
        end
    end
    return edges
end

####################################################
# Transform X matrix into a UInt8 binarized matrix
####################################################
function binarize(X, edges)
    X_bin = zeros(UInt8, size(X))
    @threads for i in 1:size(X, 2)
        X_bin[:,i] = searchsortedlast.(Ref(edges[i][1:end-1]), view(X,:,i)) .+ 1
    end
    X_bin
end

# split row ids into left and right based on best split condition
function update_set(set, best, x_bin)
    left = similar(set)
    right = similar(set)
    left_count = 0
    right_count = 0
    @inbounds for i in set
        if x_bin[i] <= best
            left_count += 1
            left[left_count] = i
        else
            right_count += 1
            right[right_count] = i
        end
    end
    resize!(left, left_count)
    resize!(right, right_count)
    return left, right
end

function update_hist!(hist_δ::Matrix{SVector{L,T}}, hist_δ²::Matrix{SVector{L,T}}, hist_𝑤::Matrix{SVector{1,T}},
    δ::Vector{SVector{L,T}}, δ²::Vector{SVector{L,T}}, 𝑤::Vector{SVector{1,T}},
    X_bin, node::TrainNode{L,T,S}) where {L,T,S}

    hist_δ .*= 0.0
    hist_δ² .*= 0.0
    hist_𝑤 .*= 0.0

    @inbounds @threads for j in node.𝑗
        @inbounds for i in node.𝑖
            hist_δ[X_bin[i,j], j] += δ[i]
            hist_δ²[X_bin[i,j], j] += δ²[i]
            hist_𝑤[X_bin[i,j], j] += 𝑤[i]
        end
    end
end

function find_split!(hist_δ::AbstractVector{SVector{L,T}}, hist_δ²::AbstractVector{SVector{L,T}}, hist_𝑤::AbstractVector{SVector{1,T}},
    params::EvoTypes, node::TrainNode{L,T,S}, info::SplitInfo{L,T,S}, edges::Vector{T}) where {L,T,S}

    # initialize tracking
    ∑δL = node.∑δ * 0
    ∑δ²L = node.∑δ² * 0
    ∑𝑤L = node.∑𝑤 * 0
    ∑δR = node.∑δ
    ∑δ²R = node.∑δ²
    ∑𝑤R = node.∑𝑤

    @inbounds for bin in 1:(length(hist_δ)-1)
        ∑δL += hist_δ[bin]
        ∑δ²L += hist_δ²[bin]
        ∑𝑤L += hist_𝑤[bin]
        ∑δR -= hist_δ[bin]
        ∑δ²R -= hist_δ²[bin]
        ∑𝑤R -= hist_𝑤[bin]

        gainL, gainR = get_gain(params.loss, ∑δL, ∑δ²L, ∑𝑤L, params.λ), get_gain(params.loss, ∑δR, ∑δ²R, ∑𝑤R, params.λ)
        gain = gainL + gainR

        if gain > info.gain && ∑𝑤L[1] >= params.min_weight + 1e-12 && ∑𝑤R[1] >= params.min_weight + 1e-12
            info.gain = gain
            info.gainL = gainL
            info.gainR = gainR
            info.∑δL = ∑δL
            info.∑δ²L = ∑δ²L
            info.∑𝑤L = ∑𝑤L
            info.∑δR = ∑δR
            info.∑δ²R = ∑δ²R
            info.∑𝑤R = ∑𝑤R
            info.cond = edges[bin]
            info.𝑖 = bin
        end # info update if gain
    end # loop on bins
end

function find_split_narrow!(hist_δ::Vector{SVector{L,T}}, hist_δ²::Vector{SVector{L,T}}, hist_𝑤::Vector{SVector{1,T}}, bins::Vector{BitSet}, X_bin, δ::Vector{SVector{L,T}}, δ²::Vector{SVector{L,T}}, 𝑤::Vector{SVector{1,T}}, ∑δ::SVector{L,T}, ∑δ²::SVector{L,T}, ∑𝑤::SVector{1,T}, params::EvoTreeRegressor, info::SplitInfo{L,T,S}, edges::Vector{T}, set::Vector{S}) where {L,T,S}

    # initialize histogram
    hist_δ .*= 0.0
    hist_δ² .*= 0.0
    hist_𝑤 .*= 0.0

    # build histogram
    @inbounds for i in set
        hist_δ[X_bin[i]] += δ[i]
        hist_δ²[X_bin[i]] += δ²[i]
        hist_𝑤[X_bin[i]] += 𝑤[i]
    end
    return
end


text:
# Integer arithmetic functions

for op ∈ (:(~), :(+), :(-))
    rename = VECTOR_SYMBOLS[op]
    @eval begin
        @inline $rename(s1::IntegerTypes) = $op($IntegerTypes)
        @vectordef $rename function Base.$op(v1) where {N,T<:IntegerTypes}
            llvmwrap(Val{$(QuoteNode(op))}(), extract_data(v1))
        end
    end
end
# @inline vnot(s1::Bool) = !s1
@inline vnot(v1::Vec{N,Bool}) where {N} = vbitwise_not(v1)
@inline vnot(v1::AbstractStructVec{N,Bool}) where {N} = SVec(vbitwise_not(extract_data(v1)))
@inline Base.:!(v1::AbstractStructVec{N,Bool}) where {N} = SVec(vbitwise_not(extract_data(v1)))

@inline vabs(s1::IntTypes) = abs(s1)
@inline function vabs(v1::Vec{N,T}) where {N,T<:IntTypes}
    # s = -vbroadcast(Vec{N,T}, signbit(v1))
    s = vright_bitshift(v1, Val{8*sizeof(T)}())
    # Note: -v1 == ~v1 + 1
    vsub(vxor(s, v1), s)
end
@inline vabs(v1::AbstractStructVec{N,T}) where {N,T<:IntTypes} = SVec(vabs(extract_data(v1)))
@inline Base.abs(v1::AbstractStructVec{N,T}) where {N,T<:IntTypes} = SVec(vabs(extract_data(v1)))

@inline vabs(v1::AbstractSIMDVector{N,T}) where {N,T<:UIntTypes} = v1
@inline Base.abs(v1::AbstractStructVec{N,T}) where {N,T<:UIntTypes} = v1

# TODO: Try T(v1>0) - T(v1<0)
#       use a shift for v1<0
#       evaluate v1>0 as -v1<0 ?

@inline vsign(s1::IntTypes) = sign(s1)
@inline vsign(v1::Vec{N,T}) where {N,T<:IntTypes} =
    vifelse(visequal(v1, vbroadcast(Vec{N,T},0)), vbroadcast(Vec{N,T},0),
        vifelse(vless(v1, vbroadcast(Vec{N,T},0)), vbroadcast(Vec{N,T},-1), vbroadcast(Vec{N,T},1)))
@inline vsign(v1::Vec{N,T}) where {N,T<:UIntTypes} =
    vifelse(visequal(v1, vbroadcast(Vec{N,T},0)), vbroadcast(Vec{N,T},0), vbroadcast(Vec{N,T},1))
    @inline vsign(v1::AbstractStructVec) = SVec(vsign(extract_data(v1)))
    @inline Base.sign(v1::AbstractStructVec) = SVec(vsign(extract_data(v1)))

@inline vsignbit(s1::IntegerTypes) = signbit(s1)
@inline vsignbit(v1::Vec{N,T}) where {N,T<:IntTypes} = vless(v1, vbroadcast(Vec{N,T}, 0))
@inline vsignbit(v1::AbstractStructVec{N,T}) where {N,T<:IntTypes} = SVec(vsignbit(extract_data(v1)))
@inline Base.signbit(v1::AbstractStructVec{N,T}) where {N,T<:IntTypes} = SVec(vsignbit(extract_data(v1)))

@inline vsignbit(v1::Vec{N,T}) where {N,T<:UIntTypes} = vbroadcast(Vec{N,Bool}, false)
@inline vsignbit(v1::AbstractStructVec{N,T}) where {N,T<:UIntTypes} = vbroadcast(SVec{N,Bool}, false)
@inline Base.signbit(v1::AbstractStructVec{N,T}) where {N,T<:UIntTypes} = vbroadcast(SVec{N,Bool}, false)

for op ∈ (:(&), :(|), :(⊻), :(+), :(-), :(*), :(÷), :(%) )
    rename = VECTOR_SYMBOLS[op]
    @eval begin
        @inline $rename(s1::IntegerTypes, s2::IntegerTypes) = $op(s1, s2)

        @vectordef $rename function Base.$op(v1, v2) where {N,T<:IntegerTypes}
            llvmwrap(Val{$(QuoteNode(op))}(), extract_data(v1), extract_data(v2))
        end

        # @inline $rename(v1::Vec{N,T}, v2::Vec{N,T}) where {N,T<:IntegerTypes} =
        #     llvmwrap(Val{$(QuoteNode(op))}(), v1, v2)
    end
end

let op = :(%)
    rename = VECTOR_SYMBOLS[op]

    @eval @vectordef $rename function Base.$op(v1, ::Type{I}) where {N, T <: IntegerTypes, I <: IntegerTypes}
        ntuple(Val(N)) do n
            VE(v1[n].value % I)
        end
    end
    @eval @vectordef $rename function Base.$op(v1, ::Type{Vec{N,I}}) where {N, T <: IntegerTypes, I <: IntegerTypes}
        ntuple(Val(N)) do n
            VE(v1[n].value % I)
        end
    end

end
@inline vrem(v1::V, ::Type{V}) where {V <: AbstractSIMDVector} = v1

@inline vcopysign(s1::IntegerTypes, s2::IntegerTypes) = copysign(s1, s2)
@inline vcopysign(v1::Vec{N,T}, v2::Vec{N,T}) where {N,T<:IntTypes} =
    vifelse(vsignbit(v2), -abs(v1), abs(v1))
@inline function vcopysign(v1::AbstractStructVec{N,T}, v2::AbstractStructVec{N,T}) where {N,T}
    SVec(vcopysign(extract_data(v1), extract_data(v2)))
end
@inline function Base.copysign(v1::AbstractStructVec{N,T}, v2::AbstractStructVec{N,T}) where {N,T}
    SVec(vcopysign(extract_data(v1), extract_data(v2)))
end
@inline vcopysign(v1::AbstractSIMDVector{N,T}, v2::AbstractSIMDVector{N,T}) where {N,T<:UIntTypes} = v1
@inline Base.copysign(v1::AbstractStructVec{N,T}, v2::AbstractStructVec{N,T}) where {N,T<:UIntTypes} = v1

@inline vflipsign(s1::IntegerTypes, s2::IntegerTypes) = flipsign(s1, s2)
@inline vflipsign(v1::Vec{N,T}, v2::Vec{N,T}) where {N,T<:IntTypes} =
    vifelse(vsignbit(v2), -v1, v1)
@inline function vflipsign(v1::AbstractStructVec{N,T}, v2::AbstractStructVec{N,T}) where {N,T<:IntTypes}
    SVec(vflipsign(extract_data(v1), extract_data(V2)))
end
@inline function Base.flipsign(v1::AbstractStructVec{N,T}, v2::AbstractStructVec{N,T}) where {N,T<:IntTypes}
    SVec(vflipsign(extract_data(v1), extract_data(V2)))
end

@inline vflipsign(v1::AbstractSIMDVector{N,T}, v2::AbstractSIMDVector{N,T}) where {N,T<:UIntTypes} = v1
@inline Base.flipsign(v1::AbstractStructVec{N,T}, v2::AbstractStructVec{N,T}) where {N,T<:UIntTypes} = v1

@inline vmax(s1::IntegerTypes, s2::IntegerTypes) = vmax(s1, s2)
@inline vmax(v1::Vec{N,T}, v2::Vec{N,T}) where {N,T<:IntegerTypes} =
    vifelse(v1>=v2, v1, v2)
@inline vmin(v1::Vec{N,T}, v2::Vec{N,T}) where {N,T<:IntegerTypes} =
    vifelse(v1>=v2, v2, v1)

vmuladd(s1::ScalarTypes, s2::ScalarTypes, s3::ScalarTypes) = muladd(s1,s2,s3)
@inline function vmuladd(v1::Vec{N,T}, v2::Vec{N,T}, v3::Vec{N,T}) where {N,T<:IntegerTypes}
    vadd(vmul(v1,v2),v3)
end
@inline function vmuladd(v1::AbstractSIMDVector{N,T}, v2::AbstractSIMDVector{N,T},
        v3::AbstractSIMDVector{N,T}) where {N,T<:IntegerTypes}
    SVec(vadd(vmul(extract_data(v1),extract_data(v2)),extract_data(v3)))
end
@inline function Base.muladd(v1::AbstractStructVec{N,T}, v2::AbstractStructVec{N,T},
        v3::AbstractStructVec{N,T}) where {N,T<:IntegerTypes}
    SVec(vadd(vmul(extract_data(v1),extract_data(v2)),extract_data(v3)))
end

vfma(s1::ScalarTypes, s2::ScalarTypes, s3::ScalarTypes) = fma(s1,s2,s3)
@inline function vfma(v1::Vec{N,T}, v2::Vec{N,T}, v3::Vec{N,T}) where {N,T<:IntegerTypes}
    vadd(vmul(v1,v2),v3)
end
@inline function vfma(v1::AbstractSIMDVector{N,T}, v2::AbstractSIMDVector{N,T},
        v3::AbstractSIMDVector{N,T}) where {N,T<:IntegerTypes}
    SVec(vadd(vmul(extract_data(v1),extract_data(v2)),extract_data(v3)))
end
@inline function Base.fma(v1::AbstractStructVec{N,T}, v2::AbstractStructVec{N,T},
        v3::AbstractStructVec{N,T}) where {N,T<:IntegerTypes}
    SVec(vadd(vmul(extract_data(v1),extract_data(v2)),extract_data(v3)))
end

# TODO: Handle negative shift counts
#       use vifelse
#       ensure vifelse is efficient
for op ∈ (:(<<), :(>>), :(>>>))
    rename = VECTOR_SYMBOLS[op]
    @eval begin
        @inline $rename(s1::IntegerTypes, s2::IntegerTypes) = $op(s1, s2)

        @vectordef $rename function Base.$op(v1, ::Val{I}) where {N,T<:IntegerTypes,I}
            llvmwrapshift(Val{$(QuoteNode(op))}(), extract_data(v1), Val{I}())
        end
        @vectordef $rename function Base.$op(v1, x2::Unsigned) where {N,T<:IntegerTypes}
            llvmwrapshift(Val{$(QuoteNode(op))}(), extract_data(v1), x2)
        end
        @vectordef $rename function Base.$op(v1, x2::Int) where {N,T<:IntegerTypes}
            llvmwrapshift(Val{$(QuoteNode(op))}(), extract_data(v1), x2)
        end
        @vectordef $rename function Base.$op(v1, x2::Integer) where {N,T<:IntegerTypes}
            llvmwrapshift(Val{$(QuoteNode(op))}(), extract_data(v1), x2)
        end
        @vectordef $rename function Base.$op(v1, v2::Vec{N,U}) where {N,T<:IntegerTypes,U<:UIntTypes}
            llvmwrapshift(Val{$(QuoteNode(op))}(), extract_data(v1), extract_data(v2))
        end
        @vectordef $rename function Base.$op(x1::T, v2) where {N,T<:IntegerTypes}
            $rename(vbroadcast(Vec{N,T}, x1), extract_data(v2))
        end

        # @inline $rename(v1::Vec{N,T}, ::Val{I}) where {N,T<:IntegerTypes,I} =
            # llvmwrapshift(Val{$(QuoteNode(op))}(), v1, Val{I}())
        # @inline $rename(v1::Vec{N,T}, x2::Unsigned) where {N,T<:IntegerTypes} =
        #     llvmwrapshift(Val{$(QuoteNode(op))}(), v1, x2)
        # @inline $rename(v1::Vec{N,T}, x2::Int) where {N,T<:IntegerTypes} =
        #     llvmwrapshift(Val{$(QuoteNode(op))}(), v1, x2)
        # @inline $rename(v1::Vec{N,T}, x2::Integer) where {N,T<:IntegerTypes} =
        #     llvmwrapshift(Val{$(QuoteNode(op))}(), v1, x2)
        # @inline $rename(v1::Vec{N,T},
        #                  v2::Vec{N,U}) where {N,T<:IntegerTypes,U<:UIntTypes} =
        #     llvmwrapshift(Val{$(QuoteNode(op))}(), v1, v2)
        # @inline $rename(v1::Vec{N,T},
        #                  v2::Vec{N,U}) where {N,T<:IntegerTypes,U<:IntegerTypes} =
        #     llvmwrapshift(Val{$(QuoteNode(op))}(), v1, v2)
        # @inline $rename(x1::T, v2::Vec{N,T}) where {N,T<:IntegerTypes} =
        #     $rename(vbroadcast(Vec{N,T}, x1), v2)
    end
end

# Promote scalars of all IntegerTypes to vectors of IntegerTypes, leaving the
# vector type unchanged

for op ∈ (
        :(==), :(!=), :(<), :(<=), :(>), :(>=),
        :(&), :(|), :(⊻), :(+), :(-), :(*),
        :copysign, :div, :flipsign, :max, :min, :rem
    )
    rename = VECTOR_SYMBOLS[op]
    @eval begin
        @inline $rename(s1::ScalarTypes, s2::ScalarTypes) = $op(s1, s2)

        @vectordef $rename function Base.$op(s1::T, v2) where {N, T <: Union{Bool,IntegerTypes}}
            $rename(vbroadcast(Vec{N,T}, s1), extract_data(v2))
        end
        @vectordef $rename function Base.$op(v1, s2::T) where {N, T <: Union{Bool,IntegerTypes}}
            $rename(extract_data(v1), vbroadcast(Vec{N,T}, s2))
        end

        # @vectordef $rename function Base.$op(s1::Bool, v2) where {N, T <: Bool}
        #     $rename(vbroadcast(Vec{N,Bool},s1), extract_data(v2))
        # end
        # @vectordef $rename function Base.$op(v2, s1::Bool) where {N, T <: Bool}
        #     $rename(vbroadcast(Vec{N,Bool},s1), extract_data(v2))
        # end
        # @inline $rename(s1::Bool, v2::Vec{N,Bool}) where {N} =
        #     $op(vbroadcast(Vec{N,Bool},s1), v2)
        # @inline $rename(s1::IntegerTypes, v2::Vec{N,T}) where {N,T<:IntegerTypes} =
        #     $op(vbroadcast(Vec{N,T}, s1), v2)
        # @inline $rename(v1::Vec{N,T}, s2::IntegerTypes) where {N,T<:IntegerTypes} =
        #     $op(v1, vbroadcast(Vec{N,T}, s2))
        #
        # @inline $rename(s1::Bool, v2::AbstractStructVec{N,Bool}) where {N} =
        #     SVec($op(Vec{N,Bool}(s1), extract_data(v2)))
        # @inline $rename(s1::IntegerTypes, v2::AbstractStructVec{N,T}) where {N,T<:IntegerTypes} =
        #     SVec($op(vbroadcast(Vec{N,T}, s1), extract_data(v2)))
        # @inline $rename(v1::AbstractStructVec{N,T}, s2::IntegerTypes) where {N,T<:IntegerTypes} =
        #     SVec($op(extract_data(v1), vbroadcast(Vec{N,T}, s2)))
    end
end
# @vectordef vifelse function Base.ifelse(c::Vec{N,Bool}, s1::T, v2) where {N,T<:IntegerTypes}
#     vifelse(c, vbroadcast(Vec{N,T}, s1), extract_data(v2))
# end
# @vectordef vifelse function Base.ifelse(c::Vec{N,Bool}, v1, s2::T) where {N,T<:IntegerTypes}
#     vifelse(c, extract_data(v1), vbroadcast(Vec{N,T}, s2))
# end
@inline function vifelse(c::Vec{N,Bool}, s1::T, v2::Vec{N,T}) where {N,T<:IntegerTypes}
    vifelse(c, vbroadcast(Vec{N,T}, s1), v2)
end
@inline function vifelse(c::Vec{N,Bool}, v1::Vec{N,T}, s2::T) where {N,T<:IntegerTypes}
    vifelse(c, v1, vbroadcast(Vec{N,T}, s2))
end
@inline function vifelse(c::AbstractSIMDVector{N,Bool}, s1::T, v2::AbstractSIMDVector{N,T}) where {N,T<:IntegerTypes}
    SVec(vifelse(extract_data(c), vbroadcast(Vec{N,T}, s1), extract_data(v2)))
end
@inline function vifelse(c::AbstractSIMDVector{N,Bool}, v1::AbstractSIMDVector{N,T}, s2::T) where {N,T<:IntegerTypes}
    SVec(vifelse(extract_data(c), extract_data(v1), vbroadcast(Vec{N,T}, s2)))
end

# @inline function vifelse(c::Vec{N,Bool}, s1::IntegerTypes, v2::Vec{N,T}) where {N,T<:IntegerTypes}
#     vifelse(c, vbroadcast(Vec{N,T}, s1), v2)
# end
# @inline function vifelse(c::Vec{N,Bool}, v1::Vec{N,T}, s2::IntegerTypes) where {N,T<:IntegerTypes}
#     vifelse(c, v1, vbroadcast(Vec{N,T}, s2))
# end

# for op ∈ (:muladd,)
#     rename = VECTOR_SYMBOLS[op]
#     @eval begin
#         @inline $rename(s1::IntegerTypes, v2::Vec{N,T},
#                 v3::Vec{N,T}) where {N,T<:IntegerTypes} =
#             $rename(vbroadcast(Vec{N,T}, s1), v2, v3)
#         @inline $rename(v1::Vec{N,T}, s2::IntegerTypes,
#                 v3::Vec{N,T}) where {N,T<:IntegerTypes} =
#             $rename(v1, vbroadcast(Vec{N,T}, s2), v3)
#         @inline $rename(s1::IntegerTypes, s2::IntegerTypes,
#                 v3::Vec{N,T}) where {N,T<:IntegerTypes} =
#             $rename(vbroadcast(Vec{N,T}, s1), vbroadcast(Vec{N,T}, s2), v3)
#         @inline $rename(v1::Vec{N,T}, v2::Vec{N,T},
#                 s3::IntegerTypes) where {N,T<:IntegerTypes} =
#             $rename(v1, v2, vbroadcast(Vec{N,T}, s3))
#         @inline $rename(s1::IntegerTypes, v2::Vec{N,T},
#                 s3::IntegerTypes) where {N,T<:IntegerTypes} =
#             $rename(vbroadcast(Vec{N,T}, s1), v2, vbroadcast(Vec{N,T}, s3))
#         @inline $rename(v1::Vec{N,T}, s2::IntegerTypes,
#                 s3::IntegerTypes) where {N,T<:IntegerTypes} =
#             $rename(v1, vbroadcast(Vec{N,T}, s2), vbroadcast(Vec{N,T}, s3))
#     end
# end

@generated function Base.:>>(a::SVec{N,I}, b::SVec{N,I}) where {N,I<:IntegerTypes}
    if I == Int64
        U = UInt64
    elseif I == Int32
        U = UInt32
    else
        throw("No method defined for Base.>>(::SVec{$N,$I}, ::SVec{$N,$I}).")
    end
    quote
        $(Expr(:meta,:inline))
        reinterpret(SVec{$N,$I}, reinterpret(SVec{$N,$U}, a) >> reinterpret(SVec{$N,$U}, b) )
    end
end
# @generated function Base.:>>(a::SVec{N,I}, b::SVec{N,I}) where {N,I<:IntegerTypes}
#     quote
#         $(Expr(:meta,:inline))
#         SVec(
#             Base.Cartesian.@ntuple $N n -> Core.VecElement(a[n] >> b[n])
#         )
#     end
# end
@generated function Base.:>>(a::SVec{N,U}, b::SVec{N,U}) where {N,U<:Union{UInt32,UInt64}}
    quote
        $(Expr(:meta,:inline))
        SVec(
            Base.Cartesian.@ntuple $N n -> Core.VecElement(a[n] >> b[n])
        )
    end
end


text:
using Latexify
using ParameterizedFunctions
using DiffEqBiological
using Test

f = @ode_def TestAlignFeedback begin
    dx = y/c_1 - x
    dy = x^c_2 - y
end c_1 c_2

g = @ode_def TestAlign2 begin
    dx = p_x - d_x*x - d_x_y*y*x
    dy = p_y - d_y*y
    dz = p_z
end p_x d_x d_x_y p_y d_y p_z

@test latexify(f) ==
raw"\begin{align}
\frac{dx}{dt} =& \frac{y}{c_{1}} - x \\
\frac{dy}{dt} =& x^{c_{2}} - y
\end{align}
"

@test latexify(f; env=:align, field=:symjac) ==
raw"\begin{align}
\frac{dx}{dt} =& -1 \cdot 1 =& c_{1}^{-1} \\
\frac{dy}{dt} =& x^{-1 + c_{2}} \cdot c_{2} =& -1 \cdot 1
\end{align}
"

@test latexalign(["d$(x)/dt" for x = f.syms], f.funcs) ==
raw"\begin{align}
\frac{dx}{dt} =& \frac{y}{c_{1}} - x \\
\frac{dy}{dt} =& x^{c_{2}} - y
\end{align}
"

@test latexalign([f,g]) == "\\begin{align}\n\\frac{dx}{dt}  &=  \\frac{y}{c_{1}} - x  &  \\frac{dx}{dt}  &=  p_{x} - d_{x} \\cdot x - d_{x\\_y} \\cdot y \\cdot x  &  \\\\\n\\frac{dy}{dt}  &=  x^{c_{2}} - y  &  \\frac{dy}{dt}  &=  p_{y} - d_{y} \\cdot y  &  \\\\\n  &    &  \\frac{dz}{dt}  &=  p_{z} \\cdot 1  & \n\\end{align}\n"

@reaction_func hill2(x, v, k) = v*x^2/(k^2 + x^2)

rn = @reaction_network TestAlignMyRnType begin
  hill2(y, v_x, k_x), 0 --> x
  p_y, 0 --> y
  (d_x, d_y), (x, y) --> 0
  (r_b, r_u), x ↔ y
end v_x k_x p_y d_x d_y r_b r_u

@test latexalign(rn; clean=true).s ==
raw"\begin{align}
\frac{dx}{dt} =& \frac{v_{x} \cdot y^{2}}{k_{x}^{2} + y^{2}} - d_{x} \cdot x - r_{b} \cdot x + r_{u} \cdot y \\
\frac{dy}{dt} =& p_{y} - d_{y} \cdot y + r_{b} \cdot x - r_{u} \cdot y
\end{align}
"

@test latexify(((1.0, 2), (3, 4)); env=:align) ==
raw"\begin{align}
1.0 =& 3 \\
2 =& 4
\end{align}
"

# @test_throws MethodError latexify(rn; bad_kwarg="should error")


text:
"""
    alpha(t::Real, lastspike, q, tau)
    alpha(t::Real, lastspike::AbstractArray{<:Real}, q::AbstractArray{<:Real}, tau::AbstractArray{<:Real})
    alpha(t::Real, lastspike::CuVecOrMat{<:Real}, q::CuVecOrMat{<:Real}, tau::CuVecOrMat{<:Real})

Evaluate an alpha synapse. Modeled as `(t - lastspike) * (q / τ) * exp(-(t - lastspike - τ) / τ) Θ(t - lastspike)`
  (where `Θ` is the Heaviside function).
Use `CuVector` instead of `Vector` for GPU support.

# Fields
- `t`: current time
- `lastspike`: last pre-synaptic spike time
- `q`: amplitude
- `tau`: time constant
"""
function alpha(t::Real, lastspike, q, tau)
    Δ = t - lastspike

    return (Δ >= 0 && Δ < Inf) * Δ * (q / tau) * exp(-(Δ - tau) / tau)
end
function alpha(t::Real, lastspike::AbstractArray{<:Real}, q::AbstractArray{<:Real}, tau::AbstractArray{<:Real})
    Δ = t .- lastspike
    I = @. Δ * (q / tau) * exp(-(Δ - tau) / tau)

    return map((δ, i) -> (δ >= 0) && (δ < Inf) ? δ * i : zero(i), Δ, I)
end
function alpha(t::Real, lastspike::CuVecOrMat{<:Real}, q::CuVecOrMat{<:Real}, tau::CuVecOrMat{<:Real})
    Δ = t .- lastspike

    return @. (Δ >= 0) * (Δ < Inf) * Δ * (q / tau) * exp(-(Δ - tau) / tau)
end


text:
function collect_ed_results(experiment_time::String, trial::Int64, day_range, hours::Int64)

    base_path = joinpath(DrWatson.datadir(), "exp_raw", "$(experiment_time)")

    Results = Dict{String, DataFrames.DataFrame}()
    names = ["cg", "pg", "pl", "pw", "rg", "ug", "il_ub", "pw_ub", "optimizer_log"]

    for n in names
        Results[n] = DataFrames.DataFrame()
    end

    for d in day_range, h in 1:hours
        folder = joinpath(base_path, "Trial-$(trial)-ED-Day-$(d)-hour-$(h)")
        for (k, v) in Results
            file_results = Feather.read(joinpath(folder, "$(k).feather"))
            Results[k] = vcat(v, file_results)
        end

    end

    return Results

end

function collect_uc_results(experiment_time::String, model::String, trial::Int64, day_range)

    base_path = joinpath(DrWatson.datadir(), "exp_raw", "$(experiment_time)")

    Results = Dict{String, DataFrames.DataFrame}()
    names = ["cg", "pg", "pl", "pw", "rg", "ug", "il_ub", "pw_ub", "optimizer_log"]

    for n in names
        Results[n] = DataFrames.DataFrame()
    end

    for d in day_range
        folder = joinpath(base_path, "Trial-$(trial)-$(model)-Day-$(d)")
        for (k, v) in Results
            file_results = Feather.read(joinpath(folder, "$(k).feather"))
            Results[k] = vcat(v, file_results)
        end

    end

    return Results

end


text:
module MatrixProductStates
using TensorNetworkTensors
using LinearAlgebra
using TensorOperations: @tensor, scalar, checked_similar_from_indices, tensorcontract
using Parameters
using Printf
using KrylovKit
using Base.Iterators: take, rest, Rest
using Base.Printf
import Base.iterate

include("auxiliary-iterators.jl")

abstract type AbstractMPS{L,T,TE,TB} end
abstract type AbstractMPO{L,T,TE,TB} end

export MPS, CanonicalMPS, MPO
export centerlink, zerosite, inner, canonicalize, canonicalize!, normalize!, normalize
export randmps
export onesiteexp, onesitevar, twositecorr, applyonesite!, entropy
include("mps.jl")

export heffmap
include("heffapplication.jl")

export dmrg
include("dmrg.jl")

export tdvp
include("tdvp.jl")

export hubbardMPO, annihilationOp, creationOp
include("hubbard.jl")

export tfisingMPO, sigmaZop
include("tfising.jl")
end # module


text:
function formdataseries(;
  refreshcomp::Bool = true,
  refreshtretcrsp::Bool = true,
  refreshtvolmcapcrsp::Bool = true,
  datapath::String=DATA_PATH,
  levname::String = LEV_NAME,
  injlsstream::Function = IN_JLS_STREAM,
  outjlsstream::Function = OUT_JLS_STREAM,
  outpath::String = OUT_PATH,
  refreshdataseries::Bool = true,
  trialrunonly::Bool = false,
  archive::Bool = true,
  retaineddataseriescolumns = RETAINED_COLS_DATA_SERIES)::DataFrame

  local lev::DataFrame
  local fyear::Int

  local dailyunlesstrial::Symbol = trialrunonly ? :monthly : :daily
  refreshdataseries = (refreshdataseries || refreshcomp)

  levname = "$levname-$dailyunlesstrial"

  if refreshdataseries
    lev = prepcomp(refreshcomp=refreshcomp)

    #supplements the market value and constructs the trailing daily vol
    lev = maketvolmcapseries!(
      lev, refreshtvolmcapcrsp=refreshtvolmcapcrsp, crsptype=dailyunlesstrial)

    lev = maketretseries!(lev, refreshtretcrsp = refreshtretcrsp, crsptype=:monthly)

    winsorizekeydata!(lev)
    constructportfolios!(lev)

    if TEST_OUTPUT
      fyear = 2015
      CSV.write("output\\$(levname)_$fyear.csv", lev[(y->y==fyear).(lev.fyear), :])
    end

    if TEST_OUTPUT
      fyear = 2014
      CSV.write("output\\$(levname)_$fyear.csv", lev[(y->y==fyear).(lev.fyear), :])
    end

    select!(lev, RETAINED_COLS_DATA_SERIES)

    outjlsstream("$datapath\\$levname.jls.lz4") do s
      serialize(s, lev)
    end

    archive && archivefile("$levname.jls.lz4") #never hurts to have a backup
  else

    try
      injlsstream("$datapath\\$levname.jls.lz4") do s
        lev = deserialize(s) #NOTE: DEBUG ONLY
      end
    catch err
      @warn "Could not load $levname. Error: $err"
      @info "Attempting to restore back-up and load from archive..."
      unarchivefile("$levname.jls.lz4")
      injlsstream("$datapath\\$levname.jls.lz4") do s
        lev = deserialize(s) #NOTE: DEBUG ONLY
      end
      @info "Successfully restored from archive."
    end
  end

  return lev
end

function winsorizekeydata!(lev::DataFrame;
    fields2winsorize::AbstractArray{Symbol}=[CQUART_FVALS; ],
    winsorizeprop=WINSORIZE_PROP, groupfield=:fyear)

  #dups will lead to race errors
  @assert length(fields2winsorize) == length(unique(fields2winsorize))

  @mpar for s ∈ fields2winsorize
    local nbefore::Int = length(unique(lev[!, s]))
    winsorize!(lev[!, s], prop = winsorizeprop)
    (length(unique(lev[!, s])) == nbefore) && (@warn "winsorization didn't do anything for field $s.
      Something is probably wrong.")
  end

  return nothing
end

# a couple of helper functions
@inline function rowbeforedatesorted(df::AbstractDataFrame, dt::Date)::Union{Nothing, DataFrameRow}
  local r::Union{DataFrameRow,Nothing}

  i::Int = searchsortedfirst(df.date, dt)

  #check if the index is valid for our purposes
  r = ((i>size(df,1)) || #this means the date or a larger date wasn't found
    (i==1)
    ? nothing : df[i-1,:])

  return r
end

#takes the last merged record for each merged comp record
function condensewithlast!(panel::DataFrame; sortedpermnodate=false,
    condense::Bool = true)::Nothing
  (!sortedpermnodate) && sort!(panel, [:permno, :date])
  panel.annflag = falses(size(panel,1))

  #now collapse the return data
  panelgrp::GroupedDataFrame = groupby(panel, :permno)
  @mpar for i ∈ 1:length(panelgrp)
    spanel::SubDataFrame = panelgrp[i]

    for sspanel ∈ groupby(spanel, :compid)
      sspanel.annflag[end] = true
    end
  end

  #now drop all other rows
  #can either free the memory or do this in place
  if condense
    filter!(r->r.annflag, panel)
    select!(panel, Not([:annflag]))
  end

  return nothing
end

function maketretseries!(lev::DataFrame;
  refreshtretcrsp::Bool = true,
  crsptype::Symbol = :daily,
  months2stale::Int = MONTHS_2_STALE_SHORT)

  crspcolumns::Vector{Symbol} = [:date, :permno, :ret, :lret,
    :ret12m, :lret12m#=, :net12m, :lnet12m=#]

  periods2stale = Month(months2stale)

  validpermnos::Vector{Int} = unique(lev.lpermno)

  crsp::DataFrame = prepcrsp(
    crspcolumns = crspcolumns,
    crsptype = :monthly,
    refreshcrsp=refreshtretcrsp,
    validpermnos = validpermnos,
    prefix = "tret") do df::DataFrame
    computenetreturns!(df)
    computetrailingreturns!(df, 12, Flret = :lret, Fnewret = :ret12m,
      months2stale = MONTHS_2_STALE_SHORT + 12)
    #computetrailingreturns!(df, 12, Flret = :lnet, Fnewret = :net12m,
    #  months2stale = MONTHS_2_STALE_SHORT + 12)
  end

  lev = mergecrsp(crsp, lev, refreshmerge=true, cachedata = false, crsptype=crsptype, prefix="tret")
  obliterate!(crsp)

  #now collapse the return data
  #***NOTE***: We are collecting the latest entry for each period
  #if periods are at irregular intervals, this is technically different than what is stated in the paper
  #but it fully avoids contaminated date announcements
  #the alternative might be to drop all such overlapping entries, or make the intervals irregular
  sort!(lev, [:permno, :date])

  #=if TEST_OUTPUT
    CSV.write("output\\maketretseries.csv", lev[(end-10_000):end, :])
    @info "maketretseries.csv written."
  end=#

  condensewithlast!(lev, sortedpermnodate=true)

  #now sort and lag
  sort!(lev, [:permno, :fyear])
  lagwithin!(lev, [:ret12m], :permno, :fyear, sorted=true)
  lagwithin!(lev, [:Lret12m], :permno, :fyear, sorted=true)

  #=if TEST_OUTPUT
    CSV.write("output\\maketretseries-last.csv", lev[(end-10_000):end, :])
    @info "maketretseries-last.csv written."
  end=#

  return lev
end

function maketvolmcapseries!(lev::DataFrame;
  refreshtvolmcapcrsp::Bool = true,
  crsptype::Symbol = :daily,
  months2stale::Int = MONTHS_2_STALE_SHORT)

  crspcolumns::Vector{Symbol} = [:date, :permno, :mktcap, :vol252d]
  periods2stale = Month(months2stale)

  validpermnos::Vector{Int} = unique(lev.lpermno)

  crsp::DataFrame = prepcrsp(
    crspcolumns = crspcolumns,
    crsptype = crsptype,
    refreshcrsp=refreshtvolmcapcrsp,
    validpermnos = validpermnos,
    prefix = "tvolmcap") do df::DataFrame
    computecrspmktcap!(df)
    computenetreturns!(df)
    computetrailingvol!(df, crsptype==:daily ? 252 : 12,
      Fret = :net, Fvol = :vol252d, months2stale = MONTHS_2_STALE_SHORT + 12)
  end

  lev = mergecrsp(crsp, lev,
    refreshmerge=true,
    cachedata = false,
    crsptype=crsptype,
    prefix="tvolmcap")

  obliterate!(crsp)

  #supplement the market value
  sort!(lev, [:permno, :date])
  lev.last = falses(size(lev,1))

  #now supplement the market data
  levgrp::GroupedDataFrame = groupby(lev, :permno)
  mkequityadded::Int = sum((ismissing).(lev.mkequity))
  @mpar for i ∈ 1:length(levgrp)
    slev::SubDataFrame = levgrp[i]

    for sslev ∈ groupby(slev, :compid)
      sslev.last[end] = true

      local Nmissing::Int = sum((ismissing).(sslev.mkequity))
      if Nmissing > 1
        (Nmissing < size(sslev,1)) && error("Inconsistent presence of mkequity")

        dt::Date = eom(sslev.fyrmonth[1])

        #find the first date before the end of the fiscal year if available
        #need to get this from slev
        local r::Union{Nothing, DataFrameRow} = rowbeforedatesorted(slev, dt)
        (!isnothing(r)) && ( dt - periods2stale ≤ r.date) && (sslev.mkequity .= r.mktcap)
      end
    end
  end

  mkequityadded -= sum((ismissing).(lev.mkequity))
  @info "Supplemented $mkequityadded mkequity records using $crsptype data"

  #now that we have patched the market value, drop all other rows
  #can either free the memory or do this in place
  lev = filter!(r->r.last, lev)

  #make the cash field
  lev.mkat = lev.lt .+ lev.mkequity
  lev.mknegcash = 1.0 .- (lev.cash ./ lev.mkat)

  #compute #flev field
  lev.mkflev = lev.netdebt ./ (lev.netdebt .+ lev.mkequity)

  #compute liab
  lev.mkliab = lev.lt ./ lev.mkat

  #make finite (note we winsorize later on)
  @mpar for s ∈ [:mknegcash, :mkflev, :mkliab]
    lev[!,s] = (f::MFloat64->coalesce(isfinite(f),false) ? f : missing).(lev[!,s])
  end

  sort!(lev, [:permno, :fyear]) #not completely sure we need this sort

  #make lagged change in assets as a control variable
  lev.lat = (log).(lev.at)
  differencewithin!(lev, :lat, :gvkey, :fyear, sorted=true)
  lev.Dat = (exp).(lev.Dlat)

  lev.bm = lev.bkequity ./ lev.mkequity

  lagwithin!(lev, [:vol252d, :mkat, :mkequity, :bm, :Dat, :op],
    :permno, :fyear, sorted=true)
  lagwithin!(lev, :Lvol252d, :permno, :fyear, sorted=true)
  #lagwithin!(lev, [:LLvol252d, :LLvol504d], :fyear, sorted=true)

  rename!(lev, [:date=>:dailyenddate, :compid=>:compidtvolmcap])
  select!(lev, Not([:permno, :last])) #dup of lpermno

  return lev
end

#creates the control variables with appropriate lagging
function makepanelmonthly(panel::DataFrame;
  crsptype::Symbol = :daily)

  local crsptvol::DataFrame
  local crsptret::DataFrame
  local crsp::DataFrame

  #load the previous crsp fields
  crsptvol = prepcrsp(
    crspcolumns = nothing,
    crsptype = crsptype,
    refreshcrsp=false,
    validpermnos = nothing,
    prefix = "tvolmcap") do df::DataFrame
      @assert false #do not refresh the crsp file here
    end
  select!(crsptvol, [:permno, :date, :vol252d])

  crsptret = prepcrsp(
    crspcolumns = nothing,
    crsptype = :monthly,
    refreshcrsp=false,
    validpermnos = nothing,
    prefix = "tret") do df::DataFrame
      @assert false #do not refresh the crsp file here
    end
  select!(crsptret, [:permno, :date, :ret12m, :ret, :lret])

  #use the dialy data to supplement the existing monthly data (new dataset is monthly)
  crsp = join(crsptvol, crsptret, on=[:permno, :date])
  obliterate!(crsptvol)
  obliterate!(crsptret)

  #this will make the lookups easier
  crsp.dyrmonth = (d->YearMonth(Year(d).value, Month(d).value)).(crsp.date)

  select!(panel, Not([:vol252d, :ret12m])) #drop these since we will work with the monthly versions
  panel = mergecrsp(crsp, panel, refreshmerge=true, cachedata = false,
    crsptype=crsptype, prefix="panel")

  #preallocate
  N::Int = size(panel,1)
  panel.L12ret12m = Vector{MFloat64}(undef, N)
  panel.L24ret12m = Vector{MFloat64}(undef, N)
  panel.fyret12m = Vector{MFloat64}(undef, N)

  panel.L12vol252d = Vector{MFloat64}(undef, N)
  panel.L24vol252d= Vector{MFloat64}(undef, N)
  panel.fyvol252d= Vector{MFloat64}(undef, N)

  panel.fyoffset = panel.dyrmonth .- panel.fyrmonth

  sort!(panel, [:permno, :dyrmonth])
  @info "beginning panel creation"
  spanels::GroupedDataFrame = groupby(panel, [:permno])
  #=@mpar =#for i ∈ 1:length(spanels)
    local spanel::SubDataFrame = spanels[i]
    local Nspanel::Int = size(spanel,1)
    local ptr::Int
    local target::YearMonth

    #this should be true always
    (Nspanel == length(Set(spanel.dyrmonth))) || error("Dates for dyrmonth are not unique w/in permno")

    for (j, r) ∈ enumerate(eachrow(spanel))

      if j > 12 #compute teh 1yr lagged controls
        target = r.dyrmonth - 12
        ptr = searchsortedfirst(spanel.dyrmonth, target)
        if spanel.dyrmonth[ptr] == target
          r.L12ret12m = spanel[ptr, :ret12m]
          r.L12vol252d = spanel[ptr, :vol252d]
        end
      end

      if j > 24 #compute teh 2yr lagged controls
        target = r.dyrmonth - 24
        ptr = searchsortedfirst(spanel.dyrmonth, target)
        if spanel.dyrmonth[ptr] == target
          r.L24ret12m = spanel[ptr, :ret12m]
          r.L24vol252d = spanel[ptr, :vol252d]
        end
      end

      if j > r.fyoffset #compute teh 2yr lagged controls
        target = r.dyrmonth - r.fyoffset
        ptr = searchsortedfirst(spanel.dyrmonth, target)
        if spanel.dyrmonth[ptr] == target
          r.fyret12m = spanel[ptr, :ret12m]
          r.fyvol252d = spanel[ptr, :vol252d]
        end
      end
    end
  end

  #println(unique(panel.fyoffset))
  #error("stop here")

  #println(describe(panel))
  @info "panel creation complete"

  return panel
end

max_stars_repo_name: clintonTE/CCA

max_issues_repo_name: clintonTE/CCA

max_forks_repo_name: clintonTE/CCA


text:
include(joinpath("dataProcessing", "anomalies_test.jl"))
include(joinpath("dataProcessing", "causeSelection_test.jl"))
include(joinpath("dataProcessing", "dataSelection_test.jl"))
include(joinpath("dataProcessing", "processOverflow_test.jl"))
include(joinpath("dataProcessing", "siteSelection_test.jl"))
include(joinpath("dataProcessing", "RainYday_test.jl"))

max_stars_repo_name: AlexLsn/CSO

max_issues_repo_name: AlexLsn/CSO

max_forks_repo_name: AlexLsn/CSO


text:
"""
Canonical polyadic N-mode tensor decomposition (CANDECOMP/PARAFAC).
"""
immutable CANDECOMP{T<:Number, N} <: TensorDecomposition{T, N}
  factors::NTuple{N, Matrix{T}} # rank X i-th tensor mode factor matrices
  lambdas::Vector{T}
  props::Dict{Symbol, Any}      # extra properties

  CANDECOMP(factors::NTuple{N, Matrix{T}}, lambdas::Vector{T}) =
    new(factors, lambdas, Dict{Symbol, Any}())

  Base.call{T,N}(::Type{CANDECOMP}, factors::NTuple{N, Matrix{T}}, lambdas::Vector{T}) =
    CANDECOMP{T,N}(factors, lambdas)
end

"""
Returns CANDECOMP rank.
"""
Base.rank(decomp::CANDECOMP) = length(decomp.lambdas)

"""
Re-composes the tensor from CANDECOMP decomposition.
"""
@generated function compose!{T,N}(dest::Array{T,N}, factors::NTuple{N, Matrix{T}}, lambdas::Vector{T})
  quote
    @nloops $N i dest begin
        elm = zero(T)
        # FIXME traversal of factors is not very efficient, would be better to have them transposed
        for j in 1:length(lambdas)
            elm += lambdas[j] * (*(@ntuple($N, k -> factors[k][i_k, j])...))
        end
        @nref($N, dest, i) = elm
    end
    dest
  end
end

@generated function compose{T,N}(factors::NTuple{N, Matrix{T}}, lambdas::Vector{T})
  quote
    compose!(Array{T}(@ntuple $N i -> size(factors[i], 1)), factors, lambdas)
  end
end

compose(decomp::CANDECOMP) = compose(decomp.factors, decomp.lambdas)

compose!{T,N}(dest::Array{T,N}, decomp::CANDECOMP{T,N}) = compose!(dest, decomp.factors, decomp.lambdas)

"""
Calculates canonical polyadic tensor decomposition (CANDECOMP/PARAFAC).

Returns:
  `CANDECOMP` object
"""
function candecomp{T,N}(tnsr::StridedArray{T,N},
                   r::Integer;
                   method::Symbol=:ALS,
                   tol::Float64=1e-5,
                   maxiter::Integer=100,
                   hosvd_init::Bool=false,
                   compute_error::Bool=false,
                   verbose::Bool=true)

    _check_tensor(tnsr, r)
    # don't use HO-SVD init if later the call would fail because of the wrong method
    verbose && info("Initializing factor matrices...")
    factors = in(method, CANDECOMP_methods) && hosvd_init ?
              hosvd(tnsr, r, compute_error=false).factors :
              _random_factors(size(tnsr), r)

    verbose && info("Applying CANDECOMP $method method...")
    res = _candecomp(Val{method}, tnsr, r, factors, tol, maxiter, verbose)
    if compute_error
      _set_rel_residue(res, tnsr)
    end
    return res
end

"Stub for non-implemented CANDECOMP algorithms."
_candecomp{T,N,S<:Symbol}(
  method::Val{S},
  tnsr::StridedArray{T,N},
  r::Integer,
  factors::Vector{Matrix{Float64}},
  tol::Float64,
  maxiter::Integer,
  verbose::Bool) = throw(ArgumentError("Unknown CANDECOMP method: $method"))

"""
Computes CANDECOMP by ALS (Alternating Least Squares) method.
"""
function _candecomp{T,N}(
    method::Type{Val{:ALS}},
    tnsr::StridedArray{T,N},
    r::Integer,
    factors::Vector{Matrix{Float64}},
    tol::Float64,
    maxiter::Integer,
    verbose::Bool)

    gram = [F'F for F in factors]
    tnsr_norm = vecnorm(tnsr)
    tnsr_flat = _col_unfold(tnsr, N)
    tnsr_size = size(tnsr)
    niters = 0
    converged = false
    resid = tnsr_norm
    lbds = Matrix{Float64}(1, r)
    V = Matrix{Float64}(div(length(tnsr), minimum(tnsr_size)), r)
    while !converged && niters < maxiter
        VB = 0
        @inbounds for i in 1:N
            idx = [N:-1:i + 1; i - 1:-1:1]
            VB = prod(k -> tnsr_size[k], idx)
            V[1:VB, :] = reduce(khatrirao, factors[idx])
            factors[i] = _row_unfold(tnsr, i) * V[1:VB, :] / reduce(.*, gram[idx])
            sum!(lbds, abs(factors[i]))
            factors[i] ./= lbds
            At_mul_B!(gram[i], factors[i], factors[i])
        end
        resid_old = resid
        resid = vecnorm(V[1:VB, :] * (factors[N] .* lbds)' - tnsr_flat)
        converged = abs(resid - resid_old) < tol * resid_old
        niters += 1
    end

    verbose && _iter_status(converged, niters, maxiter)

    return CANDECOMP((factors...), squeeze(lbds, 1))
end

"""
Computes CANDECOMP by SGSD (Simultaneous Generalized Schur Decomposition) method.
"""
function _candecomp{T,N}(
    method::Type{Val{:SGSD}},
    tnsr::StridedArray{T,N},
    r::Int,
    factors::Vector{Matrix{Float64}},
    tol::Float64,
    maxiter::Integer,
    verbose::Bool)

    ndims(tnsr) == 3 || throw(ArgumentError("This algorithm only applies to 3-mode tensors."))
    (n1, n2, n3) = size(tnsr)
    IB = [min(n1 - 1, r), (n2 == r) ? 2 : 1]
    Q = qr(factors[1], thin=false)[1]
    Z = qr(flipdim(factors[2], 2), thin=false)[1]
    q = Array(Float64, n1, n1)
    z = Array(Float64, n2, n2)

    R = zeros(size(tnsr))
    @tensor begin
        R[4,5,3] = tnsr[1,2,3] * Q[1,4] * Z[2,5]
    end

    res = vecnorm(tnsr)
    converged = false
    niters = 0
    @inbounds while !converged && niters < maxiter
        q = eye(n1)
        z = eye(n2)

        for i in 1:IB[1]
            q[:, i:n1] *= svd(q[:, i:n1]' * slice(R, :, n2 - r + i, :))[1]
        end

        @tensor begin
            R[4,2,3] = R[1,2,3] * q[1,4]
        end

        for i in r:-1:IB[2]
            z[:, 1:n2 - r + i] *= flipdim(svd(slice(R, i, :, :)' * z[:, 1:n2 - r + i])[3], 2)
        end
        Q *= q
        Z *= z

        @tensor begin
            R[4,5,3] = tnsr[1,2,3] * Q[1,4] * Z[2,5]
        end

        res_old = res
        res = vecnorm(tril(squeeze(sum(R .^ 2, 3), 3), n2 - r - 1))
        converged = abs(res - res_old) < tol
        niters += 1
    end

    verbose && _iter_status(converged, niters, maxiter)

    R = R[1:r, n2-r+1:n2, :]
    M = cat(3, eye(r, r), eye(r, r))
    @inbounds for i in r - 1:-1:1, j = i + 1:r
        d = i + 1:j - 1
        M[i, j, :] = hcat(R[j, j, :][:], R[i, i, :][:]) \ (R[i, j, :][:] - mapslices(R3 -> sum(M[i, d, 1] * (diag(R3)[d] .* M[d, j, 2])), R, [1, 2])[:])
    end

    factors[1] = Q[:, 1:r] * M[:, :, 1]
    factors[2] = Z[:, n2 - r + 1:n2] * M[:, :, 2]'
    factors[3] = _row_unfold(tnsr, 3) * khatrirao(factors[2], factors[1]) / ((factors[2]'factors[2]) .* (factors[1]'factors[1]))

    lbds = ones(1, r)
    for i in 1:3
        lbd = mapslices(vecnorm, factors[i], 1)
        factors[i] ./= lbd
        lbds .*= lbd
    end

    return CANDECOMP((factors...), squeeze(lbds, 1))
end

"""
Codes of implemented CANDECOMP methods.
"""
const CANDECOMP_methods = Symbol[t.parameters[1] for t in filter(t -> isa(t, DataType),
                                                                 [m.sig.parameters[1].parameters[1] for m in methods(_candecomp)])] |> Set{Symbol}


text:
using Documenter, ForecastAccuracy

makedocs(
    modules = [ForecastAccuracy],
    doctest  = false,
    clean    = true,
    format   = Documenter.HTML(mathengine = Documenter.MathJax()),
    sitename = "ForecastAccuracy.jl",
    authors  = "Guilherme Bodin",
    pages   = [
        "Home" => "index.md"
    ]
)

deploydocs(
    repo = "github.com/LAMPSPUC/ForecastAccuracy.jl.git",
)


text:
@show 1+1
@show 2+2
@show factorial(3)
fac = factorial
@show fac(3)
@show abs(-6)
@show 3-2(fac(2))-abs(-7)
Lion = 5
Pika = 6
@show Lion + Pika
@show a = 4
@show b = 3
@show a*b == 6
@show a*b == 2*6
nothing


text:
module SplitGraph

using TreeKnit
using TreeTools
using ProgressMeter

include("objects.jl")
include("tools.jl")
include("energy.jl")
include("likelihood.jl")

let verbose::Bool = false, vverbose::Bool = false
	global v() = verbose
	global vv() = vverbose
	global set_verbose(v) = (verbose = v)
	global set_vverbose(v) = (vverbose = v)
end

# function opttrees(γ, Trange, M, seq_lengths, t...)
# 	opttrees!(γ, Trange, M, seq_lengths, [copy(x, TreeTools.EmptyData) for x in t]...)
# end

# function opttrees(t...; kwargs...)
# 	opttrees!([copy(x, TreeTools.EmptyData) for x in t]...; kwargs...)
# end
"""
	opttrees!(t... ; kwargs...)

Return a list of MCCs for input trees.
Output:
1.
"""
function opttrees(t...;
	γ=2,
	seq_lengths=1000 * ones(Int64, length(t)),
	Trange=reverse(0.001:0.01:1.),
	M = 10,
	likelihood_sort=true,
	resolve=true,
	sa_rep = 1,
	verbose=false
)
	opttrees!(
		γ, Trange, M, seq_lengths, [copy(convert(Tree{TreeTools.EmptyData}, x)) for x in t]...;
		likelihood_sort, resolve, sa_rep, verbose
	)
end
function opttrees!(
	γ, Trange, M, seq_lengths, t::Vararg{Tree};
	likelihood_sort=true, resolve=true, sa_rep=1, verbose=false,
)
	set_verbose(verbose)

	treelist = t
	mcc = naive_mccs(treelist...)
	if length(mcc) == 1
		return mcc, 0, 0., Int64[], Float64[], Union{Missing,Float64}[]
	end
	mcc_names = TreeKnit.name_mcc_clades!(treelist, mcc)
	for (i,t) in enumerate(treelist)
		TreeKnit.reduce_to_mcc!(t, mcc)
	end
	g = trees2graph(treelist)

	# SA - Optimization
	oconfs, F, nfound = sa_opt(g, γ=γ, Trange=Trange, M=M, rep=sa_rep, resolve=resolve)
	# Computing likelihoods
	if length(oconfs) != 1
		v() && @info "Sorting $(length(oconfs)) topologically equivalent configurations."
		vv() && @info "Configurations\n $oconfs"
		vv() && @info g.labels
		oconf, L = sortconf(oconfs, treelist, g, seq_lengths, mcc_names, likelihood_sort, false)
	else
		oconf = oconfs[1]
		L = Union{Missing,Float64}[]
	end
	vv() && @info "Final configuration for this iteration: $oconf."
	vv() && @info "MCCs removed: $([mcc_names[x] for x in g.labels[.!oconf]])"
	return [mcc_names[x] for x in g.labels[.!oconf]], compute_energy(oconf,g), compute_F(oconf, g, γ), L
end

function sortconf(oconfs, trees, g::Graph, seq_lengths, mcc_names, likelihood_sort, E_sort)
	if E_sort # Only considering configurations of lowest energies
		E = [compute_energy(conf,g) for conf in oconfs]
		Emin = minimum(E)
		oconfs_ = oconfs[findall(x->x==Emin, E)]
		v() && @info "Removing ", length(oconfs) - length(oconfs_), " configurations using energy."
	else # Removing configurations where nothing is removed
		oconfs_ = oconfs[findall(c->sum(c)<length(c), oconfs)]
	end

	# Sorting the remaining configurations using Poisson likelihood ratios
	if length(oconfs_) == 1
		return oconfs_[1], Union{Missing,Float64}[]
	elseif !likelihood_sort
		v() && @info "No likelihood sort: picking a random configuration among remaining ones"
		return rand(oconfs_), Union{Missing,Float64}[]
	else
		v() && @info "Comparing $(length(oconfs_)) configurations using likelihood"
		L = Union{Missing,Float64}[]
		for conf in oconfs_
			push!(L, conf_likelihood(conf, g, seq_lengths, trees, mode=:time, v=vv()))
		end
		vv() && println("Confs: ", [[mcc_names[x] for x in g.labels[.!conf]] for conf in oconfs_])
		vv() && @info "Likelihoods: $L"
		Lmax = maximum(L)
		ismissing(Lmax) && @warn "Maximum likelihood is `missing`"
		oconfs_ = oconfs_[findall(isequal(Lmax), L)]
		if length(oconfs_) != 1 # Final sort by energy if more than one most likely conf
			E = [compute_energy(conf,g) for conf in oconfs_]
			Emin = minimum(E)
			oconfs_ = oconfs_[findall(isequal(Emin), E)]
		end
		return rand(oconfs_), L
	end
end

end # module

max_issues_repo_name: PierreBarrat/TreeKnit

max_forks_repo_name: PierreBarrat/TreeKnit


text:
# This file is a part of ValueShapes.jl, licensed under the MIT License (MIT).

using ValueShapes
using Test

using Random
using ElasticArrays, ArraysOfArrays

@testset "array_shape" begin
    @test_throws ArgumentError @inferred(ValueShapes._valshapeoftype(Vector{Int}))

    @test @inferred(valshape(rand(3))) == ArrayShape{Float64,1}((3,))
    @test @inferred(valshape(rand(3, 4, 5))) == ArrayShape{Float64,3}((3, 4, 5))

    @test @inferred(ValueShapes.default_unshaped_eltype(ArrayShape{Complex,3}((3, 4, 5)))) == Float64
    @test @inferred(ValueShapes.default_unshaped_eltype(ArrayShape{Complex{Float32},3}((3, 4, 5)))) == Float32

    @test @inferred(ValueShapes.shaped_type(ArrayShape{Real}(2, 3, 4))) == Array{Float64, 3}
    @test @inferred(ValueShapes.shaped_type(ArrayShape{Complex}(2))) == Array{Complex{Float64}, 1}
    @test @inferred(ValueShapes.shaped_type(ArrayShape{Complex{Real}}(2), Float32)) == Array{Complex{Float32}, 1}
    @test @inferred(ValueShapes.shaped_type(ArrayShape{Complex{Int16}}(2, 3))) == Array{Complex{Int16}, 2}

    @test @inferred(totalndof(ArrayShape{Float64,1}((3,)))) == 3
    @test @inferred(totalndof(ArrayShape{Complex,3}((3, 4, 5)))) == 120

    @test size(@inferred(Vector{Float64}(undef, ArrayShape{Complex}((2, 1, 3))))) == (2 * 2*1*3,)
    @test size(flatview(@inferred(VectorOfSimilarVectors{Float32}(ArrayShape{Complex}((2, 1, 3)))))) == (2 * 2*1*3, 0)

    shape = ArrayShape{Real}(2,3)
    A = @inferred(shape(undef))
    @test typeof(A) == Array{Float64,2}
    @test size(A) == (2, 3)
    @test @inferred(length(shape)) == shape.dims[1]*shape.dims[2]
    @test @inferred(length(shape)) == prod(shape.dims)
    @test @inferred(valshape(A)) == ArrayShape{Float64}(2,3)
    @test @inferred(shape([1, 2, 3, 4, 5, 6])) == [1 3 5; 2 4 6]
    @test_throws ArgumentError shape([1, 2, 3, 4, 5, 6, 7])

    @test eltype(@inferred(Vector{Float32}(undef, shape))) == Float32
    @test eltype(eltype(@inferred(VectorOfSimilarVectors{Float32}(shape)))) == Float32

    @test valshape(shape.(push!(@inferred(VectorOfSimilarVectors{Float64}(shape)), @inferred(Vector{Float64}(undef, shape))))[1]) == valshape(shape(undef))

    let
        A = collect(1:8)
        ac = ValueAccessor(ArrayShape{Real}(3), 2)
        @test @inferred(getindex(A, ac)) == [3, 4, 5]
        @test @inferred(view(A, ac)) == [3, 4, 5]
        @test @inferred(setindex!(A, [7, 2, 6], ac)) === A
        @test A[ac] == [7, 2, 6]
    end

    let
        A = collect(1:6*6)
        ac = ValueAccessor(ArrayShape{Real}(2,3), 17)
        @test @inferred(getindex(A, ac)) == [18 20 22; 19 21 23]
        @test @inferred(view(A, ac)) == [18 20 22; 19 21 23]
        @test @inferred(setindex!(A, [6 4 3; 2 1 5], ac)) === A
        @test A[ac] == [6 4 3; 2 1 5]
    end

    let
        A = collect(reshape(1:6*6, 6, 6))
        ac1 = ValueAccessor(ArrayShape{Real}(2), 2)
        ac2 = ValueAccessor(ArrayShape{Real}(3), 3)
        @test @inferred(getindex(A, ac1, ac2)) == [21 27 33; 22 28 34]
        @test @inferred(view(A, ac1, ac2)) == [21 27 33; 22 28 34]
        @test @inferred(setindex!(A, [6 4 3; 2 1 5], ac1, ac2)) === A
        @test A[ac1, ac2] == [6 4 3; 2 1 5]
    end

    let d1 = [11, 12, 13, 14], d2 = [21, 22]
        d = vcat(d1, d2)
        reshaped = shape(d)
        @test reshaped == reshape(d, (2,3))
    end

end