#! /usr/bin/env python

import numpy as np
import argparse
from matplotlib import pyplot as plt
from matplotlib import animation

def setup_data(args):
	""" Define a set of sample data """
	x = np.linspace(0,1,args.n)
	lam = np.random.normal(0, args.s, args.n)
	y = 2 + 3*x + lam
	return x,y


def test(args):
	# initialize firs state
	x_data,y_data = setup_data(args)

	testdata = np.transpose(np.array([x_data,y_data]))
	#testdata = np.array([x_data,y_data])
	
	U, s, V = np.linalg.svd(testdata, full_matrices=True, compute_uv=True)
	print U.shape
	print s.shape
	print V.shape
	
	S = np.diag(s)
	#A = np.dot(np.dot(U, S), np.transpose(V))
	Asum = np.zeros(len(x_data))
	print U.shape
	print S.shape

	print V.shape
	print np.diag(1/s)
	#Ainv = np.dot(V.dot(np.diag(1/S)), np.transpose(U))
	#Ainv = np.dot(V.dot(np.linalg.pinv(S)), np.transpose(U))
	#Anew = np.dot(V.dot(np.linalg.pinv(S)), np.transpose(U))
	#xnew = (np.linalg.pinv(A)*x_data)
	#xnew = A[:,0]*x_data
	xnew = np.dot(np.dot(V, (np.diag(1/s))), np.dot(np.transpose(U), x_data))
	print xnew.shape
	#print Ainv.shape
	#print A.shape
	test = np.dot(U, S)
	#xnew = np.transpose(Ainv*x_data)
	#xnew = np.linalg.pinv(A)*x_data
	
	#xnew = Ainv.dot(x_data)
	
	fig = plt.figure()
	
	error = np.zeros(len(y_data))

	#for ix in range(len(y_data)):
	#	error = error + (V**2)/(s**2)

	
	
	def animate(i):
		plotimage = plt.imshow(n_stepped[i], interpolation="none")
	
	#anim = animation.FuncAnimation(fig, animate,
	#		frames=len(n_stepped), interval=20, blit=False)
	#anim.save('test.gif', writer='imagemagick', fps=5)
	#anim.save('test.mp4', writer='ffmpeg', fps=5)
	plt.plot(x_data,y_data, x_data, np.transpose(xnew))
	plt.show()

	

if __name__ == '__main__':
	parser = argparse.ArgumentParser()
	parser.add_argument("-n", "--n", type=int, default=100, help="number of points")
	parser.add_argument("-s", "--s", type=int, default=0.5, help="standard deviation")
	args = parser.parse_args()
	test(args)