import sys
sys.path.append("../../")

import numpy as np
import matplotlib.pyplot as plt
import matplotlib.animation as animation

from util.io import readvalue

def positive_int(s):
	i = int(s)
	if(not i > 0):
		raise ValueError("{} <= 0".format(i))
	return i

def prepare_A(n):
	A = -2 * np.identity(n)
	A += np.roll(np.identity(n), 1, axis=0)
	A += np.roll(np.identity(n), -1, axis=0)
	return A

N = readvalue("N > ", positive_int)
M = readvalue("M > ", positive_int)
kappa = 0.05

A = prepare_A(N)
eigenvalues, eigenvectors = np.linalg.eigh(A)

#print(*zip(*enumerate(eigenvalues)))
plt.scatter(*zip(*enumerate(eigenvalues)))
plt.savefig("output/eigenvalues.png")

plt.clf()

for n in range(N):
	plt.scatter(*zip(*enumerate(eigenvectors[n])))
	plt.savefig("output/eigenvector_{}.png".format(n))
	plt.clf()

v = np.zeros(N)
v[0] = 1


fig, ax = plt.subplots()
max_eigv, = ax.plot(*zip(*enumerate(eigenvectors[0])), "bo")
scatter, = ax.plot(*zip(*enumerate(v)), "ro")

def on_tick(data):
	global v
	v = (np.identity(N) + kappa*A).dot(v)
	scatter.set_data(*zip(*enumerate(v)))

	return scatter

def run():
	for i in range(M):
		yield i


ani = animation.FuncAnimation(fig, on_tick, run, interval=1)
ani.save("output/animation.gif", dpi=80, writer='imagemagick')