### basic imports
# numpy and pandas
import numpy as np
import pandas as pd
import math

## graphics with matplotlib
import matplotlib.pyplot as plt
plt.style.use('seaborn')
%matplotlib

## keras
#from keras.datasets import mnist
#from keras import models
#from keras import layers

from tensorflow import keras


##################################################
## get mnist from keras

(train_images, train_labels), (test_images, test_labels) =  keras.datasets.mnist.load_data()

##################################################
## look at data
## train
print(train_images.shape)
print(len(train_labels))
print(train_labels)

## test
print(test_images.shape)
print(len(test_labels))
print(test_labels)

pd.Series(train_labels).value_counts()/len(train_labels)
pd.Series(test_labels).value_counts()/len(test_labels)

plt.imshow(train_images[0])

ndo=10
for i in range(ndo):
   plt.imshow(train_images[i])
   plt.title(f'It is a {train_labels[i]}')
   plt.pause(.05)
   input('go?')
##################################################
### pick an image and do svd
## http://www.rob-mcculloch.org/2022_cs/webpage/simple-svd.py

ii = 2
anim = train_images[ii]/255
plt.imshow(anim)
fig = plt.gcf()
fig.savefig('anim.pdf')

tsvd = np.linalg.svd(anim)

plt.plot(tsvd[1],'-ob')
plt.xlabel('singular value #'); plt.ylabel('singular value')
plt.title('singular values from svd decomposition of image matrix')
fig = plt.gcf()
fig.savefig('singular-values.pdf')


##################################################
## approx
#s =  len(tsvd[1])
#s = 10
s = 8 # use approx from first 8
U = tsvd[0][:,:s]
Sig = np.diag(tsvd[1][:s])
V = tsvd[2][:s]

check = U @ Sig @ V

plt.scatter(check,anim)
plt.plot(anim,anim,c='r')
plt.xlabel('approximation'); plt.ylabel('true image values')
plt.title(f'svd approx based on {s} singular values')
fig = plt.gcf()
fig.savefig('matrix-vs-approx.pdf')

plt.imshow(check)
fig = plt.gcf()
fig.savefig('image-from-approximation.pdf')