import numpy as np
import matplotlib.pyplot as plt
import matplotlib.image as mpimg
import sys

orig = mpimg.imread(sys.argv[1])

def rgb2(rgb):
    return [np.dot(rgb[...,:3], [1, 0, 0]), np.dot(rgb[...,:3], [0, 1, 0]), np.dot(rgb[...,:3], [0, 0, 1])]

def rbg2gray(rgb):
  return np.dot(rgb[...,:3], [0.299, 0.587, 0.114])

#gray = rgb2gray(orig) 

def plot_image(im):
    plt.imshow(im)
    plt.show()

def get_PCA(M, r):
    u, s, vh = np.linalg.svd(M)
    U=np.transpose(np.transpose(u)[0:r])
    Vh=vh[0:r]
    S=np.diag(s[0:r])
    projection = np.matrix(U)*np.matrix(S)*np.matrix(Vh)
    return projection

if sys.argv[2]=='p':
    print(str(len(gray))+'x'+str(len(gray[1])))
    plot_image(gray)
else:
	dim = int(sys.argv[2])
	red = get_PCA(rgb2(orig)[0],dim)
	green = get_PCA(rgb2(orig)[1], dim)
	blue = get_PCA(rgb2(orig)[2], dim)
	rbgArray = np.zeros((len(orig), len(orig[0]), 3))
	rbgArray[...,0]=red
	rbgArray[...,1]=green
	rbgArray[...,2]=blue
	plot_image(rbgArray)