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

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

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

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


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