#! /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)

	#this line doesn't work
	#testdata = np.transpose(np.array([x_data,y_data]))
	testdata = np.dstack((x_data,np.ones_like(x_data)))[0]
	print testdata.shape
	print testdata
	
	#U, s, V = np.linalg.svd(testdata, full_matrices=False, compute_uv=True)
	U, s, V = np.linalg.svd(testdata, full_matrices=False, compute_uv=True)

	S = np.diag(s)
	A = np.dot(np.dot(U, S), np.transpose(V))
	
	print "x -", x_data.shape
	print "y -", y_data.shape
	print "u -", U.shape
	print "s -", s.shape, s
	print "diag(1/s) -", np.diag(1/s).shape, np.diag(1/s)
	print "S -", S.shape, S
	print "V -", V.shape, V

	#func = np.dot(np.dot(V, np.diag(1/s)), np.dot(np.transpose(U), x_data))
	#func = np.dot(np.dot(V.T, np.linalg.pinv(s)), np.dot(U.T, y_data))
	func = np.dot(np.dot(V.T, np.diag(1/s)), np.dot(U.T, y_data))
	print "func -", func.shape
	print func
	
	#newVal = np.dot(testdata,func)
	newVal = x_data*func[0] + func[1]
	print newVal.shape

	fig = plt.figure()
	
	#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, newVal)
	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)