using Plots, Random, Distributions, LinearAlgebra

function MakeData(N, K)
	M = N/K
	Ls = Array{Float64,1}[]
	for k = 1:K
		for i = 1:M
			push!(Ls, [k,k]+0.5*randn(2) )
		end
	end
	return Ls
end

function RandomAssign(X, k)
	unif = DiscreteUniform(1, k)
	return [rand(unif) for x in X]	
end

function PlotSingleCluster(X, C, n)
	clusters = Partition(X, C)
	
	CL = [:red, :blue, :green]

	plot(legend=false, grid=false, size=(500,500), xlims=(0,5), ylims=(0,5))

	for j = 1:maximum(C)
		scatter!([c[1] for c in clusters[j]], [c[2] for c in clusters[j]], c=CL[j])		
	end

	p = plot!(legend=false, grid=false, size=(500,500), xlims=(0,5), ylims=(0,5))
	#png(p, "single$(n)")
end

function PlotCluster(X, C, n)
	CL = [:red, :blue, :green]
	clusters = Partition(X, C)

	scatter([x[1] for x in X], [x[2] for x in X], subplot=1, layout=(2,2), xlims=(0,4), ylims=(0,4))
	
	for j = 1:maximum(C)
		scatter!([c[1] for c in clusters[j]], [c[2] for c in clusters[j]], 
			subplot=j+1, c=CL[j], xlims=(0,4), ylims=(0,4))		
	end

	p = plot!(legend=false, grid=false, size=(500,500))
	#png(p, "prova$(n)")
end

function Kmeans(X, k)
	C = RandomAssign(X, k)

	for i = 1:10
		PlotSingleCluster(X, C, i)
		Z = FindRepr(X, C)	
		C = Assign(X, Z)
		println(Evaluate(X, C, Z))		
	end

	return C
end

# Argmin (see documentation for "findmin": https://docs.julialang.org/en/v1/base/collections/#Base.findmin)
Argmin(X) = findmin(X)[2]

# Inline definition of a simple function
Median(x) = sum(x)/length(x)

# Function that assign each point to the nearest representative
Assign(X, Z) = [Argmin([norm(x-z) for z in Z]) for x in X]

# Evaluate a clustering with given representatives
Evaluate(X, C, Z) = sum([norm(X[i] - Z[C[i]])^2 for i = 1:length(X)])/length(X)

# Partition the input data into cluster (list of list of points in the same cluster)
Partition(X, C) = [[X[i] for i=1:length(X) if C[i] == j] for j=1:maximum(C)]

# Compute the representative for each cluster
FindRepr(X, C) = [Median(c) for c in Partition(X, C)]

#Random.seed!(13)

k = 3
X = MakeData(300, k)
C = Kmeans(X, k)

PlotSingleCluster(X, [1 for c in C], 0)