text:
include("kernels.jl")
include("dynamicdata/DynamicData.jl")
include("corpusutils/CorpusUtils.jl")
module GeneralizedDynamicTopicModels
using GeneralizedDynamicTopicModels, Distributions, Plots, Kernels, DynamicData, CorpusUtils, SpecialFunctions
const DocumentData = Union{CorpusUtils.DocumentList, DynamicData.DynamicDataSet}
type GDTM
	m_0::Float64 #prior mean
	s_0::Float64 #prior variance
	s_x::Float64 #measurement variance
	invS #inverse variance
	eta_1::Vector{Matrix{Float64}} #natural parameters to the topic normals
	eta_2::Matrix{Matrix{Float64}}
	alpha::Float64 #alpha prior for dirichlets
	corpus::CorpusUtils.Corpus #the data
	training_doc_ids::Vector{Int64} #ids of the training documents
	validation_doc_ids::Vector{Int64} # ids of validation documents
	test_doc_ids::Vector{Int64} # ids of test documents
	T::Int64 #number of time points
	D::Int64 #number of (training) documents
	K::Int64 #number of topics
	V::Int64 #number of words
	batch_size::Int64 #size of minibatch

	lambda::Matrix{Float64} #variational parameter to the dirichlets
	phi::Vector{Matrix{Float64}} #variational parameter to the multinomials (selecting the source distribution)
	zeta::Matrix{Float64} #variational parameter for bounding the intractable expectation caused by softmax fct
	suff_stats_tk::Vector{Vector{Float64}} #suffstats
	suff_stats_tkx::Vector{Matrix{Float64}} #suffstats
	means::Vector{Matrix{Float64}} #variational means
	s::Matrix{Matrix{Float64}} #variational variance
	timestamps
	times::Vector{Float64} #time points
	visualize::Bool # switch on or off any visualization
	inducing_points::Vector{Float64} #inducing point locations for sparse GP
	Kmm::Hermitian{Float64, Matrix{Float64}} #inducing point covariance for sparse GP
	KmmInv::Hermitian{Float64, Matrix{Float64}} #inverse
	Knn::Hermitian{Float64, Matrix{Float64}} #full rank covariance for GP models
	KnnInv::Hermitian{Float64, Matrix{Float64}} #inverse
	Knm::Matrix{Float64} #cross covariance training points - inducing points for sparse GP
	KnmKmmInv::Matrix{Float64} #cross covariance x inverse inducing point covariance, to save computation
	K_tilde::Hermitian{Float64, Matrix{Float64}} #low rank approximation of full rank covariance (sparse GP)
	K_tilde_diag::Vector{Float64} #diagonal of K_tilde
	S_diags::Vector{Matrix{Float64}} #variational covariance diagonals
	Λ_diags::Vector{Matrix{Float64}} #helper for storing Λ diagonals
	μ::Vector{Matrix{Float64}} #inducing point values in sparse GP
	krn::Kernels.Kernel #kernel object for GP models
	likelihood_counts::Vector{Float64} #list of #docs seen for likelihoods (ELBO estimates)
	test_counts::Vector{Float64} #same but for test set likelihoods
	likelihoods::Vector{Float64} #measured ELBO estimates
	test_perplexities::Vector{Float64} #measured test set likelihoods
	learning_rates::Vector{Float64} #computed learning rates
	word_observation_times::Vector{Vector{Int64}}
	jitter::Diagonal{Float64}
	use_seeding::Bool
	function GDTM(corpus::CorpusUtils.Corpus, num_topics::Int64, batch_size::Int64,
					m_0::Float64, s_0::Float64, s_x::Float64, alpha::Float64, visualize::Bool;
					training_doc_ids::Vector{Int64}=Int64[], validation_doc_ids::Vector{Int64}=Int64[],
					test_doc_ids::Vector{Int64}=Int64[], seeded::Bool = false)
		n = DynamicData.length(corpus.documents)
		m = new()
		m.batch_size = batch_size
		m.timestamps = DynamicData.get_unique_timestamps(corpus.documents)
		times = map(Dates.datetime2julian,m.timestamps)
		times -= minimum(times)-1
		m.times = times
		m.T = length(times)
		m.visualize = visualize
		m.use_seeding = seeded

		#just for testing
		suggested_variance = 1./(times[end] - times[1])
		#the data
		m.corpus = corpus
		#if nothing set, assume we need only test data
		if length(training_doc_ids) == 0
			r = randperm(n)
			test_doc_num = min(round(Int64, n*.1), 10000)
			info("creating test set, using $test_doc_num for testing, $(length(r)-test_doc_num) for training")
			m.test_doc_ids = r[1:test_doc_num]
			m.training_doc_ids = r[test_doc_num+1:end]
		else
			m.training_doc_ids = training_doc_ids
		end
		if length(test_doc_ids) != 0
			m.test_doc_ids = test_doc_ids
		end
		m.D = length(m.training_doc_ids)
		m.K = num_topics
		m.V = length(corpus.lexicon)
		m.word_observation_times = Vector{Vector{Int64}}(m.V)
		[m.word_observation_times[w] = collect(CorpusUtils.used_in_times(m.corpus, m.corpus.lexicon.words[w])) for w in 1:m.V]

		#model params
		m.m_0 = m_0
		m.s_0 = s_0

		#measurement noise
		m.s_x = s_x
		#document dirichlet param
		m.alpha = alpha
		reset(m)
		Plots.plotly()
		println("done constructor")
		m
	end
end

include("utils.jl")
include("gpdtm_svi_inference.jl")
include("visualize.jl")

function reset(m::GDTM)
	#stats
	m.likelihood_counts = Float64[]
	m.test_counts = Float64[]
	m.likelihoods = Float64[]
	m.test_perplexities = Float64[]
	m.learning_rates = Float64[]

	m.timestamps = DynamicData.get_unique_timestamps(m.corpus.documents)
	times = map(Dates.datetime2julian,m.timestamps)
	times -= minimum(times)-1
	m.times = times
	m.T = length(times)
	# reset zeta parameter
	m.zeta = zeros(m.K, m.T)
	# garbage collection
	gc()
end

"""
The update step for local, i.e. document specific variational parameters. Computes the likelihood of the processed documents on the fly.
"""
function e_step(m::GDTM, data_idx::Vector{Int64})
	e_step_ll = 0
	words_seen = Int64[]
	[words_seen = union(words_seen, map(w->w.id,collect(keys(m.corpus.documents[doc_id].value)))) for doc_id in data_idx]

	#reset sufficient statistics for the minibatch
	ɸ = Array(Array{Float64, 1},m.K)
	Ξ = Array(Array{Float64, 2},m.K)
	for k in 1:m.K
		ɸ[k] = zeros(m.T)
		Ξ[k] = zeros(m.T, m.V)
	end
	timestamps_seen = IntSet()
	doc_topic_proportions = zeros(length(data_idx), m.K)
	for (i, doc_idx) in enumerate(data_idx)
		(doc_ll, λ) = doc_e_step(m, doc_idx, ɸ, Ξ, timestamps_seen)
		e_step_ll += doc_ll
		doc_topic_proportions[i,:] = λ
	end
	(e_step_ll, ɸ, Ξ, words_seen, collect(timestamps_seen), doc_topic_proportions)
end

"""
Performs the e-step on individual documents.
"""
function doc_e_step(m::GDTM, doc_idx::Int64, ɸ::Vector{Vector{Float64}}, Ξ::Vector{Matrix{Float64}}, timestamps_seen::IntSet)
	doc = m.corpus.documents[doc_idx]
	#extract wird ids and frequencies from the document
	(doc_words, freqs) = CorpusUtils.get_words_and_counts(doc.value)
	# extract timestamp of the documents
	t_doc = DynamicData.get_timestamp_index(m.corpus.documents, doc.timestamp)
	push!(timestamps_seen, t_doc)
	# do the actual inference
	(φ,λ,vi_ll) = document_inference(m, t_doc, doc_words, freqs)
	#collect sufficient statistics
	for k in 1:m.K
		for (i,w) in enumerate(doc_words)
			ɸ[k][t_doc] += freqs[i]*φ[i,k]
			Ξ[k][t_doc, w] += freqs[i]*φ[i,k]
		end
	end
	(vi_ll, λ)
end

"""
Perform variational inference on document parameters, i.e. proportions of topics in the document and single word assignment probabilities.
"""
function document_inference(m::GDTM, t_doc::Int64, doc_words::Vector{Int64}, freqs::Vector{Int64})
	last_vi_ll = -1e100
	vi_ll = 0
	converged = 1
	#get document data
	N_d = length(doc_words)
	#init phi variable for doc
	φ = ones(length(doc_words), m.K) ./ m.K
	#init lambda for doc
	λ = ones(m.K) .* m.alpha .+ N_d./m.K
	dig_λ = digamma.(λ)
	dig_λ_sum = digamma(sum(λ))

	iter = 0
	means = zeros(m.K, N_d)
	# compute means for only those words that are observed in the document from inducing point means
	for k in 1:m.K
		means[k,:] = m.KnmKmmInv[t_doc,:]'*m.μ[k][:, doc_words]
	end
	# do inference
	while converged > 1e-3
		vi_ll = 0
		for i in 1:N_d
			log_phi_sum = 0
			for k in 1:m.K
				φ[i,k] = (dig_λ[k] - dig_λ_sum) + (means[k, i] - m.zeta[k, t_doc])

				log_phi_sum = k==1 ? φ[i,k] : log_add(φ[i,k], log_phi_sum)
			end
			#normalize and exp phi
			φ[i,:] = exp(φ[i,:] .- log_phi_sum) + 1e-100
		end
		#update lambda
		λ = m.alpha .+ vec(sum(freqs.*φ,1))
		dig_λ = digamma.(λ)
		dig_λ_sum = digamma(sum(λ))

		#compute the document likelihood given the current parameters
		vi_ll = compute_document_likelihood(m, t_doc, doc_words, freqs, φ, λ, means)
		#check convergence
		converged = (last_vi_ll - vi_ll)/last_vi_ll
		last_vi_ll = vi_ll
		iter += 1
	end
	(φ, λ, vi_ll)
end

"""
Computes the document likelihood given some paramter setting.
"""
function compute_document_likelihood(m::GDTM, t_doc::Int64, doc_words::Vector{Int64}, freqs::Vector{Int64}, φ::Matrix{Float64}, λ::Vector{Float64}, means::Matrix{Float64})
	likelihood = 0
	dig_λ = digamma.(λ)
	dig_λ_sum = digamma(sum(λ))
	for k in 1:m.K
		likelihood += sum(freqs .* φ[:,k] .* (dig_λ[k] .- dig_λ_sum .+ means[k,:] .- m.zeta[k,t_doc] - log(φ[:,k])))
	end
	likelihood += sum(lgamma.(λ)) - lgamma(sum(λ)) + sum( (m.alpha .- λ) .* (dig_λ .- dig_λ_sum))
	likelihood
end

"""
Do a heldout perplexity test. Learn topic proportions (given an optimized model) on the first half of a document, then compute the second half likelihood.
Report the per-word predictive perplexity on the second half.
"""
function test(m::GDTM)
	test_ll = 0
	#set a seed, so that we are always looking at the same document parts
	srand(12345)
	total_token_count = 0.
	for doc_id in m.test_doc_ids
		test_doc = m.corpus.documents[doc_id]
		t_doc = DynamicData.get_timestamp_index(m.corpus.documents, test_doc.timestamp)
		(test_doc_words, test_freqs) = CorpusUtils.get_words_and_counts(test_doc.value)
		#expand tokens
		test_doc_tokens = Int64[]
		for (i,w) in enumerate(test_doc_words)
			[push!(test_doc_tokens, w) for j in 1:test_freqs[i]]
		end
		n = length(test_doc_tokens)
		assert(n==sum(test_freqs))
		#split document
		r = randperm(n)
		n_2 = round(Int64, n/2)
		w_1 = test_doc_tokens[r][1:n_2]
		tf_1 = ones(Int64, n_2)
		w_2 = test_doc_tokens[r][n_2+1:end]
		tf_2 = ones(Int64, length(w_2))
		# find optimal topic distribution
		(_, theta, _) = document_inference(m, t_doc, w_1, tf_1)
		#compute phis for second half
		means = zeros(m.K, length(w_2))
		for k in 1:m.K
			means[k,:] = m.KnmKmmInv[t_doc,:]'*m.μ[k][:, w_2]
		end

		phi = zeros(length(w_2), m.K)
		dig_λ = digamma.(theta)
		dig_λ_sum = digamma(sum(theta))
		for i in 1:length(w_2)
			log_phi_sum = 0
			for k in 1:m.K
				phi[i,k] = (dig_λ[k] - dig_λ_sum) + (means[k, i] - m.zeta[k, t_doc])

				log_phi_sum = k==1 ? phi[i,k] : log_add(phi[i,k], log_phi_sum)
			end
			#normalize and exp phi
			phi[i,:] = exp(phi[i,:] - log_phi_sum) + 1e-100
		end
		doc_ll = compute_document_likelihood(m, t_doc, w_2, tf_2, phi, theta, means)
		test_ll += doc_ll
		total_token_count += length(w_2)
	end
	-exp(test_ll/total_token_count)
end

"""
Performs a grid search over the parameter space for various kernels (Brownian motion, Cauchy, Ornstein-Uhlenbeck, RBF),
runs the corresponding models and stores the results.
"""
function run_param_gridsearch(batch_size::Int64, m_0::Float64, s_0::Float64, s_x::Float64, corpus::CorpusUtils.Corpus; visualize::Bool=false, num_inducing_points::Int64=25, use_seeding::Bool=false, num_topics::Vector{Int64}=[5,10], alphas::Vector{Float64}=[.5])

	inference_function = (m) -> inference_svi_gp(m, num_inducing_points)

	for alpha in alphas, k in num_topics
		m = GDTM(corpus, k, batch_size, m_0, s_0, s_x, alpha, visualize, seeded=use_seeding)
		range = m.times[end] - m.times[1]
		# effective pre-factor for maximum variance of one in the Brownian motion kernel
		suggested_variance = 1./range
		# different variance parameters for Brownian motion
		wiener_params_mult = [.5, 1., 2.]
		# keep track of best variance in terms of perplexity
		best_nu = 0.
		best_nu_ll = 1e100
		for nu in wiener_params_mult
			m.krn = Kernels.BrownianKernel(nu*suggested_variance)
			inference_function(m)
			write_stats(m, base_dir_suffix="BM")
			if best_nu_ll > m.test_perplexities[end]
				best_nu_ll = m.test_perplexities[end]
				best_nu = nu
			end
			reset(m)
		end
		# minimum lengthscale is equal to space between inducing points
		min_lengthscale = range/num_inducing_points
		lengthscales = [2*range, range, .5*range, .1*range, .05*range]

		for l in lengthscales
			# make sure that maximum variance is as suggested
			m.krn = Kernels.CauchyKernel(best_nu, l)
			inference_function(m)
			write_stats(m, base_dir_suffix="Cauchy")
			reset(m)
		end

		for l in lengthscales
			m.krn = Kernels.OrnsteinUhlenbeckKernel(best_nu, l)
			inference_function(m)
			write_stats(m, base_dir_suffix="OU")
			reset(m)
		end
		for l in lengthscales
			m.krn = Kernels.RBFKernel(best_nu, l)
			inference_function(m)
			write_stats(m, base_dir_suffix="RBF")
			reset(m)
		end
	end
end

"""
Runs a single model based on the corpus, kernel, number of topics, Dirichlet hyperparameter alpha and optional parameters provided.
Stores the result in the current directory under subdirectory "experiments" and creates topic visualizations.
The optional parameters are:
* prior_mean, prior_variance: prior mean and variance for topics
* measurement_noise: measurement noise when drawing time marginals from the stochastic process
* minibatch_size: mini batch size for stochastic algorithm
* inducing_points: number of inducing points for sparse GP
* visualize: if set to true, create probability trajectory charts for top words in topics, report learning rate, e-step (based on the mini batch) and m-step (global optimization step) log likelihoods and test set log likelihoods graphically
* use_seeding: use random documents to seed topics, for each topic at each time step, words in a random document (with appropriate time stamp) are used to artificially increase the probabilities of them in the current topic
"""
function run_model(corpus::CorpusUtils.Corpus, kernel::Kernels.Kernel, num_topics::Int64, alpha::Float64; prior_mean::Float64 = 0., prior_variance::Float64 = 10.,  measurement_noise::Float64 = .5, minibatch_size::Int64 = 256, inducing_points::Int64 = 25, visualize::Bool = false, use_seeding::Bool = false)
	m = GDTM(corpus, num_topics, minibatch_size, prior_mean, prior_variance, measurement_noise, alpha, visualize, seeded=use_seeding)
	m.krn = kernel
	inference_svi_gp(m, inducing_points)
	out_dir = write_stats(m)
	if visualize
		make_topic_charts(out_dir, corpus, true, false)
	end
end

"""
Runs all three kernels on the provided dataset for different parameter settings.
The data is expected in the following format:
#{number of timestamps}
#{seconds since epoch, i.e. unix time time stamp 1}
#{number of documents in timestamp 1}
#{types in document 1} #{type id}:#{frequency} #{type id}:#{frequency} ...
...
#{types in document N_1} #{type id}:#{frequency} #{type id}:#{frequency} ...
#{seconds since epoch, i.e. unix time time stamp 2}
#{number of documents in timestamp 2}
#{types in document 1} #{type id}:#{frequency} #{type id}:#{frequency} ...
...
#{types in document N_2} #{type id}:#{frequency} #{type id}:#{frequency} ...
#{repeat for all timestamps}
"""
function main(corpus::CorpusUtils.Corpus; minibatch_size::Int64 = 256, prior_mean::Float64 = 0., prior_variance::Float64 = 10., measurement_noise::Float64 = .5, inducing_points::Int64 = 25)
	run_param_gridsearch(minibatch_size, prior_mean, prior_variance, measurement_noise, corpus, num_inducing_points=inducing_points)
end

end # module


text:
module ParallelKMeans

using StatsBase
import Base.Threads: @spawn, @threads

export kmeans

abstract type CalculationMode end

# Single thread class to control the calculation type based on the CalculationMode
struct SingleThread <: CalculationMode
end

# Multi threaded implementation to control the calculation type based avaialble threads
struct MultiThread <: CalculationMode
    n::Int
end

# Get the number of avaialble threads for multithreading implementation
MultiThread() = MultiThread(Threads.nthreads())

"""
    pairwise!(target, x, y, mode)

Let X and Y respectively have m and n columns. Then the `pairwise!` function
computes distances between each pair of columns in X and Y and store result
in `target` array. Argument `mode` defines calculation mode, currently
following modes supported
- SingleThread()
- MultiThread()
"""
pairwise!(target, x, y) = pairwise!(target, x, y, SingleThread())

function pairwise!(target, x, y, mode::SingleThread)
    ncol = size(x, 2)

    @inbounds for k in axes(y, 1)
        for i in axes(x, 1)
            target[i, k] = (x[i, 1] - y[k, 1])^2
        end

        for j in 2:ncol
            for i in axes(x, 1)
                target[i, k] += (x[i, j] - y[k, j])^2
            end
        end
    end
    target
end

"""
    divider(n, k)

Utility function, splits 1:n sequence to k chunks of approximately same size.
"""
function divider(n, k)
    d = div(n, k)
    xz = vcat(collect((0:k-1) * d), n)
    return [t[1]:t[2] for t in zip(xz[1:end-1] .+ 1, xz[2:end])]
end

function pairwise!(target, x, y, mode::MultiThread)
    ncol = size(x, 2)
    nrow = size(x, 1)

    ranges = divider(nrow, mode.n)
    waiting_list = Task[]

    for i in 1:length(ranges) - 1
        push!(waiting_list, @spawn inner_pairwise!(target, x, y, ranges[i]))
    end

    inner_pairwise!(target, x, y, ranges[end])

    for i in 1:length(ranges) - 1
        wait(waiting_list[i])
    end

    target
end

"""
    inner_pairwise!(target, x, y, r)

Utility function for calculation of [pairwise!(target, x, y, mode)](@ref) function.
UnitRange argument `r` select subarray of original design matrix `x` that is going
to be processed.
"""
function inner_pairwise!(target, x, y, r)
    ncol = size(x, 2)

    @inbounds for k in axes(y, 1)
        for i in r
            target[i, k] = (x[i, 1] - y[k, 1])^2
        end

        for j in 2:ncol
            for i in r
                target[i, k] += (x[i, j] - y[k, j])^2
            end
        end
    end
    target
end

"""
    smart_init(X, k; init="k-means++")

This function handles the random initialisation of the centroids from the
design matrix (X) and desired groups (k) that a user supplies.

`k-means++` algorithm is used by default with the normal random selection
of centroids from X used if any other string is attempted.

A tuple representing the centroids, number of rows, & columns respecitively
is returned.
"""
function smart_init(X::Array{Float64, 2}, k::Int, mode::T = SingleThread();
        init::String="k-means++") where {T <: CalculationMode}

    n_row, n_col = size(X)

    if init == "k-means++"

        # randonmly select the first centroid from the data (X)
        centroids = zeros(k, n_col)
        rand_indices = Vector{Int}(undef, k)
        rand_idx = rand(1:n_row)
        rand_indices[1] = rand_idx
        centroids[1, :] .= X[rand_idx, :]
        centroids[k, :] .= 0.0
        distances = Array{Float64}(undef, n_row, 1)
        new_distances = Array{Float64}(undef, n_row, 1)

        # TODO: Add `colwise` function (or use it from `Distances` package)
        # compute distances from the first centroid chosen to all the other data points
        first_centroid_matrix = convert(Matrix, centroids[1, :]')

        # flatten distances
        pairwise!(distances, X, first_centroid_matrix, mode)
        distances[rand_idx] = 0.0

        for i = 2:k
            # choose the next centroid, the probability for each data point to be chosen
            # is directly proportional to its squared distance from the nearest centroid
            r_idx = wsample(1:n_row, vec(distances))
            rand_indices[i] = r_idx
            centroids[i, :] .= X[r_idx, :]

            # no need for final distance update
            i == k && break

            # compute distances from the centroids to all data points
            current_centroid_matrix = convert(Matrix, centroids[i, :]')
            # new_distances = vec(pairwise(SqEuclidean(), X, current_centroid_matrix, dims = 1))
            pairwise!(new_distances, X, first_centroid_matrix, mode)

            # and update the squared distance as the minimum distance to all centroid
            for i in 1:n_row
                distances[i, 1] = distances[i, 1] < new_distances[i, 1] ? distances[i, 1] : new_distances[i, 1]
            end
            distances[r_idx, 1] = 0.0
        end

    else
        # randomly select points from the design matrix as the initial centroids
        rand_indices = rand(1:n_row, k)
        centroids = X[rand_indices, :]
    end

    return (centroids = centroids, indices = rand_indices)
end

"""
    sum_of_squares(x, labels, centre, k)

This function computes the total sum of squares based on the assigned (labels)
design matrix(x), centroids (centre), and the number of desired groups (k).

A Float type representing the computed metric is returned.
"""
function sum_of_squares(x::Array{Float64,2}, labels::Array{Int64,1}, centre::Array)
    s = 0.0

    @inbounds for j in axes(x, 2)
        for i in axes(x, 1)
            s += (x[i, j] - centre[labels[i], j])^2
        end
    end

    return s
end

"""
    Kmeans(design_matrix, k; k_init="k-means++", max_iters=300, tol=1e-4, verbose=true)

This main function employs the K-means algorithm to cluster all examples
in the training data (design_matrix) into k groups using either the
`k-means++` or random initialisation technique for selecting the initial
centroids.

At the end of the number of iterations specified (max_iters), convergence is
achieved if difference between the current and last cost objective is
less than the tolerance level (tol). An error is thrown if convergence fails.

Details of operations can be either printed or not by setting verbose accordingly.

A tuple representing labels, centroids, and sum_squares respectively is returned.
"""
function kmeans(design_matrix::Array{Float64, 2}, k::Int, mode::T = SingleThread();
                k_init::String = "k-means++", max_iters::Int = 300, tol = 1e-4, verbose::Bool = true, init = nothing) where {T <: CalculationMode}

    n_row, n_col = size(design_matrix)
    centroids = init == nothing ? smart_init(design_matrix, k, mode, init=k_init).centroids : init

    labels = Vector{Int}(undef, n_row)
    distances = Vector{Float64}(undef, n_row)
    centroids_cnt = Vector{Int}(undef, size(centroids, 1))

    J_previous = Inf64

    nearest_neighbour = Array{Float64, 2}(undef, size(design_matrix, 1), size(centroids, 1))
    # Update centroids & labels with closest members until convergence
    for iter = 1:max_iters
        pairwise!(nearest_neighbour, design_matrix, centroids, mode)

        @inbounds for i in axes(nearest_neighbour, 1)
            labels[i] = 1
            distances[i] = nearest_neighbour[i, 1]
            for j in 2:size(nearest_neighbour, 2)
                if distances[i] > nearest_neighbour[i, j]
                    labels[i] = j
                    distances[i] = nearest_neighbour[i, j]
                end
            end
        end

        centroids .= 0.0
        centroids_cnt .= 0
        @inbounds for i in axes(design_matrix, 1)
            centroids[labels[i], 1] += design_matrix[i, 1]
            centroids_cnt[labels[i]] += 1
        end
        @inbounds for j in 2:n_col
            for i in axes(design_matrix, 1)
                centroids[labels[i], j] += design_matrix[i, j]
            end
        end
        centroids ./= centroids_cnt

        # Cost objective
        J = mean(distances)

        if verbose
            # Show progress and terminate if J stopped decreasing.
            println("Iteration $iter: Jclust = $J.")
        end

        # Final Step: Check for convergence
        if (iter > 1) & (abs(J - J_previous) < (tol * J))

            sum_squares = sum_of_squares(design_matrix, labels, centroids)

            # Terminate algorithm with the assumption that K-means has converged
            if verbose
                println("Successfully terminated with convergence.")
            end

            return (labels=labels, centroids=centroids, sum_squares=sum_squares)

        elseif (iter == max_iters) & (abs(J - J_previous) > (tol * J))
            throw(error("Failed to converge Check data and/or implementation or increase max_iter."))

        end

        J_previous = J
    end
end

end # module

ext: jl


text:
# ellipse to rotated and shifted ellipse, initial and final landmark positions observed
using Revise
using Polynomials.PolyCompat
using BridgeLandmarks
const BL = BridgeLandmarks
using Distributions
using LinearAlgebra
using Random
using FileIO
using JLD2

# using RCall
# using Random
# using DataFrames
# using DelimitedFiles
# using CSV
# using StaticArrays
#using JLD2

Random.seed!(9)
workdir = @__DIR__
cd(workdir)
outdir = workdir

# read data
dat = load("data_exp1.jld2")
xobs0 = dat["xobs0"]
xobsT = dat["xobsT"]

# set pars
p_ms = Pars_ms(δmom=0.1, cinit=0.2, γinit=2.0, σobs = 0.01)
p_ahs = Pars_ahs(δmom=0.1, cinit=0.02, γinit=0.2, db=[2.5,2.5],stdev=0.75, σobs = 0.01)
ainit = mean(norm.(diff(xobs0)))/2.0
ups = [:innov, :mala_mom, :parameter]

# run algorithm
@time landmarksmatching(xobs0,xobsT; ITER=500,pars=p_ms, outdir=outdir, ainit=ainit, updatescheme=ups)
@time landmarksmatching(xobs0,xobsT; ITER=500,pars=p_ahs, outdir=outdir, ainit=ainit, updatescheme=ups)

@time landmarksmatching(xobs0,xobsT; ITER=5,pars=p_ahs, outdir=outdir, ainit=ainit, updatescheme=ups)
plotlandmarksmatching(outdir)


text:
using Lehmann
using Printf
using Gaston
β = 40.0
rtol = 1e-8
eta = 1e-2
N = 128

dlr = DLRGrid(β = β, Euv = 1.0, isFermi = true, rtol = rtol, symmetry = :none)
τgrid = collect(LinRange(0.0, β, N))
println(τgrid)
G = Sample.SemiCircle(dlr, :τ, τgrid)
G += rand(length(G)) / 2.0 * eta
coeff = tau2dlr(dlr, G, τgrid)
coeff_sumrule = tau2dlr(dlr, G, τgrid, sumrule = π / 2)
println("sum rule: ", abs(sum(coeff) - π / 2), " versus ", abs(sum(coeff_sumrule) - π / 2))

dlr10 = DLRGrid(β = β, Euv = 100.0, isFermi = true, rtol = rtol, symmetry = :none)

Gtrue = Sample.SemiCircle(dlr, :τ, dlr10.τ)
Gfit = real.(dlr2tau(dlr, coeff, dlr10.τ))
Gfit_sumrule = real.(dlr2tau(dlr, coeff_sumrule, dlr10.τ))

@printf("%15s%30s%30s%30s\n", "tau", "true", "no sumrule", "with sumrule")
for i in 1:dlr10.size
    @printf("%15.6f%30.15f%30.15e%30.15e\n", dlr10.τ[i], Gtrue[i], abs(Gfit[i] - Gtrue[i]), abs(Gfit_sumrule[i] - Gtrue[i]))
end

set(term = "qt")
# p = plot(dlr10.τ, Gtrue, label = "original", axis = "semilogy")
xrange = (0.0, 40)
p = plot(dlr10.τ, Gfit - Gtrue, plotstyle = :linespoints, leg = Symbol("DLR no sum rule"), Axes(xrange = xrange))
plot!(dlr10.τ, Gfit_sumrule - Gtrue, plotstyle = :linespoints, leg = Symbol("DLR with sum rule"))
display(p)
readline()


text:
@inline function calc_A_third_order(x::VInt16, ::Val{N}) where N
    p = 15; r = 1; A = Int16(23170); C = Int16(-54292 >> N)
    rounding = one(x) << (p - N - 1)
    (rounding + A + (x * (x * C) >> N) >> r) >> (p - N)
end

@inline function calc_B_third_order(x::VInt16, ::Val{N}) where N
    p = 14; r = 3; B = Int16(-18199); D = Int16(58049 >> N);
    rounding = one(x) << (p - N - 1)
    (rounding + x * (B + (x * (x * D) >> N) >> r) >> N) >> (p - N)
end

@inline function calc_A_third_order(x::VInt32, ::Val{N}) where N
    p = 31; r = 1; A = Int32(1518500249); C = Int32(-3558067408 >> N)
    rounding = one(x) << (p - N - 1)
    (rounding + A + (x * (x * C) >> N) >> r) >> (p - N)
end
@inline function calc_B_third_order(x::VInt32, ::Val{N}) where N
    p = 30; r = 3; B = Int32(-1192627308); D = Int32(3804335470 >> N);
    rounding = one(x) << (p - N - 1)
    (rounding + x * (B + (x * (x * D) >> N) >> r) >> N) >> (p - N)
end

@inline function calc_A(x::VInt16, bits::B) where B <: Union{Val{6}, Val{7}}
    calc_A_third_order(x, bits)
end
@inline function calc_B(x::VInt16, bits::B) where B <: Union{Val{6}, Val{7}}
    calc_B_third_order(x, bits)
end
@inline function calc_A(x::VInt32, bits::B) where B <: Union{Val{8}}
    calc_A_third_order(x, bits)
end
@inline function calc_B(x::VInt32, bits::B) where B <: Union{Val{8}}
    calc_B_third_order(x, bits)
end


text:
@testset "approx_posterior_gp" begin
    rng = MersenneTwister(123456)
    N_cond = 3
    N_a = 5
    N_b = 6

    # Specify prior.
    f = GP(sin, Matern32Kernel())

    # Sample from prior.
    x = collect(range(-1.0, 1.0; length=N_cond))
    fx = f(x, 1e-15)
    y = rand(rng, fx)

    # Construct posterior.
    f_post = posterior(fx, y)

    # Construct optimal approximate posterior.
    f_approx_post = approx_posterior(VFE(), fx, y, fx)

    # Verify that approximate posterior ≈ posterior at various inputs.
    x_test = randn(rng, 100)
    @test mean(f_post, x_test) ≈ mean(f_approx_post, x_test)
    @test cov(f_post, x_test) ≈ cov(f_approx_post, x_test)

    # Verify that AbstractGP interface is implemented fully and consistently.
    a = collect(range(-1.0, 1.0; length=N_a))
    b = randn(rng, N_b)
    TestUtils.test_internal_abstractgps_interface(rng, f_approx_post, a, b)

    @testset "update_approx_posterior (new observation)" begin
        rng = MersenneTwister(1)
        X = rand(rng, 10)
        y = rand(rng, 10)
        Z = rand(rng, 4)

        f = GP(SqExponentialKernel())

        # online learning
        p_fx1 = approx_posterior(VFE(), f(X[1:7], 0.1), y[1:7], f(Z))
        u_p_fx1 = update_approx_posterior(p_fx1, f(X[8:10], 0.1), y[8:10])

        # batch learning
        p_fx2 = approx_posterior(VFE(), f(X, 0.1), y, f(Z))

        @test u_p_fx1.data.m_ε ≈ p_fx2.data.m_ε atol = 1e-5
        @test u_p_fx1.data.Λ_ε.U ≈ p_fx2.data.Λ_ε.U atol = 1e-5
        @test u_p_fx1.data.U ≈ p_fx2.data.U atol = 1e-5
        @test u_p_fx1.data.α ≈ p_fx2.data.α atol = 1e-5
        @test u_p_fx1.data.z ≈ p_fx2.data.z atol = 1e-5
        @test u_p_fx1.data.b_y ≈ p_fx2.data.b_y atol = 1e-5
        @test u_p_fx1.data.B_εf ≈ p_fx2.data.B_εf atol = 1e-5
        @test u_p_fx1.data.x ≈ p_fx2.data.x atol = 1e-5
        @test u_p_fx1.data.Σy ≈ p_fx2.data.Σy atol = 1e-5
    end

    @testset "update_approx_posterior (new pseudo-points)" begin
        rng = MersenneTwister(1)
        X = rand(rng, 10)
        y = rand(rng, 10)
        Z1 = rand(rng, 4)
        Z2 = rand(rng, 3)
        Z = vcat(Z1, Z2)

        f = GP(SqExponentialKernel())

        # Sequentially adding pseudo points
        p_fx1 = approx_posterior(VFE(), f(X, 0.1), y, f(Z1))
        u_p_fx1 = update_approx_posterior(p_fx1, f(Z2))

        # Adding all pseudo-points at once
        p_fx2 = approx_posterior(VFE(), f(X, 0.1), y, f(Z))

        @test u_p_fx1.data.m_ε ≈ p_fx2.data.m_ε atol = 1e-5
        @test u_p_fx1.data.Λ_ε.U ≈ p_fx2.data.Λ_ε.U atol = 1e-5
        @test u_p_fx1.data.U ≈ p_fx2.data.U atol = 1e-5
        @test u_p_fx1.data.α ≈ p_fx2.data.α atol = 1e-2 rtol = 1e-2
        @test u_p_fx1.data.z ≈ p_fx2.data.z atol = 1e-5
        @test u_p_fx1.data.b_y ≈ p_fx2.data.b_y atol = 1e-5
        @test u_p_fx1.data.B_εf ≈ p_fx2.data.B_εf atol = 1e-5
        @test u_p_fx1.data.x ≈ p_fx2.data.x atol = 1e-5
        @test u_p_fx1.data.Σy ≈ p_fx2.data.Σy atol = 1e-5
    end
end


text:
function get_parameters_drivers(obj::WindUp)
    return jcall(obj, "getParametersDrivers", List, ())
end

function modify(obj::WindUp, arg0::EstimatedMeasurement)
    return jcall(obj, "modify", void, (EstimatedMeasurement,), arg0)
end


text:
const MD_TOGGLE = true

ENV["MATRIXDEPOT_URL_REDIRECT"] = length(VERSION.prerelease) == 0 ? "1" : "1"
const basedir = tempname()
ENV["MATRIXDEPOT_DATA"] = abspath(basedir, "data")
ENV["MATRIXDEPOT_MYDEPOT"] = abspath(basedir, "myMatrixDepot")

using MatrixDepot
using Test
using LinearAlgebra
using SparseArrays

import MatrixDepot: DataError

function toggle_db()
    MatrixDepot.toggle_remote()
    MatrixDepot.init(ignoredb=true)
    MatrixDepot.update()
end

# that will download the index files if necessary and initialize internal data
MatrixDepot.init()

include("generators.jl")

@testset "MatrixDepot simulate remote matrix tests" begin
    tests = [
            "download",
            "common",
            "number",
            "property",
            ]

    tests = []

    @testset "$t" for t in tests
        tp = joinpath(@__DIR__(), "$(t).jl")
        println("running $(tp) ...")
        include(tp)
        println("finished $(tp)")
    end

    xdir = MatrixDepot.data_dir()
    xtmpdir = string(xdir, ".tmp")
    println("mv $xdir $xtmpdir")
    mv(xdir, xtmpdir, force=true)
    toggle_db()
    @test mdopen(sp(1)) != nothing
    mv(xtmpdir, xdir, force=true)
end

@testset "MatrixDepot real remote matrix tests" begin
    toggle_db()
    println("running remote.jl ...")
    include("remote.jl")
    println("finished remote.jl")
end


text:
module TestFluid  #################################################################

module MyModiaMedia #################################################################

export AbstractMedium
export density
export density_der_1, density_der_2

abstract type AbstractMedium end

struct Water <: AbstractMedium
    baseProperties::Symbol
    singleState::Bool
    data::Int64
    Water(data) = new(:BaseProperties_Water,false,data)
end

struct Air <: AbstractMedium
    baseProperties::Symbol
    singleState::Bool
    data::Float64
    extra::Float64
    Air(data,extra) = new(:BaseProperties_Air,false,data,extra)
end

struct SimpleMedium <: AbstractMedium
    baseProperties::Symbol
    singleState::Bool
    d0::Float64
    SimpleMedium(d0) = new(:BaseProperties_SimpleMedium,true,d0)
end

density(Medium::AbstractMedium, p) = error("\nError from density: No Medium defined\n")
density(Medium::Water         , p) = Medium.data*p
density(Medium::Air           , p) = Medium.data*p + Medium.extra
density(Medium::SimpleMedium  , p) = Medium.d0
density(Medium::SimpleMedium)      = Medium.d0

density_der_1(Medium::AbstractMedium, p) = 0
density_der_1(Medium::AbstractMedium)    = 0

density_der_2(Medium::Water, p) = Medium.data
density_der_2(Medium::Air  , p) = Medium.data

end
using  .MyModiaMedia

println("TestFluid: Tests features needed for Fluid models")

using  Modia
import ModiaMath

const MediumVariable() = Var(size=())

module BaseProperties
    using  Modia
    using  ..MyModiaMedia
    using  ...TestFluid: MediumVariable

    export BaseProperties_Air, BaseProperties_SimpleMedium

    @model BaseProperties_SimpleMedium begin
        Medium = MediumVariable()  # input
        p      = Float(start=1.0)  # input
        d      = Float(start=1.0)  # output
        der_d  = Float(start=1.0)  # output
    @equations begin
        d      = density(Medium)
        der_d  = 0.0
        end
    end

    @model BaseProperties_Air begin
        Medium = MediumVariable()               # input
        p      = Float(start=1.0)               # input
        d      = Float(start=1.0, state=false)  # output
        der_d  = Float(start=1.0)               # output
        der_p  = Float(start=1.0)               # output
    @equations begin
        d     = density(Medium, p)
        der_d = der(d)
        der_p = der(p)
        end
    end
end

using .BaseProperties

function makeDictFromExportedModiaModels!(modul; heading="Keys of model dictionary:")
    dict = Dict{Symbol,Any}()

    println("\n... ", heading)
    for name in names(modul)
        f = getfield(modul, Symbol(name))
        if typeof(f) == Modia.Instantiation.Model
            dict[name] = f
            println("    ", string(name))
        end
    end
    println("")
    return dict
end

const basePropertiesDict = makeDictFromExportedModiaModels!(BaseProperties, heading="Keys of `const baseProperties`:")
getBaseProperties(Medium::AbstractMedium) = basePropertiesDict[Medium.baseProperties]

const water        = MyModiaMedia.Water(100)
const air          = MyModiaMedia.Air(20, 25)
const simpleMedium = MyModiaMedia.SimpleMedium(15.0)

@model Model1 begin
    Medium1 = Var()
    Medium2 = Var()
    d1 = Float()
    d2 = Float()
    @equations begin
        d1 = density(Medium1, 1)
        d2 = density(Medium2, 1)
    end
end

@model Test1 begin
    m = Model1(Medium1=water, Medium2=air)
end

@model Model2 begin
    Medium1 = Var(size=())
    Medium2 = Var(size=())
    d1 = Float(size=())
    d2 = Float(size=())
    @equations begin
        d1 = density(Medium1, 1)
        d2 = density(Medium2, 1)
        Medium1 = Medium2
    end
end

@model Test2 begin
    m = Model2(Medium1=water)
end

@model FluidPort begin
    p      = Float(start=1.0)
    m_flow = Float(start=1.0, flow=true)
    Medium = MediumVariable()
end

@model Source begin
  Medium = AbstractMedium    # Medium MUST be redefined when instanting Source, otherwise error
  P=1
  a=FluidPort()
  d=Float(start=1.0)
  x=Float(start=999.75)
@equations begin
  a.Medium = Medium
  a.p = P
  d = density(a.Medium, a.p)
  der(x) = 1000*(d-x)
  end
end

@model Sink begin
  P=1
  a=FluidPort()
  d=Float()
  x=Float(start=999.75)
@equations begin
  a.p = P
  d = density(a.Medium, a.p)
  der(x) = 1000*(d-x)
  end
end

@model Pipe begin
    R=0.001
    a=FluidPort()
    b=FluidPort()
@equations begin
    b.Medium = a.Medium
    R*a.m_flow = a.p - b.p
    0 = a.m_flow + b.m_flow
    end
end

@model PipeSystem begin
  source=Source(P=100.0, Medium=water)
  sink  =Sink(P=10)
  pipe  =Pipe()
@equations begin
  connect(pipe.a, source.a)
  connect(pipe.b, sink.a)
  end
end

@model ParallelPipeSystem begin
  source=Source(P=100, Medium=water)
  sink=Sink(P=1)
  pipe1=Pipe()
  pipe2=Pipe(R=0.002)
@equations begin
  connect(pipe1.a, source.a)
  connect(pipe1.b, sink.a)

  connect(pipe2.a, source.a)
  connect(pipe2.b, sink.a)
  end
end

# Test of BaseProperties
@model FluidPort2 begin
    p      = Float(start=1.0)
    Medium = MediumVariable()
end

@model Volume begin
    p_start=1.0
    a = FluidPort2()
    V = Float(start=1.0)
    medium = getBaseProperties(a.Medium)(p=Float(start=p_start) )
@equations begin
    medium.Medium = a.Medium
    medium.p      = a.p
    der(V)        = 2.0
    medium.der_d*V + medium.d*der(V) = 2.0*a.p  # der(m) = der(d*V) = der(d)*V + d*der(V)
    end
end

@model Test_Volume begin
    v = Volume(a=FluidPort2(Medium=air), p_start=10.0)
    t = Float(start=0.0)
@equations begin
    der(t) = 1.0
    end
end

# ----------------------------------------------------

res = simulate(Test1, 1, logTranslation=true)
ModiaMath.plot(res, ("m.d1", "m.d2"), heading="Test1", figure=1)

res = simulate(Test2, 1, logTranslation=true)
ModiaMath.plot(res, ("m.d1", "m.d2"), heading="Test2", figure=2)

res = simulate(PipeSystem, 1, logTranslation=true)
ModiaMath.plot(res, ("sink.d", "source.d"), heading="PipeSystem", figure=3)

res = simulate(ParallelPipeSystem, 1, logTranslation=true)
ModiaMath.plot(res, ("sink.d", "source.d"), heading="ParallelPipeSystem", figure=4)

res = simulate(Test_Volume, 1, logTranslation=true, removeSingularities=false)
ModiaMath.plot(res, ("v.a.p", "v.medium.d", "v.medium.der_d"), heading="Test_Volume", figure=5)

end


text:
module CommunalHelperManualCassetteController

using ..Ahorn, Maple
using Ahorn.CommunalHelper

@mapdef Entity "CommunalHelper/ManualCassetteController" ManualCassetteController(
    x::Integer,
    y::Integer,
    startIndex::Int=0,
)

const placements = Ahorn.PlacementDict(
    "Manual Cassette Controller (Communal Helper)" => Ahorn.EntityPlacement(
        ManualCassetteController,
    ),
)

Ahorn.editingOptions(entity::ManualCassetteController) = Dict{String,Any}(
    "startIndex" => cassetteColorNames,
)

const alt = rand(1:100) == 42
const sprite = "objects/CommunalHelper/manualCassetteController/icon$(alt ? "_wacked" : "")"

function Ahorn.selection(entity::ManualCassetteController)
    x, y = Ahorn.position(entity)
    return Ahorn.getSpriteRectangle(sprite, x, y)
end

function Ahorn.render(ctx::Ahorn.Cairo.CairoContext, entity::ManualCassetteController)
    Ahorn.drawSprite(ctx, sprite, 0, 0)
end

end


text:
import NCDatasets

filename = tempname()
# The mode "c" stands for creating a new file (clobber)
ds = NCDatasets.NCDataset(filename,"c")

# define the dimension "lon" and "lat" with the size 10 and 11 resp.
ds.dim["lon"] = 26

v = NCDatasets.defVar(ds,"var_string",String,("lon",))

data = ["$(Char(Int('A')+i))" for i = 0:25]

v[:] = data
NCDatasets.close(ds)

ds = NCDatasets.NCDataset(filename)
vs = ds["var_string"].var

@test vs[:] == data

@test vs[2:3] == data[2:3]

@test vs[2] == data[2]

# ignore additional indices
@test vs[2,1,1] == data[2,1,1]

# ignore additional indices
@test vs[2:3,1,1] == data[2:3,1,1]

NCDatasets.close(ds)


text:
using Pkg
Pkg.activate(".")
using SignalBroadcastingSim
using Random
using OpenStreetMapX
using DelimitedFiles, CSV
using LightGraphs
using DataFrames
using Statistics
using Distributions, HypothesisTests
using CSV
m = OpenStreetMapX.get_map_data("maps", "torontoF.osm")
for discretize_m in [25., 50., 75.]
    p = ModelParams(discretize_m=discretize_m)
    s = Simulation(p,m)

    df = DataFrame(ix=Int[], latitude=Float64[],longitude=Float64[], OSM_nodeid=Union{Int,Nothing}[],random_node=Bool[],neighbors=String[])

    for v in s.vertex_meta
        lla = LLA(v.enu, s.m.bounds)
        push!(df, Dict( :ix => v.ix,:latitude=>lla.lat, :longitude => lla.lon, :OSM_nodeid => v.osmid, :random_node => v.randomNode, :neighbors => string(neighbors(s.g, v.ix))  ))
    end

    CSV.write("export_toronto_graph_d$(discretize_m)m.csv", df)

    dfe = DataFrame(vertex1=Int[], vertex2=Int[], distance_in_meters=Float64[])
    for e in edges(s.g)
        push!(dfe, (e.src, e.dst, s.w[e.src, e.dst]))
    end
    CSV.write("export_toronto_distances_d$(discretize_m)m.csv", dfe)
end

p = ModelParams(n_agents=0, discretize_m=50)
s = Simulation(p,m)
SignalBroadcastingSim.plot_sim(s, "road_network.html")

Random.seed!(20)
p = ModelParams(n_agents=500, discretize_m=50)
s = Simulation(p,m)
init!(s)
for i in 1:100
    SignalBroadcastingSim.step!(s)
end
SignalBroadcastingSim.plot_sim(s, "sample_200agents_discretize50m_after100steps.html")


text:
##### Energy tendencies

#####
##### First order fluxes
#####

function flux(::Energy, ::Advect, atmos, args)
    @unpack state = args
    return (state.ρu / state.ρ) * state.energy.ρe
end

function flux(::ρθ_liq_ice, ::Advect, atmos, args)
    @unpack state = args
    return (state.ρu / state.ρ) * state.energy.ρθ_liq_ice
end

function flux(::Energy, ::Pressure, atmos, args)
    @unpack state = args
    @unpack ts = args.precomputed
    return state.ρu / state.ρ * air_pressure(ts)
end

#####
##### Second order fluxes
#####

struct ViscousFlux <: TendencyDef{Flux{SecondOrder}} end
function flux(::Energy, ::ViscousFlux, atmos, args)
    @unpack state = args
    @unpack τ = args.precomputed.turbulence
    return τ * state.ρu
end

function flux(::ρθ_liq_ice, ::ViscousFlux, atmos, args)
    @unpack state, diffusive = args
    @unpack D_t = args.precomputed.turbulence
    return state.ρ * (-D_t) .* diffusive.energy.∇θ_liq_ice
end

function flux(::Energy, ::HyperdiffViscousFlux, atmos, args)
    @unpack state, hyperdiffusive = args
    return hyperdiffusive.hyperdiffusion.ν∇³u_h * state.ρu
end

function flux(::Energy, ::HyperdiffEnthalpyFlux, atmos, args)
    @unpack state, hyperdiffusive = args
    return hyperdiffusive.hyperdiffusion.ν∇³h_tot * state.ρ
end

struct DiffEnthalpyFlux <: TendencyDef{Flux{SecondOrder}} end
function flux(::Energy, ::DiffEnthalpyFlux, atmos, args)
    @unpack state, diffusive = args
    @unpack D_t = args.precomputed.turbulence
    d_h_tot = -D_t .* diffusive.energy.∇h_tot
    return d_h_tot * state.ρ
end

#####
##### Sources
#####

function source(::Energy, s::Subsidence, m, args)
    @unpack state, aux, diffusive = args
    z = altitude(m, aux)
    w_sub = subsidence_velocity(s, z)
    k̂ = vertical_unit_vector(m, aux)
    return -state.ρ * w_sub * dot(k̂, diffusive.energy.∇h_tot)
end

function source(::Energy, s::RemovePrecipitation, m, args)
    @unpack state = args
    @unpack ts = args.precomputed
    if has_condensate(ts)
        nt = remove_precipitation_sources(s, m, args)
        return nt.S_ρ_e
    else
        FT = eltype(state)
        return FT(0)
    end
end

function source(::Energy, s::WarmRain_1M, m, args)
    @unpack cache = args.precomputed.precipitation
    return cache.S_ρ_e
end

function source(::Energy, s::RainSnow_1M, m, args)
    @unpack cache = args.precomputed.precipitation
    return cache.S_ρ_e
end


text:
abstract type Integrator{dType, tType} end

abstract type DeterministicIntegrator{dType, tType} <: Integrator{dType, tType} end
abstract type StochasticIntegrator{dType, tType} <: Integrator{dType, tType} end

abstract type ODEIntegrator{dType, tType} <: DeterministicIntegrator{dType, tType} end
abstract type DAEIntegrator{dType, tType} <: DeterministicIntegrator{dType, tType} end
abstract type IODEIntegrator{dType, tType} <: DeterministicIntegrator{dType, tType} end
abstract type IDAEIntegrator{dType, tType} <: DeterministicIntegrator{dType, tType} end
abstract type PODEIntegrator{dType, tType} <: DeterministicIntegrator{dType, tType} end
abstract type PDAEIntegrator{dType, tType} <: DeterministicIntegrator{dType, tType} end

abstract type HODEIntegrator{dType, tType} <: PODEIntegrator{dType, tType} end
abstract type HDAEIntegrator{dType, tType} <: PDAEIntegrator{dType, tType} end
abstract type VODEIntegrator{dType, tType} <: IODEIntegrator{dType, tType} end
abstract type VDAEIntegrator{dType, tType} <: IDAEIntegrator{dType, tType} end

abstract type SDEIntegrator{dType, tType} <: StochasticIntegrator{dType, tType} end
abstract type PSDEIntegrator{dType, tType} <: StochasticIntegrator{dType, tType} end
abstract type SPSDEIntegrator{dType, tType} <: StochasticIntegrator{dType, tType} end

equation(integrator::Integrator) = error("equation() not implemented for ", typeof(integrator))
timestep(integrator::Integrator) = error("timestep() not implemented for ", typeof(integrator))

Base.ndims(integrator::Integrator) = ndims(equation(integrator))

abstract type Parameters{DT,TT} end

function_stages!(x::Vector{DT}, b::Vector{DT}, params::PT) where {DT, TT, PT <: Parameters{DT,TT}} = error("function_stages!() not implemented for ", PT)


text:
# License for this file: MIT (expat)
# Copyright 2017-2018, DLR Institute of System Dynamics and Control

"""
    module Modia3D.ForceElements

Elements generating forces or torques that act on objects.

# Main developers
Andrea Neumayr and Martin Otter, [DLR - Institute of System Dynamics and Control](https://www.dlr.de/sr/en)
"""
module ForceElements

export @forceElement

using  StaticArrays
import DataStructures
import ModiaMath

import Base
import Modia3D
import Modia3D.Basics

include("forces.jl")
end


text:
using Parameters, AutomotiveDrivingModels, POMDPs

include("occluded_crosswalk_env.jl")
include("helpers.jl")
include("pomdp_types.jl")
include("constant_pedestrian.jl")
include("generative_model.jl")

rng = MersenneTwister(1)
pomdp = OCPOMDP()

s0 = initialstate(pomdp, rng)
ego = s0[findfirst(EGO_ID,s0)]
println("Initial scene with $(length(s0)) cars")
println("Initial ego state: $ego")

max_stars_repo_name: MaximeBouton/AutomotivePOMDPs

max_issues_repo_name: MaximeBouton/AutomotivePOMDPs

max_forks_repo_name: MaximeBouton/AutomotivePOMDPs


text:
include("propagator_basic.jl")
include("textmap_propagator.jl")


text:
export MCMCSampler

abstract type MCMCRule end
mutable struct MCMCSampler{T<:MCMCRule} <: MonteCarloSampler
    chain_length::Int
    burn_length::Int
    seed::UInt
    rule::T
end

MCMCSampler(rule, n_samples; burn=0, seed=rand(UInt)) =
    MCMCSampler{typeof(rule)}(n_samples, burn, seed, rule)

## Cache
mutable struct MCMCSamplerCache{A<:AbstractRNG,RC} <: SamplerCache{MCMCSampler}
    rng::A
    steps_done::Int
    steps_accepted::Int
    # Might add here a field rule_cache::T, generated by some
    # create_cache_for_rule(::MCMCRule) = nothing by default
    rule_cache::RC
end

_sampler_cache(s::MCMCSampler, v, net, part) =
    MCMCSamplerCache(MersenneTwister(s.seed),
                      0, 0,
                      RuleSamplerCache(s.rule, s, v, net, part))

# Construct the cache of a specific rule.
# By default rules have no cache.
RuleSamplerCache(r::MCMCRule, s, v, net, part) = nothing
rule(s::MCMCSampler) = s.rule
rule_cache(sc::MCMCSamplerCache) = sc.rule_cache

function init_sampler!(s::MCMCSampler, net, σ, c::MCMCSamplerCache)
    c.steps_done = 0
    c.steps_accepted = 0
    set_index!(σ, rand(c.rng, 1:spacedimension(σ)))
    init_lut!(σ, net)

    while c.steps_done < s.burn_length
        markov_chain_step!(σ, s, net, c)
    end
    c.steps_done = 0

    return c
end

init_sampler_rule_cache!(rc, s::MCMCSampler, net, σ, c::MCMCSamplerCache) =
    nothing

done(s::MCMCSampler, σ, c) = c.steps_done >= s.chain_length

function samplenext!(σ, s::MCMCSampler, net, c)
    # Check termination condition, and return if verified
    done(s, σ, c) && return false
    return markov_chain_step!(σ, s, net, c)
end

"""
    markov_chain_step!(state, sampler, net, sampler_cache)

performs one step of markov chain
"""
function markov_chain_step!(σ, s::MCMCSampler, net, c)
    throw("markov_chain_step! undefined for this type")
end

## Mulithreading
function _mt_recompute_sampler_params!(samplers, s::MCMCSampler)
    nt = Threads.nthreads()
    _chain_length = Int(ceil(s.chain_length / nt))
    rng = MersenneTwister(s.seed)

    for i=1:Threads.nthreads()
        samplers[i] = MCMCSampler(s.rule,
                                   _chain_length,
                                   burn=s.burn_length,
                                   seed=rand(rng, UInt))
    end
end


text:
# This file is a part of TypeDBClient.  License is MIT: https://github.com/Humans-of-Julia/TypeDBClient.jl/blob/main/LICENSE

#
# package typedb.client.api.concept.thing;
#
# import typedb.client.api.TypeDBTransaction;
# import typedb.client.api.concept.Concept;
# import typedb.client.api.concept.type.AttributeType;
# import typedb.client.api.concept.type.RoleType;
# import typedb.client.api.concept.type.ThingType;
#
# import javax.annotation.CheckReturnValue;
# import java.util.stream.Stream;
#
# public interface Thing extends Concept {
#
#     @CheckReturnValue
#     String getIID();
#
#     @CheckReturnValue
#     ThingType getType();
#
#     @Override
#     @CheckReturnValue
#     default boolean isThing() {
#         return true;
#     }
#
#     @Override
#     @CheckReturnValue
#     Thing.Remote asRemote(TypeDBTransaction transaction);
#
#     interface Remote extends Concept.Remote, Thing {
#
#         void setHas(Attribute<?> attribute);
#
#         void unsetHas(Attribute<?> attribute);
#
#         @CheckReturnValue
#         boolean isInferred();
#
#         @CheckReturnValue
#         Stream<? extends Attribute<?>> getHas(boolean onlyKey);
#
#         @CheckReturnValue
#         Stream<? extends Attribute.Boolean> getHas(AttributeType.Boolean attributeType);
#
#         @CheckReturnValue
#         Stream<? extends Attribute.Long> getHas(AttributeType.Long attributeType);
#
#         @CheckReturnValue
#         Stream<? extends Attribute.Double> getHas(AttributeType.Double attributeType);
#
#         @CheckReturnValue
#         Stream<? extends Attribute.String> getHas(AttributeType.String attributeType);
#
#         @CheckReturnValue
#         Stream<? extends Attribute.DateTime> getHas(AttributeType.DateTime attributeType);
#
#         @CheckReturnValue
#         Stream<? extends Attribute<?>> getHas(AttributeType... attributeTypes);
#
#         @CheckReturnValue
#         Stream<? extends Relation> getRelations(RoleType... roleTypes);
#
#         @CheckReturnValue
#         Stream<? extends RoleType> getPlaying();
#     }
# }


text:
function _precompile_()
    ccall(:jl_generating_output, Cint, ()) == 1 || return nothing
    Base.precompile(Tuple{typeof(materialize),Broadcasted{DefaultArrayStyle{1}, Nothing, typeof(/), Tuple{Broadcasted{DefaultArrayStyle{1}, Nothing, typeof(/), Tuple{Vector{Float64}, Vector{Int64}}}, Float64}}})   # time: 0.0245907
    Base.precompile(Tuple{typeof(materialize),Broadcasted{DefaultArrayStyle{1}, Nothing, typeof(/), Tuple{Vector{Float64}, Float64}}})   # time: 0.0153087
    Base.precompile(Tuple{typeof(materialize!),Vector{Float64},Broadcasted{DefaultArrayStyle{1}, Nothing, typeof(/), Tuple{Vector{Float64}, Float64}}})   # time: 0.0049648
    Base.precompile(Tuple{typeof(broadcasted),Function,Vector{Float64},Float64})   # time: 0.0036101
    Base.precompile(Tuple{typeof(broadcasted),Function,Broadcasted{DefaultArrayStyle{1}, Nothing, typeof(/), Tuple{Vector{Float64}, Vector{Int64}}},Float64})   # time: 0.0034639
end


text:
"""
```
forecast(m, system, z0; enforce_zlb = false, shocks = Matrix{S}(0,0))

forecast(system, z0, shocks; enforce_zlb = false)
```

### Inputs

- `m::AbstractModel`: model object. Only needed for the method in which `shocks`
  are not provided.
- `system::System{S}`: state-space system matrices
- `kal::Kalman{S}` or `z0::Vector{S}`: result of running the Kalman filter or
  state vector in the final historical period (aka initial forecast period)

where `S<:AbstractFloat`.

### Keyword Arguments

- `cond_type::Symbol`: one of `:none`, `:semi`, or `:full`, used to determine
  how many periods to forecast ahead. If `cond_type in [:semi, :full]`, the
  forecast horizon is reduced by the number of periods of conditional
  data. Defaults to `:none`.
- `enforce_zlb::Bool`: whether to enforce the zero lower bound. Defaults to
  `false`.
- `shocks::Matrix{S}`: matrix of size `nshocks` x `shock_horizon` of shock
  innovations under which to forecast. If `shock_horizon > horizon`, the extra
  periods of shocks will be ignored; if `shock_horizon < horizon`, zeros will be
  filled in for the shocks hitting the remaining forecasted periods.
- `draw_shocks::Bool`: if `isempty(shocks)`, indicates whether to draw shocks
  according to:

  1. If `forecast_tdist_shocks(m)`, draw `horizons` many shocks from a
     `Distributions.TDist(forecast_tdist_df_val(m))`
  2. Otherwise, draw `horizons` many shocks from a
     `DegenerateMvNormal(zeros(nshocks), sqrt(system[:QQ]))`

  or to set `shocks` to a `nshocks` x `horizon` matrix of zeros. Defaults to
  `false`. If `shocks` is provided as a keyword argument, this flag has no
  effect.

### Outputs

- `states::Matrix{S}`: matrix of size `nstates` x `horizon` of forecasted states
- `obs::Matrix{S}`: matrix of size `nobs` x `horizon` of forecasted observables
- `pseudo::Matrix{S}`: matrix of size `npseudo` x `horizon` of forecasted
  pseudo-observables
- `shocks::Matrix{S}`: matrix of size `nshocks` x `horizon` of shock innovations
"""
function forecast{S<:AbstractFloat}(m::AbstractModel, system::System{S},
    z0::Vector{S}; cond_type::Symbol = :none, enforce_zlb::Bool = false,
    shocks::Matrix{S} = Matrix{S}(0, 0), draw_shocks::Bool = false)

    # Numbers of things
    nshocks = n_shocks_exogenous(m)
    horizon = forecast_horizons(m; cond_type = cond_type)

    if isempty(shocks)
        # Populate shocks matrix
        if draw_shocks
            μ = zeros(S, nshocks)
            σ = sqrt.(system[:QQ])
            dist = if forecast_tdist_shocks(m)
                # Use t-distributed shocks
                ν = forecast_tdist_df_val(m)
                DegenerateDiagMvTDist(μ, σ, ν)
            else
                # Use normally distributed shocks
                DegenerateMvNormal(μ, σ)
            end

            shocks = rand(dist, horizon)

            # Forecast without anticipated shocks
            if n_anticipated_shocks(m) > 0
                ind_ant1 = m.exogenous_shocks[:rm_shl1]
                ind_antn = m.exogenous_shocks[Symbol("rm_shl$(n_anticipated_shocks(m))")]
                ant_shock_inds = ind_ant1:ind_antn
                shocks[ant_shock_inds, :] = 0
            end
        else
            shocks = zeros(S, nshocks, horizon)
        end
    else
        # Adjust size of shocks matrix, padding with zeros or cutting off
        # periods of shocks if necessary
        shock_horizon = size(shocks, 2)
        if shock_horizon <= horizon
            shocks0 = zeros(nshocks, horizon - shock_horizon)
            shocks = hcat(shocks, shocks0)
        else
            shocks = shocks[:, 1:horizon]
        end
    end

    # Populate shocks matrix under alternative policy, if
    # user has specified a function to do so
    alt_policy = alternative_policy(m)
    if alt_policy.solve != identity &&
        alt_policy.forecast_init != identity

        shocks, z0 = alt_policy.forecast_init(m, shocks, z0, cond_type = cond_type)
    end

    # Get variables necessary to enforce the zero lower bound in the forecast
    ind_r = m.observables[:obs_nominalrate]
    ind_r_sh = m.exogenous_shocks[:rm_sh]
    zlb_value = forecast_zlb_value(m)

    forecast(system, z0, shocks; enforce_zlb = enforce_zlb,
        ind_r = ind_r, ind_r_sh = ind_r_sh, zlb_value = zlb_value)
end

function forecast{S<:AbstractFloat}(system::System{S}, z0::Vector{S},
    shocks::Matrix{S}; enforce_zlb::Bool = false, ind_r::Int = -1,
    ind_r_sh::Int = -1, zlb_value::S = 0.13/4)

    # Unpack system
    T, R, C = system[:TTT], system[:RRR], system[:CCC]
    Q, Z, D = system[:QQ], system[:ZZ], system[:DD]
    Z_pseudo, D_pseudo = system[:ZZ_pseudo], system[:DD_pseudo]

    # Setup
    nshocks = size(R, 2)
    nstates = size(T, 2)
    nobs    = size(Z, 1)
    npseudo = size(Z_pseudo, 1)
    horizon = size(shocks, 2)

    # Define our iteration function
    function iterate(z_t1, ϵ_t)
        z_t = C + T*z_t1 + R*ϵ_t

        # Change monetary policy shock to account for 0.13 interest rate bound
        if enforce_zlb
            interest_rate_forecast = getindex(D + Z*z_t, ind_r)
            if interest_rate_forecast < zlb_value
                # Solve for interest rate shock causing interest rate forecast to be exactly ZLB
                ϵ_t[ind_r_sh] = 0.
                z_t = C + T*z_t1 + R*ϵ_t
                ϵ_t[ind_r_sh] = getindex((zlb_value - D[ind_r] - Z[ind_r, :]'*z_t) / (Z[ind_r, :]' * R[:, ind_r_sh]), 1)

                # Forecast again with new shocks
                z_t = C + T*z_t1 + R*ϵ_t

                # Confirm procedure worked
                interest_rate_forecast = getindex(D + Z*z_t, ind_r)
                @assert interest_rate_forecast >= zlb_value - 0.01 "interest_rate_forecast = $interest_rate_forecast must be >= zlb_value - 0.01 = $(zlb_value - 0.01)"
            end
        end
        return z_t, ϵ_t
    end

    # Iterate state space forward
    states = zeros(S, nstates, horizon)
    states[:, 1], shocks[:, 1] = iterate(z0, shocks[:, 1])
    for t in 2:horizon
        states[:, t], shocks[:, t] = iterate(states[:, t-1], shocks[:, t])
    end

    # Apply measurement and pseudo-measurement equations
    obs    = D .+ Z*states
    pseudo = D_pseudo .+ Z_pseudo * states

    # Return forecasts
    return states, obs, pseudo, shocks
end


text:
function unmarkededges(g, x::BitArray)
    selected_edges = Vector{edgetype(g)}()

    for item in zip(1:ne(g), edges(g))
        if x[item[1]] == false
            # rem_edge!(subgraph, src(item[2]), dst(item[2]))
            push!(selected_edges, item[2])
        end
    end
    return selected_edges
end

function markededges(g, x::BitArray)
    selected_edges = Vector{edgetype(g)}()
    for item in zip(1:ne(g), edges(g))
        if x[item[1]]
            # rem_edge!(subgraph, src(item[2]), dst(item[2]))
            push!(selected_edges, item[2])
        end
    end
    return selected_edges
end

function graph_bitstring(g, x::BitArray)
    selected_edges = markededges(g, x)
    g2 = SimpleWeightedGraph(nv(g))
    for item in selected_edges
        # add_edge!(g2, src(item), dst(item), weight(item))
        add_edge!(g2, item)
    end
    return g2
end

mst_value(g) = sum([g.weights[src(edge), dst(edge)] for edge in prim_mst(g)])

max_stars_repo_name: DTUComputeTONIA/SDRLS


text:
include("Token.jl")

const RESERVEDWORDS  = Set(["programa", "declare", "escreva", "leia", "fimprog"])
const CONTROL_FLUX   = Set(["se", "entao", "senao", "fimse", "enquanto", "faca", "fimenq", "durante"])
const BOOL_VALUE     = Set(["Verdadeiro", "Falso", "V", "F"])
const TYPE_NAMES     = Set(["int", "real", "char", "texto", "boleano"])
const BLANKS         = Set([' ', '\n', '\t', '\r'])
const OPERATORS      = Set(["+", "-", "*", "/", "^", "==", "!=", "<", ">", "<=", ">=", "||", "!", "&&"])
const SEPARATORS     = Set(['(', ')', ';', ',', ':'])

isboolliteral(s::String)  = s in BOOL_VALUE
isblank(c::Char)          = c in BLANKS
isseparator(c::Char)      = c in SEPARATORS
isoperator(s::String)     = s in OPERATORS
isreservedword(s::String) = s in RESERVEDWORDS
iscontrolflux(s::String)  = s in CONTROL_FLUX
istype(s::String)         = s in TYPE_NAMES

const IDENTIFIER_REGEX   = r"[A-Za-z_]+[A-Za-z-]*[0-9]*"

mutable struct Source
    orig::String
    src::String
    curr_pos::Integer

    function Source(src)
        new(src, src, 1)
    end # function
end # struct

function next_token(s::Source, line::Integer) :: Token
    chars = collect(s.src)

    for (i, char) in enumerate(chars)
        if s.curr_pos > length(s.orig) break end # if

        pos = s.curr_pos

        res = if isblank(char)
            Token(WHITESPACE, string(char), (pos, pos), line)
        elseif char == '-' || isdigit(char)
            parse_number(s.src[i:end], pos, line)
        elseif char == '\''
            parse_char(s.src[i:end], pos, line)
        elseif char == '\"'
            parse_str(s.src[i:end], pos, line)
        elseif char == '#'
            parse_comment(s.src[i:end], pos, line)
        elseif char == ':'
            if i < length(chars)
                if chars[i+1] == '='
                    Token(ASSIGN, ":=", (pos,pos+1), line)
                else
                    Token(PUNCTUATION, string(char), (pos, pos), line)
                end # if
            end # if
        elseif isseparator(char)
            Token(PUNCTUATION, string(char), (pos, pos), line)
        else
            parse_rest(s.src[i:end], pos, line)
        end # if

        final = res.span[2]

        s.curr_pos = final + 1
        s.src = s.orig[s.curr_pos:end]

        return res
    end # for
    Token(EOF, "", (s.curr_pos, s.curr_pos), line)
end # function

function parse_number(src::String, pos::Integer, line::Integer)::Token
    chars = collect(src)
    vec_chars::Array{Char} = []

    has_dot = false

    for (i, char) in enumerate(chars)
        if isdigit(char)
            push!(vec_chars, char)
        elseif char == '-' && i == 1
            if i < length(chars)
                if isdigit(chars[i+1])
                    push!(vec_chars, char)
                else
                    return parse_rest(src, pos, line)
                end # if
            else
                return parse_rest(src, pos, line)
            end # if
        elseif char == '.'
            has_dot = true
            if i < length(chars)
                if isdigit(chars[i+1])
                    push!(vec_chars, char)
                end # if
            end # if
        else
            break
        end # if
    end # for

    len = length(vec_chars) - 1

    if has_dot
        return Token(FLOAT_NUMBER, String(vec_chars), (pos, pos+len), line)
    else
        return Token(INT_NUMBER, String(vec_chars), (pos, pos+len), line)
    end # if
end # function

function parse_char(src::String, pos::Integer, line::Integer)::Token
    chars = collect(src)

    if chars[3] != '\''
        return Token(INVALID, src[1:3], (pos, pos+2), line)
    end # if

    return Token(CHAR, src[1:3], (pos, pos+2), line)
end # function

function parse_str(src, pos, line)::Token
    chars = collect(src)
    finish = 0
    for (index, char) in enumerate(chars)
        if char == '\"' && index != 1
            finish = index
            break
        end # if
    end # for

    if last(src[1:finish]) != '\"'
        return Token(INVALID, src[1:finish], (pos, pos+finish-1), line)
    end # if

    # return Token(STRING, String(vec_chars), (pos, pos+len))
    return Token(STRING, src[1:finish], (pos, pos+finish-1), line)
end # function

function parse_rest(src, pos, line)::Token
    chars = collect(src)
    vec_chars :: Array{Char} = []

    for char in chars
        if isblank(char) || isseparator(char)
            break
        end # if
        push!(vec_chars, char)
    end # for

    str = String(vec_chars)
    len = length(str) -1

    if isreservedword(str)
        return Token(RESERVED_WORD, str, (pos, pos+len), line)
    elseif iscontrolflux(str)
        return Token(CFLUX, str, (pos, pos+len), line)
    elseif isboolliteral(str)
        return Token(BOOL, str, (pos, pos+len), line)
    elseif istype(str)
        return Token(TYPE, str, (pos, pos+len), line)
    elseif isoperator(str)
        return Token(OPERATOR, str, (pos, pos+len), line)
    elseif occursin(IDENTIFIER_REGEX, str)
        m = match(IDENTIFIER_REGEX, str)
        if len+1 == length(m.match)
            return Token(IDENTIFIER, m.match, (pos, pos+len), line)
        else
            return Token(INVALID, m.match, (pos, pos+length(m.match)-1), line)
        end # if
    else
        return Token(INVALID, str, (pos, pos+len), line)
    end # if
end # function

function parse_comment(str::String, pos, line)
    chars = collect(str)
    finish = 0
    for (i, ch) in enumerate(chars)
        if ch == '\n'
            finish = i
            break
        end # if
    end # for
    return Token(COMMENT, str[1:finish], (pos, pos+finish), line)
end # function

function tokenise(src::String)::Tokens
    i = 0
    source = Source(src)
    vec_tokens :: Array{Token} = []
    line = 1
    coll = 1
    while true
        token = next_token(source, line)
        token.col = coll

        coll += length(token.text)

        if token.id == WHITESPACE && token.text == "\n"
            line += 1
            coll = 1
        end # if
        if token.id == EOF
            push!(vec_tokens, token)
            break
        end # if
        push!(vec_tokens, token)
        i += 1
    end # while

    filter!(x -> x.id != WHITESPACE, vec_tokens)

    Tokens(vec_tokens)
end # function

max_stars_repo_name: GabrielTrettel/Portran

max_issues_repo_name: GabrielTrettel/Portran

max_forks_repo_name: GabrielTrettel/Portran