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)

	# DA COMPLETARE

	return C
end


# 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)]

Random.seed!(13)

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

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