using Flux
using Flux.Data.MNIST

labels = MNIST.labels();
images = MNIST.images(); 

using Pkg
Pkg.add("ImageShow")

using ImageShow

size(images)

size(images[8])

28*28

images[1:5]

reshape(images[1:25], 5, 5)

using Interact

images[17]

labels[17]

@manipulate for i in eachindex(images)
    hbox(images[i], hskip(1em), string("label: ", labels[i]))
end



a = vec(images[2][10, :]);

b = round.(Int, 256*images[2][10, i].val for i in 1:28);

a[1]

b[1]

using Pkg
Pkg.add("Interact")
using Interact

@manipulate for j in 1:28
    a = vec(images[2][j, :])
    b = round.(Int, 256*images[2][j, i].val for i in 1:28);
    hbox([vbox(a[i], string(" ", b[i], " " )) for i in 1:28])
end

@manipulate for j in 1:28
    a = vec(images[2][j, :])
    b = round.(Int, 256*images[2][j, i].val for i in 1:28);
    
    hbox([vbox(a[i], string(b[i])) for i in 1:28])
end

@manipulate for start in 1:28^2 - 28
    a = vec(images[2][start:start+27])
    b = round.(Int, 256*images[2][i].val for i in start:start+27);
    
    hbox([vbox(a[i], hbox(hskip(0.5em), string(b[i]))) for i in 1:28])
end

training_size = 50000

training_images = images[1:training_size]
training_labels = labels[1:training_size];

images_0 = training_images[training_labels .== 0];
images_0[1:100]

length(images_0)

function vectorize(image) 
    return round.(Int, 256*image[i].val for i in eachindex(image))
end

show(vectorize(images[2]))

matrix_0s = reduce(hcat, vectorize.(images_0));

using LinearAlgebra

U, Σ, V = svd(matrix_0s)

U * Diagonal(Σ) * V' ≈ matrix_0s

sum(Σ.^2) ≈ sum(matrix_0s.^2)

using Plots

scatter(log10.(Σ))

function normalize_image(image)
    image2 = abs.(image)
    m = maximum(image2)
    image2 = image2 ./ m
    return Gray.(reshape(image2, 28, 28))
end

using Statistics

normalize_image(mean(matrix_0s, dims=2))

big_image_matrix = reduce(hcat, vectorize.(training_images));

size(big_image_matrix)

big_image_matrix

single_digits = [ big_image_matrix[:, training_labels .== i] for i in 0:9 ]

Z = svd.(single_digits);

Us = first.(Z)

using Statistics

U[:, 1]

norm(U[:, 1])

show(U[:, 1])

@manipulate for θ in 0.0:0.01:0.12
    image = abs.(U[:, 1]) .> θ
    Gray.(reshape(image, 28, 28))
end

Us

@manipulate for digit in slider(0:9, value=0), which_sing_vec in slider(1:10, value=1)

    image = Us[digit+1][:,which_sing_vec]    
   
    normalize_image(image)
end

image = images[training_size + 1]   # was not included in the training

M = [U[:,1:3] for U in Us]  # first n sing vecs for each digit

first_sing_vecs = reduce(hcat, [U[:,1] for U in Us])

my_softmax(v) = exp.(v) / sum(exp.(v))

v = abs.(first_sing_vecs' * vec(Float64.(image)))

my_softmax(v)

M = [U[:, 1:1] for U in Us]  # first n sing vecs for each digit

my_softmax(norm.([MM' * vec(Float64.(image)) for MM in M]))

@manipulate for num_sing_vecs in 1:784

    M = [U[:, 1:num_sing_vecs] for U in Us]  # first n sing vecs for each digit

    v = my_softmax(norm.([MM' * vec(Float64.(image)) for MM in M]))
    
    hbox(v, plot(0:9, v, m=:o), findmax(v))
    
end

data = []

for num_sing_vecs in 1:15

    M = [U[:, 1:num_sing_vecs] for U in Us]  # first n sing vecs for each digit

    v = my_softmax(norm.([MM' * vec(Float64.(image)) for MM in M]))
    
    push!(data, v[4])
end

plot(data, m=:o)
scatter!([argmax(data)], [maximum(data)])
    

argmax(data)

softmax(norm.([MM' * vec(Float64.(image)) for MM in M]))

vec(images[1])