using Plots, Distributions, LinearAlgebra

function ReadMNIST(filename)
    f = open(filename)
    A = map(row -> map(x -> parse(Int, x), split(row, ',')), eachline(f))
    close(f)

    T = [a[2:end] for a in A]
    L = [a[1] for a in A]
	N = length(A[1])        
    return length(A), Int(sqrt(N-1)), T, L
end

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

function PlotDigits(row, k)
	R(X) = reshape(X, (28,28))'[end:-1:1,:]

	A = zeros(28, 28*k)
	for i = 1:k
		A[1:28, (28*(i-1)+1):(28*(i))] = R(row[i])
	end

    p = heatmap(A, subplot=1, aspect_ratio=1, xticks = nothing, yticks = nothing, 
    	         colorbar=false, c=cgrad([:lightgrey, :red, :blue]))

#    png(p, "digits_all")
end

# Inline definition of a simple function
Median(x) = sum(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)]

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

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

function Kmeans(X, k)
	C = RandomAssign(X, k)
	Z = FindRepr(X, C)	
	J0 = Evaluate(X, C, Z)
	for i = 1:100
		#PlotSingleCluster(X, C, i)
		Z = FindRepr(X, C)	
		C = Assign(X, Z)
		J = Evaluate(X, C, Z)
		if abs(J - J0) < 1e-07
			break
		end
		J0 = J
		println("$(i): $(J)")		
	end

	return FindRepr(X, C)
end

n, m, X = ReadMNIST("Test_MNIST_1_2_8.csv")
#n, m, X = ReadMNIST("mnsit_all.csv")

k = 6
Z = Kmeans(X, k)
#C = RandomAssign(X, 2)
#Z = FindRepr(X, C)			

PlotDigits(Z, k)	

#D1 = Median(X[1:100])
#D8 = Median(X[101:200])
#D2 = Median(X[201:end])
#PlotDigits([D1, D2, D8], 3)