### A Pluto.jl notebook ###
# v0.14.4

using Markdown
using InteractiveUtils

# ╔═╡ c7e41465-41ad-4f11-baec-bfaae9edbf59
begin
	using PlutoUI
	PlutoUI.TableOfContents(aside=true)
end

# ╔═╡ d0233ece-1f69-11eb-3da1-43b8ddc9fbe1
begin
	# Necessary packages
	using LinearAlgebra
	using Random
	using Statistics
	using Clustering
	using Plots
end

# ╔═╡ 00d72ad7-ea80-4954-8c75-590d4e9ccdc4
md"""
# K-means Algorithm

Data clustering is one of the main mathematical applications variety of algorithms have been developed to tackle the problem. K-means is one of the basic algorithms for data clustering.

__Prerequisites__

The reader should be familiar with basic linear algebra.

__Competences__

The reader should be able to recognise applications where K-means algorithm can be efficiently used and use it.

__Credits.__ The notebook was initially derived from M.Sc. Thesis of Ivančica Mirošević.
"""

# ╔═╡ 235e1901-eed3-4934-9214-ffdd82cf21d6
md"""
## Definitions

__Data clustering problem__ is the following: partition the given set of $m$ objects of the same type into $k$ subsets according to some criterion. Additional request may be to find the optimal $k$.

__K-means clustering problem__ is the following: partition the set  $X=\{x_{1},x_{2},\cdots ,x_{m}\}$ , where $x_{i}\in\mathbb{R}^{n}$, into $k$ _clusters_ $\pi=\{C_{1},C_{2},...,C_{k}\}$ such that

$$
J(\pi)=\sum_{i=1}^{k}\sum_{x\in
C_{i}}\| x-c_{i}\|_{2}^{2} \to \min$$

over all possible partitions. Here $c_{i}=\displaystyle\frac{1}{|C_{i}|}\sum_{x\in C_{i}} x$ is the mean of points in $C_i$ and $|C_i|$ is the cardinality of $C_i$.
"""

# ╔═╡ 26640b90-1f69-11eb-029e-4dc3a825f639
md"
## K-means clustering algorithm

1. __Initialization__: Choose initial set of $k$ means $\{c_1,\ldots,c_k\}$ (for example, by choosing randomly $k$ points from $X$).
2. __Assignment step__: Assign each point $x$ to one nearest mean $c_i$.
3. __Update step__: Compute the new means.
4. __Iteration__: Repeat Steps 2 and 3 until the assignment no longer changes.
"

# ╔═╡ 48914840-1f69-11eb-30a9-3943767bc261
md"
## First Variation clustering algorithm

A __first variation__ of a partition $\pi=\{C_1,\ldots,C_k\}$ is
a partition $\pi^{\prime}=\{C_{1}^{\prime},\cdots ,C_{k}^{\prime }\}$
obtained by moving a single point $x$ from a cluster  $C_{i}$ to a cluster $C_{j}$. Notice that $\pi$ is a first variation of itself.

A __next partition__ of the partition $\pi$ is a partition
$\mathop{\mathrm{next}}(\pi)=\mathop{\mathrm{arg min}}\limits_{\pi^{\prime}} J(\pi^{\prime})$.

We have the following algorithm:

1. Choose initial partition $\pi$.
2. Compute $\mathop{\mathrm{next}}(\pi)$
3. If $J(\mathop{\mathrm{next}}(\pi))<J(\pi)$, set $\pi=\mathop{\mathrm{next}}(\pi)$ and go to Step 2
4. Stop.
"

# ╔═╡ 2b76953c-edf2-41cf-9d75-7efb2d6a58f1
md"""
## Facts

1. The k-means clustering problem is NP-hard.

2. In the k-means algorithm, $J(\pi)$ decreases in every iteration.

3. K-means algorithm can converge to a local minimum.

4. Each iteration of the k-means algorithm requires $O(mnk)$ operations.

4. K-means algorithm is implemented in the function `kmeans()` in the package [Clustering.jl](https://github.com/JuliaStats/Clustering.jl).

5.  $J(\pi)=\mathop{\mathrm{trace}}(S_W)$, where

$$
S_{W}=\sum\limits_{i=1}^k\sum\limits_{x\in C_{i}}
(x-c_i)(x-c_{i})^{T}
=\sum_{i=1}^k\frac{1}{2|C_{i}|}\sum_{x\in C_{i}}\sum_{y \in C_{i}}
(x-y)(x-y)^{T}.$$

Let $c$ denote the mean of $X$. Then $S_W=S_T-S_B$, where

$$
\begin{aligned}
S_{T}&=\sum_{x\in X}(x-c)(x-c)^{T} =
\frac{1}{2m}\sum_{i=1}^m\sum_{j=1}^m
(x_{i}-x_{j})(x_{i}-x_{j})^{T}, \\
S_{B}&=\sum_{i=1}^k|C_{i}|(c_{i}-c)(c_{i}-c)^{T} =
\frac{1}{2m}\sum_{i=1}^k\sum_{j=1}^k|C_{i}||C_{j}|
(c_{i}-c_{j})(c_{i}-c_{j})^{T}.
\end{aligned}$$

6. In order to try to avoid convergence to local minima, the k-means algorithm can be enhanced with first variation by adding the following steps:
    1. Compute $\mathop{\mathrm{next}}(\pi)$.
    2. If $J(\mathop{\mathrm{next}}(\pi))<J(\pi)$, set $\pi=\mathop{\mathrm{next}}(\pi)$ and go to Step 2.

"""

# ╔═╡ 86d55833-ee5a-4f73-b0bf-e8b6e6d65617
function myKmeans(X::Vector{T}, k::Int) where T
    # X is Array of Arrays
    m,n=length(X),length(X[1])
    C=Vector{Int}(undef,m)
    # Choose random k means among X
    c=X[randperm(m)[1:k]]
    # This is just to start the while loop
    cnew=copy(c)
    cnew[1]=cnew[1].+(1.0,1.0)
    # Loop
    iterations=0
    while cnew!=c
        iterations+=1
        cnew=copy(c)
        # Assignment
        for i=1:m
            C[i]=findmin([norm(X[i].-c[j]) for j=1:k])[2]
        end
        # Update
        for j=1:k
          c[j]=(mean([x[1] for x in X[C.==j]]),mean([x[2] for x in X[C.==j]]))
        end
    end
    C,c,iterations
end

# ╔═╡ e15fa4cf-930a-4065-b839-c2ec944afcfa
md"""

## Examples

### Random clusters

We generate $k$ random clusters around points with integer coordinates.
"""

# ╔═╡ bacdc83d-1acd-46df-8d06-3e0279978a33
begin
	# Generate points as Tuple()
	k=5
	Random.seed!(1235)
	centers= [Tuple(rand(-5:5,2)) for i=1:k]
	# Number of points in cluster
	sizes=rand(10:50,k)
	csizes=cumsum(sizes)
	# X is array of arrays
	X=Vector{Tuple{Float64,Float64}}(undef,sum(sizes))
	X[1:csizes[1]]=[centers[1].+Tuple((rand(2).-0.5)/2) for i=1:sizes[1]]
	for j=2:k
		X[csizes[j-1]+1:csizes[j]]=[centers[j].+Tuple((rand(2).-0.5)/2) for i=1:sizes[j]]
	end
	centers, sizes
end

# ╔═╡ 99aba610-1f75-11eb-1b5e-2f2ee211ed7f
X

# ╔═╡ fceab134-0921-459b-8d91-952a44abb07b
begin
	# Plot
	scatter(X,label="Points",title="Points")
	scatter!(centers,markershape = :hexagon, ms = 6,label="Centers")
end

# ╔═╡ f559b870-608a-426e-a23b-b66ef6d4b47f
# Plot the solution
function plotKmeansresult(C::Vector,c::Vector,X::Vector)
    scatter()
    # Clusters
    for j=1:k
        scatter!(X[findall(C.==j)],label="Cluster $j")
    end
    # Means
    scatter!(c,markershape=:hexagon,ms=6,color=:red,label="Centers")
	plot!(title="Computed clusters")
end

# ╔═╡ ba84fe97-4dd0-47d3-b980-f5e5a8de6fa3
X

# ╔═╡ 310a4957-cc58-424b-845c-3ebaf7db77f3
md"""
__What happens?__

We see that the algorithm, although simple, for this example
converges to a local minimum.

Let us try the function `kmeans()` from the package `Clustering.jl`.
The inputs to `kmeans()` are:

* a matrix whose columns are points,
* number of clusters we are looking for, and,
* optionally, the method to compute initial means.

If we choose `init=:rand`, the results are similar. If we choose
`init=:kmpp`, which is the default, the results are better, but convergence to a local minimum is still possible.

__Run the clustering several times!__

```
seeding_algorithm(s::Symbol) =
    s == :rand ? RandSeedAlg() :
    s == :kmpp ? KmppAlg() :
    s == :kmcen ? KmCentralityAlg() :
    error("Unknown seeding algorithm $s")
```
"""

# ╔═╡ f9af8a75-48fd-4e82-b43e-61867bead6f9
methods(kmeans)

# ╔═╡ 7e869699-8eb1-4e38-905d-a9448d037841
begin
	Xₘ=transpose([[x[1] for x in X] [x[2] for x in X]])
	output=kmeans(Matrix(Xₘ),k,init=:kmpp)
end

# ╔═╡ 4f8152ee-ce86-4df6-81d2-be2df33df980
fieldnames(KmeansResult)

# ╔═╡ 95fb2e66-e811-47f4-a772-8270d0c5d83b
output.centers

# ╔═╡ 26b7bac0-1f72-11eb-16fc-c1d74920c462
output.assignments

# ╔═╡ 7bb0e9ce-72dd-4623-ab3d-a6d18fb9dd15
# We need to modify the plotting function
function plotKmeansresult(out::KmeansResult,X::AbstractArray)
    k=size(out.centers,2)
    scatter(aspect_ratio=1)
    # Clusters
    for j=1:k
        scatter!(X[1,findall(out.assignments.==j)], X[2,findall(out.assignments.==j)], label="Cluster $j")
    end
    # Means
    scatter!(out.centers[1,:], out.centers[2,:], markershape=:hexagon,ms=6,color=:red,label="Centers")
end

# ╔═╡ 8e5b89ce-e6b5-4473-bd98-a2285331418c
begin
	C,c,iterations=myKmeans(X,k)
	plotKmeansresult(C,c,X)
end

# ╔═╡ 406bb1b2-a8bb-4cd2-95c3-3975f99e3f32
begin
	out=kmeans(Xₘ,k,init=:kmpp)
	plotKmeansresult(out,Xₘ)
end

# ╔═╡ 2445e75a-77f9-4c27-b400-44d8b49fc03b
md"""
### Concentric rings

The k-means algorithm works well if clusters can be separated by hyperplanes. In this example it is not the case.
"""

# ╔═╡ 21bdc615-76e8-46cc-ac86-310d28f6cd25
begin
	# Number of rings, try also k=3
	k₁=2
	# Center
	Random.seed!(5361)
	center=[rand(-5:5);rand(-5:5)]
	# Radii
	radii=randperm(10)[1:k₁]
	# Number of points in circles
	sizes₁=rand(1000:2000,k₁)
	center,radii,sizes₁
end

# ╔═╡ 9a4060b7-ba5a-401f-90f2-b9780ed93ddb
begin
	# Generate points
	X₁=Array{Float64}(undef,2,sum(sizes₁))
	csizes₁=cumsum(sizes₁)
	# Random angles
	ϕ=2*π*rand(sum(sizes₁))
	for i=1:csizes₁[1]
		X₁[:,i]=center+radii[1]*[cos(ϕ[i]);sin(ϕ[i])] + (rand(2).-0.5)/50
	end
	for j=2:k₁
		for i=csizes₁[j-1]+1:csizes₁[j]
			X₁[:,i]=center+radii[j]*[cos(ϕ[i]);sin(ϕ[i])] + (rand(2).-0.5)/50
		end
	end
	scatter(X₁[1,:],X₁[2,:],title="Concentric Rings", aspect_ratio=1,label="Points")
end

# ╔═╡ 0be24dd5-f3b0-4daa-aad8-5dbcdfca9313
begin
	out₁=kmeans(X₁,k₁,init=:rand)
	plotKmeansresult(out₁,X₁)
end

# ╔═╡ 1637ce15-f55a-4d66-8a00-45a555bf5d61

# ╔═╡ Cell order:
# ╟─c7e41465-41ad-4f11-baec-bfaae9edbf59
# ╟─00d72ad7-ea80-4954-8c75-590d4e9ccdc4
# ╟─235e1901-eed3-4934-9214-ffdd82cf21d6
# ╟─26640b90-1f69-11eb-029e-4dc3a825f639
# ╟─48914840-1f69-11eb-30a9-3943767bc261
# ╟─2b76953c-edf2-41cf-9d75-7efb2d6a58f1
# ╠═d0233ece-1f69-11eb-3da1-43b8ddc9fbe1
# ╠═86d55833-ee5a-4f73-b0bf-e8b6e6d65617
# ╟─e15fa4cf-930a-4065-b839-c2ec944afcfa
# ╠═bacdc83d-1acd-46df-8d06-3e0279978a33
# ╠═99aba610-1f75-11eb-1b5e-2f2ee211ed7f
# ╠═fceab134-0921-459b-8d91-952a44abb07b
# ╠═f559b870-608a-426e-a23b-b66ef6d4b47f
# ╠═8e5b89ce-e6b5-4473-bd98-a2285331418c
# ╠═ba84fe97-4dd0-47d3-b980-f5e5a8de6fa3
# ╟─310a4957-cc58-424b-845c-3ebaf7db77f3
# ╠═f9af8a75-48fd-4e82-b43e-61867bead6f9
# ╠═7e869699-8eb1-4e38-905d-a9448d037841
# ╠═4f8152ee-ce86-4df6-81d2-be2df33df980
# ╠═95fb2e66-e811-47f4-a772-8270d0c5d83b
# ╠═26b7bac0-1f72-11eb-16fc-c1d74920c462
# ╠═7bb0e9ce-72dd-4623-ab3d-a6d18fb9dd15
# ╠═406bb1b2-a8bb-4cd2-95c3-3975f99e3f32
# ╟─2445e75a-77f9-4c27-b400-44d8b49fc03b
# ╠═21bdc615-76e8-46cc-ac86-310d28f6cd25
# ╠═9a4060b7-ba5a-401f-90f2-b9780ed93ddb
# ╠═0be24dd5-f3b0-4daa-aad8-5dbcdfca9313
# ╠═1637ce15-f55a-4d66-8a00-45a555bf5d61

src/Module C - Applications/L09 K-means Algorithm.jl


# ------------------------------------------------------------------------------------------
# # Arrays and Loops
#
# We'll cover:
# 1. Array literals
# 2. Concatenation
# 3. For loops
# 4. Comprehensions
# 5. Element types
# 6. Dequeues
# 7. The `bang!` convention
# 8. Variable names vs. copies
#
# # Arrays
#
# Julia has highly efficient multidimensional arrays, both constructed and indexed with
# square brackets.
#
# Syntax: <br>
# ```julia
# [item1, item2, ...]
# ```
# ------------------------------------------------------------------------------------------

squares = [1, 4, 9, 15, 25, 36, 49, 64]

squares[1]

squares[1:3]

squares[end]

squares[4] = 16

squares

typeof(squares)

# ------------------------------------------------------------------------------------------
# ## Concatenation
#
# If, instead of commas, you just use spaces, then the values are concatenated horizontally.
# ------------------------------------------------------------------------------------------

cubes = [1, 8, 27, 64, 125, 216, 343, 512]

powers = [1:8 squares cubes]

powers[4, 2]

powers[:, 3]

powers[7, :]

typeof(powers)

# ------------------------------------------------------------------------------------------
# Semicolon separators perform vertical concatenation:
# ------------------------------------------------------------------------------------------

[squares; cubes]

# ------------------------------------------------------------------------------------------
# Whereas commas would simply create an array of arrays:
# ------------------------------------------------------------------------------------------

nested_powers = [[1,2,3,4,5,6,7,8], squares, cubes]

nested_powers[2]

# ------------------------------------------------------------------------------------------
# Horizontal and vertical concatenation can be used together to as a simple syntax for
# matrix literals:
# ------------------------------------------------------------------------------------------

[1 3 5; 2 4 6]

# ------------------------------------------------------------------------------------------
# # Loops
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# Of course, we could construct this programmatically with a for-loop. The syntax for a
# `for` loop is
#
# ```julia
# for *var* in *loop iterable*
#     *loop body*
# end
# ```
#
#
# ------------------------------------------------------------------------------------------

A = fill(0, (8, 3)) # Allocate an 8x3 matrix to store the values into
for pow in 1:3
    for value in 1:8
        A[value, pow] = value ^ pow
    end
end
A

A == powers

# ------------------------------------------------------------------------------------------
# ## Array Comprehensions
# ------------------------------------------------------------------------------------------

squares = [value^2 for value in 1:8]

cubes = [value^3 for value in 1:8]

powers = [value^pow for value in 1:8, pow in 1:3]

# ------------------------------------------------------------------------------------------
# # The element type
#
# Note that every time an array prints out, it is displaying its element type and
# dimensionality, for example `Array{Int64, 2}`. This describes what it can store — and thus
# what it can return upon indexing.
# ------------------------------------------------------------------------------------------

typeof(powers)

typeof(powers[1, 1])

# ------------------------------------------------------------------------------------------
# Further, the array will try to convert any new values assigned into it to its element
# type:
# ------------------------------------------------------------------------------------------

powers[1, 1] = 1.6

powers[1, 1] = -5.0 # This can be losslessly converted to an integer

powers

# ------------------------------------------------------------------------------------------
# Arrays that have an exact and concrete element type are generally significantly faster, so
# Julia will try to find an amenable element type for you in its literal construction
# syntax:
# ------------------------------------------------------------------------------------------

fortytwosarray = [42, 42.0, 4.20e1, 4.20f1, 84//2, 0x2a]

for x in fortytwosarray
    show(x)
    println("\tisa $(typeof(x))")
end

# ------------------------------------------------------------------------------------------
# The `Any` array can be helpful for disabling these behaviors and allowing all kinds of
# different objects:
# ------------------------------------------------------------------------------------------

anyfortytwos = Any[42, 42.0, 4.20e1, 4.20f1, 84//2, 0x2a]

anyfortytwos[1] = "FORTY TWO"
anyfortytwos

# ------------------------------------------------------------------------------------------
# # Vectors as dequeues
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# One-dimensional arrays can be appended to and have items removed from them:
# ------------------------------------------------------------------------------------------

fib = [1, 1, 2, 3, 5, 8, 13]

push!(fib, 21)

push!(fib, sum(fib[end-1:end]))

pop!(fib)

pushfirst!(fib, 0)

popfirst!(fib)

# ------------------------------------------------------------------------------------------
# ## Aside: why so shouty!?
#
# Why are there exclamations in all the above function calls? `push!(fib, ...)`,
# `pop!(fib)`, etc.?
# ------------------------------------------------------------------------------------------

push!

# ------------------------------------------------------------------------------------------
# This is entirely a convention. It's a signal to the caller that one of the argument is
# going to be _mutated_. It's perhaps easiest to demonstrate with an example:
# ------------------------------------------------------------------------------------------

A = rand(0:10, 10)

sort(A)

A

sort!(A)

A # That changed A!

# ------------------------------------------------------------------------------------------
# ## Aside: names vs. copies
#
# Remember that variables are just names we give our objects. So watch what happens if we
# give the same object two different names:
# ------------------------------------------------------------------------------------------

fibonacci = [1, 1, 2, 3, 5, 8, 13]

some_numbers = fibonacci
some_numbers[1] = 404
some_numbers

fibonacci

fibonacci[1] = 1
some_numbers = copy(fibonacci)
some_numbers[2] = 404
fibonacci

.nbexports/introductory-tutorials/intro-to-julia/long-version/020 Arrays and Loops.jl


## Here is the code from unit 1

# This treats everything as a dense matrix :(
A = [0 1 1 1 1 0 1 0 0 0
     1 0 0 0 0 0 0 0 0 0
     0 0 0 0 1 1 1 0 0 0
     1 1 0 0 1 1 0 1 0 0
     1 1 1 1 0 0 0 0 1 0
     0 0 1 0 0 0 1 0 1 1
     0 0 0 0 0 1 0 0 0 1
     0 0 0 0 0 0 0 0 1 0
     0 0 0 0 0 0 0 1 0 0
     0 0 0 0 0 0 0 0 0 0]
d = sum(A,dims=2)
##
using Printf
for i=1:10
  for j=1:10
    if i==j
      @printf("1 & ")
    elseif A[j,i] != 0
      @printf("-\\alpha/%i & ", d[j])
    else
      @printf("0 & ")
    end
  end
  println("\\\\")
end
## Build the I - alpha D^+ A'
alpha = 0.85
M = zeros(size(A)...)
for i=1:10
  for j=1:10
    if i==j
      M[i,j] = 1
    elseif A[j,i] != 0
      M[i,j] = -alpha*A[j,i]/d[j]
    else
      M[i,j] = 0
    end
  end
end
M[:,end] .= -alpha/10
M[end,end] += 1
M[:,1]
##
b = ones(size(A,1))/size(A,1)*(1-alpha)
x = M\b

##
using DelimitedFiles
using SparseArrays
ei,ej = copy.(eachcol(Int.(readdlm("wiki-simple.edges"))))
xy = readdlm("wiki-simple.xy")
pages = readlines("wiki-simple.nodes")
n = size(xy,1)
A = sparse(ei,ej,1,n,n)

## Simple PageRank
function simplepagerank(A,α,v)
  @assert(0 ≤ α < 1, "needs probably α")
  @assert(all(vi -> vi ≥ 0, v), "needs non-negative v")
  v = v ./ sum(v) # we can normalize for them.
  d = vec(sum(A,dims=2)) # compute the degrees
  x = copy(v) # start of with v
  nsteps = 2*ceil(Int,log(eps(1.0))/log(α)) # upper bound on steps
  for i=1:nsteps
    x = α*(A'*(x./d)) .+ (1-α).*v
  end
  return x/sum(x) # renormalize to probability
end
pr = simplepagerank(A,0.85,ones(n)./n)


using LinearAlgebra, StatsBase, TimerOutputs, DataFrames, StatsPlots, BenchmarkTools

function structured_P(L::Int; p::Float64 = 0.45, r::Float64 = 0.01)::Matrix{Float64}
    q = 1 - p - r
    P = diagm(fill(r,L)) + diagm(-1=>fill(q,L-1)) + diagm(1 => fill(p,L-1))
    P[1,1] = 1-p
    P[L,L] = 1-q
    return P
end

structured_π(L::Int; p::Float64 = 0.45, r::Float64 = 0.01)::Vector{Float64} = begin
    q = 1 - p - r
    [(p/q)^i  for i in 1:L] * (q-p) / p / (1-(p/q)^L) #Explicit expression (birth death)
end;

# Method 1: Linear Algebra
function method_1(input_matrix, L)
    # Transpose both sides, add a row to the bottom to represent the additional requirement
    return [I - input_matrix'; ones(1, L)] \ [zeros(L); 1]
end

# Method 2: Take a high matrix power
function method_2(input_matrix, L)
    matrix_power = 10^5 # arbitrary
    # Return the first column - the 'j' doesn't matter
    return (input_matrix^(matrix_power))[1, :]
end

# Method 3: Eigen trickery
function method_3(input_matrix, L)
    eigens = eigen(input_matrix')
    correct_vector = eigens.vectors[:, findfirst(isapprox(1), eigens.values)]
    return correct_vector * UniformScaling(1 / sum(correct_vector))
end

# Method 4: Random sample
function method_4(input_matrix, L)
    repeats = 10^5 # arbitrary
    states = 1:L
    current_value = rand(states)
    chain = zeros(L) # we don't need the order, just the counts
    for i in 1:repeats
        # Take a weighted sample based on the current state
        current_value = sample(states, Weights(view(input_matrix, current_value, :)))
        chain[current_value] += 1
    end
    # Apply the 1/n
    return map(x->x/repeats, chain)
end

# The L's given by the task
L_given = [2:5..., 10:10:50..., 100:100:500..., 1000]

# Get the benchmarks warmed up and ready to go
suite = BenchmarkGroup()
for method in [method_1, method_2, method_3, method_4]
    suite[String(Symbol(method))] = BenchmarkGroup()
end

# Prepare data structures
norms_df = DataFrame(a=L_given)
approx_methods_for_df = zip('b':'e', [method_1, method_2, method_3, method_4])

for approx_method ∈ approx_methods_for_df
    # Benchmark each L for each function, add it to suite
    for L in L_given
        suite[String(Symbol(approx_method[2]))][L] = @benchmarkable $(approx_method[2])(structured_P($L), $L)
    end

    # Put the error into the df for plotting
    norms_df[!, string(approx_method[1])] = map(L->norm(approx_method[2](structured_P(L), L) - structured_π(L)), L_given)
end

# Run our thing
results = run(suite, seconds = 2, samples = 5)

# Plot the DataFrames

# Make sure none of the values break the graph
inds = (norms_df.c .> 0) .& (norms_df.d .> 0) .& (norms_df.b .> 0) .& (norms_df.e .> 0)
# Plot the absolute errors of each method
@show @df norms_df plot(
    :a[inds], [:b[inds] :c[inds] :d[inds] :e[inds]],
    xaxis=:log10,
    xtickfont = font(5, "Courier"),
    xlabel = "L",
    ylabel = "Absolute Error",
    yscale=:log10,
    label = ["Method 1 (Linear Algebra)" "Method 2 (Matrix Power)" "Method 3 (Eigenvector)" "Method 4 (Random Seq)"],
    title = "Demonstration of Approximation Methods"
)

# Prepare the benchmarks for plotting (load their median times into a dataframe)
benchmarks_df = DataFrame(L=L_given)
for method in keys(results)
    benchmarks_df[!, String(Symbol(method))] = map(L->median(results[method][L]).time, L_given)
end

# Plot the benchmarks
@show @df benchmarks_df plot(
    :L, [:method_1 :method_2 :method_3 :method_4],
    xaxis=:log10,
    xtickfont = font(5, "Courier"),
    xlabel = "L",
    ylabel = "Time (in one unit or another)",
    yscale=:log10,
    label = ["Method 1 (Linear Algebra)" "Method 2 (Matrix Power)" "Method 3 (Eigenvector)" "Method 4 (Random Seq)"],
    title = "Demonstration of Approximation Methods"
)


"""
	normcdf(x)

Probablity that Z ≤ x, where Z is a standard normal random variable.
"""
normcdf(x::Real) = erfc(-x / √two(x)) / 2

"""
	normcdf(a, b)

Probablity that a ≤ Z ≤ b, where Z is a standard normal random variable.
WARNING: Silently returns a negative value if a > b.
"""
normcdf(a::Real, b::Real) = erf(a / √two(a), b / √two(b)) / 2

"""
	normcdfinv(x)

Inverse of normcdf.
"""
normcdfinv(x::Real) = -√two(x) * erfcinv(2x)

"""
	lognormcdf(a, b)

Computes log(normcdf(a, b)), but retaining accuracy.
"""
lognormcdf(a::Real, b::Real) = logerf(a / √two(a), b / √two(b)) - log(two(a + b))

"""
	tnmean(a)

Mean of the standard normal distribution,
truncated to the interval (a, +∞).
"""
tnmean(a::Real) = sqrt(two(a)/π) / erfcx(a/√two(a))

"""
	tnvar(a)

Variance of the standard normal distribution,
truncated to the interval (a, +∞).
WARNING: Fails for very very large values of `a`.
"""
function tnvar(a::Real)
	μ = tnmean(a)
	return one(μ) - (μ - a) * μ
end

"""
	tnstd(a)

Standard deviation of the standard normal distribution,
truncated to the interval (a, +∞).
"""
tnstd(a::Real) = √tnvar(a)

"""
	mills(x::Real)

Mills ratio of the standard normal distribution.
Defined as (1 - cdf(x)) / pdf(x).
"""
mills(x::Real) = erfcx(x / √two(x)) * √(π/two(x))


# Returns an array of positive prime numbers less than or equal to lim
function sieve(lim :: Int)
    if lim < 2 return [] end
    limi :: Int = (lim - 1) ÷ 2 # calculate the required array size
    isprime :: Array{Bool} = trues(limi)
    llimi :: Int = (isqrt(lim) - 1) ÷ 2 # and calculate maximum root prime index
    result :: Array{Int} = [2]  #Initial array
    for i in 1:limi
        if isprime[i]
            p = i + i + 1 # 2i + 1
            if i <= llimi
                for j = (p*p-1)>>>1:p:limi # quick shift/divide in case LLVM doesn't optimize divide by 2 away
                    isprime[j] = false
                end
            end
            push!(result, p)
        end
    end
    return result
end


@doc raw"""
    ChiSquared()

The Chi-Squared base function is given by:

```math
f(\mathbf{x}, \mathbf{y}) = \sum_i \frac{(x_i - y_i)^2}{x_i + y_i}
```
"""
struct ChiSquared <: PreMetric end

@inline function base_aggregate(::ChiSquared, s::T, x::T, y::T) where {T}
    x == y == zero(T) ? s : s + (x-y)^2/(x+y)
end


# coding: utf-8

# In[1]:

#another fractal script via recursion
#this script heavily involves analytic geometry
#a good material to help you rewind high school geometry
# https://youthconway.com/handouts/analyticgeo.pdf
using Plots

# In[2]:

#simple solution to get coefficients of the equation
function get_line_params(x1,y1,x2,y2)

    slope=(y1-y2)/(x1-x2)
    intercept=y1-slope*x1

    return slope,intercept

end

#compute euclidean distance
function euclidean_distance(point1,point2)
    return √((point1[1]-point2[1])^2+(point1[2]-point2[2])^2)
end

#standard solution to quadratic equation
function solve_quadratic_equation(A,B,C)
    x1=(-B+√(B^2-4*A*C))/(2*A)
    x2=(-B-√(B^2-4*A*C))/(2*A)
    return [x1,x2]
end

#analytic geometry to compute target datapoints
function get_datapoint(pivot,measure,length,direction="inner")

    #for undefined slope
    if pivot[1]==measure[1]
        y1=pivot[2]+length
        y2=pivot[2]-length
        x1=pivot[1]
        x2=pivot[1]

    #for general cases
    else

        #get line equation
        slope,intercept=get_line_params(pivot[1],pivot[2],
                                   measure[1],measure[2],)

        #solve quadratic equation
        A=1
        B=-2*pivot[1]
        C=pivot[1]^2-length^2/(slope^2+1)
        x1,x2=solve_quadratic_equation(A,B,C)

        #get y from line equation
        y1=slope*x1+intercept
        y2=slope*x2+intercept

    end

    if direction=="inner"

        #take the one between pivot and measure points
        if euclidean_distance((x1,y1),measure)<euclidean_distance((x2,y2),measure)
            datapoint=(x1,y1)
        else
            datapoint=(x2,y2)
        end

    else

        #take the one farther away from measure points
        if euclidean_distance((x1,y1),measure)>euclidean_distance((x2,y2),measure)
            datapoint=(x1,y1)
        else
            datapoint=(x2,y2)
        end

    end

    return datapoint

end

#recursively compute the coordinates of koch curve data points
#to effectively connect the data points,the best choice is to use turtle
#it would be too difficult to connect the dots via analytic geometry
function koch_snowflake(base1,base2,base3,n,arr)

    #base case
    if n==0
        return

    else

        #find mid point
        #midpoint between base1 and base2 has to satisfy two conditions
        #it has to be on the same line as base1 and base2
        #assume this line follows y=kx+b
        #the midpoint is (x,kx+b)
        #base1 is (α,kα+b),base2 is (δ,kδ+b)
        #the euclidean distance between midpoint and base1 should be
        #half of the euclidean distance between base1 and base2
        #(x-α)^2+(kx+b-kα-b)^2=((α-δ)^2+(kα+b-kδ-b)^2)/4
        #apart from x,everything else in the equation is constant
        #this forms a simple quadratic solution to get two roots
        #one root would be between base1 and base2 which yields midpoint
        #and the other would be farther away from base2
        #this function solves the equation via (-B+(B^2-4*A*C)^0.5)/(2*A)
        #alternatively,you can use scipy.optimize.root
        #the caveat is it does not offer both roots
        #a wrong initial guess could take you to the wrong root
        midpoint=get_datapoint(base1,base2,euclidean_distance(base1,base2)/2)

        #compute the top point of a triangle
        #the computation is similar to midpoint
        #the euclidean distance between triangle_top and midpoint should be
        #one third of the distance between base3 and midpoint
        triangle_top=get_datapoint(midpoint,base3,
                                   euclidean_distance(midpoint,base3)/3,
                                   "outer")

        #two segment points divide a line into three equal parts
        #the computation is almost the same as midpoint
        #the euclidean distance between segment1 and base1
        #should be one third of the euclidean distance between base2 and base1
        segment1=get_datapoint(base1,base2,euclidean_distance(base1,base2)/3)
        segment2=get_datapoint(base2,base1,euclidean_distance(base1,base2)/3)

        #compute the nearest segment point of the neighboring line
        segment_side_1=get_datapoint(base1,base3,euclidean_distance(base1,base3)/3)
        segment_side_2=get_datapoint(base2,base3,euclidean_distance(base2,base3)/3)

        append!(arr,[segment1,segment2,triangle_top])

        #recursion
        koch_snowflake(base1,segment1,segment_side_1,n-1,arr)
        koch_snowflake(segment1,triangle_top,segment2,n-1,arr)
        koch_snowflake(triangle_top,segment2,segment1,n-1,arr)
        koch_snowflake(segment2,base2,segment_side_2,n-1,arr)
    end
    arr
end

# In[3]:

#set data points
point1=(0.0,0.0)
point2=(3.0,0.0)
point3=(1.5,1.5*√(3))

n=4

# In[4]:

#collect coordinates
coordinates=[point1,point2,point3]
append!(coordinates,koch_snowflake(point1,point2,point3,n,[]))
append!(coordinates,koch_snowflake(point3,point1,point2,n,[]))
append!(coordinates,koch_snowflake(point2,point3,point1,n,[]));

# In[5]:

#viz
#visually u can tell some of the data points are miscalculated
#purely caused by floating point arithmetic
gr(size=(250,250))
scatter([i[1] for i in coordinates],
        [i[2] for i in coordinates],markersize=2,
        legend=false,showaxis=false,ticks=false,grid=false)


# This file is a part of Julia. License is MIT: https://julialang.org/license

function laplace_iter_devec(u, dx2, dy2, Niter, N)
    uout = copy(u)
    for iter = 1:Niter
        for i = 2:N-1
            for j = 2:N-1
                uout[i,j] = ( (u[i-1,j]+u[i+1,j])*dy2 + (u[i,j-1]+u[i,j+1])*dx2 ) * (1./(2*(dx2+dy2)))
            end
        end
        u, uout = uout, u
    end
    return u
end

function laplace_devec()
    N = 150
    u = zeros(N, N)
    u[1,:] = 1
    Niter = 2^10
    dx2 = dy2 = 0.1*0.1
    u = laplace_iter_devec(u, dx2, dy2, Niter, N)
end

function laplace_iter_vec(u, dx2, dy2, Niter, N)
    for i = 1:Niter
        u[2:N-1, 2:N-1] = ((u[1:N-2, 2:N-1] + u[3:N, 2:N-1])*dy2 + (u[2:N-1,1:N-2] + u[2:N-1, 3:N])*dx2) * (1./ (2*(dx2+dy2)))
    end
    return u
end

function laplace_vec()
    N = 150
    u = zeros(N,N)
    u[1,:] = 1
    Niter = 2^10
    dx2 = dy2 = 0.1*0.1
    u = laplace_iter_vec(u, dx2, dy2, Niter, N)
end

function laplace_iter_vec_sub(u, dx2, dy2, Niter, N)
    for i = 1:Niter
        u[2:N-1, 2:N-1] = ((view(u, 1:N-2, 2:N-1) + view(u,3:N, 2:N-1))*dy2 + (view(u,2:N-1,1:N-2) + view(u,2:N-1, 3:N))*dx2) * (1./ (2*(dx2+dy2)))
    end
    return u
end

function laplace_vec_sub()
    N = 150
    u = zeros(N,N)
    u[1,:] = 1
    Niter = 2^10
    dx2 = dy2 = 0.1*0.1
    u = laplace_iter_vec_sub(u, dx2, dy2, Niter, N)
end


using Turing
#using Plots; gr()

# NOTE: chiange-point model is taken from http://www.mit.edu/~ilkery/papers/GibbsSampling.pdf
@model change_point(x) = begin
  N = length(x)
  n ~ Uniform(1, N)
  lam1 ~ Gamma(2, 1)
  lam2 ~ Gamma(2, 1)
  for i = 1:floor(Int, n)
    x[i] ~ Poisson(lam1)
  end
  for i = floor(Int, n)+1:N
    x[i] ~ Poisson(lam2)
  end
end

_N = 50
_n = floor(Int, rand(Uniform(1, _N)))
_lam1 = rand(Gamma(2, 1))
_lam2 = rand(Gamma(2, 1))
x = Vector{Int}(undef, _N)
for i = 1:_n
  x[i] = rand(Poisson(_lam1))
end
for i = _n+1:_N
  x[i] = rand(Poisson(_lam2))
end
println("N = $_N, n = $_n, lam1 = $_lam1, lam2 = $_lam2")

#plot(x, linetype=[:scatter], lab="data point")
#plot!([_n], linetype=:vline, lab="change point")

modelf = change_point(x)

chn = sample(modelf, SMC(2000))

lam1_ = mean(chn[:lam1]); lam2_ = mean(chn[:lam2]); n_ = mean(chn[:n])

#plot!([i <= n_ ? lam1_ : lam2_ for i = 1:_N], linetype=:steppre, lab="inference result")
#title!("Change-point model (raw data & inference results)")
#xlabel!("n"); ylabel!("counts")


### A Pluto.jl notebook ###
# v0.12.21

using Markdown
using InteractiveUtils

# ╔═╡ 1f62d332-8303-11eb-31e8-4b34a42c4b9f
begin
  using Pkg
    Pkg.activate("MLJ_env", shared=true)

    using DataFrames
    using PlutoUI
    using Test
    using Printf
    using Plots
    using LinearAlgebra
    using Random
    using Statistics
	using StatsBase
	using Distributions
	using Images
	using ImageIO
end

# ╔═╡ f7ad6d8c-8467-11eb-081d-2ba27071eec5
using Colors #, ImageMetadata

# ╔═╡ d67867cc-8302-11eb-3d40-d5a50515cdb9
md"""
## C04w04: Fitting Gaussian Mixture Models with EM

In this assignment we will
  - implement the EM algorithm for a Gaussian mixture model
  - apply your implementation to cluster images
  - explore clustering results and interpret the output of the EM algorithm
"""

# ╔═╡ 53e17eba-8303-11eb-2471-eb8521bf88d2
md"""
### Implementing the EM algorithm for Gaussian mixture models

In this section, we will implement the EM algorithm taking the following steps:
  - Provide a log likelihood function for this model.
  - Implement the EM algorithm.
  - Create some synthetic data.
  - Visualize the progress of the parameters during the course of running EM.
  - Visualize the convergence of the model.
"""

# ╔═╡ 68c8f63a-8303-11eb-0ae3-c3f39e9acf8d
md"""
#### Log likelihood

We provide a function to calculate log likelihood for mixture of Gaussians. The log likelihood quantifies the probability of observing a given set of data under a particular setting of the parameters in our model.

We will use this to assess the convergence of our EM algorithm; specifically, we will keep looping through EM update steps until the log likehood ceases to increase at a certain rate.
"""

# ╔═╡ 84b9158c-8303-11eb-2ce4-c109ce3ff51a
function log_sum_exp(z)
    """
    Compute log(Σᵢ exp(zᵢ)) for some array z.
    """
    max_z = maximum(z)
    max_z + log(sum(exp.(z .- max_z)))
end

# ╔═╡ c01ba57e-8303-11eb-34e7-dfb911bac2f2
function log_likelihood(data, weights, μ, σ)
    """
    Compute the log likelihood of the data for a Gaussian mixture model
	with the given parameters.
    """
    nclusters, ndim = length(μ), size(data[1])[1]
	DT = eltype(data[1])
    llh = zero(DT)

    for d ∈ data
        z = zeros(DT, nclusters)
        for k ∈ 1:nclusters
            ## Compute: (x - μ)ᵀ ×  Σ⁻¹ × (x - μ)
            δ = reshape(d .- μ[k], :, 1) # make sure 2-dim!
            exp_term = dot(δ', inv(σ[k]) * δ)

            ## Compute log likelihood contribution for this data point
			## and this cluster
            z[k] += log(weights[k])
            z[k] -= 1/2. * (ndim * log(2. * π) + log(det(σ[k])) + exp_term)
    	end
		## Increment llh contribution of this data point across all clusters
		llh += log_sum_exp(z)
  	end
    llh
end

# ╔═╡ 835e2f8e-8444-11eb-3f3c-d534adaf4bd6
function approx(actᵥ::T, expᵥ::T; ϵ=1e-8) where T <: Real
	abs(actᵥ - expᵥ) ≤ ϵ
end

# ╔═╡ b0655cea-8445-11eb-1fbc-9350beaed253
function approx(actᵥ::AbstractArray{T}, expᵥ::AbstractArray{T};
		ϵ=1e-8) where T <: Real
	all(t -> approx(t[1], t[2]; ϵ), zip(actᵥ, expᵥ))
end

# ╔═╡ 2d4fb59a-8449-11eb-3f02-a5ece0576a9b
function approx(actᵥ::AbstractArray{T}, expᵥ::AbstractArray{T};
		ϵ=1e-8) where {T <: AbstractArray{T1} where T1 <: Real}
	all(t -> approx(t[1], t[2]; ϵ), zip(actᵥ, expᵥ))
end

# ╔═╡ 580675c0-830a-11eb-16e4-b963a391b138
md"""
#### E-step: assign cluster responsibilities, given current parameters

The first step in the EM algorithm is to compute cluster responsibilities.

Let $r_{ik}$ denote the responsibility of cluster $k$ for data point $i$.
Note that cluster responsibilities are fractional parts. *Cluster responsibilities for a single data point $i$ should sum to 1*.

$$r_{i1} + r_{i2} + \ldots + r_{iK} = 1$$

To figure how much a cluster is responsible for a given data point, we compute the likelihood of the data point under the  particular cluster assignment, multiplied by the weight of the cluster.

For data point $i$ and cluster $k$, this quantity is

$$r_{ik} \propto \pi_k N(x_i | \mu_k, \Sigma_k)$$

where $N(x_i | \mu_k, \Sigma_k)$ is the Gaussian distribution for cluster $k$ (with mean $\mu_k$ and covariance $\Sigma_k$).

We used $\propto$ because the quantity $N(x_i | \mu_k, \Sigma_k)$ is not yet the responsibility we want.
 To ensure that all responsibilities over each data point add up to 1, we add the normalization constant in the denominator:

$$r_{ik} = \frac{\pi_k N(x_i | \mu_k, \Sigma_k)}{\sum_{k=1}^{K} \pi_k N(x_i | \mu_k, \Sigma_k)}$$

Complete the following function that computes $r_{ik}$ for all data points $i$ and clusters $k$.
"""

# ╔═╡ 4d69bcec-830c-11eb-0c80-314c8cded6db
function compute_responsibilities(data, weights, μ, σ)
    """
    E-step fo EM Algorithm:
    - compute responsibilities, given the current parameters
    """
	@assert length(data) > 0

    ndata, nclusters = length(data), length(μ)
    resp = zeros(eltype(data[1]), ndata, nclusters)

    ## Update resp. matrix so that resp[i, k] is the responsibility of
	## cluster k for data point i.
    for i ∈ 1:ndata
        for k ∈ 1:nclusters
			d = MvNormal(μ[k], σ[k])
			resp[i, k] = weights[k] * pdf(d, data[i])
		end
	end
    ## Add up responsibilities over each data point and normalize
    row_sums = reshape(sum(resp, dims=2), :, 1)
    resp ./= row_sums
end

# ╔═╡ 31835154-830d-11eb-09fd-196046609868
begin
  # const ϵ = 1e-8
  resp₀ = compute_responsibilities([[1., 2.], [-1., -2.]],
	[0.3, 0.7],                         # weights
    [[0., 0.], [1., 1.]],               # μ
	[[1.5 0.; 0. 2.5], [1. 1.; 1. 2.]]  # σ
  )

  @test size(resp₀) == (2, 2)
  @test approx(resp₀, [0.10512733 0.89487267; 0.46468164 0.53531836])
end

# ╔═╡ 430e58a8-8310-11eb-2a53-ed2608f43122
md"""
#### M-step: Update parameters, given current cluster responsibilities

Once the cluster responsibilities are computed, we update the parameters (*weights*, *means*, and *covariances*) associated with the clusters.

**Computing soft counts**. Before updating the parameters, we first compute what is known as "soft counts".

The soft count of a cluster is the sum of all cluster responsibilities for that cluster:

$$N^{\text{soft}}_k = r_{1k} + r_{2k} + \ldots + r_{Nk} = \sum_{i=1}^{N} r_{ik}$$

where we loop over data points. Note that, unlike k-means, we must loop over every single data point in the dataset. This is because all clusters are represented in all data points, to a varying degree.

"""

# ╔═╡ 42ef055c-8310-11eb-1ae5-d5ae8bf22c1f
"""
  Compute the total responsibility assigned to each cluster, which will be useful when
  implementing M-steps below. In the lectures this is called N^soft
"""
compute_soft_counts(resp) = sum(resp; dims=1)

# ╔═╡ 93988ed8-8310-11eb-15f4-554b537eacd6
md"""
**Updating weights.** The cluster weights show us how much each cluster is represented over all data points.

The weight of cluster $k$ is given by the ratio of the soft count
$N^{\text{soft}}_{k}$ to the total number of data points $N$:

$$\hat{\pi}_k = \frac{N^{\text{soft}}_{k}}{N}$$

Notice that $N$ is equal to the sum over the soft counts $N^{\text{soft}}_{k}$ of all clusters.

Complete the following function:
"""

# ╔═╡ 937da226-8310-11eb-1966-6bcc8ff3cb47
function compute_weights(counts)
    nclusters, n = length(counts), sum(counts)
    weights = zeros(eltype(counts), nclusters)

	for k ∈ 1:nclusters
        ## Update the weight for cluster k using the M-step update rule for
		## the cluster weight: π̂ₖ
        weights[k] = counts[k] / n
	end
    weights
end

# ╔═╡ 93653664-8310-11eb-23f2-d5236b8bc0ed
begin
	resp₁ = compute_responsibilities([[1., 2.], [-1.,-2.], [0., 0.]],
		[0.3, 0.7],                         # weights
		[[0., 0.], [1., 1.]],               # μ
		[[1.5  0.; 0. 2.5], [1. 1.; 1. 2.]] # σ
  	)
	counts₁ = compute_soft_counts(resp₁)
	weights₁ = compute_weights(counts₁)

	@test approx(weights₁, [0.27904865942515705 0.720951340574843])
end

# ╔═╡ 93490854-8310-11eb-3de5-5b24918d8840
with_terminal() do
  println(" - counts: $(counts₁)")
  println(" - weigths: $(weights₁)")
  println("Checkpoint passed!")
end

# ╔═╡ 42d1c078-8310-11eb-2262-bb5ae92a2b27
md"""
**Updating means**. The mean of each cluster is set to the [weighted average](https://en.wikipedia.org/wiki/Weighted_arithmetic_mean) of all data points, weighted by the cluster responsibilities:

$$\hat{\mu}_k = \frac{1}{N_k^{\text{soft}}} \sum_{i=1}^N r_{ik}x_i$$

Complete the following function:
"""

# ╔═╡ 1b27764e-8315-11eb-3854-f9e5b591744a
function compute_means(data, resp, counts)
    nclusters, ndata = length(counts), length(data)
    means = [zeros(eltype(data[1]), length(data[1])) for _ ∈ 1:nclusters]

    for k in 1:nclusters
        ## Update means for cluster k using the M-step update rule
		## for the mean variables.
        ## This will assign the variable means[k] to be our estimate for μ̂ₖ.
        weighted_sum = sum(
          [resp[i, k] * data[i] for i ∈ 1:ndata]
        )
        means[k] = weighted_sum / counts[k]
	end
    means
end

# ╔═╡ 34b1bb5e-8317-11eb-0260-5f1d3132a7ba
begin
	data_tmp = [[1., 2.], [-1.,-2.]]
	resp₂ = compute_responsibilities(data_tmp,
		[0.3, 0.7],                         # weights
		[[0., 0.], [1., 1.]],               # μ
		[[1.5  0.; 0. 2.5], [1. 1.; 1. 2.]] # σ
	)
	counts₂ = compute_soft_counts(resp₂)
	means₂ = compute_means(data_tmp, resp₂, counts₂)
	exp_means₂ = [[-0.6310085, -1.262017], [0.25140299, 0.50280599]]
	@test approx(means₂[1], exp_means₂[1])
	@test approx(means₂[2], exp_means₂[2])
end

# ╔═╡ 1ad0cbe4-8315-11eb-341c-dbce64acb374
with_terminal() do
  println(" - counts: $(counts₂)")
  println(" - means: $(means₂)")
  println("Checkpoint passed!")
end

# ╔═╡ 1a99af8a-8315-11eb-0f74-eb97741a40ca
md"""

**Updating covariances**.  The covariance of each cluster is set to the weighted average of all [outer products](https://people.duke.edu/~ccc14/sta-663/LinearAlgebraReview.html), weighted by the cluster responsibilities:

$$\hat{\Sigma}_k = \frac{1}{N^{\text{soft}}_k}\sum_{i=1}^N r_{ik} (x_i - \hat{\mu}_k)(x_i - \hat{\mu}_k)^T$$

The "outer product" in this context refers to the matrix product

$$(x_i - \hat{\mu}_k)(x_i - \hat{\mu}_k)^T.$$
Letting $(x_i - \hat{\mu}_k)$ to be $d \times 1$ column vector, this product is a $d \times d$ matrix. Taking the weighted average of all outer products gives us the covariance matrix, which is also $d \times d$.

Reminder:
The inner product (of two column vectors $1 \times n$) denoted as $<v, w> = v^{t}w$ is a scalar, whereas the outer product of the same two column vectors denoted as $v \otimes w = vw^{t}$ is a matrix.

Complete the following function:

"""

# ╔═╡ 986e362c-8317-11eb-0eef-cd1de11cd4b2
function compute_covariances(data, resp, counts, μ)
    nclusters, ndim, ndata = length(counts), size(data[1])[1], length(data)
    covs = [zeros(eltype(data[1]), ndim, ndim) for _ ∈ 1:nclusters]

    for k ∈ 1:nclusters
        ## Update covariances for cluster k using the M-step update rule
		## for covariance variables.
        ## This will assign the variable covariances[k] to be the estimate for Σ̂ₖ.
        w_sum = zeros(eltype(data[1]), ndim, ndim)
        for i ∈ 1:ndata
            w_sum += resp[i, k] * ((data[i] - μ[k]) * (data[i] - μ[k])')
		end
		covs[k] = (1. / counts[k]) * w_sum
	end
    covs
end

# ╔═╡ 57c76562-8373-11eb-0e43-256ec189f399
begin
	## deps on prev test cell for resp₂, counts₂, means₂
	covs₂ = compute_covariances(data_tmp, resp₂, counts₂, means₂)

	exp_covs₂ = [
		[0.60182827 1.20365655; 1.20365655 2.4073131],
		[0.93679654 1.87359307; 1.87359307 3.74718614]
	]

	@test approx(covs₂[1], exp_covs₂[1])
	@test approx(covs₂[2], exp_covs₂[2])
end

# ╔═╡ a12bc08e-8371-11eb-29f0-e14f2c17c1d1
with_terminal() do
  println(" - covariances: $(covs₂)")
  println("Checkpoint passed!")
end

# ╔═╡ c59d7e92-8373-11eb-09d1-55814d7e9f9e
md"""
#### The EM algorithm

Now let us write a function that takes initial parameter estimates and runs EM.
"""

# ╔═╡ c5808488-8373-11eb-3d9a-51c955524390
function em(data, init_μ, init_σ, init_weights;
    maxiter=1000, ϵ=1e-4)
    ## Make copies of initial parameters, which we will update during each iteration
    μ, σ = copy(init_μ), copy(init_σ)
    weights = copy(init_weights)

    ## Infer dimensions of dataset and the number of clusters
    ndata, ndim, nclusters = length(data), size(data[1])[1], length(μ)

    ## Initialize some useful variables
    resp = zeros(eltype(data[1]), ndata, nclusters)
    llh = log_likelihood(data, weights, μ, σ)
    llh_trace = [llh]
	n_iter = 1
    for it ∈ 1:maxiter
		n_iter += 1
        it % 5 == 0 && println("Iteration $(it) / $(maxiter)")

        ## E-step: compute responsibilities
        resp = compute_responsibilities(data, weights, μ, σ)

        ## M-step
        ## Calc. the total resp. assigned to each cluster, useful when
        ## implementing M-steps below. In the lectures this is called N^{soft}
        counts = compute_soft_counts(resp)

        ## update weight of cluster k using M-step update rule for cluster weight π̂ₖ
        weights = compute_weights(counts)

        ## update means of cluster k using the M-step update rule for the mean vars.
        ## assign the variable μ[k] to be our estimate for μ̂ₖ.
        μ = compute_means(data, resp, counts)

        ## update covs ofcluster k using the M-step update rule for cov vars.
        ## assign the variable σ[k] to be the estimate for Σ̂ₖ.
        σ = compute_covariances(data, resp, counts, μ)

        ## Compute the loglikelihood at this iteration
        llh_latest = log_likelihood(data, weights, μ, σ)
        push!(llh_trace, llh_latest)

        ## Check for convergence in log-likelihood and store
        (llh_latest - llh) < ϵ && llh_latest > -Base.Inf && break
        llh = llh_latest
  	end
  	Dict{Symbol, Any}(
    	:weights => weights,
    	:means => μ,
    	:covs => σ,
    	:loglik => llh_trace,
    	:resp => resp
  	)
end

# ╔═╡ c56791bc-8373-11eb-00e3-230f82206892
md"""
#### Testing the implementation on the simulated data

To help us develop and test our implementation, we will generate some observations from a mixture of Gaussians and then run our EM algorithm to discover the mixture components.
We'll begin with a function to generate the data, and a quick plot to visualize its output for a 2-dimensional mixture of three Gaussians.

Now we will create a function to generate data from a mixture of Gaussians model.
"""

# ╔═╡ a790272e-8375-11eb-2a40-df02f3d1ca24
function generate_mog_data(ndata, μ, σ, weights; seed=42)
    """
    Creates a list of data points
    """
    nclusters = length(weights)
    data = []

	Random.seed!(seed)
	rng = MersenneTwister(seed)

    for i ∈ 1:ndata
        ## Use sample (from StatsBase) and weights to pick a:
		## 1 ≤ cluster_id < nclusters
		k = sample(rng, 1:nclusters, Weights(weights), 1; replace=false)[1]

		## Use np.random.multivariate_normal to create data from this cluster
		x = rand(MvNormal(μ[k], σ[k]))  ## ????

		# np.random.multivariate_normal(means[k], covariances[k])
        push!(data, x)
	end
	data
end

# ╔═╡ 00261792-8441-11eb-3969-55f797a05e23

# ╔═╡ a740cde6-8375-11eb-0f2a-efc1f23f8f38
begin
  init_μ = [
	[5., 0.], # mean of cluster 1
	[1., 1.], # mean of cluster 2
	[0., 5.]  # mean of cluster 3
  ]

  init_σ =[
	[.5 0.; 0 .5],       # covariance of cluster 1
	[.92 .38; .38 .91],  # covariance of cluster 2
	[.5 0.; 0 .5]]       # covariance of cluster 3

  init_weights = [1/4., 1/2., 1/4.]  # weights of each cluster

  data = generate_mog_data(100, init_μ, init_σ, init_weights)
end

# ╔═╡ a72ae382-8375-11eb-15fc-c750a071dccf
scatter(map(t -> t[1], data), map(t -> t[2], data), color=[:lightblue], lw=3)

# ╔═╡ b3e522d6-844c-11eb-2aa5-750b3ee4c212
md"""
**Checkpoint** for *em* function
"""

# ╔═╡ 9aee7716-83d7-11eb-3f3c-d534adaf4bd6
 begin
 	m_data = [[0.31845368, 5.03182099], [1.45029339, 2.15717573],
 [1.02645431, 4.76042375], [ 5.23493625, -0.81138851], [ 4.65710743, -0.08532351],
 [0.43746553, 4.93778385], [-0.12756373,  5.14296293], [4.18200738, 0.24819247],
 [ 1.65646672, -0.87991369], [-0.69508006,  5.03848976], [6.01348014, 1.87722666],
 [ 0.20099208, -0.93213203], [-0.38268222,  0.26365374], [ 5.03394554, -1.66740986],
 [ 1.53889119, -0.12731715], [0.81448617, 1.63543709], [1.7219243 , 0.28460715],
 [4.90757786, 0.56716088], [0.60512066, 4.87555125], [ 1.30867865, -0.58711045],
 [0.19570654, 5.73836612], [3.97078181, 0.39328234], [0.47190777, 1.47637309],
 [1.32157017, 0.56486368], [3.89988459, 0.92433204], [0.11205488, 6.4072012 ],
 [0.40717315, 1.67134485], [0.56231532, 5.01349864],
 [1.82113515, 0.78892223], [-0.77578232,  4.62650898], [ 0.45032469, -1.01099583],
 [2.25024431, 0.64832921], [1.47157789, 1.68387047], [-0.0420993 ,  0.50923158],
 [1.18348882, 0.56978228], [5.52201954, 0.46015202], [0.69012115, 1.40085221],
 [0.02126064, 4.76228496], [5.32547107, 1.46856369], [2.58605452, 2.11944448],
 [0.182161  , 4.44577711], [ 0.11836846, -0.33361633], [0.66083532, 1.03879141],
 [1.91659334, 0.63739381], [0.28487706, 4.5665366 ], [0.4764496 , 5.72581327],
 [3.10453287, 2.19997423], [-0.85376105,  4.17389402], [0.08542074, 4.38781614],
 [0.21418408, 5.60131432], [-0.29932464,  4.62574683], [2.25835904, 1.1068076 ],
 [ 4.77358263, -0.06002466], [0.73093444, 1.40220082], [0.39250487, 0.57133562],
 [0.8949198, 1.9300788], [-0.25767001,  4.11851619], [0.55042767, 0.21762904],
 [-0.34493741,  4.57098776], [4.71632515, 0.63293438], [1.67820052, 1.29111918],
 [ 1.4862169 , -0.07691204], [5.46625127, 0.01184258], [-1.33846048, -0.20047555],
 [0.01458455, 5.5685916 ], [1.49605116, 4.82940436], [ 5.75963176, -0.13066323],
 [1.70380591, 1.31376129], [ 4.76328545, -0.33344549], [ 0.8945073 , -1.13036243],
 [1.61279831, 0.66792909], [-0.20028537,  5.0349852 ], [-0.79147695,  4.72403614],
 [1.33337808, 1.427949  ], [1.81545643, 2.37052768], [1.02894019, 3.21262401],
 [-0.1914201 ,  4.03020414], [1.45589226, 0.85329139], [2.23603443, 0.87710097],
 [0.32608726, 0.09240739], [-0.18786611,  5.85099597], [4.29931854, 1.35428737],
 [0.49849313, 0.68444377], [-0.8336478 ,  5.81745508], [-0.57795091, -1.6674263 ],
 [1.83292993, 0.7963827 ],[1.01736806, 0.8455941 ], [0.48493064, 0.45675926],
 [-0.47669162,  3.91542167],[0.96494312, 0.65663175], [1.04999151, 0.38350349],
 [1.17458921, 0.73066238],[0.04638356, 4.86196943], [1.66742479, 1.54181187],
 [0.94372389, 1.32117047],[ 5.63522868, -0.0138849 ],[0.52699679, 5.37301739],
 [4.80426637, 0.8017483], [-0.00321541,  0.2083691],[1.38949165, 0.95020545]
	];

 	m_μs = [[0.19570654, 5.73836612], [6.01348014, 1.87722666],
		[0.52699679, 5.37301739]];

	m_σs =[[ 3.41361848  -1.91940532; -1.91940532  4.68785076],
		[ 3.41361848 -1.91940532; -1.91940532  4.68785076],
		[ 3.41361848 -1.91940532; -1.91940532  4.68785076]];

	m_weights = [0.3333333333333333, 0.3333333333333333, 0.3333333333333333];

	m_results = em(m_data, m_μs, m_σs, m_weights);
	:done
end

# ╔═╡ 4b50e4b0-8444-11eb-3969-55f797a05e23
begin
	@test approx(m_results[:means], Array{Float64,1}[
			Float64[0.021382851931258, 4.947729003491],
			Float64[4.942392352282, 0.3136531084758],
			Float64[1.0818112539095, 0.73903507897]])

	@test approx(m_results[:weights], Float64[0.3007102300634,
			0.17993710074063, 0.5193526691959])

	@test approx(m_results[:covs], Array{Float64,2}[
			[0.2932613984949 0.05048454589297; 0.05048454589297 0.352815373349],
			[0.3556437022527 -0.014948748777286; -0.014948748777286 0.6669502500144],
			[0.6711499192647 0.33058964530721; 0.33058964530721 0.9042972427651]])

	@test approx(m_results[:weights], [0.3007102300634,
			0.17993710074063, 0.5193526691959])

	@test approx(m_results[:loglik][1:5], [-541.3161249736, -372.1355927523,
			-366.9935696838, -365.6599199088, -364.3348647749, -362.87960462340,
			-361.3750403053])
end

# ╔═╡ a7135b54-8375-11eb-2d63-392c4fa94aea
md"""
**Now we will fit a mixture of Gaussians** to this data using our implementation of the EM algorithm. As with k-means, it is important to ask how we obtain an initial configuration of mixing weights and component parameters.

In this simple case, we'll take three random points to be the initial cluster means, use the empirical covariance of the data to be the initial covariance in each cluster (a clear overestimate), and set the initial mixing weights to be uniform across clusters.
"""

# ╔═╡ a6fc6796-8375-11eb-02af-01cee06a06dd
begin
	seed₂ = 2010
	Random.seed!(seed₂)
	rng = MersenneTwister(seed₂)

	## Initialization of parameters
	const K = 3
	chosen = sample(rng, 1:length(data), K; replace=false)
  	#        np.random.choice(length(data), K, replace=False)
  	@assert length(chosen) == K

  	init_μ₁ = [data[x] for x in chosen];
  	init_σ₁ = [cov(data) for _ ∈ 1:K]; # cov(data, rowvar=0) => col. is an observ.
  	init_w = [1.0 / K for _ ∈ 1:K];

  	## Run EM
  	results = em(data, init_μ₁, init_σ₁, init_w);
	:done
end

# ╔═╡ a6b33828-8375-11eb-027a-130ef7a14128
md"""
**Note**. Like k-means, EM is prone to converging to a local optimum. In practice, you may want to run EM multiple times with different random initialization. We have omitted multiple restarts to keep the assignment reasonably short.

Our algorithm returns a dictionary with five elements:
  - 'loglik': a record of the log likelihood at each iteration
  - 'resp': the final responsibility matrix
  - 'means': a list of K means
  - 'covs': a list of K covariance matrices
  - 'weights': the weights corresponding to each model component

"""

# ╔═╡ 4db52910-844d-11eb-3969-55f797a05e23
results[:means]

# ╔═╡ 6f4476bc-844d-11eb-1fbc-9350beaed253
results[:weights]

# ╔═╡ 5ab06896-844d-11eb-3f3c-d534adaf4bd6
md"""
*For the following quiz questions, please round your answer to the nearest thousandth (3 decimals)*

**Quiz Question**: What is the weight that EM assigns to the first component after running the above codeblock?
"""

# ╔═╡ a69ff718-8375-11eb-3965-89ecd2088d97
round(results[:weights][1]; digits=3)

# ╔═╡ a687f546-8375-11eb-3435-1fa822f71ff1
md"""
**Quiz Question**: Using the same set of results, obtain the mean that EM assigns the second component. What is the mean in the first dimension?
"""

# ╔═╡ a66ede62-8375-11eb-3746-758cb7efc8a8
## Answer:   the second component / first dim
round(results[:means][2][1], digits=3)

# ╔═╡ a657ed74-8375-11eb-1d97-1bd417fa42ee
md"""
**Quiz Question**: Using the same set of results, obtain the covariance that EM assigns the third component. What is the variance in the first dimension?
"""

# ╔═╡ a63e00e4-8375-11eb-10db-897f94a0dc32
# Answer:   covs third component / first dim
results[:covs][3, 1]

# ╔═╡ c54c0456-8373-11eb-1cb6-11bf225bc6b0
md"""
### Plot progress of parameters

One useful feature of testing our implementation on low-dimensional simulated data is that we can easily visualize the results.

We will use the following `plot_contours` function to visualize the Gaussian components over the data at three different points in the algorithm's execution:

  1. At initialization (using initial_mu, initial_cov, and initial_weights)
  2. After running the algorithm to completion
  3. After just 12 iterations (using parameters estimates returned when setting `maxiter=12`)
"""

# ╔═╡ 9ba36afc-83d7-11eb-32d4-4d4d0df107d4
function bivariate_normal(X, Y;
    σ_x=1.0, σ_y=1.0, μ_x=0.0, μ_y=0.0, σ_xy=0.0)
    Xμ = X .- μ_x
    Yμ = Y .- μ_y

    ρ = σ_xy / (σ_x * σ_y)
    z = Xμ.^2 ./ σ_x^2 + Yμ.^2 ./ σ_y^2 - 2 * ρ * Xμ .* Yμ ./ (σ_x * σ_y)

  	den = 2. * π * σ_x * σ_y * √(1 - ρ^2)
	exp.(-z / (2. * (1. - ρ^2))) / den
end

# ╔═╡ cd5eb61e-83d9-11eb-1df3-9389091527e8
function meshgrid(x, y)
  @assert length(x) == length(y)
  n = length(x)

  vx, vy = copy(x), copy(y)
  for _ ∈ 1:n-1
    vx = hcat(vx, x)
    vy = hcat(vy, y)
  end
  vx, vy = reshape(vx, :, n)', reshape(vy, :, n)'
  Matrix(vx), Matrix(vy)
end

# ╔═╡ 5a708ed0-83d8-11eb-147c-abf3f80f8c92
function plot_contours(data, means, covs, title)
    δ = 0.025
    k = length(means)
    X = Y = range(-3.0; stop=9.0, step=δ)

	col = [:green, :red, :indigo]
	parms = []
	scatter([_x[1] for _x ∈ data], [_y[2] for _y ∈ data], legend=false, title=title)

	for i ∈ 1:K
		μ, σ = means[i], covs[i]
		σ_x, σ_y = √(σ[1, 1]), √(σ[2, 2])
		σ_xy = σ[1, 2] / (σ_x * σ_y)

		push!(parms, (σ_x=σ_x, σ_y=σ_y, μ_x=μ[1], μ_y=μ[2], σ_xy=σ_xy))
	end

	contour!(X, Y, (x, y) -> bivariate_normal(x, y; parms[1]...), color=col[1])
  	contour!(X, Y, (x, y) -> bivariate_normal(x, y; parms[2]...), color=col[2])
  	contour!(X, Y, (x, y) -> bivariate_normal(x, y; parms[3]...), color=col[3])
end

# ╔═╡ 9b4a0a5e-83d7-11eb-3f02-a5ece0576a9b
plot_contours(data, init_μ₁, init_σ₁, "Initial clusters")

# ╔═╡ e9ae52fe-845e-11eb-15c1-65ebc1f4df8e
begin
	# Parameters after running EM to convergence
	plot_contours(data, results[:means], results[:covs], "Final clusters")
end

# ╔═╡ 3ab0e342-845f-11eb-321d-61972f4e6bd4
md"""
Fill in the following code block to visualize the set of parameters we get after running EM for 12 iterations.
"""

# ╔═╡ 4321d874-845f-11eb-107c-c32e7c640947
begin
	res12 = em(data, init_μ₁, init_σ₁, init_w);
	plot_contours(data, res12[:means], res12[:covs], "Clusters after 12 iterations")
end

# ╔═╡ ffd6d3a2-845f-11eb-1308-39662d5cf06e
md"""
**Quiz Question**: Plot the loglikelihood that is observed at each iteration. Is the loglikelihood plot monotonically increasing, monotonically decreasing, or neither [multiple choice]?
"""

# ╔═╡ a3ac0552-845f-11eb-2283-6defb781f7c7
plot(1:length(results[:loglik]), results[:loglik], lw=3,
	xlabel="Iteration",
	ylabel="Log-likelihood",
	legend=false)

# ╔═╡ 125d76fa-8460-11eb-1d49-57e4eea41e89
md"""
 - The loglikelihood is monotonically increasing
"""

# ╔═╡ c72c03ba-8469-11eb-2fc0-47dab53b1ce9
md"""
## Fitting a Gaussian mixture model for image data¶
"""

# ╔═╡ 075f9b50-8464-11eb-35b3-c945b1e49ccf
begin
	const IMG_DIR = "../../ML_UW_Spec/data/images"
	images = []

	for dir_entry ∈ cd(readdir, IMG_DIR)
		println("Found $(dir_entry)")
		sdir = join([IMG_DIR, dir_entry], "/")

		for img_file ∈ readdir(sdir)
			push!(images, load(join([sdir, img_file], "/")))
		end
	end
end

# ╔═╡ a27f740e-8465-11eb-171f-3ff85b79e29b
length(images)

# ╔═╡ 77bd32d4-8466-11eb-2713-6b11add92ba3
mosaicview(images[1:12]...; nrow=3, ncol=4)

# ╔═╡ fb990ba8-8466-11eb-13be-7d3f33c0576c
mosaicview(images[400:411]...; nrow=3, ncol=4)

# ╔═╡ 36448e26-8467-11eb-12eb-23daf58ea752
mosaicview(images[700:711]...; nrow=3, ncol=4)

# ╔═╡ d1b3f548-8466-11eb-304e-119b7c99959e
mosaicview(images[1000:1011]...; nrow=3, ncol=4)

# ╔═╡ 8c8995d0-8467-11eb-3efb-d1a7961fdefa
img_data = [
	[mean(red.(img)), mean(green.(img)), mean(blue.(img))] for img ∈ images
]

# ╔═╡ bcc9335c-8469-11eb-1492-d310a071b09a
md"""
We need to come up with initial estimates for the mixture weights and component parameters. Let's take three images to be our initial cluster centers, and let's initialize the covariance matrix of each cluster to be diagonal with each element equal to the sample variance from the full data. As in our test on simulated data, we'll start by assuming each mixture component has equal weight.

This may take a few minutes to run.
"""

# ╔═╡ 7b6cc3d8-846e-11eb-315c-5544ae999cc8
var_col(img_data, i::Integer) = var(x[i] for x ∈ img_data)

# ╔═╡ daffa734-8469-11eb-35ea-376de2c6e1e3
begin
	seed₃ = 2010
	Random.seed!(seed₃)
	rng₃ = MersenneTwister(seed₃)

	# Initalize parameters
	const KC = 4
	const NC = 3  # 3 colors R, G, B
	n_images = length(images)

	im_init_μs = [img_data[x] for x ∈ sample(rng₃, 1:n_images, KC; replace=false)]
	@assert length(im_init_μs) == KC

	id_mat = Matrix{Float64}(I, NC, NC);
	im_cov = id_mat .* vec([var_col(img_data, c) for c in 1:NC]);

	im_init_σs = [im_cov for _ ∈ 1:KC];     # [cov, cov, cov, cov]
	im_init_ws = [1.0 / KC for _ ∈ 1:KC]; # [1/4., 1/4., 1/4., 1/4.]

	## Run our EM algorithm on the image data using the above initializations.
	## This should converge in about 125 iterations
	res_img = em(img_data, im_init_μs, im_init_σs, im_init_ws);
	:done
end

# ╔═╡ 7fb68656-8471-11eb-3969-55f797a05e23
function _plot(data_x, data_y)
  plot(data_x, data_y, lw=4, xlabel="Iteration", ylabel="Log-likelihood'")
end

# ╔═╡ ab34bece-8471-11eb-3f3c-d534adaf4bd6
begin
	llh_ = res_img[:loglik]
	_plot(1:length(llh_), llh_)
end

# ╔═╡ cccb8bd0-8471-11eb-1fbc-9350beaed253
md"""
The log likelihood increases so quickly on the first few iterations that we can barely see the plotted line. Let's plot the log likelihood after the first three iterations to get a clearer view of what's going on:
"""

# ╔═╡ d63fd518-8471-11eb-147c-abf3f80f8c92
_plot(3:length(llh_), llh_[3:end])

# ╔═╡ d624eb68-8471-11eb-32d4-4d4d0df107d4
md"""
#### Evaluating uncertainty

Next we'll explore the evolution of cluster assignment and uncertainty. Remember that the EM algorithm represents uncertainty about the cluster assignment of each data point through the responsibility matrix. Rather than making a 'hard' assignment of each data point to a single cluster, the algorithm computes the responsibility of each cluster for each data point, where the responsibility corresponds to our certainty that the observation came from that cluster.

We can track the evolution of the responsibilities across iterations to see how these 'soft' cluster assignments change as the algorithm fits the Gaussian mixture model to the data; one good way to do this is to plot the data and color each point according to its cluster responsibilities. Our data are three-dimensional, which can make visualization difficult, so to make things easier we will plot the data using only two dimensions, taking just the [R G], [G B] or [R B] values instead of the full [R G B] measurement for each observation.

"""

# ╔═╡ 37bd479e-8472-11eb-3106-e3d88fb54de4
# TODO...

# ╔═╡ 37a19df0-8472-11eb-2662-6157d71dd17f
md"""
#### Interpreting each cluster

Let's dig into the clusters obtained from our EM implementation. Recall that our goal in this section is to cluster images based on their RGB values. We can evaluate the quality of our clustering by taking a look at a few images that 'belong' to each cluster. We hope to find that the clusters discovered by our EM algorithm correspond to different image categories - in this case, we know that our images came from four categories ('cloudy sky', 'rivers', 'sunsets', and 'trees and forests'), so we would expect to find that each component of our fitted mixture model roughly corresponds to one of these categories.

If we want to examine some example images from each cluster, we first need to consider how we can determine cluster assignments of the images from our algorithm output. This was easy with k-means - every data point had a 'hard' assignment to a single cluster, and all we had to do was find the cluster center closest to the data point of interest. Here, our clusters are described by probability distributions (specifically, Gaussians) rather than single points, and our model maintains some uncertainty about the cluster assignment of each observation.

One way to phrase the question of cluster assignment for mixture models is as follows: how do we calculate the distance of a point from a distribution? Note that simple Euclidean distance might not be appropriate since (non-scaled) Euclidean distance doesn't take direction into account.  For example, if a Gaussian mixture component is very stretched in one direction but narrow in another, then a data point one unit away along the 'stretched' dimension has much higher probability (and so would be thought of as closer) than a data point one unit away along the 'narrow' dimension.

In fact, the correct distance metric to use in this case is known as [Mahalanobis distance](https://en.wikipedia.org/wiki/Mahalanobis_distance). For a Gaussian distribution, this distance is proportional to the square root of the negative log likelihood. This makes sense intuitively - reducing the Mahalanobis distance of an observation from a cluster is equivalent to increasing that observation's probability according to the Gaussian that is used to represent the cluster. This also means that we can find the cluster assignment of an observation by taking the Gaussian component for which that observation scores highest. We'll use this fact to find the top examples that are 'closest' to each cluster.

__Quiz Question:__ Calculate the likelihood (score) of the first image in our data set (`images[0]`) under each Gaussian component through a call to `multivariate_normal.pdf`.  Given these values, what cluster assignment should we make for this image? Hint: don't forget to use the cluster weights.
"""

# ╔═╡ 3787dc3a-8472-11eb-129f-b5e34bda1328
res_img[:weights]

# ╔═╡ 37702d1a-8472-11eb-1df3-9389091527e8
res_img[:means]

# ╔═╡ d606647c-8471-11eb-3f02-a5ece0576a9b
# d = MvNormal(μ[k], σ[k])
#			resp[i, k] = weights[k] * pdf(d, data[i])

# ╔═╡ d92f92e4-8472-11eb-17cf-eb0a1c444ee9
begin
	proba_img = Float64[]
	for k in 1:KC
		d = MvNormal(res_img[:means][k], res_img[:covs][k])
		push!(proba_img, res_img[:weights][k] * pdf(d, img_data[1]))
	end
	proba_img
end

# ╔═╡ d9132afa-8472-11eb-38c5-39d44a1e57f7
md"""
Now we calculate cluster assignments for the entire image dataset using the result of running EM for 60 iterations above.
"""

# ╔═╡ d8f3fd56-8472-11eb-3856-99f30a89c6d3
begin
	p_weights = res_img[:weights]
	p_means = res_img[:means]
	p_covs = res_img[:covs]
	rgb = img_data

	_N = length(img_data)   ## number of images
	_K = length(p_means)    ## number of clusters

	assignments = zeros(Float64, _N)
	probs = zeros(Float64, _N)

	for i ∈ 1:_N
    	## Compute the score of data point i under each Gaussian component:
    	p = [
			p_weights[k] * pdf(MvNormal(p_means[k], p_covs[k]), rgb[i]) for k ∈ 1:_K
    	]

    	## Compute assign. of each data point to a given cluster based on the scores:
    	assignments[i] = argmax(p)
    	## For data point i, store the corresponding score under this cluster assign:
    	probs[i] = maximum(p)
	end

	df_assign = DataFrame(:cluster_num => Int.(assignments), :probs => probs, :image => images);
	:done
end

# ╔═╡ 9a3f1a20-8475-11eb-1fbc-9350beaed253
last(df_assign, 5)

# ╔═╡ 9a2134e2-8475-11eb-3f3c-d534adaf4bd6
md"""
We'll use the 'assignment' DataFrame to find the top images from each cluster by sorting the datapoints within each cluster by their score under that cluster (stored in probs).

Create a function that returns the top 5 images assigned to a given category in our data.
"""

# ╔═╡ 9a00824c-8475-11eb-3969-55f797a05e23
function get_top_images(df, cluster; k=5)
    images_in_cluster = df[df.cluster_num .== cluster, :]
	top_images = sort(images_in_cluster, [:probs], rev=true)[1:k, :]
    top_images.image
end

# ╔═╡ 17ac70be-8476-11eb-3f02-a5ece0576a9b
mosaicview(get_top_images(df_assign, 1)...; nrow=2, ncol=3)

# ╔═╡ 8b254648-8477-11eb-1df3-9389091527e8
mosaicview(get_top_images(df_assign, 2)...; nrow=2, ncol=3)

# ╔═╡ 8b068190-8477-11eb-147c-abf3f80f8c92
mosaicview(get_top_images(df_assign, 3)...; nrow=2, ncol=3)

# ╔═╡ 8ae806ca-8477-11eb-32d4-4d4d0df107d4
mosaicview(get_top_images(df_assign, 4)...; nrow=2, ncol=3)

# ╔═╡ a5dc1b88-8477-11eb-129f-b5e34bda1328
md"""
These look pretty good! Our algorithm seems to have done a good job overall at 'discovering' the four categories that from which our image data was drawn. It seems to have had the most difficulty in distinguishing between rivers and cloudy skies, probably due to the similar color profiles of images in these categories; if we wanted to achieve better performance on distinguishing between these categories, we might need a richer representation of our data than simply the average [R G B] values for each image.
"""

# ╔═╡ Cell order:
# ╟─d67867cc-8302-11eb-3d40-d5a50515cdb9
# ╠═1f62d332-8303-11eb-31e8-4b34a42c4b9f
# ╟─53e17eba-8303-11eb-2471-eb8521bf88d2
# ╟─68c8f63a-8303-11eb-0ae3-c3f39e9acf8d
# ╠═84b9158c-8303-11eb-2ce4-c109ce3ff51a
# ╠═c01ba57e-8303-11eb-34e7-dfb911bac2f2
# ╠═835e2f8e-8444-11eb-3f3c-d534adaf4bd6
# ╠═b0655cea-8445-11eb-1fbc-9350beaed253
# ╠═2d4fb59a-8449-11eb-3f02-a5ece0576a9b
# ╟─580675c0-830a-11eb-16e4-b963a391b138
# ╠═4d69bcec-830c-11eb-0c80-314c8cded6db
# ╠═31835154-830d-11eb-09fd-196046609868
# ╟─430e58a8-8310-11eb-2a53-ed2608f43122
# ╟─42ef055c-8310-11eb-1ae5-d5ae8bf22c1f
# ╟─93988ed8-8310-11eb-15f4-554b537eacd6
# ╠═937da226-8310-11eb-1966-6bcc8ff3cb47
# ╠═93653664-8310-11eb-23f2-d5236b8bc0ed
# ╠═93490854-8310-11eb-3de5-5b24918d8840
# ╟─42d1c078-8310-11eb-2262-bb5ae92a2b27
# ╠═1b27764e-8315-11eb-3854-f9e5b591744a
# ╠═34b1bb5e-8317-11eb-0260-5f1d3132a7ba
# ╠═1ad0cbe4-8315-11eb-341c-dbce64acb374
# ╟─1a99af8a-8315-11eb-0f74-eb97741a40ca
# ╠═986e362c-8317-11eb-0eef-cd1de11cd4b2
# ╠═57c76562-8373-11eb-0e43-256ec189f399
# ╠═a12bc08e-8371-11eb-29f0-e14f2c17c1d1
# ╟─c59d7e92-8373-11eb-09d1-55814d7e9f9e
# ╠═c5808488-8373-11eb-3d9a-51c955524390
# ╟─c56791bc-8373-11eb-00e3-230f82206892
# ╠═a790272e-8375-11eb-2a40-df02f3d1ca24
# ╠═00261792-8441-11eb-3969-55f797a05e23
# ╠═a740cde6-8375-11eb-0f2a-efc1f23f8f38
# ╠═a72ae382-8375-11eb-15fc-c750a071dccf
# ╟─b3e522d6-844c-11eb-2aa5-750b3ee4c212
# ╠═9aee7716-83d7-11eb-3f3c-d534adaf4bd6
# ╠═4b50e4b0-8444-11eb-3969-55f797a05e23
# ╟─a7135b54-8375-11eb-2d63-392c4fa94aea
# ╠═a6fc6796-8375-11eb-02af-01cee06a06dd
# ╟─a6b33828-8375-11eb-027a-130ef7a14128
# ╠═4db52910-844d-11eb-3969-55f797a05e23
# ╠═6f4476bc-844d-11eb-1fbc-9350beaed253
# ╟─5ab06896-844d-11eb-3f3c-d534adaf4bd6
# ╠═a69ff718-8375-11eb-3965-89ecd2088d97
# ╟─a687f546-8375-11eb-3435-1fa822f71ff1
# ╠═a66ede62-8375-11eb-3746-758cb7efc8a8
# ╟─a657ed74-8375-11eb-1d97-1bd417fa42ee
# ╠═a63e00e4-8375-11eb-10db-897f94a0dc32
# ╟─c54c0456-8373-11eb-1cb6-11bf225bc6b0
# ╠═9ba36afc-83d7-11eb-32d4-4d4d0df107d4
# ╠═cd5eb61e-83d9-11eb-1df3-9389091527e8
# ╠═5a708ed0-83d8-11eb-147c-abf3f80f8c92
# ╠═9b4a0a5e-83d7-11eb-3f02-a5ece0576a9b
# ╠═e9ae52fe-845e-11eb-15c1-65ebc1f4df8e
# ╟─3ab0e342-845f-11eb-321d-61972f4e6bd4
# ╠═4321d874-845f-11eb-107c-c32e7c640947
# ╟─ffd6d3a2-845f-11eb-1308-39662d5cf06e
# ╠═a3ac0552-845f-11eb-2283-6defb781f7c7
# ╟─125d76fa-8460-11eb-1d49-57e4eea41e89
# ╟─c72c03ba-8469-11eb-2fc0-47dab53b1ce9
# ╠═075f9b50-8464-11eb-35b3-c945b1e49ccf
# ╠═a27f740e-8465-11eb-171f-3ff85b79e29b
# ╠═77bd32d4-8466-11eb-2713-6b11add92ba3
# ╠═fb990ba8-8466-11eb-13be-7d3f33c0576c
# ╠═36448e26-8467-11eb-12eb-23daf58ea752
# ╠═d1b3f548-8466-11eb-304e-119b7c99959e
# ╠═f7ad6d8c-8467-11eb-081d-2ba27071eec5
# ╠═8c8995d0-8467-11eb-3efb-d1a7961fdefa
# ╟─bcc9335c-8469-11eb-1492-d310a071b09a
# ╠═7b6cc3d8-846e-11eb-315c-5544ae999cc8
# ╠═daffa734-8469-11eb-35ea-376de2c6e1e3
# ╠═7fb68656-8471-11eb-3969-55f797a05e23
# ╠═ab34bece-8471-11eb-3f3c-d534adaf4bd6
# ╟─cccb8bd0-8471-11eb-1fbc-9350beaed253
# ╠═d63fd518-8471-11eb-147c-abf3f80f8c92
# ╟─d624eb68-8471-11eb-32d4-4d4d0df107d4
# ╠═37bd479e-8472-11eb-3106-e3d88fb54de4
# ╟─37a19df0-8472-11eb-2662-6157d71dd17f
# ╠═3787dc3a-8472-11eb-129f-b5e34bda1328
# ╠═37702d1a-8472-11eb-1df3-9389091527e8
# ╠═d606647c-8471-11eb-3f02-a5ece0576a9b
# ╠═d92f92e4-8472-11eb-17cf-eb0a1c444ee9
# ╟─d9132afa-8472-11eb-38c5-39d44a1e57f7
# ╠═d8f3fd56-8472-11eb-3856-99f30a89c6d3
# ╠═9a3f1a20-8475-11eb-1fbc-9350beaed253
# ╟─9a2134e2-8475-11eb-3f3c-d534adaf4bd6
# ╠═9a00824c-8475-11eb-3969-55f797a05e23
# ╠═17ac70be-8476-11eb-3f02-a5ece0576a9b
# ╠═8b254648-8477-11eb-1df3-9389091527e8
# ╠═8b068190-8477-11eb-147c-abf3f80f8c92
# ╠═8ae806ca-8477-11eb-32d4-4d4d0df107d4
# ╟─a5dc1b88-8477-11eb-129f-b5e34bda1328


"""
    hermit(X, Y, Yp)

Computes interpolating polynomial passing through the points whoose coordinates are provided in `X[]` and `Y[]`,
as well as the values of differentiated function in `Yp[]`, using Hermite interpolation method.
Returns coefficiants of the Hermite polynomial.

    hermit(X::Array{Float64}, Y::Array::Float64, Yp::Array{Float64})

### OUTPUT:
  * `P` - coefficiants of the interpolating polynomial of the type `Array{Float64}`,
"""
function hermit(X::Array{Float64}, Y::Array{Float64}, Yp::Array{Float64})
    n = length(X)
    Z = zeros(n*2)
    Q = zeros(2*n+1, 2*n+1)
    for i in 1:n
        Z[2*i-1] = X[i]
        Z[2*i] = X[i]
        Q[2*i-1, 1] = Y[i]
        Q[2*i, 1] = Y[i]
        Q[2*i, 2] = Yp[i]

        if i > 1
            Q[2*i, 2] = (Q[2*i, 1] - Q[2*i-1, 1])/(Z[2*i] - Z[2*i-1])
        end
    end

    for i in 3:2*n+2
        for j in 3:i
            Q[i, j] = (Q[i, j-1] - Q[i-1, j-1])/(Z[i] - Z[(i+1)-j])
        end
    end

    P = []

    for i in i:2*n+1
        push!(P, Q[i, i])
    end

    return P
end


export minimumjerk, minimumjerkf, functionform

"""
Retorna os coeficientes do polinômio de trajetória com base no princípio do minimum jerk considerando o tempo inicial sendo 0(zero).
Entradas
x0 -> posição inicial
v0 -> velocidade inicial
a0 -> aceleração inicial
t0 -> tempo inicial
xf -> posição final
vf -> velocidade final
af -> aceleração final
tf -> tempo para chegar ao ponto final
"""
function minimumjerk(x0::T, v0::T, a0::T, t0::T, xf::T, vf::T, af::T, tf::T)  where {T<:AbstractFloat}
    b=[ 1 t0 t0^2 t0^3 t0^4 t0^5;
    0 1 2*t0 3*(t0^2) 4*(t0^3) 5*(t0^4);
    0 0 2 6*t0 12*(t0^2) 20*(t0^3);
    1 tf tf^2 tf^3 tf^4 tf^5;
    0 1 2*tf 3*(tf^2) 4*(tf^3) 5*(tf^4);
    0 0 2 6*tf 12*(tf^2) 20*(tf^3)]
    c = [x0, v0, a0, xf, vf, af]
    a = inv(b)*c
    return a
end

"""
Retorna os coeficientes do polinômio de trajetória com base no princípio do minimum jerk considerando o tempo inicial sendo 0(zero).
Entradas
x0 -> posição inicial
v0 -> velocidade inicial
a0 -> aceleração inicial
xf -> posição final
vf -> velocidade final
af -> aceleração final
T -> tempo para chegar ao ponto final
"""
function minimumjerk(x0::T, v0::T, a0::T, xf::T, vf::T, af::T, tf::T) where T<:AbstractFloat
    t0 = 0
    minimumjerk(x0,v0,a0,t0,xf,vf,af,tf)
end

function minimumjerkf(x0::T, v0::T, a0::T, t0::T, xf::T, vf::T, af::T, tf::T) where T<:AbstractFloat
    functionform(minimumjerk(x0,v0,a0,t0,xf,vf,af,tf))
end

function minimumjerkf(x0::T, v0::T, a0::T, xf::T, vf::T, af::T, tf::T) where T<:AbstractFloat
    t0 = 0.0
    functionform(minimumjerk(x0,v0,a0,t0,xf,vf,af,tf))
end

# Retorna as funções de posição, velocidade, aceleração e jerk a parti dos coeficientes do polinômio do mínimo jerk
# Entradas
# a -> vetor com os coeficientes do polinômio
function functionform(a::Vector)
    @eval out1 = $(Meta.parse("t -> $(a[1]) + $(a[2])*t + $(a[3])*(t^2) + $(a[4])*(t^3) + $(a[5])*(t^4) + $(a[6])*(t^5)"))
    @eval out2 = $(Meta.parse("t -> $(a[2]) + 2*$(a[3])*t + 3*$(a[4])*(t^2) + 4*$(a[5])*(t^3) + 5*$(a[6])*(t^4)"))
    @eval out3 = $(Meta.parse("t ->  2*$(a[3]) + 6*$(a[4])*t + 12*$(a[5])*(t^2) + 20*$(a[6])*(t^3)"))
    @eval out4 = $(Meta.parse("t -> 6*$(a[4]) + 24*$(a[5])*t + 60*$(a[6])*(t^2)"))
    return out1, out2, out3, out4
end

# Retorna as funções de posição, velocidade, aceleração e jerk
# partindo de um vetor de posições finais de cada junta e do tempo final
# essa função considera que as condições iniciais (psição, velocidade e aceleração) são
# nulas em cada junta
function minimumjerk(final_positions::Vector{T},tend::T) where T<:AbstractFloat
    xr = Function[]
    vr = Function[]
    ar = Function[]
    jr = Function[]
    for i in final_positions
        x, v, a, j = minimumjerkf(0.0, 0.0, 0.0,0.0, i, 0.0, 0.0, tend)
        push!(xr,x)
        push!(vr,v)
        push!(ar,a)
        push!(jr,j)
    end
    xr, vr, ar, jr
end