text:
module SymbolicRegression

include("operators.jl")
include("hyperparams.jl")

# Types
export Population,
    PopMember,
    HallOfFame,
    Options,

    #Functions:
    RunSR,
    SRCycle,
    calculateParetoFrontier,
    countNodes,
    printTree,
    stringTree,
    evalTreeArray,

    #Operators:
    plus,
    sub,
    mult,
    square,
    cube,
    pow,
    div,
    logm,
    logm2,
    logm10,
    sqrtm,
    neg,
    greater,
    relu,
    logical_or,
    logical_and

using Optim
using Printf: @printf
using Random: shuffle!, randperm
using Distributed

const maxdegree = 2

# The Val{i} optimizes it into a branching statement (https://discourse.julialang.org/t/meta-programming-an-if-else-statement-of-user-defined-length/53525)
function BINOP!(x::Array{Float32, 1}, y::Array{Float32, 1}, ::Val{i}, ::Val{clen}, options::Options) where {i,clen}
    op = options.binops[i]
    # broadcast!(op, x, x, y)
    @inbounds @simd for j=1:clen
        x[j] = op(x[j], y[j])
    end
end

function UNAOP!(x::Array{Float32, 1}, ::Val{i}, ::Val{clen}, options::Options) where {i,clen}
    op = options.unaops[i]
    @inbounds @simd for j=1:clen
        x[j] = op(x[j])
    end
end

# Sum of square error between two arrays
function SSE(x::Array{Float32}, y::Array{Float32})::Float32
    diff = (x - y)
    return sum(diff .* diff)
end
function SSE(x::Nothing, y::Array{Float32})::Float32
    return 1f9
end

# Sum of square error between two arrays, with weights
function SSE(x::Array{Float32}, y::Array{Float32}, w::Array{Float32})::Float32
    diff = (x - y)
    return sum(diff .* diff .* w)
end
function SSE(x::Nothing, y::Array{Float32}, w::Array{Float32})::Float32
    return Nothing
end

# Mean of square error between two arrays
function MSE(x::Nothing, y::Array{Float32})::Float32
    return 1f9
end

# Mean of square error between two arrays
function MSE(x::Array{Float32}, y::Array{Float32})::Float32
    return SSE(x, y)/size(x)[1]
end

# Mean of square error between two arrays
function MSE(x::Nothing, y::Array{Float32}, w::Array{Float32})::Float32
    return 1f9
end

# Mean of square error between two arrays
function MSE(x::Array{Float32}, y::Array{Float32}, w::Array{Float32})::Float32
    return SSE(x, y, w)/sum(w)
end

function id(x::Float32)::Float32
    x
end

function debug(verbosity, string...)
    verbosity > 0 ? println(string...) : nothing
end

function getTime()::Integer
    return round(Integer, 1e3*(time()-1.6e9))
end

# Define a serialization format for the symbolic equations:
mutable struct Node
    #Holds operators, variables, constants in a tree
    degree::Integer #0 for constant/variable, 1 for cos/sin, 2 for +/* etc.
    val::Union{Float32, Integer} #Either const value, or enumerates variable
    constant::Bool #false if variable
    op::Integer #enumerates operator (separately for degree=1,2)
    l::Union{Node, Nothing}
    r::Union{Node, Nothing}

    Node(val::Float32) = new(0, val, true, 1, nothing, nothing)
    Node(val::Integer) = new(0, val, false, 1, nothing, nothing)
    Node(op::Integer, l::Node) = new(1, 0.0f0, false, op, l, nothing)
    Node(op::Integer, l::Union{Float32, Integer}) = new(1, 0.0f0, false, op, Node(l), nothing)
    Node(op::Integer, l::Node, r::Node) = new(2, 0.0f0, false, op, l, r)

    #Allow to pass the leaf value without additional node call:
    Node(op::Integer, l::Union{Float32, Integer}, r::Node) = new(2, 0.0f0, false, op, Node(l), r)
    Node(op::Integer, l::Node, r::Union{Float32, Integer}) = new(2, 0.0f0, false, op, l, Node(r))
    Node(op::Integer, l::Union{Float32, Integer}, r::Union{Float32, Integer}) = new(2, 0.0f0, false, op, Node(l), Node(r))
end

# Copy an equation (faster than deepcopy)
function copyNode(tree::Node)::Node
   if tree.degree == 0
       return Node(tree.val)
   elseif tree.degree == 1
       return Node(tree.op, copyNode(tree.l))
    else
        return Node(tree.op, copyNode(tree.l), copyNode(tree.r))
   end
end

# Count the operators, constants, variables in an equation
function countNodes(tree::Node)::Integer
    if tree.degree == 0
        return 1
    elseif tree.degree == 1
        return 1 + countNodes(tree.l)
    else
        return 1 + countNodes(tree.l) + countNodes(tree.r)
    end
end

# Count the max depth of a tree
function countDepth(tree::Node)::Integer
    if tree.degree == 0
        return 1
    elseif tree.degree == 1
        return 1 + countDepth(tree.l)
    else
        return 1 + max(countDepth(tree.l), countDepth(tree.r))
    end
end

# Convert an equation to a string
function stringTree(tree::Node, options::Options)::String
    if tree.degree == 0
        if tree.constant
            return string(tree.val)
        else
            if options.useVarMap
                return varMap[tree.val]
            else
                if options.printZeroIndex
                    return "x$(tree.val - 1)"
                else
                    return "x$(tree.val)"
                end
            end
        end
    elseif tree.degree == 1
        return "$(options.unaops[tree.op])($(stringTree(tree.l, options)))"
    else
        return "$(options.binops[tree.op])($(stringTree(tree.l, options)), $(stringTree(tree.r, options)))"
    end
end

# Print an equation
function printTree(tree::Node, options::Options)
    println(stringTree(tree, options))
end

# Return a random node from the tree
function randomNode(tree::Node)::Node
    if tree.degree == 0
        return tree
    end
    a = countNodes(tree)
    b = 0
    c = 0
    if tree.degree >= 1
        b = countNodes(tree.l)
    end
    if tree.degree == 2
        c = countNodes(tree.r)
    end

    i = rand(1:1+b+c)
    if i <= b
        return randomNode(tree.l)
    elseif i == b + 1
        return tree
    end

    return randomNode(tree.r)
end

# Count the number of unary operators in the equation
function countUnaryOperators(tree::Node)::Integer
    if tree.degree == 0
        return 0
    elseif tree.degree == 1
        return 1 + countUnaryOperators(tree.l)
    else
        return 0 + countUnaryOperators(tree.l) + countUnaryOperators(tree.r)
    end
end

# Count the number of binary operators in the equation
function countBinaryOperators(tree::Node)::Integer
    if tree.degree == 0
        return 0
    elseif tree.degree == 1
        return 0 + countBinaryOperators(tree.l)
    else
        return 1 + countBinaryOperators(tree.l) + countBinaryOperators(tree.r)
    end
end

# Count the number of operators in the equation
function countOperators(tree::Node)::Integer
    return countUnaryOperators(tree) + countBinaryOperators(tree)
end

# Randomly convert an operator into another one (binary->binary;
# unary->unary)
function mutateOperator(tree::Node, options::Options)::Node
    if countOperators(tree) == 0
        return tree
    end
    node = randomNode(tree)
    while node.degree == 0
        node = randomNode(tree)
    end
    if node.degree == 1
        node.op = rand(1:options.nuna)
    else
        node.op = rand(1:options.nbin)
    end
    return tree
end

# Count the number of constants in an equation
function countConstants(tree::Node)::Integer
    if tree.degree == 0
        return convert(Integer, tree.constant)
    elseif tree.degree == 1
        return 0 + countConstants(tree.l)
    else
        return 0 + countConstants(tree.l) + countConstants(tree.r)
    end
end

# Randomly perturb a constant
function mutateConstant(
        tree::Node, T::Float32,
        options::Options)::Node
    # T is between 0 and 1.

    if countConstants(tree) == 0
        return tree
    end
    node = randomNode(tree)
    while node.degree != 0 || node.constant == false
        node = randomNode(tree)
    end

    bottom = 0.1f0
    maxChange = options.perturbationFactor * T + 1.0f0 + bottom
    factor = maxChange^Float32(rand())
    makeConstBigger = rand() > 0.5

    if makeConstBigger
        node.val *= factor
    else
        node.val /= factor
    end

    if rand() > options.probNegate
        node.val *= -1
    end

    return tree
end

# Evaluate an equation over an array of datapoints
function evalTreeArray(tree::Node, cX::Array{Float32, 2}, options::Options)::Union{Array{Float32, 1}, Nothing}
    clen = size(cX)[1]
    if tree.degree == 0
        if tree.constant
            return fill(tree.val, clen)
        else
            return copy(cX[:, tree.val])
        end
    elseif tree.degree == 1
        cumulator = evalTreeArray(tree.l, cX, options)
        if cumulator === nothing
            return nothing
        end
        op_idx = tree.op
        UNAOP!(cumulator, Val(op_idx), Val(clen), options)
        @inbounds for i=1:clen
            if isinf(cumulator[i]) || isnan(cumulator[i])
                return nothing
            end
        end
        return cumulator
    else
        cumulator = evalTreeArray(tree.l, cX, options)
        if cumulator === nothing
            return nothing
        end
        array2 = evalTreeArray(tree.r, cX, options)
        if array2 === nothing
            return nothing
        end
        op_idx = tree.op
        BINOP!(cumulator, array2, Val(op_idx), Val(clen), options)
        @inbounds for i=1:clen
            if isinf(cumulator[i]) || isnan(cumulator[i])
                return nothing
            end
        end
        return cumulator
    end
end

# Score an equation
function scoreFunc(X::Array{Float32, 2}, y::Array{Float32, 1}, baseline::Float32, tree::Node, options::Options)::Float32
    prediction = evalTreeArray(tree, X, options)
    if prediction === nothing
        return 1f9
    end
    if options.weighted
        mse = MSE(prediction, y, weights)
    else
        mse = MSE(prediction, y)
    end
    return mse / baseline + countNodes(tree)*options.parsimony
end

# Score an equation with a small batch
function scoreFuncBatch(X::Array{Float32, 2}, y::Array{Float32, 1}, baseline::Float32, tree::Node, options::Options)::Float32
    # options.batchSize
    batch_idx = randperm(size(X)[1])[1:options.batchSize]
    batch_X = X[batch_idx, :]
    prediction = evalTreeArray(tree, batch_X, options)
    if prediction === nothing
        return 1f9
    end
    size_adjustment = 1f0
    batch_y = y[batch_idx]
    if options.weighted
        batch_w = weights[batch_idx]
        mse = MSE(prediction, batch_y, batch_w)
        size_adjustment = 1f0 * size(X)[1] / options.batchSize
    else
        mse = MSE(prediction, batch_y)
    end
    return size_adjustment * mse / baseline + countNodes(tree)*options.parsimony
end

# Add a random unary/binary operation to the end of a tree
function appendRandomOp(tree::Node, options::Options, nfeatures::Int)::Node
    node = randomNode(tree)
    while node.degree != 0
        node = randomNode(tree)
    end

    choice = rand()
    makeNewBinOp = choice < options.nbin/(options.nuna + options.nbin)
    if rand() > 0.5
        left = Float32(randn())
    else
        left = rand(1:nfeatures)
    end
    if rand() > 0.5
        right = Float32(randn())
    else
        right = rand(1:nfeatures)
    end

    if makeNewBinOp
        newnode = Node(
            rand(1:options.nbin),
            left,
            right
        )
    else
        newnode = Node(
            rand(1:options.nuna),
            left
        )
    end
    node.l = newnode.l
    node.r = newnode.r
    node.op = newnode.op
    node.degree = newnode.degree
    node.val = newnode.val
    node.constant = newnode.constant
    return tree
end

# Insert random node
function insertRandomOp(tree::Node, options::Options, nfeatures::Int)::Node
    node = randomNode(tree)
    choice = rand()
    makeNewBinOp = choice < options.nbin/(options.nuna + options.nbin)
    left = copyNode(node)

    if makeNewBinOp
        right = randomConstantNode(nfeatures)
        newnode = Node(
            rand(1:options.nbin),
            left,
            right
        )
    else
        newnode = Node(
            rand(1:options.nuna),
            left
        )
    end
    node.l = newnode.l
    node.r = newnode.r
    node.op = newnode.op
    node.degree = newnode.degree
    node.val = newnode.val
    node.constant = newnode.constant
    return tree
end

# Add random node to the top of a tree
function prependRandomOp(tree::Node, options::Options, nfeatures::Int)::Node
    node = tree
    choice = rand()
    makeNewBinOp = choice < options.nbin/(options.nuna + options.nbin)
    left = copyNode(tree)

    if makeNewBinOp
        right = randomConstantNode(nfeatures)
        newnode = Node(
            rand(1:options.nbin),
            left,
            right
        )
    else
        newnode = Node(
            rand(1:options.nuna),
            left
        )
    end
    node.l = newnode.l
    node.r = newnode.r
    node.op = newnode.op
    node.degree = newnode.degree
    node.val = newnode.val
    node.constant = newnode.constant
    return node
end

function randomConstantNode(nfeatures::Int)::Node
    if rand() > 0.5
        val = Float32(randn())
    else
        val = rand(1:nfeatures)
    end
    newnode = Node(val)
    return newnode
end

# Return a random node from the tree with parent
function randomNodeAndParent(tree::Node, parent::Union{Node, Nothing})::Tuple{Node, Union{Node, Nothing}}
    if tree.degree == 0
        return tree, parent
    end
    a = countNodes(tree)
    b = 0
    c = 0
    if tree.degree >= 1
        b = countNodes(tree.l)
    end
    if tree.degree == 2
        c = countNodes(tree.r)
    end

    i = rand(1:1+b+c)
    if i <= b
        return randomNodeAndParent(tree.l, tree)
    elseif i == b + 1
        return tree, parent
    end

    return randomNodeAndParent(tree.r, tree)
end

# Select a random node, and replace it an the subtree
# with a variable or constant
function deleteRandomOp(tree::Node, options::Options, nfeatures::Int)::Node
    node, parent = randomNodeAndParent(tree, nothing)
    isroot = (parent === nothing)

    if node.degree == 0
        # Replace with new constant
        newnode = randomConstantNode(nfeatures)
        node.l = newnode.l
        node.r = newnode.r
        node.op = newnode.op
        node.degree = newnode.degree
        node.val = newnode.val
        node.constant = newnode.constant
    elseif node.degree == 1
        # Join one of the children with the parent
        if isroot
            return node.l
        elseif parent.l == node
            parent.l = node.l
        else
            parent.r = node.l
        end
    else
        # Join one of the children with the parent
        if rand() < 0.5
            if isroot
                return node.l
            elseif parent.l == node
                parent.l = node.l
            else
                parent.r = node.l
            end
        else
            if isroot
                return node.r
            elseif parent.l == node
                parent.l = node.r
            else
                parent.r = node.r
            end
        end
    end
    return tree
end

# Simplify tree
function combineOperators(tree::Node, options::Options)::Node
    # NOTE: (const (+*-) const) already accounted for. Call simplifyTree before.
    # ((const + var) + const) => (const + var)
    # ((const * var) * const) => (const * var)
    # ((const - var) - const) => (const - var)
    # (want to add anything commutative!)
    # TODO - need to combine plus/sub if they are both there.
    if tree.degree == 0
        return tree
    elseif tree.degree == 1
        tree.l = combineOperators(tree.l, options)
    elseif tree.degree == 2
        tree.l = combineOperators(tree.l, options)
        tree.r = combineOperators(tree.r, options)
    end

    top_level_constant = tree.degree == 2 && (tree.l.constant || tree.r.constant)
    if tree.degree == 2 && (options.binops[tree.op] === mult || options.binops[tree.op] === plus) && top_level_constant
        op = tree.op
        # Put the constant in r. Need to assume var in left for simplification assumption.
        if tree.l.constant
            tmp = tree.r
            tree.r = tree.l
            tree.l = tmp
        end
        topconstant = tree.r.val
        # Simplify down first
        below = tree.l
        if below.degree == 2 && below.op == op
            if below.l.constant
                tree = below
                tree.l.val = options.binops[op](tree.l.val, topconstant)
            elseif below.r.constant
                tree = below
                tree.r.val = options.binops[op](tree.r.val, topconstant)
            end
        end
    end

    if tree.degree == 2 && options.binops[tree.op] === sub && top_level_constant
        # Currently just simplifies subtraction. (can't assume both plus and sub are operators)
        # Not commutative, so use different op.
        if tree.l.constant
            if tree.r.degree == 2 && options.binops[tree.r.op] === sub
                if tree.r.l.constant
                    #(const - (const - var)) => (var - const)
                    l = tree.l
                    r = tree.r
                    simplified_const = -(l.val - r.l.val) #neg(sub(l.val, r.l.val))
                    tree.l = tree.r.r
                    tree.r = l
                    tree.r.val = simplified_const
                elseif tree.r.r.constant
                    #(const - (var - const)) => (const - var)
                    l = tree.l
                    r = tree.r
                    simplified_const = l.val + r.r.val #plus(l.val, r.r.val)
                    tree.r = tree.r.l
                    tree.l.val = simplified_const
                end
            end
        else #tree.r.constant is true
            if tree.l.degree == 2 && options.binops[tree.l.op] === sub
                if tree.l.l.constant
                    #((const - var) - const) => (const - var)
                    l = tree.l
                    r = tree.r
                    simplified_const = l.l.val - r.val#sub(l.l.val, r.val)
                    tree.r = tree.l.r
                    tree.l = r
                    tree.l.val = simplified_const
                elseif tree.l.r.constant
                    #((var - const) - const) => (var - const)
                    l = tree.l
                    r = tree.r
                    simplified_const = r.val + l.r.val #plus(r.val, l.r.val)
                    tree.l = tree.l.l
                    tree.r.val = simplified_const
                end
            end
        end
    end
    return tree
end

# Simplify tree
function simplifyTree(tree::Node, options::Options)::Node
    if tree.degree == 1
        tree.l = simplifyTree(tree.l, options)
        if tree.l.degree == 0 && tree.l.constant
            return Node(options.unaops[tree.op](tree.l.val))
        end
    elseif tree.degree == 2
        tree.l = simplifyTree(tree.l, options)
        tree.r = simplifyTree(tree.r, options)
        constantsBelow = (
             tree.l.degree == 0 && tree.l.constant &&
             tree.r.degree == 0 && tree.r.constant
        )
        if constantsBelow
            return Node(options.binops[tree.op](tree.l.val, tree.r.val))
        end
    end
    return tree
end

# Define a member of population by equation, score, and age
mutable struct PopMember
    tree::Node
    score::Float32
    birth::Integer

    PopMember(t::Node, score::Float32) = new(t, score, getTime())

end

function PopMember(X::Array{Float32, 2}, y::Array{Float32, 1}, baseline::Float32, t::Node, options::Options)
    PopMember(t, scoreFunc(X, y, baseline, t, options))
end

# Check if any binary operator are overly complex
function flagBinOperatorComplexity(tree::Node, op::Int, options::Options)::Bool
    if tree.degree == 0
        return false
    elseif tree.degree == 1
        return flagBinOperatorComplexity(tree.l, op, options)
    else
        if tree.op == op
            overly_complex = (
                    ((options.bin_constraints[op][1] > -1) &&
                     (countNodes(tree.l) > options.bin_constraints[op][1]))
                      ||
                    ((options.bin_constraints[op][2] > -1) &&
                     (countNodes(tree.r) > options.bin_constraints[op][2]))
                )
            if overly_complex
                return true
            end
        end
        return (flagBinOperatorComplexity(tree.l, op, options) || flagBinOperatorComplexity(tree.r, op, options))
    end
end

# Check if any unary operators are overly complex
function flagUnaOperatorComplexity(tree::Node, op::Int, options::Options)::Bool
    if tree.degree == 0
        return false
    elseif tree.degree == 1
        if tree.op == op
            overly_complex = (
                      (options.una_constraints[op] > -1) &&
                      (countNodes(tree.l) > options.una_constraints[op])
                )
            if overly_complex
                return true
            end
        end
        return flagUnaOperatorComplexity(tree.l, op, options)
    else
        return (flagUnaOperatorComplexity(tree.l, op, options) || flagUnaOperatorComplexity(tree.r, op, options))
    end
end

# Go through one simulated options.annealing mutation cycle
#  exp(-delta/T) defines probability of accepting a change
function iterate(X::Array{Float32, 2}, y::Array{Float32, 1}, baseline::Float32, member::PopMember, T::Float32, curmaxsize::Integer, frequencyComplexity::Array{Float32, 1}, options::Options)::PopMember
    prev = member.tree
    tree = prev
    #TODO - reconsider this
    if options.batching
        beforeLoss = scoreFuncBatch(X, y, baseline, prev, options)
    else
        beforeLoss = member.score
    end

    nfeatures = size(X)[2]

    mutationChoice = rand()
    #More constants => more likely to do constant mutation
    weightAdjustmentMutateConstant = min(8, countConstants(prev))/8.0
    cur_weights = copy(options.mutationWeights) .* 1.0
    cur_weights[1] *= weightAdjustmentMutateConstant
    n = countNodes(prev)
    depth = countDepth(prev)

    # If equation too big, don't add new operators
    if n >= curmaxsize || depth >= options.maxdepth
        cur_weights[3] = 0.0
        cur_weights[4] = 0.0
    end
    cur_weights /= sum(cur_weights)
    cweights = cumsum(cur_weights)

    successful_mutation = false
    #TODO: Currently we dont take this \/ into account
    is_success_always_possible = true
    attempts = 0
    max_attempts = 10

    #############################################
    # Mutations
    #############################################
    while (!successful_mutation) && attempts < max_attempts
        tree = copyNode(prev)
        successful_mutation = true
        if mutationChoice < cweights[1]
            tree = mutateConstant(tree, T, options)

            is_success_always_possible = true
            # Mutating a constant shouldn't invalidate an already-valid function

        elseif mutationChoice < cweights[2]
            tree = mutateOperator(tree, options)

            is_success_always_possible = true
            # Can always mutate to the same operator

        elseif mutationChoice < cweights[3]
            if rand() < 0.5
                tree = appendRandomOp(tree, options, nfeatures)
            else
                tree = prependRandomOp(tree, options, nfeatures)
            end
            is_success_always_possible = false
            # Can potentially have a situation without success
        elseif mutationChoice < cweights[4]
            tree = insertRandomOp(tree, options, nfeatures)
            is_success_always_possible = false
        elseif mutationChoice < cweights[5]
            tree = deleteRandomOp(tree, options, nfeatures)
            is_success_always_possible = true
        elseif mutationChoice < cweights[6]
            tree = simplifyTree(tree, options) # Sometimes we simplify tree
            tree = combineOperators(tree, options) # See if repeated constants at outer levels
            return PopMember(tree, beforeLoss)

            is_success_always_possible = true
            # Simplification shouldn't hurt complexity; unless some non-symmetric constraint
            # to commutative operator...

        elseif mutationChoice < cweights[7]
            tree = genRandomTree(5, options, nfeatures) # Sometimes we generate a new tree completely tree

            is_success_always_possible = true
        else # no mutation applied
            return PopMember(tree, beforeLoss)
        end

        # Check for illegal equations
        for i=1:options.nbin
            if successful_mutation && flagBinOperatorComplexity(tree, i, options)
                successful_mutation = false
            end
        end
        for i=1:options.nuna
            if successful_mutation && flagUnaOperatorComplexity(tree, i, options)
                successful_mutation = false
            end
        end

        attempts += 1
    end
    #############################################

    if !successful_mutation
        return PopMember(copyNode(prev), beforeLoss)
    end

    if options.batching
        afterLoss = scoreFuncBatch(X, y, baseline, tree, options)
    else
        afterLoss = scoreFunc(X, y, baseline, tree, options)
    end

    if options.annealing
        delta = afterLoss - beforeLoss
        probChange = exp(-delta/(T*options.alpha))
        if options.useFrequency
            oldSize = countNodes(prev)
            newSize = countNodes(tree)
            probChange *= frequencyComplexity[oldSize] / frequencyComplexity[newSize]
        end

        return_unaltered = (isnan(afterLoss) || probChange < rand())
        if return_unaltered
            return PopMember(copyNode(prev), beforeLoss)
        end
    end
    return PopMember(tree, afterLoss)
end

# Create a random equation by appending random operators
function genRandomTree(length::Integer, options::Options, nfeatures::Int)::Node
    tree = Node(1.0f0)
    for i=1:length
        tree = appendRandomOp(tree, options, nfeatures)
    end
    return tree
end

# A list of members of the population, with easy constructors,
#  which allow for random generation of new populations
mutable struct Population
    members::Array{PopMember, 1}
    n::Integer

    Population(pop::Array{PopMember, 1}) = new(pop, size(pop)[1])
    Population(pop::Array{PopMember, 1}, npop::Integer) = new(pop, npop)

end

function Population(X::Array{Float32, 2}, y::Array{Float32, 1}, baseline::Float32, npop::Integer, options::Options, nfeatures::Int)
    Population([PopMember(X, y, baseline, genRandomTree(3, options, nfeatures), options) for i=1:npop], npop)
end

function Population(X::Array{Float32, 2}, y::Array{Float32, 1}, baseline::Float32, npop::Integer, nlength::Integer, options::Options, nfeatures::Int)
    Population([PopMember(X, y, baseline, genRandomTree(nlength, options, nfeatures), options) for i=1:npop], npop)
end

# Sample 10 random members of the population, and make a new one
function samplePop(pop::Population, options::Options)::Population
    idx = rand(1:pop.n, options.ns)
    return Population(pop.members[idx])
end

# Sample the population, and get the best member from that sample
function bestOfSample(pop::Population, options::Options)::PopMember
    sample = samplePop(pop, options)
    best_idx = argmin([sample.members[member].score for member=1:sample.n])
    return sample.members[best_idx]
end

function finalizeScores(X::Array{Float32, 2}, y::Array{Float32, 1}, baseline::Float32, pop::Population, options::Options)::Population
    need_recalculate = options.batching
    if need_recalculate
        @inbounds @simd for member=1:pop.n
            pop.members[member].score = scoreFunc(X, y, baseline, pop.members[member].tree, options)
        end
    end
    return pop
end

# Return best 10 examples
function bestSubPop(pop::Population; topn::Integer=10)::Population
    best_idx = sortperm([pop.members[member].score for member=1:pop.n])
    return Population(pop.members[best_idx[1:topn]])
end

# Pass through the population several times, replacing the oldest
# with the fittest of a small subsample
function regEvolCycle(X::Array{Float32, 2}, y::Array{Float32, 1}, baseline::Float32, pop::Population, T::Float32, curmaxsize::Integer,
                      frequencyComplexity::Array{Float32, 1}, options::Options)::Population
    # Batch over each subsample. Can give 15% improvement in speed; probably moreso for large pops.
    # but is ultimately a different algorithm than regularized evolution, and might not be
    # as good.
    if options.fast_cycle
        shuffle!(pop.members)
        n_evol_cycles = round(Integer, pop.n/options.ns)
        babies = Array{PopMember}(undef, n_evol_cycles)

        # Iterate each ns-member sub-sample
        @inbounds Threads.@threads for i=1:n_evol_cycles
            best_score = Inf32
            best_idx = 1+(i-1)*options.ns
            # Calculate best member of the subsample:
            for sub_i=1+(i-1)*options.ns:i*options.ns
                if pop.members[sub_i].score < best_score
                    best_score = pop.members[sub_i].score
                    best_idx = sub_i
                end
            end
            allstar = pop.members[best_idx]
            babies[i] = iterate(X, y, baseline, allstar, T, curmaxsize, frequencyComplexity, options)
        end

        # Replace the n_evol_cycles-oldest members of each population
        @inbounds for i=1:n_evol_cycles
            oldest = argmin([pop.members[member].birth for member=1:pop.n])
            pop.members[oldest] = babies[i]
        end
    else
        for i=1:round(Integer, pop.n/options.ns)
            allstar = bestOfSample(pop, options)
            baby = iterate(X, y, baseline, allstar, T, curmaxsize, frequencyComplexity, options)
            #printTree(baby.tree)
            oldest = argmin([pop.members[member].birth for member=1:pop.n])
            pop.members[oldest] = baby
        end
    end

    return pop
end

# Cycle through regularized evolution many times,
# printing the fittest equation every 10% through
function SRCycle(X::Array{Float32, 2}, y::Array{Float32, 1}, baseline::Float32,
        pop::Population,
        ncycles::Integer,
        curmaxsize::Integer,
        frequencyComplexity::Array{Float32, 1};
        verbosity::Integer=0,
        options::Options
       )::Population

    allT = LinRange(1.0f0, 0.0f0, ncycles)
    for iT in 1:size(allT)[1]
        if options.annealing
            pop = regEvolCycle(X, y, baseline, pop, allT[iT], curmaxsize, frequencyComplexity, options)
        else
            pop = regEvolCycle(X, y, baseline, pop, 1.0f0, curmaxsize, frequencyComplexity, options)
        end

        if verbosity > 0 && (iT % verbosity == 0)
            bestPops = bestSubPop(pop)
            bestCurScoreIdx = argmin([bestPops.members[member].score for member=1:bestPops.n])
            bestCurScore = bestPops.members[bestCurScoreIdx].score
            debug(verbosity, bestCurScore, " is the score for ", stringTree(bestPops.members[bestCurScoreIdx].tree, options))
        end
    end

    return pop
end

# Get all the constants from a tree
function getConstants(tree::Node)::Array{Float32, 1}
    if tree.degree == 0
        if tree.constant
            return [tree.val]
        else
            return Float32[]
        end
    elseif tree.degree == 1
        return getConstants(tree.l)
    else
        both = [getConstants(tree.l), getConstants(tree.r)]
        return [constant for subtree in both for constant in subtree]
    end
end

# Set all the constants inside a tree
function setConstants(tree::Node, constants::Array{Float32, 1})
    if tree.degree == 0
        if tree.constant
            tree.val = constants[1]
        end
    elseif tree.degree == 1
        setConstants(tree.l, constants)
    else
        numberLeft = countConstants(tree.l)
        setConstants(tree.l, constants)
        setConstants(tree.r, constants[numberLeft+1:end])
    end
end

# Proxy function for optimization
function optFunc(x::Array{Float32, 1}, X::Array{Float32, 2}, y::Array{Float32, 1}, baseline::Float32, tree::Node, options::Options)::Float32
    setConstants(tree, x)
    return scoreFunc(X, y, baseline, tree, options)
end

# Use Nelder-Mead to optimize the constants in an equation
function optimizeConstants(X::Array{Float32, 2}, y::Array{Float32, 1}, baseline::Float32, member::PopMember, options::Options)::PopMember
    nconst = countConstants(member.tree)
    if nconst == 0
        return member
    end
    x0 = getConstants(member.tree)
    f(x::Array{Float32,1})::Float32 = optFunc(x, X, y, baseline, member.tree, options)
    if size(x0)[1] == 1
        algorithm = Newton
    else
        algorithm = NelderMead
    end

    try
        result = optimize(f, x0, algorithm(), Optim.Options(iterations=100))
        # Try other initial conditions:
        for i=1:options.nrestarts
            tmpresult = optimize(f, x0 .* (1f0 .+ 5f-1*randn(Float32, size(x0)[1])), algorithm(), Optim.Options(iterations=100))
            if tmpresult.minimum < result.minimum
                result = tmpresult
            end
        end

        if Optim.converged(result)
            setConstants(member.tree, result.minimizer)
            member.score = convert(Float32, result.minimum)
            member.birth = getTime()
        else
            setConstants(member.tree, x0)
        end
    catch error
        # Fine if optimization encountered domain error, just return x0
        if isa(error, AssertionError)
            setConstants(member.tree, x0)
        else
            throw(error)
        end
    end
    return member
end

# List of the best members seen all time
mutable struct HallOfFame
    members::Array{PopMember, 1}
    exists::Array{Bool, 1} #Whether it has been set

    # Arranged by complexity - store one at each.
end

function HallOfFame(options::Options)
    actualMaxsize = options.maxsize + maxdegree
    HallOfFame([PopMember(Node(1f0), 1f9) for i=1:actualMaxsize], [false for i=1:actualMaxsize])
end

# Check for errors before they happen
function testConfiguration(options::Options)
    test_input = LinRange(-100f0, 100f0, 99)

    try
        for left in test_input
            for right in test_input
                for binop in options.binops
                    test_output = binop.(left, right)
                end
            end
            for unaop in options.unaops
                test_output = unaop.(left)
            end
        end
    catch error
        @printf("\n\nYour configuration is invalid - one of your operators is not well-defined over the real line.\n\n\n")
        throw(error)
    end
end

function calculateParetoFrontier(X::Array{Float32, 2}, y::Array{Float32, 1},
                                 hallOfFame::HallOfFame, options::Options)
    # Dominating pareto curve - must be better than all simpler equations
    dominating = PopMember[]
    actualMaxsize = options.maxsize + maxdegree
    for size=1:actualMaxsize
        if hallOfFame.exists[size]
            member = hallOfFame.members[size]
            if options.weighted
                curMSE = MSE(evalTreeArray(member.tree, X, options), y, weights)
                member.score = curMSE
            else
                curMSE = MSE(evalTreeArray(member.tree, X, options), y)
                member.score = curMSE
            end
            numberSmallerAndBetter = 0
            for i=1:(size-1)
                if options.weighted
                    hofMSE = MSE(evalTreeArray(hallOfFame.members[i].tree, X, options), y, weights)
                else
                    hofMSE = MSE(evalTreeArray(hallOfFame.members[i].tree, X, options), y)
                end
                if (hallOfFame.exists[size] && curMSE > hofMSE)
                    numberSmallerAndBetter += 1
                end
            end
            betterThanAllSmaller = (numberSmallerAndBetter == 0)
            if betterThanAllSmaller
                push!(dominating, member)
            end
        end
    end
    return dominating
end

function RunSR(X::Array{Float32, 2}, y::Array{Float32, 1},
               niterations::Integer, options::Options)

    testConfiguration(options)

    if options.weighted
        avgy = sum(y .* weights)/sum(weights)
        baselineMSE = MSE(y, convert(Array{Float32, 1}, ones(size(X)[1]) .* avgy), weights)
    else
        avgy = sum(y)/size(X)[1]
        baselineMSE = MSE(y, convert(Array{Float32, 1}, ones(size(X)[1]) .* avgy))
    end

    nfeatures = size(X)[2]

    # 1. Start a population on every process
    allPops = Future[]
    # Set up a channel to send finished populations back to head node
    channels = [RemoteChannel(1) for j=1:options.npopulations]
    bestSubPops = [Population(X, y, baselineMSE, 1, options, nfeatures) for j=1:options.npopulations]
    hallOfFame = HallOfFame(options)
    actualMaxsize = options.maxsize + maxdegree
    frequencyComplexity = ones(Float32, actualMaxsize)
    curmaxsize = 3
    if options.warmupMaxsize == 0
        curmaxsize = options.maxsize
    end

    for i=1:options.npopulations
        future = @spawnat :any Population(X, y, baselineMSE, options.npop, 3, options, nfeatures)
        push!(allPops, future)
    end

    # # 2. Start the cycle on every process:
    @sync for i=1:options.npopulations
        @async allPops[i] = @spawnat :any SRCycle(X, y, baselineMSE, fetch(allPops[i]), options.ncyclesperiteration, curmaxsize, copy(frequencyComplexity)/sum(frequencyComplexity), verbosity=options.verbosity, options=options)
    end
    println("Started!")
    cycles_complete = options.npopulations * niterations
    if options.warmupMaxsize != 0
        curmaxsize += 1
        if curmaxsize > options.maxsize
            curmaxsize = options.maxsize
        end
    end

    last_print_time = time()
    num_equations = 0.0
    print_every_n_seconds = 5
    equation_speed = Float32[]

    for i=1:options.npopulations
        # Start listening for each population to finish:
        @async put!(channels[i], fetch(allPops[i]))
    end

    while cycles_complete > 0
        @inbounds for i=1:options.npopulations
            # Non-blocking check if a population is ready:
            if isready(channels[i])
                # Take the fetch operation from the channel since its ready
                cur_pop = take!(channels[i])
                bestSubPops[i] = bestSubPop(cur_pop, topn=options.topn)

                #Try normal copy...
                bestPops = Population([member for pop in bestSubPops for member in pop.members])

                for member in cur_pop.members
                    size = countNodes(member.tree)
                    frequencyComplexity[size] += 1
                    if member.score < hallOfFame.members[size].score
                        hallOfFame.members[size] = deepcopy(member)
                        hallOfFame.exists[size] = true
                    end
                end

                # Dominating pareto curve - must be better than all simpler equations
                dominating = calculateParetoFrontier(X, y, hallOfFame, options)
                open(options.hofFile, "w") do io
                    println(io,"Complexity|MSE|Equation")
                    for member in dominating
                        println(io, "$(countNodes(member.tree))|$(member.score)|$(stringTree(member.tree, options))")
                    end
                end
                cp(options.hofFile, options.hofFile*".bkup", force=true)

                # Try normal copy otherwise.
                if options.migration
                    for k in rand(1:options.npop, round(Integer, options.npop*options.fractionReplaced))
                        to_copy = rand(1:size(bestPops.members)[1])
                        cur_pop.members[k] = PopMember(
                            copyNode(bestPops.members[to_copy].tree),
                            bestPops.members[to_copy].score)
                    end
                end

                if options.hofMigration && size(dominating)[1] > 0
                    for k in rand(1:options.npop, round(Integer, options.npop*options.fractionReplacedHof))
                        # Copy in case one gets used twice
                        to_copy = rand(1:size(dominating)[1])
                        cur_pop.members[k] = PopMember(
                           copyNode(dominating[to_copy].tree), dominating[to_copy].score
                        )
                    end
                end

                # TODO: Turn off this async when debugging - any errors in this code
                #         are silent.
                # begin
                @async begin
                    allPops[i] = @spawnat :any let
                        tmp_pop = SRCycle(X, y, baselineMSE, cur_pop, options.ncyclesperiteration, curmaxsize, copy(frequencyComplexity)/sum(frequencyComplexity), verbosity=options.verbosity, options=options)
                        @inbounds @simd for j=1:tmp_pop.n
                            if rand() < 0.1
                                tmp_pop.members[j].tree = simplifyTree(tmp_pop.members[j].tree, options)
                                tmp_pop.members[j].tree = combineOperators(tmp_pop.members[j].tree, options)
                                if options.shouldOptimizeConstants
                                    tmp_pop.members[j] = optimizeConstants(X, y, baselineMSE, tmp_pop.members[j], options)
                                end
                            end
                        end
                        tmp_pop = finalizeScores(X, y, baselineMSE, tmp_pop, options)
                        tmp_pop
                    end
                    put!(channels[i], fetch(allPops[i]))
                end

                cycles_complete -= 1
                cycles_elapsed = options.npopulations * niterations - cycles_complete
                if options.warmupMaxsize != 0 && cycles_elapsed % options.warmupMaxsize == 0
                    curmaxsize += 1
                    if curmaxsize > options.maxsize
                        curmaxsize = options.maxsize
                    end
                end
                num_equations += options.ncyclesperiteration * options.npop / 10.0
            end
        end
        sleep(1e-3)
        elapsed = time() - last_print_time
        #Update if time has passed, and some new equations generated.
        if elapsed > print_every_n_seconds && num_equations > 0.0
            # Dominating pareto curve - must be better than all simpler equations
            current_speed = num_equations/elapsed
            average_over_m_measurements = 10 #for print_every...=5, this gives 50 second running average
            push!(equation_speed, current_speed)
            if length(equation_speed) > average_over_m_measurements
                deleteat!(equation_speed, 1)
            end
            average_speed = sum(equation_speed)/length(equation_speed)
            curMSE = baselineMSE
            lastMSE = curMSE
            lastComplexity = 0
            if options.verbosity > 0
                @printf("\n")
                @printf("Cycles per second: %.3e\n", round(average_speed, sigdigits=3))
                cycles_elapsed = options.npopulations * niterations - cycles_complete
                @printf("Progress: %d / %d total iterations (%.3f%%)\n", cycles_elapsed, options.npopulations * niterations, 100.0*cycles_elapsed/(options.npopulations*niterations))
                @printf("Hall of Fame:\n")
                @printf("-----------------------------------------\n")
                @printf("%-10s  %-8s   %-8s  %-8s\n", "Complexity", "MSE", "Score", "Equation")
                @printf("%-10d  %-8.3e  %-8.3e  %-.f\n", 0, curMSE, 0f0, avgy)
            end

            actualMaxsize = options.maxsize + maxdegree
            for size=1:actualMaxsize
                if hallOfFame.exists[size]
                    member = hallOfFame.members[size]
                    if options.weighted
                        curMSE = MSE(evalTreeArray(member.tree, X, options), y, weights)
                    else
                        curMSE = MSE(evalTreeArray(member.tree, X, options), y)
                    end
                    numberSmallerAndBetter = 0
                    for i=1:(size-1)
                        if options.weighted
                            hofMSE = MSE(evalTreeArray(hallOfFame.members[i].tree, X, options), y, weights)
                        else
                            hofMSE = MSE(evalTreeArray(hallOfFame.members[i].tree, X, options), y)
                        end
                        if (hallOfFame.exists[size] && curMSE > hofMSE)
                            numberSmallerAndBetter += 1
                        end
                    end
                    betterThanAllSmaller = (numberSmallerAndBetter == 0)
                    if betterThanAllSmaller
                        delta_c = size - lastComplexity
                        delta_l_mse = log(curMSE/lastMSE)
                        score = convert(Float32, -delta_l_mse/delta_c)
                        if options.verbosity > 0
                            @printf("%-10d  %-8.3e  %-8.3e  %-s\n" , size, curMSE, score, stringTree(member.tree, options))
                        end
                        lastMSE = curMSE
                        lastComplexity = size
                    end
                end
            end
            debug(options.verbosity, "")
            last_print_time = time()
            num_equations = 0.0
        end
    end
    return hallOfFame
end

end #module SR


text:
"""
file with structs
"""
module BasicStructs
using Parameters

"""
main struct with configuration and an algorithm function
"""
@with_kw struct StepDefStruct
    mainFun # main function to Be invoked over data  stored in DataList

end #StepDefStruct

"""
data about single array and its meta data
"""
@with_kw struct DataItem
    patientId::String # used for identification of patient
    voxX ::Float64 # Voxelspacing x
    voxY ::Float64 # Voxelspacing y
    voxZ ::Float64 # Voxelspacing z
    isCt ::Bool # if true we ware dealing with computer tomography image
    isPET::Bool # if true it is PET image
    isMRI::Bool # if true it is MRI image
    isMask::Bool # if true it is mask
    isAlgoOutput::Bool # if true it is saved output of our algorithm
    array # 3 dimensional array with our data
end #DataItem

"""
list of data items plus some meta data
"""
@with_kw struct DataList
    patientId ::String
    listOfData ::Vector<DataItem>
    mvspx ::Float64 # Voxelspacing x

end #DataList

"""
holding data needed for correct display
"""
@with_kw struct ForDiplayConf
    mvspx ::Float64 # Voxelspacing x

end #DataItem

end#BasicStructs

ext: jl


text:
push!(LOAD_PATH, "src/")
push!(LOAD_PATH, "../src/")
using Documenter, Solids
makedocs(sitename="Solids.jl")

deploydocs(
    repo = "github.com/plut/Solids.jl.git",
)


text:
# # Nonlinear Cart-Pole
#
#md # [![](https://img.shields.io/badge/show-nbviewer-579ACA.svg)](@__NBVIEWER_ROOT_URL__/models/Cart-Pole.ipynb)
#
# The model is a version of the cart-pole introduced by Barto, Sutton, and Anderson in [1].

# ## Model
#
# The dynamics of the cart-pole system are described as follows:
#
# ```math
# \begin{aligned}
# \ddot{x} &= \dfrac{u + mlω^2sin(θ)}{m_t} - \dfrac{ml(g sin(θ)- cos(θ))
#             (\dfrac{u + mlω^2sin(θ)}{m_t})}{l(\dfrac{4}{3} - \dfrac{m cos(θ)^2}{m_t})}
#             \dfrac{cos(θ)}{m_t}    \\
# \ddot{θ} &= \dfrac{(g sin(θ)- cos(θ))(\dfrac{u + mlω^2sin(θ)}{m_t})}{l(\dfrac{4}{3}
#             - \dfrac{m cos(θ)^2}{m_t})} \dfrac{cos(θ)}{m_t}
# \end{aligned}
# ```
# where ``u ∈ {−10,10}`` is the input force, which either pushes the cart left
# or right, ``g = 9.8`` is gravity, ``m = 0.1`` is the pole’s mass, ``l = 0.5``
# is half the pole’s length , ``m_t = 1.1`` is the total mass, ``x`` is the
# position of the cart, θ is the angle of the pendulum with respect to the
# positive y-axis, ``v = \dot{x}`` is the linear velocity of the cart, and
# ``ω = \dot{θ}`` is the angular velocity of the pendulum. The controller has four
# inputs ``(x, \dot{x}, \theta, ̇\dot{\theta})``, four layers with ``[24, 48, 12, 2]`` neurons
# respectively, and two outputs. The two outputs are then compared, and the
# input sent to the plant depends on which output index has the greatest value.
#  Thus, as an example if ``\textrm{output}_1 > \textrm{output}_2`` then the input force supplied
# to the plant is 10.  However if ``\textrm{output}_1 < \textrm{output}_2`` then the input supplied
# to the plant is -10.

using ClosedLoopReachability

@taylorize function cartpole!(du, u, p, t)
    local f, m, l, mt, g = 10, 0.1, 0.5, 1.1, 9.8
    sinθ = sin(u[3])
    cosθ = cos(u[3])
    aux = (f + m*l*u[4]^2*sinθ) / mt
    aux2 = l*(4/3 - m*cosθ^2/mt)
    aux3 = (g*sinθ- cosθ) * aux / aux2
    aux4 = cosθ/mt
    aux5 = m*l*aux3

    du[1] = u[2]
    du[2] = aux - aux5 * aux4
    du[3] = u[4]
    du[4] = aux3 * aux4

    return du
end

# define the initial-value problem
X₀ = Hyperrectangle(low=[-0.05, -0.05, -0.05, -0.05], high=[0.05, 0.05, 0.05, 0.05])

prob = @ivp(x' = cartpole!(x), dim: 4, x(0) ∈ X₀)

# solve it
## sol = solve(prob, T=1.0, alg=TMJets(maxsteps=20_000, abstol=1e-10));

# ## Specifications
#
# For this benchmark, the verification objective is to demonstrate that the
# pole will eventually reach the upward position and that it will remain there.
#  In other words, the goal is to achieve a value of ``θ = 0`` and stay there.
#  Some other specifications to be met are, for at least 12 seconds,
# ``x ∈ [-2.4,2.4]`` and ``θ ∈ [-15,15]`` degrees. The initial conditions for
# all state variables were chosen uniformly at random between ``[-0.05, 0.05]``.

# ## Results

using MAT
path = joinpath(@modelpath("Cart-Pole", "CartPole_Controller.mat"))
controller = read_nnet_mat(path, key="controller")

# ## References

# [^BSA83]: A. G. Barto, R. S. Sutton, and C. W. Anderson. *Neuronlike adaptive
#           elements that can solve difficult learning control problems.*
#           [IEEE Transactions on Systems, Man, and Cybernetics, SMC-13(5):834–846,
#           Sep. 1983](https://ieeexplore.ieee.org/abstract/document/6313077).
#


text:
struct GraphStructure
    m::Int64                      # Number fo grid cells
    nbs::Vector{Vector{Int64}}   # list of neighbors for each grid cell
end

function showGraphStructure(io::IO, obj::GraphStructure; prefix::String = "")

    println(io, prefix, "GraphStructure")
    println(io, prefix, "m :\t", obj.m)
    println(io, prefix, "nbs :\t", typeof(obj.nbs), "[", length(obj.nbs), "]")

end

function Base.show(io::IO, obj::GraphStructure)

    showGraphStructure(io, obj)

end


text:
# Struct representing the distance function used by the algorithm
struct DistanceFunction
    # String representing the distance type
    type::String

    # Default costructor
    DistanceFunction() = new("euclidean")
    DistanceFunction(s::String) = new(s)
end

# Used by the Slave to calculate distance betweeen 2 points' coordinates
function calculateDistance(f::DistanceFunction, a::Vector, b::Vector)
    if f.type=="sqeuclidean"
        return sqeuclidean(a,b)
    else
        return euclidean(a,b)
    end
end


text:
################################
# Solving the Replicator Equator equation,
# appearing in Corrolary 2.2 of
# https://www.biorxiv.org/content/10.1101/623330v1
# with a functor model
################################

cd(@__DIR__)
using HighDimPDE
using Random
using Test
using Flux
using Revise
using PyPlot
using UnPack
using Functors
plotting = true

tspan = (0f0,5f-1)
dt = 1f-2 # time step
μ(X,p,t) = 0f0 # advection coefficients
σ(X,p,t) = 1f-1 # diffusion coefficients
d = 5
ss0 = 5f-2 #std g0
U = 1f0
u_domain = repeat([-U,U]', d, 1)

##############################
####### Neural Network #######
##############################
batch_size = 1000
train_steps = 10000
K = 100

struct mymodel
        sigma_pred
        a_pred
end
@functor mymodel
(m::mymodel)(x) = m.a_pred ./ prod(m.sigma_pred) .* exp.(-5f-1 * sum((x ./ m.sigma_pred).^2,dims=1) )

sigma_pred = repeat([ss0],d)
a_pred = [1f0]
fn = mymodel(sigma_pred, a_pred)

opt = Flux.Optimiser(ExpDecay(0.1,
                1,
                1000,
                1e-6),
                ADAM() )#optimiser
alg = DeepSplitting(fn, K=K, opt = opt, mc_sample = UniformSampling(u_domain[:,1], u_domain[:,2]) )

##########################
###### PDE Problem #######
##########################
g(x) = Float32((2*π)^(-d/2)) * ss0^(- Float32(d) * 5f-1) * exp.(-5f-1 *sum(x .^2f0 / ss0, dims = 1)) # initial condition
m(x) = - 5f-1 * sum(x.^2, dims=1)
vol = prod(u_domain[:,2] - u_domain[:,1])
f(y, z, v_y, v_z, ∇v_y, ∇v_z, t) = max.(0f0, v_y) .* (m(y) .- vol * max.(0f0, v_z) .* m(z) ) # nonlocal nonlinear part of the

# defining the problem
prob = PIDEProblem(g, f, μ, σ, tspan,
                    u_domain = u_domain
                    )
# solving
@time xgrid,ts,sol = solve(prob,
                alg,
                dt,
                verbose = true,
                abstol=1f-6,
                maxiters = train_steps,
                batch_size=batch_size,
                use_cuda = true
                )

###############################
######### Plotting ############
###############################
if plotting
        clf()
        fig, ax = plt.subplots(1,2, sharey = true)

        xgrid1 = collect((-U:5f-3:U))
        xgrid = [vcat(x, fill(0f0,d-1)) for x in xgrid1]

        # Analytic sol
        function _SS(x, t, p)
                d = length(x)
                MM = σ(x, p, t) * ones(d)
                SSt = MM .* ((MM .* sinh.(MM *t) .+ ss0 .* cosh.( MM * t)) ./ (MM .* cosh.(MM * t ) .+ ss0 .* sinh.(MM * t)))
                return SSt
        end

        function uanal(x, t, p)
                d = length(x)
                return (2*π)^(-d/2) * prod(_SS(x, t, p) .^(-1/2)) * exp(-0.5 *sum(x .^2 ./ _SS(x, t, p)) )
        end

        for t in collect(tspan[1]: dt : tspan[2])
                ys = uanal.(xgrid, t, Ref(Dict()))
                ax[2].plot(xgrid1, reduce(hcat,ys)[:], label = "t = $t")
        end
        ax[2].set_title("Analytical solution")
        gcf()

        #Deepsplitting sol
        for i in 1:length(sol)
                ax[1].scatter(xgrid1, reduce(vcat,sol[i].(xgrid)), s = .2, label="t = $(dt * (i-1))")
        end
        gcf()

        ax[1].set_title("DeepSplitting")

        for _a in ax
                _a.legend()
        end
        gcf()
        savefig("hamel_$(d)d.pdf")

        #####
        # other DimensionMismatch
        #####
        if false
                dx = 0.05
                x = u_domain[1,1]:dx:u_domain[1,2]
                plt.contourf(x,x,g.(repeat(x,2)))
        end
end
gcf()


text:
#AnalogMTK: Base plot generation tools
#-------------------------------------------------------------------------------

#==Useful constants
===============================================================================#
const linlin = cons(:a, xyaxes=set(xscale=:lin, yscale=:lin))
const linlog = cons(:a, xyaxes=set(xscale=:lin, yscale=:log))
const loglin = cons(:a, xyaxes=set(xscale=:log, yscale=:lin))
const loglog = cons(:a, xyaxes=set(xscale=:log, yscale=:log))

const dfltline = cons(:a, line=set(style=:solid, color=:blue, width=2))

#==Validation
===============================================================================#
#Validate x, y pair simple vectors
function validatexy_vec(x, y)
	sz = size(x); sz_y = size(y)
	if sz != sz_y
		throw("x & y sizes do not match")
	elseif length(sz) > 1
		throw("Dimension of data too large (Simple vector expected).")
	end
	return
end

#Validate x, y pair for 1 parameter sweep (2D array)
function validatexy_1param(x, y)
	sz = size(x); sz_y = size(y)
	if sz != sz_y
		throw("x & y sizes do not match")
	elseif length(sz) > 2
		throw("Dimension of data too large (2D arrays expected).")
	end
	return
end

#==Waveform builders
===============================================================================#
function _wfrm__(x::Vector, y::Vector, sweepid::String)
	validatexy_vec(x, y)
	return DataF1(x,y)
end
_wfrm__(x::AbstractVector, y::Vector, sweepid::String) = _wfrm__(collect(x), y, sweepid)
function _wfrm__(x::Array{Float64,2}, y::Array{Float64,2}, sweepid::String)
	validatexy_1param(x, y)
	N = size(x, 1)
	wfrm = fill(DataRS{DataF1}, PSweep(sweepid, collect(1:N))) do i
		return DataF1(x[i,:], y[i,:]) #Returns from implicit "do" function
	end
	return wfrm
end

"""`waveform(x, y; sweepid="i")`

Create a waveform object with a given string to identify the sweep.
"""
waveform(x, y; sweepid::String="i") = _wfrm__(x, y, sweepid)

#Last line


text:
mutable struct QPStatus
    iters::Int
end

QPStatus() = QPStatus(-1)

function reset!(qps::QPStatus)
    qps.iters = 0
end
function next!(qps::QPStatus)
    qps.iters += 1
end

""" `pp = BoxConstrainedQP`
    Type for problem
    min_x 1/2xᵀGx + cᵀx, s.t dᵢ ≤ xᵢ ∀ n-m<i
    where `m=size(d,1), n=size(c,1)`.
    Only d .== 0 is allowed at the moment.
    Stores `d,G,c,sol,μλ` and some temporary variables
    Where `sol` is the solution after calling `solve!(qp)`
    G is factorization of the input `G` where
    G.M is the original input `G`
"""

struct BoxConstrainedQP{T, GT<:AbstractCholeskySpecial{T}, VT<:AbstractVector{T}}
    c::VT
    d::VT
    m::Int
    G::GT
    μλ::VT
    λdual::VT
    tmp1::VT # size(G,1)
    tmp2::VT # size(G,1)
    status::QPStatus
end

function BoxConstrainedQP(G::AbstractMatrix{T}, c::VT, d::VT; semidefinite=true, ϵ = sqrt(eps(T))) where {T, VT<:AbstractVector{T}, MT<:AbstractMatrix{T}}
    m = size(d,1)
    n = size(G,1) - m

    @assert issymmetric(G)
    @assert all(isequal(0), d) # For now we only allow positive constraints

    GF = if semidefinite
            F = cholesky(Hermitian(G + I*ϵ))
            CholeskySpecialShifted(F, G, ϵ)
        else
            F = cholesky(Hermitian(G))
            CholeskySpecial(F, G)
        end

    bQP = BoxConstrainedQP{T,typeof(GF),VT}(
        c,d,m,
        GF,
        similar(c),  # μλ
        similar(d),  # λdual
        similar(c),  # tmp1
        similar(c),  # tmp2
        QPStatus())
    return bQP
end

# Dispatch to create factor with correct shift
new_factor(M1, bQP::BoxConstrainedQP{T,GT}) where {T, GT<:CholeskySpecial} =
    cholesky(Hermitian(M1))
new_factor(M1, bQP::BoxConstrainedQP{T,GT}) where {T, GT<:CholeskySpecialShifted} =
    cholesky(Hermitian(M1 + I*bQP.G.shift))

function run_smartstart(bQP)
    n = size(bQP.G,1) - bQP.m
    M1 = copy(bQP.G.M)
    count = 0
    for i=(n+1):size(bQP.G,1)
        if bQP.c[i] > 0
            count += 1
            M1[:,i] .= 0
            M1[i,:] .= 0
        end
    end
    DEBUG && println("started with: $count active")
    empty!(bQP.G.idx)
    for i=(n+1):size(bQP.G,1)
        if bQP.c[i] > 0
            M1[i,i] = 1
            push!(bQP.G.idx, i)
        end
    end

    F = new_factor(M1, bQP)

    bQP.G.F.UL .= F.UL
    #
    # for i = 1:length(bQP.d)
    #     if bQP.c[n+i] > 0
    #         count += 1
    #         addactive!(bQP, i)
    #     end
    # end
    # println("started with: $count active")
    return
end

## Simple QP functionality
"""
Solve min 1/2 pᵀGp+gₖᵀp, s.t pᵢ=0, ∀ i ∈ Wᵢ
where gₖ = Gx+c
return optimal p and lagrange multipliers
"""
function solveEqualityQP(bQP::BoxConstrainedQP{T}, x::AbstractVector{T}) where T
    G = bQP.G.M
    μλ = bQP.μλ
    λ = bQP.λdual
    g = bQP.tmp1
    tmp = bQP.tmp2

    mul!(g, G, x)    # g = G*x
    g .+= bQP.c      # g = Gx+c
    μλ .= .-g
    # This may throw on semidefinite, is caught in findDescent
    ldiv!(bQP.G, μλ) # G*μλ = -g
    for i in bQP.G.idx # set elements in working set to zero
        μλ[i] = zero(T)
    end

    mul!(tmp, G, μλ)
    tmp .+= g # tmp = G*μλ + g
    tmpidx = sort!([bQP.G.idx...]) # Make sure λ is in predictable order
    for (i,j) in enumerate(tmpidx)
        λ[i] = tmp[j]
    end

    return μλ, view(λ, 1:length(bQP.G.idx))
end

"""
Solve min pᵀGp+gₖᵀp, s.t pᵢ=0, ∀ i ∈ Wᵢ
where gₖ = Gx+c
return optimal p and lagrange multipliers
"""
function solveEqualityOrInfiniteDescentQP(bQP::BoxConstrainedQP{T}, x::AbstractVector{T}) where T
    G = bQP.G.M
    μλ = bQP.μλ
    λ = bQP.λdual
    g = bQP.tmp1
    tmp = bQP.tmp2

    mul!(g, G, x)    # g = G*x
    g .+= bQP.c      # g = Gx+c
    μλ .= .-g

    _, projection = ldiv2!(bQP.G, μλ) # G*μλ = -g, or projection i.e. infinite descent
    for i in bQP.G.idx # set elements in working set to zero
        DEBUG && projection && @assert abs(μλ[i]) < 1e-10
        μλ[i] = zero(T)
    end

    if !projection # Fix lambda
        mul!(tmp, G, μλ)
        tmp .+= g # tmp = G*μλ + g
        tmpidx = sort!([bQP.G.idx...]) # Make sure λ is in predictable order
        for (i,j) in enumerate(tmpidx)
            λ[i] = tmp[j]
        end
    else
        DEBUG && println("g:", g)
        DEBUG && println("μλ:", μλ)
        DEBUG && println("dot(μλ, g): $(dot(μλ, g))")
        @assert dot(μλ, g) < 0 # Make sure we got descent direction
    end

    return μλ, view(λ, 1:length(bQP.G.idx)), projection
end

function findInfiniteDescent(bQP::BoxConstrainedQP{T}, x) where T
    G = bQP.G.M
    μλ = bQP.μλ
    λ = bQP.λdual
    g = bQP.tmp1
    tmp = bQP.tmp2

    mul!(g, G, x)    # g = G*x
    g .+= bQP.c      # g = Gx+c
    μλ .= zero(T)   # We want to solve G*μλ = 0, with projection from -g
    g .= .- g       # g = - g
    ldiv!(bQP.G, μλ, x0=g) # G*μλ = 0, projected from -g
    for i in bQP.G.idx # set elements in working set to zero
        DEBUG && @assert abs(μλ[i]) < 1e-10
        μλ[i] = zero(T)
    end
    DEBUG && println("g:", -g)
    DEBUG && println("μλ:", μλ)
    DEBUG && println("dot(μλ, g): $(dot(μλ, -g))")
    @assert dot(μλ, -g) < 0 # Make sure we got descent direction (g = -g)

    return μλ, view(λ, 1:length(bQP.G.idx))
end

function findblocking(bQP::BoxConstrainedQP{T}, xk, pk::AbstractVector{T}) where T
    m = bQP.m    # Number of inequality constraints
    Wk = bQP.G.idx .- m        # Indices in λ
    inotinW = setdiff(1:m, Wk) # The indices for λ not active
    μendi = length(xk)-m       # Where inequality starts in xk, pk
    minα = typemax(T)
    mini = -1
    for i in inotinW
        if pk[μendi+i] < 0 # aᵢᵀpₖ < 0
            # Note: bᵢ = 0
            v = -xk[μendi+i]/pk[μendi+i]
            if  v < minα
                minα = v
                mini = i
            end
        end
    end
    return minα, mini
end

function deleteactive!(bQP::BoxConstrainedQP, i)
    active = sort!([bQP.G.idx...]) # Make sure λ is in predictable order
    addrowcol!(bQP.G, active[i])
end

function addactive!(bQP::BoxConstrainedQP, i)
    deleterowcol!(bQP.G, (size(bQP.G,1) - bQP.m)  + i) # Inequality constraints are at end of QP.G
end

function findDescent(bQP, xk)
    local pk, λi, infinitedescent
    DEBUG && println("Finding descent")
    try
        # Expect solveEqualityQP to throw if min does not exist
        pk, λi = solveEqualityQP(bQP, xk)
        DEBUG && println("Found descent")
        infinitedescent = false
    catch
        # Find infinite descnet direction
        pk, λi = findInfiniteDescent(bQP, xk)
        DEBUG && println("Found infinite descent")
        infinitedescent = true
    end

    return pk, λi, infinitedescent
end

function findDescent(bQP::BoxConstrainedQP{T,GT}, xk) where {T, GT<:CholeskySpecialShifted}
    DEBUG && println("Finding descent")
    pk, λi, infinitedescent = solveEqualityOrInfiniteDescentQP(bQP, xk)

    DEBUG && !infinitedescent && println("Found descent")
    DEBUG &&  infinitedescent && println("Found infinite descent")
    # DEBUG && println("Finding descent")
    # try
    #     # Expect solveEqualityQP to throw if min does not exist
    #     pk, λi = solveEqualityQP(bQP, xk)
    #     DEBUG && println("Found descent")
    #     infinitedescent = false
    # catch
    #     # Find infinite descnet direction
    #     pk, λi = findInfiniteDescent(bQP, xk)
    #     DEBUG && println("Found infinite descent")
    #     infinitedescent = true
    # end

    return pk, λi, infinitedescent
end
"""
Active set method for QP as in Numerical Optimization, Nocedal Wright
    min_y,λ   1/2[y;λ]ᵀG[y;λ] + cᵀ[y;λ], s.t λ≥0
    where y ∈ R^m, λ ∈ R^n
"""
function solve!(bQP::BoxConstrainedQP{T}) where T
    #xk = fill(zero(T), m+n) # Feasible point
    # TODO better initial guess?
    xk = T.(abs.(randn(size(bQP.c,1))))
    Wk = bQP.G.idx
    for i in Wk
        xk[i] = zero(T) # Make sure we start feasible with respect to initial guess
    end
    done = false
    while !done
        next!(bQP.status) # Update iteration count and similar
        #println(bQP.status.iters)
        DEBUG && println("working set: ", Wk)
        pk, λi, infinitedescent = findDescent(bQP, xk)
        #println("infinitedescent: $infinitedescent")
        #println("norm xk: $(norm(xk)) ")
        #println("norm pk: $(norm(pk)) ")
        DEBUG && println("xk: ", xk)
        DEBUG && println("pk: ", pk)
        #println("λi: ", λi)
        if !infinitedescent && norm(pk) ≤ sqrt(eps(T))
            if all(v -> v ≥ 0, λi)
                # Optimal solution
                done = true
            else
                # Remove active constraint from Wi
                (v, idx) = findmin(λi)
                deleteactive!(bQP, idx)
            end
        else # p ≠ 0
            minα, mini = findblocking(bQP, xk, pk)
            DEBUG && println("mini: $mini")
            if minα < 0
                error("Unexpected error: minα less than 0")
            end
            # If infinite descent, go all the way, otherwise max 1
            if infinitedescent
                DEBUG && println("takning infinite descent direction, α=$minα")
                if minα == typemax(T) || mini <= 0
                    error("Dual seems to be unbounded")
                end
                xk .= xk .+ minα.*pk
                addactive!(bQP, mini)
            else
                α = min(one(T), minα)
                xk .= xk .+ α.*pk
                if minα ≤ 1
                    # mini is blocking constraint
                    addactive!(bQP, mini)
                else
                    #No blocking constraint, continue
                end
            end
        end
    end
    return xk
end

max_stars_repo_name: mfalt/QPDAS

max_issues_repo_name: mfalt/QPDAS

max_forks_repo_name: mfalt/QPDAS


text:
include("my_transformation.jl")
using SatelliteToolbox

"""
test for eceftoseof
"""
#=
satpos = [6771000 0 0]
satvel = [0 8000 0]

vin = [0. 0. 0.]
vout = eceftoseof(vin, satpos, satvel, "pos")
println(vout)
=#

"""
test for seoftoscsf
"""

vin = [1. 0. 0.]
qua = SatelliteToolbox.Quaternion(cos(π/4), sin(π/4), 0., 0.)
vout = seoftoscsf(vin, qua)
println(vout)

max_stars_repo_name: yui-project/ACSim

max_issues_repo_name: yui-project/ACSim

max_forks_repo_name: yui-project/ACSim


text:
struct RMPTrajectory{M,S,m}
    xm::Vector{SVector{m,S}}
    vm::Vector{SVector{m,S}}
    am::Vector{SVector{m,S}}
    robot_coord_rep::CoordinateRep
    CM::Vector{Chart}

    xns::Vector{Vector}
    vns::Vector{Vector}
    task_coord_reps::Vector{CoordinateRep}
    CNs::Vector{Vector{Chart}}

    Time::S
    dt::S
end

function RMPTrajectory(xm_in, vm_in, Time, dt, task::TaskX, CM::Chart{I,M}, CN,
        robot_coord_rep::ChartRep, task_coord_rep=ChartRep(), log_tasks=false) where {M,I}

    am, CN = single_task_acceleration(xm_in, vm_in, task, CM, CN, robot_coord_rep)
    xm0, vm0, am0 = [xm_in], [vm_in], [am]

    xns, vns, CNs = [], [], [CN]
    if log_tasks
        if task_coord_rep == ChartRep()
            xn, vn = task_differential_map_chart(xm_in, vm_in, task, CM, CN)
            push!(xns, xn)
            push!(vns, vn)
        else
            xne, vne = task_differential_map_chart_emb(xm_in, vm_in, task, CM, CN)
            push!(xns, xne)
            push!(vns, vne)
        end
    end

    RMPTrajectory{M,eltype(xm_in),dim(M)}(xm0, vm0, am0, robot_coord_rep, [CM],
        xns, vns, [task_coord_rep], [CNs], Time, dt)
end

function RMPTrajectory(xme_in, vme_in, Time, dt, task::TaskX, CM::Chart{I,M}, CN,
    robot_coord_rep::EmbRep, task_coord_rep=ChartRep(), log_tasks=false) where {M,I}

    CM = choose_chart_emb(xme_in, CM)
    ame, CN = single_task_acceleration(xme_in, vme_in, task, CM, CN, robot_coord_rep)
    xm0, vm0, am0 = [xme_in], [vme_in], [ame]

    xns, vns, CNs = [], [], [CN]
    if log_tasks
        if task_coord_rep == ChartRep()
            xn, vn = task_differential_map_emb_chart(xme_in, vme_in, task, CM, CN)
            push!(xns, xn)
            push!(vns, vn)
        else
            xne, vne = task_differential_map_emb(xme_in, vme_in, task, CM, CN)
            push!(xns, xne)
            push!(vns, vne)
        end
    end

    RMPTrajectory{M,eltype(xme_in),embdim(M)}(xm0, vm0, am0, robot_coord_rep, [CM],
        xns, vns, [task_coord_rep], [CNs], Time, dt)
end

function RMPTrajectory(xm_in, vm_in, Time, dt, tasks::TaskXList, CM::Chart{I,M}, CNs::ChartList,
        robot_coord_rep::ChartRep, task_coord_reps=fill(ChartRep(), length(tasks)),
        log_tasks=false) where {M,I}

    am, CNs = multiple_task_acceleration(xm_in, vm_in, tasks, CM, CNs,
        robot_coord_rep, log_task_chart=true)
    xm0, vm0, am0 = [xm_in], [vm_in], [am]

    xns, vns = [], []
    if log_tasks
        for i = 1:length(tasks)
            if task_coord_reps[i] == ChartRep()
                xn, vn = task_differential_map_chart(xm_in, vm_in, tasks[i], CM, CNs[i])
                push!(xns, [xn])
                push!(vns, [vn])
            else
                xne, vne = task_differential_map_chart_emb(xm_in, vm_in, tasks[i], CM, CNs[i])
                push!(xns, [xne])
                push!(vns, [vne])
            end
        end
    end

    RMPTrajectory{M,eltype(xm_in),dim(M)}(xm0, vm0, am0, robot_coord_rep, [CM],
        xns, vns, task_coord_reps, [CNs], Time, dt)
end

function RMPTrajectory(xme_in, vme_in, Time, dt, tasks::TaskXList, CM::Chart{I,M}, CNs::ChartList,
    robot_coord_rep::EmbRep, task_coord_reps=fill(ChartRep(), length(tasks)),
    log_tasks=false) where {M,I}

    CM = choose_chart_emb(xme_in, CM)
    ame, CNs = multiple_task_acceleration(xme_in, vme_in, tasks, CM, CNs,
        robot_coord_rep, log_task_chart=true)
    xm0, vm0, am0 = [xme_in], [vme_in], [ame]

    xns, vns = [], []
    if log_tasks
        for i = 1:length(tasks)
            if task_coord_reps[i] == ChartRep()
                xn, vn = task_differential_map_emb_chart(xme_in, vme_in, tasks[i], CM, CNs[i])
                push!(xns, [xn])
                push!(vns, [vn])
            else
                xne, vne = task_differential_map_emb(xme_in, vme_in, tasks[i], CM, CNs[i])
                push!(xns, [xne])
                push!(vns, [vne])
            end
        end
    end

    RMPTrajectory{M,eltype(xme_in),embdim(M)}(xm0, vm0, am0, robot_coord_rep, [CM],
        xns, vns, task_coord_reps, [CNs], Time, dt)
end

function rk4step(x, v, a, dt, f, args...)
    k1x = x
    k1v = v
    k1a = a

    k2x = k1x + dt*k1v/2
    k2v = k1v + dt*k1a/2
    k2a = f(k2x, k2v, args...)

    k3x = k1x + dt*k2v/2
    k3v = k1v + dt*k2a/2
    k3a = f(k3x, k3v, args...)

    k4x = k1x + dt*k3v
    k4v = k1v + dt*k3a
    k4a = f(k4x, k4v, args...)

    x_next = k1x + dt*(k1v + 2*(k2v + k3v) + k4v)/6
    v_next = k1v + dt*(k1a + 2*(k2a + k3a) + k4a)/6

    x_next, v_next
end

function rk4step_timedep(x, v, a, t, dt, f, args...)
    k1x = x
    k1v = v
    k1a = a

    k2x = k1x + dt*k1v/2
    k2v = k1v + dt*k1a/2
    k2a = f(k2x, k2v, t + dt/2, args...)

    k3x = k1x + dt*k2v/2
    k3v = k1v + dt*k2a/2
    k3a = f(k3x, k3v, t + dt/2, args...)

    k4x = k1x + dt*k3v
    k4v = k1v + dt*k3a
    k4a = f(k4x, k4v, t + dt, args...)

    x_next = k1x + dt*(k1v + 2*(k2v + k3v) + k4v)/6
    v_next = k1v + dt*(k1a + 2*(k2a + k3a) + k4a)/6

    x_next, v_next
end

propagate_task(xm, vm, task, CM, CNs, Time, dt) =
    propagate_task(xm, vm, task, CM, CNs, Time, dt, ChartRep())

function propagate_task(xm, vm, task::TaskX, CM, CN, Time, dt,
        robot_coord_rep::ChartRep; task_coord_rep=ChartRep(), log_tasks=false)
    Nsteps = Int(cld(Time, dt))
    traj = RMPTrajectory(xm, vm, Time, dt, task, CM, CN, robot_coord_rep, task_coord_rep, log_tasks)
    am, CN = traj.am[1], traj.CNs[1][1]

    for i = 1:Nsteps-1
        xm, vm = rk4step(xm, vm, am, dt,
            (x, v, args...) -> single_task_acceleration(x, v, args...)[1],
            task, CM, CN, robot_coord_rep)

        CM_next = choose_chart_chart(xm, CM)
        if CM_next != CM
            xm, vm = chart_transition_differential(xm, vm, CM, CM_next)
            CM = CM_next
        end
        am, CN = single_task_acceleration(xm, vm, task, CM, CN, robot_coord_rep)
        push!(traj.xm, xm)
        push!(traj.vm, vm)
        push!(traj.am, am)
        push!(traj.CM, CM)

        if log_tasks
            if task_coord_rep == ChartRep()
                xn, vn = task_differential_map_chart(xm, vm, task, CM, CN)
                push!(traj.xns, xn)
                push!(traj.vns, vn)
            else
                xne, vne = task_differential_map_chart_emb(xm, vm, task, CM, CN)
                push!(traj.xns, xne)
                push!(traj.vns, vne)
            end
        end
        push!(traj.CNs, [CN])
    end

    return traj
end

function propagate_task(xme, vme, task::TaskX, CM, CN, Time, dt,
        robot_coord_rep::EmbRep; task_coord_rep=ChartRep(), log_tasks=false)
    Nsteps = Int(cld(Time, dt))
    traj = RMPTrajectory(xme, vme, Time, dt, task, CM, CN, robot_coord_rep, task_coord_rep, log_tasks)
    ame, CM, CN = traj.am[1], traj.CM[1], traj.CNs[1][1]
    xm, vm, am = emb_to_chart_differential(xme, vme, ame, CM)

    for i = 1:Nsteps-1
        xm, vm = rk4step(xm, vm, am, dt,
            (x, v, args...) -> single_task_acceleration(x, v, args...)[1],
            task, CM, CN, ChartRep())
        CM_next = choose_chart_chart(xm, CM)
        if CM_next != CM
            xm, vm = chart_transition_differential(xm, vm, CM, CM_next)
            CM = CM_next
        end
        am, CN = single_task_acceleration(xme, vme, task, CM, CN, ChartRep())
        xme, vme, ame = chart_to_emb_differential(xm, vm, am, CM)
        push!(traj.xm, xme)
        push!(traj.vm, vme)
        push!(traj.am, ame)
        push!(traj.CM, CM)

        if log_tasks
            if task_coord_rep == ChartRep()
                xn, vn = task_differential_map_emb_chart(xme, vme, task, CM, CN)
                push!(traj.xns, xn)
                push!(traj.vns, vn)
            else
                xne, vne = task_differential_map_emb(xme, vme, task, CM, CN)
                push!(traj.xns, xne)
                push!(traj.vns, vne)
            end
        end
        push!(traj.CNs, [CN])
    end

    return traj
end

propagate_tasks(xm, vm, tasks, CM, CNs, Time, dt) =
    propagate_tasks(xm, vm, tasks, CM, CNs, Time, dt, ChartRep())

function propagate_tasks(xm, vm, tasks::TaskXList, CM, CNs::ChartList, Time, dt,
        robot_coord_rep::ChartRep; task_coord_reps=fill(ChartRep(), length(tasks)), log_tasks=false)
    Nsteps = Int(cld(Time, dt))
    Ntasks = length(tasks)
    traj = RMPTrajectory(xm, vm, Time, dt, tasks, CM, CNs,
        robot_coord_rep, task_coord_reps, log_tasks)
    xm, vm, am, CNs = traj.xm[1], traj.vm[1], traj.am[1], traj.CNs[1]

    for i = 1:Nsteps-1
        xm, vm = rk4step(xm, vm, am, dt,
            (x, v, args...) -> multiple_task_acceleration(x, v, args...)[1],
            tasks, CM, CNs, robot_coord_rep)

        CM_next = choose_chart_chart(xm, CM)
        if CM_next != CM
            xm, vm = chart_transition_differential(xm, vm, CM, CM_next)
            CM = CM_next
        end
        am, CNs = multiple_task_acceleration(xm, vm, tasks, CM, CNs,
            robot_coord_rep, log_task_chart=true)
        push!(traj.xm, xm)
        push!(traj.vm, vm)
        push!(traj.am, am)
        push!(traj.CM, CM)

        if log_tasks
            for j in 1:Ntasks
                if task_coord_reps[j] == ChartRep()
                    xn, vn = task_differential_map_chart(xm, vm, tasks[j], CM, CNs[j])
                    push!(traj.xns[j], xn)
                    push!(traj.vns[j], vn)
                else
                    xne, vne = task_differential_map_chart_emb(xm, vm, tasks[j], CM, CNs[j])
                    push!(traj.xns[j], xne)
                    push!(traj.vns[j], vne)
                end
            end
            push!(traj.CNs, CNs)
        end
    end

    return traj
end

function propagate_tasks(xme, vme, tasks::TaskXList, CM, CNs::ChartList, Time, dt,
        robot_coord_rep::EmbRep; task_coord_reps=fill(ChartRep(), length(tasks)), log_tasks=false)
    Nsteps = Int(cld(Time, dt))
    Ntasks = length(tasks)
    traj = RMPTrajectory(xme, vme, Time, dt, tasks, CM, CNs,
        robot_coord_rep, task_coord_reps, log_tasks)
    ame, CM, CN = traj.am[1], traj.CM[1], traj.CNs[1][1]
    xm, vm, am = emb_to_chart_differential(xme, vme, ame, CM)

    for i = 1:Nsteps-1
        xm, vm = rk4step(xm, vm, am, dt,
            (x, v, args...) -> multiple_task_acceleration(x, v, args...)[1],
            tasks, CM, CNs, ChartRep())
        CM_next = choose_chart_chart(xm, CM)
        if CM_next != CM
            xm, vm = chart_transition_differential(xm, vm, CM, CM_next)
            CM = CM_next
        end
        am, CN = multiple_task_acceleration(xm, vm, tasks, CM, CNs, ChartRep())
        xme, vme, ame = chart_to_emb_differential(xm, vm, am, CM)
        push!(traj.xm, xme)
        push!(traj.vm, vme)
        push!(traj.am, ame)
        push!(traj.CM, CM)

        if log_tasks
            for j in 1:Ntasks
                if task_coord_reps[j] == ChartRep()
                    xn, vn = task_differential_map_emb_chart(xme, vme, tasks[j], CM, CNs[j])
                    push!(traj.xns[j], xn)
                    push!(traj.vns[j], vn)
                else
                    xne, vne = task_differential_map_emb(xme, vme, tasks[j], CM, CNs[j])
                    push!(traj.xns[j], xne)
                    push!(traj.vns[j], vne)
                end
            end
            push!(traj.CNs, CNs)
        end
    end

    return traj
end

struct RMPTreeTrajectory{M,S,m}
    xm::Vector{SVector{m,S}}
    vm::Vector{SVector{m,S}}
    am::Vector{SVector{m,S}}
    robot_coord_rep::CoordinateRep
    CM::Vector{Chart}
    root::TreeNode{M}

    Time::S
    dt::S
end

function RMPTreeTrajectory(xm_in, vm_in, Time, dt, root::TreeNode{M}, CM,
        robot_coord_rep::ChartRep, args...; log_tasks=false, time_dep=false) where M <: Manifold

    log_tasks && clear_logs!(root)
    if !time_dep
        am = task_acceleration(xm_in, vm_in, root, CM, args..., robot_coord_rep; log_tasks)
    else
        am = task_acceleration(xm_in, vm_in, 0., root, CM, args..., robot_coord_rep; log_tasks)
    end
    xm0, vm0, am0 = [xm_in], [vm_in], [am]

    RMPTreeTrajectory{M,eltype(xm_in),dim(M)}(xm0, vm0, am0, robot_coord_rep, [CM], root,
        Time, dt)
end

function RMPTreeTrajectory(xme_in, vme_in, Time, dt, root::TreeNode{M}, CM,
        robot_coord_rep::EmbRep, args...; log_tasks=false, time_dep=false) where M <: Manifold

    log_tasks && clear_logs!(root)
    CM = choose_chart_emb(xme_in, CM)
    if !time_dep
        ame = task_acceleration(xme_in, vme_in, root, CM, args..., robot_coord_rep; log_tasks)
    else
        ame = task_acceleration(xme_in, vme_in, 0., root, CM, args..., robot_coord_rep; log_tasks)
    end
    xm0, vm0, am0 = [xme_in], [vme_in], [ame]

    RMPTreeTrajectory{M,eltype(xme_in),embdim(M)}(xm0, vm0, am0, robot_coord_rep, [CM], root,
        Time, dt)
end

function propagate_tasks(xm, vm, root::TreeNode{M}, CM, Time, dt,
        robot_coord_rep::ChartRep, args...; log_tasks=false, time_dep=false) where M <: Manifold
    Nsteps = Int(cld(Time, dt))
    set_charts_chart!(root, xm, CM)
    traj = RMPTreeTrajectory(xm, vm, Time, dt, root, CM, robot_coord_rep, args...;
        log_tasks, time_dep)
    am = traj.am[1]
    t = 0.

    for i = 1:Nsteps-1
        if !time_dep
            xm, vm = rk4step(xm, vm, am, dt,
                (x, v, args...) -> task_acceleration(x, v, args...),
                root, CM, args..., robot_coord_rep)
        else
            xm, vm = rk4step_timedep(xm, vm, am, t, dt,
                (x, v, args...) -> task_acceleration(x, v, args...),
                root, CM, args..., robot_coord_rep)
        end

        CM_next = choose_chart_chart(xm, CM)
        if CM_next != CM
            xm, vm = chart_transition_differential(xm, vm, CM, CM_next)
            CM = CM_next
        end
        if !time_dep
            am = task_acceleration(xm, vm, root, CM, args..., robot_coord_rep; log_tasks)
        else
            am = task_acceleration(xm, vm, t, root, CM, args..., robot_coord_rep; log_tasks)
        end
        push!(traj.xm, xm)
        push!(traj.vm, vm)
        push!(traj.am, am)
        push!(traj.CM, CM)
        t += dt
    end

    return traj
end

function propagate_tasks(xme, vme, root::TreeNode{M}, CM, Time, dt,
        robot_coord_rep::EmbRep, args...; log_tasks=false, time_dep=false) where M <: Manifold
    Nsteps = Int(cld(Time, dt))
    set_charts_emb!(root, xme, CM)
    traj = RMPTreeTrajectory(xme, vme, Time, dt, root, CM, robot_coord_rep, args...;
        log_tasks, time_dep)
    ame, CM = traj.am[1], traj.CM[1]
    xm, vm, am = emb_to_chart_differential(xme, vme, ame, CM)
    t = 0.

    for i = 1:Nsteps-1
        if !time_dep
            xm, vm = rk4step(xm, vm, am, dt,
                (x, v, args...) -> task_acceleration(x, v, args...),
                root, CM, args..., ChartRep())
        else
            xm, vm = rk4step_timedep(xm, vm, am, t, dt,
                (x, v, args...) -> task_acceleration(x, v, args...),
                root, CM, args..., ChartRep())
        end
        CM_next = choose_chart_chart(xm, CM)
        if CM_next != CM
            xm, vm = chart_transition_differential(xm, vm, CM, CM_next)
            CM = CM_next
        end
        if !time_dep
            am = task_acceleration(xm, vm, root, CM, args..., ChartRep(); log_tasks)
        else
            am = task_acceleration(xm, vm, t, root, CM, args..., ChartRep(); log_tasks)
        end
        xme, vme, ame = chart_to_emb_differential(xm, vm, am, CM)
        push!(traj.xm, xme)
        push!(traj.vm, vme)
        push!(traj.am, ame)
        push!(traj.CM, CM)
        t += dt
    end

    return traj
end


text:
################################################################################
#
#  Compute unit group from class group context
#
################################################################################

function _unit_group(O::NfOrd, c::ClassGrpCtx)
  u = UnitGrpCtx{FacElem{nf_elem, AnticNumberField}}(O)
  _unit_group_find_units(u, c)
  return u
end

# TH:
# Current strategy
# ================
# Compute some "random" kernel elements using the transformation matrix to first
# find r independent units, where r is the unit rank.
# In the second round, try to enlarge the unit group with some random kernel
# elements.
function _unit_group_find_units_with_transform(u::UnitGrpCtx, x::ClassGrpCtx)
  @vprint :UnitGroup 1 "Processing ClassGrpCtx to find units ... \n"

  @vprint :UnitGroup 1 "Relation module  $(x.M)\n"

  O = order(u)

  #ker = transpose(ker)

  K = nf(order(x.FB.ideals[1]))
  r = unit_rank(O)
  r1, r2 = signature(O)

  A = u.units

  j = 0

  dim_ker = length(x.M.rel_gens.rows) # dimension of the kernel

  total_dim = length(x.M.bas_gens.rows) + dim_ker

  kelem = fmpz[ 0 for i in 1:total_dim ]

  trafos = x.M.trafo

  MAX_RND_RD_1 = 2*r

  time_indep = 0.0
  time_add_dep_unit = 0.0
  time_kernel = 0.0
  time_torsion = 0.0

  while(length(A) < r)

    @assert length(kelem) == length(x.R_gen) + length(x.R_rel)

    for i in 1:length(kelem)
      kelem[i] = 0
    end

    if j > MAX_RND_RD_1
      return 0
    end
    @vprint :UnitGroup 1 "Found $(length(A)) independent units so far ($(r - length(A)) left to find)\n"
    j = j + 1

    for i in 1:dim_ker
      kelem[end - i + 1] = rand(0:1)
    end

    time_kernel += @elapsed for i in length(trafos):-1:1
      apply_right!(kelem, trafos[i])
    end

    _make_row_primitive!(kelem)

    @assert length(kelem) == length(x.R_gen) + length(x.R_rel)

    y = FacElem(vcat(x.R_gen, x.R_rel), kelem)

    time_torsion += @elapsed is_tors, p = istorsion_unit(y, false, u.tors_prec)
    u.tors_prec = max(p, u.tors_prec)
    if is_tors
      continue
    end

    @vprint :UnitGroup 1 "Exponents are of bit size $(maximum([ nbits(o) for o in kelem]))\n"

    time_indep += @elapsed _add_unit(u, y)

  end

  @vprint :UnitGroup 1 "Found $r linear independent units \n"

  @vprint :UnitGroup 1 "Independent unit time: $time_indep\n"
  @vprint :UnitGroup 1 "Adding dependent unit time: $time_add_dep_unit\n"
  @vprint :UnitGroup 1 "Torsion test time: $time_torsion\n"
  @vprint :UnitGroup 1 "Kernel time: $time_kernel\n"

  u.full_rank = true

  u.units = reduce(u.units, u.tors_prec)

  not_larger = 0

  @vprint :UnitGroup 1 "Enlarging unit group by adding more kernel basis elements ...\n"
  while not_larger < 3

    @assert length(kelem) == length(x.R_gen) + length(x.R_rel)

    for i in 1:length(kelem)
      kelem[i] = 0
    end

    for i in 1:dim_ker
      kelem[end - i + 1] = rand(-2:2)
    end

    time_kernel += @elapsed for i in length(trafos):-1:1
      apply_right!(kelem, trafos[i])
    end

    @assert length(kelem) == length(x.R_gen) + length(x.R_rel)

    y = FacElem(vcat(x.R_gen, x.R_rel), kelem)

    @hassert :UnitGroup 2 _isunit(y)
    @vprint :UnitGroup 2 "Reduce 1st...\n"
    y = reduce_mod_units([y], u)[1]

    @vprint :UnitGroup 2 "Test if kernel element yields torsion unit ... \n"
    @v_do :UnitGroup 2 pushindent()
    time_torsion += @elapsed is_tors, p = istorsion_unit(y, false, u.tors_prec)
    u.tors_prec = max(p, u.tors_prec)
    if is_tors
      @v_do :UnitGroup 2 popindent()
      @vprint :UnitGroup 2 "Element is torsion unit\n"
      not_larger = not_larger + 1
      continue
    end
    @v_do :UnitGroup 2 popindent()

    @v_do :UnitGroup 2 pushindent()
    time_add_dep_unit += @elapsed m = _add_dependent_unit(u, y)::Bool
    @v_do :UnitGroup 2 popindent()

    if !m
      not_larger = not_larger + 1
    else
      not_larger = 0
    end
  end
  #final reduction ...
  u.units = reduce(u.units, u.tors_prec)

  u.tentative_regulator = regulator(u.units, 64)

  @vprint :UnitGroup 1 "Finished processing\n"
  @vprint :UnitGroup 1 "Regulator of current unit group is $(u.tentative_regulator)\n"
  @vprint :UnitGroup 1 "-"^80 * "\n"
  @vprint :UnitGroup 1 "Independent unit time: $time_indep\n"
  @vprint :UnitGroup 1 "Adding dependent unit time: $time_add_dep_unit\n"
  @vprint :UnitGroup 1 "Torsion test time: $time_torsion\n"
  @vprint :UnitGroup 1 "Kernel time: $time_kernel\n"
  @vtime_add :UnitGroup 1 x :unit_hnf_time time_kernel
  return 1
end

function _unit_group_find_units(u::UnitGrpCtx, x::ClassGrpCtx)
  @vprint :UnitGroup 1 "Processing ClassGrpCtx to find units ... \n"
  @vprint :UnitGroup 1 "Relation module $(x.M)\n"

  O = order(u)

  K = nf(order(x.FB.ideals[1]))
  r = unit_rank(O)

  if r == 0
    Ar = ArbField(u.indep_prec, false)
    u.tentative_regulator = Ar(1)
    u.regulator = Ar(1)
    u.regulator_precision = u.indep_prec
    u.full_rank = true
    return 1
  end

  # I am not allowed to do this before the other block
  if nrows(x.M.rel_gens) == 0
    @vprint :UnitGroup 1 "No additional relations. Going back ...\n"
    return 0
  end

  r1, r2 = signature(O)

  A = u.units

  j = 0

  MAX_RND_RD_1 = 2*r

  time_indep = 0.0
  time_add_dep_unit = 0.0
  time_kernel = 0.0
  time_torsion = 0.0

  j = 0

  not_larger = 0

  @vprint :UnitGroup 1 "Enlarging unit group by adding kernel elements ...\n"

  while not_larger < 5

    add_units = Int[]
    rel = SMat{fmpz}()
    for jj in 1:min(div(nrows(x.M.rel_gens), 10)+1, 20)
      xj = rand(1:nrows(x.M.rel_gens))
      if xj in add_units
        continue
      end
      push!(add_units, xj)
      push!(rel, x.M.rel_gens[xj])
    end

    @vprint :UnitGroup 1 "Saturating ...\n"
    time_kernel += @elapsed k, d = solve_dixon_sf(x.M.bas_gens, rel)
    @vtime_add_elapsed :UnitGroup 1 x :saturate_time s = saturate(hcat(k, (-d)*identity_matrix(SMat, FlintZZ, k.r)))
    @vprint :UnitGroup 1 "Done\n"

    ge = vcat(x.R_gen[1:k.c], x.R_rel[add_units])
    for i=1:s.r
      @vprint :UnitGroup 1 "Looking at the $(i)-th element in the saturation ...\n"
      y = FacElem(ge[s[i].pos], s[i].values)
      @hassert :UnitGroup 2 _isunit(y)
      @vprint :UnitGroup 1 "(It really is a unit.)\n"

      @hassert :UnitGroup 9000 denominator(minpoly(evaluate(y))) == 1

      if u.full_rank
        @vprint :UnitGroup 2 "have full rank, can reduce unit first...\n"
        y = reduce_mod_units([y], u)[1]
      else
        @vprint :UnitGroup 2 "no full rank, cannot reduce unit first...\n"
      end

      @vprint :UnitGroup 1 "Exponents are of bit size $(maximum([ nbits(o) for o in values(y.fac)]))\n"

      @vprint :UnitGroup 1 "Test if kernel element yields torsion unit ... \n"
      @v_do :UnitGroup 2 pushindent()

      time_torsion += @elapsed is_tors, p = istorsion_unit(y, false, u.tors_prec)

      u.tors_prec = max(p, u.tors_prec)
      if is_tors
        @v_do :UnitGroup 2 popindent()
        @vprint :UnitGroup 1 "Element is torsion unit\n"
        not_larger += 1
        if u.full_rank && not_larger > 5
          break
        end
        continue
      end
      @vprint :UnitGroup 1 "Element is non-torsion unit\n"
      @v_do :UnitGroup 2 popindent()

      @v_do :UnitGroup 2 pushindent()

      if u.full_rank
        time_add_dep_unit += @elapsed m = _add_dependent_unit(u, y)::Bool
        if m
          @vprint :UnitGroup 1 "improved reg, reg is $(regulator(u.units, 16))\n"
        end
      else
        old_len = length(u.units)
        time_add_dep_unit += @elapsed _add_unit(u, y)
        m = old_len != length(u.units)
        if m
          u.units = reduce(u.units)
        end
        u.full_rank = length(u.units) == r
        if u.full_rank
          @vprint :UnitGroup 1 "reached full rank, reg is $(regulator(u.units, 16))\n"
        end
      end
      @v_do :UnitGroup 2 popindent()

      if !m
        not_larger = not_larger + 1
        if u.full_rank && not_larger > 5
          break
        end
      else
        not_larger = 0
      end
    end
  end

  #final reduction ...
  u.units = reduce(u.units, u.tors_prec)

  if u.full_rank
    u.tentative_regulator = regulator(u.units, 64)
  end

  @vprint :UnitGroup 1 "Finished processing\n"
  if u.full_rank
    @vprint :UnitGroup 1 "Regulator of current unit group is $(u.tentative_regulator)\n"
  else
    @vprint :UnitGroup 1 "current rank is $(length(u.units)) need $r\n"
  end
  @vprint :UnitGroup 1 "-"^80 * "\n"
  @vprint :UnitGroup 1 "Independent unit time: $time_indep\n"
  @vprint :UnitGroup 1 "Adding dependent unit time: $time_add_dep_unit\n"
  @vprint :UnitGroup 1 "Torsion test time: $time_torsion\n"
  @vprint :UnitGroup 1 "Kernel time: $time_kernel\n"

  @vtime_add :UnitGroup 1 x :unit_hnf_time time_kernel
  if u.full_rank
    return 1
  else
    return 0
  end
end


text:
"""
    uri(mlf::MLFlow, endpoint="", query=missing)

Retrieves an URI based on `mlf`, `endpoint`, and, optionally, `query`.

# Examples
```@example
MLFlowClient.uri(mlf, "experiments/get", Dict(:experiment_id=>10))
```
"""
function uri(mlf::MLFlow, endpoint="", query=missing)
    u = URI("$(mlf.baseuri)/api/$(mlf.apiversion)/mlflow/$(endpoint)")
    !ismissing(query) && return URI(u; query=query)
    u
end

"""
    mlfget(mlf, endpoint; kwargs...)

Performs a HTTP GET to a specifid endpoint. kwargs are turned into GET params.
"""
function mlfget(mlf, endpoint; kwargs...)
    apiuri = uri(mlf, endpoint, kwargs)
    headers = ["Content-Type: application/json"]
    try
        response = HTTP.get(apiuri, headers)
        return JSON.parse(String(response.body))
    catch e
        throw(e)
    end
end

"""
    mlfpost(mlf, endpoint; kwargs...)

Performs a HTTP POST to the specified endpoint. kwargs are converted to JSON and become the POST body.
"""
function mlfpost(mlf, endpoint; kwargs...)
    apiuri = uri(mlf, endpoint)
    headers = ["Content-Type: application/json"]
    body = JSON.json(kwargs)
    try
        response = HTTP.post(apiuri, headers, body)
        return JSON.parse(String(response.body))
    catch e
        throw(e)
    end
end

"""
    generatefilterfromparams(filter_params::AbstractDict{K,V}) where {K,V}

Generates a `filter` string from `filter_params` dictionary.

# Arguments
- `filter_params`: dictionary to use for filter generation.

# Returns
A string that can be passed as `filter` to [`searchruns`](@ref).

# Examples

```@example
generatefilterfromparams(Dict("paramkey1" => "paramvalue1", "paramkey2" => "paramvalue2"))
```
"""
function generatefilterfromparams(filter_params::AbstractDict{K,V}) where {K,V}
    length(filter_params) > 0 || return ""
    # NOTE: may have issues with escaping.
    filters = ["param.\"$(k)\" = \"$(v)\"" for(k, v) ∈ filter_params]
    join(filters, " and ")
end

max_stars_repo_name: JuliaAI/MLFlowClient

max_issues_repo_name: JuliaAI/MLFlowClient

max_forks_repo_name: JuliaAI/MLFlowClient


text:
# Sema
function restore_nns_annotation(x::Sema, v::AnnotationValue, rng::SourceRange,
                                spec::CXXScopeSpec)
    return RestoreNestedNameSpecifierAnnotation(x, v, rng, spec)
end

function restore_nns_annotation(x::Sema, tok::Token, spec::CXXScopeSpec)
    val = getAnnotationValue(tok)
    rng = getAnnotationRange(tok)
    return restore_nns_annotation(x, val, rng, spec)
end

Base.isempty(x::LookupResult) = empty(x)


text:
include("./produce_data.jl")
using NPZ
using HDF5

function dist_logbin(sizes, start_index, increment_log_bin)
    dist_size = normalize!(float(fit(Histogram, sizes[start_index:end], 1:maximum(sizes))))
    return logbin(dist_size.edges[1][1:end-1], dist_size.weights, increment_factor=increment_log_bin)
end

function main(;
        check_equilibration::Bool = true,
        minimum_number_avalanches::Int = Int(1e4),
        rolling_window::Int = Int(1e3),
        increment_log_bin::Float64 = sqrt(1.1)
    )
    # create data
    subsample, data = produce_data_avalanches_sandpile()

    if check_equilibration
        println("# determine initial equilibration period (will be discarded) from mean size of avalanches")
        mean_size = rollmean(data["size"]["full"], rolling_window)
        start_index = floor(Int, 1.5 * findfirst(mean_size .> mean(mean_size[floor(Int, length(mean_size)/2):end])))
    else
        start_index = 1
    end

    if length(data["size"]["full"]) - start_index > minimum_number_avalanches
        println("# enough data ... create output file")
        filename_out = @sprintf("%s/data/data_avalanche_sandpile.h5", path_project)
        fid = h5open(filename_out, "w")
        close(fid)
        println("# full sample")
        sizes = data["size"]["full"]
        x,P = dist_logbin(sizes, start_index, increment_log_bin)

        datasetname = "size/full"
        h5write(filename_out, datasetname, hcat(x,P))
        h5writeattr(filename_out, datasetname, Dict("description"=>"avalanche size distribution of full system (x,P)"))

        durations = data["duration"]["full"]
        x,P = dist_logbin(durations, start_index, increment_log_bin)
        datasetname = "duration/full"
        h5write(filename_out, datasetname, hcat(x,P))
        h5writeattr(filename_out, datasetname, Dict("description"=>"avalanche duration distribution of full system (x,P)"))

        # subsamples
        println("# subsamples")
        for (i,l) in enumerate(subsample["l"])
            x,P = dist_logbin(data["size"]["random"][i], start_index, increment_log_bin)
            datasetname = @sprintf("size/random/%d",l)
            h5write(filename_out, datasetname, hcat(x,P))
            h5writeattr(filename_out, datasetname, Dict("description"=>@sprintf("avalanche size distribution of randomly subsampled (n=%dx%d) systems (x,P)", l,l)))

            x,P = dist_logbin(data["size"]["window"][i], start_index, increment_log_bin)
            datasetname = @sprintf("size/window/%d",l)
            h5write(filename_out, datasetname, hcat(x,P))
            h5writeattr(filename_out, datasetname, Dict("description"=>@sprintf("avalanche size distribution of finite-field-of-view subsampled (%dx%d in center) systems (x,P)", l,l)))
        end
    end
end

# call function
main()