#!/usr/bin/python

# Makes the sandpile movies based on
# Kadanoff et al. paper

import copy
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.animation as animation
from mpl_toolkits.mplot3d import Axes3D
import avalanche

def plotting(i, ups, surf):
	ax.clear()
	ax.grid(False)
	ax.view_init(azim = 45+i,elev = 70)
	surf = ax.plot_surface(x_data, y_data, heights_array[i], rstride=1, cstride=1,
			cmap='Oranges',linewidth=0)
	print i
	return surf

nx = 50
ny = 50
grain_number = 2000000
sigma_critical = 10
slide_rule = 2
calc_slope_rule = 0
movie = True
last_plot = True

pile = avalanche.avalanche2D(nx = nx, ny = ny, sigma_critical = sigma_critical, 
				slide_rule = slide_rule, boundary = [0,0,0,0], calc_slope_rule = calc_slope_rule)

if (movie or last_plot):
	#plt.ion()
	fig = plt.figure()
	ax = fig.add_subplot(111, projection='3d')
	ax.axis('off')
	x_data, y_data = np.meshgrid(np.arange(pile.heights.shape[1]),
				np.arange(pile.heights.shape[0]))
	heights_array = []


for grain in range(grain_number):

	height, slope = pile.throw_sand()

	if (movie and grain%(grain_number/100) == 0):
		print 'Grain percentage: {0}'.format(grain/grain_number*100.0)
		heights_copy = copy.copy(pile.heights)
		heights_array.append(heights_copy)
		#ax.clear()
		#ax.grid(False)
		#ax.view_init(azim = 180+45,elev = 70)
		#surf = ax.plot_surface(x_data, y_data, pile.heights, rstride=1, cstride=1,
		#	cmap='Oranges',linewidth=0)
		#plt.draw()
	
	if (slope > pile.sigma_critical):
		pile.avalanche(critical_height = False)

if last_plot:
	ax.grid(False)
	ax.view_init(azim = 45,elev = 70)
	surf = ax.plot_surface(x_data, y_data, pile.heights, rstride=1, cstride=1,
		cmap='Oranges',linewidth=0)
	plt.show()

if movie:
	surf = ax.plot_surface(x_data, y_data, heights_array[0], rstride=1, cstride=1,cmap='Oranges',linewidth=0)
	ani = animation.FuncAnimation(fig, plotting, interval = 100, frames=100,
					fargs = (heights_array[0], surf), blit = False)
	ani.save('pile_movie_grains{0}_slide{1}_slope{2}.mp4'.format(grain_number,slide_rule,calc_slope_rule),
			 codec="libx264", bitrate=-1)
