'''
Much thanks to this person for the tutorial on image processing
http://nbviewer.ipython.org/gist/frankcleary/4d2bd178708503b556b0
'''

import matplotlib.pyplot as plt
import numpy as np
import time
from PIL import Image

#This image is hard coded to my desktop
img = Image.open('/Users/briand/Desktop/briand/head_shot.jpg').convert('LA')
imgmat = np.array(list(img.getdata(band=0)), float)
imgmat.shape = (img.size[1], img.size[0])
imgmat = np.matrix(imgmat)





#plt.imshow(imgmat, cmap='gray');

U, sig, Vt = np.linalg.svd(imgmat)


#Plot the spectrum of the image and the total sum-squares captured by the k-th value

norm = math.sqrt(sum(sig*sig))
energy_k = [math.sqrt(k)/norm for k in cumsum(sig*sig)]

plt.figure()
ax1 = plt.subplot(211)
plt.plot(sig)
plt.title('Kth Singular Value')
plt.tick_params(axis='x',which='both',bottom='off',top='off',labelbottom='off')


ax2 = plt.subplot(212)
plt.plot(energy_k)
plt.title('Normalized Cumulative Energy of Kth Singular Value')

ax2.set_xlabel('Kth Singular Value')



#Now work on an approximate image
def plotApprox(U, sig, Vt, k):
    reconstimg = np.matrix(U[:, :k]) * np.diag(sig[:k]) * np.matrix(Vt[:k, :])
    plt.imshow(reconstimg, cmap='gray');
    
def plotSeq(U, sig, Vt, seq_k):
    fig = plt.figure()
    for i,k in enumerate(seq_k[:8]):
        ax = fig.add_subplot(2,4,i+1)
        plotApprox(U, sig , Vt, k)
        plt.title('Rank {}'.format(k))
        ax.axes.get_xaxis().set_visible(False)
        ax.axes.get_yaxis().set_visible(False)
        


plotSeq(U,sig,Vt,[1,2,3,4,10,25,50,200])

#plotApprox(U, sig, Vt, 10)