# high level functions used in the project
using PyPlot
using DelimitedFiles

function run_fit_procedure(description::String, data::Array{Int,2}, algo_fit_spec,
                            nr_repetitions::Int, mcmc_data_file::String,
                            top_n_fraction::Float64, rl_module_path::String, opts...)
    start_time = now()
    println("start: $start_time")
    println(description)
    data_string = RL_Fit.to_short_string(RL_Fit.get_data_stats(data))
    println(data_string)
        flush(stdout)
    println("run_fit_procedure")

    mcmc_samples = RL_Fit.run_parallel_Metropolis(data, algo_fit_spec, rl_module_path, opts...; nr_of_repetitions=nr_repetitions)
    nr_map_samples = Int(ceil(size(mcmc_samples)[1] * top_n_fraction))
    println("sampling data size: $(size(mcmc_samples)), top_N: $nr_map_samples")
    top_MAP, top_n_avgMAP, top_n_min, top_n_max, top_n_stdMAP, top_n_samples = RL_Fit.get_MAP(mcmc_samples, nr_map_samples)
    println("result timestamp: $(now())")
    println("top_MAP: $(round.(top_MAP,digits=4)) (Max Likelihood estimate (top 1 sample). LL includes Prior LL)")
    println("top_n_avgMAP: $(round.(top_n_avgMAP,digits=4)) . (Avg over top $nr_map_samples samples. LL includes Prior LL)")
    println("top_n_min: $(round.(top_n_min,digits=4))")
    println("top_n_max: $(round.(top_n_max,digits=4))")
    # evaluate using top_n_avgMAP parameter. get the LL without the prior (penalty)
    LL_noPrior_avgMAP, all_signals_avg = algo_fit_spec.rl_impl(data, top_n_avgMAP[1:end-1], opts...)
    println("LL top_n_avgMAP=$(round(LL_noPrior_avgMAP, digits=4))")
    # evaluate using top_MAP parameter. get the LL without the prior (penalty)
    LL_noPrior_MAP, all_signals = algo_fit_spec.rl_impl(data, top_MAP[1:end-1], opts...)
    println("LL top_MAP=$(round(LL_noPrior_MAP, digits=4))")

    flush(stdout)
    end_time = now()
    println("end of fitting: $end_time, duration: $(end_time-start_time)")
    top_MAP, top_n_avgMAP, top_n_min, top_n_max, top_n_stdMAP, top_n_samples, mcmc_samples, LL_noPrior_MAP, LL_noPrior_avgMAP, all_signals
end

function _write_subj_data_file(subj_data_path::String , data_stats::Data_stats ,subj_signals::Step_by_step_signals, pop_signals::Step_by_step_signals, subj_fit_mcmc_figure)
    # writes the data to the file system.
    # the directory subj_data_path is created if it does not yet exist.
    if !ispath(subj_data_path)
        mkpath(subj_data_path)
    end

    #learning signals when per-subject (posterior) params are used
    path = "$(subj_data_path)params_subj.csv"
    writedlm(path, subj_signals.param_vector)

    # always write the action_selection_prob
    writedlm("$(subj_data_path)action_selection_prob_subj.csv", subj_signals.action_selection_prob)

    if subj_fit_mcmc_figure != nothing
    # keep track of the MCMC samples for later visualizations
    # writedlm("$(subj_data_path)/mcmc_samples_subj.csv", subj_fit_mcmc_samples)
    # plot
        subj_fit_mcmc_figure.savefig("$(subj_data_path)MCMC.png")   # save the figure to file
    end

    if length(subj_signals.RPE)>0
      path = "$(subj_data_path)RPE_subj.csv"
        d = hcat(subj_signals.RPE, subj_signals.action_selection_prob)
        writedlm(path, d)
    end
    if length(subj_signals.SPE)>0
        path = "$(subj_data_path)SPE_subj.csv"
        d = hcat(subj_signals.SPE, subj_signals.action_selection_prob)
        writedlm(path, d)
    end

    # learning signals when population (prior) params are used:
    path = "$(subj_data_path)params_pop.csv"
    writedlm(path, pop_signals.param_vector)

    # always write the action_selection_prob
    writedlm("$(subj_data_path)action_selection_prob_pop.csv", pop_signals.action_selection_prob)

    if length(pop_signals.RPE)>0
        path = "$(subj_data_path)RPE_pop.csv"
          d = hcat(pop_signals.RPE, pop_signals.action_selection_prob)
          writedlm(path, d)

    end
    if length(pop_signals.SPE)>0
        path = "$(subj_data_path)SPE_pop.csv"
        writedlm(path, hcat(pop_signals.SPE, pop_signals.action_selection_prob))
    end

    if haskey(pop_signals.results_dict, "RPE_QV")
        path = "$(subj_data_path)RPE_QV_pop.csv"
        writedlm(path, pop_signals.results_dict["RPE_QV"])
    end
    if haskey(subj_signals.results_dict, "RPE_QV")
        path = "$(subj_data_path)RPE_QV_subj.csv"
        writedlm(path, subj_signals.results_dict["RPE_QV"])
    end

    if haskey(pop_signals.results_dict, "trialByTrial_Vs")
        path = "$(subj_data_path)StateValue_pop.csv"
        writedlm(path, pop_signals.results_dict["trialByTrial_Vs"])
    end
    if haskey(subj_signals.results_dict, "trialByTrial_Vs")
        path = "$(subj_data_path)StateValue_subj.csv"
        writedlm(path, subj_signals.results_dict["trialByTrial_Vs"])
    end

    ###############################
    # Store V-values (if there are any in the results_dict)
    if haskey(pop_signals.results_dict, "StateValue_V1_record")
        path = "$(subj_data_path)StateValue_V1_pop.csv"
        writedlm(path, pop_signals.results_dict["StateValue_V1_record"])
    end
    if haskey(subj_signals.results_dict, "StateValue_V1_record")
        path = "$(subj_data_path)StateValue_V1_subj.csv"
        writedlm(path, subj_signals.results_dict["StateValue_V1_record"])
    end
    if haskey(pop_signals.results_dict, "StateValue_V2_record")
        path = "$(subj_data_path)StateValue_V2_pop.csv"
        writedlm(path, pop_signals.results_dict["StateValue_V2_record"])
    end
    if haskey(subj_signals.results_dict, "StateValue_V2_record")
        path = "$(subj_data_path)StateValue_V2_subj.csv"
        writedlm(path, subj_signals.results_dict["StateValue_V2_record"])
    end
    if haskey(pop_signals.results_dict, "StateValue_V3_record")
        path = "$(subj_data_path)StateValue_V3_pop.csv"
        writedlm(path, pop_signals.results_dict["StateValue_V3_record"])
    end
    if haskey(subj_signals.results_dict, "StateValue_V3_record")
        path = "$(subj_data_path)StateValue_V3_subj.csv"
        writedlm(path, subj_signals.results_dict["StateValue_V3_record"])
    end
    ##############################

    if haskey(pop_signals.results_dict, "Surprise")
        path = "$(subj_data_path)Surprise_pop.csv"
        writedlm(path, pop_signals.results_dict["Surprise"])
    end
    if haskey(subj_signals.results_dict, "Surprise")
        path = "$(subj_data_path)Surprise_subj.csv"
        writedlm(path, subj_signals.results_dict["Surprise"])
    end
    if haskey(pop_signals.results_dict, "gamma_surprise")
        path = "$(subj_data_path)gamma_surprise_pop.csv"
        writedlm(path, pop_signals.results_dict["gamma_surprise"])
    end
    if haskey(subj_signals.results_dict, "gamma_surprise")
        path = "$(subj_data_path)gamma_surprise_subj.csv"
        writedlm(path, subj_signals.results_dict["gamma_surprise"])
    end
    ##############################
    if haskey(pop_signals.results_dict, "DeltaPolicy")
        path = "$(subj_data_path)DeltaPolicy_pop.csv"
        writedlm(path, pop_signals.results_dict["DeltaPolicy"])
    end
    if haskey(subj_signals.results_dict, "DeltaPolicy")
        path = "$(subj_data_path)DeltaPolicy_subj.csv"
        writedlm(path, subj_signals.results_dict["DeltaPolicy"])
    end
    #
    if haskey(pop_signals.results_dict, "PolicyParametersImbalance")
        path = "$(subj_data_path)PolicyParametersImbalance_pop.csv"
        writedlm(path, pop_signals.results_dict["PolicyParametersImbalance"])
    end
    if haskey(subj_signals.results_dict, "PolicyParametersImbalance")
        path = "$(subj_data_path)PolicyParametersImbalance_subj.csv"
        writedlm(path, subj_signals.results_dict["PolicyParametersImbalance"])
    end
    #
    if haskey(pop_signals.results_dict, "PolicyImbalance")
        path = "$(subj_data_path)PolicyImbalance_pop.csv"
        writedlm(path, pop_signals.results_dict["PolicyImbalance"])
    end
    if haskey(subj_signals.results_dict, "PolicyImbalance")
        path = "$(subj_data_path)PolicyImbalance_subj.csv"
        writedlm(path, subj_signals.results_dict["PolicyImbalance"])
    end
    #
    if haskey(pop_signals.results_dict, "PolicyParametersImbalanceAllstates")
        path = "$(subj_data_path)PolicyParametersImbalanceAllstates_pop.csv"
        writedlm(path, pop_signals.results_dict["PolicyParametersImbalanceAllstates"])
    end
    if haskey(subj_signals.results_dict, "PolicyParametersImbalanceAllstates")
        path = "$(subj_data_path)PolicyParametersImbalanceAllstates_subj.csv"
        writedlm(path, subj_signals.results_dict["PolicyParametersImbalanceAllstates"])
    end
    #
    if haskey(pop_signals.results_dict, "trial_by_trial_p_sa_max")
        path = "$(subj_data_path)action_selection_prob_MAX_pop.csv"
        writedlm(path, pop_signals.results_dict["trial_by_trial_p_sa_max"])
    end
    if haskey(subj_signals.results_dict, "trial_by_trial_p_sa_max")
        path = "$(subj_data_path)action_selection_prob_MAX_subj.csv"
        writedlm(path, subj_signals.results_dict["trial_by_trial_p_sa_max"])
    end

    # LL of subject
    path = "$(subj_data_path)LL_pop.csv"
    writedlm(path, pop_signals.LL)

    path = "$(subj_data_path)LL_subj.csv"
    writedlm(path, subj_signals.LL)

    # BIC of subject
    path = "$(subj_data_path)BIC_pop.csv"
    BIC_pop = get_BIC(pop_signals.LL, length(pop_signals.param_vector), data_stats.nr_actions)
    writedlm(path, BIC_pop)

    path = "$(subj_data_path)BIC_subj.csv"
    BIC_subj = get_BIC(subj_signals.LL, length(subj_signals.param_vector), data_stats.nr_actions)
    writedlm(path, BIC_subj)

    #
    if haskey(pop_signals.results_dict, "PolicyParametersDiffChosenVsUnchosen")
        path = "$(subj_data_path)PolicyParametersDiffChosenVsUnchosen_pop.csv"
        writedlm(path, pop_signals.results_dict["PolicyParametersDiffChosenVsUnchosen"])
    end
    if haskey(subj_signals.results_dict, "PolicyParametersDiffChosenVsUnchosen")
        path = "$(subj_data_path)PolicyParametersDiffChosenVsUnchosen_subj.csv"
        writedlm(path, subj_signals.results_dict["PolicyParametersDiffChosenVsUnchosen"])
    end
    ##############################

    # textfile describing unstructured data for info only
    path = "$(subj_data_path)description.txt"
    f = open(path, "w")
    write(f,"DateTime of fitting: $(now())\n")
    write(f,"DATA_STATS\n")
    write(f,to_string(data_stats))
    write(f,"\nSUBJECT_SIGNALS\n")
    write(f, to_string(subj_signals))
    write(f,"\nPOPULATION_SIGNALS\n")
    write(f, to_string(pop_signals))
    close(f)
end

function run_single_subj_fit_write_csv(
    data_path::String,
    base_path::String,
    id_subject::Int64,
    fit_max_episodes_per_subj::Int64,
    learner_function_V::Function,
    population_prior,
    population_map_params,
    burn_in::Int64,
    nr_mcmc_samples::Int64,
    inter_collect_interval::Int64,
    nr_of_repetitions::Int64,
    top_n_fraction::Float64,
    descritpion::String,
    rl_module_path::String;
    write_mcmc_picture_file = true)
    # a higher level script that writes all the information we currently
    # want to the file system.

    data_single_subj, data_stats = load_SARSPEI(data_path; subject_id = id_subject,
                                            max_episodes_per_subj=fit_max_episodes_per_subj)

    mSpec = Metropolis_Spec(learner_function_V, population_prior, burn_in,nr_mcmc_samples, inter_collect_interval)
    mcmc_samples = run_parallel_Metropolis(data_single_subj, mSpec, rl_module_path; nr_of_repetitions=nr_of_repetitions)

    nr_map_samples = Int(ceil(size(mcmc_samples)[1] * top_n_fraction))
    top_MAP, top_n_avgMAP, top_n_min, top_n_max, top_n_stdMAP, top_n_samples = get_MAP(mcmc_samples, nr_map_samples)

    map_val_subj = deepcopy(top_MAP)

    println("subj: $id_subject, population prior, map_val_subj: $(round.(map_val_subj,digits=4))")
    flush(stdout)
    # for evaluation (unlike fitting) load ALL Episodes
    data_single_subj, data_stats = load_SARSPEI(data_path; subject_id = id_subject)
    LL, subj_signals = learner_function_V(data_single_subj, map_val_subj)
    subj_signals = subj_signals[id_subject]

    LL_pop, pop_signals = learner_function_V(data_single_subj, population_map_params)
    pop_signals = pop_signals[id_subject]

    subj_fit_mcmc_figure = nothing
    if write_mcmc_picture_file
        fig_title = "Subj:$id_subject, $(RL_Fit.to_short_string(data_stats)), $learner_function_V, LL_sub: $(round(LL,digits=2)), LL_pop: $(round(LL_pop,digits=2))"
        subj_fit_mcmc_figure = RL_Fit.plot_result(mcmc_samples, population_prior, fig_title)#, top_n_fraction = 1.)
    end

    subj_data_path = "$(base_path)subj_$(id_subject)/"
    _write_subj_data_file(subj_data_path, data_stats,subj_signals, pop_signals, subj_fit_mcmc_figure)
    plt.close("all")
    subj_signals, pop_signals, mcmc_samples, data_stats
end

function get_BIC(LL, nr_params, nr_actions)
    BIC = -2. * LL + nr_params * log(nr_actions)
end

function get_BIClogevidence(LL, nr_params, nr_actions)
    BIClogevidence = LL - nr_params * log(nr_actions) / 2.
    # or
    # BIClogevidence = - get_BIC(LL, nr_params, nr_actions) / 2.
end


using Gnuplot
x = 0:0.3:4
a = exp.(- x)
b =  exp.(- x.^2)
# theme
function gp_theme(bgcolor = "#212946"; borderc ="white",
        ticsc = "white", ylabc = "white", gridst = 4,
        xlabc = "white", gridc = "#2A3459", labc = "white")
    bgcp1 = "set object rectangle from screen 0,0 to screen 1,1"
    bgcp2 = " behind fc '$(bgcolor)' fs solid noborder"
    bgcolor = bgcp1*bgcp2
    @gp :- "set border lw 1 lc '$(borderc)'"
    @gp :- "set tics textcolor rgb '$(ticsc)'"
    @gp :- bgcolor
    @gp :- "set ylabel 'y' textcolor '$(ylabc)'"
    @gp :- "set xlabel 'x' textcolor '$(xlabc)'"
    @gp :- "set grid ls 1 lc '$(gridc)' dt $(gridst)"
    @gp :- "set key t r textcolor '$(labc)'"
end

# plot
@gp x a "w linespoints lw 1 lc '#08F7FE' pt 7 t 'e^{-x}'"
@gp :- x b "w linespoints lw 1 lc '#FFE64D' pt 7 t 'e^{-x^2}'"
# glow effect
for i  in 1:10
    @gp :- x a "w l lw $(1 + 1.05*i) lc '#F508F7FE' t ''"
    @gp :- x b "w l lw $(1 + 1.05*i) lc '#F5FFE64D' t ''" :-
end
## fill between
@gp(:-,x, a, "w filledcu y=0 lw 1 lc '#08F7FE' t ''",
    "set style fill transparent solid 0.08 noborder")
@gp(:-, x, a, b, "w filledcu lw 1 lc '#FFE64D' t ''")
gp_theme("black")
save(term="pngcairo size 600,400", output="glow1.png") # hide


# Copyright (c) 2015, Intel Corporation
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
#     * Redistributions of source code must retain the above copyright notice,
#       this list of conditions and the following disclaimer.
#     * Redistributions in binary form must reproduce the above copyright
#       notice, this list of conditions and the following disclaimer in the
#       documentation and/or other materials provided with the distribution.
#     * Neither the name of Intel Corporation nor the names of its contributors
#       may be used to endorse or promote products derived from this software
#       without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

# FIXME: This should be moved into the Latte core as it uses special internal
# handling.
export ReshapeLayer

function ReshapeEnsemble(name::Symbol, input::AbstractEnsemble, shape)
    neurons = Array(ReshapeNeuron, shape...)
    for i in 1:length(neurons)
        neurons[i] = ReshapeNeuron(0.0, 0.0)
    end
    ReshapeEnsemble(name, neurons, [], TrainTest, 1, input, shape)
end

function init(ensemble::ReshapeEnsemble, net::Net)
    resized = (ensemble.shape..., batch_size(net))
    for t in 1:net.time_steps
        input = get_buffer(net, symbol(ensemble.input.name, :value), t)
        ∇ = get_buffer(net, symbol(ensemble.input.name, :∇), t)
        set_buffer(net, symbol(ensemble.name, :value), reshape(input, resized), t)
        set_buffer(net, symbol(ensemble.name, :∇), reshape(∇, resized), t)
    end
end

function ReshapeLayer(name::Symbol, net::Net, input_ensemble, shape)
    ens = ReshapeEnsemble(name, input_ensemble, shape)
    add_ensemble(net, ens)
    ens
end


using KernelAbstractions
using Test

import KernelAbstractions.NDIteration: NDRange, StaticSize, DynamicSize

@kernel function index(A)
    I  = @index(Global, NTuple)
    @show A[I...]
end
kernel = index(CPU(), DynamicSize(), DynamicSize())
iterspace = NDRange{1, StaticSize{(128,)}, StaticSize{(8,)}}();
ctx = KernelAbstractions.mkcontext(kernel, 1, nothing, iterspace, Val(KernelAbstractions.NoDynamicCheck()))

@test KernelAbstractions.Cassette.overdub(ctx, KernelAbstractions.__index_Global_NTuple, CartesianIndex(1)) == (1,)


################################################################################
#
#   IMPLEMENTATIONS OF DIFFERENT LATTICE MODIFICATION FUNCTIONS
#   (MOST OF THEM ALSO WORK FOR UNITCELLS)
#
#   STRUCTURE OF THE FILE
#
#   1) CONNECTIONS ADDING REMOVING
#       - Add a new connection
#       - remove connections (based on indices)
#       - remove connections (based on strength)
#
#   2) SITES ADDING REMOVING
#       - add a new site
#       - remove a site (based on index)
#
#   3) CONNECTION STRENGTH MODIFICATION
#       - all connections
#       - mapping of strengths
#       - evaluate strengths
#       - optimize connections
#
#   4) REAL SPACE ROTATION / SCALING / SHIFTING
#       - Rotation around some axis
#       - Scaling along some axis
#       - Global (isotropic) scaling
#       - Global (isotropic) scaling to mean / minimum bond length
#       - Shifting along some axis
#       - Shifting along some lattice vector
#       - Shifting center of lattice to origin
#
#   5) LATTICE VECTOR OPTIMIZATION
#
################################################################################

################################################################################
#
#   1) CONNECTIONS ADDING REMOVING
#       - Add a new connection
#       - remove connections (based on indices)
#       - remove connections (based on strength)
#
################################################################################

#-------------------------------------------------------------------------------
#
#  Add connection to given unitcell / lattice objects
#
#-------------------------------------------------------------------------------
"""
    addConnection!(unitcell::Unitcell, index_from::Int64, index_to::Int64, strength [, wrap=-1; overwrite::Bool=false])
    addConnection!(lattice::Lattice,   index_from::Int64, index_to::Int64, strength [, wrap=-1; overwrite::Bool=false])

Function to add a connection to either a `Unitcell` object or a `Lattice` object.
Necessary parameters include the indices of sites between which the connection takes place as well as the strength of the connection.
Additionally (in case of a unitcell or periodic boundaries), a wrap can be specified in terms of lattice vectors.
If no wrap is specified, the wrap will be chosen automatically as without wrapping around periodic boundaries.

The additional option `overwrite` specifies if one wants to simply add the connection
or if one wants to first search if there is a connection that already satisfies the new connection (in which case nothing would be added).

NOTE 1: This function always tries to add two connections, one going forward, one going backward.

NOTE 2: This function modifies the given object and does not create a copy.

# Examples

```julia-repl
julia> addConnection!(unitcell, 1, 3, "tx", (0,1))

julia> addConnection!(lattice, 31, 27, "tx")

```
"""
function addConnection!(unitcell::Unitcell, index_from::Int64, index_to::Int64, strength, wrap=-1; overwrite::Bool=false)

    # maybe default wrap
    if wrap == -1
        if length(unitcell.lattice_vectors) == 3
            wrap = (0,0,0)
        elseif length(unitcell.lattice_vectors) == 2
            wrap = (0,0)
        else
            wrap = (0)
        end
    end

    # construct the returning wrap
    if length(wrap) == 3
        wrap_return = (-wrap[1],-wrap[2],-wrap[3])
    elseif length(wrap) == 2
        wrap_return = (-wrap[1],-wrap[2])
    else
        wrap_return = (-wrap)
    end
    # construct the returning strength
    if typeof(strength) == Complex
        strength_return = deepcopy(conj(strength))
    else
        strength_return = deepcopy(strength)
    end

    # construct the connection as an array
    connection_1 = Any[index_from; index_to; strength; wrap]
    # construct the returning connection as an array
    connection_2 = Any[index_to; index_from; strength_return; wrap_return]

    # if not overwrite, check if the connection exists
    if !overwrite
        # report if found
        found_c1 = false
        found_c2 = false
        # check in all connections
        for c in unitcell.connections
            if  Int(c[1]) == Int(connection_1[1]) && Int(c[2]) == Int(connection_1[2]) && c[3] == connection_1[3] && c[4] == connection_1[4]
                found_c1 = true
            end
            if  Int(c[1]) == Int(connection_2[1]) && Int(c[2]) == Int(connection_2[2]) && c[3] == connection_2[3] && c[4] == connection_2[4]
                found_c2 = true
            end
        end
        # only add if not added already
        if !found_c1
            push!(unitcell.connections, connection_1);
        end
        if !found_c2
            push!(unitcell.connections, connection_2);
        end
    # otherwise, just add the connections
    else
        push!(unitcell.connections, connection_1);
        push!(unitcell.connections, connection_2);
    end

    # return nothing
    return nothing
end
function addConnection!(lattice::Lattice, index_from::Int64, index_to::Int64, strength, wrap=-1; overwrite::Bool=false)

    # maybe default wrap
    if wrap == -1
        if length(lattice.lattice_vectors) == 3
            wrap = (0,0,0)
        elseif length(lattice.lattice_vectors) == 2
            wrap = (0,0)
        else
            wrap = (0)
        end
    end

    # construct the returning wrap
    if length(wrap) == 3
        wrap_return = (-wrap[1],-wrap[2],-wrap[3])
    elseif length(wrap) == 2
        wrap_return = (-wrap[1],-wrap[2])
    else
        wrap_return = (-wrap)
    end
    # construct the returning strength
    if typeof(strength) == Complex
        strength_return = deepcopy(conj(strength))
    else
        strength_return = deepcopy(strength)
    end

    # construct the connection as an array
    connection_1 = Any[index_from; index_to; strength; wrap]
    # construct the returning connection as an array
    connection_2 = Any[index_to; index_from; strength_return; wrap_return]

    # if not overwrite, check if the connection exists
    if !overwrite
        # report if found
        found_c1 = false
        found_c2 = false
        # check in all connections
        for c in lattice.connections
            if  Int(c[1]) == Int(connection_1[1]) && Int(c[2]) == Int(connection_1[2]) && c[3] == connection_1[3] && c[4] == connection_1[4]
                found_c1 = true
            end
            if  Int(c[1]) == Int(connection_2[1]) && Int(c[2]) == Int(connection_2[2]) && c[3] == connection_2[3] && c[4] == connection_2[4]
                found_c2 = true
            end
        end
        # only add if not added already
        if !found_c1
            push!(lattice.connections, connection_1);
        end
        if !found_c2
            push!(lattice.connections, connection_2);
        end
    # otherwise, just add the connections
    else
        push!(lattice.connections, connection_1);
        push!(lattice.connections, connection_2);
    end

    # return nothing
    return nothing
end
export addConnection!

#-------------------------------------------------------------------------------
#
#   Remove interactions with given strengths of the unitcell / lattice
#
#-------------------------------------------------------------------------------
"""
    removeConnections!(unitcell::Unitcell, index_from::Int64, index_to::Int64 [, strength ])
    removeConnections!(lattice::Lattice,   index_from::Int64, index_to::Int64 [, strength ])

Function to remove connections from either a `Unitcell` object or a `Lattice` object based on the sites that they connect.
Optionally, one can pass a connection strength to also filter based on this strength. If no strength is passed, all connnections
between the two sites will be removed.

    removeConnections!(unitcell::Unitcell  [, strength])
    removeConnections!(lattice::Lattice    [, strength])

Function to remove connections from either a `Unitcell` object or a `Lattice` object only based on the strength of connections.
If no strength is passed, all connnections with Float64 strength `0.0` will be removed.

NOTE 1: This function always tries to delte both connections, one going forward, one going backward.

NOTE 2: This function modifies the given object and does not create a copy.

# Examples

```julia-repl
julia> removeConnections!(unitcell, 1, 3, "tx")

julia> removeConnections!(unitcell, 1, 3)

julia> removeConnections!(lattice, 23, 73)

julia> removeConnections!(lattice, "tx")

```
"""
function removeConnections!(lattice::Lattice,   index_from::Int64, index_to::Int64, strength)
    # use the julia filter function to filter out all elements which have the strength passed and the fitting indices
    filter!(x->!(x[3]==strength && Int(x[1])==index_from && Int(x[2])==index_to) && !(x[3]==strength && Int(x[2])==index_from && Int(x[1])==index_to), lattice.connections)
    # return nothing
    return nothing
end
function removeConnections!(unitcell::Unitcell, index_from::Int64, index_to::Int64, strength)
    # use the julia filter function to filter out all elements which have the strength passed and the fitting indices
    filter!(x->!(x[3]==strength && Int(x[1])==index_from && Int(x[2])==index_to) && !(x[3]==strength && Int(x[2])==index_from && Int(x[1])==index_to), unitcell.connections)
    # return nothing
    return nothing
end
function removeConnections!(lattice::Lattice,   index_from::Int64, index_to::Int64)
    # use the julia filter function to filter out all elements which have the fitting indices
    filter!(x->!(Int(x[1])==index_from && Int(x[2])==index_to) && !(Int(x[2])==index_from && Int(x[1])==index_to), lattice.connections)
    # return nothing
    return nothing
end
function removeConnections!(unitcell::Unitcell, index_from::Int64, index_to::Int64)
    # use the julia filter function to filter out all elements which have the fitting indices
    filter!(x->!(Int(x[1])==index_from && Int(x[2])==index_to) && !(Int(x[2])==index_from && Int(x[1])==index_to), unitcell.connections)
    # return nothing
    return nothing
end
function removeConnections!(lattice::Lattice,   strength=0.0)
    # use the julia filter function to filter out all elements which have the strength passed
    filter!(x->x[3]!=strength, lattice.connections)
    # return nothing
    return nothing
end
function removeConnections!(unitcell::Unitcell, strength=0.0)
    # use the julia filter function to filter out all elements which have the strength passed
    filter!(x->x[3]!=strength, unitcell.connections)
    # return nothing
    return nothing
end

export removeConnections!

################################################################################
#
#   2) SITES ADDING REMOVING
#       - add a new site
#       - remove a site (based on index)
#
################################################################################

"""
    addSite!(unitcell::Unitcell, position::Array{Float64,1})
    addSite!(lattice::Lattice,   position::Array{Float64,1})

Function to add a site to either a `Unitcell` object or a `Lattice` object which is located at
the real space position `position`. The function returns the index of the newly added site for further use.

NOTE 1: The newly added site is not connected to any other sites yet.

NOTE 2: This function modifies the given object and does not create a copy.

# Examples

```julia-repl
julia> addSite!(unitcell, Float64[0.0, 0.5])
5

julia> addSite!(lattice, Float64[0.0, 0.5, 1.0])
187
```
"""
function addSite!(lattice::Lattice, position::Array{Float64,1})
    # push the site into the positions list
    push!(lattice.positions, position)
    # return the site index
    return length(lattice.positions)
end
function addSite!(unitcell::Unitcell, position::Array{Float64,1})
    # push the site into the positions list
    push!(unitcell.basis, position)
    # return the site index
    return length(unitcell.basis)
end
export addSite!

"""
    removeSite!(unitcell::Unitcell, index::Int64)
    removeSite!(lattice::Lattice,   index::Int64)

Function to remove a site from either a `Unitcell` object or a `Lattice` object with index `index`.
The function furhtermore removes all connections that are connected to the site and relabels all
other connections as well to account for the index shift introduced by the removal.

NOTE: This function modifies the given object and does not create a copy.

# Examples

```julia-repl
julia> removeSite!(unitcell, 2)

julia> removeSite!(lattice, 128)

```
"""
function removeSite!(lattice::Lattice, index::Int64)
    # delete the element from the positions list
    deleteat!(lattice.positions, index)
    # delete all connections that touch the removed site
    deleteat!(lattice.connections, [c[1]==index || c[2]==index for c in lattice.connections])
    # now shift all indices in the connections array by 1
    for (i,c) in enumerate(lattice.connections)
        if Int(c[1]) > index
            c[1] = Int(c[1]) - 1
        end
        if Int(c[2]) > index
            c[2] = Int(c[2]) - 1
        end
    end
    # return nothing
    return nothing
end
function removeSite!(unitcell::Unitcell, index::Int64)
    # delete the element from the positions list
    deleteat!(unitcell.basis, index)
    # delete all connections that touch the removed site
    deleteat!(unitcell.connections, [c[1]==index || c[2]==index for c in unitcell.connections])
    # now shift all indices in the connections array by 1
    for (i,c) in enumerate(unitcell.connections)
        if Int(c[1]) > index
            c[1] = Int(c[1]) - 1
        end
        if Int(c[2]) > index
            c[2] = Int(c[2]) - 1
        end
    end
    # return nothing
    return nothing
end
export removeSite!

################################################################################
#
#   3) CONNECTION STRENGTH MODIFICATION
#       - all connections
#       - mapping of strengths
#       - evaluate strengths
#       - optimize connections
#
################################################################################

#-----------------------------------------------------------------------------------------------------------------------------
#
#   Set / Map interaction strengths of the unitcell / lattice
#
#-----------------------------------------------------------------------------------------------------------------------------
"""
    setAllInteractionStrengths!(unitcell::Unitcell, strength)
    setAllInteractionStrengths!(lattice::Lattice,   strength)

Function to overwrite ALL connection strength values of either a `Unitcell` object or a `Lattice` object
with the value `strength`.

NOTE: This function modifies the given object and does not create a copy.

# Examples

```julia-repl
julia> setAllInteractionStrengths!(unitcell, "tx")

julia> setAllInteractionStrengths!(lattice, 1.0)

```
"""
function setAllInteractionStrengths!(unitcell::Unitcell, strength)
    # go through all connections and modify their strength
    for c in unitcell.connections
        c[3] = strength
    end
    # return nothing
    return nothing
end
function setAllInteractionStrengths!(lattice::Lattice,   strength)
    # go through all connections and modify their strength
    for c in lattice.connections
        c[3] = strength
    end
    # return nothing
    return nothing
end
export setAllInteractionStrengths!

"""
    mapInteractionStrengths!(unitcell::Unitcell, mapping::Dict [; replace_in_strings::Bool=false, evaluate::Bool=false])
    mapInteractionStrengths!(lattice::Lattice,   mapping::Dict [; replace_in_strings::Bool=false, evaluate::Bool=false])

Function to map connection strength values of either a `Unitcell` object or a `Lattice` object
by replacing them along keys of a dictonary `mapping`.

Furthermore, to be able to resolve strengths like `"tx*ty"` (`String` valued) to a `Float64` value one can use the following two steps:
1) chosing `replace_in_strings=true` to also search for mapping keys in `String` connection strengths.
2) chosing `evaluate=true` to evaluate every `String` connection strength by `eval(parse(.))` after all replacement has taken place.

NOTE: This function modifies the given object and does not create a copy.

# Examples

Example mapping could read

    mapping = Dict("tx"=>1.0, "ty"=>1.0, "tz"=>2.0)

which modifies all `String` valued connection strengths `"tx"`, `"ty"` and `"tz"` to be `Float64` with values `1.0` and `2.0`

```julia-repl
julia> mapInteractionStrengths!(unitcell, mapping)

julia> mapInteractionStrengths!(lattice, mapping, replace_in_strings=true)

```
"""
function mapInteractionStrengths!(unitcell::Unitcell, mapping::Dict; replace_in_strings::Bool=false, evaluate::Bool=false)
    # get the old strengths
    old_strengths = keys(mapping)
    # iterate for replacement
    for old_strength in old_strengths
        # new strength is given by mapping
        new_strength = mapping[old_strength]
        # check all connections that are identical
        for c in unitcell.connections
            if c[3] == old_strength
                c[3] = new_strength
            end
        end
        # check for string replacement
        if replace_in_strings
            for c in unitcell.connections
                if typeof(c[3]) == String && contains(c[3], old_strength)
                    c[3] = replace(c[3], old_strength, new_strength)
                end
            end
        end
    end
    # maybe even evaluate
    if evaluate
        for c in unitcell.connections
            if typeof(c[3]) == String
                c[3] = eval(parse(c[3]))
            end
        end
    end
    # return nothing
    return nothing
end
function mapInteractionStrengths!(lattice::Lattice,   mapping::Dict; replace_in_strings::Bool=false, evaluate::Bool=false)
    # get the old strengths
    old_strengths = keys(mapping)
    # iterate for replacement
    for old_strength in old_strengths
        # new strength is given by mapping
        new_strength = mapping[old_strength]
        # check all connections that are identical
        for c in lattice.connections
            if c[3] == old_strength
                c[3] = new_strength
            end
        end
        # check for string replacement
        if replace_in_strings
            for c in lattice.connections
                if typeof(c[3]) == String && contains(c[3], old_strength)
                    c[3] = replace(c[3], old_strength, new_strength)
                end
            end
        end
    end
    # maybe even evaluate
    if evaluate
        for c in lattice.connections
            if typeof(c[3]) == String
                c[3] = eval(parse(c[3]))
            end
        end
    end
    # return nothing
    return nothing
end
export mapInteractionStrengths!

"""
    evaluateInteractionStrengths!(unitcell::Unitcell [, parameters::Dict])
    evaluateInteractionStrengths!(lattice::Lattice   [, parameters::Dict])

Function to evaluate connection strength values of either a `Unitcell` object or a `Lattice` object
by first optionally replacing them along keys of a dictonary `parameters` and then calling `eval(parse(.))`.

NOTE: This function modifies the given object and does not create a copy.

# Examples

Example parameters could read

    parameters = Dict("tx"=>1.0, "ty"=>1.0, "tz"=>2.0)

which modifies all `String` valued connection strengths `"tx"`, `"ty"` and `"tz"` to be `Float64` with values `1.0` and `2.0`

```julia-repl
julia> evaluateInteractionStrengths!(unitcell)

julia> evaluateInteractionStrengths!(lattice, parameters)

```
"""
function evaluateInteractionStrengths!(unitcell::Unitcell, parameters::Dict=Dict())
    # just call the mapping functions
    mapInteractionStrengths!(unitcell, parameters, replace_in_strings=true, evaluate=true)
end
function evaluateInteractionStrengths!(lattice::Lattice,   parameters::Dict=Dict())
    # just call the mapping functions
    mapInteractionStrengths!(lattice, parameters, replace_in_strings=true, evaluate=true)
end
export evaluateInteractionStrengths!

#-----------------------------------------------------------------------------------------------------------------------------
#
#   Optimize the lattice / unitcell by collapsing redundant connections into fewer connections and modify given object
#
#-----------------------------------------------------------------------------------------------------------------------------

"""
    optimizeConnections!(unitcell::Unitcell)
    optimizeConnections!(lattice::Lattice)

Function to optimize connections of either a `Unitcell` object or a `Lattice` object
by checking which ones can be compressed and put together.

NOTE: This function modifies the given object and does not create a copy.

# Examples

```julia-repl
julia> optimizeConnections!(unitcell)

```
"""
function optimizeConnections!(lattice::Lattice)
    # build up new connections
    connections_new = Array{Any,1}[]
    # go through all old connections to check for redundant entries
    for c in lattice.connections
        # check if the connection already is present in the new list
        found = false
        for c_new in connections_new
            if c[1] == c_new[1] && c[2] == c_new[2] && c[4] == c_new[4]
                # found a redundent connection
                found = true
                # add connection strength
                if typeof(c[3]) == String || typeof(c_new[3]) == String
                    c_new[3] = "$(c_new[3])+$(c[3])"
                else
                    c_new[3] = c_new[3] + c[3]
                end
                # break the inner loop
                break
            end
        end
        # if not, add the current connection
        if !found
            push!(connections_new, c)
        end
    end
    # check if the connection strength is 0
    connections_new_2 = Array{Any,1}[]
    for c in connections_new
        # check if close to 0 connection strength
        if typeof(c[3]) == Float64 || typeof(c[3]) == Int64
            if abs(c[3]) > 1e-18
                push!(connections_new_2, c)
            else
                # do not push! c[3] too small!
            end
        elseif isnull(tryparse(Float64,c[3]))
            # not a float and cannot be parsed to float, just push in list as it is
            push!(connections_new_2, c)
        else
            # check if the strength is non-zero
            if abs(parse(Float64,c[3])) > 1e-18
                push!(connections_new_2, c)
            else
                # do not push! c[3] too small!
            end
        end
    end
    # overwrite the connections in the given lattice
    lattice.connections = connections_new_2
    # return nothing
    return nothing
end
function optimizeConnections!(unitcell::Unitcell)
    # build up new connections
    connections_new = Array{Any,1}[]
    # go through all old connections to check for redundant entries
    for c in unitcell.connections
        # check if the connection already is present in the new list
        found = false
        for c_new in connections_new
            if c[1] == c_new[1] && c[2] == c_new[2] && c[4] == c_new[4]
                # found a redundent connection
                found = true
                # add connection strength
                if typeof(c[3]) == String || typeof(c_new[3]) == String
                    c_new[3] = "$(c_new[3])+$(c[3])"
                else
                    c_new[3] = c_new[3] + c[3]
                end
                # break the inner loop
                break
            end
        end
        # if not, add the current connection
        if !found
            push!(connections_new, c)
        end
    end
    # check if the connection strength is 0
    connections_new_2 = Array{Any,1}[]
    for c in connections_new
        # check if close to 0 connection strength
        if typeof(c[3]) == Float64 || typeof(c[3]) == Int64
            if abs(c[3]) > 1e-18
                push!(connections_new_2, c)
            else
                # do not push! c[3] too small!
            end
        elseif isnull(tryparse(Float64,c[3]))
            # not a float and cannot be parsed to float, just push in list as it is
            push!(connections_new_2, c)
        else
            # check if the strength is non-zero
            if abs(parse(Float64,c[3])) > 1e-18
                push!(connections_new_2, c)
            else
                # do not push! c[3] too small!
            end
        end
    end
    # overwrite the connections in the given lattice
    unitcell.connections = connections_new_2
    # return nothing
    return nothing
end
export optimizeConnections!

################################################################################
#
#   4) REAL SPACE ROTATION / SCALING / SHIFTING
#       - Rotation around some axis
#       - Scaling along some axis
#       - Global (isotropic) scaling
#       - Global (isotropic) scaling to mean / minimum bond length
#       - Shifting along some axis
#       - Shifting along some lattice vector
#       - Shifting center of lattice to origin
#
################################################################################

###################
#   ROTATION
###################

# ROTATION AROUND X AXIS
function rotateAroundXAxis!(unitcell::Unitcell, angle::Float64)
    # rotate the lattice vectors
    for lv in unitcell.lattice_vectors
        # rotate and overwrite simulateously all components
        lv[2], lv[3] = cos(angle)*lv[2] + sin(angle)*lv[3]  ,   -sin(angle)*lv[2] + cos(angle)*lv[3]
    end
    # rotate the basis sites
    for bs in unitcell.basis
        # rotate and overwrite simulateously all components
        bs[2], bs[3] = cos(angle)*bs[2] + sin(angle)*bs[3]  ,   -sin(angle)*bs[2] + cos(angle)*bs[3]
    end
end
function rotateAroundXAxis!(lattice::Lattice, angle::Float64)
    # rotate the lattice vectors
    for lv in lattice.lattice_vectors
        # rotate and overwrite simulateously all components
        lv[2], lv[3] = cos(angle)*lv[2] + sin(angle)*lv[3]  ,   -sin(angle)*lv[2] + cos(angle)*lv[3]
    end
    # rotate the basis sites
    for p in lattice.positions
        # rotate and overwrite simulateously all components
        p[2], p[3] = cos(angle)*p[2] + sin(angle)*p[3]  ,   -sin(angle)*p[2] + cos(angle)*p[3]
    end
end
export rotateAroundXAxis!

function rotateAroundXAxisDeg!(unitcell::Unitcell, angle::Float64)
    rotateAroundXAxis!(unitcell, angle*pi/180.0)
end
function rotateAroundXAxisDeg!(lattice::Lattice, angle::Float64)
    rotateAroundXAxis!(lattice, angle*pi/180.0)
end
export rotateAroundXAxisDeg!

# ROTATION AROUND Y AXIS
function rotateAroundYAxis!(unitcell::Unitcell, angle::Float64)
    # rotate the lattice vectors
    for lv in unitcell.lattice_vectors
        # rotate and overwrite simulateously all components
        lv[3], lv[1] = cos(angle)*lv[3] + sin(angle)*lv[1]  ,   -sin(angle)*lv[3] + cos(angle)*lv[1]
    end
    # rotate the basis sites
    for bs in unitcell.basis
        # rotate and overwrite simulateously all components
        bs[3], bs[1] = cos(angle)*bs[3] + sin(angle)*bs[1]  ,   -sin(angle)*bs[3] + cos(angle)*bs[1]
    end
end
function rotateAroundYAxis!(lattice::Lattice, angle::Float64)
    # rotate the lattice vectors
    for lv in lattice.lattice_vectors
        # rotate and overwrite simulateously all components
        lv[3], lv[1] = cos(angle)*lv[3] + sin(angle)*lv[1]  ,   -sin(angle)*lv[3] + cos(angle)*lv[1]
    end
    # rotate the basis sites
    for p in lattice.positions
        # rotate and overwrite simulateously all components
        p[3], p[1] = cos(angle)*p[3] + sin(angle)*p[1]  ,   -sin(angle)*p[3] + cos(angle)*p[1]
    end
end
export rotateAroundYAxis!

function rotateAroundYAxisDeg!(unitcell::Unitcell, angle::Float64)
    rotateAroundYAxis!(unitcell, angle*pi/180.0)
end
function rotateAroundYAxisDeg!(lattice::Lattice, angle::Float64)
    rotateAroundYAxis!(lattice, angle*pi/180.0)
end
export rotateAroundYAxisDeg!

# ROTATION AROUND Z AXIS
function rotateAroundZAxis!(unitcell::Unitcell, angle::Float64)
    # rotate the lattice vectors
    for lv in unitcell.lattice_vectors
        # rotate and overwrite simulateously all components
        lv[1], lv[2] = cos(angle)*lv[1] + sin(angle)*lv[2]  ,   -sin(angle)*lv[1] + cos(angle)*lv[2]
    end
    # rotate the basis sites
    for bs in unitcell.basis
        # rotate and overwrite simulateously all components
        bs[1], bs[2] = cos(angle)*bs[1] + sin(angle)*bs[2]  ,   -sin(angle)*bs[1] + cos(angle)*bs[2]
    end
end
function rotateAroundZAxis!(lattice::Lattice, anglTODOe::Float64)
    # rotate the lattice vectors
    for lv in lattice.lattice_vectors
        # rotate and overwrite simulateously all components
        lv[1], lv[2] = cos(angle)*lv[1] + sin(angle)*lv[2]  ,   -sin(angle)*lv[1] + cos(angle)*lv[2]
    end
    # rotate the basis sites
    for p in lattice.positions
        # rotate and overwrite simulateously all components
        p[1], p[2] = cos(angle)*p[1] + sin(angle)*p[2]  ,   -sin(angle)*p[1] + cos(angle)*p[2]
    end
end
export rotateAroundZAxis!

function rotateAroundZAxisDeg!(unitcell::Unitcell, angle::Float64)
    rotateAroundZAxis!(unitcell, angle*pi/180.0)
end
function rotateAroundZAxisDeg!(lattice::Lattice, angle::Float64)
    rotateAroundZAxis!(lattice, angle*pi/180.0)
end
export rotateAroundZAxisDeg!

###################
#   SCALING
###################

# SCALING ALONG X AXIS
function scaleAlongXAxis!(unitcell::Unitcell, factor::Float64)
    # scale all lattice vectors
    for l in unitcell.lattice_vectors
        l[1] *= factor
    end
    # scale all positions
    for p in unitcell.basis
        p[1] *= factor
    end
end
function scaleAlongXAxis!(lattice::Lattice, factor::Float64)
    # scale all lattice vectors
    for l in lattice.lattice_vectors
        l[1] *= factor
    end
    # scale all positions
    for p in lattice.positions
        p[1] *= factor
    end
end
export scaleAlongXAxis!

# SCALING ALONG Y AXIS
function scaleAlongYAxis!(unitcell::Unitcell, factor::Float64)
    # scale all lattice vectors
    for l in unitcell.lattice_vectors
        l[2] *= factor
    end
    # scale all positions
    for p in unitcell.basis
        p[2] *= factor
    end
end
function scaleAlongYAxis!(lattice::Lattice, factor::Float64)
    # scale all lattice vectors
    for l in lattice.lattice_vectors
        l[2] *= factor
    end
    # scale all positions
    for p in lattice.positions
        p[2] *= factor
    end
end
export scaleAlongYAxis!

# SCALING ALONG Z AXIS
function scaleAlongZAxis!(unitcell::Unitcell, factor::Float64)
    # scale all lattice vectors
    for l in unitcell.lattice_vectors
        l[3] *= factor
    end
    # scale all positions
    for p in unitcell.basis
        p[3] *= factor
    end
end
function scaleAlongZAxis!(lattice::Lattice, factor::Float64)
    # scale all lattice vectors
    for l in lattice.lattice_vectors
        l[3] *= factor
    end
    # scale all positions
    for p in lattice.positions
        p[3] *= factor
    end
end
export scaleAlongZAxis!

# SCALING ISOTROPIC
function scaleIsotropic!(unitcell::Unitcell, factor::Float64)
    # scale all lattice vectors
    for l in unitcell.lattice_vectors
        l .*= factor
    end
    # scale all positions
    for p in unitcell.basis
        p .*= factor
    end
end
function scaleIsotropic!(lattice::Lattice, factor::Float64)
    # scale all lattice vectors
    for l in lattice.lattice_vectors
        l .*= factor
    end
    # scale all positions
    for p in lattice.positions
        p .*= factor
    end
end
export scaleIsotropic!

# SCALING TO MEAN BOND LENGTH
function scaleToMeanBondLength!(unitcell::Unitcell, bond_length::Float64=1.0)
    # get all bond lengths
    bond_lengths = zeros(length(unitcell.connections))
    for i in 1:length(bond_lengths)
        # get the connection
        c = unitcell.connections[i]
        # get the connecting vector
        c_vector = unitcell.basis[Int(c[2])] .- unitcell.basis[Int(c[1])]
        for (j,l) in enumerate(unitcell.lattice_vectors)
            c_vector .+= l.*c[4][j]
        end
        # get how long it is
        bond_lengths[i] = sqrt(sum(c_vector.*c_vector))
    end
    # calculate the mean bond length
    mean_bond_length = mean(bond_lengths)
    # calculate how the scale factor looks like
    factor = bond_length / mean_bond_length
    # scale the object
    scaleIsotropic!(unitcell, factor)
end
function scaleToMeanBondLength!(lattice::Lattice, bond_length::Float64=1.0)
    # get all bond lengths
    bond_lengths = zeros(length(lattice.connections))
    for i in 1:length(bond_lengths)
        # get the connection
        c = lattice.connections[i]
        # get the connecting vector
        c_vector = lattice.positions[Int(c[2])] .- lattice.positions[Int(c[1])]
        for (j,l) in enumerate(lattice.lattice_vectors)
            c_vector .+= l.*c[4][j]
        end
        # get how long it is
        bond_lengths[i] = sqrt(sum(c_vector.*c_vector))
    end
    # calculate the mean bond length
    mean_bond_length = mean(bond_lengths)
    # calculate how the scale factor looks like
    factor = bond_length / mean_bond_length
    # scale the object
    scaleIsotropic!(lattice, factor)
end
export scaleToMeanBondLength!

# SCALING TO MINIMUM BOND LENGTH
function scaleToMinimumBondLength!(unitcell::Unitcell, bond_length::Float64=1.0)
    # get all bond lengths
    bond_lengths = zeros(length(unitcell.connections))
    for i in 1:length(bond_lengths)
        # get the connection
        c = unitcell.connections[i]
        # get the connecting vector
        c_vector = unitcell.basis[Int(c[2])] .- unitcell.basis[Int(c[1])]
        for (j,l) in enumerate(unitcell.lattice_vectors)
            c_vector .+= l.*c[4][j]
        end
        # get how long it is
        bond_lengths[i] = sqrt(sum(c_vector.*c_vector))
    end
    # calculate the minimum bond length
    minimum_bond_length = minimum(bond_lengths)
    # calculate how the scale factor looks like
    factor = bond_length / minimum_bond_length
    # scale the object
    scaleIsotropic!(unitcell, factor)
end
function scaleToMinimumBondLength!(lattice::Lattice, bond_length::Float64=1.0)
    # get all bond lengths
    bond_lengths = zeros(length(lattice.connections))
    for i in 1:length(bond_lengths)
        # get the connection
        c = lattice.connections[i]
        # get the connecting vector
        c_vector = lattice.positions[Int(c[2])] .- lattice.positions[Int(c[1])]
        for (j,l) in enumerate(lattice.lattice_vectors)
            c_vector .+= l.*c[4][j]
        end
        # get how long it is
        bond_lengths[i] = sqrt(sum(c_vector.*c_vector))
    end
    # calculate the minimum bond length
    minimum_bond_length = minimum(bond_lengths)
    # calculate how the scale factor looks like
    factor = bond_length / minimum_bond_length
    # scale the object
    scaleIsotropic!(lattice, factor)
end
export scaleToMinimumBondLength!

###################
#   SHIFTING
###################

# SHIFT ALONG X AXIS
function shiftAlongXAxis!(unitcell::Unitcell, offset::Float64)
    # shift all positions
    for p in unitcell.basis
        p[1] += offset
    end
end
function shiftAlongXAxis!(lattice::Lattice, offset::Float64)
    # shift all positions
    for p in lattice.positions
        p[1] += offset
    end
end
export shiftAlongXAxis!

# SHIFT ALONG Y AXIS
function shiftAlongYAxis!(unitcell::Unitcell, offset::Float64)
    # shift all positions
    for p in unitcell.basis
        p[2] += offset
    end
end
function shiftAlongYAxis!(lattice::Lattice, offset::Float64)
    # shift all positions
    for p in lattice.positions
        p[2] += offset
    end
end
export shiftAlongYAxis!

# SHIFT ALONG Z AXIS
function shiftAlongZAxis!(unitcell::Unitcell, offset::Float64)
    # shift all positions
    for p in unitcell.basis
        p[3] += offset
    end
end
function shiftAlongZAxis!(lattice::Lattice, offset::Float64)
    # shift all positions
    for p in lattice.positions
        p[3] += offset
    end
end
export shiftAlongZAxis!

# SHIFT BY VECTOR
function shiftByVector!(unitcell::Unitcell, vector::Array{Float64,1})
    # shift all positions
    for p in unitcell.basis
        p .+= vector
    end
end
function shiftByVector!(lattice::Lattice, vector::Array{Float64,1})
    # shift all positions
    for p in lattice.positions
        p .+= vector
    end
end
export shiftByVector!

# SHIFT ALONG ANY LATTICE VECTOR
function shiftAlongLatticeVector!(unitcell::Unitcell, lattice_vector::Int64, offset::Float64)
    # define a vector for shifting
    vector = unitcell.lattice_vectors[lattice_vector] .* (offset / sqrt(sum(unitcell.lattice_vectors[lattice_vector].*unitcell.lattice_vectors[lattice_vector])))
    # shift by that vector
    shiftByVector!(unitcell, vector)
end
function shiftAlongLatticeVector!(lattice::Lattice, lattice_vector::Int64, offset::Float64)
    # define a vector for shifting
    vector = lattice.lattice_vectors[lattice_vector] .* (offset / sqrt(sum(lattice.lattice_vectors[lattice_vector].*lattice.lattice_vectors[lattice_vector])))
    # shift by that vector
    shiftByVector!(lattice, vector)
end
export shiftAlongLatticeVector!

# SHIFT ALONG A1
function shiftAlongA1!(unitcell::Unitcell, offset::Float64)
    # shift along the first lattice vector
    shiftAlongLatticeVector!(unitcell, 1, offset)
end
function shiftAlongA1!(lattice::Lattice, offset::Float64)
    # shift along the first lattice vector
    shiftAlongLatticeVector!(lattice, 1, offset)
end
export shiftAlongA1!

# SHIFT ALONG A2
function shiftAlongA2!(unitcell::Unitcell, offset::Float64)
    # shift along the second lattice vector
    shiftAlongLatticeVector!(unitcell, 2, offset)
end
function shiftAlongA2!(lattice::Lattice, offset::Float64)
    # shift along the second lattice vector
    shiftAlongLatticeVector!(lattice, 2, offset)
end
export shiftAlongA2!

# SHIFT ALONG A3
function shiftAlongA3!(unitcell::Unitcell, offset::Float64)
    # shift along the third lattice vector
    shiftAlongLatticeVector!(unitcell, 3, offset)
end
function shiftAlongA3!(lattice::Lattice, offset::Float64)
    # shift along the third lattice vector
    shiftAlongLatticeVector!(lattice, 3, offset)
end
export shiftAlongA3!

# SHIFT ALONG ANY LATTICE VECTOR (Relative)
function shiftAlongLatticeVectorRelative!(unitcell::Unitcell, lattice_vector::Int64, offset::Float64)
    # define a vector for shifting
    vector = unitcell.lattice_vectors[lattice_vector] .* offset
    # shift by that vector
    shiftByVector!(unitcell, vector)
end
function shiftAlongLatticeVectorRelative!(lattice::Lattice, lattice_vector::Int64, offset::Float64)
    # define a vector for shifting
    vector = lattice.lattice_vectors[lattice_vector] .* offset
    # shift by that vector
    shiftByVector!(lattice, vector)
end
export shiftAlongLatticeVectorRelative!

# SHIFT ALONG A1 (Relative)
function shiftAlongA1Relative!(unitcell::Unitcell, offset::Float64)
    # shift along the first lattice vector
    shiftAlongLatticeVectorRelative!(unitcell, 1, offset)
end
function shiftAlongA1Relative!(lattice::Lattice, offset::Float64)
    # shift along the first lattice vector
    shiftAlongLatticeVectorRelative!(lattice, 1, offset)
end
export shiftAlongA1Relative!

# SHIFT ALONG A2 (Relative)
function shiftAlongA2Relative!(unitcell::Unitcell, offset::Float64)
    # shift along the second lattice vector
    shiftAlongLatticeVectorRelative!(unitcell, 2, offset)
end
function shiftAlongA2Relative!(lattice::Lattice, offset::Float64)
    # shift along the second lattice vector
    shiftAlongLatticeVectorRelative!(lattice, 2, offset)
end
export shiftAlongA2Relative!

# SHIFT ALONG A3 (Relative)
function shiftAlongA3Relative!(unitcell::Unitcell, offset::Float64)
    # shift along the third lattice vector
    shiftAlongLatticeVectorRelative!(unitcell, 3, offset)
end
function shiftAlongA3Relative!(lattice::Lattice, offset::Float64)
    # shift along the third lattice vector
    shiftAlongLatticeVectorRelative!(lattice, 3, offset)
end
export shiftAlongA3Relative!

# SHIFT CENTER TO ORIGIN
function shiftCenterToOrigin!(unitcell::Unitcell)
    # look for the central point
    center = unitcell.basis .* 0.0
    for p in unitcell.basis
        center .+= p
    end
    center ./= length(unitcell.basis)
    # shift by negative of that vector
    shiftByVector!(unitcell, center.*(-1))
end
function shiftCenterToOrigin!(lattice::Lattice)
    # look for the central point
    center = lattice.positions .* 0.0
    for p in lattice.positions
        center .+= p
    end
    center ./= length(lattice.positions)
    # shift by negative of that vector
    shiftByVector!(lattice, center.*(-1))
end
export shiftCenterToOrigin!

################################################################################
#
#   5) LATTICE VECTOR OPTIMIZATION
#
################################################################################

# Function to optimize lattice vectors
function optimizeLatticeVectors2D!(unitcell::Unitcell; verbose::Bool=false)
    # init the loop by assuming one refolded a vector
    refolded_one_vector = true
    # as long as one vector was refolded, repeat this loop
    while refolded_one_vector
        # init inside the loop by assuming one did not refold a vector
        refolded_one_vector = false
        # try a1
        if dot(unitcell.lattice_vectors[1],unitcell.lattice_vectors[1]) > dot(unitcell.lattice_vectors[1].-unitcell.lattice_vectors[2], unitcell.lattice_vectors[1].-unitcell.lattice_vectors[2])
            # print
            if verbose
                println("a1 --> a1 - a2")
            end
            # change the lattice vector
            unitcell.lattice_vectors[1] = unitcell.lattice_vectors[1] .- unitcell.lattice_vectors[2]
            # change all connections
            for c in unitcell.connections
                if c[4][1] != 0
                    c[4] = (c[4][1], c[4][2]+c[4][1])
                end
            end
            # did refold
            refolded_one_vector = true
        end
        if dot(unitcell.lattice_vectors[1],unitcell.lattice_vectors[1]) > dot(unitcell.lattice_vectors[1].+unitcell.lattice_vectors[2], unitcell.lattice_vectors[1].+unitcell.lattice_vectors[2])
            # print
            if verbose
                println("a1 --> a1 + a2")
            end
            # change the lattice vector
            unitcell.lattice_vectors[1] = unitcell.lattice_vectors[1] .+ unitcell.lattice_vectors[2]
            # change all connections
            for c in unitcell.connections
                if c[4][1] != 0
                    c[4] = (c[4][1], c[4][2]-c[4][1])
                end
            end
            # did refold
            refolded_one_vector = true
        end
        # try a2
        if dot(unitcell.lattice_vectors[2],unitcell.lattice_vectors[2]) > dot(unitcell.lattice_vectors[2].-unitcell.lattice_vectors[1], unitcell.lattice_vectors[2].-unitcell.lattice_vectors[1])
            # print
            if verbose
                println("a2 --> a2 - a1")
            end
            # change the lattice vector
            unitcell.lattice_vectors[2] = unitcell.lattice_vectors[2] .- unitcell.lattice_vectors[1]
            # change all connections
            for c in unitcell.connections
                if c[4][2] != 0
                    c[4] = (c[4][1]+c[4][2], c[4][2])
                end
            end
            # did refold
            refolded_one_vector = true
        end
        if dot(unitcell.lattice_vectors[2],unitcell.lattice_vectors[2]) > dot(unitcell.lattice_vectors[2].+unitcell.lattice_vectors[1], unitcell.lattice_vectors[2].+unitcell.lattice_vectors[1])
            # print
            if verbose
                println("a2 --> a2 + a1")
            end
            # change the lattice vector
            unitcell.lattice_vectors[2] = unitcell.lattice_vectors[2] .+ unitcell.lattice_vectors[1]
            # change all connections
            for c in unitcell.connections
                if c[4][2] != 0
                    c[4] = (c[4][1]-c[4][2], c[4][2])
                end
            end
            # did refold
            refolded_one_vector = true
        end
    end
end

# Function to optimize lattice vectors
function optimizeLatticeVectors3D!(unitcell::Unitcell; verbose::Bool=false)
    # init the loop by assuming one refolded a vector
    refolded_one_vector = true
    # as long as one vector was refolded, repeat this loop
    while refolded_one_vector
        # init inside the loop by assuming one did not refold a vector
        refolded_one_vector = false
        # try all combinations of the remaining two lattice vectors
        for comb in Array{Int64,1}[[1,0], [-1,0], [0,1], [0,-1], [1,1], [1,-1], [-1,1], [-1,-1]]
            # try a1
            a1p = unitcell.lattice_vectors[1] .+ comb[1].*unitcell.lattice_vectors[2] .+ comb[2].*unitcell.lattice_vectors[3]
            # check if this is closer
            if dot(unitcell.lattice_vectors[1],unitcell.lattice_vectors[1]) > dot(a1p, a1p)
                # print
                if verbose
                    println("a1 --> a1 + $(comb[1])*a2 + $(comb[2])*a3")
                end
                # change the lattice vector
                unitcell.lattice_vectors[1] = a1p
                # change all connections
                for c in unitcell.connections
                    c[4] = (c[4][1], c[4][2]-c[4][1]*comb[1], c[4][3]-c[4][1]*comb[2])
                end
                # did refold
                refolded_one_vector = true
            end
            # try a2
            a2p = unitcell.lattice_vectors[2] .+ comb[1].*unitcell.lattice_vectors[1] .+ comb[2].*unitcell.lattice_vectors[3]
            # check if this is closer
            if dot(unitcell.lattice_vectors[2],unitcell.lattice_vectors[2]) > dot(a2p, a2p)
                # print
                if verbose
                    println("a2 --> a2 + $(comb[1])*a1 + $(comb[2])*a3")
                end
                # change the lattice vector
                unitcell.lattice_vectors[2] = a2p
                # change all connections
                for c in unitcell.connections
                    c[4] = (c[4][1]-c[4][2]*comb[1], c[4][2], c[4][3]-c[4][2]*comb[2])
                end
                # did refold
                refolded_one_vector = true
            end
            # try a3
            a3p = unitcell.lattice_vectors[3] .+ comb[1].*unitcell.lattice_vectors[1] .+ comb[2].*unitcell.lattice_vectors[2]
            # check if this is closer
            if dot(unitcell.lattice_vectors[3],unitcell.lattice_vectors[3]) > dot(a3p, a3p)
                # print
                if verbose
                    println("a3 --> a3 + $(comb[1])*a1 + $(comb[2])*a2")
                end
                # change the lattice vector
                unitcell.lattice_vectors[3] = a3p
                # change all connections
                for c in unitcell.connections
                    c[4] = (c[4][1]-c[4][3]*comb[1], c[4][2]-c[4][3]*comb[2], c[4][3])
                end
                # did refold
                refolded_one_vector = true
            end
        end
    end
end

# general function
function optimizeLatticeVectors!(unitcell::Unitcell; verbose::Bool=false)
    if length(unitcell.lattice_vectors) == 2
        return optimizeLatticeVectors2D!(unitcell, verbose=verbose)
    elseif length(unitcell.lattice_vectors) == 3
        return optimizeLatticeVectors3D!(unitcell, verbose=verbose)
    else
        # DO NOTHING
    end
end
export optimizeLatticeVectors!


using BITEModel
using Base.Test
const BM = BITEModel

## Forward model
################

plotyes = false
verbose = true

global gid, hs1d, ivs1d, hs2d, ivs2d, glaciertype, # this is a bit ugly but
       gl, gb, pp, pm, pn                          # needed because of
                                                   # include("../scripts/forward.jl")

@testset "Forward model with update_cache=$update_cache" for update_cache in [false,true]
    # Synthetic (run by including scripts/forward.jl)
    ###########
    glaciertype = :synth
    plotyes = false
    include("../scripts/forward.jl")
    Base.Test.@inferred BM.fwdm(gb, pp, pm, pn, gb.bdot-gb.dhdt, gb.fsl, gb.temp)
    Base.Test.@inferred BM.make_bands(gl, pp, pm, pn)
    fwdsol = BM.fwdm(gb, pp, pm, pn, gb.bdot-gb.dhdt, gb.fsl, gb.temp)

    hs1d_ = [0.0, 167.025, 182.466, 189.223, 188.814, 183.506, 175.65,
             166.279, 151.378, 141.485, 131.018, 119.681, 109.229, 98.3605,
             86.7481, 76.5545, 67.8539, 63.5029, 57.3474, 50.484, 1.0]
    ivs1d_ = [0.0, 3.3608, 5.31713, 6.91192, 8.26612, 9.03612, 9.67009,
              10.0828, 10.5786, 11.0246, 10.9542, 10.8031, 10.8764, 10.6826,
              9.68466, 8.8367, 7.85266, 6.36658, 4.69378, 2.97981, 0.0]
    @test norm(hs1d_-fwdsol.hs1d)<2e-3
    @test norm(ivs1d_-fwdsol.ivs1d)<5e-4

    # test 2D at some random indices
    @test length(fwdsol.hs2d)==15921
    inds = [8179, 15516, 3523, 5846, 3236, 13314, 4959, 11, 12505, 6528,
            11823, 13553, 1043, 1420, 4210, 1991, 34, 4056, 5848, 15447,
            13452, 390, 8071, 9023, 12655, 3717, 7865, 8238, 14551, 4782]
    hs2d_ = [176.314, 0.0, 129.79, 172.317, 108.296, 0.0, 0.0, 0.0,
             57.2049, 0.0, 104.003, 0.0, 0.0, 0.0, 55.4563, 0.0, 0.0, 161.232,
             174.862, 0.0, 109.54, 0.0, 95.0336, 222.902, 119.131, 82.9983,
             98.8966, 230.184, 0.0, 132.357]
    ivs2d_ = [8.22742, 0.0, 6.17785, 7.93149, 6.84787, 0.0, 0.0, 0.0, 0.0,
              0.0, 7.88523, 0.0, 0.0, 0.0, 7.28846, 0.0, 0.0, 5.7638,
              7.93149, 0.0, 4.89699, 0.0, 0.0, 6.58108, 6.01924,
              7.96476, 6.96876, 6.70248, 0.0, 8.22056]
    @test norm(hs2d_-fwdsol.hs2d[inds])<1e-3
    @test norm(ivs2d_-fwdsol.ivs2d[inds])<5e-4

    # # ITMIX (run by hand)
    # #####################
    # glacier=:Urumqi # pretty fast this one
    # pl_kws = Dict{Symbol,Any}(:use_glogem => false)
    # gid = ITMIXGlacier(glacier)
    # gl,gb,pp,pm,pn,pl = BM.init_forward(gid, verbose;
    #                                     update_cache=update_cache, pl_kws...)
    # hs1d, taus1d, ivs1d, hs2d, taus2d, ivs2d, taus1d_l, vol_ratio, gb =
    #     BM.fwdm(gb, pp, pm, pn, gb.bdot-gb.dhdt, gb.fsl, gb.temp)
    # Base.Test.@inferred BM.fwdm(gb, pp, pm, pn, gb.bdot-gb.dhdt, gb.fsl, gb.temp)
    # Base.Test.@inferred BM.make_bands(gl, pp, pm, pn)

    # hs1d_ = [0.0, 25.4951, 29.2112, 32.2141, 32.3665, 30.4775, 29.6186,
    #          29.2802, 27.4351, 28.3908, 29.75, 37.2109, 39.2642,
    #          40.3066, 42.8661, 44.964 , 48.2977, 53.8874, 56.7994,
    #          57.8052, 51.7617, 40.3284, 31.8477, 5.63301]
    # ivs1d_ = [0.0, 1.47265, 2.2546, 3.01979, 3.92654, 4.90726,
    #           4.87954, 4.17233, 3.82581, 3.51174, 3.2171, 2.61505, 2.57666,
    #           2.68062, 2.84157, 3.30573, 3.93735, 4.08862, 4.08278, 4.08057,
    #           4.2603, 4.31557, 2.45511, 0.00525341]
    # @test norm(hs1d_-hs1d)<5e-4
    # @test norm(ivs1d-ivs1d_)<5e-4

    # # test 2D at some random indices
    # @test length(hs2d)==133331
    # inds = [85082, 118470, 7808, 31087, 78719, 96341, 104214, 9715,
    #         91721, 131435, 49463, 127799, 36643, 131306, 71573, 88897,
    #         110465, 58186, 93094, 29907, 105007, 65361, 118730, 67826,
    #         118921, 66973, 114737, 84487, 19245, 126314]
    # hs2d_ = [0.0, 54.1206, 0.0, 0.0, 39.7876, 0.0, 0.0, 39.0787,
    #          23.5233, 0.0, 0.0, 19.2427, 0.0, 0.0, 22.237, 50.1744,
    #          0.0, 71.186, 0.0, 35.3194, 27.0517, 77.9176, 0.0, 21.425,
    #          0.0, 0.0, 62.8376, 82.3543, 43.1886, 28.9674]
    # ivs2d_ = [0.0, 1.73392, 0.0, 0.0, 3.50925, 0.0, 0.0, 1.3902,
    #           5.06693, 0.0, 0.0, 3.04547, 0.0, 0.0, 2.57997, 3.94191,
    #           0.0, 2.76635, 0.0, 1.6151, 6.13513, 3.14324, 0.0,
    #           2.31482, 0.0, 0.0, 1.87241, 4.19182, 1.40529, 1.68439]
    # @test norm(hs2d_-hs2d[inds])<5e-4
    # @test norm(ivs2d_-ivs2d[inds])<5e-4
end

# make sure this runs
include("../scripts/forward-advanced.jl")

## Inverse model
################
@testset "Inverse model" begin
    gid = SyntheticGlacier(:bench)
    gl,gb,pp,pm,pn,pl = BM.init_forward(gid)
    (theta0, logposterior, logposterior1d, logprior, logprior_debug,
     loglikelihood, loglikelihood1d, fwdm_fn, fwdm1d_fn,
     pmcmc_defaults, fit_target) =
         BM.init_inverse(gb, pp, pm, pn)
    Base.Test.@inferred fwdm_fn(theta0.th0)
    Base.Test.@inferred logprior(theta0.th0)
    Base.Test.@inferred loglikelihood(theta0.th0)
    Base.Test.@inferred logposterior(theta0.th0)

    pmcmc = BM.MCMCNum(;pmcmc_defaults...,
                       niter=5*10^4,
                       nthin = 50)

    # sample posterior
    srand(0)
    thetas, theta_expect, theta_mode, accept_ratio, blobs, lpost, mc_hs1d, mc_ivs1d,
      mh2d, sh2d, miv2d, siv2d, min_h2d, max_h2d, min_iv2d, max_iv2d = mcmc(logposterior, theta0, pmcmc; verbose=true)
    @test round(accept_ratio,5) == 0.11801

    @test norm(theta_expect- [1.81352, 1.40765, -0.29361, -0.170199,
                              -0.0292185, -0.15876, -0.3106, -0.165316, -0.0345219, 0.172756,
                              0.406927, 0.162265, -0.144996, 0.315479, 0.16641, -0.00986009,
                              0.0872772, -0.109316, 0.711055, -0.239801, -0.121702, -0.0558777,
                              -0.219639]) < 1e-5
    # test 2D at some random indices
    inds = [8179, 15516, 3523, 5846, 3236, 13314, 4959, 11, 12505, 6528,
            11823, 13553, 1043, 1420, 4210, 1991, 34, 4056, 5848, 15447,
            13452, 390, 8071, 9023, 12655, 3717, 7865, 8238, 14551, 4782]

    @test norm(mh2d[inds]-[195.185, 0.0, 112.593, 175.143, 110.134, 0.0, 0.0, 0.0, 48.5617,
                           0.0, 115.334, 0.0, 0.0, 0.0, 48.79, 0.0, 0.0, 127.909, 173.234,
                           0.0, 86.9006, 0.0, 80.6748, 169.704, 104.964, 85.447, 86.3577,
                           180.154, 0.0, 147.823])<2e-3

    @test norm(sh2d[inds]-[14.5854, 0.0, 21.0239, 18.2294, 11.4231, 0.0, 0.0, 0.0, 9.05472,
                           0.0, 10.1193, 0.0, 0.0, 0.0, 19.4267, 0.0, 0.0, 24.128, 21.1017,
                           0.0, 16.3925, 0.0, 15.0425, 27.2756, 19.6065, 9.78661, 33.134,
                           32.304, 0.0, 11.914])<2e-4

    @test norm(miv2d[inds]-[11.8817, 0.0, 8.42797, 10.3403, 10.106,
                            0.0, 0.0, 0.0, 0.0, 0.0, 12.5615, 0.0, 0.0, 0.0, 16.4126, 0.0,
                            0.0, 7.73834, 10.3403, 0.0, 7.61675, 0.0, 0.0, 7.85299, 8.40403,
                            14.5605, 16.3398, 8.09658, 0.0, 12.651])<1e-4
end


# Struct for combining useful policy parameters together
@with_kw mutable struct PolicyParams{Pol <: Policy, T2 <: AbstractSpace}
    π::Pol
    space::T2 = action_space(π)
    π_explore = π
    π⁻ = nothing
end
PolicyParams(π::T) where {T <: Policy} = PolicyParams(π=π)

abstract type NetworkPolicy <: Policy end

# Fixing bug with gpu deepcopies
function Base.deepcopy(x::NetworkPolicy)
    invoke(deepcopy, Tuple{Any}, (x |> cpu))  |> device(x)
end

device(π::N) where N<:NetworkPolicy = π.device
actor(π::N) where N<:NetworkPolicy = π
critic(π::N) where N<:NetworkPolicy = π

# Call policies as functions calls the value function
(pol::NetworkPolicy)(x...) = value(pol, x...)

layers(π::Chain) = π.layers

function polyak_average!(to, from, τ=1f0)
    to_data = Flux.params(to).order.data
    from_data, from_device = Flux.params(from).order.data, device(from)
    device_match = from_device == device(to)
    for i = 1:length(to_data)
        if device_match
            copyto!(to_data[i], τ.*from_data[i] .+ (1f0-τ).*to_data[i])
        else
            copyto!(to_data[i], τ.*from_data[i] .+ (1f0-τ).*from_device(to_data[i]))
        end
    end
end

function Base.copyto!(to, from)
    for i = 1:length(Flux.params(to).order.data)
        copyto!(Flux.params(to).order.data[i], Flux.params(from).order.data[i])
    end
end

## Network for representing continous functions (value or policy)
mutable struct ContinuousNetwork <: NetworkPolicy
    network
    output_dim
    device
    function ContinuousNetwork(network, output_dim=nothing, dev=nothing)
        if isnothing(output_dim)
            try
                # Flux v0.12+
                output_dim = size(last(network.layers).weight,1)
            catch
                # Flux v0.11
                output_dim = size(last(network.layers).W,1)
            end
        end
        new(network, output_dim, device(network))
    end
end

Flux.@functor ContinuousNetwork

Flux.trainable(π::ContinuousNetwork) = Flux.trainable(π.network)

layers(π::ContinuousNetwork) = π.network.layers

POMDPs.action(π::ContinuousNetwork, s) = value(π, s)

POMDPs.value(π::ContinuousNetwork, s) = mdcall(π.network, s, π.device)

POMDPs.value(π::ContinuousNetwork, s, a) = mdcall(π.network, vcat(s,a), π.device)

action_space(π::ContinuousNetwork) = ContinuousSpace(π.output_dim)

## Network for representing a discrete set of outputs (value or policy)
# NOTE: Incoming actions (i.e. arguments) are all assumed to be one hot encoding. Outputs are discrete actions taken form outputs
mutable struct DiscreteNetwork <: NetworkPolicy
    network
    outputs
    logit_conversion
    always_stochastic
    device
    DiscreteNetwork(network, outputs, logit_conversion=softmax, always_stochastic=false, dev=nothing) = new(network, cpu(outputs), logit_conversion, always_stochastic, device(network))
end

Flux.@functor DiscreteNetwork

Flux.trainable(π::DiscreteNetwork) = Flux.trainable(π.network)

layers(π::DiscreteNetwork) = π.network.layers

POMDPs.value(π::DiscreteNetwork, s) = mdcall(π.network, s, π.device)

POMDPs.value(π::DiscreteNetwork, s, a_oh) = sum(value(π, s) .* a_oh, dims = 1)

POMDPs.action(π::DiscreteNetwork, s) = π.always_stochastic ? exploration(π, s)[1] : π.outputs[mapslices(argmax, value(π, s), dims=1)]

function Flux.onehotbatch(π::DiscreteNetwork, a)
    ignore() do
        a_oh = Flux.onehotbatch(a[:] |> cpu, π.outputs) |> device(a)
        length(a) == 1 ? dropdims(a_oh, dims=2) : a_oh
    end
end

logits(π::DiscreteNetwork, s) = π.logit_conversion(value(π, s))

categorical_logpdf(probs, a_oh) = log.(sum(probs .* a_oh, dims = 1) .+ eps(Float32))

function exploration(π::DiscreteNetwork, s; kwargs...)
    ps = logits(π, s)
    a = π.outputs[mapslices((v)->rand(Categorical(v)), ps, dims=1)]
    a, categorical_logpdf(ps, Flux.onehotbatch(π, a))
end

function Distributions.logpdf(π::DiscreteNetwork, s, a)
    # If a does not seem to be a one-hot encoding then we encode it
    Zygote.ignore() do
        size(a,1) == 1 && (a = Flux.onehotbatch(π, a))
    end
    return categorical_logpdf(logits(π, s), a)
end

function Distributions.entropy(π::DiscreteNetwork, s)
    ps = logits(π, s)
    sum(ps .* log.(ps .+ eps(Float32)), dims=1)
end

action_space(π::DiscreteNetwork) = DiscreteSpace(length(π.outputs), π.outputs)

## Double Network architecture
mutable struct DoubleNetwork{T1, T2} <: NetworkPolicy
    N1::T1
    N2::T2
end

Flux.@functor DoubleNetwork

Flux.trainable(π::DoubleNetwork) = (Flux.trainable(π.N1)..., Flux.trainable(π.N2)...)

layers(π::DoubleNetwork) = unique((layers(π.N1)..., layers(π.N2)...))

device(π::DoubleNetwork) = device(π.N1) == device(π.N2) ? device(π.N1) : error("Mismatched devices")

POMDPs.value(π::DoubleNetwork, s) = (value(π.N1, s), value(π.N2, s))

POMDPs.value(π::DoubleNetwork, s, a) = (value(π.N1, s, a), value(π.N2, s, a))

POMDPs.action(π::DoubleNetwork, s) = (action(π.N1, s), action(π.N2, s))

exploration(π::DoubleNetwork, s; kwargs...) = (exploration(π.N1, s; kwargs...), exploration(π.N2, s; kwargs...))

Distributions.logpdf(π::DoubleNetwork, s, a) = (logpdf(π.N1, s, a), logpdf(π.N2, s, a))

Distributions.entropy(π::DoubleNetwork, s) = (entropy(π.N1, s), entropy(π.N2, s))

action_space(π::DoubleNetwork) = action_space(π.N1)

## Actor Critic Architecture
mutable struct ActorCritic{TA, TC} <: NetworkPolicy
    A::TA # actor
    C::TC # critic
end

Flux.@functor ActorCritic

Flux.trainable(π::ActorCritic) = (Flux.trainable(π.A)..., Flux.trainable(π.C)...)

layers(π::ActorCritic) = unique((layers(π.A)..., layers(π.C)...))

device(π::ActorCritic) = device(π.A) == device(π.C) ? device(π.A) : error("Mismatched devices")

POMDPs.value(π::ActorCritic, s) = value(π.C, s)

POMDPs.value(π::ActorCritic, s, a) = value(π.C, s, a)

POMDPs.action(π::ActorCritic, s) = action(π.A, s)

exploration(π::ActorCritic, s; kwargs...) = exploration(π.A, s; kwargs...)

Distributions.logpdf(π::ActorCritic, s, a) = logpdf(π.A, s, a)

Distributions.entropy(π::ActorCritic, s) = entropy(π.A, s)

action_space(π::ActorCritic) = action_space(π.A)

actor(π::AC) where AC<:ActorCritic = π.A
critic(π::AC) where AC<:ActorCritic = π.C

## Network for concatentaing a latent vector to states
mutable struct LatentConditionedNetwork <: NetworkPolicy
    policy
    z
end

device(π::LatentConditionedNetwork) = device(π.policy)

Flux.@functor LatentConditionedNetwork

Flux.trainable(π::LatentConditionedNetwork) = Flux.trainable(π.policy)

layers(π::LatentConditionedNetwork) = layers(π.policy)

POMDPs.action(π::LatentConditionedNetwork, s) = action(π.policy, vcat(s, π.z))
POMDPs.value(π::LatentConditionedNetwork, s, args...) = value(π.policy, vcat(s, π.z), args...)

action_space(π::LatentConditionedNetwork) = action_space(π.policy)
actor(π::LatentConditionedNetwork) = actor(π.policy)
critic(π::LatentConditionedNetwork) = critic(π.policy)

## Gaussian Policy
mutable struct GaussianPolicy <: NetworkPolicy
    μ::ContinuousNetwork
    logΣ::ContinuousNetwork
    always_stochastic::Bool
    GaussianPolicy(μ::ContinuousNetwork, logΣ::ContinuousNetwork, always_stochastic=false) = new(μ, logΣ, always_stochastic)
    GaussianPolicy(μ::ContinuousNetwork, logΣ::AbstractArray, always_stochastic=false) = new(μ, ContinuousNetwork(Chain(ConstantLayer(logΣ)), length(logΣ)), always_stochastic)
end

Flux.@functor GaussianPolicy

Flux.trainable(π::GaussianPolicy) = (Flux.trainable(π.μ)..., Flux.trainable(π.logΣ)...)

layers(π::GaussianPolicy) = (layers(π.μ)..., layers(π.logΣ))

device(π::GaussianPolicy) = device(π.μ) == device(π.logΣ) ? device(π.μ) : error("Mismatched devices")

POMDPs.action(π::GaussianPolicy, s) = π.always_stochastic ? exploration(π, s)[1] : π.μ(s)

function gaussian_logpdf(μ, logΣ, a)
    σ² = exp.(logΣ).^2
    sum(-((a .- μ).^2) ./ (2 .* σ²) .-  0.9189385332046727f0 .- logΣ, dims = 1) # 0.9189385332046727f0 = log(sqrt(2π))
end

function exploration(π::GaussianPolicy, s; kwargs...)
    μ, logΣ = π.μ(s), π.logΣ(s)
    σ = exp.(logΣ)
    ϵ = Zygote.ignore(() -> randn(Float32, size(μ)...) |> device(s))
    a = ϵ.*σ .+ μ
    a, gaussian_logpdf(μ, logΣ, a)
end

Distributions.logpdf(π::GaussianPolicy, s, a) = gaussian_logpdf(π.μ(s), π.logΣ(s), a)

Distributions.entropy(π::GaussianPolicy, s) = 1.4189385332046727f0 .+ sum(π.logΣ(s)) # 1.4189385332046727 = 0.5 + 0.5 * log(2π)

action_space(π::GaussianPolicy) = action_space(π.μ)

## Squashed Gaussian policy
mutable struct SquashedGaussianPolicy <: NetworkPolicy
    μ::ContinuousNetwork
    logΣ::ContinuousNetwork
    ascale::Float32
    always_stochastic::Bool
    SquashedGaussianPolicy(μ, logΣ, ascale=1f0, always_stochastic=false) = new(μ, logΣ, ascale, always_stochastic)
    SquashedGaussianPolicy(μ, logΣ::Array, ascale=1f0, always_stochastic=false) = new(μ, ContinuousNetwork(Chain(ConstantLayer(logΣ)), length(logΣ)), ascale, always_stochastic)
end

Flux.@functor SquashedGaussianPolicy

Flux.trainable(π::SquashedGaussianPolicy) = (Flux.trainable(π.μ)..., Flux.trainable(π.logΣ)...)

layers(π::SquashedGaussianPolicy) = unique((layers(π.μ)..., layers(π.logΣ)...))

device(π::SquashedGaussianPolicy) = device(π.μ) == device(π.logΣ) ? device(π.μ) : error("Mismatched devices")

POMDPs.action(π::SquashedGaussianPolicy, s) = π.always_stochastic ? exploration(π,s)[1] : π.ascale .* tanh.(π.μ(s))

function squashed_gaussian_σ(logΣ)
    LOG_STD_MIN = -5
    LOG_STD_MAX = 2
    # logΣ = LOG_STD_MIN .+ 0.5f0 .* (LOG_STD_MAX - LOG_STD_MIN) * (logΣ .+ 1)
    logΣ = clamp.(logΣ, LOG_STD_MIN, LOG_STD_MAX)
    exp.(logΣ)
end

# a is  untanh'd
function squashed_gaussian_logprob(μ, logΣ, a)
    σ² = squashed_gaussian_σ(logΣ).^2
    sum(-((a .- μ).^2) ./ (2 .* σ²) .- 0.9189385332046727f0 .- logΣ .- 2*(log(2f0) .- a .- softplus.(-2 .* a)), dims=1)
end

function exploration(π::SquashedGaussianPolicy, s; kwargs...)
    μ, logΣ = π.μ(s), π.logΣ(s)
    σ = squashed_gaussian_σ(logΣ)
    ϵ = Zygote.ignore(() -> randn(Float32, size(μ)...) |> device(s))
    a_pretanh = ϵ.*σ .+ μ
    π.ascale .* tanh.(a_pretanh), squashed_gaussian_logprob(μ ,logΣ, a_pretanh)
end

Distributions.logpdf(π::SquashedGaussianPolicy, s, a) = squashed_gaussian_logprob(π.μ(s), π.logΣ(s), atanh.(clamp.(a ./ π.ascale, -1f0 + 1f-5, 1f0 - 1f-5)))

Distributions.entropy(π::SquashedGaussianPolicy, s) = 1.4189385332046727f0 .+ sum(π.logΣ(s), dims=1) # 1.4189385332046727 = 0.5 + 0.5 * log(2π) #TODO: This doesn't account for the squash

action_space(π::SquashedGaussianPolicy) = action_space(π.μ)

## Distribution policy -> For state-independent policies
mutable struct DistributionPolicy <: Policy
    distribution
end

layers(π::DistributionPolicy) = ()

device(π::DistributionPolicy) = cpu
actor(π::DistributionPolicy) = π
POMDPs.action(π::DistributionPolicy, s) = exploration(π, s)[1]

function exploration(π::DistributionPolicy, s; kwargs...)
    B = ndims(s) > 1 ? (size(s)[end]) : () # NOTE: this is a hack and doesn't work if the state space has more than one dim
    a = rand(π.distribution, B...)
    if a isa Float64 || a isa AbstractArray{Float64}
        a = Float32.(a)
    end
    a |> device(s), logpdf(π, s, a)
end

Distributions.logpdf(π::DistributionPolicy, s, a) = Float32.(reshape([logpdf(π.distribution, a)...], 1, :)) |> device(s)

function Distributions.entropy(π::DistributionPolicy, s)
    B = ndims(s) > 1 ? (size(s)[end]) : () # NOTE: this is a hack and doesn't work if the state space has more than one dim
    Float32(entropy(π.distribution)) * ones(Float32, 1, B...)
end

action_space(π::DistributionPolicy) = ContinuousSpace(length(π.distribution))

## Mixed policy
mutable struct MixedPolicy <: Policy
    ϵ::Function
    policy
end

MixedPolicy(ϵ::Real, policy) = MixedPolicy((i) -> ϵ, policy)
ϵGreedyPolicy(ϵ, actions) = MixedPolicy(ϵ, DistributionPolicy(ObjectCategorical(actions)))

function exploration(π::MixedPolicy, s; π_on, i,)
    ϵ = π.ϵ(i)
    x = rand() < ϵ ? action(π.policy, s) : action(π_on, s)

    # Turn the action into an array if it is a value
    !(x isa AbstractArray || x isa Tuple) && (x=fill(x, 1))

    logp1 = Base.log(ϵ) .+ logpdf(π.policy, s, x)
    logp2 = Base.log(1-ϵ) .+ logpdf(π_on, s, x)

    x, logsumexp(vcat(logp1, logp2), dims=1)
end

## Exploration policy with Gaussian noise
@with_kw mutable struct GaussianNoiseExplorationPolicy <: Policy
    σ::Function = (i) -> 0.01f0
    a_min::Vector{Float32} = [-Inf32]
    a_max::Vector{Float32} = [Inf32]
    ϵ_min::Float32 = -Inf32
    ϵ_max::Float32 = Inf32
end

GaussianNoiseExplorationPolicy(σ::Real; kwargs...) = GaussianNoiseExplorationPolicy(σ = (i) -> σ; kwargs...)
GaussianNoiseExplorationPolicy(σ::Function; kwargs...) = GaussianNoiseExplorationPolicy(σ = σ; kwargs...)

function exploration(π::GaussianNoiseExplorationPolicy, s; π_on, i)
    a = action(π_on, s) |> cpu
    ϵ = randn(Float32, size(a)...)*π.σ(i)
    clamp.(a .+ clamp.(ϵ, π.ϵ_min, π.ϵ_max), π.a_min, π.a_max) |> device(s), NaN
end

## use exploration policy for first N timesteps, then revert to base policy
@with_kw mutable struct FirstExplorePolicy <: Policy
    N::Int64 # Number of steps to explore for
    initial_policy::Policy # the policy to use for the first N steps
    after_policy::Union{Nothing, Policy} = nothing # the policy to use after the first N steps. Nothing means you will use on-policy
end

FirstExplorePolicy(N::Int64, initial_policy::Policy) = FirstExplorePolicy(N, initial_policy, nothing)

function exploration(π::FirstExplorePolicy, s; π_on, i)
    if i < π.N
        return action(π.initial_policy, s), NaN
    elseif isnothing(π.after_policy)
        return action(π_on, s), NaN
    else
        return exploration(π.after_policy, s, π_on=π_on, i=i)
    end
end


using UnPack

@testset "Parameters" begin
    # First constructor
    p = Parameters(0.4, 1.4737, 1e-5, Float32)
    @unpack ϕ, kT, N, τ, seed = p
    @test ϕ == 0.4f0
    @test kT == 1.4737f0
    @test N == 512
    @test τ == 0.00001f0
    @test seed == 0

    # Second constructor, define number of particles
    num_particles = 256
    p1 = Parameters(0.4, 1.4737, num_particles, 1e-5, Float32)
    @unpack ϕ, kT, N, τ, seed = p1
    @test ϕ == 0.4f0
    @test kT == 1.4737f0
    @test N == num_particles
    @test τ == 0.00001f0
    @test seed == 0

    # Third constructor, define number of particles and random seed
    num_particles = 256
    new_seed = 804533
    p2 = Parameters{Float32}(0.4, 1.4737, num_particles, 1e-5, new_seed)
    @unpack ϕ, kT, N, τ, seed = p2
    @test ϕ == 0.4f0
    @test kT == 1.4737f0
    @test N == num_particles
    @test τ == 0.00001f0
    @test seed == new_seed
end
@testset "SimulationSystem" begin
    # Create the parameters and simulation system from it
    p = Parameters(0.4, 1.4737, 1e-5, Float64)
    # 3 is the dimension of the system, a 3D simulation system
    dims = 3
    simul_sys = SimulationSystem(p, dims, Float64)
    @unpack params, ρ, boxl, rc, system, energy = simul_sys
    # Check the parameters
    @unpack ϕ, kT, N, τ, seed = p
    @test ϕ == 0.4
    @test kT == 1.4737
    @test N == 512
    @test τ == 0.00001
    @test seed == 0
    # Now check the system parameters
    density = 6.0 * params.ϕ / π
    @test density ≈ ρ
    @test boxl == cbrt(params.N / density)
    @test rc == boxl * 0.5
    @test energy == 0.0
    # Check that the positions and forces are correctly initialized
    total_size = (512, dims)
    @test size(system.positions) == total_size
    @test size(system.forces) == total_size
    @test all(system.positions .== 0.0)
    @test all(system.forces .== 0.0)
end


using Random

mutable struct Experience{O,A}
    observation::O
    action::A
    reward::Float64
    value::Float64
end

mutable struct FixedSizeBuffer{O,A} <: AbstractBuffer
    experiences::Vector{Experience{O,A}}
    buffersize::Int
    tail::Int
    len::Int
    sample_order::Vector{Int}
    last_sample::Int

    function FixedSizeBuffer{O,A}(buffersize::Int, rng::MersenneTwister) where {O,A}
        v = Vector{Experience{O,A}}(undef, buffersize)

        new(v, buffersize, 1, 0, randperm(rng, buffersize), 1)
    end
end

function add!(buffer::FixedSizeBuffer{O,A}, obs::O, action::A, reward::Float64) where {O,A}
    buffer.experiences[buffer.tail] = Experience{O,A}(obs, action, reward, 0.0)
    buffer.tail = mod1(buffer.tail + 1, buffer.buffersize)
    buffer.len = buffer.len == buffer.buffersize ? buffer.buffersize : buffer.len + 1

    nothing
end

Base.length(buffer::FixedSizeBuffer) = buffer.len

# Should i conform to the AbstractArray interface?
get(buffer::FixedSizeBuffer, i::Int) = buffer.experiences[mod1(i, buffer.len)]

# Preciso de um mecanismo para iterar pelas últimas n experiências (para atualizar seus valores)
struct BufferIterator{O,A}
    buffer::FixedSizeBuffer{O,A}
    len::Int
end

function Base.iterate(iter::BufferIterator, state=1)
    state > iter.len && return nothing

    return (get(iter.buffer, iter.buffer.tail - state), state + 1)
end

Base.eltype(::Type{BufferIterator{O,A}}) where {O,A} = Experience{O,A}

Base.length(iter::BufferIterator) = iter.len

take_last(buffer::FixedSizeBuffer, n::Int) = BufferIterator(buffer, n)

# Preciso de um mecanismo para amostrar um minibatch do buffer

function sample!(buffer::FixedSizeBuffer{O,A}, batchsize::Int) where {O,A}
    batch = Vector{Experience{O,A}}(undef, batchsize)

        for i in 1:batchsize
        batch[i] = get(buffer, buffer.last_sample + i)
    end

    buffer.last_sample = mod1(buffer.last_sample + batchsize, buffer.buffersize)

    return batch
end


### pop! functions
import Base.pop!

# Include reconstruct procedures for Compressed Stacks
include("reconstruct.jl")

## Functions that empty (pop) the element in partially compressed blocks
function empty_first!(stack::CompressedStack, index::Int, lvl::Int)
  pop!(stack.first_partial[lvl])
  if lvl > 1
    if length(stack.first_partial[lvl]) == 0
      empty_first!(stack, index, lvl-1)
    else
      new_index = stack.first_partial[lvl][end].last
      propagate_first!(stack, new_index, lvl-1)
    end
  end
end

function empty_second!(stack::CompressedStack, index::Int, lvl::Int)
  if !isempty(stack.first_partial[lvl])
    empty_first!(stack, index, lvl)
  else
    pop!(stack.second_partial[lvl])
    if lvl > 1
      if length(stack.second_partial[lvl]) == 0
        empty_second!(stack, index, lvl-1)
      else
        new_index = stack.second_partial[lvl][end].last
        propagate_second!(stack, new_index, lvl-1)
      end
    end
  end
end

## Functions to propagate the index of the element that have been popped
function propagate_first!(stack::CompressedStack, index::Int, lvl::Int)
  for i in 1:lvl
    update_top!(stack.first_partial[i], index)
  end
end

function propagate_second!(stack::CompressedStack, index::Int, lvl::Int)
  if !isempty(stack.first_partial[lvl])
    propagate_first!(stack, index, lvl)
  else
    update_top!(stack.second_partial[lvl], index)
    if lvl > 1
      propagate_second!(stack, index, lvl - 1)
    end
  end
end

## Functions to pop from first/second explicit then from partial
function pop_first!(stack::CompressedStack)
  index = top(stack)
  elt = pop!(stack.first_explicit)
  stack.pop_action(stack, elt.data)
  if stack.depth > 1
    if isempty(stack.first_explicit)
      empty_first!(stack, index, stack.depth - 1)
    else
      new_index = stack.first_explicit[end].index
      propagate_first!(stack, new_index, stack.depth - 1)
    end
  end
end

function pop_second!(stack::CompressedStack)
  index = top(stack)
  elt = pop!(stack.second_explicit)
  stack.pop_action(stack, elt.data)
  if stack.depth > 1
    if isempty(stack.second_explicit)
      empty_second!(stack, index, stack.depth - 1)
    else
      new_index = stack.second_explicit[end].index
      propagate_second!(stack, new_index, stack.depth - 1)
    end
  end
end

## redefinition of pop! for Compressed Stacks
function pop!(stack::CompressedStack)
  if !isempty(stack.first_explicit)
    pop_first!(stack)
  else
    if isempty(stack.second_explicit)
      # Reconstruct the compressed stack with the first available signature
      # Implementation is in reconstruct.jl
      reconstruct!(stack)
    end
    pop_second!(stack)
  end
end

####################
## redefinition of pop! for Normal Stacks
function pop!(stack::NormalStack)
  before = position(stack.input)
  elt = pop!(stack.data)
  stack.pop_action(stack, elt.data)
  after = position(stack.input)
  if before != after
    println("Before and after are different!")
  end
end


# R2_IRI_temporal_include.jl

include("../tcm_code/tcm_transactionsToCounts.jl")

include("../tcm_code/tcm_estimateMNL.jl")
include("../tcm_code/tcm_predictMNL.jl")

include("../tcm_code/tcm_estimateMMNL.jl")
include("../tcm_code/tcm_predictMMNL.jl")

include("../tcm_code/tcm_estimateHALOMNL.jl")
include("../tcm_code/tcm_predictHALOMNL.jl")

include("../tcm_code/tcm_estimateRanking.jl")
include("../tcm_code/tcm_predictRanking.jl")

include("../tcm_code/tcm_estimateForest.jl")

include("../tcm_code/tcm_predictForest.jl")

include("../tcm_code/tcm_evaluateKL.jl")

include("../tcm_code/tcm_leafStats.jl")
include("../tcm_code/tcm_rankingStats.jl")

include("R2_IRI_temporal_helperfn.jl")