text:
#= """
    init_model()

Initialize the model, loading the data from external `npy` files, specified in the params file.
All prior data as been included previously, and is globaly accessed by the function.

Returns an `dp_parallel_sampling` (e.g. the main data structure) with the configured parameters and data.
"""
function init_model()
    if random_seed != nothing
        @eval @everywhere seed!($random_seed)
    end
    if use_verbose
        println("Loading and distributing data:")
        @time data = distribute(Float32.(load_data(data_path, prefix = data_prefix)))
    else
        data = distribute(Float32.(load_data(data_path, prefix = data_prefix)))
    end
    total_dim = size(data,2)
    model_hyperparams = model_hyper_params(hyper_params,α,total_dim)

    labels = distribute(rand(1:initial_clusters,(size(data,2))) .+ ((outlier_mod > 0) ? 1 : 0))
    labels_subcluster = distribute(rand(1:2,(size(data,2))))
    group = local_group(model_hyperparams,data,labels,labels_subcluster,local_cluster[],Float32[])
    return dp_parallel_sampling(model_hyperparams,group)
end =#

"""
    init_model(all_data)

Initialize the model, from `all_data`, should be `Dimensions X Samples`, type `Float32`
All prior data as been included previously, and is globaly accessed by the function.

Returns an `dp_parallel_sampling` (e.g. the main data structure) with the configured parameters and data.
"""
function init_model_from_data(all_data)
    if random_seed != nothing
        @eval @everywhere Random.seed!($random_seed)
    end
    if use_verbose
        println("Loading and distributing data:")
        @time data = distribute(all_data)
    else
        data = distribute(all_data)
    end

    total_dim = size(data,2)
    model_hyperparams = model_hyper_params(hyper_params,α,total_dim)
    labels = distribute(rand(1:initial_clusters,(size(data,2))) .+ ((outlier_mod > 0) ? 1 : 0))
    labels_subcluster = distribute(rand(1:2,(size(data,2))))
    group = local_group(model_hyperparams,data,labels,labels_subcluster,local_cluster[],Float32[])
    return dp_parallel_sampling(model_hyperparams,group)
end

"""
    init_model(all_data, init_labels)

Initialize the model, from `all_data`, should be `Dimensions X Samples`, type `Float32`
All prior data as been included previously, and is globaly accessed by the function.

Returns an `dp_parallel_sampling` (e.g. the main data structure) with the configured parameters and data.
"""
function init_model_from_data(all_data, init_labels, unique_init_labels)
    if random_seed != nothing
        @eval @everywhere Random.seed!($random_seed)
    end
    if use_verbose
        println("Loading and distributing data:")
        @time data = distribute(all_data)
    else
        data = distribute(all_data)
    end

    total_dim = size(data,2)
    model_hyperparams = model_hyper_params(hyper_params,α,total_dim)

    init_labels_unmapped = Vector{Int64}()
    for i=1:length(init_labels)
        ind_unmapped = findall(x->x==init_labels[i], unique_init_labels)[1]
        push!(init_labels_unmapped, ind_unmapped)
    end

    labels = distribute(init_labels_unmapped)

    labels_subcluster = distribute(rand(1:2,(size(data,2))))
    group = local_group(model_hyperparams, data, labels, labels_subcluster, local_cluster[], Float32[], Int64[], 0)
    return dp_parallel_sampling(model_hyperparams,group)
end

"""
    init_first_clusters!(dp_model::dp_parallel_sampling, initial_cluster_count::Int64))

Initialize the first clusters in the model, according to the number defined by initial_cluster_count

Mutates the model.
"""
function init_first_clusters!(dp_model::dp_parallel_sampling, initial_cluster_count::Int64, unique_init_labels::AbstractArray{Int64,1})
    #if outlier_mod > 0   # usually this does not hold
    #    push!(dp_model.group.local_clusters, create_outlier_local_cluster(dp_model.group,outlier_hyper_params))
    #end

    for i=1:initial_cluster_count
        push!(dp_model.group.local_clusters, create_first_local_cluster(dp_model.group, i))
        push!(dp_model.group.labels_mapping, unique_init_labels[i])
    end

    dp_model.group.max_orig_lbl = Int64(maximum(unique_init_labels))

    @sync update_suff_stats_posterior!(dp_model.group)
    sample_clusters!(dp_model.group, false)
    broadcast_cluster_params([create_thin_cluster_params(x) for x in dp_model.group.local_clusters],[1.0])
end

"""
    dp_parallel(all_data::AbstractArray{Float32,2},
        local_hyper_params::distribution_hyper_params,
        α_param::Float32,
         iters::Int64 = 100,
         init_clusters::Int64 = 1,
         seed = nothing,
         verbose = true,
         save_model = false,
         burnout = 15,
         gt = nothing,
         max_clusters = Inf,
         outlier_weight = 0,
         outlier_params = nothing)

Run the model.
# Args and Kwargs
 - `all_data::AbstractArray{Float32,2}` a `DxN` array containing the data
 - `local_hyper_params::distribution_hyper_params` the prior hyperparams
 - `α_param::Float32` the concetration parameter
 - `iters::Int64` number of iterations to run the model
 - `init_clusters::Int64` number of initial clusters
 - `seed` define a random seed to be used in all workers, if used must be preceeded with `@everywhere using random`.
 - `verbose` will perform prints on every iteration.
 - `save_model` will save a checkpoint every 25 iterations.
 - `burnout` how long to wait after creating a cluster, and allowing it to split/merge
 - `gt` Ground truth, when supplied, will perform NMI and VI analysis on every iteration.
 - `max_clusters` limit the number of cluster
 - `outlier_weight` constant weight of an extra non-spliting component
 - `outlier_params` hyperparams for an extra non-spliting component

# Return values
dp_model, iter_count , nmi_score_history, liklihood_history, cluster_count_history
 - `dp_model` The DPMM model inferred
 - `iter_count` Timing for each iteration
 - `nmi_score_history` NMI score per iteration (if gt suppled)
 - `likelihood_history` Log likelihood per iteration.
 - `cluster_count_history` Cluster counts per iteration.
"""
function dp_parallel(all_data::AbstractArray{Float32,2},
        local_hyper_params::distribution_hyper_params,
        α_param::Float32,
         iters::Int64 = 100,
         init_clusters::Int64 = 1,
         seed = nothing,
         verbose = true,
         save_model = false,
         burnout = 15,
         gt = nothing,
         max_clusters = Inf,
         outlier_weight = 0,
         outlier_params = nothing)
    global iterations = iters
    global random_seed = seed
    global hyper_params = local_hyper_params
    global α = α_param
    global use_verbose = verbose
    global should_save_model = save_model
    global burnout_period = burnout
    global max_num_of_clusters = max_clusters
    global outlier_mod = outlier_weight
    global outlier_hyper_params = outlier_params

    init_labels = outlier_params  # this is a workaround to pass additional varibale to python wrapper
    unique_init_labels = sort(unique(init_labels))
        global initial_clusters = length(unique_init_labels)
    println("===")
    println("initial_clusters: " * string(initial_clusters))
    println("===")

    dp_model = init_model_from_data(all_data, init_labels, unique_init_labels)
    global leader_dict = get_node_leaders_dict()
    init_first_clusters!(dp_model, initial_clusters, unique_init_labels)
    if use_verbose
        println("Node Leaders:")
        println(leader_dict)
    end
    global ground_truth = gt
    @eval @everywhere global hard_clustering = $hard_clustering
    return run_model(dp_model, 1)
end

"""
    fit(all_data::AbstractArray{Float32,2},local_hyper_params::distribution_hyper_params,α_param::Float32;
       iters::Int64 = 100, init_clusters::Int64 = 1,seed = nothing, verbose = true, save_model = false, burnout = 20, gt = nothing, max_clusters = Inf, outlier_weight = 0, outlier_params = nothing)

Run the model (basic mode).
# Args and Kwargs
 - `all_data::AbstractArray{Float32,2}` a `DxN` array containing the data
 - `local_hyper_params::distribution_hyper_params` the prior hyperparams
 - `α_param::Float32` the concetration parameter
 - `iters::Int64` number of iterations to run the model
 - `init_clusters::Int64` number of initial clusters
 - `seed` define a random seed to be used in all workers, if used must be preceeded with `@everywhere using random`.
 - `verbose` will perform prints on every iteration.
 - `save_model` will save a checkpoint every 25 iterations.
 - `burnout` how long to wait after creating a cluster, and allowing it to split/merge
 - `gt` Ground truth, when supplied, will perform NMI and VI analysis on every iteration.
 - `max_clusters` limit the number of cluster
 - `outlier_weight` constant weight of an extra non-spliting component
 - `outlier_params` hyperparams for an extra non-spliting component

# Return Values
 - `labels` Labels assignments
 - `clusters` Cluster parameters
 - `weights` The cluster weights, does not sum to `1`, but to `1` minus the weight of all uninstanistaed clusters.
 - `iter_count` Timing for each iteration
 - `nmi_score_history` NMI score per iteration (if gt suppled)
 - `likelihood_history` Log likelihood per iteration.
 - `cluster_count_history` Cluster counts per iteration.
 - `sub_labels` Sub labels assignments

# Example:
```julia
julia> x,y,clusters = generate_gaussian_data(10000,2,6,100.0)
...

julia> hyper_params = DPMMSubClusters.niw_hyperparams(1.0,
                  zeros(2),
                  5,
                  [1 0;0 1])
DPMMSubClusters.niw_hyperparams(1.0f0, Float32[0.0, 0.0], 5.0f0, Float32[1.0 0.0; 0.0 1.0])

julia> ret_values= fit(x,hyper_params,10.0, iters = 100, verbose=false)

...

julia> unique(ret_values[1])
6-element Array{Int64,1}:
 3
 6
 1
 2
 5
 4
```
"""
#function fit(all_data::AbstractArray{Float32,2},local_hyper_params::distribution_hyper_params,α_param::Float32, init_labels::AbstractArray{Int64,1};
#        iters::Int64 = 100, init_clusters::Int64 = 1, seed = nothing, verbose = true, save_model = false, burnout = 20, gt = nothing, max_clusters = Inf, outlier_weight = 0, outlier_params = nothing)
#        dp_model, iter_count , nmi_score_history, liklihood_history, cluster_count_history = dp_parallel(all_data, local_hyper_params,α_param, init_labels, iters,init_clusters, seed,verbose, save_model,burnout,gt, max_clusters, outlier_weight, outlier_params)
#        println("irit's run")
#        return Array(dp_model.group.labels), [x.cluster_params.cluster_params.distribution for x in dp_model.group.local_clusters], dp_model.group.weights,iter_count , nmi_score_history, liklihood_history, cluster_count_history,Array(dp_model.group.labels_subcluster)
#end

#  ---------------------- THIS IS THE FIT FUNCTION THAT IS CALLED WHEN RUNNING THE DP CODE: -----------------------------------------------------------------
#fit(all_data::AbstractArray, local_hyper_params::distribution_hyper_params, α_param, init_labels::AbstractArray;
#        iters = 100, init_clusters::Number = 1,
#        seed = nothing, verbose = true,
#        save_model = false,burnout = 20, gt = nothing, max_clusters = Inf, outlier_weight = 0, outlier_params = nothing) =
#    fit(Float32.(all_data), local_hyper_params, Float32(α_param), Int64.(init_labels), iters = Int64(iters),
#        init_clusters=Int64(init_clusters), seed = seed, verbose = verbose,
#        save_model = save_model, burnout = burnout, gt = gt, max_clusters = max_clusters, outlier_weight = outlier_weight, outlier_params = outlier_params)
#  ----------------------------------------------------------------------------------------------------------------------------------------------------------

function fit(all_data::AbstractArray{Float32,2},local_hyper_params::distribution_hyper_params,α_param::Float32;
    iters::Int64 = 100, init_clusters::Int64 = 1, seed = nothing, verbose = true, save_model = false, burnout = 20, gt = nothing, max_clusters = Inf, outlier_weight = 0, outlier_params::AbstractArray{Int64,1} = nothing)

    psi = local_hyper_params.ψ
    println("Is psi positive definite?: ", minimum(eigvals(psi))>0)
    #println("typeof(psi): ", typeof(psi))

    dp_model, iter_count , nmi_score_history, liklihood_history, cluster_count_history = dp_parallel(all_data, local_hyper_params,α_param, iters,init_clusters, seed,verbose, save_model,burnout,gt, max_clusters, outlier_weight, outlier_params)
    println("irit's run")

    mappings = dp_model.group.labels_mapping
    original_labels = [mappings[x] for x in dp_model.group.labels]
    # old: labels were returned as: Array(dp_model.group.labels)

    print("Labels_mapping: ")
    println(dp_model.group.labels_mapping)

    return original_labels, [x.cluster_params.cluster_params.distribution for x in dp_model.group.local_clusters], dp_model.group.weights, iter_count , nmi_score_history, liklihood_history, cluster_count_history, mappings, Array(dp_model.group.labels_subcluster)
end

#  ---------------------- THIS IS THE FIT FUNCTION THAT IS CALLED WHEN RUNNING THE DP CODE: -----------------------------------------------------------------
fit(all_data::AbstractArray,local_hyper_params::distribution_hyper_params,α_param;
        iters = 100, init_clusters::Number = 1,
        seed = nothing, verbose = true,
        save_model = false,burnout = 20, gt = nothing, max_clusters = Inf, outlier_weight = 0, outlier_params::AbstractArray = nothing) =
    fit(Float32.(all_data),local_hyper_params,Float32(α_param),iters = Int64(iters),
        init_clusters=Int64(init_clusters), seed = seed, verbose = verbose,
        save_model = save_model, burnout = burnout, gt = gt, max_clusters = max_clusters, outlier_weight = outlier_weight, outlier_params = Int64.(outlier_params))
#  ----------------------------------------------------------------------------------------------------------------------------------------------------------

#= """
    fit(all_data::AbstractArray{Float32,2},α_param::Float32;
        iters::Int64 = 100, init_clusters::Int64 = 1,seed = nothing, verbose = true, save_model = false,burnout = 20, gt = nothing, max_clusters = Inf, outlier_weight = 0, outlier_params = nothing)

Run the model (basic mode) with default `NIW` prior.
# Args and Kwargs
 - `all_data::AbstractArray{Float32,2}` a `DxN` array containing the data
 - `α_param::Float32` the concetration parameter
 - `iters::Int64` number of iterations to run the model
 - `init_clusters::Int64` number of initial clusters
 - `seed` define a random seed to be used in all workers, if used must be preceeded with `@everywhere using random`.
 - `verbose` will perform prints on every iteration.
 - `save_model` will save a checkpoint every 25 iterations.
 - `burnout` how long to wait after creating a cluster, and allowing it to split/merge
 - `gt` Ground truth, when supplied, will perform NMI and VI analysis on every iteration.
 - `outlier_weight` constant weight of an extra non-spliting component
 - `outlier_params` hyperparams for an extra non-spliting component

# Return Values
 - `labels` Labels assignments
 - `clusters` Cluster parameters
 - `weights` The cluster weights, does not sum to `1`, but to `1` minus the weight of all uninstanistaed clusters.
 - `iter_count` Timing for each iteration
 - `nmi_score_history` NMI score per iteration (if gt suppled)
 - `likelihood_history` Log likelihood per iteration.
 - `cluster_count_history` Cluster counts per iteration.
 - `sub_labels` Sub labels assignments

# Example:
```julia
julia> x,y,clusters = generate_gaussian_data(10000,2,6,100.0)
...

julia> ret_values= fit(x,10.0, iters = 100, verbose=false)

...

julia> unique(ret_values[1])
6-element Array{Int64,1}:
 3
 6
 1
 2
 5
 4
```
"""
function fit(all_data::AbstractArray{Float32,2},α_param::Float32;
        iters::Int64 = 100, init_clusters::Int64 = 1,seed = nothing, verbose = true, save_model = false,burnout = 20, gt = nothing, max_clusters = Inf, outlier_weight = 0, outlier_params = nothing)
    data_dim = size(all_data,1)
    cov_mat = Matrix{Float32}(I, data_dim, data_dim)
    local_hyper_params = niw_hyperparams(1,zeros(Float32,data_dim),data_dim+3,cov_mat)
    dp_model,iter_count , nmi_score_history, liklihood_history, cluster_count_history = dp_parallel(all_data, local_hyper_params,α_param,iters,init_clusters, seed,verbose,save_model,burnout,gt, max_clusters,outlier_weight, outlier_params)
    return Array(dp_model.group.labels), [x.cluster_params.cluster_params.distribution for x in dp_model.group.local_clusters], dp_model.group.weights,iter_count , nmi_score_history, liklihood_history, cluster_count_history, Array(dp_model.group.labels_subcluster)
end

fit(all_data::AbstractArray, α_param;
        iters = 100, init_clusters = 1,
        seed = nothing, verbose = true,
        save_model = false,burnout = 20, gt = nothing, max_clusters = Inf, outlier_weight = 0, outlier_params = nothing) =
    fit(Float32.(all_data),Float32(α_param),iters = Int64(iters),
        init_clusters=Int64(init_clusters), seed = seed, verbose = verbose,
        save_model = save_model, burnout = burnout, gt = gt, max_clusters = max_clusters, outlier_weight = outlier_weight, outlier_params = outlier_params)

"""
    dp_parallel(model_params::String; verbose = true, save_model = true,burnout = 5, gt = nothing)

Run the model in advanced mode.
# Args and Kwargs
 - `model_params::String` A path to a parameters file (see below)
 - `verbose` will perform prints on every iteration.
 - `save_model` will save a checkpoint every `X` iterations, where `X` is specified in the parameter file.
 - `burnout` how long to wait after creating a cluster, and allowing it to split/merge
 - `gt` Ground truth, when supplied, will perform NMI and VI analysis on every iteration.

# Return values
dp_model, iter_count , nmi_score_history, liklihood_history, cluster_count_history
 - `dp_model` The DPMM model inferred
 - `iter_count` Timing for each iteration
 - `nmi_score_history` NMI score per iteration (if gt suppled)
 - `likelihood_history` Log likelihood per iteration.
 - `cluster_count_history` Cluster counts per iteration.
"""
function dp_parallel(model_params::String; verbose = true, gt = nothing)
    include(model_params)
    global use_verbose = verbose
    dp_model = init_model()
    global leader_dict = get_node_leaders_dict()
    global should_save_model = enable_saving
    global ground_truth = gt
    global burnout_period = burnout_period
    global max_num_of_clusters = max_clusters
    init_first_clusters!(dp_model, initial_clusters)
    if use_verbose
        println("Node Leaders :")
        println(leader_dict)
    end
    @eval @everywhere global hard_clustering = $hard_clustering
    return run_model(dp_model, 1 ,model_params)
end =#

function run_model(dp_model, first_iter, model_params="none", prev_time = 0)
    start_time= time()
    iter_count = []
    liklihood_history = []
    v_score_history = []
    nmi_score_history = []
    global ground_truth
    cur_parr_count = 10
    cluster_count_history = []

    @sync for i in (nworkers()== 0 ? procs() : workers())
        @spawnat i set_parr_worker(dp_model.group.labels,cur_parr_count)
    end

    for i=first_iter:iterations
        # plot_group(dp_model.group)

        final = false
        no_more_splits = false
        if i >= iterations - argmax_sample_stop #We assume the cluters k has been setteled by now, and a low probability random split can do dmg
            final = true
        end
        if i >= iterations - split_stop || length(dp_model.group.local_clusters) >= max_num_of_clusters
            no_more_splits = true
        end

        prev_time = time()

        # -------------- IRIT PRINTS ------------------
        # if mod(i, 50) == 0
        #     println("--")
        #     # Final label counts and sub-clusters:
        #     vec = dp_model.group.local_clusters
        #     labels_res = dp_model.group.labels
        #     c = maximum(labels_res)
        #     label_counts = zeros(Int(c))
        #     for w=1:length(labels_res)
        #         l = Int(labels_res[w])
        #         label_counts[l] = label_counts[l] + 1
        #     end
        #     #for w=1:length(label_counts)
        #     #    orig_lbl = dp_model.group.labels_mapping[w]
        #     #    println("dp label: " * string(w) * "; final label: " * string(orig_lbl) * "; count: " * string(label_counts[w]))   # * ", sub-cluster cnts:" * ", " * string(vec[w].cluster_params.cluster_params_l.suff_statistics.N) * ", " * string(vec[w].cluster_params.cluster_params_r.suff_statistics.N)
        #     #end
        #     #println("--")
        #     #print("Labels_mapping: ")
        #     #println(dp_model.group.labels_mapping)
        #     #println("--")
        # end
        # -------------- IRIT PRINTS - DONE ------------

        group_step(dp_model.group, no_more_splits, final, i==1)
        iter_time = time() - prev_time
        push!(iter_count, iter_time)

        push!(cluster_count_history,length(dp_model.group.local_clusters))

        if ground_truth != nothing
            group_labels = Array(dp_model.group.labels)
            push!(v_score_history, varinfo(Int.(ground_truth),group_labels))
            push!(nmi_score_history, mutualinfo(Int.(ground_truth),group_labels,normed=true))
        else
            push!(v_score_history, "no gt")
            push!(nmi_score_history, "no gt")
        end
        if use_verbose
            push!(liklihood_history,calculate_posterior(dp_model))
            println("Iteration : " * string(i) * " || Clusters count: " *
                string(cluster_count_history[end]) *
                " || Log posterior: " * string(liklihood_history[end]) *
                " || Vi score: " * string(v_score_history[end]) *
                " || NMI score: " * string(nmi_score_history[end]) *
                " || Iter Time:" * string(iter_time) *
                 " || Total time:" * string(sum(iter_count)))
        else
            push!(liklihood_history,1)
        end
        # if length(dp_model.group.local_clusters) > cur_parr_count
        #     cur_parr_count += max(20,length(dp_model.group.local_clusters))
        #     @sync for i in (nworkers()== 0 ? procs() : workers())
        #         @spawnat i set_parr_worker(dp_model.group.labels,cur_parr_count)
        #     end
        # end
        if i % model_save_interval == 0 && should_save_model
            println("Saving Model:")
            # save_time = time()
            @time save_model(dp_model,save_path, save_file_prefix, i, time() - start_time, model_params)
            # start_time += time() - save_time
        end
    end

    # Update the weights after all changes:
    sample_clusters!(dp_model.group, false)

    return dp_model, iter_count , nmi_score_history, liklihood_history, cluster_count_history
end

"""
    run_model_from_checkpoint(filename)

Run the model from a checkpoint created by it, `filename` is the path to the checkpoint.
Only to be run when using the advanced mode, note that the data must be in the same path as previously.

# Example:
```julia
julia> dp = run_model_from_checkpoint("checkpoint__50.jld2")
Loading Model:
  1.073261 seconds (2.27 M allocations: 113.221 MiB, 2.60% gc time)
Including params
Loading data:
  0.000881 seconds (10.02 k allocations: 378.313 KiB)
Creating model:
Node Leaders:
Dict{Any,Any}(2=>Any[2, 3])
Running model:
...
```
"""
function run_model_from_checkpoint(filename)
    println("Loading Model:")
    @time @load filename group hyperparams iter total_time global_params
    println("Including params")
    include(global_params)
    println("Loading data:")
    @time data = distribute((Float32.(load_data(data_path, prefix = data_prefix))))
    println("Creating model:")
    # @time begin
    group.labels = distribute(group.labels)
    group.labels_subcluster = distribute(group.labels_subcluster)
    dp_model = create_model_from_saved_data(group, data, hyperparams)
    # end
    global leader_dict = get_node_leaders_dict()
    println("Node Leaders:")
    println(leader_dict)
    @eval @everywhere global hard_clustering = $hard_clustering
    println("Running model:")
    return run_model(dp_model, iter+1, global_params, total_time)
end

function save_model(model, path, filename, iter, total_time, global_params)
    filename = path * filename * "_"*string(iter)*".jld2"
    hyperparams = model.model_hyperparams
    group = create_pts_less_group(model.group)
    @save filename group hyperparams iter total_time global_params
end

function calculate_posterior(model::dp_parallel_sampling)
    log_posterior = logabsgamma(model.model_hyperparams.α)[1] - logabsgamma(size(model.group.points,2)+model.model_hyperparams.α)[1]
    for cluster in model.group.local_clusters
        if cluster.cluster_params.cluster_params.suff_statistics.N == 0
            continue
        end
        log_posterior += log_marginal_likelihood(cluster.cluster_params.cluster_params.hyperparams,
            cluster.cluster_params.cluster_params.posterior_hyperparams,
            cluster.cluster_params.cluster_params.suff_statistics)
        log_posterior += log(model.model_hyperparams.α) + logabsgamma(cluster.cluster_params.cluster_params.suff_statistics.N)[1]
    end
    return log_posterior
end

function set_parr_worker(labels,cluster_count)
    global glob_parr = zeros(Float32,size(localpart(labels),1),cluster_count)
end

"""
    cluster_statistics(points,labels, clusters)

Provide avg statsitcs of probabiliy and likelihood for given points, labels and clusters

# Args and Kwargs
 - `points` a `DxN` array containing the data
 - `labels` points labels
 - `clusters` vector of clusters distributions

# Return values
avg_ll, avg_prob
 - `avg_ll` each cluster avg point ll
 - `avg_prob` each cluster avg point prob

# Example:
```julia
julia> dp = run_model_from_checkpoint("checkpoint__50.jld2")
Loading Model:
  1.073261 seconds (2.27 M allocations: 113.221 MiB, 2.60% gc time)
Including params
Loading data:
  0.000881 seconds (10.02 k allocations: 378.313 KiB)
Creating model:
Node Leaders:
Dict{Any,Any}(2=>Any[2, 3])
Running model:
...
```
"""
function cluster_statistics(points,labels, clusters)
    parr = zeros(Float32,length(labels), length(clusters))
    tic = time()
    for (k,cluster) in enumerate(clusters)
        log_likelihood!(reshape((@view parr[:,k]),:,1), points,cluster)
    end
    log_likelihood_array = copy(parr)
    log_likelihood_array[isnan.(log_likelihood_array)] .= -Inf #Numerical errors arent fun
    max_log_prob_arr = maximum(log_likelihood_array, dims = 2)
    log_likelihood_array .-= max_log_prob_arr
    map!(exp,log_likelihood_array,log_likelihood_array)
    # println("lsample log cat2" * string(log_likelihood_array))
    sum_prob_arr = sum(log_likelihood_array, dims =[2])
    log_likelihood_array ./=  sum_prob_arr
    avg_ll = zeros(length(clusters))
    avg_prob = zeros(length(clusters))
    for i=1:length(clusters)
        avg_ll[i] = sum(parr[labels .== i,i]) / sum(labels .== i)
        avg_prob[i] = sum(log_likelihood_array[labels .== i,i]) / sum(labels .== i)
    end
    return avg_ll, avg_prob
end


text:
import Test.@testset
import Test.@test

using LightGraphs

push!(LOAD_PATH, pwd())
using Coils

const tol = 1e-10

@testset "Solving a system" begin
    g, vertex_positions = cuboid_system([1, 1, 1], [2, 2, 2])
    poi = [[0,0,0]]
    Bgoal(x) = [0, 0, 50]
    simpleloops, simpleloopscurrents = solve_system(g, vertex_positions, poi, Bgoal)

    field = field_loops(poi[1], simpleloops, simpleloopscurrents, vertex_positions)
    @test abs(field[3] - 50) < tol
end

if !(("TRAVIS" in keys(ENV)) & Sys.isapple())
    # do not execute this test on OSX in Travis - there is a problem with
    # matplotlib there
    @testset "Example code with plotting" begin
        include("example.jl")
    end
end

ext: jl


text:
# Plot data points:
using Plots
"""
    plot_data!(plt,X,y)

# Examples

```julia-repl
using BayesLaplace, Plots
X, y = toy_data_linear(100)
plt = plot()
plot_data!(plt, hcat(X...)', y)
```

"""
function plot_data!(plt,X,y)
    Plots.scatter!(plt, X[y.==1.0,1],X[y.==1.0,2], color=1, clim = (0,1), label="y=1")
    Plots.scatter!(plt, X[y.==0.0,1],X[y.==0.0,2], color=0, clim = (0,1), label="y=0")
end

# Plot contour of posterior predictive:
using Plots
"""
    plot_contour(X,y,𝑴;clegend=true,title="",length_out=50,type=:laplace,zoom=0,xlim=nothing,ylim=nothing)

Generates a contour plot for the posterior predictive surface.

# Examples

```julia-repl
using BayesLaplace, Plots
import BayesLaplace: predict
using NNlib: σ
X, y = toy_data_linear(100)
X = hcat(X...)'
β = [1,1]
𝑴 =(β=β,)
predict(𝑴, X) = σ.(𝑴.β' * X)
plot_contour(X, y, 𝑴)
```

"""
function plot_contour(X,y,𝑴;clegend=true,title="",length_out=50,type=:laplace,zoom=0,xlim=nothing,ylim=nothing)

    # Surface range:
    if isnothing(xlim)
        xlim = (minimum(X[:,1]),maximum(X[:,1])).+(zoom,-zoom)
    else
        xlim = xlim .+ (zoom,-zoom)
    end
    if isnothing(ylim)
        ylim = (minimum(X[:,2]),maximum(X[:,2])).+(zoom,-zoom)
    else
        ylim = ylim .+ (zoom,-zoom)
    end
    x_range = collect(range(xlim[1],stop=xlim[2],length=length_out))
    y_range = collect(range(ylim[1],stop=ylim[2],length=length_out))

    # Predictions:
    if type==:laplace
        Z = [predict(𝑴,[x, y])[1] for x=x_range, y=y_range]
    else
        Z = [plugin(𝑴,[x, y])[1] for x=x_range, y=y_range]
    end

    # Plot:
    plt = contourf(
        x_range, y_range, Z';
        legend=clegend, title=title, linewidth=0,
        xlim=xlim,
        ylim=ylim
    )
    plot_data!(plt,X,y)

end

# Helper function to predict from network trained for binary classification and producing logits as output:
import BayesLaplace: predict
predict(𝑴::Flux.Chain, X::AbstractArray) = Flux.σ.(𝑴(X))


text:
function getIndustryPortfolios()
    # download data
    dataUrl = "http://mba.tuck.dartmouth.edu/pages/faculty/ken.french/ftp/30_Industry_Portfolios_daily_TXT.zip"
    (data, dataNames, varnames) = readFamaFrenchRaw(dataUrl)

    # pick value weighted sub-part
    selectedData = data[1]

    # get correct column names
    TimeSeries.rename(selectedData, split(varnames))

    # fix encoding of missing values
    dataVals = selectedData.values
    dataVals[dataVals .== -99.99] = 0

    cleanData = TimeSeries.TimeArray(selectedData.timestamp, dataVals, split(varnames))

    fileName = joinpath(Pkg.dir("EconDatasets"), "data", "IndustryPfs.csv")
    writetimearray(cleanData, fileName)
end


text:
export RK4,RK4_naive
export step!

"""
    RK4(arguments;realdata)

Explicit Runge-Kutta fourth order solver.

Construct an object of type `TimeSolver` to be used in `Problem(model, initial, param; solver::TimeSolver)`

Arguments can be either
0. an object of type `AbstractModel`;
1. an `Array` of size `(N,datasize)` where `N` is the number of collocation points and `datasize` the number of equations solved;
2. `(param,datasize)` where `param is a `NamedTuple` containing a key `N`, and `datasize` a integer (optional, by default `datasize=2`).

The keyword argument `realdata` is optional, and determines whether pre-allocated vectors are real- or complex-valued.
By default, they are either determined by the model or the type of the array in case `0.` and `1.`, complex-valued in case `2.`.

"""
struct RK4 <: TimeSolver

    U1 :: Array
    dU :: Array
    label :: String

    function RK4( U :: Array; realdata=nothing )
        U1 = copy(U)
        dU = copy(U)
        if realdata==true
            U1 = real.(U1);dU = real.(dU)
        end
        if realdata==false
            U1 = complex.(U1);dU = complex.(dU)
        end
        new( U1, dU, "RK4")
    end

    function RK4( model :: AbstractModel; realdata=nothing )
        U=model.mapto(Init(x->0*x,x->0*x))
        RK4( U; realdata=realdata)
    end
    function RK4( param::NamedTuple, datasize=2::Int; realdata=nothing )
        RK4( zeros(Complex{Float64}, (param.N,datasize)) ; realdata=realdata)
    end
end

function step!(s  :: RK4,
                m :: AbstractModel,
                U  ,
                dt )

    s.U1 .= U
    m.f!( s.U1 )
    s.dU .= s.U1

    s.U1 .= U .+ dt/2 .* s.U1
    m.f!( s.U1 )
    s.dU .+= 2 .* s.U1

    s.U1 .= U .+ dt/2 .* s.U1
    m.f!( s.U1 )
    s.dU .+= 2 .* s.U1

    s.U1 .= U .+ dt .* s.U1
    m.f!( s.U1 )
    s.dU .+= s.U1

    U .+= dt/6 .* s.dU

end

"""
    RK4_naive()

Runge-Kutta fourth order solver.

A naive version of `RK4`, without argument since no pre-allocation is performed.

"""
struct RK4_naive <: TimeSolver

    label :: String

    function RK4_naive() new("RK4 (naive)") end
end

function step!(s  :: RK4_naive,
               m :: AbstractModel ,
               U  ,
               dt )

    U0 = copy(U)
    m.f!( U0 )
    U1 = copy(U0)

    U0 .= U .+ dt/2 .* U1
    m.f!( U0 )
    U2 = copy(U0)

    U0 .= U .+ dt/2 .* U2
    m.f!( U0 )
    U3 = copy(U0)

    U0 .= U .+ dt .* U3
    m.f!( U0 )
    U4 = copy(U0)

    U .+= dt/6 .* (U1 + 2*U2 + 2*U3 + U4 )

end


text:
module GeneticsMakie

using CairoMakie
using Makie.GeometryBasics
using DataFrames
using SnpArrays
using Statistics
using Distributions

include("parsegtf.jl")
include("plotgenes.jl")
include("plotisoforms.jl")
include("plotld.jl")
include("mungesumstats.jl")
include("plotlocus.jl")
include("plotqq.jl")
include("plotgwas.jl")
include("plotrg.jl")
include("gwas.jl")
include("genome.jl")

end


text:
using GalacticNOMAD, GalacticOptim
using Test

@testset "GalacticNOMAD.jl" begin
    rosenbrock(x, p) = (p[1] - x[1])^2 + p[2] * (x[2] - x[1]^2)^2
    x0 = zeros(2)
    _p = [1.0, 100.0]
    l1 = rosenbrock(x0, _p)

    f = OptimizationFunction(rosenbrock)

    prob = OptimizationProblem(f, x0, _p)
    sol = GalacticOptim.solve(prob, NOMADOpt())
    @test 10 * sol.minimum < l1

    prob = OptimizationProblem(f, x0, _p; lb=[-1.0, -1.0], ub=[1.5, 1.5])
    sol = GalacticOptim.solve(prob, NOMADOpt())
    @test 10 * sol.minimum < l1
end


text:
# Use baremodule to shave off a few KB from the serialized `.ji` file
baremodule Check_jll
using Base
using Base: UUID
import JLLWrappers

JLLWrappers.@generate_main_file_header("Check")
JLLWrappers.@generate_main_file("Check", UUID("491db154-c145-5abe-9c32-446728d60cce"))
end  # module Check_jll


text:
# The whitespace in this file is weird. Read in and out as csv to normalize whitespace

dataset_path = joinpath(dirname(@__FILE__),
                        "climate-model-simulation-crashes.data.orig")

data = readdlm(dataset_path)
writecsv("$dataset_path.custom", data[:, 1:end - 1])

max_stars_repo_name: CoventryResearch/uci-data

max_issues_repo_name: CoventryResearch/uci-data

max_forks_repo_name: CoventryResearch/uci-data


text:
using ITensors
using ITensorGLMakie

include(joinpath("..", "src", "visualize_dangling.jl"))

i = Index(2, "i")

Zp = ITensor(i)
Zm = ITensor(i)
Xp = ITensor(i)
Xm = ITensor(i)

fig_Zp = visualize_dangling(Zp; vertex_labels="Z+", tags=true)
fig_Zm = visualize_dangling(Zm; vertex_labels="Z-", tags=true)
fig_Xp = visualize_dangling(Xp; vertex_labels="X+", tags=true)
fig_Xm = visualize_dangling(Xm; vertex_labels="X-", tags=true)

fig_ZpXp = visualize_dangling([Zp, Xp]; vertex_labels=["Z+", "X+"], tags=true)

Z = ITensor(i', i)
fig_ZXp = visualize_dangling([Z, Xp]; vertex_labels=["Z+", "X+"], tags=true)

# using Makie
# save("figure.png", fig)


text:
# ------------------------------------------------------------------------------------------
# # Learning: algorithms, objectives, and assumptions
# (c) Deniz Yuret 2019
#
# In this notebook we will analyze three classic learning algorithms.
# * **Perceptron:** ([Rosenblatt, 1957](https://en.wikipedia.org/wiki/Perceptron)) a neuron
# model trained with a simple algorithm that updates model weights using the input when the
# prediction is wrong.
# * **Adaline:** ([Widrow and Hoff, 1960](https://en.wikipedia.org/wiki/ADALINE)) a neuron
# model trained with a simple algorithm that updates model weights using the error
# multiplied by the input (aka least mean square (LMS), delta learning rule, or the Widrow-
# Hoff rule).
# * **Softmax classification:** ([Cox,
# 1958](https://en.wikipedia.org/wiki/Multinomial_logistic_regression)) a multiclass
# generalization of the logistic regression model from statistics (aka multinomial LR,
# softmax regression, maxent classifier etc.).
#
# We will look at these learners from three different perspectives:
# * **Algorithm:** First we ask only **how** the learner works, i.e. how it changes after
# observing each example.
# * **Objectives:** Next we ask **what** objective guides the algorithm, whether it is
# optimizing a particular objective function, and whether we can use a generic *optimization
# algorithm* instead.
# * **Assumptions:** Finally we ask **why** we think this algorithm makes sense, what prior
# assumptions does this imply and whether we can use *probabilistic inference* for optimal
# learning.
# ------------------------------------------------------------------------------------------

using Knet, Plots, Statistics, LinearAlgebra, Random
Base.argmax(a::KnetArray) = argmax(Array(a))
Base.argmax(a::AutoGrad.Value) = argmax(value(a))
ENV["COLUMNS"] = 72

# ------------------------------------------------------------------------------------------
# ### Data
# ------------------------------------------------------------------------------------------

include(Knet.dir("data/mnist.jl"))
xtrn, ytrn, xtst, ytst = mnist()
ARRAY = Array{Float32}
xtrn, xtst = ARRAY(mat(xtrn)), ARRAY(mat(xtst))
onehot(y) = (m=ARRAY(zeros(maximum(y),length(y))); for i in 1:length(y); m[y[i],i]=1; end; m)
ytrn, ytst = onehot(ytrn), onehot(ytst);

println.(summary.((xtrn, ytrn, xtst, ytst)));

NTRN,NTST,XDIM,YDIM = size(xtrn,2), size(xtst,2), size(xtrn,1), size(ytrn,1)

# ------------------------------------------------------------------------------------------
# ### Prediction
# ------------------------------------------------------------------------------------------

# Model weights
w = ARRAY(randn(YDIM,XDIM))

# Class scores
w * xtrn

# Predictions
[ argmax(w * xtrn[:,i]) for i in 1:NTRN ]'

# Correct answers
[ argmax(ytrn[:,i]) for i in 1:NTRN ]'

# Accuracy
acc(w,x,y) = mean(argmax(w * x, dims=1) .== argmax(y, dims=1))
acc(w,xtrn,ytrn), acc(w,xtst,ytst)

# ------------------------------------------------------------------------------------------
# ## Algorithms
# ------------------------------------------------------------------------------------------

# Training loop
function train(algo,x,y,T=2^20)
    w = ARRAY(zeros(size(y,1),size(x,1)))
    nexamples = size(x,2)
    nextprint = 1
    for t = 1:T
        i = rand(1:nexamples)
        algo(w, x[:,i], y[:,i])  # <== this is where w is updated
        if t == nextprint
            println((iter=t, accuracy=acc(w,x,y), wnorm=norm(w)))
            nextprint = min(2t,T)
        end
    end
    w
end

# ------------------------------------------------------------------------------------------
# ### Perceptron
# ------------------------------------------------------------------------------------------

function perceptron(w,x,y)
    guess = argmax(w * x)
    class = argmax(y)
    if guess != class
        w[class,:] .+= x
        w[guess,:] .-= x
    end
end

# (iter = 1048576, accuracy = 0.8950333333333333, wnorm = 1321.2463f0) in 7 secs
@time wperceptron = train(perceptron,xtrn,ytrn);

# ------------------------------------------------------------------------------------------
# ### Adaline
# ------------------------------------------------------------------------------------------

function adaline(w,x,y; lr=0.0001)
    error = w * x - y
    w .-= lr * error * x'
end

# (iter = 1048576, accuracy = 0.8523, wnorm = 1.2907721f0) in 31 secs with lr=0.0001
@time wadaline = train(adaline,xtrn,ytrn);

# ------------------------------------------------------------------------------------------
# ### Softmax classifier
# ------------------------------------------------------------------------------------------

function softmax(w,x,y; lr=0.01)
    probs = exp.(w * x)
    probs ./= sum(probs)
    error = probs - y
    w .-= lr * error * x'
end

# (iter = 1048576, accuracy = 0.9242166666666667, wnorm = 26.523603f0) in 32 secs with lr=0.01
@time wsoftmax = train(softmax,xtrn,ytrn);

# ------------------------------------------------------------------------------------------
# ## Objectives
# ------------------------------------------------------------------------------------------

# Training via optimization
function optimize(loss,x,y; lr=0.1, iters=2^20)
    w = Param(ARRAY(zeros(size(y,1),size(x,1))))
    nexamples = size(x,2)
    nextprint = 1
    for t = 1:iters
        i = rand(1:nexamples)
        L = @diff loss(w, x[:,i], y[:,i])
        ∇w = grad(L,w)
        w .-= lr * ∇w
        if t == nextprint
            println((iter=t, accuracy=acc(w,x,y), wnorm=norm(w)))
            nextprint = min(2t,iters)
        end
    end
    w
end

# ------------------------------------------------------------------------------------------
# ### Perceptron minimizes the score difference between the correct class and the prediction
# ------------------------------------------------------------------------------------------

function perceptronloss(w,x,y)
    score = w * x
    guess = argmax(score)
    class = argmax(y)
    score[guess] - score[class]
end

# (iter = 1048576, accuracy = 0.8908833333333334, wnorm = 1322.4888f0) in 62 secs with lr=1
@time wperceptron2 = optimize(perceptronloss,xtrn,ytrn,lr=1);

# ------------------------------------------------------------------------------------------
# ### Adaline minimizes the squared error
# ------------------------------------------------------------------------------------------

function quadraticloss(w,x,y)
    0.5 * sum(abs2, w * x - y)
end

# (iter = 1048576, accuracy = 0.8498333333333333, wnorm = 1.2882874f0) in 79 secs with lr=0.0001
@time wadaline2 = optimize(quadraticloss,xtrn,ytrn,lr=0.0001);

# ------------------------------------------------------------------------------------------
# ### Softmax classifier maximizes the probabilities of correct answers
# (or minimizes negative log likelihood, aka cross-entropy or softmax loss)
# ------------------------------------------------------------------------------------------

function negloglik(w,x,y)
    probs = exp.(w * x)
    probs = probs / sum(probs)
    class = argmax(y)
    -log(probs[class])
end

# (iter = 1048576, accuracy = 0.9283833333333333, wnorm = 26.593485f0) in 120 secs with lr=0.01
@time wsoftmax2 = optimize(negloglik,xtrn,ytrn,lr=0.01);

max_stars_repo_name: grenkoca/JuliaTutorials

max_issues_repo_name: grenkoca/JuliaTutorials

max_forks_repo_name: grenkoca/JuliaTutorials


text:
#jl #! format: off
# # ACOPF with [PowerSimulations.jl](https://github.com/NREL-SIIP/PowerSimulations.jl) using [PowerModels.jl](https://github.com/lanl-ansi/PowerModels.jl)

# **Originally Contributed by**: Clayton Barrows

# ## Introduction

# PowerSimulations.jl supports non-linear AC optimal power flow through a deep integration
# with [PowerModels.jl](https://github.com/lanl-ansi/PowerModels.jl). This example shows a
# single multi-period optimization of economic dispatch with a full representation of
# AC optimal power flow.

# ## Dependencies
using SIIPExamples
using PowerSystems
using PowerSimulations
using PowerSystemCaseBuilder

# We can use the a [TAMU synthetic ERCOT dataset](https://electricgrids.engr.tamu.edu/electric-grid-test-cases/).
# The TAMU data format relies on a folder containing `.m` or `.raw` files and `.csv`
# files for the time series data. We have provided a parser for the TAMU data format with
# the `TamuSystem()` function. A version of the system with only one week of time series
# is included in PowerSystemCaseBuilder.jl, we can use that version here:
sys = build_system(PSYTestSystems, "tamu_ACTIVSg2000_sys")
transform_single_time_series!(sys, 2, Hour(1))

# Since we'll be doing non-linear optimization, we need a solver that supports non-linear
# problems. Ipopt is quite good.
using Ipopt
solver = optimizer_with_attributes(Ipopt.Optimizer)

# In the [OperationsProblem example](https://nbviewer.jupyter.org/github/NREL-SIIP/SIIPExamples.jl/blob/master/notebook/3_PowerSimulations_examples/01_operations_problems.ipynb)
# we defined a unit-commitment problem with a copper plate representation of the network.
# Here, we want do define an economic dispatch (linear generation decisions) with an ACOPF
# network representation.
# So, starting with the network, we can select from _almost_ any of the endpoints on this
# tree:
print_tree(PowerSimulations.PM.AbstractPowerModel)

# For now, let's just choose a standard ACOPF formulation.
ed_template = OperationsProblemTemplate(QCLSPowerModel)
set_device_model!(ed_template, ThermalStandard, ThermalDispatch)
set_device_model!(ed_template, PowerLoad, StaticPowerLoad)
#set_device_model!(ed_template, FixedAdmittance, StaticPowerLoad) #TODO add constructor for shunts in PSI

# Now we can build a 4-hour economic dispatch / ACOPF problem with the TAMU data.
problem = OperationsProblem(
    ed_template,
    sys,
    horizon = 1,
    optimizer = solver,
    balance_slack_variables = true,
)
build!(problem, output_dir = mktempdir())

# And solve it ... (it's infeasible)
solve!(problem)


text:
cd(@__DIR__)
using Pkg
Pkg.activate("")
using Revise, PAAV, Plots, DataFrames, Statistics, StatsPlots
Random.seed!(502241)
experiment = Experiment(n_trials=10^4, populate_visicon=conjunctive_set,
    set_size=30)
run_condition!(experiment; topdown_weight=.6)

pyplot()
df = DataFrame(experiment.data)
@df df histogram(:rt, group=(:target_present,:response),
    layout=(1,3), xlims=(1,4) ,grid=false, color=:grey, size=(600,300))

by(df, [:target_present,:response], :rt=>mean,:rt=>std)

experiment = Experiment(n_trials=10^4, populate_visicon=conjunctive_set,
    set_size=30)
run_condition!(experiment; topdown_weight=1.0, bottomup_weight=0.0, Δτ=.15, noise=.5)

df = DataFrame(experiment.data)
@df df histogram(:rt, group=(:target_present,:response),
    layout=(1,3), xlims=(1,4), grid=false, color=:grey, size=(600,300))

by(df, [:target_present,:response], :rt=>mean,:rt=>std, :rt=>length)


text:
# Autogenerated wrapper script for KAShim_jll for armv7l-linux-gnueabihf
export libkashim

using LibUV_jll
## Global variables
PATH = ""
LIBPATH = ""
LIBPATH_env = "LD_LIBRARY_PATH"

# Relative path to `libkashim`
const libkashim_splitpath = ["lib", "libkashim.so"]

# This will be filled out by __init__() for all products, as it must be done at runtime
libkashim_path = ""

# libkashim-specific global declaration
# This will be filled out by __init__()
libkashim_handle = C_NULL

# This must be `const` so that we can use it with `ccall()`
const libkashim = "libkashim.so"

"""
Open all libraries
"""
function __init__()
    global artifact_dir = abspath(artifact"KAShim")

    # Initialize PATH and LIBPATH environment variable listings
    global PATH_list, LIBPATH_list
    # We first need to add to LIBPATH_list the libraries provided by Julia
    append!(LIBPATH_list, [joinpath(Sys.BINDIR, Base.LIBDIR, "julia"), joinpath(Sys.BINDIR, Base.LIBDIR)])
    # From the list of our dependencies, generate a tuple of all the PATH and LIBPATH lists,
    # then append them to our own.
    foreach(p -> append!(PATH_list, p), (LibUV_jll.PATH_list,))
    foreach(p -> append!(LIBPATH_list, p), (LibUV_jll.LIBPATH_list,))

    global libkashim_path = normpath(joinpath(artifact_dir, libkashim_splitpath...))

    # Manually `dlopen()` this right now so that future invocations
    # of `ccall` with its `SONAME` will find this path immediately.
    global libkashim_handle = dlopen(libkashim_path)
    push!(LIBPATH_list, dirname(libkashim_path))

    # Filter out duplicate and empty entries in our PATH and LIBPATH entries
    filter!(!isempty, unique!(PATH_list))
    filter!(!isempty, unique!(LIBPATH_list))
    global PATH = join(PATH_list, ':')
    global LIBPATH = join(LIBPATH_list, ':')

    # Add each element of LIBPATH to our DL_LOAD_PATH (necessary on platforms
    # that don't honor our "already opened" trick)
    #for lp in LIBPATH_list
    #    push!(DL_LOAD_PATH, lp)
    #end
end  # __init__()


text:
using AxisKeys
using AxisSets
using Dates
using Documenter
using FeatureTransforms
using FeatureTransforms.TestUtils
using Impute
using Missings
using OrderedCollections
using ReadOnlyArrays
using Statistics
using Test
using TimeZones

using AxisSets:
    Pattern,
    KeyAlignmentError,
    KeyedDataset,
    constraintmap,
    dimpaths,
    flatten,
    validate
using Impute: ThresholdError

@testset "AxisSets.jl" begin
    include("flatten.jl")
    include("patterns.jl")
    include("dataset.jl")
    include("indexing.jl")
    include("functions.jl")
    include("impute.jl")
    include("featuretransforms.jl")

    # The doctests fail on x86, so only run them on 64-bit hardware & Julia 1.6
    Sys.WORD_SIZE == 64 && v"1.6" <= VERSION < v"1.7" && doctest(AxisSets)
end


text:
@info "Starting..."
using LinearAlgebra, StaticArrays, DifferentialEquations, Makie, ODEInterfaceDiffEq
using BSON: @save, @load
include("system.jl")
include("util.jl")
include("plotting.jl")

@info "Defining ODE system..."
N = M = 3;
p = (k_m=0.2, l₀ = 1.0)
t = (0.0, 60.0)

u₀ = uniform_grid(M, N, p[:l₀])
for i in 1:N
    u₀[3,i,1] += 0.3*p[:l₀]
    u₀[3,i,M] -= 0.3*p[:l₀]
end

@info "Displaying initial positions."
display(plot_lindo(u₀))

@info "Integrating ODE system..."
prob = ODEProblem(rigid!, u₀, t, p)
@time sol = solve(prob, progress=true, force_dtmin=true, DifferentialEquations.TRBDF2())

@save "$(pwd())/data/$(N)x$(M).bson" sol
@info "Saved solution to $(pwd())/data/$(N)x$(M).bson"

@info "Rendering..."
render(sol)
@info "Done."