text:
module Fitter

using ..Models: DataParameters, ClusterParameters, GaussianParameters, NtlParameters, DpParameters
using ..Models: FiniteTopicModelParameters, FiniteTopicModelSufficientStatistics
using ..Models: DataSufficientStatistics, MixtureSufficientStatistics, GaussianSufficientStatistics
using ..Models: MultinomialParameters, MultinomialSufficientStatistics, GeometricArrivals
using ..Models: HmmSufficientStatistics, StationaryHmmSufficientStatistics, NonstationaryHmmSufficientStatistics
using ..Models: Model, Mixture, ClusterSufficientStatistics
using ..Models: ParametricArrivalsClusterParameters
using ..Models: PitmanYorArrivals, SufficientStatistics
using ..Models: ArrivalDistribution, AuxillaryVariables, ArrivalsAuxillaryVariables
using ..Models: GaussianAuxillaryVariables, MultinomialAuxillaryVariables
using ..Models: FiniteTopicModelAuxillaryVariables, GeometricAuxillaryVariables
using ..Models: DpAuxillaryVariables
using ..Samplers: Sampler, GibbsSampler, SequentialMonteCarlo, SequentialImportanceSampler
using ..Samplers: MetropolisWithinGibbsSampler, MetropolisHastingsSampler
using ..Utils: logbeta, logmvbeta, log_multinomial_coeff, gumbel_max, isnothing, compute_normalized_weights
using ..Utils: effective_sample_size, hist
using ..Utils: mvt_logpdf, logfactorial
using ..Utils: hasproperty

using Distributions
using LinearAlgebra
using Statistics
using ProgressBars
using ProgressMeter
using SparseArrays
using SpecialFunctions
using DataStructures

asserting() = true # when set to true, this will enable all `@assertion`s

macro assertion(test)
    esc(:(if $(@__MODULE__).asserting()
        @assert($test)
    end))
end

function random_initial_assignment!(auxillary_variables::AuxillaryVariables{A, D}, sufficient_stats::SufficientStatistics,
                                    cluster_assignments::Vector{Int64}, ::Model, data) where
                                    {A, D <: Union{GaussianAuxillaryVariables, MultinomialAuxillaryVariables}}
    nothing
end

function random_initial_assignment!(auxillary_variables::AuxillaryVariables{A, D}, sufficient_stats::SufficientStatistics,
                                    cluster_assignments::Vector{Int64}, model::Model, data) where {A, D <: FiniteTopicModelAuxillaryVariables}
    num_documents = length(cluster_assignments)
    num_tokens = size(data)[1]
    num_topics = model.data_parameters.num_topics
    for document = 1:num_documents
        cluster = cluster_assignments[document]
        for token = 1:num_tokens
            num_instances = data[token, document]
            for instance = 1:num_instances
                new_topic = rand(1:num_topics)
                auxillary_variables.data.token_topic_assignments[document][token][instance] = new_topic
                # update sufficient statistics
                update_token_topic_statistics!(document, token, new_topic, cluster, auxillary_variables, sufficient_stats,
                                               update_type="add")
            end
        end
    end
end

function random_initial_assignment!(instances::Matrix, data::Matrix, sufficient_stats::SufficientStatistics{C, D},
                                    model::Mixture, auxillary_variables::AuxillaryVariables{A, F},
                                    sampler::MetropolisWithinGibbsSampler) where {C, D, A, F}
    n = sufficient_stats.n
    iteration = 1
    cluster = 1
    clusters = []
    for observation = 1:n
        coin_flip = reshape(rand(Binomial(1, 0.1), 1), 1)[1]
        if coin_flip === 1 || length(clusters) === 0
            cluster = observation
            append!(clusters, cluster)
        else
            cluster = rand(clusters)
        end
        add_observation!(observation, cluster, instances, iteration, data, sufficient_stats, model, auxillary_variables)
    end
    random_initial_assignment!(auxillary_variables, sufficient_stats, vec(instances[:, iteration]), model, data)
end

function random_initial_assignment!(instances::Matrix, data::Matrix, sufficient_stats::SufficientStatistics{C, D},
                                    model::Mixture, auxillary_variables::AuxillaryVariables{A, F},
                                    sampler::MetropolisHastingsSampler) where {C, D, A, F}
    n = sufficient_stats.n
    iteration = 1
    cluster = 1
    clusters = []
    for observation = 1:n
        coin_flip = reshape(rand(Binomial(1, 0.1), 1), 1)[1]
        if coin_flip === 1 || length(clusters) === 0
            cluster = observation
            append!(clusters, cluster)
        else
            clusters_to_sample = clusters[clusters .< observation]
            clusters_to_sample = clusters_to_sample[maximum([1, end - sampler.proposal_radius]):end]
            cluster = rand(clusters_to_sample)
        end
        add_observation!(observation, cluster, instances, iteration, data, sufficient_stats, model, auxillary_variables)
    end
    random_initial_assignment!(auxillary_variables, sufficient_stats, vec(instances[:, iteration]), model, data)
end

function deterministic_initial_assignment!(instances, data, sufficient_stats, model::Mixture, auxillary_variables,
                                           sampler::MetropolisWithinGibbsSampler, chain)
    n = sufficient_stats.n
    iteration = 1
    cluster = 1
    for observation = 1:n
        if sampler.assignment_types[chain] === "all same cluster"
            cluster = 1
        elseif sampler.assignment_types[chain] === "all different clusters"
            cluster = observation
        else
            error("Error in assignment type")
        end
        add_observation!(observation, cluster, instances, iteration, data, sufficient_stats, model, auxillary_variables)
    end
    random_initial_assignment!(auxillary_variables, sufficient_stats, vec(instances[:, iteration]), model, data)
end

function prepare_cluster_parameter_array(model::Mixture{C, D}, num_total_iterations, num_chains) where {C <: NtlParameters, D}
    return Array{Float64}(undef, 2, num_total_iterations, num_chains)
end

function prepare_cluster_parameter_array(model::Mixture{C, D}, num_total_iterations, num_chains) where {C <: DpParameters, D}
    return Array{Float64}(undef, num_total_iterations, num_chains)
end

function prepare_cluster_parameter_array(model::Mixture{C, D}, num_total_iterations) where {C <: NtlParameters, D}
    return Array{Float64}(undef, 2, num_total_iterations)
end

function prepare_cluster_parameter_array(model::Mixture{C, D}, num_total_iterations) where {C <: DpParameters, D}
    return Array{Float64}(undef, num_total_iterations)
end

function fit(data::Matrix{T}, model, sampler::MetropolisWithinGibbsSampler) where {T <: Real}
    sample_cluster_parameter_posterior = model.cluster_parameters.sample_parameter_posterior
    if hasproperty(model.cluster_parameters, :arrival_distribution)
        sample_arrival_parameter_posterior = model.cluster_parameters.arrival_distribution.sample_parameter_posterior
    else
        sample_arrival_parameter_posterior = false
    end
    num_total_iterations = sampler.num_iterations + sampler.num_burn_in
    n = size(data)[2]
    num_chains = length(sampler.assignment_types)
    output_assignments = Array{Int64}(undef, n, num_total_iterations, num_chains)
    output_log_likelihoods = Array{Float64}(undef, num_total_iterations, num_chains)
    if sample_arrival_parameter_posterior
        output_arrival_parameter_posterior = Array{Float64}(undef, num_total_iterations, num_chains)
    end
    if sample_cluster_parameter_posterior
        output_cluster_parameter_posterior = prepare_cluster_parameter_array(model, num_total_iterations, num_chains)
    end

    original_cluster_parameters = deepcopy(model.cluster_parameters)

    for chain = 1:num_chains
        model.cluster_parameters = deepcopy(original_cluster_parameters)
        assignments = Array{Int64}(undef, n, num_total_iterations)
        log_likelihoods = Vector{Float64}(undef, num_total_iterations)
        sufficient_stats = prepare_sufficient_statistics(model, data)
        auxillary_variables = prepare_auxillary_variables(model, data)
        if sample_arrival_parameter_posterior
            arrival_parameter_posterior = Vector{Float64}(undef, num_total_iterations)
        end
        if sampler.assignment_types[chain] === "random"
            random_initial_assignment!(assignments, data, sufficient_stats, model, auxillary_variables, sampler)
        else
            deterministic_initial_assignment!(assignments, data, sufficient_stats, model, auxillary_variables, sampler, chain)
        end
        log_likelihoods[1] = compute_joint_log_likelihood(sufficient_stats, model, auxillary_variables)
        if sample_arrival_parameter_posterior
            arrival_parameter_posterior[1] = sample_arrival_posterior(sufficient_stats, auxillary_variables, model.cluster_parameters)
        end
        if sample_cluster_parameter_posterior
            sample!(model, sufficient_stats, output_cluster_parameter_posterior, 1, chain)
        end
        @showprogress for iteration = 2:num_total_iterations
            assignments[:, iteration] = assignments[:, iteration-1]
            sample!(assignments, iteration, data, sufficient_stats, model, auxillary_variables, sampler)
            gibbs_sample!(auxillary_variables, sufficient_stats, model, data, assignments[:, iteration])
            log_likelihoods[iteration] = compute_joint_log_likelihood(sufficient_stats, model, auxillary_variables)
            if sample_arrival_parameter_posterior
                arrival_parameter_posterior[iteration] = sample_arrival_posterior(sufficient_stats, auxillary_variables,
                                                                                model.cluster_parameters)
            end
            if sample_cluster_parameter_posterior
                sample!(model, sufficient_stats, output_cluster_parameter_posterior, iteration, chain)
            end
        end
        output_assignments[:, :, chain] = assignments[:, :]
        output_log_likelihoods[:, chain] = log_likelihoods[:]
        if sample_arrival_parameter_posterior
            output_arrival_parameter_posterior[:, chain] = arrival_parameter_posterior[:]
        end
    end
    output_iterations = Vector((sampler.num_burn_in+1):num_total_iterations)
    output_iterations = output_iterations[output_iterations .% sampler.skip .=== 0]
    mcmc_output = Dict{String, Array}("assignments" => output_assignments[:, output_iterations, :],
                                      "log likelihood" => output_log_likelihoods[output_iterations, :])
    if sample_arrival_parameter_posterior
        mcmc_output["arrival posterior"] = output_arrival_parameter_posterior[output_iterations, :]
    end
    if sample_cluster_parameter_posterior
        mcmc_output["cluster parameter"] = output_cluster_parameter_posterior
    end
    return mcmc_output
end

function sample!(model::Mixture{C, D}, sufficient_stats, cluster_parameter_posterior, iteration, chain) where {C <: DpParameters, D}
    previous_alpha = model.cluster_parameters.alpha
    n = sufficient_stats.n

    # Propose new alpha from truncated normal centered around previous alpha
    untruncated_proposal_distribution = Normal(previous_alpha, 1)
    proposal_distribution = truncated(untruncated_proposal_distribution, 0, n)
    proposed_alpha = rand(proposal_distribution, 1)[1]
    proposal_log_weight = logpdf(proposal_distribution, proposed_alpha)

    # Compute log weight of previous alpha
    untruncated_previous_distribution = Normal(proposed_alpha, 1)
    previous_distribution = truncated(untruncated_previous_distribution, 0, n)
    previous_log_weight = logpdf(previous_distribution, previous_alpha)

    # Compute the log acceptance ratio
    model.cluster_parameters.alpha = proposed_alpha
    log_acceptance_ratio = (compute_assignment_log_likelihood(sufficient_stats, model.cluster_parameters, nothing) - proposal_log_weight)
    model.cluster_parameters.alpha = previous_alpha
    log_acceptance_ratio -= (compute_assignment_log_likelihood(sufficient_stats, model.cluster_parameters, nothing) - previous_log_weight)

    log_uniform_variate = log(rand(Uniform(0, 1), 1)[1])
    if log_uniform_variate < log_acceptance_ratio
        model.cluster_parameters.alpha = proposed_alpha
    end

    cluster_parameter_posterior[iteration, chain] = model.cluster_parameters.alpha
end

function sample!(model::Mixture{C, D}, sufficient_stats, cluster_parameter_posterior, iteration, chain) where {C <: NtlParameters, D}
    previous_parameters = model.cluster_parameters.prior
    n = sufficient_stats.n
    num_parameters = length(previous_parameters)

    # Propose new alpha from independent truncated normals centered around previous parameters
    proposed_parameters = Vector{Float64}(undef, num_parameters)
    proposal_log_weight = 0
    previous_log_weight = 0

    for i = 1:num_parameters
        previous_parameter = previous_parameters[i]
        untruncated_proposal_distribution = Normal(previous_parameter, 1)
        proposal_distribution = truncated(untruncated_proposal_distribution, 0, n)
        proposed_parameter = rand(proposal_distribution, 1)[1]
        proposed_parameters[i] = proposed_parameter
        proposal_log_weight += logpdf(proposal_distribution, proposed_parameter)

        # Compute log weight of previous parameter
        untruncated_previous_distribution = Normal(proposed_parameter, 1)
        previous_distribution = truncated(untruncated_previous_distribution, 0, n)
        previous_log_weight += logpdf(previous_distribution, previous_parameter)
    end

    # Compute the log acceptance ratio
    model.cluster_parameters.prior = proposed_parameters
    log_acceptance_ratio = (compute_assignment_log_likelihood(sufficient_stats, model.cluster_parameters, nothing) - proposal_log_weight)
    model.cluster_parameters.prior = previous_parameters
    log_acceptance_ratio -= (compute_assignment_log_likelihood(sufficient_stats, model.cluster_parameters, nothing) - previous_log_weight)

    log_uniform_variate = log(rand(Uniform(0, 1), 1)[1])
    if log_uniform_variate < log_acceptance_ratio
        model.cluster_parameters.prior = proposed_parameters
    end

    cluster_parameter_posterior[:, iteration, chain] = deepcopy(model.cluster_parameters.prior[:])
end

function sample_arrival_posterior(sufficient_stats, auxillary_variables, cluster_parameters::NtlParameters)
    posterior_parameters = compute_arrival_distribution_posterior(sufficient_stats, cluster_parameters)
    beta_distribution = Beta(posterior_parameters[1], posterior_parameters[2])
    return rand(beta_distribution, 1)[1]
end

function sample_arrival_posterior(sufficient_stats, auxillary_variables, cluster_parameters::DpParameters)
    return NaN
end

function sample!(instances, iteration, data, sufficient_stats, model, auxillary_variables, sampler::GibbsSampler)
    n = sufficient_stats.n
    for observation = 1:n
        gibbs_sample!(observation, instances, iteration, data, sufficient_stats, model, auxillary_variables)
    end
end

function propose_cluster(previous_cluster, observation, sufficient_stats, sampler::MetropolisHastingsSampler)
    clusters = findall(sufficient_stats.cluster.clusters[1:observation-1])
    #clusters = get_clusters(sufficient_stats.cluster)
    if !(previous_cluster in clusters)
        push!(clusters, previous_cluster)
    end
    if previous_cluster !== observation
        push!(clusters, observation)
    end
    sort!(clusters)
    num_clusters = length(clusters)
    center_index = findfirst(clusters .=== previous_cluster)
    start_index = maximum([1, center_index - sampler.proposal_radius])
    stop_index = minimum([num_clusters, center_index + sampler.proposal_radius])
    clusters = clusters[start_index:stop_index]
    if !(observation in clusters)
        push!(clusters, observation)
    end
    cluster = rand(clusters)
    return (cluster, -log(length(clusters)))
end

function sample!(instances, iteration, data, sufficient_stats, model, auxillary_variables, sampler::MetropolisHastingsSampler)
    n = sufficient_stats.n
    num_accepted = 0
    num_proposals = 0
    for observation = 1:n
        previous_cluster = instances[observation, iteration]
        if (observation > previous_cluster) || (observation === previous_cluster && sufficient_stats.cluster.num_observations[previous_cluster] === 1)
        #if true
            num_proposals += 1
            remove_observation!(observation, instances, iteration, data, sufficient_stats, model, auxillary_variables)

            previous_cluster_log_weight = data_log_predictive(observation, previous_cluster, sufficient_stats.data, model.data_parameters,
                                                              data, auxillary_variables)
            if previous_cluster === observation
                previous_cluster_log_weight += new_cluster_log_predictive(sufficient_stats, model.cluster_parameters,
                                                                          auxillary_variables, observation)
            else
                previous_cluster_log_weight += compute_cluster_log_predictive(observation, previous_cluster, sufficient_stats, model.cluster_parameters)
                previous_cluster_log_weight += existing_cluster_log_predictive(sufficient_stats, model.cluster_parameters,
                                                                               auxillary_variables, observation,
                                                                               previous_cluster)
            end

            # Propose a new cluster for the current observation
            (proposed_cluster, proposal_log_weight) = propose_cluster(previous_cluster, observation, sufficient_stats, sampler)
            # Compute the joint likelihood of the proposal
            proposed_cluster_log_weight = data_log_predictive(observation, proposed_cluster, sufficient_stats.data, model.data_parameters,
                                                              data, auxillary_variables)
            if proposed_cluster === observation
                proposed_cluster_log_weight += new_cluster_log_predictive(sufficient_stats, model.cluster_parameters,
                                                                          auxillary_variables, observation)
            else
                proposed_cluster_log_weight += compute_cluster_log_predictive(observation, proposed_cluster, sufficient_stats, model.cluster_parameters)
                proposed_cluster_log_weight += existing_cluster_log_predictive(sufficient_stats, model.cluster_parameters,
                                                                               auxillary_variables, observation,
                                                                               proposed_cluster)
            end

            (_, previous_cluster_proposal_log_weight) = propose_cluster(proposed_cluster, observation, sufficient_stats, sampler)

            log_acceptance_ratio = proposed_cluster_log_weight + previous_cluster_proposal_log_weight
            log_acceptance_ratio -= (previous_cluster_log_weight + proposal_log_weight)

            log_uniform_variate = log(rand(Uniform(0, 1), 1)[1])
            if log_uniform_variate < log_acceptance_ratio
                final_cluster = proposed_cluster
                num_accepted += 1
            else
                final_cluster = previous_cluster
            end
            add_observation!(observation, final_cluster, instances, iteration, data, sufficient_stats, model,
                             auxillary_variables)
        end
    end
    acceptance_rate = num_accepted/num_proposals
    if sampler.adaptive
        if acceptance_rate > 0.3
            sampler.proposal_radius += 1
        elseif acceptance_rate < 0.3
            sampler.proposal_radius = maximum([sampler.proposal_radius - 1, 1])
        end
    end
end

function gibbs_sample!(auxillary_variables::AuxillaryVariables, sufficient_stats::SufficientStatistics,
                       cluster_parameters::NtlParameters)
    nothing
end

function gibbs_sample!(auxillary_variables::AuxillaryVariables, sufficient_stats::SufficientStatistics,
                       cluster_parameters::DpParameters)
    nothing
end

function gibbs_sample!(auxillary_variables::AuxillaryVariables, sufficient_stats::SufficientStatistics,
                       data_parameters::Union{GaussianParameters, MultinomialParameters},
                       data::Matrix{T}, cluster_assignments) where {T <: Real}
    nothing
end

function sample_topic(sufficient_stats, data_parameters, cluster, token)
    topic_parameter = data_parameters.topic_parameter
    word_parameter = data_parameters.word_parameter
    num_topics = data_parameters.num_topics
    cluster_topic_frequencies = sufficient_stats.data.cluster_topic_frequencies
    topic_token_frequencies = sufficient_stats.data.topic_token_frequencies
    topics = Vector(1:num_topics)
    topic_log_weights = Array{Float64}(undef, num_topics)
    num_tokens = size(topic_token_frequencies)[1]
    for topic = topics
        cluster_log_weight = log(cluster_topic_frequencies[topic, cluster] + topic_parameter)
        cluster_log_weight -= log(sum(cluster_topic_frequencies[:, cluster]) + num_topics*topic_parameter)
        token_log_weight = log(topic_token_frequencies[token, topic] + word_parameter)
        token_log_weight -= log(sum(topic_token_frequencies[:, topic]) + num_tokens*word_parameter)
        topic_log_weights[topic] = cluster_log_weight + token_log_weight
    end
    return gumbel_max(topics, topic_log_weights)
end

function update_token_topic_statistics!(document, token, topic, cluster, auxillary_variables, sufficient_stats;
                                        update_type="add")
    if update_type === "add"
        sufficient_stats.data.cluster_topic_frequencies[topic, cluster] += 1
        sufficient_stats.data.topic_token_frequencies[token, topic] += 1
        auxillary_variables.data.document_topic_frequencies[topic, document] += 1
    elseif update_type === "remove"
        sufficient_stats.data.cluster_topic_frequencies[topic, cluster] -= 1
        sufficient_stats.data.topic_token_frequencies[token, topic] -= 1
        auxillary_variables.data.document_topic_frequencies[topic, document] -= 1
    end
end

function gibbs_sample!(auxillary_variables, sufficient_stats::SufficientStatistics,
                       data_parameters::FiniteTopicModelParameters, data::Matrix{T},
                       cluster_assignments::Vector{Int64}) where {T <: Real}
    num_documents = length(cluster_assignments)
    num_tokens = size(data)[1]
    for document = 1:num_documents
        cluster = cluster_assignments[document]
        for token = 1:num_tokens
            num_instances = data[token, document]
            for instance = 1:num_instances
                previous_topic = auxillary_variables.data.token_topic_assignments[document][token][instance]
                # update sufficient statistics
                update_token_topic_statistics!(document, token, previous_topic, cluster, auxillary_variables, sufficient_stats,
                                               update_type="remove")
                # sample new topic
                (new_topic, _) = sample_topic(sufficient_stats, data_parameters, cluster, token)
                auxillary_variables.data.token_topic_assignments[document][token][instance] = new_topic
                # update sufficient statistics
                update_token_topic_statistics!(document, token, new_topic, cluster, auxillary_variables, sufficient_stats,
                                               update_type="add")
            end
        end
    end
end

function gibbs_sample!(auxillary_variables::AuxillaryVariables, sufficient_stats::SufficientStatistics, model::Model, data,
                       cluster_assignments::Vector{Int64})
    gibbs_sample!(auxillary_variables, sufficient_stats, model.cluster_parameters)
    gibbs_sample!(auxillary_variables, sufficient_stats, model.data_parameters, data, cluster_assignments)
end

function gibbs_sample!(observation, instances, iteration, data, sufficient_stats, model::Mixture, auxillary_variables)
    remove_observation!(observation, instances, iteration, data, sufficient_stats, model, auxillary_variables)
    (cluster, _) = gibbs_proposal(observation, data, sufficient_stats, model, auxillary_variables)
    add_observation!(observation, cluster, instances, iteration, data, sufficient_stats, model, auxillary_variables)
end

function gibbs_proposal(observation::Int64, data::Matrix{T}, sufficient_stats::SufficientStatistics,
                        model::Mixture, auxillary_variables) where {T<:Real}
    data_parameters = model.data_parameters
    cluster_parameters = model.cluster_parameters
    clusters = get_clusters(sufficient_stats.cluster)
    num_clusters = length(clusters)

    choices = Array{Int64}(undef, num_clusters+1)
    choices[1:num_clusters] = clusters
    choices[num_clusters+1] = observation

    weights = Array{Float64}(undef, num_clusters+1)
    weights[1:num_clusters] = compute_existing_cluster_log_predictives(observation, clusters,
                                                                       sufficient_stats, cluster_parameters,
                                                                       auxillary_variables)
    weights[num_clusters+1] = new_cluster_log_predictive(sufficient_stats, cluster_parameters, auxillary_variables,
                                                         observation)
    weights += compute_data_log_predictives(observation, choices, data, sufficient_stats.data, data_parameters, auxillary_variables)

    return gumbel_max(choices, weights)
end

function propose(observation, data, particles, particle, sufficient_stats, model::Mixture, auxillary_variables)
    return gibbs_proposal(observation, data, sufficient_stats, model, auxillary_variables)
end

function prepare_arrivals_auxillary_variables(::DpParameters, data::Matrix{T}) where
                                             {T <: Real}
    return DpAuxillaryVariables()
end

function prepare_arrivals_auxillary_variables(::NtlParameters, data::Matrix{T}) where
                                             {T <: Real}
    return GeometricAuxillaryVariables()
end

function prepare_data_auxillary_variables(::GaussianParameters, data::Matrix{T}) where {T <: Real}
    return GaussianAuxillaryVariables()
end

function prepare_data_auxillary_variables(::MultinomialParameters, data::Matrix{T}) where {T <: Real}
    return MultinomialAuxillaryVariables()
end

function prepare_data_auxillary_variables(data_parameters::FiniteTopicModelParameters, data::Matrix{T}) where {T <: Real}
    num_topics = data_parameters.num_topics
    num_documents = size(data)[2]
    num_words = size(data)[1]
    #                              Vector(Document => Dict(word => Dict(token => assignment)))
    word_topic_assignments = Vector{Dict{Int64, Dict{Int64, Int64}}}(undef, num_documents)
    for document = 1:num_documents
        word_dicts = []
        for word = 1:num_words
            token_assignments = [(token, 0) for token in 1:data[word, document]]
            token_assignments = Dict{Int64, Int64}(token_assignments)
            push!(word_dicts, (word, token_assignments))
        end
        document_dict = Dict{Int64, Dict{Int64, Int64}}(word_dicts)
        word_topic_assignments[document] = document_dict
    end
    document_topic_frequencies = zeros(Int64, num_topics, num_documents)
    data_aux_variables = FiniteTopicModelAuxillaryVariables(word_topic_assignments, document_topic_frequencies)
    return data_aux_variables
end

function prepare_auxillary_variables(model, data)
    arrival_aux_variables = prepare_arrivals_auxillary_variables(model.cluster_parameters, data)
    data_aux_variables = prepare_data_auxillary_variables(model.data_parameters, data)
    auxillary_variables = AuxillaryVariables(arrival_aux_variables, data_aux_variables)
    return auxillary_variables
end

function prepare_sufficient_statistics(model, data)
    n = size(data)[2]
    num_assigned = 0
    cluster_sufficient_stats = prepare_cluster_sufficient_statistics(model, n)
    data_sufficient_stats = prepare_data_sufficient_statistics(data, model.data_parameters)
    return SufficientStatistics(n, num_assigned, cluster_sufficient_stats, data_sufficient_stats)
end

function prepare_data_sufficient_statistics(data::Matrix{Float64}, data_parameters::GaussianParameters)
    dim = size(data)[1]
    n = size(data)[2]
    data_means = zeros(Float64, dim, n+1)
    data_precision_quadratic_sums = zeros(Float64, n+1)
    posterior_means = Matrix{Float64}(undef, dim, n+1)
    for i = 1:(n+1)
        posterior_means[:, i] = data_parameters.prior_mean
    end
    posterior_covs = Matrix{Float64}(undef, dim*dim, n+1)
    for i = 1:(n+1)
        posterior_covs[:, i] = vec(data_parameters.prior_covariance)
    end
    return GaussianSufficientStatistics(data_means, data_precision_quadratic_sums, posterior_means, posterior_covs)
end

function prepare_data_sufficient_statistics(data::Matrix{Int64}, data_parameters::MultinomialParameters)
    dim = size(data)[1]
    n = size(data)[2]
    counts = zeros(Int64, dim, n+1)
    posterior_dirichlet_scale = Matrix{Float64}(undef, dim, n+1)
    for i = 1:(n+1)
        posterior_dirichlet_scale[:, i] = vec(data_parameters.prior_dirichlet_scale)
    end
    return MultinomialSufficientStatistics(counts, posterior_dirichlet_scale)
end

function prepare_data_sufficient_statistics(data::Matrix{Int64}, data_parameters::FiniteTopicModelParameters)
    num_topics = data_parameters.num_topics
    num_tokens = size(data)[1]
    num_documents = size(data)[2]
    cluster_topic_frequencies = zeros(Int64, num_topics, num_documents)
    cluster_topic_posterior = zeros(Int64, num_topics, num_documents)
    topic_token_frequencies = zeros(Int64, num_tokens, num_topics)
    topic_token_posterior = zeros(Int64, num_tokens, num_topics)
    return FiniteTopicModelSufficientStatistics(cluster_topic_frequencies, cluster_topic_posterior,
                                                topic_token_frequencies, topic_token_posterior)
end

function prepare_cluster_sufficient_statistics(::Mixture, n::Int64)
    cluster_num_observations = Vector{Int64}(zeros(Int64, n))
    cumulative_num_observations = FenwickTree{Int64}(n)
    clusters = BitArray(undef, n)
    clusters .= false
    cluster_sufficient_stats = MixtureSufficientStatistics(cluster_num_observations, cumulative_num_observations, clusters)
    return cluster_sufficient_stats
end

function compute_cluster_data_log_likelihood(cluster::Int64, sufficient_stats::SufficientStatistics, data_parameters::FiniteTopicModelParameters,
                                             auxillary_variables)
    dirichlet_posterior = vec(sufficient_stats.data.cluster_topic_posterior[:, cluster])
    num_topics = length(dirichlet_posterior)
    dirichlet_prior = data_parameters.topic_parameter*ones(Float64, num_topics)
    counts = convert(Vector{Int64}, round.(dirichlet_posterior - dirichlet_prior))
    log_likelihood = logmvbeta(dirichlet_posterior) - logmvbeta(dirichlet_prior)
    return log_likelihood
end

function compute_cluster_data_log_likelihood(cluster::Int64, sufficient_stats::SufficientStatistics,
                                             data_parameters::MultinomialParameters, auxillary_variables)
    dirichlet_posterior = sufficient_stats.data.posterior_dirichlet_scale[:, cluster]
    dirichlet_prior = data_parameters.prior_dirichlet_scale
    counts = convert(Vector{Int64}, round.(dirichlet_posterior - dirichlet_prior))
    log_likelihood = logmvbeta(dirichlet_posterior) - logmvbeta(dirichlet_prior)
    return log_likelihood
end

function compute_cluster_data_log_likelihood(cluster::Int64, sufficient_stats::SufficientStatistics,
                                             data_parameters::GaussianParameters, auxillary_variables)
    posterior_mean = vec(sufficient_stats.data.posterior_means[:, cluster])
    dim = length(posterior_mean)
    posterior_covariance = reshape(sufficient_stats.data.posterior_covs[:, cluster], dim, dim)
    prior_mean = data_parameters.prior_mean
    prior_covariance = data_parameters.prior_covariance
    prior_precision = inv(prior_covariance)

    data_precision_quadratic_sum = sufficient_stats.data.data_precision_quadratic_sums[cluster]
    posterior_quadratic = transpose(posterior_mean) * posterior_covariance * posterior_mean
    prior_quadratic = transpose(prior_mean) * prior_precision * prior_mean

    log_likelihood = -(1/2)*(data_precision_quadratic_sum + prior_quadratic - posterior_quadratic)

    data_log_determinant = log(abs(det(data_parameters.data_covariance)))
    prior_log_determinant = log(abs(det(prior_covariance)))
    posterior_log_determinant = log(abs(det(posterior_covariance)))

    n = sufficient_stats.cluster.num_observations[cluster]
    log_likelihood += log(2*pi)*dim*(n + 1)/2 - (n/2)data_log_determinant - (1/2)prior_log_determinant + (1/2)posterior_log_determinant
    return log_likelihood
end

function compute_auxillary_data_log_likelihood(sufficient_stats, data_parameters::Union{GaussianParameters, MultinomialParameters},
                                               auxillary_variables)
    return 0.
end

function compute_auxillary_data_log_likelihood(sufficient_stats, data_parameters::FiniteTopicModelParameters, auxillary_variables)
    log_likelihood = 0.
    num_words = data_parameters.num_words
    dirichlet_prior = data_parameters.word_parameter*ones(Int64, num_words)
    num_topics = data_parameters.num_topics
    for topic = 1:num_topics
        dirichlet_posterior = vec(sufficient_stats.data.topic_token_posterior[:, topic])
        topic_log_likelihood = logmvbeta(dirichlet_posterior) - logmvbeta(dirichlet_prior)
        log_likelihood += topic_log_likelihood
    end
    return log_likelihood
end

function compute_data_log_likelihood(sufficient_stats::SufficientStatistics, data_parameters::DataParameters, auxillary_variables)
    clusters = get_clusters(sufficient_stats.cluster)
    log_likelihood = 0
    for cluster = clusters
        cluster_data_log_likelihood = compute_cluster_data_log_likelihood(cluster, sufficient_stats, data_parameters, auxillary_variables)
        log_likelihood += cluster_data_log_likelihood
    end
    aux_log_likelihood = compute_auxillary_data_log_likelihood(sufficient_stats, data_parameters, auxillary_variables)
    log_likelihood += aux_log_likelihood
    return log_likelihood
end

function compute_cluster_assignment_log_likelihood(cluster::Int64, cluster_sufficient_stats::MixtureSufficientStatistics,
                                                   cluster_parameters::NtlParameters)
    if cluster > 1
        psi_posterior = compute_stick_breaking_posterior(cluster, cluster_sufficient_stats, cluster_parameters.prior)
        return logbeta(psi_posterior) - logbeta(cluster_parameters.prior)
    else
        return 0
    end
end

function compute_arrivals_log_likelihood(sufficient_stats::SufficientStatistics, cluster_parameters::NtlParameters,
                                         auxillary_variables)
    phi_posterior = compute_arrival_distribution_posterior(sufficient_stats, cluster_parameters)
    log_likelihood = logbeta(phi_posterior) - logbeta(cluster_parameters.arrival_distribution.prior)
    @assert (log_likelihood < Inf)
    return log_likelihood
end

function compute_assignment_log_likelihood(sufficient_stats::SufficientStatistics{C, D},
                                           cluster_parameters::DpParameters,
                                           auxillary_variables::Union{AuxillaryVariables, Nothing}) where
                                           {C <: MixtureSufficientStatistics, D}
    num_observations = sufficient_stats.cluster.num_observations
    num_observations = num_observations[num_observations .> 0]
    num_clusters = length(num_observations)
    alpha = cluster_parameters.alpha
    log_likelihood = sum(loggamma.(num_observations)) + num_clusters*log(alpha)
    n = sufficient_stats.n
    log_likelihood += loggamma(alpha) - loggamma(n + alpha)
    return log_likelihood
end

function compute_assignment_log_likelihood(sufficient_stats::SufficientStatistics{C, D},
                                           cluster_parameters::NtlParameters,
                                           auxillary_variables) where
                                           {C <: MixtureSufficientStatistics, D}
    log_likelihood = 0
    clusters = get_clusters(sufficient_stats.cluster)
    for cluster = clusters
        log_likelihood += compute_cluster_assignment_log_likelihood(cluster, sufficient_stats.cluster,
                                                                    cluster_parameters)
        @assert (log_likelihood < Inf)
    end
    log_likelihood += compute_arrivals_log_likelihood(sufficient_stats, cluster_parameters, auxillary_variables)
    @assert (log_likelihood < Inf)
    return log_likelihood
end

function compute_joint_log_likelihood(sufficient_stats::SufficientStatistics, model, auxillary_variables)
    data_parameters = model.data_parameters
    cluster_parameters = model.cluster_parameters
    data_log_likelihood = compute_data_log_likelihood(sufficient_stats, data_parameters, auxillary_variables)
    @assert (!isnan(data_log_likelihood) & (data_log_likelihood < Inf))
    assignment_log_likelihood = compute_assignment_log_likelihood(sufficient_stats, cluster_parameters, auxillary_variables)
    @assert (!isnan(assignment_log_likelihood) & (assignment_log_likelihood < Inf))
    log_likelihood = data_log_likelihood + assignment_log_likelihood
    return log_likelihood
end

function update_data_sufficient_statistics!(sufficient_stats::SufficientStatistics{C,D}, cluster::Int64,
                                            observation::Int64, model::Model, data, auxillary_variables;
                                            update_type::String="add") where {C, D<:GaussianSufficientStatistics}
    datum = vec(data[:, observation])
    cluster_mean = vec(sufficient_stats.data.data_means[:, cluster])
    data_precision_quadratic_sum = sufficient_stats.data.data_precision_quadratic_sums[cluster]
    n = sufficient_stats.cluster.num_observations[cluster]
    datum_precision_quadratic_sum = transpose(datum)*model.data_parameters.data_precision*datum
    if update_type == "add"
        cluster_mean = (cluster_mean.*(n-1) + datum)./n
        data_precision_quadratic_sum += datum_precision_quadratic_sum
    elseif update_type == "remove"
        if n > 0
            cluster_mean = (cluster_mean.*(n+1) - datum)./n
            data_precision_quadratic_sum -= datum_precision_quadratic_sum
        else
            cluster_mean .= 0
            data_precision_quadratic_sum = 0
        end
    else
        message = "$update_type is not a supported update type."
        throw(ArgumentError(message))
    end
    sufficient_stats.data.data_means[:, cluster] = cluster_mean
    sufficient_stats.data.data_precision_quadratic_sums[cluster] = data_precision_quadratic_sum
end

function update_data_sufficient_statistics!(sufficient_stats::SufficientStatistics{C, D}, cluster::Int64,
                                            observation::Int64, model::Model, data, auxillary_variables;
                                            update_type::String="add") where {C, D <: FiniteTopicModelSufficientStatistics}
    doc_topic_frequencies = vec(auxillary_variables.data.document_topic_frequencies[:, observation])
    if update_type === "add"
        sufficient_stats.data.cluster_topic_frequencies[:, cluster] += doc_topic_frequencies
    elseif update_type === "remove"
        sufficient_stats.data.cluster_topic_frequencies[:, cluster] -= doc_topic_frequencies
    else
        message = "$update_type is not a supported update type."
        throw(ArgumentError(message))
    end
end

function update_data_sufficient_statistics!(sufficient_stats::SufficientStatistics{C,D}, cluster::Int64,
                                            observation::Int64, ::Model, data, auxillary_variables;
                                            update_type::String="add") where {C, D<:MultinomialSufficientStatistics}
    datum = vec(data[:, observation])
    cluster_counts = vec(sufficient_stats.data.total_counts[:, cluster])
    if update_type === "add"
        cluster_counts += datum
    elseif update_type === "remove"
        cluster_counts -= datum
    else
        message = "$update_type is not a supported update type."
        throw(ArgumentError(message))
    end
    sufficient_stats.data.total_counts[:, cluster] = cluster_counts
end

function update_cluster_sufficient_statistics!(sufficient_stats::SufficientStatistics{C, D}, cluster::Int64;
                                               update_type::String="add") where {C <: MixtureSufficientStatistics, D}
    if update_type === "add"
        sufficient_stats.cluster.num_observations[cluster] += 1
        inc!(sufficient_stats.cluster.cumulative_num_observations, cluster, 1)
        if sufficient_stats.cluster.num_observations[cluster] === 1
            sufficient_stats.cluster.clusters[cluster] = true
        end
        sufficient_stats.num_assigned += 1
    elseif update_type === "remove"
        sufficient_stats.cluster.num_observations[cluster] -= 1
        dec!(sufficient_stats.cluster.cumulative_num_observations, cluster, 1)
        if sufficient_stats.cluster.num_observations[cluster] === 0
            sufficient_stats.cluster.clusters[cluster] = false
        end
        sufficient_stats.num_assigned -= 1
    else
        message = "$update_type is not a supported update type."
        throw(ArgumentError(message))
    end
end

function update_cluster_stats_new_birth_time!(num_observations::Vector{Int64}, new_time::Int64, old_time::Int64)
    num_observations[new_time] = num_observations[old_time]
    num_observations[old_time] = 0
end

function update_cluster_stats_new_birth_time!(cluster_sufficient_stats::MixtureSufficientStatistics,
                                              new_time::Int64, old_time::Int64)
    num_observations = cluster_sufficient_stats.num_observations[old_time]
    update_cluster_stats_new_birth_time!(cluster_sufficient_stats.num_observations, new_time, old_time)
    cluster_sufficient_stats.clusters[new_time] = true
    cluster_sufficient_stats.clusters[old_time] = false
    dec!(cluster_sufficient_stats.cumulative_num_observations, old_time, num_observations)
    inc!(cluster_sufficient_stats.cumulative_num_observations, new_time, num_observations)
end

function update_data_stats_new_birth_time!(data_sufficient_stats::GaussianSufficientStatistics,
                                           new_time::Int64, old_time::Int64, data_parameters::GaussianParameters)
    data_sufficient_stats.data_means[:, new_time] = data_sufficient_stats.data_means[:, old_time]
    data_sufficient_stats.data_means[:, old_time] .= 0
    data_sufficient_stats.posterior_means[:, new_time] = data_sufficient_stats.posterior_means[:, old_time]
    data_sufficient_stats.posterior_covs[:, new_time] = data_sufficient_stats.posterior_covs[:, old_time]
    data_sufficient_stats.posterior_means[:, old_time] = data_parameters.prior_mean
    data_sufficient_stats.posterior_covs[:, old_time] = vec(data_parameters.prior_covariance)
    data_sufficient_stats.data_precision_quadratic_sums[new_time] = data_sufficient_stats.data_precision_quadratic_sums[old_time]
    data_sufficient_stats.data_precision_quadratic_sums[old_time] = 0
end

function update_data_stats_new_birth_time!(data_sufficient_stats::MultinomialSufficientStatistics,
                                           new_time::Int64, old_time::Int64, data_parameters::MultinomialParameters)
    data_sufficient_stats.total_counts[:, new_time] = data_sufficient_stats.total_counts[:, old_time]
    data_sufficient_stats.total_counts[:, old_time] .= 0
    data_sufficient_stats.posterior_dirichlet_scale[:, new_time] = data_sufficient_stats.posterior_dirichlet_scale[:, old_time]
    data_sufficient_stats.posterior_dirichlet_scale[:, old_time] = data_parameters.prior_dirichlet_scale
end

function update_data_stats_new_birth_time!(data_sufficient_stats::FiniteTopicModelSufficientStatistics,
                                           new_time::Int64, old_time::Int64, data_parameters::FiniteTopicModelParameters)
    data_sufficient_stats.cluster_topic_frequencies[:, new_time] = data_sufficient_stats.cluster_topic_frequencies[:, old_time]
    data_sufficient_stats.cluster_topic_frequencies[:, old_time] .= 0
    data_sufficient_stats.cluster_topic_posterior[:, new_time] = data_sufficient_stats.cluster_topic_posterior[:, old_time]
    data_sufficient_stats.cluster_topic_posterior[:, old_time] .= 0
end

function update_cluster_birth_time_remove!(cluster::Int64, iteration::Int64, instances::Matrix{Int64},
                                           sufficient_stats::SufficientStatistics,
                                           data_parameters::DataParameters)
    assigned_to_cluster = instances[:, iteration] .=== cluster
    new_cluster_time = findfirst(assigned_to_cluster)
    instances[assigned_to_cluster, iteration] .= new_cluster_time
    update_cluster_stats_new_birth_time!(sufficient_stats.cluster, new_cluster_time, cluster)
    update_data_stats_new_birth_time!(sufficient_stats.data, new_cluster_time, cluster, data_parameters)
    return new_cluster_time
end

function remove_observation!(observation::Int64, instances::Matrix{Int64}, iteration::Int64, data::Matrix{T},
                             sufficient_stats::SufficientStatistics, model::Mixture,
                             auxillary_variables::AuxillaryVariables) where {T <: Real}
    n = sufficient_stats.n
    cluster = instances[observation, iteration]

    instances[observation, iteration] = n+1
    update_cluster_sufficient_statistics!(sufficient_stats, cluster, update_type="remove")
    update_data_sufficient_statistics!(sufficient_stats, cluster, observation, model, data, auxillary_variables,
                                       update_type="remove")
    compute_data_posterior_parameters!(cluster, sufficient_stats, model.data_parameters)
    if (cluster === observation) && (sufficient_stats.cluster.num_observations[cluster] > 0)
        cluster = update_cluster_birth_time_remove!(cluster, iteration, instances, sufficient_stats,
                                                    model.data_parameters)
    end
end

function update_cluster_birth_time_add!(cluster::Int64, observation::Int64, iteration::Int64, instances::Matrix{Int64},
                                        sufficient_stats::SufficientStatistics, data_parameters)
    assigned_to_cluster = (instances[:, iteration] .=== cluster)
    old_cluster_time = cluster
    cluster = observation
    instances[assigned_to_cluster, iteration] .= cluster
    update_cluster_stats_new_birth_time!(sufficient_stats.cluster, cluster, old_cluster_time)
    update_data_stats_new_birth_time!(sufficient_stats.data, cluster, old_cluster_time, data_parameters)
    return cluster
end

function add_observation!(observation::Int64, cluster::Int64, instances::Matrix{Int64}, iteration::Int64,
                          data::Matrix{T}, sufficient_stats::SufficientStatistics, model::Mixture,
                          auxillary_variables::AuxillaryVariables) where
                          {T <: Real}
    if observation < cluster
        cluster = update_cluster_birth_time_add!(cluster, observation, iteration, instances, sufficient_stats,
                                                 model.data_parameters)
    end
    instances[observation, iteration] = cluster
    update_cluster_sufficient_statistics!(sufficient_stats, cluster, update_type="add")
    update_data_sufficient_statistics!(sufficient_stats, cluster, observation, model, data, auxillary_variables,
                                       update_type="add")
    compute_data_posterior_parameters!(cluster, sufficient_stats, model.data_parameters)
end

function compute_arrival_distribution_posterior(sufficient_stats::SufficientStatistics{C, D},
                                                cluster_parameters::NtlParameters) where
                                                {C <: Union{MixtureSufficientStatistics, HmmSufficientStatistics}, D}
    num_assigned = sufficient_stats.num_assigned
    num_clusters = length(get_clusters(sufficient_stats.cluster))
    @assert num_assigned >= num_clusters
    phi_posterior = deepcopy(cluster_parameters.arrival_distribution.prior)
    phi_posterior[1] += num_clusters - 1
    phi_posterior[2] += num_assigned - num_clusters
    return phi_posterior
end

### Mixture cluster log predictives

function new_cluster_log_predictive(sufficient_stats::SufficientStatistics{C, D}, cluster_parameters::DpParameters,
                                    auxillary_variables::AuxillaryVariables, ::Int64) where {C <: MixtureSufficientStatistics, D}
    return log(cluster_parameters.alpha)
end

function new_cluster_log_predictive(sufficient_stats::SufficientStatistics{C, D}, cluster_parameters::NtlParameters,
                                    auxillary_variables, observation::Int64) where
                                    {C <: MixtureSufficientStatistics, D}
    num_complement = compute_num_complement(observation, sufficient_stats.cluster)
    psi_posterior = deepcopy(cluster_parameters.prior)
    psi_posterior[2] += num_complement
    phi_posterior = compute_arrival_distribution_posterior(sufficient_stats, cluster_parameters)
    return log(phi_posterior[1]) - log(sum(phi_posterior)) + logbeta(psi_posterior) - logbeta(cluster_parameters.prior)
end

function existing_cluster_log_predictive(sufficient_stats::SufficientStatistics{C, D},
                                         cluster_parameters::NtlParameters,
                                         auxillary_variables, observation::Int64, cluster::Int64) where
                                         {C <: MixtureSufficientStatistics, D}
    phi_posterior = compute_arrival_distribution_posterior(sufficient_stats, cluster_parameters)
    return log(phi_posterior[2]) - log(sum(phi_posterior))
end

### End mixture cluster log predictives

function compute_num_complement(cluster::Int64, cluster_suff_stats::MixtureSufficientStatistics;
                                missing_observation::Union{Int64, Nothing}=nothing)
    sum_of_older_clusters = sum(cluster_suff_stats.num_observations[1:(cluster-1)])
    num_complement = sum_of_older_clusters - (cluster - 1)
    if missing_observation !== nothing && missing_observation < cluster
        num_complement += 1
    end
    return num_complement
end

function compute_stick_breaking_posterior(cluster::Int64, cluster_suff_stats::MixtureSufficientStatistics, beta_prior::Vector{T};
                                          missing_observation::Union{Int64, Nothing}=nothing) where {T <: Real}
    posterior = deepcopy(beta_prior)
    num_obs = cluster_suff_stats.num_observations[cluster]
    @assert num_obs >= 0
    if num_obs > 0
        posterior[1] += cluster_suff_stats.num_observations[cluster] - 1
    end
    posterior[2] += compute_num_complement(cluster, cluster_suff_stats, missing_observation=missing_observation)
    @assert posterior[1] > 0
    @assert posterior[2] > 0
    return posterior
end

function compute_stick_breaking_posterior(cluster::Int64, sufficient_stats::SufficientStatistics,
                                          ntl_parameters::NtlParameters;
                                          missing_observation::Union{Int64, Nothing}=nothing)
    posterior = compute_stick_breaking_posterior(cluster, sufficient_stats.cluster,
                                                 ntl_parameters.prior, missing_observation=missing_observation)
    return posterior
end

function compute_cluster_log_predictive(observation, cluster, sufficient_stats, cluster_parameters::NtlParameters;
                                        psi_parameters::Matrix=Array{Int64}(undef, 2, 0),
                                        logbetas::Vector=Array{Float64}(undef, 0))
    log_predictive = 0
    # Clusters younger than the current cluster
    clusters = get_clusters(sufficient_stats.cluster)

    if cluster < observation
        if cluster > 1
            if size(psi_parameters)[2] === 0
                cluster_psi = compute_stick_breaking_posterior(cluster, sufficient_stats, cluster_parameters,
                                                               missing_observation=observation)
            else
                cluster_psi = vec(psi_parameters[:, cluster])
            end
            log_predictive += log(cluster_psi[1]) - log(sum(cluster_psi))
        end

        # clusters younger than the current cluster
        younger_clusters_range = Vector{Int64}((cluster+1):(observation-1))
        younger_clusters_mask = clusters .∈ Ref(younger_clusters_range)
        younger_clusters = clusters[younger_clusters_mask]
        for younger_cluster = younger_clusters
            if size(psi_parameters)[2] === 0
                younger_cluster_psi = compute_stick_breaking_posterior(younger_cluster, sufficient_stats, cluster_parameters,
                                                                       missing_observation=observation)
            else
                younger_cluster_psi = vec(psi_parameters[:, younger_cluster])
            end
            log_predictive += log(younger_cluster_psi[2]) - log(sum(younger_cluster_psi))
        end
    else # observation < cluster
        cluster_num_obs = sufficient_stats.cluster.num_observations[cluster]
        cluster_old_psi = compute_stick_breaking_posterior(cluster, sufficient_stats, cluster_parameters,
                                                           missing_observation=observation)
        younger_clusters_log_predictive = 0
        # clusters younger than the observation
        younger_clusters_range = Vector{Int64}((observation+1):(cluster-1))
        younger_clusters_mask = clusters .∈ Ref(younger_clusters_range)
        younger_clusters = clusters[younger_clusters_mask]

        for younger_cluster = younger_clusters
            if size(psi_parameters)[2] === 0
                younger_cluster_old_psi = compute_stick_breaking_posterior(younger_cluster, sufficient_stats, cluster_parameters,
                                                                           missing_observation=observation)
            else
                younger_cluster_old_psi = vec(psi_parameters[:, younger_cluster])
            end
            younger_cluster_new_psi = deepcopy(younger_cluster_old_psi)
            younger_cluster_new_psi[2] += cluster_num_obs
            if length(logbetas) === 0
                younger_clusters_log_predictive += logbeta(younger_cluster_new_psi) - logbeta(younger_cluster_old_psi)
            else
                younger_clusters_log_predictive += logbeta(younger_cluster_new_psi) - logbetas[younger_cluster]
            end
        end

        log_predictive = younger_clusters_log_predictive

        if observation > 1
            new_num_complement = compute_num_complement(observation, sufficient_stats.cluster, missing_observation=observation)
            cluster_new_psi = deepcopy(cluster_old_psi)
            cluster_new_psi[1] += 1
            cluster_new_psi[2] = new_num_complement + cluster_parameters.prior[2]
        else
            cluster_new_psi = [1, 1]
        end

        if length(logbetas) === 0
            cluster_log_predictive = logbeta(cluster_new_psi) - logbeta(cluster_old_psi)
        else
            cluster_log_predictive = logbeta(cluster_new_psi) - logbetas[cluster]
        end

        log_predictive += cluster_log_predictive
    end
    return log_predictive
end

function compute_cluster_log_predictives(observation::Int64, clusters::Vector{Int64},
                                         sufficient_stats, ntl_parameters::NtlParameters)
    # Not strictly the number of observations
    n = maximum(clusters)
    n = maximum([n, observation])

    clusters = get_clusters(sufficient_stats.cluster)
    num_clusters = length(clusters)
    log_weights = Array{Float64}(undef, num_clusters)
    psi_parameters = zeros(Float64, 2, n)
    logbetas = zeros(Float64, n)

    for cluster = clusters
        if cluster > 1
            cluster_psi = compute_stick_breaking_posterior(cluster, sufficient_stats, ntl_parameters,
                                                           missing_observation=observation)
            psi_parameters[:, cluster] = cluster_psi
            logbetas[cluster] = logbeta(cluster_psi)
        end
    end

    for (i, cluster) = enumerate(clusters)
        log_weights[i] = compute_cluster_log_predictive(observation, cluster, sufficient_stats, ntl_parameters,
                                                        psi_parameters=psi_parameters, logbetas=logbetas)
    end
    return log_weights
end

function compute_existing_cluster_log_predictives(observation::Int64, clusters::Vector{Int64},
                                                  sufficient_stats::SufficientStatistics,
                                                  cluster_parameters::NtlParameters,
                                                  auxillary_variables)
    log_weights = compute_cluster_log_predictives(observation, clusters, sufficient_stats, cluster_parameters)
    for (i, cluster) in enumerate(clusters)
        log_weights[i] += existing_cluster_log_predictive(sufficient_stats, cluster_parameters, auxillary_variables, observation, cluster)
    end
    return log_weights
end

function compute_existing_cluster_log_predictives(::Int64, clusters::Vector{Int64}, sufficient_stats, ::DpParameters,
                                                  auxillary_variables)
    num_observations = deepcopy(sufficient_stats.cluster.num_observations[clusters])
    num_observations[num_observations .< 1] .= 1
    log_weights = log.(num_observations)
    return log_weights
end

function compute_data_posterior_parameters!(cluster::Int64, sufficient_stats::SufficientStatistics,
                                            data_parameters::GaussianParameters)
    data_mean = vec(sufficient_stats.data.data_means[:, cluster])
    n = sufficient_stats.cluster.num_observations[cluster]
    prior_mean = data_parameters.prior_mean
    prior_cov = data_parameters.prior_covariance
    prior_precision = inv(prior_cov)
    if n > 0
        data_precision = inv(data_parameters.data_covariance)
        cov = inv(prior_precision + n*data_precision)
        posterior_mean = cov*(prior_cov*prior_mean + n*data_precision*data_mean)
    else
        posterior_mean = prior_mean
        cov = prior_cov
    end
    sufficient_stats.data.posterior_means[:, cluster] = vec(posterior_mean)
    sufficient_stats.data.posterior_covs[:, cluster] = vec(cov)
end

function compute_data_posterior_parameters!(cluster::Int64, sufficient_stats, data_parameters::MultinomialParameters)
    cluster_counts = sufficient_stats.data.total_counts[:, cluster]
    posterior_parameters = (cluster_counts + data_parameters.prior_dirichlet_scale)
    sufficient_stats.data.posterior_dirichlet_scale[:, cluster] = posterior_parameters
end

function compute_data_posterior_parameters!(cluster::Int64, sufficient_stats, data_parameters::FiniteTopicModelParameters)
    sufficient_stats.data.cluster_topic_posterior[:, :] = sufficient_stats.data.cluster_topic_frequencies .+ data_parameters.topic_parameter
    sufficient_stats.data.topic_token_posterior[:, :] = sufficient_stats.data.topic_token_frequencies .+ data_parameters.word_parameter
end

function data_log_predictive(observation::Int64, cluster::Int64,
                             data_sufficient_stats::FiniteTopicModelSufficientStatistics,
                             data_parameters::FiniteTopicModelParameters, data, auxillary_variables)
    document_topic_frequencies = vec(auxillary_variables.data.document_topic_frequencies[:, observation])
    cluster_topic_posterior = vec(data_sufficient_stats.cluster_topic_posterior[:, cluster])
    n = sum(document_topic_frequencies)
    posterior = DirichletMultinomial(n, cluster_topic_posterior)
    return logpdf(posterior, document_topic_frequencies)
end

function data_log_predictive(observation::Int64, cluster::Int64,
                             data_sufficient_stats::GaussianSufficientStatistics,
                             data_parameters::GaussianParameters, data, auxillary_variables)
    datum = vec(data[:, observation])
    posterior_mean = vec(data_sufficient_stats.posterior_means[:, cluster])
    dim = length(posterior_mean)
    posterior_cov = reshape(data_sufficient_stats.posterior_covs[:, cluster], dim, dim)
    posterior = MvNormal(posterior_mean, data_parameters.data_covariance + posterior_cov)
    return logpdf(posterior, datum)
end

function data_log_predictive(observation::Int64, cluster::Int64,
                             data_sufficient_stats::MultinomialSufficientStatistics,
                             ::MultinomialParameters, data, auxillary_variables)
    datum = vec(data[:, observation])
    posterior_scale = vec(data_sufficient_stats.posterior_dirichlet_scale[:, cluster])
    n = sum(datum)
    posterior = DirichletMultinomial(n, posterior_scale)
    log_predictive = logpdf(posterior, datum)
    return log_predictive
end

function compute_data_log_predictives(observation::Int64, clusters::Vector{Int64}, data::Matrix{T},
                                      data_sufficient_stats::DataSufficientStatistics,
                                      data_parameters::DataParameters, auxillary_variables) where {T <: Real}
    log_predictive = Array{Float64}(undef, length(clusters))
    for (index, cluster) = enumerate(clusters)
        log_predictive[index] = data_log_predictive(observation, cluster, data_sufficient_stats, data_parameters, data, auxillary_variables)
    end
    return log_predictive
end

get_clusters(cluster_sufficient_stats::ClusterSufficientStatistics) = findall(cluster_sufficient_stats.clusters)

function update_cluster_stats_new_birth_time!(num_observations::Vector{Int64}, new_time::Int64, old_time::Int64)
    num_observations[new_time] = num_observations[old_time]
    num_observations[old_time] = 0
end

function compute_cluster_log_predictives(observation, cluster::Int64, previous_cluster, cluster_sufficient_stats,
                                         cluster_parameters::NtlParameters)
    return reshape(compute_cluster_log_predictives(observation, [cluster], previous_cluster, cluster_sufficient_stats,
                                                   cluster_parameters), 1)[1]
end

function compute_existing_cluster_log_predictives(observation::Int64, cluster::Int64, cluster_sufficient_stats,
                                                  cluster_parameters)
    result = compute_existing_cluster_log_predictives(observation, [cluster], cluster_sufficient_stats,
                                                      cluster_parameters)
    return reshape(result, 1)[1]
end

end

ext: jl

max_stars_repo_name: realseanla/ntl-stick-breaking-julia

max_issues_repo_name: realseanla/ntl-stick-breaking-julia

max_forks_repo_name: realseanla/ntl-stick-breaking-julia


text:
# Julia wrapper for header: /usr/include/scip/cons_cumulative.h
# Automatically generated using Clang.jl wrap_c

function SCIPincludeConshdlrCumulative(scip)
    ccall((:SCIPincludeConshdlrCumulative, libscip), SCIP_RETCODE, (Ptr{SCIP_},), scip)
end

function SCIPcreateConsCumulative(scip, cons, name, nvars, vars, durations, demands, capacity, initial, separate, enforce, check, propagate, _local, modifiable, dynamic, removable, stickingatnode)
    ccall((:SCIPcreateConsCumulative, libscip), SCIP_RETCODE, (Ptr{SCIP_}, Ptr{Ptr{SCIP_CONS}}, Cstring, Cint, Ptr{Ptr{SCIP_VAR}}, Ptr{Cint}, Ptr{Cint}, Cint, UInt32, UInt32, UInt32, UInt32, UInt32, UInt32, UInt32, UInt32, UInt32, UInt32), scip, cons, name, nvars, vars, durations, demands, capacity, initial, separate, enforce, check, propagate, _local, modifiable, dynamic, removable, stickingatnode)
end

function SCIPcreateConsBasicCumulative(scip, cons, name, nvars, vars, durations, demands, capacity)
    ccall((:SCIPcreateConsBasicCumulative, libscip), SCIP_RETCODE, (Ptr{SCIP_}, Ptr{Ptr{SCIP_CONS}}, Cstring, Cint, Ptr{Ptr{SCIP_VAR}}, Ptr{Cint}, Ptr{Cint}, Cint), scip, cons, name, nvars, vars, durations, demands, capacity)
end

function SCIPsetHminCumulative(scip, cons, hmin)
    ccall((:SCIPsetHminCumulative, libscip), SCIP_RETCODE, (Ptr{SCIP_}, Ptr{SCIP_CONS}, Cint), scip, cons, hmin)
end

function SCIPgetHminCumulative(scip, cons)
    ccall((:SCIPgetHminCumulative, libscip), Cint, (Ptr{SCIP_}, Ptr{SCIP_CONS}), scip, cons)
end

function SCIPsetHmaxCumulative(scip, cons, hmax)
    ccall((:SCIPsetHmaxCumulative, libscip), SCIP_RETCODE, (Ptr{SCIP_}, Ptr{SCIP_CONS}, Cint), scip, cons, hmax)
end

function SCIPgetHmaxCumulative(scip, cons)
    ccall((:SCIPgetHmaxCumulative, libscip), Cint, (Ptr{SCIP_}, Ptr{SCIP_CONS}), scip, cons)
end

function SCIPgetVarsCumulative(scip, cons)
    ccall((:SCIPgetVarsCumulative, libscip), Ptr{Ptr{SCIP_VAR}}, (Ptr{SCIP_}, Ptr{SCIP_CONS}), scip, cons)
end

function SCIPgetNVarsCumulative(scip, cons)
    ccall((:SCIPgetNVarsCumulative, libscip), Cint, (Ptr{SCIP_}, Ptr{SCIP_CONS}), scip, cons)
end

function SCIPgetCapacityCumulative(scip, cons)
    ccall((:SCIPgetCapacityCumulative, libscip), Cint, (Ptr{SCIP_}, Ptr{SCIP_CONS}), scip, cons)
end

function SCIPgetDurationsCumulative(scip, cons)
    ccall((:SCIPgetDurationsCumulative, libscip), Ptr{Cint}, (Ptr{SCIP_}, Ptr{SCIP_CONS}), scip, cons)
end

function SCIPgetDemandsCumulative(scip, cons)
    ccall((:SCIPgetDemandsCumulative, libscip), Ptr{Cint}, (Ptr{SCIP_}, Ptr{SCIP_CONS}), scip, cons)
end

function SCIPcheckCumulativeCondition(scip, sol, nvars, vars, durations, demands, capacity, hmin, hmax, violated, cons, printreason)
    ccall((:SCIPcheckCumulativeCondition, libscip), SCIP_RETCODE, (Ptr{SCIP_}, Ptr{SCIP_SOL}, Cint, Ptr{Ptr{SCIP_VAR}}, Ptr{Cint}, Ptr{Cint}, Cint, Cint, Cint, Ptr{UInt32}, Ptr{SCIP_CONS}, UInt32), scip, sol, nvars, vars, durations, demands, capacity, hmin, hmax, violated, cons, printreason)
end

function SCIPnormalizeCumulativeCondition(scip, nvars, vars, durations, demands, capacity, nchgcoefs, nchgsides)
    ccall((:SCIPnormalizeCumulativeCondition, libscip), SCIP_RETCODE, (Ptr{SCIP_}, Cint, Ptr{Ptr{SCIP_VAR}}, Ptr{Cint}, Ptr{Cint}, Ptr{Cint}, Ptr{Cint}, Ptr{Cint}), scip, nvars, vars, durations, demands, capacity, nchgcoefs, nchgsides)
end

function SCIPsplitCumulativeCondition(scip, nvars, vars, durations, demands, capacity, hmin, hmax, split)
    ccall((:SCIPsplitCumulativeCondition, libscip), SCIP_RETCODE, (Ptr{SCIP_}, Cint, Ptr{Ptr{SCIP_VAR}}, Ptr{Cint}, Ptr{Cint}, Cint, Ptr{Cint}, Ptr{Cint}, Ptr{Cint}), scip, nvars, vars, durations, demands, capacity, hmin, hmax, split)
end

function SCIPpresolveCumulativeCondition(scip, nvars, vars, durations, hmin, hmax, downlocks, uplocks, cons, delvars, nfixedvars, nchgsides, cutoff)
    ccall((:SCIPpresolveCumulativeCondition, libscip), SCIP_RETCODE, (Ptr{SCIP_}, Cint, Ptr{Ptr{SCIP_VAR}}, Ptr{Cint}, Cint, Cint, Ptr{UInt32}, Ptr{UInt32}, Ptr{SCIP_CONS}, Ptr{UInt32}, Ptr{Cint}, Ptr{Cint}, Ptr{UInt32}), scip, nvars, vars, durations, hmin, hmax, downlocks, uplocks, cons, delvars, nfixedvars, nchgsides, cutoff)
end

function SCIPpropCumulativeCondition(scip, presoltiming, nvars, vars, durations, demands, capacity, hmin, hmax, cons, nchgbds, initialized, explanation, cutoff)
    ccall((:SCIPpropCumulativeCondition, libscip), SCIP_RETCODE, (Ptr{SCIP_}, SCIP_PRESOLTIMING, Cint, Ptr{Ptr{SCIP_VAR}}, Ptr{Cint}, Ptr{Cint}, Cint, Cint, Cint, Ptr{SCIP_CONS}, Ptr{Cint}, Ptr{UInt32}, Ptr{UInt32}, Ptr{UInt32}), scip, presoltiming, nvars, vars, durations, demands, capacity, hmin, hmax, cons, nchgbds, initialized, explanation, cutoff)
end

function SCIPrespropCumulativeCondition(scip, nvars, vars, durations, demands, capacity, hmin, hmax, infervar, inferinfo, boundtype, bdchgidx, relaxedbd, explanation, result)
    ccall((:SCIPrespropCumulativeCondition, libscip), SCIP_RETCODE, (Ptr{SCIP_}, Cint, Ptr{Ptr{SCIP_VAR}}, Ptr{Cint}, Ptr{Cint}, Cint, Cint, Cint, Ptr{SCIP_VAR}, Cint, SCIP_BOUNDTYPE, Ptr{SCIP_BDCHGIDX}, Cdouble, Ptr{UInt32}, Ptr{SCIP_RESULT}), scip, nvars, vars, durations, demands, capacity, hmin, hmax, infervar, inferinfo, boundtype, bdchgidx, relaxedbd, explanation, result)
end

function SCIPvisualizeConsCumulative(scip, cons)
    ccall((:SCIPvisualizeConsCumulative, libscip), SCIP_RETCODE, (Ptr{SCIP_}, Ptr{SCIP_CONS}), scip, cons)
end

function SCIPsetSolveCumulative(scip, solveCumulative)
    ccall((:SCIPsetSolveCumulative, libscip), SCIP_RETCODE, (Ptr{SCIP_}, Ptr{Cvoid}), scip, solveCumulative)
end

function SCIPsolveCumulative(scip, njobs, ests, lsts, objvals, durations, demands, capacity, hmin, hmax, timelimit, memorylimit, maxnodes, solved, infeasible, unbounded, error)
    ccall((:SCIPsolveCumulative, libscip), SCIP_RETCODE, (Ptr{SCIP_}, Cint, Ptr{Cdouble}, Ptr{Cdouble}, Ptr{Cdouble}, Ptr{Cint}, Ptr{Cint}, Cint, Cint, Cint, Cdouble, Cdouble, Clonglong, Ptr{UInt32}, Ptr{UInt32}, Ptr{UInt32}, Ptr{UInt32}), scip, njobs, ests, lsts, objvals, durations, demands, capacity, hmin, hmax, timelimit, memorylimit, maxnodes, solved, infeasible, unbounded, error)
end

function SCIPcreateWorstCaseProfile(scip, profile, nvars, vars, durations, demands)
    ccall((:SCIPcreateWorstCaseProfile, libscip), SCIP_RETCODE, (Ptr{SCIP_}, Ptr{SCIP_PROFILE}, Cint, Ptr{Ptr{SCIP_VAR}}, Ptr{Cint}, Ptr{Cint}), scip, profile, nvars, vars, durations, demands)
end

function SCIPcomputeHmin(scip, profile, capacity)
    ccall((:SCIPcomputeHmin, libscip), Cint, (Ptr{SCIP_}, Ptr{SCIP_PROFILE}, Cint), scip, profile, capacity)
end

function SCIPcomputeHmax(scip, profile, capacity)
    ccall((:SCIPcomputeHmax, libscip), Cint, (Ptr{SCIP_}, Ptr{SCIP_PROFILE}, Cint), scip, profile, capacity)
end


text:
using MacroTools
import ..LambdaCalculus.Named

export @lambda, @λ

function meta_freevars(expr)
    if isexpr(expr, :$)
        @assert length(expr.args) == 1
        return Set{Symbol}()
    end

    @match expr begin
        (f_(arg1_, args__)) => mapreduce(meta_freevars, ∪, [f, arg1, args...])
        (arg_Symbol -> body_) => filter(!isequal(arg), meta_freevars(body))
        (name_Symbol) => Set([name])
        other_ => error("unhandled expression: $other")
    end
end

function meta_convert1(expr, Γ::NamingContext, level::Int = 0)
    if isexpr(expr, :$)
        @assert length(expr.args) == 1
        return :(shift($level, $(esc(expr.args[1]))))
    end

    @match expr begin
        (f_(arg1_, args__)) => mapfoldl(e -> meta_convert1(e, Γ, level),
                                        (e, arg) -> :(App($e, $arg)),
                                        [f, arg1, args...])
        (arg_Symbol -> body_) => :(Abs($(meta_convert1(body, pushfirst(Γ, arg), level + 1))))
        (name_Symbol) => :(Var($(findfirst(isequal(name), Γ))))
        other_ => error("unhandled expression: $other")
    end
end

function meta_convert(expr)
    expr = MacroTools.prewalk(unblock ∘ rmlines, expr)
    meta_convert1(expr, NamingContext(meta_freevars(expr)))
end

function meta_convert(expr, Γ::NamingContext)
    expr = MacroTools.prewalk(unblock ∘ rmlines, expr)
    meta_convert1(expr, Γ)
end

"Convert a (well-formed) Julia expression to a `Term`."
macro lambda(expr)
    meta_convert(expr)
end

macro lambda(Γ, expr)
    if Γ.head == :vect
        meta_convert(expr, NamingContext(collect(Symbol, Γ.args)))
    else
        error("invalid context specification: $Γ")
    end
end

"Convert a (well-formed) Julia expression to a `Term`."
macro λ(expr)
    meta_convert(expr)
end

macro λ(Γ, expr)
    if Γ.head == :vect
        meta_convert(expr, NamingContext(collect(Symbol, Γ.args)))
    else
        error("invalid context specification: $Γ")
    end
end


text:
using MbedTLS
const hmacstate = Ref{MbedTLS.MD{true}}()

function hmac(s1,s2,s3,s4)
    if !isdefined(hmacstate, :x)
        return ""
    else
        MbedTLS.reset!(hmacstate[])
        for s in (s1, s2, s3, s4)
            write(hmacstate[], s)
        end
        # Take the digest (returned as a byte array) and convert it to hex string representation
        digest = MbedTLS.finish!(hmacstate[])
        hexdigest = Vector{UInt8}(length(digest)*2)
        for i = 1:length(digest)
            b = digest[i]
            d = b >> 4
            hexdigest[2i-1] = UInt8('0')+d+39*(d>9)
            d = b & 0xf
            hexdigest[2i] = UInt8('0')+d+39*(d>9)
        end
        return String(hexdigest)
    end
end


text:
"""
```
combined_optimizer(fcn::Function,
                   x0::Vector,
                   args...;
                   xtol::Real            = 1e-32, # default from Optim.jl
                   ftol::Float64         = 1e-14, # Default from csminwel
                   grtol::Real           = 1e-8,  # default from Optim.jl
                   iterations::Int       = 1000,
                   store_trace::Bool     = false,
                   show_trace::Bool      = false,
                   extended_trace::Bool  = false,
                   verbose::Symbol       = :none,
                   rng::AbstractRNG      = MersenneTwister(),
                   autodiff::Bool        = false,
                   neighbor!::Function   = Optim.default_neighbor!,
                   temperature::Function = Optim.log_temperature,
                   max_cycles::Int       = 4,
                   kwargs...)

```
This routine alternates between LBFGS and simulated annealing.
LBFGS is good for quickly moving in the correct direction, while simulated annealing
is used to escape local minima.
"""
function combined_optimizer(fcn::Function,
                            x0::Vector,
                            args...;
                            xtol::Real            = 1e-32, # default from Optim.jl
                            ftol::Float64         = 1e-14, # default from csminwel
                            grtol::Real           = 1e-8,  # default from Optim.jl
                            iterations::Int       = 1000,
                            store_trace::Bool     = false,
                            show_trace::Bool      = false,
                            extended_trace::Bool  = false,
                            verbose::Symbol       = :none,
                            rng::AbstractRNG      = MersenneTwister(),
                            autodiff::Bool        = false,
                            neighbor!::Function   = Optim.default_neighbor!,
                            temperature::Function = Optim.log_temperature,
                            max_cycles::Int       = 4,
                            kwargs...)

    cycle      = 0
    converged  = false
    x_opt      = x0
    use_sa_out = false
    f_opt      = fcn(x_opt)
    out_sa     = nothing
    out_lbfgs  = nothing

    while cycle < max_cycles && converged == false

        # first, run LBFGS
        println(verbose, :low, "Running L-BFGS...")
        out_lbfgs = Optim.optimize(fcn, x_opt, lbfgs(),
                   Optim.Options(autodiff=autodiff, g_tol = grtol, f_tol = ftol, x_tol = xtol,
                   iterations = iterations, store_trace = store_trace, show_trace = show_trace,
                   extended_trace = extended_trace))

        # store relevant information from the optimizer
        minimum_lbfgs   = out_lbfgs.minimum
        minimizer_lbfgs = out_lbfgs.minimizer

        # second, run simulated annealing with the trace disabled
        println(verbose, :low, "Running simulated annealing...")

        # simulated annealing gets more iterations, no trace printout
        sa_iterations = min(iterations * 10, 250)
        out_sa = Optim.optimize(fcn, x_opt,
                                method = SimulatedAnnealing(neighbor!   = neighbor!,
                                                            temperature = temperature),
                                iterations = sa_iterations, store_trace = store_trace,
                                show_trace = false,
                                extended_trace = extended_trace)

        minimum_sa   = out_sa.minimum
        minimizer_sa = out_sa.minimizer

        # reject a move to a worse spot
        use_sa_out = minimum_sa < minimum_lbfgs
        x_opt      = use_sa_out ? minimizer_sa : minimizer_lbfgs
        cycle_best = use_sa_out ? minimum_sa : minimum_lbfgs

        # check convergence
        rel_diff  = abs(f_opt - cycle_best)/abs(f_opt)
        converged = rel_diff < ftol
        cycle    += 1

        println(verbose, :low, "cycle: $cycle")
        println(verbose, :low, "relative function difference: $(round(rel_diff,5))")

        f_opt = cycle_best
    end

    println("optimization complete. cycles: $cycle")
    out = use_sa_out ? out_sa : out_lbfgs
    return out
end


text:
using Test

include("LogicoW.jl")

@testset "substw" begin
 @test substw([:x,:y],[:P :x :y],[[:x,:y]]) == [:P :x :y]
 @test substw([:x,:y],[:P :x :y],[[:x],[:y]]) == [:P :x :y]
 @test substw([:x,:y],[:P :x :y],[:a,:b]) == [:P :a :b]
 @test substw([:x,:y],[:P :x :y],[[:a],[:b]]) == [:P :a :b]
 @test substw([:x,:y],[:P :x :y],[[:a :b],[:c :d]]) == [:P :a :b :c :d]
 @test substw([:x,:y],[:P :x :k :y],[[:a :b],[:c :d]]) == [:P :a :b :k :c :d]
 @test substw([:x,:y],[:P :x :y],[[],[:c :d]]) == [:P :c :d]
 @test substw([:x,:y],[:P :x :y],[[:a :b],[]]) == [:P :a :b]

end

@testset "unifyttw Tests" begin

 @test unifyttw([:x,:y],[:x :c :y], [:P :a :b :d :e :f]) == [:x,:y]
 @test unifyttw([:x,:y],[:x :c :y], [:P :a :b :c :d :e :f]) == [[:P :a :b],[:e :f]]
 @test unifyttw([:x,:y],[:x :c :y], [:P :a :b :c]) == [[:P :a :b],[]]
 @test unifyttw([:x,:y],[:x :c :y], [:P :c :d :e]) == [[:P],[:d :e]]

 @test unifyttw([:x,:y],[:x :y], [:P :a :b :c :d :e :f]) == [[:P :a :b :c :d :e :f],[]]

 @test unifyttw([:x,:y],[:x :c :y], [:a :c :d]) == [[:a],[:d]]
 @test unifyttw([:x,:y],[:x :c :y], [:c :d]) == [[],[:d]]
 @test unifyttw([:x,:y],[:x :c :y], [:a :c])== [[:a],[]]
 @test unifyttw([:x,:y],[:P :x :c :y], [:P :a :c :d]) == [[:a],[:d]] # or [:a,:d]
 @test unifyttw([:x,:y],[:P :x :c :y], [:P :a :b :c :d]) == [[:a :b],[:d]] # or [[:a :b],:d]
 @test unifyttw([:x,:y],[:P :x :c :y], [:P :a :b :c :d :e :f]) == [[:a :b],[:e :f]]

 @test unifyttw([:x,:y],[:P :x :c :y], [:P :c]) == [[],[]]
 @test unifyttw([:x,:y],[:P :x :c :y], [:P :a :c]) == [[:a],[]]
 @test unifyttw([:x,:y],[:P :x :c :y], [:P :c :b]) == [[],[:b]]
end

max_stars_repo_name: sazare/cheapViews

max_issues_repo_name: sazare/cheapViews

max_forks_repo_name: sazare/cheapViews


text:
@testset "Poly" begin

  K = PadicField(2, 100)
  Kx, x = PolynomialRing(K, "x")
  L, gL = eisenstein_extension(x^2+2, "a")

  @testset "Fun Factor" for F in [K, L]
    Fx, x = PolynomialRing(F, "x")
    f = x^5
    for i = 0:4
      c = K(rand(FlintZZ, 1:100))
      f += c*x^i
    end
    u = 1
    for i = 1:5
      c = K(rand(FlintZZ, 1:100))
      u += 2*c*x^i
    end

    g = f*u
    u1, f1 = @inferred Hecke.fun_factor(g)
    @test u == u1
    @test f1 == f
  end

  @testset "Gcd" for F in [K, L]
    Fx, x = PolynomialRing(F, "x")
    f = (2*x+1)*(x+1)
    g = x^3+1
    gg = @inferred gcd(f, g)
    @test gg == x+1

    f = (2*x+1)*(x+1)
    g = (2*x+1)*(x+2)
    @test gcd(f, g) == 2*x+1

    f = (x + 1//K(2)) * (2*x^2+x+1)
    g = 2*x+1
    @test gcd(f, g) == g
  end

  @testset "Gcdx" for F in [K, L]
    Fx, x = PolynomialRing(F, "x")
    f = (2*x+1)*(x+1)
    g = x^3+1
    d, u, v = gcdx(f, g)
    @test d == gcd(f, g)
    @test u*f + v*g == d

    f = (2*x+1)*(x+1)
    g = (2*x+1)*(x+2)
    d, u, v = @inferred gcdx(f, g)
    @test gcd(f, g) == d
    @test d == u*f + v*g

    f = (x + 1//K(2)) * (2*x^2+x+1)
    g = 2*x+1
    d, u, v = gcdx(f, g)
    @test g == d
    @test u*f + v*g == d
  end

  @testset "Hensel" for F in [K, L]
    Fx, x = PolynomialRing(F, "x")
    f = (x+1)^3
    g = (x^2+x+1)
    h = x^2 +2*x + 8
    ff = f*g*h
    lf = @inferred Hecke.Hensel_factorization(ff)
    @test prod(values(lf)) == ff
  end

  @testset "Slope Factorization" for F in [K, L]
    Fx, x = PolynomialRing(F, "x")
    f = prod(x-2^i for i = 1:5)
    lf = @inferred Hecke.slope_factorization(f)
    @test prod(keys(lf)) == f
    @test all(x -> isone(degree(x)), keys(lf))
    @test length(Hecke.slope_factorization(2*x+1)) == 1
  end

  @testset "Resultant" begin
    R, x = PolynomialRing(PadicField(853, 2), "x")
    a = 4*x^5 + x^4 + 256*x^3 + 192*x^2 + 48*x + 4
    b = derivative(a)
    rab = @inferred resultant(a, b)
    @test rab == det(sylvester_matrix(a, b))
  end
end


text:
immutable Reword<:RelEx hi::Word; lo::Word end
==(x::Reword, y::Reword) = x.hi == y.hi && x.lo == y.lo

rx(hi::Word, lo::Word) = Reword(hi, lo)

immutable BotR<:RelEx end
immutable TopR<:RelEx end

==(x::BotR, y::BotR) = true
==(x::TopR, y::TopR) = true

botr() = BotR()
topr() = TopR()

immutable AltR{S<:RelEx,T<:RelEx}<:RelEx a::S; b::T end
immutable CatR{S<:RelEx,T<:RelEx}<:RelEx a::S; b::T end

==(x::AltR, y::AltR) = x.a == y.a && x.b == y.b
==(x::CatR, y::CatR) = x.a == y.a && x.b == y.b

alt(x1::RelEx, x2::RelEx, xs...) = (isempty(xs) ? AltR(x1, x2)
                                    : AltR(x1, alt(x2, xs...)))

cat(x1::RelEx, x2::RelEx, xs...) = (isempty(xs) ? CatR(x1, x2)
                                    : CatR(x1, cat(x2, xs...)))

immutable StarR{T<:RelEx}<:RelEx a::T end
immutable PlusR{T<:RelEx}<:RelEx a::T end

==(x::StarR, y::StarR) = x.a == y.a
==(x::PlusR, y::PlusR) = x.a == y.a

star(x::RelEx) = StarR(x)
plus(x::RelEx) = PlusR(x)


text:
#Solves the predictor corrector euler scheme - pg327 of Bruti-Liberati 2007
#Strong order 0.5, with some degree of implicitness

#Prototyping predictor corrector code
struct PCEproto <: Solver end
struct PCE <: Solver end

"""
optimized predictor corrector integrate
ab! : bias & diffusion function
x0 : Initial state
t : times to use for time evolution
noise : Noise object
callback : a function that runs on t, x, W
"""
function integrate(a!::Function, b!::Function, x0::AbstractVector{T}, t::AbstractVector{<:Real},
    noise::Noise, pce::PCEproto, callback::Function; θ::Float64=0.5, η::Float64 = 0., bbprime!::Function = (t,x,bbp)-> nothing) where {T<:Number}

    n = length(x0)
    bbp = zeros(T, n)

    function abar!(t, x, ff)
        a!(t, x, ff)
        bbprime!(t, x, bbp)
        ff .= ff - η*bbp
        return nothing
    end
    W = W0(noise)
    ΔW = copy(W)
    x = copy(x0)

    f = Vector{T}(n)
    g = Matrix{T}(n, noise.m)

    xbar = copy(x0)
    fbar = copy(f)
    gbar = copy(g)

    callback(t[1], x, W)
    for i in 1:length(t)-1
        Δt = t[i+1]-t[i]
        randn!(noise, ΔW)
        ΔW .*= sqrt(Δt)

        abar!(t[i], x, f)
        b!(t[i], x, g)

        xbar .= copy(x)
        xbar .+= f*Δt + g*ΔW

        abar!(t[i+1], xbar, fbar)
        b!(t[i+1], xbar, gbar)

        x .+= (θ*fbar + (1-θ)*f)*Δt + (η*gbar + (1-η)*g)*ΔW
        W .+= ΔW
        callback(t, x, W)
    end
end

function integrate(a!::Function, b!::Function, x0::AbstractVector{T}, t::AbstractVector{<:Real},
    noise::Noise, pce::PCE, callback::Function; θ::Float64=0.5, η::Float64 = 0., bbprime!::Function = (t,x,bbp)-> nothing) where {T<:Number}

    n = length(x0)
    bbp = zeros(T, n)

    function abar!(t, x, ff)
        a!(t, x, ff)
        bbprime!(t, x, bbp)
        bbp .*= -η
        ff .+=  bbp
        return nothing
    end

    W = W0(noise)
    ΔW = copy(W)
    x = copy(x0)

    fΔt = Vector{T}(n)
    g = Matrix{T}(n, noise.m)
    gΔW = Vector{T}(n)

    xbar = copy(x)

    #save the first data point
    callback(t[1], x, W)

    for i in 1:length(t)-1
        Δt = t[i+1]-t[i]
        randn!(noise, ΔW)
        ΔW .*= sqrt(Δt)

        abar!(t[i], x, fΔt)
        b!(t[i], x, g)

        fΔt .*= Δt
        A_mul_B!(gΔW, g, ΔW)

        xbar = copy(x)
        xbar .+= fΔt
        xbar .+= gΔW

        fΔt .*= (1-θ)
        gΔW .*= (1-η)

        x .+= fΔt
        x .+= gΔW

        abar!(t[i+1], xbar, fΔt)
        b!(t[i+1], xbar, g)

        fΔt .*= (Δt*θ) #fΔbar*θ
        A_mul_B!(gΔW, g, ΔW)
        gΔW .*= η #gbar*η
        x .+= fΔt
        x .+= gΔW

        W .+= ΔW
        callback(t, x, W)
    end
    return nothing
end

max_stars_repo_name: onoderat/SDEs

max_issues_repo_name: onoderat/SDEs

max_forks_repo_name: onoderat/SDEs


text:
module RunTests

println("Testing")

include("gradientstoragetests.jl")
include("leastsquares.jl")
include("lasso.jl")
include("loggertests.jl")

end


text:
## Analysis and Display

function circleShape(h, k, r)
    θ = LinRange(0, 2 * π, 500)
    h .+ r * sin.(θ), k .+ r * cos.(θ)
end

function histogram2D(states::Vector{Any}, sz::Int, zoom::Int)
    x,y=getxy(states)
    xbins = range(1, sz;step=1 / zoom)
    ybins = range(1, sz;step=1 / zoom)
    histplot = histogram2d(x, y, nbins=xbins, ybins, aspect_ratio=:equal, show_empty_bins=true)
    return histplot
end

function histogram2D(states::Vector{Any}, sz::Int, zoom::Int, d::ArrayDD, truestate::StateFlatBg)
    x,y=getxy(states)
    xbins = range(1, sz;step=1 / zoom)
    ybins = range(1, sz;step=1 / zoom)
    fig = heatmap(d.data, color=:greys)
    histogram2d!(fig, x, y, nbins=xbins, ybins, aspect_ratio=:equal, show_empty_bins=false)
    dcircle = 0.5
    for nn = 1:truestate.n
        plot!(fig, circleShape(truestate.x[nn], truestate.y[nn], dcircle), color=:blue)
    end
    return fig
end

function histogram2D(states::Vector{Any}, sz::Int, zoom::Int, truestate::StateFlatBg)
    x,y=getxy(states)
    xbins = range(1, sz;step=1 / zoom)
    ybins = range(1, sz;step=1 / zoom)
    fig=histogram2d(x, y, nbins=xbins, ybins, aspect_ratio=:equal, show_empty_bins=false)
    dcircle = 0.5
    for nn = 1:truestate.n
        plot!(fig, circleShape(truestate.x[nn], truestate.y[nn], dcircle), color=:blue)
    end
    histogram2d!(fig, x, y, nbins=xbins, ybins, aspect_ratio=:equal, show_empty_bins=false,yflip=true)
    # scatter!(fig, x, y)
    return fig
end

#"Show model and data vs chain step"
function showoverlay(states::Vector{Any},data::BAMFData,psf::MicroscopePSFs.PSF)
    len=length(states)
    tmp=deepcopy(data)
    sz=data.sz
    d=Array{Float32}(undef,(sz,sz*size(tmp.data,3),len))
    m=Array{Float32}(undef,(sz,sz*size(tmp.data,3),len))
    println(size(d))
    dataim=reshape(data.data,(sz,sz*size(data.data,3)))

    println(states[1])
    for nn=1:len
        d[:,:,nn]=dataim
        genmodel!(states[nn],psf,tmp)
        m[:,:,nn]=reshape(tmp.data,(sz,sz*size(data.data,3)))
    end

    globmax=maximum((maximum(m),maximum(d)))
    d=d./globmax
    m=m./globmax
    out=cat(dims=1,d,m,d-m)
    println(size(out))
    return out
end

function plotstate(truestate::StateFlatBg,foundstate::StateFlatBg)
    fig=plot()
    dcircle = 0.5
    for nn = 1:truestate.n
        plot!(fig, circleShape(truestate.x[nn], truestate.y[nn], dcircle), color=:blue,yflip=true)
    end

    for nn = 1:foundstate.n
        plot!(fig, circleShape(foundstate.x[nn], foundstate.y[nn], dcircle), color=:red,yflip=true)
    end

    return fig
end

function plotstate(truestate::StateFlatBg,foundstate::StateFlatBg_Results; fig::Plots.Plot{Plots.GRBackend}=plot())
    #fig=plot()
    dcircle = 0.5
    for nn = 1:truestate.n
        plot!(fig, circleShape(truestate.x[nn], truestate.y[nn], dcircle), color=:blue,yflip=true)
    end

    for nn = 1:foundstate.n
        σ=(foundstate.σ_x[nn]+foundstate.σ_x[nn])/2f0
        plot!(fig, circleShape(foundstate.x[nn], foundstate.y[nn], σ), color=:red,yflip=true)
    end

    return fig
end


text:
using PhilipsHue, Colors

println("loaded PhilipsHue")

# these will have to be set to valid values before you run the test

B = PhilipsHueBridge("192.168.1.2"      , "PHtB7h1lkkDScH9Ln4JJpeke3V3lIqZMV9Ry5PkQ")

lightnumbers = getlightnumbers(B)

firstlight = first(lightnumbers)

println("Bridge config: \n\t"       , getbridgeconfig(B))
println("Bridge IP: \n\t"           , getIP())
println("Bridge initialized?: \n\t" , isinitialized(B))
println("Light config: \n\t"        , getlights(B))
println("Get light config: \n\t"    , getlight(B, firstlight))
println("Get light config: \n\t"    , getlight(B, 4))

println("Set first light on: \n\t"  , setlight(B, firstlight, Dict("on" => true)))
println("Set first light off: \n\t" , setlight(B, firstlight, Dict("on" => false)))
println("Set first light  \n\t"     , setlight(B, firstlight, Dict("on" => true, "sat" => 123, "bri" => 243, "hue" => 123)))

# random all lights
for i in 0:20
    setlights(B, Dict("bri" => rand(0:255), "hue" => rand(1:65000), "sat" => rand(1:255)))
    sleep(0.5)
end

# for a random light, set RGB color

# Thanks, @ScottPJones!
for r in 0:0.25:1, g in 0:0.25:1, b in 0:0.25:1
    setlight(B, lightnumbers[rand(1:end)], Colors.RGB(r, g, b))
    sleep(.25)
end

setlights(B, Dict("bri" => 255))


text:
include("7.1 Inverse hyperbolic sine/.jl")
include("7.2 Inverse hyperbolic cosine/.jl")
include("7.3 Inverse hyperbolic tangent/.jl")
include("7.5 Inverse hyperbolic secant/.jl")

max_stars_repo_path: src/rules/7 Inverse hyperbolic functions/.jl


text:
include("test_abFunc.jl")
include("test_abGrad.jl")
include("test_bFunc.jl")
include("test_bGrad.jl")
include("test_exactdmd.jl")
include("test_projb.jl")
include("test_setup.jl")


text:
# Random Fourier Features
# from: Rahimi,Recht, "Random Features for Large-Scale Kernel Machines", NIPS 2007
#
# Danny Perry (dperry@cs.utah.edu)
# Oct 2014
#
# Gaussian Kernel Approximation

using Hadamard
using Distributions

# A set of random fourier functions
type RandomFourierFunctions{T}
	directions::Array{T,2} # random direction
	offsets::Array{T,1} # random shift in [0,2*pi]
end

# project the input data into the approximate kernel space
#
# @param X - input data
# @param RFFs - the random fourier functions
function RFFProject!{T}(z::Array{T,2}, X::Array{T,2}, RFFs::RandomFourierFunctions{T})
	n = size(X,1)
	D = size(RFFs.directions,1)
	z[:,:] = convert(Array{T,2}, (sqrt(2/D)) * cos( (X * RFFs.directions') .+ ones(n,1) * RFFs.offsets' ) )
end
function RFFProject{T}(X::Array{T,2}, RFFs::RandomFourierFunctions{T})
	z = zeros(T, size(X,1), size(RFFs.directions,1))
	RFFProject!(z, X, RFFs)
	return z
end
# use the sin-cos form instead
function RFFProjectSinCos!{T}(z::Array{T,2}, X::Array{T,2}, RFFs::RandomFourierFunctions{T})
	n = size(X,1)
	D = 2*size(RFFs.directions,1)
	z[:,2:2:end] = X * RFFs.directions'
	z[:,1:2:end] = convert( Array{T,2}, sin( z[:,2:2:end] ) )
	z[:,2:2:end] = convert( Array{T,2}, cos( z[:,2:2:end] ) )
	z[:,:] *= (sqrt(2/D))
end
function RFFProjectSinCos{T}(X::Array{T,2}, RFFs::RandomFourierFunctions{T})
	z = zeros(T, size(X,1), 2*size(RFFs.directions,1))
	RFFProjectSinCos!(z, X, RFFs)
	return z
end

# @param sigma - Gaussian kernel parameter
# @param dimension - dimension of input data
# @param num_functions - how many functions to generate
function GenerateFunctionsGaussian(T, sigma, input_dimension, num_functions)
	# draw directions from a normal distribution with appropriate sigma corresponding to the kernel:
	w = randn(num_functions, input_dimension) * (1/sigma)

	# draw random offsets from a uniform distribution in [0,2*pi]
	b = rand(T, (num_functions,)) * 2 * pi

	# construct the random function struct:
	RandomFourierFunctions{T}(w,b)
end
# use the sin-cos form instead
function GenerateFunctionsGaussianSinCos(T, sigma, input_dimension, num_functions)
	if mod(num_functions,2) != 0
		error("num_functions must be an even number.")
	end
	# draw directions from a normal distribution with appropriate sigma corresponding to the kernel:
	w = randn(div(num_functions,2), input_dimension) * (1/sigma)

	# draw random offsets from a uniform distribution in [0,2*pi]
	b = zeros(1)

	# construct the random function struct:
	RandomFourierFunctions{T}(w,b)
end

# A version that draws from the probability distribution with bias towards the data
# @param sigma - Gaussian kernel parameter
# @param dimension - dimension of input data
# @param num_functions - how many functions to generate
# @param Xsample - subsample of data to select random basis
# @param Ksample - subsample of data to select random basis as kernel matrix
# @param centered - whether or not the data should be centered
function GenerateFunctionsGaussianBiased(sigma, input_dimension, num_functions, Xsample, Ksample, centered)
	w = zeros(num_functions, input_dimension)
	b = zeros(num_functions)

	samples = num_functions
	runs = 100
	# take $(samples) and pick the closest to the eigenvectors

	for k=1:samples:num_functions
		nsamples = samples
		if k + samples > num_functions
			nsamples = (num_functions-k+1)
		end
		dists = Float64[]
		rffs = RandomFourierFunctions[]
		for i=1:runs
			ws = randn(nsamples, input_dimension) * (1/sigma)
			bs = rand(nsamples) * 2 * pi
			push!(rffs, RandomFourierFunctions(ws,bs))
			zs = RFFProject( Xsample, rffs[end] )
			if centered
				zs  = zs - ones(size(zs,1),1)*mean(zs,1)
			end
			Ktmp = zs * zs'
			#push!(dists, vecnorm(Ktmp - Ksample))
			push!(dists, norm(Ktmp - Ksample))
		end
		minval,mini = findmin(dists)
		w[k:k+nsamples-1,:] = rffs[mini].directions
		b[k:k+nsamples-1] = rffs[mini].offsets
	end

	# draw directions from a normal distribution with appropriate sigma corresponding to the kernel:
	#w = randn(num_functions, input_dimension) * (1/sigma)
	# draw random offsets from a uniform distribution in [0,2*pi]
	#b = rand(num_functions) * 2 * pi

	# construct the random function struct:
	RandomFourierFunctions(w,b)
end

# Fast Food - approximate kernel space
# adapted from matlab implementation found here:
# http://www.mathworks.com/matlabcentral/fileexchange/49142-fastfood-kernel-expansions

# A set of fast food functions
type FastFood{T}
	b::Array{Int64,2}
	g::Array{T,2}
	p::Array{Int64,2}
	s::Array{T,2}
end

# @param dimension - dimension of input data
# @param num_functions - how many functions to generate
function GenerateFastFood(T, num_functions, dimension)

	d = dimension
	n = num_functions

	l = int(ceil(log2(d)))
	d = 2^l
	k = int(ceil(n/d))
	n = d*k

	# B matrix - diagonal of uniform {-1,1} values
	b = ((rand(d,k) .> .5)*2) .- 1

	# G matrix - diagonal of Guassian values
	g = randn(d,k)

	# Π matrix - diagonal permutation matrix
	Π = zeros(Int64,d,k) #randperm(d,k)

	if true
		# S matrix - scaling matrix
		gfro = 1.0/sqrt(vecnorm(g))
		s = zeros(k,d)
		gammad = Gamma(d-1,1)
		for i=1:k
			Π[:,i] = randperm(d)
			s[i,:] = rand(gammad, d)
		end
		tmp = gfro * s

		s = zeros(n, 1);
		for i = 1:k
				s[((i-1)*d+1):(i*d)] = tmp[i]
		end

		return FastFood{T}(b,g,Π,s) # to go
	else
		for i=1:k
			Π[:,i] = randperm(d)
		end
	end

	FastFood{T}(b,g,Π) # to go
end

# @param X - input data
# @param ff - fastfood parameters
# @param sigma - Gaussian kernel parameter
function FFProject{T}(X::Array{T,2}, ff::FastFood{T}, sigma)
	if !isa(sigma,T)
		sigma = convert(T,sigma)
	end
	m = size(X,2) # num data points
	d = size(X,1) # dimension of data
	k = size(ff.b,2) # factor of num of FF dimensions

	# pad with zeros until d = 2^l
	l = int(ceil(log2(d)))
	d = 2^l
	if d == size(X,1)
		X = X
	else
		X2 = zeros(T, d,m)
		X2[1:size(X,1),:] = X
		X = X2
	end
	n = d*k

	THT = zeros(T, n,m)

	for i=1:k
		tmp = copy(X)
		for j=1:m
			tmp[:,j] .*= ff.b[:,i]
		end
		tmp = fwht_natural(tmp,1)
		tmp = tmp[ff.p[:,i],:]
		for j=1:m
			tmp[:,j] .*= ff.g[:,i]*d
		end
    ind = ((i-1)*d+1):(i*d)
		THT[ind, :] = fwht_natural(tmp,1)
	end
	for i=1:m
		THT[:,i] .*= ff.s
	end

	THT ./= sigma
	phi = [cos(THT); sin(THT)] ./ sqrt(n)
end


text:
using ArgParse
s=ArgParseSettings()
@add_arg_table s begin
    "--part2"
        help="Do part 2"
        action=:store_true
    "--test"
        help="Run test data"
        action=:store_true
end
parsed_args = parse_args(ARGS, s)

function costtomove(n)
    parsed_args["part2"] ? div(n * (n + 1), 2) : n
end

function findcost(data, aim)

    sum([costtomove(abs(c - aim)) for c in data])
end

function findminimum(data)
    total = sum(data)
    attempt = div(total, length(data))
    costplus = findcost(data, attempt + 1)
    cost = findcost(data, attempt)
    costminus = findcost(data, attempt - 1)
    while true
        if cost <= costplus && cost <= costminus
            return cost
        end
        if costminus <= cost
            costplus = cost
            cost = costminus
            attempt = attempt - 1
            costminus = findcost(data, attempt - 1)
        else
            costminus = cost
            cost = costplus
            attempt = attempt + 1
            costminus = findcost(data, attempt  + 1)
        end
    end
end

test_suffix = parsed_args["test"] ? "_test" : ""
lines = readlines("input/day_7$test_suffix.txt")
data = [parse(Int32, crab) for crab in split(lines[1],",")]
answer = findminimum(data)
println("The answer is: $answer")


text:
function read(fn::AbstractString, bands = nothing)
    isfile(fn) || error("File not found.")
    dataset = ArchGDAL.unsafe_read(fn)

    if bands === nothing
        nbands = ArchGDAL.nraster(dataset)
        bands = 1:nbands
    end
    A = ArchGDAL.read(dataset, bands)
    am = get_affine_map(dataset)

    # nodata masking
    # A = Array{Union{Missing, eltype(A)}}(A)
    mask = falses(size(A))

    for i = bands
        band = ArchGDAL.getband(dataset, i)
        maskflags = mask_flags(band)

        # All values are valid, skip masking
        if :GMF_ALL_VALID in maskflags
            @debug "No masking"
            continue
        # Mask is valid for all bands
        elseif :GMF_PER_DATASET in maskflags
            @debug "Mask for each band"
            maskband = ArchGDAL.getmaskband(band)
            m = ArchGDAL.read(maskband) .== 0
            mask[:,:,i] = m
        # Alpha layer
        elseif :GMF_ALPHA in maskflags
            @warn "Dataset has band $i with an Alpha band, which is unsupported for now."
            continue
        # Nodata values
        elseif :GMF_NODATA in maskflags
            @debug "Flag NODATA"
            nodata = get_nodata(band)
            mask[:, :, i] = A[:,:,i] .== nodata
        else
            @warn "Unknown/unsupported mask."
        end
    end

    if any(mask)
        A = Array{Union{Missing,eltype(A)}}(A)
        A[mask] .= missing
    end

    wkt = ArchGDAL.getproj(dataset)
    ArchGDAL.destroy(dataset)
    GeoArray(A, am, wkt)
end

"""
- rast(fn::AbstractString, bands = nothing): construct `GeoArray` object from file
- rast(ga::GeoArray; vals::AbstractArray{T}): Reconstruct `GeoArray` object
"""
rast(fn::AbstractString, bands = nothing) = read(fn, bands)
function rast(ga::GeoArray; vals::AbstractArray{T}) where T<: Real
    GeoArray(vals, ga.f, ga.crs)
end

"""
    st_read(file::String, options...)
    st_read(files::Array{String,1}, options)

# Arguments:
- `options`: other parameters to `ArchGDAL.read(dataset, options...)`.
"""
function st_read(file::String, options...)
    ArchGDAL.read(file) do dataset
        ArchGDAL.read(dataset, options...)
    end
end

# read multiple tiff files and cbind
function st_read(files::Array{String,1}, options)
    # bands = collect(bands)
    # bands = collect(Int32, bands)
    res = map(file -> st_read(file, options...), files)
    vcat(res...)
end

export st_read, rast