### A Pluto.jl notebook ###
# v0.17.2

using Markdown
using InteractiveUtils

# ╔═╡ 3cd07b90-9e21-4a74-8d99-a6375beb7d1e
begin
	using Revise
	using Pkg
	Pkg.activate("../")

	using BayesianModels
	using DelimitedFiles
	using ForwardDiff
	using LinearAlgebra
	using Plots
	using Zygote

	include("./plotting.jl")
end

# ╔═╡ 6958bde2-c50f-4afd-8d08-7dfbfd97169a
md"""
# Learning and Modeling the Languages Phonetics

*[Lucas Ondel](https://lucasondel.github.io/), 8 March 2022*
"""

# ╔═╡ b83311ab-ca4f-432e-8627-5207517a81db
md"""## Lingua Libre Formants data

We use the [Lingua Libre Formants](https://github.com/lucasondel/LinguaLibreFormants) data for our analysis.

### Getting the data

We download the data and store it in a local file.
"""

# ╔═╡ d1588a1c-babe-4223-bc6d-f353314ed2c0
localfile = "formants.tsv"

# ╔═╡ 262e0395-814f-4315-b25e-21c18a7a9a5f
begin
	# Database URL
	local dataurl = "https://raw.githubusercontent.com/lucasondel/LinguaLibreFormants/main/data/formants.tsv"

	# Download the data if it is not done already
	if ! isfile(localfile)
		run(`wget $dataurl`)
	end
end

# ╔═╡ 4bbdd3d2-eb62-4889-9c8a-20141c0b140a
md"""
### Data preparation

In a nutshell:
1. we extract the list of languages, speakers and phones
2. organized the data as a $3 \times N$ matrix where 3 is for 3 first formants and $N$ is the number of data points
3. for each data point we extract the annotation (language, speaker, phone)
"""

# ╔═╡ ee1d1c45-70f5-4fe7-a8e2-9134d45a2b95
begin
	local rawdata = readdlm(localfile)

	speakers = Set()
	languages = Set()
	phones = Set()
	phones_per_lang = Dict()
	labels = []
	data = []
	for (lang, spk, sex, ipaphone, f1, f2, f3, duration) in eachrow(rawdata)
		# Remove the duration symbol
		ipaphone = replace(ipaphone, "ː" => "")

		lang, spk, ipaphone = Symbol(lang), Symbol(spk), Symbol(ipaphone)
		push!(data, vcat(f1, f2, f3))

		phone = ipaphone
		spk = (lang, sex, spk)

		l = get(phones_per_lang, lang, Set())
		push!(l, phone)
		phones_per_lang[lang] = l

		push!(labels, (lang, spk, phone))
		push!(speakers, spk)
		push!(languages, lang)
		push!(phones, phone)
	end
	data = hcat(data...)
	labels = vcat(labels...)

	# Mean-variance normalization of the data
	local x̄ = sum(data, dims=2) / size(data, 2)
	local x̄² = sum(data .^ 2, dims=2) / size(data, 2)
	local var = x̄² - (x̄ .^ 2)
	data = (data .- x̄) ./ sqrt.(var)

	# Transform the languages, speakers and phone sets to ordered collections
	languages = collect(languages)
	speakers = collect(speakers)
	phones = collect(phones)

	# Build the element <-> id mappings
	lang2id = Dict(lang => i for (i, lang) in enumerate(languages))
	id2lang = Dict(i => lang for (i, lang) in enumerate(languages))
	spk2id = Dict(spk => i for (i, spk) in enumerate(speakers))
	id2spk = Dict(i => spk for (i, spk) in enumerate(speakers))
	phone2id = Dict(phone => i for (i, phone) in enumerate(phones))
	id2phone = Dict(i => phone for (i, phone) in enumerate(phones))

	# Transform the labels to use the id rather than symbols
	labels = [(lang2id[lang], spk2id[spk], phone2id[phone])
			  for (lang, spk, phone) in labels]

	# Group the data for each triplet (lang, speaker, phone)
	groups = Dict()
	for (label, sample) in zip(labels, eachcol(data))
		list = get(groups, label, [])
		push!(list, sample)
		groups[label] = list
	end
	groups = tuple((k => hcat(groups[k]...) for k in keys(groups))...)
end

# ╔═╡ 474abeb1-82f8-4a74-a03c-be1cf55df840
md"Some important constants:"

# ╔═╡ 31e56fed-e4f0-4c71-abb2-5c2eff5a9a47
D = size(data, 1)

# ╔═╡ 61cc1ab2-a615-4a9b-91f9-bc5909e29297
N = size(data, 2)

# ╔═╡ e08cc384-49dc-4585-a2ab-fbec9fcfe9d2
L = length(languages)

# ╔═╡ b8e7f92f-0065-46fe-ac6b-5adf29b3379f
S = length(speakers)

# ╔═╡ ef71661f-338e-4d57-86ed-fc2c308a744f
P = length(phones)

# ╔═╡ c01de0cd-43f8-49cc-a8c9-ffbbc182d9e9
begin
	md"""
	### Data statistics

	| \# languages | \# speakers  | \# phones |
	|:------------:|:------------:|:---------:|
	| $L           | $S           | $P        |

	"""
end

# ╔═╡ 8b535f17-7a55-42da-ba7a-c161bec9fd28
md"""
## Model

### Embeddings
1.  $\boldsymbol{\lambda} \sim p(\boldsymbol{\lambda}) \triangleq \mathcal{N}( \mathbf{0}, \mathbf{I})$, $\boldsymbol{\lambda}$ → language embedding of dimension $Q_\lambda$
2.  $\boldsymbol{\sigma} \sim p(\boldsymbol{\sigma}) \triangleq \mathcal{N}( \mathbf{0}, \mathbf{I})$, $\boldsymbol{\sigma}$ → speaker embedding of dimension $Q_\sigma$
3.  $\boldsymbol{\pi} \sim p(\boldsymbol{\pi}) \triangleq \mathcal{N}( \mathbf{0}, \mathbf{I})$ $\boldsymbol{\pi}$ → phone embedding of dimension $Q_\pi$
"""

# ╔═╡ 5175712e-7f27-459e-8e6a-9bed325d6b4e
md"""
### Subspace construction
1.  $\mathscr{L}_i \sim p(\mathscr{L}) \triangleq \mathcal{N}(\mathbf{0}, \mathbf{I})$, $\{\mathscr{L}_0, \dots, \mathscr{L}_{Q_\lambda} \}$ → Language hyper-subspace bases of dimension $(Q_\pi + 1) \times K$
2.  $\mathscr{S}_j \sim p(\mathscr{S}) \triangleq \mathcal{N}(\mathbf{0}, \mathbf{I})$, $\{\mathscr{S}_0, \dots, \mathscr{S}_{Q_\sigma} \}$ → Speaker hyper-subspace bases $(Q_\pi + 1) \times K$
3.  $\mathscr{P} = \sum_{i=0}^{Q_\lambda} \lambda_i \mathscr{L}_i + \sum_{j=0}^{Q_\sigma} \sigma_j \mathscr{S}_j$ where $\lambda_0 = \sigma_0 = 1$, → Phonetic subspace bases of dimension $(Q_\pi + 1) \times K$
"""

# ╔═╡ 763b5f95-db02-4ede-b632-848a646232d9
md"""
### Phone model
1.  $\boldsymbol{\psi} = \mathscr{P}^\top \begin{bmatrix} \boldsymbol{\pi} \\ 1 \end{bmatrix}$, $\boldsymbol{\psi} \in \mathbb{R}^K$
2.  $\boldsymbol{\eta} = f(\boldsymbol{\psi})$
3.  $\mathbf{x}_n \sim p(\mathbf{x} | \boldsymbol{\eta}) \triangleq \mathcal{N}(\boldsymbol{\mu}, \boldsymbol{\Lambda}^{-1})$, where $\boldsymbol{\eta}$ are the natural parameters of the multivariate Normal distribution
"""

# ╔═╡ 841a3f0e-8975-49d2-959e-afe832ed41e4
md"""
### Parameter function mapping

The function $f$ transforms the parameters $\boldsymbol{\psi}$ from the real space into the natural parameter space. It is defined as follows:

```math
	\begin{align}
	\boldsymbol{\psi} &= \begin{bmatrix}
		 \boldsymbol{\psi}^{(\boldsymbol{\Lambda\mu})} \\
		 \boldsymbol{\psi}^{(\mathbf{L})}_1 \\
		 \boldsymbol{\psi}^{(\mathbf{L})}_2
	\end{bmatrix} \\
	\boldsymbol{\Lambda\mu} &= \boldsymbol{\psi}^{(\boldsymbol{\Lambda\mu})} \\
	\text{diag}(\mathbf{L}) &= \exp \{ \boldsymbol{\psi}^{(\mathbf{L})}_1 \} \\
	\text{vech}^{\dagger}(\mathbf{L}) &= \boldsymbol{\psi}^{(\mathbf{L})}_2 \\
	\boldsymbol{\Lambda} &= \mathbf{L} \mathbf{L}^\top
	\end{align}
```
where $\text{vech}^{\dagger}$ is the half-vectorization of a matrix **not including the diagonal**
"""

# ╔═╡ b0761304-49df-4626-a677-488295dc786d
function f(ψ::AbstractVector)
		L = length(ψ)
		D = Int((-(3/2) + sqrt((9/4) + 2L)) )

		Λμ = ψ[1:D]
		A = diagm(exp.(1e-3 .+ ψ[D+1:2*D]))
		L = A + CompressedLowerTriangular(D, 1, ψ[2*D+1:end])
		L⁻¹ = inv(L)

		Λ = L * L'
		μᵀΛμ = ((L⁻¹' * L⁻¹) * Λμ)' * Λμ
		logdetΛ = 2*sum(ψ[D+1:2*D])

		Λμ, μᵀΛμ, Λ, logdetΛ
	end

# ╔═╡ 4154a267-326d-495c-aa60-156560aa5f72
Qλ = 2

# ╔═╡ 4c69229d-a272-49eb-9dc1-484c53bad338
pλ = Normal(Qλ)

# ╔═╡ eebb75b8-8ece-4e26-972b-ed962ce1dd42
Qσ = 10

# ╔═╡ 3a303018-4560-4553-96f6-1600b0e1d6c6
pσ = Normal(Qσ)

# ╔═╡ b629f310-dbf7-4b6e-9c3d-199a31a1bfb5
Qπ = 2

# ╔═╡ 51d10c40-5dce-47e2-aeba-958beadda7b8
pπ = Normal(Qπ)

# ╔═╡ c352df19-b653-43f3-b490-046447cea85a
phonemodel = NormalModel((pλ=pλ, pσ=pσ, pπ=pπ))

# ╔═╡ c47dcaf4-049a-45e1-853f-d589283b2038
# The dimension of the parameter space is determined
# by the dimension of the observed space
K = Int(2D + (D^2 - D) / 2)

# ╔═╡ add6cb90-47bd-40e8-bbbb-34938714a79b
pL = JointNormalFixedCov(zeros(Qπ + 1, K))

# ╔═╡ 5415acd0-5e5f-40db-97eb-91b659a1ac5b
pS = JointNormalFixedCov(zeros(Qπ + 1, K))

# ╔═╡ b69c63d7-3f45-45d2-8278-05bc2fb5f3f1
md"""
### Hyper-parameters

| name          | type            | value|  description                  |
|:--------------|:----------------|:----:|:------------------------------|
| ``Q_\lambda`` | hyper-parameter | $Qλ  | language embedding dimension  |
| ``Q_\sigma``  | hyper-parameter | $Qσ  | speaker embedding dimension   |
| ``Q_\pi``     | hyper-parameter | $Qπ  | phone embedding dimension     |
| ``D``         | constant        | $D   | observed space dimension      |
| ``K``         | constant        | $K   | phone model natural parameter dimension|
| ``L`` 		| constant 		  | $L   | number of languages in the data |
| ``S``         | constant        | $S   | number of speakers in  the data |
| ``P``         | constant        | $P   | number of phones in the data  |
"""

# ╔═╡ 946af84f-a5f4-4780-a8c5-13d4acbb9a80
md"""
### Model parameters summary

```math
\begin{align}

\boldsymbol{\theta} &= \Big\{ \{\mathbf{L}_i\}, \{\mathbf{S}_j \}, \{\boldsymbol{\lambda}_l \}, \{\boldsymbol{\sigma}_s \}, \{ \boldsymbol{\pi}_p \} \Big\} \\
\{\mathbf{L}_i \}, \; &\forall i \in \{1, \dots, Q_\lambda \} \\
\{\mathbf{S}_j \}, \;  &\forall j \in \{1, \dots, Q_\sigma \} \\
\{\mathbf{\boldsymbol{\lambda}}_l \}, \; &\forall l \in \{1, \dots, L\} \\
\{\mathbf{\boldsymbol{\sigma}}_s \}, \; &\forall s \in \{1, \dots, S\} \\
\{\mathbf{\boldsymbol{\pi}_p} \}, \; &\forall p \in \{1, \dots, P\}
\end{align}
```
"""

# ╔═╡ f2f5e843-3505-421e-9c9c-71a564387bb2
# Complete model: hierarchical generalized subspace model.
hgsm = GSMPhonetic(f, pL, pS, phonemodel)

# ╔═╡ 1dd9bfa7-b7c9-4075-9c6c-2f33871d5e0a
md"""
### Variational posterior

```math
\begin{align}
	q(\boldsymbol{\theta}) &= \prod_{m}q(\mathbf{L}_m) \prod_{n} q(\mathbf{S}_n) \prod_{l} q(\boldsymbol{\lambda}_l) \prod_{s} q(\boldsymbol{\sigma}_s) \prod_{p} q(\boldsymbol{\pi}_p) \\
	q(\mathbf{L}_m) &= \delta(\mathbf{L}_m - \boldsymbol{\mu}^{L}_m) \\
	q(\mathbf{S}_n) &=  \delta(\mathbf{S}_n - \boldsymbol{\mu}^{S}_n) \\
	q(\boldsymbol{\lambda}_l) &= \mathcal{N}(\boldsymbol{\mu}^{\lambda}_l, \boldsymbol{\Sigma}^{\lambda}_l) \\
	q(\boldsymbol{\sigma}_s) &= \mathcal{N}(\boldsymbol{\mu}^{\sigma}_s, \boldsymbol{\Sigma}^{\sigma}_s) \\
		q(\boldsymbol{\pi}_p) &= \mathcal{N}(\boldsymbol{\mu}^{\pi}_p, \boldsymbol{\Sigma}^{\pi}_p)
\end{align}
```
"""

# ╔═╡ 6f1225eb-8df1-4991-8292-c0b5f2c45c10
q₀L = [JointNormalFixedCov(i == Qλ + 1 ? zeros(Qπ + 1, K) : randn(Qπ + 1, K) / K)
	   for i = 1:(Qλ+1)]

# ╔═╡ 349ef71a-55f8-4c9a-8216-d6fbce8d42e2
q₀S = [JointNormalFixedCov(j == Qσ + 1 ? zeros(Qπ + 1, K) : randn(Qπ + 1, K) / K)
	   for j = 1:(Qσ+1)]

# ╔═╡ 27fa667c-4a76-4fed-9afc-7482dc1339a0
q₀λ = [Normal(zeros(Qλ), Matrix(I, Qλ, Qλ)) for lang in languages]

# ╔═╡ 3a5cadd7-d15b-4da3-8a0b-9e7678514ff0
q₀σ = [Normal(zeros(Qσ), Matrix(I, Qσ, Qσ)) for spk in speakers]

# ╔═╡ e343d7f3-9a50-4952-b429-40b982b2b49f
q₀π = [Normal(zeros(Qπ), Matrix(I, Qπ, Qπ)) for phone in phones]

# ╔═╡ 42ccdf0d-8a5f-404d-8965-cccc7f85d164
q₀θ = (q₀L, q₀S, q₀λ, q₀σ, q₀π)

# ╔═╡ ffe22379-2f91-4b54-8f74-413ac6d2ce08
md"""### Minimum Divergence"""

# ╔═╡ 4d374801-703c-4fb0-b77d-bbb41fbfd54b
function renorm_hyperbases(qM, qh, ph)
	μh = map(μ, qh)
	stdh = stdparams.(qh, μh)

	μh₀ = sum(μh) / length(qh)
	m₀, Σ₀ = stdparams(ph, μh₀)
	L₀ = cholesky(Σ₀).L
	L₀⁻¹ = inv(L₀)

	new_qh = [Normal(L₀⁻¹ * (m - m₀), Symmetric(L₀⁻¹ * Σ * L₀⁻¹')) for (m, Σ) in stdh]

	B = vec.(getproperty.(qM, (:M)))
	dims = size(qM[1].M)
	M, b = hcat(B[1:end-1]...), B[end]
	new_M = reshape.(eachcol(M * L₀), dims...)
	new_b = reshape(b + M * m₀, dims...)

	new_qM = [JointNormalFixedCov(V) for V in [new_M..., new_b]]
	new_qM, new_qh
end

# ╔═╡ 546fd344-42e2-4881-9607-da99db000c33
function renorm_bases(qL, qS, qh, ph)
	μh = map(μ, qh)
	stdh = stdparams.(qh, μh)

	μh₀ = sum(μh) / length(qh)
	m₀, Σ₀ = stdparams(ph, μh₀)
	L₀ = cholesky(Σ₀).L
	L₀⁻¹ = inv(L₀)

	new_qh = [Normal(L₀⁻¹ * (m - m₀), Symmetric(L₀⁻¹ * Σ * L₀⁻¹')) for (m, Σ) in stdh]

	Ms = getproperty.(qL, (:M))
	new_Ms = map(Ms) do M
		W, b = M[1:end-1,:], M[end:end, :]
		vcat(L₀' * W, b + m₀' * W)
	end
	new_qL = [JointNormalFixedCov(V) for V in new_Ms]

	Ms = getproperty.(qS, (:M))
	new_Ms = map(Ms) do M
		W, b = M[1:end-1,:], M[end:end, :]
		vcat(L₀' * W, b + m₀' * W)
	end
	new_qS = [JointNormalFixedCov(V) for V in new_Ms]

	new_qL, new_qS, new_qh
end

# ╔═╡ 8b939634-e3f2-4a29-91de-4ed56787bc78
function minimum_divergence(hgsm, qθ)
	qL, qS, qλ, qσ, qπ = qθ

	qL2, qλ2 = renorm_hyperbases(qL, qλ, hgsm.model.pθ.pλ)
	qS2, qσ2 = renorm_hyperbases(qS, qσ, hgsm.model.pθ.pσ)
	qL2, qS2, qπ2 = renorm_bases(qL2, qS2, qπ, hgsm.model.pθ.pπ)

	qL2, qS2, qλ2, qσ2, qπ2
end

# ╔═╡ 85e89c27-0371-42a9-8b3c-4fc677d1f7a7
md"""
### Training

"""

# ╔═╡ 2822d23b-197d-4dcc-85d0-3df653fcf3b6
begin
	qθ_md = q₀θ

	kls_md = []
	loss_md = []
	for t in 1:100000
		μ = getstats(hgsm, qθ_md)
		(l, (∇μ,)) = withgradient(μ -> elbo(hgsm, groups, qθ_md, μ) / N, μ)

		qθ_md = newposterior(hgsm, qθ_md, ∇μ; lrate=1e-2)

		if t % 100 == 0
			qθ_md = minimum_divergence(hgsm, qθ_md)
		end

		push!(kls_md, kldiv(hgsm, qθ_md, getstats(hgsm, qθ_md)))
		push!(loss_md, l)
	end
	loss_md
end

# ╔═╡ 5162a02c-b803-4be5-bff8-801ec24001d3
plot(loss_md[1000:1:end], legend=false, linealpha=0.5)

# ╔═╡ 05c5717a-c0fa-47a4-9de9-c747b6fb3f93
md"""### Visualization"""

# ╔═╡ 4274c7f4-38d5-4bf0-87ae-417d715d9b3b
begin
	local qλ = qθ_md[3]
	local p = plot(legend=false)
	for (i, q) in enumerate(qλ)
		x, diagxxᵀ, hxxᵀ = unpack(q, μ(q))
		m = x
		Σ = diagm(diagxxᵀ) + CompressedSymmetric(Qλ, 1, hxxᵀ) - x * x'
		m = [m[1], m[2]]
		Σ = Σ[1:2,1:2]
		annotate!(m[1], m[2], id2lang[i])
		plotnormal!(m, Σ, σ=2, fillalpha=0.4, linealpha=0, label=false)
	end
	p
end

# ╔═╡ aab42b7b-b269-4e6f-af74-f09ec9c6b862
begin
	local qσ = qθ_md[4]
	local p = plot(size=(800,800), legend=false, xlims=(-3, 3), ylims=(-3, 3))
	for (i, q) in enumerate(qσ)
		x, diagxxᵀ, hxxᵀ = unpack(q, μ(q))
		m = x
		Σ = diagm(diagxxᵀ) + CompressedSymmetric(Qσ, 1, hxxᵀ) - x * x'
		m = [m[1], m[2]]
		Σ = [Σ[1,1] Σ[1,2]; Σ[2, 1] Σ[2,2]] #Σ[1:2,1:2]
		#annotate!(m[1], m[2], speakers[i][1])
		color = speakers[i][2] == "male" ? :blue : :green
		plotnormal!(m, Σ, σ=2, fillcolor=color, fillalpha=0.4, linealpha=0, label=false)
	end
	p
end

# ╔═╡ b9bba539-ba31-420e-a718-a952db80d4c2
begin
	local qπ = qθ_md[5]
	local d1, d2 = 1, 2
	local p = plot(size=(600, 500), legend=false, xlabel="π₁", ylabel="π₂")
	for (i, q) in enumerate(qπ)
		x, diagxxᵀ, hxxᵀ = unpack(q, μ(q))
		m = x
		Σ = diagm(diagxxᵀ) + CompressedSymmetric(Qπ, 1, hxxᵀ) - x * x'
		m = [m[d1], m[d2]]
		Σ = [Σ[d1, d1] Σ[d1, d2]; Σ[d2, d1] Σ[d2, d2]]
		annotate!(m[1], m[2], id2phone[i])
		plotnormal!(m, Σ, σ=2, fillcolor=:blue, fillalpha=0.1, linealpha=0, label=false)
	end
	p
end

# ╔═╡ f7d8e930-3d7b-4135-bb64-a8b17367f683
qL, qS, qλ, qσ, qπ = qθ_md

# ╔═╡ cd347af2-ad3e-4e7d-b471-54b668511e69
function hgsm_η(hgsm::GSMPhonetic, qL, qS, λ, σ, π)
	L = getproperty.(qL, (:M,))
	S = getproperty.(qS, (:M,))
	P = sum(L .* vcat(λ, 1)) + sum(S .* vcat(σ, 1))
	ψ = P' * vcat(π, 1)
	Λμ, μᵀΛμ, Λ, logdetΛ = hgsm.f(ψ)
	vcat(Λμ, -(1/2)*diag(Λ), -vech(Λ, 1))
end

# ╔═╡ 7d457821-67c0-47e0-9c00-7bd2a6560cc8
hgsm_η(hgsm, qL, qS, qλ[1].μ, qσ[1].μ, qπ[1].μ)

# ╔═╡ bffb14d2-1173-4511-80b2-0676af0cd3f8
function hgsm_μ(hgsm::GSMPhonetic, η)
	L = length(η)
	D = Int((-(3/2) + sqrt((9/4) + 2L)))

	Λμ = η[1:D]
	diagΛ = -2 * η[D+1:2*D]
	vechΛ = - η[2*D+1:end]
	Λ = diagm(diagΛ) + CompressedSymmetric(D, 1, vechΛ)

	Σ = inv(Symmetric(Λ))
	μ = Σ * Λμ
	xxᵀ = Σ + μ * μ'
	vcat(μ, diag(xxᵀ), vech(xxᵀ, 1))
end

# ╔═╡ 3a4921ab-d8e4-463b-ae41-d9ac4b034c41
hgsm_μ(hgsm, [0., 0., -1, -1, 0.5])

# ╔═╡ c3f6fc47-2843-498d-8d86-0c99e8638ff4
hgsm_μ(hgsm, hgsm_η(hgsm, qL, qS, qλ[1].μ, qσ[1].μ, qπ[1].μ))

# ╔═╡ 5d54df87-0d8a-4371-93e9-cc20fd659974
function hgsm_A(hgsm::GSMPhonetic, η)
	L = length(η)
	D = Int((-(3/2) + sqrt((9/4) + 2L)))

	Λμ = η[1:D]
	diagΛ = -2 * η[D+1:2*D]
	vechΛ = - η[2*D+1:end]
	Λ = diagm(diagΛ) + CompressedSymmetric(D, 1, vechΛ)
	Σ = inv(Λ)

	-(1/2) * logdet(Λ) + (1/2) * Λμ' * (Σ * Λμ)
end

# ╔═╡ f0d90bdc-aa48-4edf-a673-e6dc6d27e2e0
function trajectory_λ(hgsm, qL, qS, λ₁, λ₂, σ, π; step=1e-2)
	traj = [λ₁]
	dist = 0
	H = η -> ForwardDiff.hessian(η -> hgsm_A(hgsm, η), η)
	J = λ -> ForwardDiff.jacobian(λ -> hgsm_η(hgsm, qL, qS, λ, σ, π), λ)

	η₂ = hgsm_η(hgsm, qL, qS, λ₂, σ, π)

	while norm(traj[end] - λ₂) > 1e-1
		λ = traj[end]
		ηₜ = hgsm_η(hgsm, qL, qS, λ, σ, π)
		μₜ = hgsm_μ(hgsm, ηₜ)
		Jλ = J(λ)
		∇λ = (1/2) * Jλ' * (η₂ - ηₜ)
		#∇λ = -Jλ' * (ηₜ - (η₂ + μₜ))
		#∇λ = (λ₂ - λ)
		∇λ = ∇λ ./ norm(∇λ)

		G = Jλ' * Symmetric(H(ηₜ)) * Jλ
		v = ∇λ #inv(G) * ∇λ
		v = step * v ./ norm(v)
		dist += sqrt((step * ∇λ)' * G * (step * ∇λ))
		#dist += sqrt((step * ∇λ)' * G * (step * ∇λ))

		push!(traj, λ + v)
	end
	hcat(traj...), dist
end

# ╔═╡ 0d8c54e2-3677-4d5d-adc1-b6b3f482e409
t1, dist1 = trajectory_λ(hgsm, qL, qS, qλ[1].μ, qλ[6].μ, qσ[1].μ, qπ[1].μ)

# ╔═╡ b4c44e20-ffbc-4376-8cc6-3d9b1f64af9a
t2, dist2 = trajectory_λ(hgsm, qL, qS, qλ[6].μ, qλ[1].μ, qσ[1].μ, qπ[1].μ)

# ╔═╡ b317dbf6-2a60-400f-96d5-addd740be016
begin
	# Manual sorting so that languages from the same family appear together
	sorted_languages = [:spa, :ita, :fra, :rus, :pol, :deu]

	dist_lang = zeros(L, L)
	local idmap = Dict(i => findall(l -> l == lang, languages)[1] for (i, lang) in enumerate(sorted_languages))
	for i = 1:L, j = 1:L
		x, y = idmap[i], idmap[j]
		for phone in phones_per_lang[languages[x]]
			p = phone2id[phone]
			println(i, ": ", languages[x], " ", j, ": ", languages[y], " ", phones[p])
			traj, Δ = trajectory_λ(hgsm, qL, qS, qλ[x].μ, qλ[y].μ, zeros(Qσ), qπ[p].μ)
			dist_lang[i, j] += Δ / length(phones_per_lang[languages[x]])
		end
	end
end

# ╔═╡ 8d416ab7-3ae7-4984-a7e6-1f56692d1b7b
begin
	mdstr = "| language | # speakers | phones |\n"
	mdstr *= "|:---------|:-----------:|:-------|\n"
	for lang in sorted_languages
		mdstr *= "| $lang | "

		ns = 0
		for (l, spk) in speakers
			if l == lang ns += 1end
		end
		mdstr *= "$ns | "

		for phone in sort(collect(phones_per_lang[lang]))
			mdstr *= "$phone "
		end
		mdstr *= "|\n"
	end
	Markdown.parse(mdstr)
end

# ╔═╡ 69777c50-240b-4be7-a824-b5dd0eb5fb7c
languages

# ╔═╡ 00d323cd-b928-49cd-a590-0e8852ffb264
begin
	heatmap(
		dist_lang,
		xticks=(1:6, sorted_languages),
		yticks=(1:6, sorted_languages),
		tickfont = (14,),
		xmirror=true,
		yflip=true,
		c=:grays
	)
end

# ╔═╡ 4d44afe7-b0dc-40b1-a00d-73ad346e0576
dist_lang

# ╔═╡ cea584da-2d0e-4c4f-a964-2f7254bf74aa
languages

# ╔═╡ 295da023-8cb9-4b13-99f8-fa7690c20b5c
begin
	local d1, d2 = 1, 2
	local qλ = qθ_md[3]
	local p = plot(legend=false)
	for (i, q) in enumerate(qλ)
		x, diagxxᵀ, hxxᵀ = unpack(q, μ(q))
		m = x
		Σ = diagm(diagxxᵀ) + CompressedSymmetric(Qλ, 1, hxxᵀ) - x * x'
		m = [m[d1], m[d2]]
		Σ = [Σ[d1, d1] Σ[d1, d2]; Σ[d2, d1] Σ[d2, d2]]
		annotate!(m[1], m[2], id2lang[i])
		plot!(t1[d1, :], t1[d2, :], linecolor=:black)
		plot!(t2[d1, :], t2[d2, :], linecolor=:black)
		plotnormal!(m, Σ, σ=2, fillalpha=0.4, linealpha=0, label=false)
	end
	p
end

# ╔═╡ d15d2ae4-01b4-4312-b921-2fac2d943036
function hgsm_llh(hgsm::GSMPhonetic, X, qL, qS, λ, σ, π)
	L = getproperty.(qL, (:M,))
	S = getproperty.(qS, (:M,))
	P = sum(L .* vcat(λ, 1)) + sum(S .* vcat(σ, 1))
	ψ = P' * vcat(π, 1)
	#Λμ, μᵀΛμ, Λ, logdetΛ = hgsm.f(ψ)
	#vcat(Λμ, -(1/2)*diag(Λ), -vech(Λ, 1))
	loglikelihood(hgsm.model, X, hgsm.f(ψ))
end

# ╔═╡ fa30a558-5bed-4ce1-8be2-ed55a3d7a9f8
function hgsm_stdparams(hgsm::GSMPhonetic, qL, qS, λ, σ, π)
	L = getproperty.(qL, (:M,))
	S = getproperty.(qS, (:M,))
	P = sum(L .* vcat(λ, 1)) + sum(S .* vcat(σ, 1))
	ψ = P' * vcat(π, 1)
	Λμ, μᵀΛμ, Λ, logdetΛ = hgsm.f(ψ)
	Σ = inv(Λ)
	μ = Σ * Λμ
	Σ, μ
end

# ╔═╡ a36c7fd9-6e47-4386-a9e0-cb9a9dc6a644
id2spk

# ╔═╡ 0f26901f-e592-4768-ac67-ee526a640ede
phone2id

# ╔═╡ 8509deeb-f156-4350-a21a-d1d70949091d
lang2id

# ╔═╡ 696d1d7d-68df-45c3-b5d9-0cae311e553a
repeat([2], 10)

# ╔═╡ 6836772d-67c1-475f-8d99-e11fc7f752a0
begin
	local rawdata = readdlm(localfile)
	local data = []

	#pdataraw = Dict("œ" => [], "ɛ" => [])
	#pdataraw = Dict("ã" => [], "õ" => [])
	#pdataraw = Dict("ẽ" => [], "õ" => [])
	pdataraw = Dict("$p" => [] for p in phones)
	testphones = collect(keys(pdataraw))
	for (lang, spk, sex, ipaphone, f1, f2, f3, duration) in eachrow(rawdata)
		# Remove the duration symbol
		ipaphone = replace(ipaphone, "ː" => "")

		lang, spk, ipaphone = lang, spk, ipaphone

		#if lang == "fra" && spk == "spk8" && sex == "male" && ipaphone in keys(pdataraw)
		if ipaphone in keys(pdataraw)
			push!(pdataraw[ipaphone], vcat(f1, f2, f3))
		end
		push!(data, vcat(f1, f2, f3))
	end

	pdata = Dict()
	for (k, v) in pdataraw
		pdata[k] = hcat(v...)
	end
	data = hcat(data...)

	local x̄ = sum(data, dims=2) / size(data, 2)
	local x̄² = sum(data .^ 2, dims=2) / size(data, 2)
	local var = x̄² - (x̄ .^ 2)
	for (k, v) in pdata
		pdata[k] = (v .- x̄) ./ sqrt.(var)
	end
	pdata
end

# ╔═╡ 9c9d2747-b6db-4957-84be-3c327c4dc1df
begin
	local spk = 3
	local lang = 4
	#local testphones = ["ẽ", "õ"] #["œ", "ɛ"]
	p = plot(xlims=(-3, 3), ylims=(-3, 3))
	for phone in testphones
		i = phone2id[Symbol(phone)]
		Σ, m = hgsm_stdparams(hgsm, qL, qS, qλ[lang].μ, qσ[spk].μ, qπ[i].μ)
		plotnormal!(m[1:2], Σ[1:2,1:2], σ=2, fillalpha=0.4, linealpha=0, label="$phone")
	end

	for phone in testphones
		scatter!(pdata["$phone"][1,1:10], pdata["$phone"][2,:], label="$phone")
	end
	p
end

# ╔═╡ d711ad75-3fc6-4726-a506-962aaccb2de7
begin
	local p = plot()
	for phone in testphones
		scatter!(pdata["$phone"][1,1:10], pdata["$phone"][2,1:10], label="$phone")
	end
	p
end

# ╔═╡ c80d9947-e02e-4fb0-bcb3-e29a4702688e
begin
	local spk = 3
	local i = phone2id[Symbol(testphones[1])]
	local j = phone2id[Symbol(testphones[2])]
	local errors = zeros(L)

	for lang in 1:L

		llh1 = vcat(
			hgsm_llh(hgsm, pdata[testphones[1]], qL, qS, qλ[lang].μ, qσ[spk].μ, qπ[i].μ),
			hgsm_llh(hgsm, pdata[testphones[1]], qL, qS, qλ[lang].μ, qσ[spk].μ, qπ[j].μ)
		)

		llh2 = vcat(
			hgsm_llh(hgsm, pdata[testphones[2]], qL, qS, qλ[lang].μ, qσ[spk].μ, qπ[i].μ),
			hgsm_llh(hgsm, pdata[testphones[2]], qL, qS, qλ[lang].μ, qσ[spk].μ, qπ[j].μ)
		)
		errors[lang] = ( sum(llh1[1,:] .> llh1[2,:]) + sum(llh2[1, :] .< llh2[2, :]) ) / (size(llh1, 2) + size(llh2, 2))
	end
	errors
end

# ╔═╡ 43692f77-f402-41c7-b22e-0b8a42bf8451
speakers

# ╔═╡ 6d1e2c27-4c08-4324-af85-45ed1c4950e0
local rawdata = readdlm(localfile)

# ╔═╡ Cell order:
# ╟─6958bde2-c50f-4afd-8d08-7dfbfd97169a
# ╠═3cd07b90-9e21-4a74-8d99-a6375beb7d1e
# ╟─b83311ab-ca4f-432e-8627-5207517a81db
# ╠═d1588a1c-babe-4223-bc6d-f353314ed2c0
# ╠═262e0395-814f-4315-b25e-21c18a7a9a5f
# ╟─4bbdd3d2-eb62-4889-9c8a-20141c0b140a
# ╠═ee1d1c45-70f5-4fe7-a8e2-9134d45a2b95
# ╟─474abeb1-82f8-4a74-a03c-be1cf55df840
# ╠═31e56fed-e4f0-4c71-abb2-5c2eff5a9a47
# ╠═61cc1ab2-a615-4a9b-91f9-bc5909e29297
# ╠═e08cc384-49dc-4585-a2ab-fbec9fcfe9d2
# ╠═b8e7f92f-0065-46fe-ac6b-5adf29b3379f
# ╠═ef71661f-338e-4d57-86ed-fc2c308a744f
# ╟─c01de0cd-43f8-49cc-a8c9-ffbbc182d9e9
# ╟─8d416ab7-3ae7-4984-a7e6-1f56692d1b7b
# ╟─8b535f17-7a55-42da-ba7a-c161bec9fd28
# ╠═4c69229d-a272-49eb-9dc1-484c53bad338
# ╠═3a303018-4560-4553-96f6-1600b0e1d6c6
# ╠═51d10c40-5dce-47e2-aeba-958beadda7b8
# ╟─5175712e-7f27-459e-8e6a-9bed325d6b4e
# ╠═add6cb90-47bd-40e8-bbbb-34938714a79b
# ╠═5415acd0-5e5f-40db-97eb-91b659a1ac5b
# ╟─763b5f95-db02-4ede-b632-848a646232d9
# ╠═c352df19-b653-43f3-b490-046447cea85a
# ╟─841a3f0e-8975-49d2-959e-afe832ed41e4
# ╠═b0761304-49df-4626-a677-488295dc786d
# ╟─b69c63d7-3f45-45d2-8278-05bc2fb5f3f1
# ╠═4154a267-326d-495c-aa60-156560aa5f72
# ╠═eebb75b8-8ece-4e26-972b-ed962ce1dd42
# ╠═b629f310-dbf7-4b6e-9c3d-199a31a1bfb5
# ╠═c47dcaf4-049a-45e1-853f-d589283b2038
# ╟─946af84f-a5f4-4780-a8c5-13d4acbb9a80
# ╠═f2f5e843-3505-421e-9c9c-71a564387bb2
# ╟─1dd9bfa7-b7c9-4075-9c6c-2f33871d5e0a
# ╠═6f1225eb-8df1-4991-8292-c0b5f2c45c10
# ╠═349ef71a-55f8-4c9a-8216-d6fbce8d42e2
# ╠═27fa667c-4a76-4fed-9afc-7482dc1339a0
# ╠═3a5cadd7-d15b-4da3-8a0b-9e7678514ff0
# ╠═e343d7f3-9a50-4952-b429-40b982b2b49f
# ╠═42ccdf0d-8a5f-404d-8965-cccc7f85d164
# ╟─ffe22379-2f91-4b54-8f74-413ac6d2ce08
# ╠═4d374801-703c-4fb0-b77d-bbb41fbfd54b
# ╠═546fd344-42e2-4881-9607-da99db000c33
# ╠═8b939634-e3f2-4a29-91de-4ed56787bc78
# ╟─85e89c27-0371-42a9-8b3c-4fc677d1f7a7
# ╠═2822d23b-197d-4dcc-85d0-3df653fcf3b6
# ╠═5162a02c-b803-4be5-bff8-801ec24001d3
# ╟─05c5717a-c0fa-47a4-9de9-c747b6fb3f93
# ╠═4274c7f4-38d5-4bf0-87ae-417d715d9b3b
# ╠═aab42b7b-b269-4e6f-af74-f09ec9c6b862
# ╠═b9bba539-ba31-420e-a718-a952db80d4c2
# ╠═f7d8e930-3d7b-4135-bb64-a8b17367f683
# ╠═cd347af2-ad3e-4e7d-b471-54b668511e69
# ╠═7d457821-67c0-47e0-9c00-7bd2a6560cc8
# ╠═bffb14d2-1173-4511-80b2-0676af0cd3f8
# ╠═3a4921ab-d8e4-463b-ae41-d9ac4b034c41
# ╠═c3f6fc47-2843-498d-8d86-0c99e8638ff4
# ╠═5d54df87-0d8a-4371-93e9-cc20fd659974
# ╠═f0d90bdc-aa48-4edf-a673-e6dc6d27e2e0
# ╠═0d8c54e2-3677-4d5d-adc1-b6b3f482e409
# ╠═b4c44e20-ffbc-4376-8cc6-3d9b1f64af9a
# ╠═b317dbf6-2a60-400f-96d5-addd740be016
# ╠═69777c50-240b-4be7-a824-b5dd0eb5fb7c
# ╠═00d323cd-b928-49cd-a590-0e8852ffb264
# ╠═4d44afe7-b0dc-40b1-a00d-73ad346e0576
# ╠═cea584da-2d0e-4c4f-a964-2f7254bf74aa
# ╠═295da023-8cb9-4b13-99f8-fa7690c20b5c
# ╠═d15d2ae4-01b4-4312-b921-2fac2d943036
# ╠═fa30a558-5bed-4ce1-8be2-ed55a3d7a9f8
# ╠═9c9d2747-b6db-4957-84be-3c327c4dc1df
# ╠═d711ad75-3fc6-4726-a506-962aaccb2de7
# ╠═c80d9947-e02e-4fb0-bcb3-e29a4702688e
# ╠═a36c7fd9-6e47-4386-a9e0-cb9a9dc6a644
# ╠═0f26901f-e592-4768-ac67-ee526a640ede
# ╠═8509deeb-f156-4350-a21a-d1d70949091d
# ╠═696d1d7d-68df-45c3-b5d9-0cae311e553a
# ╠═6836772d-67c1-475f-8d99-e11fc7f752a0
# ╠═43692f77-f402-41c7-b22e-0b8a42bf8451
# ╠═6d1e2c27-4c08-4324-af85-45ed1c4950e0

examples/Learning and Modeling Languages Phonetics v2.jl


"""
```julia
produce_orbitdiagram(
	ds::ContinuousDynamicalSystem, plane, i::Int, p_index, pvalues; kwargs...
)
```
Produce an orbit diagram (sometimes wrongly called bifurcation diagram)
for the `i` variable(s) of the given continuous
system by computing Poincaré surfaces of section using `plane`
for the given parameter values (see [`poincaresos`](@ref)).

`i` can be `Int` or `AbstractVector{Int}`.
If `i` is `Int`, returns a vector of vectors. Else
it returns a vector of vectors of vectors.
Each entry are the points at each parameter value.

## Keyword Arguments
* `printparams::Bool = false` : Whether to print the parameter used
  during computation in order to keep track of running time.
* `direction, warning, Ttr, rootkw, diffeq` :
  Propagated into [`poincaresos`](@ref).
* `u0 = nothing` : Specify an initial state. If `nothing`, the previous state after each
  parameter is used to seed the new initial condition at the new parameter
  (with the very first state being the system's state). This makes convergence to the
  attractor faster, necessitating smaller `Ttr`. Otherwise `u0` can be a standard state,
  or a vector of states, so that a specific state is used for each parameter.

## Description
For each parameter, a PSOS reduces the system from a flow to a map. This then allows
the formal computation of an "orbit diagram" for the `i` variable
of the system, just like it is done in [`orbitdiagram`](@ref).

The parameter change is done as `p[p_index] = value` taking values from `pvalues`
and thus you must use a parameter container that supports this
(either `Array`, `LMArray`, dictionary or other).

See also [`poincaresos`](@ref), [`orbitdiagram`](@ref).
"""
function produce_orbitdiagram(
        ds::CDS{IIP, S, D}, plane, idxs, p_index, pvalues;
        tfinal::Real = 100.0, direction = -1, printparams = false, warning = true,
        Ttr = 0.0, u0 = nothing, rootkw = (xrtol = 1e-6, atol = 1e-6),
        diffeq = NamedTuple(), kwargs...
    ) where {IIP, S, D}

    if !isempty(kwargs)
        @warn DIFFEQ_DEP_WARN
        diffeq = NamedTuple(kwargs)
    end

    i = typeof(idxs) <: Int ? idxs : SVector{length(idxs), Int}(idxs...)

    _check_plane(plane, D)
    typeof(u0) <: Vector{<:AbstractVector} && @assert length(u0)==length(p)
    integ = integrator(ds; diffeq)
    planecrossing = PlaneCrossing(plane, direction > 0)
    p0 = ds.p[p_index]
    output = Vector{typeof(ds.u0[i])}[]
	plane_distance = (t) -> planecrossing(integ(t))

    for (n, p) in enumerate(pvalues)
        integ.p[p_index] = p
        printparams && println("parameter = $p")
        if typeof(u0) <: Vector{<:AbstractVector}
            st = u0[n]
        elseif isnothing(u0)
            st = integ.u
        else
            st = u0
        end
        reinit!(integ, st)
		Ttr != 0 && step!(integ, Ttr)
		data = _poincaresos(integ, plane_distance, planecrossing, tfinal+Ttr, i, rootkw)
		push!(output, data)
        warning && length(output[end]) == 0 && @warn "For parameter $p $PSOS_ERROR"
    end
    # Reset the parameter of the system:
    ds.p[p_index] = p0
    return output
end


@doc raw"""
`SDE`: Stratonovich Stochastic Differential Equation

Defines a stochastic differential initial value problem
```math
\begin{align*}
\dq (t) &= v(t, q(t)) \, dt + B(t, q(t)) \circ dW , & q(t_{0}) &= q_{0} ,
\end{align*}
```
with drift vector field ``v``, diffusion matrix ``B``,
initial conditions ``q_{0}``, the dynamical variable ``q``
taking values in ``\mathbb{R}^{d}``, and the m-dimensional Wiener process W

### Fields

* `d`:  dimension of dynamical variable ``q`` and the vector field ``v``
* `m`:  dimension of the Wiener process
* `n`:  number of initial conditions
* `ns`: number of sample paths
* `v`:  function computing the deterministic vector field
* `B`:  function computing the d x m diffusion matrix
* `t₀`: initial time
* `q₀`: initial condition for dynamical variable ``q`` (may be a random variable itself)

The functions `v` and `B`, providing the drift vector field and diffusion matrix,
`v(t, q, v)` and `B(t, q, B; col=0)`, where `t` is the current time, `q` is the
current solution vector, and `v` and `B` are the variables which hold the result
of evaluating the vector field ``v`` and the matrix ``B`` on `t` and `q` (if col==0),
or the column col of the matrix B (if col>0).

### Example

```julia
    function v(λ, t, q, v)
        v[1] = λ*q[1]
        v[2] = λ*q[2]
    end

    function B(μ, t, q, B; col=0)
        if col==0 #whole matrix
            B[1,1] = μ*q[1]
            B[2,1] = μ*q[2]
        elseif col==1
            #just first column
        end
    end

    t₀ = 0.
    q₀ = [1., 1.]
    λ  = 2.
    μ  = 1.

    v_sde = (t, q, v) -> v(λ, t, q, v)
    B_sde = (t, q, B) -> B(μ, t, q, B)

    sde = SDE(v_sde, B_sde, t₀, q₀)
```
"""
struct SDE{dType <: Number, tType <: Number, vType <: Function, BType <: Function, N} <: Equation{dType, tType}
    d::Int
    m::Int
    n::Int
    ns::Int
    v::vType
    B::BType
    t₀::tType
    q₀::Array{dType, N}           #Initial condition: N=1 - single deterministic, N=2 - single random or multiple deterministic, N=3 - multiple random
    periodicity::Vector{dType}

    function SDE{dType,tType,vType,BType,N}(d, m, n, ns, v, B, t₀, q₀; periodicity=[]) where {dType <: Number, tType <: Number, vType <: Function, BType <: Function, N}

        @assert dType == eltype(q₀)
        @assert tType == typeof(t₀)
        @assert ndims(q₀) == N
        @assert d == size(q₀,1)

        if ns==1 && n==1
            # single sample path and single initial condition, therefore N=1
            @assert N == 1
        elseif ns>1 && n==1
            # single initial condition, but may be random (N=2) or deterministic (N=1)
            @assert N ∈ (1,2)
            if N==2
                @assert ns == size(q₀,2)
            end
        elseif ns==1 && n>1
            # multiple  deterministic initial conditions, so N=2
            @assert N == 2
            @assert n == size(q₀,2)
        elseif ns>1 && n>1
            # either multiple random initial conditions (N=3),
            # or multiple deterministic initial conditions (N=2)
            @assert N ∈ (2,3)

            if N==2
                @assert n == size(q₀,2)
            else
                @assert ns == size(q₀,2)
                @assert n == size(q₀,3)
            end
        end

        if !(length(periodicity) == d)
            periodicity = zeros(dType, d)
        end

        new(d, m, n, ns, v, B, t₀, q₀, periodicity)
    end
end

function SDE(m::Int, ns::Int, v::VT, B::BT, t₀::TT, q₀::DenseArray{DT,1}; periodicity=[]) where {DT,TT,VT,BT}
    # A 1D array q₀ contains a single deterministic initial condition, so n=1, but we still need to specify
    # the number of sample paths ns
    SDE{DT, TT, VT, BT, 1}(size(q₀, 1), m, 1, ns, v, B, t₀, q₀, periodicity=periodicity)
end

# A 2-dimensional matrix q0 can represent a single random initial condition with ns>1 and n=1,
# or a set of deterministic initial conditions with n>1 (for which we can have both ns=1 and ns>1)
# The function below assumes q0 to represent a single random initial condition (n=1, ns=size(q₀, 2))
function SDE(m::Int, v::VT, B::BT, t₀::TT, q₀::DenseArray{DT,2}; periodicity=[]) where {DT,TT,VT,BT}
    SDE{DT, TT, VT, BT, 2}(size(q₀, 1), m, 1, size(q₀, 2), v, B, t₀, q₀, periodicity=periodicity)
end

# On the other hand, the function below assumes q₀ represents multiple deterministic initial conditions
# (n=size(q₀, 2)), but these initial conditions may be run an arbitrary number ns of sample paths, so ns has to be explicitly specified
function SDE(m::Int, ns::Int, v::VT, B::BT, t₀::TT, q₀::DenseArray{DT,2}; periodicity=[]) where {DT,TT,VT,BT}
    SDE{DT, TT, VT, BT, 2}(size(q₀, 1), m, size(q₀, 2), ns, v, B, t₀, q₀, periodicity=periodicity)
end

# # OLD FUNCTION FOR A 2D matrix q₀
# # The argument IC specifies whether there are multiple deterministic initial conditions
# # (IC=true, so n>1) or a single random one (default IC=false, so n=1)
# function SDE(m::Int, v::VT, B::BT, t₀::TT, q₀::DenseArray{DT,2}; periodicity=[], IC=false) where {DT,TT,VT,BT}
#     if IC==true
#         SDE{DT, TT, VT, BT, 2}(size(q₀, 1), m, size(q₀, 2), 1, v, B, t₀, q₀, periodicity=periodicity)
#     else
#         SDE{DT, TT, VT, BT, 2}(size(q₀, 1), m, 1, size(q₀, 2), v, B, t₀, q₀, periodicity=periodicity)
#     end
# end

function SDE(m::Int, v::VT, B::BT, t₀::TT, q₀::DenseArray{DT,3}; periodicity=[]) where {DT,TT,VT,BT}
    # A 3D array q₀ contains multiple random initial condition, so n=size(q₀,3) and ns=size(q₀,2)
    SDE{DT, TT, VT, BT, 3}(size(q₀, 1), m, size(q₀,3), size(q₀,2), v, B, t₀, q₀, periodicity=periodicity)
end

function SDE(m::Int, ns::Int, v::VT, B::BT, q₀::DenseArray{DT,1}; periodicity=[]) where {DT,VT,BT}
    SDE(m, ns, v, B, zero(DT), q₀, periodicity=periodicity)
end

# Assumes q0 represents a single random initial condition (n=1, ns=size(q₀, 2))
function SDE(m::Int, v::VT, B::BT, q₀::DenseArray{DT,2}; periodicity=[]) where {DT,VT,BT}
    SDE(m, v, B, zero(DT), q₀, periodicity=periodicity)
end

# Assumes q₀ represents multiple deterministic initial conditions (n=size(q₀, 2))
function SDE(m::Int, ns::Int, v::VT, B::BT, q₀::DenseArray{DT,2}; periodicity=[]) where {DT,VT,BT}
    SDE(m, ns, v, B, zero(DT), q₀, periodicity=periodicity)
end

function SDE(m::Int, v::VT, B::BT, q₀::DenseArray{DT,3}; periodicity=[]) where {DT,VT,BT}
    SDE(m, v, B, zero(DT), q₀, periodicity=periodicity)
end

Base.hash(sde::SDE, h::UInt) = hash(sde.d, hash(sde.m, hash(sde.n, hash(sde.ns, hash(sde.v, hash(sde.B, hash(sde.t₀, hash(sde.q₀, hash(sde.periodicity, h)))))))))

Base.:(==)(sde1::SDE, sde2::SDE) = (
                                sde1.d == sde2.d
                             && sde1.m == sde2.m
                             && sde1.n == sde2.n
                             && sde1.ns == sde2.ns
                             && sde1.v == sde2.v
                             && sde1.B == sde2.B
                             && sde1.t₀ == sde2.t₀
                             && sde1.q₀ == sde2.q₀
                             && sde1.periodicity == sde2.periodicity)

function Base.similar(sde::SDE{DT,TT,VT,UT}, q₀::DenseArray{DT,1}) where {DT, TT, VT, UT}
    similar(sde, sde.t₀, q₀)
end

# Assumes q0 represents a single random initial condition (n=1, ns=size(q₀, 2))
function Base.similar(sde::SDE{DT,TT,VT,UT}, q₀::DenseArray{DT,2}) where {DT, TT, VT, UT}
    similar(sde, sde.t₀, q₀)
end

function Base.similar(sde::SDE{DT,TT,VT,UT}, q₀::DenseArray{DT,3}) where {DT, TT, VT, UT}
    similar(sde, sde.t₀, q₀)
end

function Base.similar(sde::SDE{DT,TT,VT,BT}, t₀::TT, q₀::DenseArray{DT,1}) where {DT, TT, VT, BT}
    @assert sde.d == size(q₀,1)
    SDE(sde.m, sde.ns, sde.v, sde.B, t₀, q₀, periodicity=sde.periodicity)
end

# Assumes q0 represents a single random initial condition (n=1, ns=size(q₀, 2))
function Base.similar(sde::SDE{DT,TT,VT,BT}, t₀::TT, q₀::DenseArray{DT,2}) where {DT, TT, VT, BT}
    @assert sde.d == size(q₀,1)
    SDE(sde.m, sde.ns, sde.v, sde.B, t₀, q₀, periodicity=sde.periodicity)
end

function Base.similar(sde::SDE{DT,TT,VT,BT}, t₀::TT, q₀::DenseArray{DT,3}) where {DT, TT, VT, BT}
    @assert sde.d == size(q₀,1)
    SDE(sde.m, sde.ns, sde.v, sde.B, t₀, q₀, periodicity=sde.periodicity)
end

function Base.similar(sde::SDE{DT,TT,VT,UT}, q₀::DenseArray{DT,1}, ns::Int) where {DT, TT, VT, UT}
    similar(sde, sde.t₀, q₀, ns)
end

# Assumes q₀ represents multiple deterministic initial conditions (n=size(q₀, 2))
function Base.similar(sde::SDE{DT,TT,VT,UT}, q₀::DenseArray{DT,2}, ns::Int) where {DT, TT, VT, UT}
    similar(sde, sde.t₀, q₀, ns)
end

function Base.similar(sde::SDE{DT,TT,VT,BT}, t₀::TT, q₀::DenseArray{DT,1}, ns::Int) where {DT, TT, VT, BT}
    @assert sde.d == size(q₀,1)
    SDE(sde.m, ns, sde.v, sde.B, t₀, q₀, periodicity=sde.periodicity)
end

# Assumes q₀ represents multiple deterministic initial conditions (n=size(q₀, 2))
function Base.similar(sde::SDE{DT,TT,VT,BT}, t₀::TT, q₀::DenseArray{DT,2}, ns::Int) where {DT, TT, VT, BT}
    @assert sde.d == size(q₀,1)
    SDE(sde.m, ns, sde.v, sde.B, t₀, q₀, periodicity=sde.periodicity)
end

Base.ndims(sde::SDE) = sde.d


# %% Interactive bifurcations for 1D energy balance
using DrWatson
@quickactivate "NonlinearDynamicsTextbook"
include(srcdir("colorscheme.jl"))
using DynamicalSystems, OrdinaryDiffEq, Random
using InteractiveDynamics: record_interaction
import GLMakie

D = 25
omega_vec = randn(D)

ds = Systems.kuramoto(D)
diffeq = (alg = Tsit5(), adaptive = false, dt = 0.05)
integ = integrator(ds; diffeq...)

fig = GLMakie.Figure()
display(fig)
cmap = :curl
axku = GLMakie.Axis(fig[1,1])

# Other static elements
GLMakie.hidedecorations!(axku)
axku.aspect = GLMakie.DataAspect()
GLMakie.lines!(axku, cos.(0:0.001:2π+0.002), sin.(0:0.001:2π+0.002);color = :black)

# Plot balls
phases = GLMakie.Observable(copy(integ.u))
balls = GLMakie.lift(u -> [GLMakie.Point2f0(cos(φ), sin(φ)) for φ in u], phases)
GLMakie.scatter!(axku, balls;
    markersize  = 12, color = 1:D, colormap = cmap, strokewidth=2, strokecolor = :black
)

# Plot arrows
arrows_end = GLMakie.lift(b -> 0.8b, balls)
arrows_start = [GLMakie.Point2f0(0,0) for i in 1:D]
GLMakie.arrows!(axku, arrows_start, arrows_end;
    color = 1:D, colormap = (cmap, 0.5), arrowsize = 10, linewidth = 4,
)

# Plot mean field Vector
# add sliders and buttons
fig[2, 1] = controllayout = GLMakie.GridLayout(tellwidth = false)
run = controllayout[1, 1] = GLMakie.Button(fig; label = "run")
Ks = 0.0:0.1:10.0
kslider = GLMakie.labelslider!(fig, "K =", Ks; sliderkw = Dict(:startvalue => ds.p.K))
controllayout[1, 2] = kslider.layout

# run button functionality
isrunning = GLMakie.Observable(false)
GLMakie.on(run.clicks) do clicks; isrunning[] = !isrunning[]; end
GLMakie.on(run.clicks) do clicks
    @async while isrunning[]
        GLMakie.isopen(fig.scene) || break # ensures computations stop if closed window
        step!(integ)
        phases[] = integ.u
        yield()
        sleep(0.001) # or `yield()` instead
    end
end

GLMakie.on(kslider.slider.value) do val
    integ.p.K = val
    u_modified!(integ, true)
end

# %%
record_interaction(fig, string(@__FILE__)[1:end-2]*"mp4"; total_time = 15, sleep_time = 2)


using RayTracer

@testset "creating_a_ray" begin
    origin = RayTracer.point(1,2,3)
    direction = RayTracer.vector(4,5,6)
    r = RayTracer.ray(origin, direction)
    @test r.origin ≈ origin
    @test r.direction ≈ direction
end

@testset "computing_point_from_distance" begin
    origin = RayTracer.point(2,3,4)
    direction = RayTracer.vector(1,0,0)
    r = RayTracer.ray(origin, direction)
    @test RayTracer.position(r,0) ≈ RayTracer.point(2,3,4)
    @test RayTracer.position(r,1) ≈ RayTracer.point(3,3,4)
    @test RayTracer.position(r,-1) ≈ RayTracer.point(1,3,4)
    @test RayTracer.position(r,2.5) ≈ RayTracer.point(4.5,3,4)
end

@testset "ray_intersecting_two_points_shpere" begin
    origin = RayTracer.point(0,0,-5)
    direction = RayTracer.vector(0,0,1)
    r = RayTracer.ray(origin, direction)
    s = RayTracer.sphere()
    xs = RayTracer.intersect(s,r)
    @test length(xs) == 2
    @test xs[1].t ≈ 4
    @test xs[2].t ≈ 6
end

@testset "ray_intersecting_tangent_point" begin
    origin = RayTracer.point(0,1,-5)
    direction = RayTracer.vector(0,0,1)
    r = RayTracer.ray(origin, direction)
    s = RayTracer.sphere()
    xs = RayTracer.intersect(s,r)
    @test length(xs) == 2
    @test xs[1].t ≈ 5
    @test xs[2].t ≈ 5
end

@testset "ray_missing_intersection_sphere" begin
    origin = RayTracer.point(0,2,-5)
    direction = RayTracer.vector(0,0,1)
    r = RayTracer.ray(origin, direction)
    s = RayTracer.sphere()
    xs = RayTracer.intersect(s,r)
    @test length(xs) == 0
end

@testset "ray_intersecting_center_shpere" begin
    origin = RayTracer.point(0,0,0)
    direction = RayTracer.vector(0,0,1)
    r = RayTracer.ray(origin, direction)
    s = RayTracer.sphere()
    xs = RayTracer.intersect(s,r)
    @test length(xs) == 2
    @test xs[1].t ≈ -1
    @test xs[2].t ≈ 1
end

@testset "ray_intersecting_front_sphere" begin
    origin = RayTracer.point(0,0,5)
    direction = RayTracer.vector(0,0,1)
    r = RayTracer.ray(origin, direction)
    s = RayTracer.sphere()
    xs = RayTracer.intersect(s,r)
    @test length(xs) == 2
    @test xs[1].t ≈ -6
    @test xs[2].t ≈ -4
end

@testset "intersection_creation" begin
    s = RayTracer.sphere()
    i = RayTracer.intersection(3,s)
    @test i.t ≈ 3
    @test i.object === s
end

@testset "aggregating_intersections" begin
    s = RayTracer.sphere()
    i1 = RayTracer.intersection(1,s)
    i2 = RayTracer.intersection(2,s)
    xs = [i1,i2]
    @test length(xs) == 2
    @test xs[1].t ≈ 1
    @test xs[2].t ≈ 2
end

@testset "intersect_has_obhect_on_intersection" begin
    origin = RayTracer.point(0,0,-5)
    direction = RayTracer.vector(0,0,1)
    r = RayTracer.ray(origin, direction)
    s = RayTracer.sphere()
    xs = RayTracer.intersect(s,r)
    @test length(xs) == 2
    @test xs[1].object === s
    @test xs[2].object === s
end

@testset "hit_when_all_t_are_positive" begin
    s = RayTracer.sphere()
    i1 = RayTracer.intersection(1,s)
    i2 = RayTracer.intersection(2,s)
    xs = [i2,i1]
    i = RayTracer.hit(xs)
    @test i === i1
end

@testset "hit_when_one_t_is_positive" begin
    s = RayTracer.sphere()
    i1 = RayTracer.intersection(-1,s)
    i2 = RayTracer.intersection(1,s)
    xs = [i2,i1]
    i = RayTracer.hit(xs)
    @test i === i2
end

@testset "hit_when_all_t_are_negative" begin
    s = RayTracer.sphere()
    i1 = RayTracer.intersection(-2,s)
    i2 = RayTracer.intersection(-1,s)
    xs = [i2,i1]
    i = RayTracer.hit(xs)
    @test i === nothing
end

@testset "hit_lowest_positive_t" begin
    s = RayTracer.sphere()
    i1 = RayTracer.intersection(5,s)
    i2 = RayTracer.intersection(7,s)
    i3 = RayTracer.intersection(-3,s)
    i4 = RayTracer.intersection(2,s)
    xs = [i1,i2,i3,i4]
    i = RayTracer.hit(xs)
    @test i === i4
end

@testset "translating_ray" begin
    origin = RayTracer.point(1,2,3)
    direction = RayTracer.vector(0,1,0)
    r = RayTracer.ray(origin, direction)
    m = RayTracer.translation(3,4,5)
    r2 = RayTracer.transform(r,m)
    @test r2.origin ≈ RayTracer.point(4,6,8)
    @test r2.direction ≈ RayTracer.vector(0,1,0)
end

@testset "scaling_ray" begin
    origin = RayTracer.point(1,2,3)
    direction = RayTracer.vector(0,1,0)
    r = RayTracer.ray(origin, direction)
    m = RayTracer.scaling(2,3,4)
    r2 = RayTracer.transform(r,m)
    @test r2.origin ≈ RayTracer.point(2,6,12)
    @test r2.direction ≈ RayTracer.vector(0,3,0)
end

@testset "sphere_default_transformation" begin
    s = RayTracer.sphere()
    @test s.transform ≈ Matrix(I,4,4)
end

@testset "changing_sphere_transformation" begin
    s = RayTracer.sphere()
    t = RayTracer.translation(2,3,4)
    RayTracer.set_transform!(s,t)
    @test s.transform ≈ t
end

@testset "intersecting_scaled_sphere_with_ray" begin
    origin = RayTracer.point(0,0,-5)
    direction = RayTracer.vector(0,0,1)
    r = RayTracer.ray(origin, direction)
    s = RayTracer.sphere()
    RayTracer.set_transform!(s,RayTracer.scaling(2,2,2))
    xs = RayTracer.intersect(s,r)
    @test length(xs) == 2
    @test xs[1].t ≈ 3
    @test xs[2].t ≈ 7
end

@testset "intersecting_translated_sphere_with_ray" begin
    origin = RayTracer.point(0,0,-5)
    direction = RayTracer.vector(0,0,1)
    r = RayTracer.ray(origin, direction)
    s = RayTracer.sphere()
    RayTracer.set_transform!(s,RayTracer.translation(5,0,0))
    xs = RayTracer.intersect(s,r)
    @test length(xs) == 0
end


@testset "frac_hour" begin
    @test frac_hour(Minute, 1+1/60) == Hour(1) + Minute(1)
    @test frac_hour(1+1/60) == Hour(1) + Minute(1)
end

@testset "moving_average" begin
    GPPd = @pipe sin.((1:365).*π./365) .+ 0.5 .* rand(365) |> allowmissing(_)
    GPPd[100:120] .= missing
    GPPd_smooth = moving_average(GPPd, 15)
    #plot(1:365, GPPd)
    #plot!(1:365, GPPd_smooth)
    @test length(GPPd_smooth) == length(GPPd)
    @test GPPd_smooth[1] == mean(GPPd[1:8])
    @test GPPd_smooth[8] == mean(GPPd[1:15])
    @test isfinite(GPPd_smooth[100])
    @test ismissing(GPPd_smooth[101])
end

@testset "get_nonoverlapping_periods" begin
    dfo = DataFrame(
        p_start = DateTime(2021) .+ Hour.(SA[0,4,5,10]),
        p_end = DateTime(2021) .+ Hour.(SA[2,5,6,11]),
        )
    dfno = get_nonoverlapping_periods(dfo)
    @test dfno.p_start == DateTime(2021) .+ Hour.(SA[0,4,10])
    @test dfno.p_end == DateTime(2021) .+ Hour.(SA[2,6,11])
    #
    # extend last row
    dfo = DataFrame(
        p_start = DateTime(2021) .+ Hour.(SA[0,4,5]),
        p_end = DateTime(2021) .+ Hour.(SA[2,5,6]),
        )
    dfno = get_nonoverlapping_periods(dfo)
    @test dfno.p_start == DateTime(2021) .+ Hour.(SA[0,4])
    @test dfno.p_end == DateTime(2021) .+ Hour.(SA[2,6])
end

@testset "set_datetime_ydh!" begin
    dfo = DataFrame(year = 2021, doy = repeat(1:3, inner = 48), hour = repeat(0.0:0.5:23.5, outer = 3))
    df = copy(dfo)
    df2 = set_datetime_ydh!(df)
    @test df2 === df
    @test df2.datetime[1] == ZonedDateTime(DateTime(2021), tz"UTC")
    @test all(Dates.year.(df2.datetime) .== dfo.year)
    @test all(Dates.dayofyear.(df2.datetime) .== dfo.doy)
    @test all(Dates.hour.(df2.datetime) .+ Dates.minute.(df2.datetime)/60 .== dfo.hour)
end


function complex_test()
    c1 = 2 + 3im
    c2 = 4 + 5im
    c3 = c1 + c2
    @assert(c3 == (6 + 8im))
    c4 = c1 + 3
    @assert(c4 == (5 + 3im))
    c5 = c1 + 4im
    @assert(c5 == (2 + 7im))
    c6 = c3 - 2.3
    @assert(c6 == (3.7 + 8im))
end

if abspath(PROGRAM_FILE) == @__FILE__
    complex_test()
    println("OK")
end


using Graphs
using NamedGraphs
using DataGraphs

g = NamedDimDataGraph{String,String}(grid((2, 2)); dims=(2, 2))

for v in vertices(g)
  g[v] = "V$v"
end
for e in edges(g)
  g[e] = "E$e"
end

gg = g ⊔ g

@show has_vertex(gg, 1, 1, 1)
@show has_vertex(gg, 2, 1, 1)
@show has_edge(gg, (1, 1, 1) => (1, 1, 2))
@show has_edge(gg, (2, 1, 1) => (2, 1, 2))
@show nv(gg) == 2nv(g)
@show ne(gg) == 2ne(g)

gg = [g; g]

@show has_vertex(gg, 1, 1)
@show has_vertex(gg, 2, 1)
@show has_vertex(gg, 3, 1)
@show has_vertex(gg, 4, 1)
@show has_edge(gg, (1, 1) => (1, 2))
@show has_edge(gg, (3, 1) => (3, 2))
@show nv(gg) == 2nv(g)
@show ne(gg) == 2ne(g)

gg = [g;; g]

@show has_vertex(gg, 1, 1)
@show has_vertex(gg, 1, 2)
@show has_vertex(gg, 1, 3)
@show has_vertex(gg, 1, 4)
@show has_edge(gg, (1, 1) => (1, 2))
@show has_edge(gg, (1, 3) => (1, 4))
@show nv(gg) == 2nv(g)
@show ne(gg) == 2ne(g)


"""
`FDGZProblem` - Initalize the data structure
"""
function FDGZProblem(ϵ, n)
    return FDGZProblem(0, 1, ϵ, n,  DirichletBC(-1), DirichletBC(1))
end

function FDGZProblem(a, b, ϵ, n, bcl::TBCL, bcr::TBCR) where {TBCL<:DirichletBC, TBCR<:DirichletBC}
    Δx = (b-a)/(n+1);
    x = LinRange(a, b, n+2)[2:end-1];
    F = zeros(n);

    return FDGZProblem(a, b, Δx, n, x, ϵ, F, bcl, bcr)
end


using LightGraphs, SimpleWeightedGraphs

function bfs_depth(input_data, start_position, go_condition, stop_condition=(x->false))
  status = zeros(Int32, size(input_data)...)
  queue::Array{NamedTuple{(:position, :depth),Tuple{CartesianIndex, Int32}},1} = []
  if typeof(start_position) <: Tuple
    start_position = CartesianIndex(start_position)
  end
  push!(queue, (position=start_position, depth=0))
  while length(queue) > 0
    (position, depth) = popfirst!(queue)
    #println(position)
    if status[position] != 0
      continue
    end
    status[position] = depth
    if ! stop_condition(input_data[position])
      for movement in [CartesianIndex(-1,0), CartesianIndex(1,0), CartesianIndex(0,-1), CartesianIndex(0,1)]
        new_pos = position + movement
        if go_condition(input_data[new_pos])
          push!(queue, (position=new_pos, depth=depth + 1))
        end
      end
    end
  end
  return status
end

function generate_distance_graph(distances)
  println("distances: $distances")
  tokens = keys(distances)
  my_tokens = Set()
  for token in tokens
    push!(my_tokens, token[1])
    push!(my_tokens, token[2])
  end
  #g = path_graph(maximum(my_tokens))
  mg = SimpleWeightedDiGraph(maximum(my_tokens))
  for token in tokens
    println("Adding: $(token[1]), $(token[2]), $(distances[token])")
    add_edge!(mg, token[1], token[2], distances[token])
  end
  return mg
end

function dijkstra(distances, start_token)
  mg = generate_distance_graph(distances)

  ds = dijkstra_shortest_paths(mg, start_token)
  println(ds.dists)

  return ds.dists

end


#=

  Labeled dice and Building blocks problems in Julia ConstraintSolver.jl

  Labeled dice
  --------------
  From Jim Orlin "Colored letters, labeled dice: a logic puzzle"
  http://jimorlin.wordpress.com/2009/02/17/colored-letters-labeled-dice-a-logic-puzzle/
  """
  My daughter Jenn bough a puzzle book, and showed me a cute puzzle.  There
  are 13 words as follows:  BUOY, CAVE, CELT, FLUB, FORK, HEMP, JUDY,
  JUNK, LIMN, QUIP, SWAG, VISA, WISH.

  There are 24 different letters that appear in the 13 words.  The question
  is:  can one assign the 24 letters to 4 different cubes so that the
  four letters of each word appears on different cubes.  (There is one
  letter from each word on each cube.)  It might be fun for you to try
  it.  I'll give a small hint at the end of this post. The puzzle was
  created by Humphrey Dudley.
  """

  Building blocks
  ---------------
  From http://brownbuffalo.sourceforge.net/BuildingBlocksClues.html
  """
  Each of four alphabet blocks has a single letter of the alphabet on each
  of its six sides. In all, the four blocks contain every letter but
  Q and Z. By arranging the blocks in various ways, you can spell all of
  the words listed below. Can you figure out how the letters are arranged
  on the four blocks?

  BAKE ONYX ECHO OVAL
  GIRD SMUG JUMP TORN
  LUCK VINY LUSH WRAP
  """

  Note: This is a somewhat generalized model for solving both
        Building blocks and Labeled Dice problems.

  Model created by Hakan Kjellerstrand, hakank@gmail.com
  See also my Julia page: http://www.hakank.org/julia/

=#
using ConstraintSolver, JuMP
using Cbc, GLPK, Ipopt
const CS = ConstraintSolver
include("constraints_utils.jl")

function block_problem(problem,print_solutions=true,all_solutions=true)

    cbc_optimizer = optimizer_with_attributes(Cbc.Optimizer, "logLevel" => 0)
    glpk_optimizer = optimizer_with_attributes(GLPK.Optimizer)
    ipopt_optimizer = optimizer_with_attributes(Ipopt.Optimizer)

    model = Model(optimizer_with_attributes(CS.Optimizer,   "all_solutions"=> all_solutions,
                                                            # "all_optimal_solutions"=>true,
                                                            "logging"=>[],

                                                            # "traverse_strategy"=>:BFS,
                                                            "traverse_strategy"=>:DFS,
                                                            # "traverse_strategy"=>:DBFS,

                                                            "branch_split"=>:Smallest,
                                                            # "branch_split"=>:Biggest,
                                                            # "branch_split"=>:InHalf,

                                                            # https://wikunia.github.io/ConstraintSolver.jl/stable/options/#branch_strategy-(:Auto)
                                                            "branch_strategy" => :IMPS, # default
                                                            # "branch_strategy" => :ABS, # Activity Based Search
                                                            "activity.decay" => 0.999, # default 0.999
                                                            "activity.max_probes" => 3, # default, 10
                                                            "activity.max_confidence_deviation" => 3, # default 20

                                                            # "simplify"=>false,
                                                            # "simplify"=>true, # default

                                                            "time_limit"=>6,

                                                            # "backtrack" => false, # default true
                                                            # "backtrack_sorting" => false, # default true

                                                            # "lp_optimizer" => cbc_optimizer,
                                                            # "lp_optimizer" => glpk_optimizer,
                                                            # "lp_optimizer" => ipopt_optimizer,
                                        ))
    # a,b,c,d,e,f,g,h,i,j,k,l,m,n,o,p,q,r,s,t,u,v,w,x,y,z = 1:26
    println("problem ", problem[:name])
    num_cubes = problem[:num_cubes]
    num_sides = problem[:num_sides]
    letters = split(problem[:letters],"")
    num_chars = length(letters)

    # letters = problem[:letters]
    words = problem[:words]
    num_words = length(words)

    d = Dict(zip(letters,1:length(letters)))
    d_rev = [k=>v for (v,k) in d] # reversed dict

    cube_len = num_cubes*num_sides
    @variable(model, 1 <= cube[1:cube_len] <= num_cubes, Int)

    # each letters in a word must be on a different die
    for i in 1:num_words
      word = [d[w] for w in split(words[i,:][1],"")]
      this_die = @variable(model, [1:length(word)],CS.Integers(1:num_cubes))
      for i in 1:length(word)
          my_element(model,word[i],cube,this_die[i])
      end
      @constraint(model, this_die in CS.AllDifferent())
    end

    # there must be exactly NumSides (6) letters of each die (num_cubes)
    global_cardinality_count(model, cube, [num_sides for _ in 1:num_cubes])

    # simple symmetry breaking: place first letter on cube 1
    @constraint(model, cube[1] == 1)

    # Solve the problem
    optimize!(model)

    status = JuMP.termination_status(model)
    # println("status:$status")
    if status == MOI.OPTIMAL
        num_sols = MOI.get(model, MOI.ResultCount())
        println("num_sols:$num_sols\n")
        if print_solutions
            for sol in 1:num_sols
                println("\nsolution #$sol")
                cube_val = convert.(Integer,JuMP.value.(cube; result=sol))
                println("cube:$cube_val")
                println("words:")
                for word in words
                    print(word,": ")
                    for w in split(word,"")
                        dd = d[w]
                        print("$w: $(cube_val[dd]) ")
                    end
                    println()
                end
                println("\ndice:")
                dice = [ [] for _ in 1:num_cubes]
                for (c,i) in collect(zip(letters,cube_val))
                    push!(dice[i],c)
                end
                join.(sort.(unique.(dice))," ").|>println
            end
        end
    else
        println("status:$status")
    end

    return status
end

function block_problems(problem)
    all_block_problems = Dict(
      :labeled_dice =>  Dict(
           :name => "labeled_dice",
           # number of cubes
           :num_cubes => 4,
           # number of sides of of a cube
           :num_sides => 6,
           # the letters to use
           :letters => "abcdefghijklmnopqrstuvwy",
           # the words to place
           :words => ["buoy", "cave", "celt", "flub",
                      "fork", "hemp", "judy","junk",
                      "limn", "quip", "swag", "visa",
                      "wish"]
    ),

    :building_blocks => Dict(
           :name => "building_blocks",
           :num_cubes => 4,
           :num_sides => 6,
           # Note: 'f' is not in any word but it needed to get 24 chars
           :letters => "abcdefghijklmnoprstuvwxy",
           :words => ["bake", "onyx", "echo", "oval",
                      "gird", "smug", "jump", "torn",
                      "luck", "viny", "lush", "wrap"]
        )

   )

   return all_block_problems[problem]

end

@time block_problem(block_problems(:labeled_dice),true,true)
println()
@time block_problem(block_problems(:building_blocks),true,true)


using Knet
using CUDArt
using CUSPARSE
using Base.Test

a = rand(Float32,5,3)
b = convert(SparseMatrixCSC{Float32}, sprand(1000,3,.005))
b[10,:] = 1
c = a*b'
sc = sparse(c)
da = CudaArray(a)
db = CudaSparseMatrixCSC(b)
cc = similar(sc)
rand!(cc.nzval)
dc1 = CudaSparseMatrixCSR(cc)
Knet.A_mul_Bt_slow!(dc1, da, db)
sc1 = Knet.to_host2(dc1)
c1 = full(sc1)
@test @show c == c1
dc2 = CudaSparseMatrixCSR(cc)
Knet.A_mul_Bt!(dc2, da, db)
sc2 = Knet.to_host2(dc2)
c2 = full(sc2)
@test @show c == c2

T=1000
I=1000
J=20
K=10000

foo1(c,a,b)=(for t=1:T; Knet.A_mul_Bt_slow!(c,a,b); end)
foo2(c,a,b)=(for t=1:T; Knet.A_mul_Bt!(c,a,b); end)

a = rand(Float32, I, J)
b = convert(SparseMatrixCSC{Float32}, sprand(K,J,1/K))
c = a*b'
sc = sparse(c)
da = CudaArray(a)
db = CudaSparseMatrixCSC(b)
cc = similar(sc)
rand!(cc.nzval)
dc1 = CudaSparseMatrixCSR(cc)
dc2 = CudaSparseMatrixCSR(cc)
@date foo1(dc1, da, db)
@date foo1(dc1, da, db)
@date foo2(dc2, da, db)
@date foo2(dc2, da, db)

sc1 = Knet.to_host2(dc1)
c1 = full(sc1)
@test @show c == c1

sc2 = Knet.to_host2(dc2)
c2 = full(sc2)
@test @show c == c2


function update_histogram!(histogram, bin)
    if bin > length(histogram)
        push!(histogram, 1)
    else
        histogram[bin] += 1
    end

    return histogram
end

###########################################################################
# version as this would usually be written in Turing, with the slight
# modification that we need to associate G₀ with z (to calculate marginals of
# y for "unobserved clusters")
@model function imm(y, α, ::Type{T}=Array{Float64, 1}) where {T}
    N = length(y)

    K = 0
    nk = Vector{Int}()
    z = Vector{Int}(undef, N)

    G₀ = Normal()

    for n = 1:N
        z[n] ~ ChineseRestaurantProcess(DP(α, G₀), nk)
        nk = update_histogram!(nk, z[n])
        K = max(K, z[n])
    end

    μ = T(undef, K)
    for k = 1:K
        μ[k] ~ G₀
    end

    for n = 1:N
        y[n] ~ Normal(μ[z[n]], 1.0)
    end
end

function test_imm()
    model_imm = imm(data_neal[5:7], α_neal)
    graph_imm = trackdependencies(model_imm)

    # we leave out the μs, because there might be 1--3 of them
    @testdependencies(model_imm, z[1], z[2], z[3], y[1], y[2], y[3])
    # @test_nothrow sample(model_imm, Gibbs(AutoConditional(:z), HMC(0.01, 10, :μ)), 2)

    # for comparison:
    # sample(model_imm, Gibbs(MH(:z => filldist(Categorical(9), 9)), HMC(0.01, 10, :μ)), 2)

    # Analytic tests
    z = [graph_imm[20].value, graph_imm[37].value, graph_imm[53].value]
    μ = [v.value for v in values(graph_imm.statements)
         if v isa AutoGibbs.Assumption && DynamicPPL.subsumes(@varname(μ), v.vn)]
    y = graph_imm[2].value
    K = graph_imm[56].value
    N = graph_imm[7].value

    CRP(h) = ChineseRestaurantProcess(DirichletProcess(α_neal), h)

    # exploiting exchangeability:
    # 𝓅(zₙ | z₋ₙ, μ, yₙ) ∝ 𝓅(zₙ | z₋ₙ) 𝓅(yₙ | zₙ, μ)
    function cond_collapsed(n, k)
        # nk is the histogram of z without z[k]
        nk = zeros(Int, N)
        for i in eachindex(z)
            i != k && (nk[z[i]] += 1)
        end

        l = logpdf(CRP(nk), z[k])

        if k <= K
            # 𝓅(yₙ | zₙ = k, μ)
            l += logpdf(Normal(μ[k]), y[n])
        else # k == K + 1
            # 𝓅(yₙ | zₙ = K + 1, μ) = ∫ 𝓅(yₙ | m) 𝓅(m) dm = pdf(G₀, yₙ)
            l += logpdf(Normal(), y[n])
        end

        return exp(l)
    end

    p_z1_coll = cond_collapsed.(1, 1:1)
    p_z2_coll = cond_collapsed.(2, 1:2)
    p_z3_coll = cond_collapsed.(3, 1:3)
    analytic_conditionals = [@varname(z[1]) => Categorical(p_z1_coll ./ sum(p_z1_coll)),
                             @varname(z[2]) => Categorical(p_z2_coll ./ sum(p_z2_coll)),
                             @varname(z[3]) => Categorical(p_z3_coll ./ sum(p_z3_coll))]
    @info "normal IMM analytic conditionals (collapsed)" analytic_conditionals

    θ = AutoGibbs.sampled_values(graph_imm)

    local calculated_conditionals
    @test_nothrow calculated_conditionals = conditionals(graph_imm, @varname(z))
    @info "normal IMM calculated conditionals" Dict(vn => cond(θ) for (vn, cond) in calculated_conditionals)

    for (vn, analytic_conditional) in analytic_conditionals
        # @show vn => probs(calculated_conditionals[vn]), probs(analytic_conditional)
        @test issimilar(calculated_conditionals[vn](θ), analytic_conditional)
    end
end

###########################################################################
@model function imm_oneshot(y, α, ::Type{T}=Array{Float64, 1}) where {T}
    N = length(y)
    G₀ = Normal()
    z ~ AutoGibbs.CRP(N, α, G₀)
    K = maximum(z)

    μ = T(undef, K)
    for k = 1:K
        μ[k] ~ G₀
    end

    for n = 1:N
        y[n] ~ Normal(μ[z[n]], 1.0)
    end
end

function test_imm_oneshot()
    model_imm_oneshot = imm_oneshot(data_neal[5:7], α_neal)
    graph_imm_oneshot = trackdependencies(model_imm_oneshot)

    # we leave out the μs, because there might be 1--3 of them
    @testdependencies(model_imm_oneshot, z, y[1], y[2], y[3])
    # @test_nothrow sample(model_imm_oneshot, Gibbs(AutoConditional(:z), HMC(0.01, 10, :μ)), 2)

    # for comparison:
    # sample(model_imm_oneshot, Gibbs(MH(:z => filldist(Categorical(9), 9)), HMC(0.01, 10, :μ)), 2)

    # Analytic tests
    z = graph_imm_oneshot[10].value
    μ = [v.value for v in values(graph_imm_oneshot.statements)
         if v isa AutoGibbs.Assumption && DynamicPPL.subsumes(@varname(μ), v.vn)]
    y = graph_imm_oneshot[2].value
    K = graph_imm_oneshot[11].value
    N = graph_imm_oneshot[7].value

    CRP(h) = ChineseRestaurantProcess(DirichletProcess(α_neal), h)

    # exploiting exchangeability:
    # 𝓅(zₙ | z₋ₙ, μ, yₙ) ∝ 𝓅(zₙ | z₋ₙ) 𝓅(yₙ | zₙ, μ)
    function cond_collapsed(n, k)
        # nk is the histogram of z without z[k]
        nk = zeros(Int, N)
        for i in eachindex(z)
            i != k && (nk[z[i]] += 1)
        end

        l = logpdf(CRP(nk), z[k])

        if k <= K
            # 𝓅(yₙ | zₙ = k, μ)
            l += logpdf(Normal(μ[k]), y[n])
        else # k == K + 1
            # 𝓅(yₙ | zₙ = K + 1, μ) = ∫ 𝓅(yₙ | m) 𝓅(m) dm = pdf(G₀, yₙ)
            l += logpdf(Normal(), y[n])
        end

        return exp(l)
    end

    p_z1_coll = cond_collapsed.(1, 1:1)
    p_z2_coll = cond_collapsed.(2, 1:2)
    p_z3_coll = cond_collapsed.(3, 1:3)
    analytic_conditionals = [@varname(z[1]) => Categorical(p_z1_coll ./ sum(p_z1_coll)),
                             @varname(z[2]) => Categorical(p_z2_coll ./ sum(p_z2_coll)),
                             @varname(z[3]) => Categorical(p_z3_coll ./ sum(p_z3_coll))]
    @info "normal IMM analytic conditionals (collapsed)" analytic_conditionals

    # θ = AutoGibbs.sampled_values(graph_imm_oneshot)

    # local calculated_conditionals
    # @test_nothrow calculated_conditionals = conditionals(graph_imm_oneshot, @varname(z))
    # @info "one-shot IMM calculated conditionals" Dict(vn => cond(θ) for (vn, cond) in calculated_conditionals)

    # for (vn, analytic_conditional) in analytic_conditionals
    #     # @show vn => probs(calculated_conditionals[vn]), probs(analytic_conditional)
    #     @test issimilar(calculated_conditionals[vn](θ), analytic_conditional)
    # end
end

###########################################################################
@model function imm_manual(y, α, ::Type{T}=Array{Float64, 1}) where {T}
    N = length(y)

    K = 0
    nk = Vector{Int}()
    z = Vector{Int}(undef, N)

    G₀ = Normal()

    for n = 1:N
        probs = [nk; α] ./ (n + α - 1)
        z[n] ~ Categorical(probs)
        nk = update_histogram!(nk, z[n])
        K = max(K, z[n])
    end

    μ = T(undef, K)
    for k = 1:K
        μ[k] ~ G₀
    end

    for n = 1:N
        y[n] ~ Normal(μ[z[n]], 1.0)
    end
end

function test_imm_manual()
    model_imm_manual = imm_manual(data_neal[5:7], α_neal)
    graph_imm_manual = trackdependencies(model_imm_manual)

    @testdependencies(model_imm_manual, z[1], z[2], z[3], y[1], y[2], y[3])
    # @test_nothrow sample(model_imm_manual, Gibbs(AutoConditional(:z), HMC(0.01, 10, :μ)), 2)

    calculated_conditionals = conditionals(graph_imm_manual, @varname(z))

    # # log 𝓅(z₁, ..., zₙ)
    # function logpdf_crp(z)
    #     N = length(z)
    #     l = 0.0
    #     nk = Int[]
    #     K = 0
    #     log_α = log(α)

    #     for n = 1:N
    #         if z[n] <= K
    #             l += log(z[n]) - log(n + α - 1)
    #             nk[z[n]] += 1
    #         else
    #             l += log_α - log(n + α - 1)
    #             push!(nk, 1)
    #             K += 1
    #         end

    #     end
    #     return l
    # end

    # # 𝓅(zₙ = k| z₁, ..., zₙ₋₁, μ, yₙ) ∝ (∏_{i = z ≥ n} 𝓅(zᵢ | z₁,...,zᵢ)) 𝓅(yₙ | zₙ, μ)
    # function cond(n, k)
    #     # 𝓅(zₙ = k | z₁, ..., zₙ₋₁)
    #     l = logpdf_crp([j == n ? k : z[j] for j = 1:n])

    #     if k <= K
    #         # 𝓅(yₙ | zₙ = k, μ)
    #         l += logpdf(Normal(μ[k]), y[n])
    #     else
    #         # 𝓅(yₙ | zₙ = K + 1, μ) = ∫ 𝓅(yₙ | m) 𝓅(m) dm = pdf(G₀, yₙ)
    #         l += logpdf(Normal(), y[n])
    #     end

    #     return exp(l)
    # end

    # p_z1 = cond.(1, 1:1)
    # p_z2 = cond.(2, 1:2)
    # p_z3 = cond.(3, 1:3)
    # analytic_conditionals = [@varname(z[1]) => Categorical(p_z1 ./ sum(p_z1)),
    #                          @varname(z[2]) => Categorical(p_z2 ./ sum(p_z2)),
    #                          @varname(z[3]) => Categorical(p_z3 ./ sum(p_z3))]
    # @info "IMM analytic" analytic_conditionals

    θ = AutoGibbs.sampled_values(graph_imm_manual)

    local calculated_conditionals
    @test_nothrow calculated_conditionals = conditionals(graph_imm_manual, @varname(z))
    @info "manual IMM calculated conditionals" Dict(vn => cond(θ) for (vn, cond) in calculated_conditionals)
end

##########################################################################
#########################################################################
### TEST TOGGLES
# test_imm()
# test_imm_oneshot()
# test_imm_manual()


module SimpleFilters

using TinyModia
using DifferentialEquations
using ModiaPlot

SimpleModel = Model(
    T = 0.2,
    x = Var(init=0.5),
    equation = :[T * der(x) + x = 2]
)

# @showModel(SimpleModel)

model = @instantiateModel(SimpleModel)
simulate!(model, Tsit5(), stopTime = 5, requiredFinalStates = [1.9999996962023168])
plot(model, ["x"])

LowPassFilter = Model(
    T = 0.2,
    u = input,
    y = output | Var(:x),
    x = Var(init=0.0u"V"),
    equations = :[T * der(x) + x = u],
)

# @showModel(LowPassFilter)

HighPassFilter = LowPassFilter | Model(
        y = Var(:(-x + u)),
    )

# @showModel(HighPassFilter)

# setLogMerge(true)

LowAndHighPassFilter = LowPassFilter | Model(
    y = nothing,
    low = output | Var(:x),
    high = output | Var(:(-x + u)),
)

setLogMerge(false)

# @showModel(LowAndHighPassFilter)

TestLowAndHighPassFilter = LowAndHighPassFilter | Model(
    u = :(sin( (time+1u"s")*u"1/s/s" * time)*u"V"),
    x = Var(init=0.2u"V")
)

# @showModel(TestLowAndHighPassFilter)

model = @instantiateModel(TestLowAndHighPassFilter, log=false, logCode=false)
simulate!(model, Tsit5(), stopTime = 5, requiredFinalStates = [-0.22633046061014014])
plot(model, ["low", "high"])

TwoFilters = Model(
    high = HighPassFilter,
    low = LowPassFilter,
)

# @showModel(TwoFilters)

BandPassFilter = Model(
    u = input,
    y = output,
    high = HighPassFilter | Map(T=0.5, x=Var(init=0.1u"V")),
    low = LowPassFilter | Map(x=Var(init=0.2u"V")),
    equations = :[
        high.u = u,
        low.u = high.y,
        y = low.y]
)

# @showModel(BandPassFilter)

TestBandPassFilter = BandPassFilter | Model(
    u = :(sin( (0.05*time+1u"s")*u"1/s/s" * time)*u"V"),
)

# @showModel(TestBandPassFilter)

setLogMerge(false)

model = @instantiateModel(TestBandPassFilter, logDetails=false)
simulate!(model, Tsit5(), stopTime = 50, requiredFinalStates = [-0.25968254453053435, -0.6065869606784539])
plot(model, ["u", "y"])

end


immutable AliasTable <: AbstractCategoricalSampler
    accept::Vector{Float64}
    alias::Vector{Int}
    isampler::RandIntSampler
end

function AliasTable(probs)

    n = length(probs)
    accept = float64(probs*n)

    alias = Array(Int,n)
    larges = Array(Int,0)
    smalls = Array(Int,0)

    for i = 1:n
        acci = accept[i]
        if acci > 1.0
            push!(larges,i)
        elseif acci < 1.0
            push!(smalls,i)
        end
    end
    while !isempty(larges) && !isempty(smalls)
        s = pop!(smalls)
        l = pop!(larges)
        alias[s] = l
        accept[l] = (accept[l] - 1.0) + accept[s]
        if accept[l] > 1
            push!(larges,l)
        else
            push!(smalls,l)
        end
    end

    # this loop should be redundant, except for rounding
    for s = smalls
        accept[s] = 1.0
    end

    AliasTable(accept, alias, RandIntSampler(n))
end

function rand(a::AliasTable)
    i = rand(a.isampler)
    u = rand()
    u < a.accept[i] ? i : a.alias[i]
end

Base.show(io::IO, a::AliasTable) = @printf io "AliasTable with %d entries" length(a.accept)


include("Parsers.jl")
include("Utils.jl")
using DataInterpolations

abstract type AbstractDiff  end

struct AnalyticalDeriv <: AbstractDiff end
struct FiniteDiff <: AbstractDiff end
struct TotalVariationalDerivativative <: AbstractDiff end

function differentiate(sol,o::TotalVariationalDerivativative)
    # interp = LagrangeInterpolation(munge(sol.u),sol.t)
    interp = QuadraticInterpolation(munge(sol.u),sol.t)
    y_reconstruct = interp.(sol.t)
    v = [DataInterpolations.derivative(interp,k) for k in sol.t]
    w = munge(v)
    if length(sol.u[1]) == 3
        fig = plot(w[1,:],w[2,:],w[3,:])
    elseif length(sol.u[1]) == 2
        fig = plot(w[1,:],w[2,:])
    end
    display(fig)
    return v
end

function differentiate(sol,o::FiniteDiff)
    q = sol.u[2:end] .- sol.u[1:end-1]
    q = q./(sol.t[2] - sol.t[1])
    w = munge(q)
    if length(sol.u[1]) == 3
        fig = plot(w[1,:],w[2,:],w[3,:])
    elseif length(sol.u[1]) == 2
        fig = plot(w[1,:],w[2,:])
    end
    display(fig)
    return w
end

function differentiate(sol::Matrix,t,o::TotalVariationalDerivativative)
    interp = QuadraticInterpolation(sol,t)
    y_reconstruct = interp.(t)
    v = [DataInterpolations.derivative(interp,k) for k in t]
    w = munge(v)
    if size(sol,1) == 3
        fig = plot(w[1,:],w[2,:],w[3,:])
    elseif size(sol,1) == 2
        fig = plot(w[1,:],w[2,:])
    end
    display(fig)
    return v
end

function differentiate(sol::Matrix,t,o::FiniteDiff)
    q = sol[:,2:end] .- sol[:,1:end-1]
    q = q./(t[2] - t[1])
    w = q
    if size(sol,1) == 3
        fig = plot(w[1,:],w[2,:],w[3,:])
    elseif size(sol,1) == 2
        fig = plot(w[1,:],w[2,:])
    end
    display(fig)
    return w
end

#--------------- The analytical derivative is computed for testing---------------------#

function differentiate(sol,ds::LorenzSystem,o::AnalyticalDeriv)
    function lorenz(u) # Modify when changing parameters
        u1 = 10*(u[2] - u[1])
        u2 = u[1]*(28 - u[3]) - u[2]
        u3 = u[1]*u[2] - 8/3 * u[3]
        return [u1 u2 u3]
    end
    deriv = []
    for i = 1:length(sol.u)
        append!(deriv,lorenz(sol.u[i]))
    end
    deriv = reshape(deriv,(3,length(sol.u)))
    fig = plot(deriv[1,:],deriv[2,:],deriv[3,:],title="Lorenz - Derivative")
    display(fig)
    return deriv

end

function differentiate(sol,ds::LotkaVolterra,o::AnalyticalDeriv)
    function lk(u) # Modify when changing parameters
        u1 = 0.7*u[1] - 0.3*u[1]*u[2]
        u2 = -0.3*u[2] + 0.4*u[1]*u[2]
        return [u1 u2]
    end
    deriv = []
    for i = 1:length(sol.u)
        append!(deriv,lk(sol.u[i]))
    end
    deriv = reshape(deriv,(2,length(sol.u)))
    fig = plot(deriv[1,:],deriv[2,:],title="Lotka Volterra - Derivative")
    display(fig)
    return deriv
end


"""
    test_constvel(; kwargs...)

Advects a gaussian concentration c0(x, y, t) with a constant velocity flow
u(x, y) = uvel and v(x, y) = vvel and compares the final state with
cfinal = c0(x-uvel*tfinal, y-vvel*tfinal)
"""
function test_constvel(stepper, dt, nsteps)

  nx, Lx = 128, 2π
  uvel, vvel = 0.2, 0.1
  u(x, y) = uvel
  v(x, y) = vvel

  prob = TracerAdvDiff.Problem(; nx=nx, Lx=Lx, kap=0.0, u=u, v=v, dt=dt, stepper=stepper, steadyflow=true)
  sol, cl, v, p, g = prob.sol, prob.clock, prob.vars, prob.params, prob.grid

  x, y = gridpoints(g)

  σ = 0.1
  c0func(x, y) = 0.1*exp.(-(x.^2+y.^2)/(2σ^2))

  c0 = c0func.(x, y)
  tfinal = nsteps*dt
  cfinal = @. c0func(x - uvel*tfinal, y - vvel*tfinal)

  TracerAdvDiff.set_c!(prob, c0)

  stepforward!(prob, nsteps)
  TracerAdvDiff.updatevars!(prob)

  isapprox(cfinal, v.c, rtol=g.nx*g.ny*nsteps*1e-12)
end

"""
    test_timedependenttvel(; kwargs...)

Advects a gaussian concentration c0(x, y, t) with a time-varying velocity flow
u(x, y, t) = uvel and v(x, y, t) = vvel*sign(-t+tfinal/2) and compares the final
state with cfinal = c0(x-uvel*tfinal, y)
"""
function test_timedependentvel(stepper, dt, tfinal)

  nx, Lx = 128, 2π
  nsteps = round(Int, tfinal/dt)

  if !isapprox(tfinal, nsteps*dt, rtol=1e-12)
      error("tfinal is not multiple of dt")
  end

  uvel, vvel = 0.2, 0.1
  u(x, y, t) = uvel
  v(x, y, t) = t <= tfinal/2 ? vvel : -vvel

  prob = TracerAdvDiff.Problem(; nx=nx, Lx=Lx, kap=0.0, u=u, v=v, dt=dt, stepper=stepper)
  sol, cl, v, p, g = prob.sol, prob.clock, prob.vars, prob.params, prob.grid
  x, y = gridpoints(g)

  σ = 0.1
  c0func(x, y) = @. 0.1*exp(-(x^2+y^2)/(2σ^2))

  c0 = @. c0func(x, y)
  tfinal = nsteps*dt
  cfinal = @. c0func(x - uvel*tfinal, y)

  TracerAdvDiff.set_c!(prob, c0func)

  stepforward!(prob, nsteps)
  TracerAdvDiff.updatevars!(prob)

  isapprox(cfinal, v.c, rtol=g.nx*g.ny*nsteps*1e-12)
end

"""
    test_diffusion(; kwargs...)

Diffuses a gaussian concentration c0(x, y, t) and compares the final state with
the analytic solution of the heat equation, cfinal
"""
function test_diffusion(stepper, dt, tfinal; steadyflow = true)

  nx  = 128
  Lx  = 2π
  kap = 0.01
  nsteps = round(Int, tfinal/dt)

  if !isapprox(tfinal, nsteps*dt, rtol=1e-12)
      error("tfinal is not multiple of dt")
  end

  grid = TwoDGrid(nx, Lx)
  prob = TracerAdvDiff.Problem(; steadyflow=steadyflow, grid=grid, nx=nx,
    Lx=Lx, kap=kap, dt=dt, stepper=stepper)
  sol, cl, v, p, g = prob.sol, prob.clock, prob.vars, prob.params, prob.grid
  x, y = gridpoints(g)

  c0ampl, σ = 0.1, 0.1
  c0func(x, y) = @. c0ampl*exp(-(x^2+y^2)/(2σ^2))

  c0 = @. c0func.(x, y)
  tfinal = nsteps*dt
  σt = sqrt(2*kap*tfinal + σ^2)
  cfinal = @. c0ampl*σ^2/σt^2 * exp(-(x^2+y^2)/(2*σt^2))

  TracerAdvDiff.set_c!(prob, c0)

  stepforward!(prob, nsteps)
  TracerAdvDiff.updatevars!(prob)

  isapprox(cfinal, v.c, rtol=g.nx*g.ny*nsteps*1e-12)
end

"""
    test_hyperdiffusion(; kwargs...)

Diffuses a gaussian concentration c0(x, y, t) using hyperdiffusivity and
compares the final state with the analytic solution of the heat equation, cfinal
"""
function test_hyperdiffusion(stepper, dt, tfinal; steadyflow = true)

    nx  = 128
    Lx  = 2π
    kap = 0.0   # no diffusivity
    eta = kap   # no diffusivity
   kaph = 0.01  # hyperdiffusivity coeff
  nkaph = 1     # nkaph=1 converts hyperdiffusivity to plain diffusivity
                # so we can compare with the analytic solution of heat equation

  nsteps = round(Int, tfinal/dt)

  if !isapprox(tfinal, nsteps*dt, rtol=1e-12)
      error("tfinal is not multiple of dt")
  end

   g = TwoDGrid(nx, Lx)
  x, y = gridpoints(g)

  u, v = zero(x), zero(x) #0*x, 0*x

  vs = TracerAdvDiff.Vars(g)
  pr = TracerAdvDiff.ConstDiffSteadyFlowParams(eta, kap, kaph, nkaph, u, v)
  eq = TracerAdvDiff.Equation(pr, g)
  prob = FourierFlows.Problem(eq, stepper, dt, g, vs, pr)

  c0ampl, σ = 0.1, 0.1
  c0func(x, y) = @. c0ampl*exp(-(x^2+y^2)/(2σ^2))

  c0 = @. c0func(x, y)
  tfinal = nsteps*dt
  σt = sqrt(2*kaph*tfinal + σ^2)
  cfinal = @. c0ampl*σ^2/σt^2 * exp(-(x^2+y^2)/(2*σt^2))

  TracerAdvDiff.set_c!(prob, c0)

  stepforward!(prob, nsteps)
  TracerAdvDiff.updatevars!(prob)

  isapprox(cfinal, vs.c, rtol=g.nx*g.ny*nsteps*1e-12)
end


# Undergraduate Student: Arturo Burgos
# Professor: João Rodrigo Andrade
# Federal University of Uberlândia - UFU, Fluid Mechanics Laboratory - MFLab, Block 5P, Uberlândia, MG, Brazil

using Printf # Just the same as import a module
using Statistics
using CPUTime

# A = rand(1:50,3,3)
# B = rand(1:50,3,3)

A = [1 2; 3 4]
println("\n\n\nThe A Matrix is: \n\n",A)

B = [1 2; 3 4]
println("\n\nThe B Matrix is: \n\n",B)

R = zeros(2,2) # Here I create a zero matrix just to get the multiply of A and B

D = A*B # This is the vectorized way for this problem, like MATLAB
C = A.*B # This multiplies point by point

println("\n\nThe D Matrix is: \n\n",D)
println("\n\nThe C Matrix is: \n\n",C)

# println("\n\n",D[:,1])

function multivec(mtx1,mtx2,mtx3)

    mtx3 = mtx1*mtx2

    return mtx3 # Returns a value to the variable
end

R = multivec(A,B,R);
println("\n\nThe R Matrix is: \n\n",R)
println("\n\n")

@time @CPUtime multivec(A,B,R)


using LinearAlgebra
import Flux
import Flux: σ

gate(h, n) = (1:h) .+ h*(n-1)
gate(x::AbstractVector, h, n) = @view x[gate(h,n)]
gate(x::AbstractMatrix, h, n) = @view x[gate(h,n),:]

struct CachedRecur{C, M <: AbstractVecOrMat}
   cell  :: C
   init  :: M
   state :: M
end

Flux.@functor CachedRecur

Flux.trainable(m::CachedRecur) = (m.cell, m.init)

struct CachedLSTMCell{M <: AbstractVecOrMat, V <: AbstractVecOrMat}
   Wi :: M
   Wh :: M
   b  :: V
   cache :: NamedTuple{(:g, :Wix), NTuple{2,M}}
end

Flux.@functor CachedLSTMCell

Flux.trainable(m::CachedLSTMCell) = (m.Wi, m.Wh, m.b)

function CachedLSTMCell(in::Integer, out::Integer; batchsize::Integer,
                        init = Flux.glorot_uniform)
   b = init(4out)
   b[gate(out, 2)] .= 1
   O = Flux.zeros(4out, batchsize)
   cache = (g = O, Wix = copy(O))
   return CachedLSTMCell(init(4out, in), init(4out, out), b, cache)
end

function (m::CachedLSTMCell)((h, c), x)
   Wi, Wh, b, o, batch = m.Wi, m.Wh, m.b, size(h, 1), axes(x, 2)
   @views g, Wix = m.cache.g[:,batch], m.cache.Wix[:,batch]
   g .= mul!(Wix, Wi, x) .+ mul!(g, Wh, h) .+ b
   @views ifo, cell = g[1:(3o),:], g[(3o+1):end,:]
   ifo .= σ.(ifo)
   cell .= tanh.(cell)
   input, forget, output = gate(ifo, o, 1), gate(ifo, o, 2), gate(ifo, o, 3)
   @. c = forget * c + input * cell
   @. h = output * tanh(c)
   return (h, c), h
end


using Statistics
"""
    CrossEntropyLoss(dims=1)
    (l::CrossEntropyLoss)(scores, answers; average=true)
Calculates negative log likelihood error on your predicted scores.
`answers` should be integers corresponding to correct class indices.
If an answer is 0, loss from that answer will not be included.
This is the masking feature when you are working with unequal length sequences.

if dims==1,
    First dimension is assumed to be predicted logits.
* size(scores) = D,[BATCH,T1,T2,...]
* size(answers)= [BATCH,T1,T2,...]

elseif dims==2,
    Last dimension is assumed to be predicted logits.
* size(scores) = [BATCH,T1,T2,...],D
* size(answers)= [BATCH,T1,T2,...]
"""
struct CrossEntropyLoss <: Loss
    dims::Integer
end
CrossEntropyLoss(;dims=1) = CrossEntropyLoss(dims)
@inline (l::CrossEntropyLoss)(y,answers::Array{<:Integer}; average=true) = nllmask(y, answers; dims=l.dims, average=average)

"""
    BCELoss(average=true)
    (l::BCELoss)(scores, answers)
    Computes binary cross entropy given scores(predicted values) and answer labels. answer values should be {0,1}, then it returns negative of
    mean|sum(answers * log(p) + (1-answers)*log(1-p)) where p is equal to 1/(1 + exp.(scores)). see `LogisticLoss`.
"""
struct BCELoss <: Loss end
@inline (l::BCELoss)(y,answers::Array{<:Integer})=bce(y,answers)

"""
    LogisticLoss(average=true)
    (l::LogisticLoss)(scores, answers)
    Computes logistic loss given scores(predicted values) and answer labels. answer values should be {-1,1}, then it returns mean|sum(log(1 +
    exp(-answers*scores))). See also `BCELoss`.
"""
struct LogisticLoss <: Loss end
@inline (l::LogisticLoss)(y,answers::Array{<:Integer}) = logistic(y,answers)

"""
    SigmoidCrossEntropyLoss()
    (l::SigmoidCrossEntropyLoss)(scores, labels)

Measures the error in discrete classification tasks in which each class is independent and not mutually exclusive.
`labels` should be same size with scores where existing classes pointed by one others zero.

if dims==1, First dimension is assumed to be predicted logits.

elseif dims==2, Last dimension is assumed to be predicted logits.
"""
struct SigmoidCrossEntropyLoss <: Loss
    dims::Int
end
SigmoidCrossEntropyLoss(;dims=1) = SigmoidCrossEntropyLoss(dims)
function (l::SigmoidCrossEntropyLoss)(x, z; average=true)
    y = sum(relu.(x) .- x .* z - log.(sigm.(abs.(x))), dims=:)
    average ? y ./ (prod(size(x))/prod(size(x)[l.dims])) : y
end

####
#### Utils
####
@inline function nllmask(y,a::AbstractArray{<:Integer}; dims=1, average=true)
    indices = findindices(y, a, dims=dims)
    lp = logp(y,dims=dims)[indices]
    average ? -mean(lp) : -sum(lp)
end

function findindices(y,a::AbstractArray{<:Integer}; dims=1)
    n       = length(a)
    nonmask = a .> 0
    indices = Vector{Int}(undef,sum(nonmask))
    if dims == 1                   # instances in first dimension
        y1 = size(y,1)
        y2 = div(length(y),y1)
        if n != y2; throw(DimensionMismatch()); end
        k = 1
        @inbounds for (j,v) in enumerate(nonmask)
            !v && continue
            indices[k] = (j-1)*y1 + a[j]
            k += 1
        end
    elseif dims == 2               # instances in last dimension
        y2 = size(y,ndims(y))
        y1 = div(length(y),y2)
        if n != y1; throw(DimensionMismatch()); end
        k = 1
        @inbounds for (j,v) in enumerate(nonmask)
            !v && continue
            indices[k] = (a[j]-1)*y1 + j
            k += 1
        end
    else
        error("findindices only supports dims = 1 or 2")
    end
    return indices
end