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


export EFTfitterModel
export to_correlation_matrix

export get_observables, get_parameters, get_measurements, get_correlations
export get_nuisance_correlations, get_total_covariance, get_covariances, get_measurement_distributions

"""
    struct EFTfitterModel

This is the central type for using EFTfitter.
It comprises all information necessary for performing an analysis.
Only active `Measurement` and `Correlation` objects will be considered.

Fields:
* `parameters::BAT.NamedTupleDist`
* `measurements::NamedTuple{<:Any, <:Tuple{Vararg{Measurement}}}`
* `measurementdistributions::NamedTuple{<:Any, <:Tuple{Vararg{MeasurementDistribution}}}`
* `correlations::NamedTuple{<:Any, <:Tuple{Vararg{Correlation}}}`
* `nuisances::Union{NamedTuple{<:Any, <:Tuple{Vararg{NuisanceCorrelation}}}, Nothing}`

Constructors:

```julia
EFTfitterModel(
    parameters::BAT.NamedTupleDist,
    measurements::NamedTuple{<:Any, <:Tuple{Vararg{AbstractMeasurement}}},
    correlations::NamedTuple{<:Any, <:Tuple{Vararg{AbstractCorrelation}}},
    nuisances::Union{NamedTuple{<:Any, <:Tuple{Vararg{NuisanceCorrelation}}}, Nothing} = nothing
)
```

Examples:

```julia
model = EFTfitterModel(parameters, measurements, correlations) # no nuisance correlations

model = EFTfitterModel(parameters, measurements, correlations, nuisances) # with nuisance correlations
)
```
"""
struct EFTfitterModel
    parameters::BAT.NamedTupleDist
    measurements::NamedTuple{<:Any, <:Tuple{Vararg{Measurement}}}
    measurementdistributions::NamedTuple{<:Any, <:Tuple{Vararg{MeasurementDistribution}}}
    correlations::NamedTuple{<:Any, <:Tuple{Vararg{Correlation}}}
    nuisances::Union{NamedTuple{<:Any, <:Tuple{Vararg{NuisanceCorrelation}}}, Nothing}
end

function EFTfitterModel(
    parameters::BAT.NamedTupleDist,
    measurements::NamedTuple{<:Any, <:Tuple{Vararg{AbstractMeasurement}}},
    correlations::NamedTuple{<:Any, <:Tuple{Vararg{AbstractCorrelation}}},
    nuisances::Union{NamedTuple{<:Any, <:Tuple{Vararg{NuisanceCorrelation}}}, Nothing} = nothing
)
    measurement_vec, measurement_keys = unpack(measurements)
    correlation_vec, uncertainty_keys = unpack(correlations)

    # convert elements of MeasurementDistribution to Measurement for each bin
    binned_measurements, binned_measurement_keys = convert_to_bins(measurement_vec, measurement_keys)
    # use only active measurements/bins
    active_measurements, active_measurement_keys, corrs = only_active_measurements(binned_measurements, binned_measurement_keys, correlation_vec)
    # use only active uncertainties and correlations
    active_measurements, active_correlations, uncertainty_keys = only_active_uncertainties(active_measurements, corrs, uncertainty_keys)

    correlation_nt = namedtuple(uncertainty_keys, active_correlations)
    measurement_nt = namedtuple(active_measurement_keys, active_measurements)
    meas_dists_nt  = create_distributions(measurements, uncertainty_keys)
    nuisances_nt   = only_active_nuisances(nuisances, active_measurement_keys, uncertainty_keys)

    params = add_nuisance_parameters(parameters, nuisances_nt)

    return EFTfitterModel(params, measurement_nt, meas_dists_nt, correlation_nt, nuisances_nt)
end
#=============================================================#
"""
    get_parameters(m::EFTfitterModel)

Returns model parameters.
"""
get_parameters(m::EFTfitterModel) = m.parameters

"""
    get_measurements(m::EFTfitterModel)

Returns a `NamedTuple` with the `Measurement`s in the `EFTfitterModel`.
"""
get_measurements(m::EFTfitterModel) = m.measurements

"""
    get_measurement_distributions(m::EFTfitterModel)

Returns a `NamedTuple` with the `MeasurementDistribution`s in the `EFTfitterModel`.
"""
get_measurement_distributions(m::EFTfitterModel) = m.measurementdistributions

"""
    get_correlations(m::EFTfitterModel)

Returns a `NamedTuple` with the `Correlations`s in the `EFTfitterModel`.
"""
get_correlations(m::EFTfitterModel) = m.correlations

"""
    get_nuisance_correlations(m::EFTfitterModel)

Returns a `NamedTuple` with the `NuisanceCorrelation`s in the `EFTfitterModel`.
"""
get_nuisance_correlations(m::EFTfitterModel) = m.nuisances

"""
    get_observables(m::EFTfitterModel)

Returns a `NamedTuple` with the `Observable`s in the `EFTfitterModel`.
Note: The upper and lower limits are ignored and for each unique `Function`s only one `Observable` is returned.
"""
function  get_observables(model::EFTfitterModel)
    meas = get_measurements(model)
    obs = unique(Observable.([m.observable.func for m in values(meas)]))
    obs_names = [string(o.func) for o in obs]
    observables_nt = namedtuple(obs_names, obs)
end

#========================================================#

function unpack(nt::NamedTuple{N,T}) where {N,T}
     return collect(nt), collect(keys(nt))
end

function convert_to_bins(
    measurements::Array{<:AbstractMeasurement, 1},
    measurement_keys::Array{Symbol, 1}
)
    binned_measurements = reduce(vcat, convert_to_bins.(measurements))
    binned_meas_keys = reduce(vcat, keys_of_bins.(measurements, measurement_keys))

    if isa(binned_measurements, Measurement) # when only one measurement is active
        return [binned_measurements], [binned_meas_keys]
    end

    return binned_measurements, binned_meas_keys
end

function convert_to_bins(m::Measurement)
    return m
end

function convert_to_bins(md::MeasurementDistribution)
    nbins = length(md.value)
    uncertainties = [[u[i] for u in md.uncertainties] for i in 1:nbins]

    unc_nt = namedtuple.(Ref(keys(md.uncertainties)), uncertainties)
    return Measurement.(md.observable, md.value, unc_nt, md.active)
end

function only_active_measurements(
    measurements::Array{Measurement, 1},
    measurement_keys::Array{Symbol, 1},
    correlations::Array{<:AbstractCorrelation, 1}
)
    active_idxs = filter(x->measurements[x].active, 1:length(measurements))
    measurements = measurements[active_idxs]
    measurement_keys = measurement_keys[active_idxs]

    corrs = get_correlations.(correlations, Ref(active_idxs))

    return measurements, measurement_keys, corrs
end

#TODO: rename?
function get_correlations(
    correlation::Correlation,
    idxs::Array{<:Integer, 1}
)
    n=length(idxs)
    # ignore size of given correlation matrix if active==false
    if correlation.active == false
        return Correlation(Symmetric(zeros(n, n)), false)
    end

    return Correlation(Symmetric(correlation.matrix[idxs, idxs]), correlation.active)
end

function get_correlations(
    correlation::NoCorrelation,
    idxs::Array{<:Integer, 1}
)
    n = length(idxs)
    return Correlation(Symmetric(Matrix{Float64}(I, n, n)), correlation.active)
end

#--------------------
function create_distributions(
    m::NamedTuple{<:Any, <:Tuple{Vararg{AbstractMeasurement}}},
    uncertainty_keys::Vector{Symbol}
)
    active_idxs = [i for i in 1:length(m) if (isa(m[i], MeasurementDistribution) && any(m[i].active))]
    dists = [only_active_bins(m[i], uncertainty_keys) for i in active_idxs]
    dists_keys = [keys(m)[i] for i in active_idxs]

    if isempty(dists)
        return NamedTuple()
    else
        return namedtuple(dists_keys, dists)
    end
end

function only_active_bins(md::MeasurementDistribution, uncertainty_keys::Vector{Symbol})
    obs = md.observable[md.active]
    vals = md.value[md.active]

    unc = [md.uncertainties[u][md.active] for u in uncertainty_keys]
    uncertainties = namedtuple(uncertainty_keys, unc)
    bin_names = md.bin_names[md.active]

    return MeasurementDistribution(obs, vals, uncertainties=uncertainties, bin_names=bin_names)
end

function keys_of_bins(m::Measurement, key::Symbol)
    return key
end

function keys_of_bins(md::MeasurementDistribution, key::Symbol)
    return [Symbol(String(key)*"_"*String(b)) for b in md.bin_names]
end

function only_active_uncertainties(
    measurements::Array{Measurement, 1},
    correlations::Array{<:AbstractCorrelation, 1},
    uncertainty_keys::Array{Symbol, 1}
)
    active_idxs = filter(x->correlations[x].active, 1:length(correlations))
    active_correlations = correlations[active_idxs]
    uncertainty_keys = uncertainty_keys[active_idxs]

    measurements = [@set m.uncertainties = select(m.uncertainties, uncertainty_keys) for m in measurements]

    return measurements, active_correlations, uncertainty_keys
end

function only_active_nuisances(
    nuisances::Nothing,
    active_meas_keys::Array{Symbol, 1},
    unc_keys::Array{Symbol, 1}
)
    return nothing
end

function only_active_nuisances(
    nuisances::NamedTuple{<:Any, <:Tuple{Vararg{NuisanceCorrelation}}},
    active_meas_keys::Array{Symbol, 1},
    unc_keys::Array{Symbol, 1}
)
    nuisance_vec, nuisance_keys = unpack(nuisances)

    active_nuisances = NuisanceCorrelation[]
    active_nuisance_keys = Symbol[]

    for (nui, nui_k) in zip(nuisance_vec, nuisance_keys)
        if check_nuisance(nui, active_meas_keys, unc_keys)
            push!(active_nuisances, nui)
            push!(active_nuisance_keys, nui_k)
        end
    end

    if length(active_nuisances) > 0
        return namedtuple(active_nuisance_keys, active_nuisances)
    else
        return nothing
    end
end

function check_nuisance(
    nuisance::NuisanceCorrelation,
    active_meas_keys::Array{Symbol, 1},
    unc_keys::Array{Symbol, 1}
)
    if nuisance.meas1 == nuisance.meas2
        return false
    elseif !(nuisance.unc_key in unc_keys)
        return false
    elseif !(nuisance.meas1 in active_meas_keys)
        return false
    elseif !(nuisance.meas2 in active_meas_keys)
        return false
    end

    return true
end

function get_total_covariance(m::EFTfitterModel)
    covs = get_covariances(m)
    total_cov = Symmetric(sum(covs))

    return total_cov
end

function get_covariances(m::EFTfitterModel)
    unc_values = [[Float64(meas.uncertainties[u]) for meas in m.measurements] for u in keys(m.correlations)]
    corrs = [c.matrix for c in m.correlations]

    covs = [Symmetric(σ*ρ*σ) for (σ, ρ) in zip(diagm.(unc_values), corrs)]
end

function Nbins(measurements::NamedTuple{<:Any, <:Tuple{Vararg{AbstractMeasurement}}})
    sum([length(m.value) for m in measurements])
end

function get_indices(measurements, symb)
    meas_vec, meas_keys = unpack(measurements)
    names = keys_of_bins(meas_vec, meas_keys)
    ind = [findfirst(x -> x==symb, names)]

    if ind[1] != nothing
        return ind
    else
        ind = [findfirst(x -> x==Symbol(String(symb)*"_bin$i"), names) for i in 1:length(measurements[symb].value)]
    end
end

"""
    to_correlation_matrix(
        measurements::NamedTuple{<:Any, <:Tuple{Vararg{AbstractMeasurement}}},
        correlations::Tuple{Symbol,Symbol, Union{<:Real, Array{<:Real, 2}}}...
    )

    Returns a `Matrix{Float64}` as the correlation matrix for the measurements with all diagonal-elements being unity.
    The correlations are specified by passing tuples of two `Symbols` (the keys of the `Measurements`) with a value or matrix for the correlation coefficients.
    If the matrix is not positive-definite, a warning is shown but the matrix is still returned.

    Example:
    ```julia
    dist_corr = [1.0 0.5 0.0;
                 0.5 1.0 0.0;
                 0.0 0.0 1.0]

    another_corr_matrix = to_correlation_matrix(
        measurements,
        (:Meas1, :Meas2, 0.4),
        (:Meas1, :MeasDist, 0.1),
        (:MeasDist, :MeasDist, dist_corr),
        (:MeasDist_bin2, :MeasDist_bin3, 0.3),
    )
    ```
"""
function to_correlation_matrix(
    measurements::NamedTuple{<:Any, <:Tuple{Vararg{AbstractMeasurement}}},
    correlations::Any...#Tuple{Symbol,Symbol, Union{<:Real, Array{<:Real, 2}}}... #TODO: typing
)
    nmeas = Nbins(measurements)
    corr = Matrix{Float64}(I, nmeas, nmeas)

    for c in correlations #TODO: split into functions
        if isa(c, Tuple{Symbol,Symbol, Union{<:Real, Array{<:Real, 2}}})
            i = get_indices(measurements, c[1])
            j = get_indices(measurements, c[2])

            corr[i, j] .= c[3]
            corr[j, i] .= c[3]

        elseif isa(c, Tuple{Any, Union{<:Real, Array{<:Real, 2}}}) #TODO: add testing
            i = reduce(vcat, get_indices.(Ref(measurements), c[1]))
            corr[i, i] .= c[2]
        end
    end

    corr[diagind(corr)] .= 1

    if !isposdef(corr)
        @warn "The correlation matrix $corr is not positive semidefinite"
    end

    return corr
end

function has_nuisance_correlations(m::EFTfitterModel)
    m.nuisances == nothing ? (return false) : (return true)
end

function add_nuisance_parameters(
    parameters::NamedTupleDist,
    nuisances_nt::NamedTuple{<:Any, <:Tuple{Vararg{NuisanceCorrelation}}}
)
    param_keys = [keys(parameters)...,  keys(nuisances_nt)...]

    nuisances = [nui.prior for nui in nuisances_nt]
    param_dists = [values(parameters)...,  nuisances...]
    nt = namedtuple(param_keys, param_dists)

    return BAT.NamedTupleDist(nt)
end

function add_nuisance_parameters(
    parameters::NamedTupleDist,
    nuisances_nt::Nothing
)
    return parameters
end

# Todo: rename?
function keys_of_bins(
    measurements::Array{T, 1},
    measurement_keys::Array{Symbol, 1}
) where T <: AbstractMeasurement
    keys = reduce(vcat, keys_of_bins.(measurements, measurement_keys))
    return keys
end


# Naive implementation

struct DQExperience{N <: Real,T <: Real, A<:AbstractArray{T}}
    s::A
    a::N
    r::T
    sp::A
    done::Bool
end

function Base.convert(::Type{DQExperience{Int32, Float32, C}}, x::DQExperience{A, B, C}) where {A, B, C}
    return DQExperience{Int32, Float32, C}(convert(C, x.s),
                                            convert(Int32, x.a),
                                            convert(Float32, x.r),
                                            convert(C, x.sp),
                                            x.done)
end

mutable struct PrioritizedReplayBuffer{N<:Integer, T<:AbstractFloat,CI,Q,A<:AbstractArray{T}}
    max_size::Int64
    batch_size::Int64
    rng::AbstractRNG
    α::Float32
    β::Float32
    ϵ::Float32
    _curr_size::Int64
    _idx::Int64
    _priorities::Vector{T}
    _experience::Vector{DQExperience{N,T,Q}}

    _s_batch::A
    _a_batch::Vector{CI}
    _r_batch::Vector{T}
    _sp_batch::A
    _done_batch::Vector{T}
    _weights_batch::Vector{T}
end

function PrioritizedReplayBuffer(env::AbstractEnvironment{OV},
                                max_size::Int64,
                                batch_size::Int64;
                                rng::AbstractRNG = MersenneTwister(0),
                                α::Float32 = 6f-1,
                                β::Float32 = 4f-1,
                                ϵ::Float32 = 1f-3) where {OV}
    s_dim = obs_dimensions(env)
    experience = Vector{DQExperience{Int32, Float32, OV}}(undef, max_size)
    priorities = Vector{Float32}(undef, max_size)
    _s_batch = zeros(Float32, s_dim..., batch_size)
    _a_batch = [CartesianIndex(0,0) for i=1:batch_size]
    _r_batch = zeros(Float32, batch_size)
    _sp_batch = zeros(Float32, s_dim..., batch_size)
    _done_batch = zeros(Float32, batch_size)
    _weights_batch = zeros(Float32, batch_size)
    return PrioritizedReplayBuffer(max_size, batch_size, rng, α, β, ϵ, 0, 1, priorities, experience,
                _s_batch, _a_batch, _r_batch, _sp_batch, _done_batch, _weights_batch)
end

is_full(r::PrioritizedReplayBuffer) = r._curr_size == r.max_size

max_size(r::PrioritizedReplayBuffer) = r.max_size

function add_exp!(r::PrioritizedReplayBuffer, expe::DQExperience, td_err::T=abs(expe.r)) where T
    @assert td_err + r.ϵ > 0.
    priority = (td_err + r.ϵ)^r.α
    r._experience[r._idx] = expe
    r._priorities[r._idx] = priority
    r._idx = mod1((r._idx + 1),r.max_size)
    if r._curr_size < r.max_size
        r._curr_size += 1
    end
end

function update_priorities!(r::PrioritizedReplayBuffer, indices::Vector{Int64}, td_errors::V) where V <: AbstractArray
    new_priorities = (abs.(td_errors) .+ r.ϵ).^r.α
    @assert all(new_priorities .> 0f0)
    r._priorities[indices] = new_priorities
end

function StatsBase.sample(r::PrioritizedReplayBuffer)
    @assert r._curr_size >= r.batch_size
    @assert r.max_size >= r.batch_size # could be checked during construction
    sample_indices = sample(r.rng, 1:r._curr_size, Weights(r._priorities[1:r._curr_size]), r.batch_size, replace=false)
    return get_batch(r, sample_indices)
end

function get_batch(r::PrioritizedReplayBuffer, sample_indices::Vector{Int64})
    @assert length(sample_indices) == size(r._s_batch)[end]
    for (i, idx) in enumerate(sample_indices)
        @inbounds begin
            r._s_batch[.., i] = vec(r._experience[idx].s)
            r._a_batch[i] = CartesianIndex(r._experience[idx].a, i)
            r._r_batch[i] = r._experience[idx].r
            r._sp_batch[.., i] = vec(r._experience[idx].sp)
            r._done_batch[i] = r._experience[idx].done
            r._weights_batch[i] = r._priorities[idx]
        end
    end
    p = r._weights_batch ./ sum(r._priorities[1:r._curr_size])
    weights = (r._curr_size * p).^(-r.β)
    return r._s_batch, r._a_batch, r._r_batch, r._sp_batch, r._done_batch, sample_indices, weights
end

function populate_replay_buffer!(replay::PrioritizedReplayBuffer,
                                 env::AbstractEnvironment,
                                 action_indices;
                                 max_pop::Int64=replay.max_size, max_steps::Int64=100,
                                 policy::Policy = RandomPolicy(env.problem))
    o = reset!(env)
    done = false
    step = 0
    for t=1:(max_pop - replay._curr_size)
        a = action(policy, o)
        ai = action_indices[a]
        op, rew, done, info = step!(env, a)
        exp = DQExperience(o, ai, Float32(rew), op, done)
        add_exp!(replay, exp, abs(Float32(rew))) # assume initial td error is r
        o = op
        # println(o, " ", action, " ", rew, " ", done, " ", info) #TODO verbose?
        step += 1
        if done || step >= max_steps
            o = reset!(env)
            done = false
            step = 0
        end
    end
    @assert replay._curr_size >= replay.batch_size
end


module YaoLang

using LinearAlgebra
using YaoAPI

include("runtime/locations.jl")
include("runtime/generic_circuit.jl")

module Compiler

using TimerOutputs
const to = TimerOutput()
timings() = (TimerOutputs.print_timer(to); println())
enable_timings() = (TimerOutputs.enable_debug_timings(Compiler); return)

using ExprTools
using IRTools
using IRTools.Inner
using YaoAPI

using IRTools.Inner:
    IR, Block, BasicBlock, Statement, Variable, block, blocks, tab, print_stmt, printargs
using Base.Meta: ParseError

import ..YaoLang
using YaoLang: AbstractLocations, merge_locations, Locations, CtrlLocations, Circuit, PrimitiveCircuit

include("compiler/parse.jl")
include("compiler/ir.jl")
include("compiler/print.jl")
include("compiler/codegen.jl")
include("compiler/compiler.jl")
include("compiler/reflection.jl")
include("compiler/primitive.jl")
include("compiler/utils.jl")
include("compiler/kwdefs.jl")
include("compiler/validation.jl")
include("compiler/trace.jl")

function __init__()
    TimerOutputs.reset_timer!(to)
    # not sure why this doesn't work inside the module
    IRTools.Inner.printers[:quantum] = function (io, ex)
        get(printers, ex.args[1], print)(io, ex)
    end
end
end

using .Compiler
export @device, @primitive, @code_yao, @quantum
include("runtime/primitives.jl")

end # module


"""
# TauPy

Use the Python package ObsPy to calculate seismic travel times and ray paths for
teleseismic phases.

## Phases defined only by event depth and distance

For this case, functions return vectors of `Phase`s, which only contain details of each
phase in terms of distance from the event and radius.  Two methods are available:

## Phases defined by geographic coordinates

Here, functions give vectors of `PhaseGeog`s, which include information about the
geographic location within the Earth of the event, station, and potentially ray
path between them.

## Functions

In the case of either 1D or geographical phases, two functions are available:

- `travel_time` returns `Phase` objects which contain only time and slowness information.
- `path` returns `Phase`s which also contain the ray path.

Each have two methods, one where only depth and distance are specified, and the other
for when one wants also to include geographic information.

There is also the convenience function `turning_depth`, which gives the turning
depth of a `Phase` or `PhaseGeog` when the path has been calculated, and throws
an error when it has not.
"""
module TauPy

export
    AbstractPhase,
    Phase,
    PhaseGeog,
    path,
    travel_time,
    turning_depth

using PyCall

function __init__()
    # This is just to make sure geographiclib is installed for geographic ray paths
    pyimport_conda("geographiclib", "geographiclib", "conda-forge")
    copy!(Taup, pyimport_conda("obspy.taup", "obspy", "conda-forge"))
    for m in AVAILABLE_MODELS
        global MODEL[m] = Taup.TauPyModel(m)
        global RADIUS[m] = MODEL[m].model.radius_of_planet
    end
end

const Taup = PyNULL()

const TauPyFloat = Float64
"All available models for us in TauPy"
const AVAILABLE_MODELS = ("1066a", "1066b", "ak135", # "ak135f", # FIXME: Doesn't work
                          "ak135f_no_mud", "herrin",
                          "iasp91", "jb", "prem", "pwdk", "sp6")
"Default model used by TauPy"
const DEFAULT_MODEL = "ak135"
"Dictionary of pre-instantiated models, where keys are the model names"
const MODEL = Dict{String,PyObject}()
"Dictionary of the radius (km) of models"
const RADIUS = Dict{String,TauPyFloat}()

"""
    AbstractPhase

Abstract type from which `Phase` and `PhaseGeog` subtype.  Methods which may take
either `Phase`s or `PhaseGeog`s should specify their arguments as `AbstractPhase`s.
"""
abstract type AbstractPhase end

function Base.:(==)(p1::AbstractPhase, p2::AbstractPhase)
    p1 === p2 && return true
    typeof(p1) == typeof(p2) || return false
    for f in fieldnames(typeof(p1))
        if getfield(p1, f) isa AbstractVector
            length(getfield(p1, f)) == length(getfield(p2, f)) || return false
            all((a,b) -> a==b, zip(getfield(p1, f), getfield(p2, f))) ||
                return false
        else
            getfield(p1, f) == getfield(p2, f) || return false
        end
    end
    true
end

include("phase.jl")
include("phase_geog.jl")
include("memo.jl")

"""
    available_models()

Return the models which are available for use in TauPy.
"""
available_models() = AVAILABLE_MODELS

"""
    turning_depth(p::AbstractPhase) -> depth

Return the turning `depth` (km) for a phase which has had its ray path calculated.
Throw an error if the path has not been calculated.
"""
turning_depth(p::AbstractPhase) = length(p.radius) > 0 ? RADIUS[p.model] - minimum(p.radius) :
    error("No path calculated for phase")

"""
    _call_taup(taup_function, model, args...; cache=true) -> arrivals

Generic routine to call the Obspy taup module and process the arguments into a vector
of `Phase`s or `PhaseGeog`s.

`taup_function` is a `Symbol` giving the name of the obspy.taup function to call.
"""
function _call_taup(func::Symbol, model, args...; cache=true)
    key = cache_args_key(func, model, args...)
    arr = if cache && haskey(RAY_CACHE, key)
        get_cache(key)
    else
        arr = _phases_from_arrivals(getproperty(MODEL[model], func)(args...), model, args...)
        cache && update_cache!(key, arr)
        arr
    end
end

end # module


import java.security.KeyPairGenerator;
import java.security.KeyPair;
import java.security.PublicKey;
import java.security.PrivateKey;

public class Test {
    public static void main(String args[]) throws Exception {
	KeyPairGenerator kpg = KeyPairGenerator.getInstance("RSA");
	kpg.initialize(512); // 512 is the keysize
	KeyPair kp = kpg.generateKeyPair();

	PrivateKey PRIVKEY(H){L} privkey = (PRIVKEY(H){L})kp.getPrivate();
	PublicKey PUBKEY(H){L} pubkey = (PUBKEY(H){L})kp.getPublic();

	byte[]{L} bytes = new byte[10];
	bytes[0] = ({H})((byte)42);
	try {
	    byte[] out = CryptoLibrary.encrypt(pubkey, bytes);
	} catch (Exception e) {}
    }
}


@inline function allequal(x)
    length(x) < 2 && return true
    first_el = x[1]
    i = 2
    @inbounds for i=2:length(x)
        x[i] == first_el || return false
    end
    return true
end

"""
    vcat(A::SnpArray...;des="tmp_vcat_" * string(vcat_counter))

Concatenate SnpArray along dimension 1. (i.e. stack persons)
Only creates ".bed" file. Use vcat(A::SnpData...) for ".fam" and ".bim" files.
"""

arr_vcat_counter = 1
function Base.vcat(A::SnpArray...; des="tmp_vcat_arr_" * string(arr_vcat_counter))
    global arr_vcat_counter
    if des == "tmp_vcat_arr_" * string(arr_vcat_counter)
        arr_vcat_counter = arr_vcat_counter + 1
    end
    num_rows = [size(x)[1] for x in A]
    num_cols = [size(x)[2] for x in A]

    @assert allequal(num_cols) "number of columns are not the same"
    des_rows = sum(num_rows)
    des_cols = num_cols[1]
    cum_rows = Base.vcat([0], cumsum(num_rows))

    desbedfile = des * ".bed"

    makestream(desbedfile, "w+") do io
        write(io, 0x1b6c)
        write(io, 0x01)
        write(io, Matrix{UInt8}(undef, (des_rows + 3) >> 2, des_cols))
    end
    des_arr = SnpArray(desbedfile, des_rows, "r+")

    for i in 1:length(A)
        des_arr[(cum_rows[i] + 1): cum_rows[i + 1], :] .= @view A[i][:, :]
    end
    des_arr
end

"""
    hcat(A...;des="tmp_hcat_" * string(hcat_counter))

Concatenate SnpData along dimension 2. (i.e. stack locations)
Only creates ".bed" file. Use hcat(A::SnpData...) for ".fam" and ".bim" files.
"""
arr_hcat_counter = 1
function Base.hcat(A::SnpArray...; des="tmp_hcat_arr_" * string(arr_hcat_counter))
    global arr_hcat_counter
    if des == "tmp_hcat_arr_"*string(arr_vcat_counter)
        arr_hcat_counter = arr_hcat_counter + 1
    end
    num_rows = [size(x)[1] for x in A]
    num_cols = [size(x)[2] for x in A]

    @assert allequal(num_rows) "number of rows are not the same"
    des_rows = num_rows[1]
    des_cols = sum(num_cols)
    cum_cols = Base.vcat([0], cumsum(num_cols))

    desbedfile = des * ".bed"

    makestream(desbedfile, "w+") do io
        write(io, 0x1b6c)
        write(io, 0x01)
        write(io, Matrix{UInt8}(undef, (des_rows + 3) >> 2, des_cols))
    end
    des_arr = SnpArray(desbedfile, des_rows, "r+")

    des_arr.data .= hcat([s.data for s in A]...)
    des_arr
end

"""
   hvcat(rows::Tuple{Vararg{Int}}, values...; des="tmp_hvcat_arr_" * string(arr_hvcat_counter))

Horizontal and vertical concatenation in one call.
Only creates ".bed" file. Use hvcat(rows::Tuple{Vararg{Int}}, values::SnpData...) for ".fam" and ".bim" files.
"""
arr_hvcat_counter = 1
function Base.hvcat(rows::Tuple{Vararg{Int}}, A::SnpArray...; des="tmp_hvcat_arr_" * string(arr_hvcat_counter))
    global arr_hvcat_counter
    if des == "tmp_hvcat_arr_"*string(arr_hvcat_counter)
        arr_hvcat_counter = arr_hvcat_counter + 1
    end

    num_block_rows = length(rows)

    num_cols = 0
    for i = 1:rows[1]
        num_cols += size(A[i], 2)
    end

    num_rows = 0
    a = 1
    for i = 1:num_block_rows
        num_rows += size(A[a], 1)
        a += rows[i]
    end

    desbedfile = des * ".bed"

    makestream(desbedfile, "w+") do io
        write(io, 0x1b6c)
        write(io, 0x01)
        write(io, Matrix{UInt8}(undef, (num_rows +3) >> 2, num_cols))
    end
    des_arr = SnpArray(desbedfile, num_rows, "r+")

    a = 1
    r = 1
    for i = 1:num_block_rows
        c = 1
        szi = size(A[a], 1)
        for j = 1:rows[i]
            Aj = A[a+j-1]
            szj = size(Aj, 2)
            if size(Aj, 1) != szi
                throw(ArgumentError("mismatched height in block row $(i) (expected $szi, got $(size(Aj, 1)))"))
            end
            if c-1+szj > num_cols
                throw(ArgumentError("block row $(i) has mismatched number of columns (expected $num_cols, got $(c-1+szj)"))
            end
            des_arr[r:r-1+szi, c:c-1+szj] .= @view Aj[:,:]
            c += szj
        end
        if c != num_cols+1
            throw(ArgumentError("block row $(i) has mismatched number of columns (expected $num_cols, got $(c-1))"))
        end
        r += szi
        a += rows[i]
    end
    des_arr
end


function get_initial_conditions_template(model::OperationModel)
    # This is done to avoid passing the duals but also not re-allocating the PTDF when it
    # exists
    network_model = NetworkModel(
        get_network_formulation(model.template);
        use_slacks=get_use_slacks(get_network_model(model.template)),
        PTDF=get_PTDF(get_network_model(model.template)),
    )

    ic_template = ProblemTemplate(network_model)
    for device_model in values(model.template.devices)
        base_model = get_initial_conditions_device_model(model, device_model)
        base_model.use_slacks = device_model.use_slacks
        base_model.time_series_names = device_model.time_series_names
        base_model.attributes = device_model.attributes
        set_device_model!(ic_template, base_model)
    end
    for device_model in values(model.template.branches)
        base_model = get_initial_conditions_device_model(model, device_model)
        base_model.use_slacks = device_model.use_slacks
        base_model.time_series_names = device_model.time_series_names
        base_model.attributes = device_model.attributes
        set_device_model!(ic_template, base_model)
    end
    ic_template.services = model.template.services
    return ic_template
end

function build_initial_conditions_model!(model::T) where {T <: OperationModel}
    model.internal.ic_model_container = deepcopy(get_optimization_container(model))
    ic_settings = deepcopy(model.internal.ic_model_container.settings)
    main_problem_horizon = get_horizon(ic_settings)
    # TODO: add an interface to allow user to configure initial_conditions problem
    model.internal.ic_model_container.JuMPmodel = _make_jump_model(ic_settings)
    template = get_initial_conditions_template(model)
    model.internal.ic_model_container.settings = ic_settings
    set_horizon!(ic_settings, min(INITIALIZATION_PROBLEM_HORIZON, main_problem_horizon))
    init_optimization_container!(
        model.internal.ic_model_container,
        get_network_formulation(get_template(model)),
        get_system(model),
    )
    TimerOutputs.@timeit BUILD_PROBLEMS_TIMER "Build Initialization $(get_name(model))" begin
        build_impl!(model.internal.ic_model_container, template, get_system(model))
    end
    return
end