from functools import partial

import numpy as np
import sympy
import matplotlib.pyplot as plt
from sympy.plotting import plot_implicit
from matplotlib import animation

from carry import Carry

def make_animation(fig, frames, **kwargs):
	fig.tight_layout()
	#why isn't this how the function is exposed by default
	return lambda func: animation.FuncAnimation(fig, func, frames, init_func=lambda: None, **kwargs)

def poly_from_array(array):
	ret = 0
	for i, row in enumerate(array):
		for j, val in enumerate(row):
			ret += val*(x**i * y**j)
	return ret

x, y = sympy.symbols("x y")


xy_2 = Carry([[2,-1],[-1,0]])
xy_3 = Carry([[3,-1],[-1,0]])
x2y_3 = Carry([[3,-2],[-1,0]])
x3y_4 = Carry([[4,-3],[-1,0]])

laplace = Carry([[0,-1,0],[-1,4,-1],[0,-1,0]])
laplace3 = Carry([[0,0,0],[-1,3,-1],[0,-1,0]])

almost_folium = Carry([[0,0,-1], [0,2,0], [-1,0,0]])
folium = Carry([[0,0,0,-1], [0,2,0,0], [0,0,0,0], [-1,0,0,0]])
folium3 = Carry([[0,0,0,-1], [0,3,0,0], [0,0,0,0], [-1,0,0,0]])
folium4 = Carry([	[ 0, 0,0,-1, 0], 
					[ 0, 0,0, 0,-1], 
					[ 0, 0,4, 0, 0], 
					[-1, 0,0, 0, 0], 
					[ 0,-1,0, 0, 0]])

def triangle_spread(n):
	ret = np.zeros((n + 1, n + 1), dtype=np.int32)
	ret[0, 0] = -1
	ret[1, 1] = 3
	ret[n, 0] = -1
	ret[0, n] = -1
	return ret

triangle1 = Carry(triangle_spread(1))
triangle2 = Carry(triangle_spread(2))
triangle3 = Carry(triangle_spread(3)) #factorable!
triangle4 = Carry(triangle_spread(4))

tri3_rot = Carry([	[ 0, 0,-1, 0, 0], 
					[ 0, 0, 3, 0, 0], 
					[-1, 0, 0, 0,-1] ])

tri3_rot_real = Carry([	[ 0, 0, 0,-1, 0, 0, 0], 
						[ 0, 0, 0, 0, 0, 0, 0], 
						[ 0, 0, 0, 3, 0, 0, 0], 
						[-1, 0, 0, 0, 0, 0,-1] ])

tri3_rot2 = Carry([	[ 0,-1, 0], 
					[ 0, 3, 0], 
					[-1, 0,-1] ])

tri3_tall = Carry([	[ 0,-1, 0], 
					[ 0, 0, 0], 
					[ 0, 3, 0], 
					[-1, 0,-1] ])

tri3_tall2 = Carry([	[ 0,-1, 0], 
						[ 0, 3, 0], 
						[ 0, 0, 0], 
						[-1, 0,-1] ])

def start_anim(dims=100, carry=xy_2, frames=None, interval=200, inter='add', i_val=1, center=(0,0)):
	if center == (0,0) and carry.over_pos != center:
		center = (dims // 2, dims // 2)

	if inter == 'add':
		next_func = partial(carry.add, val=i_val, center=center)
	elif inter == 'mult':
		if i_val == 1:
			raise ValueError(f"too small value {i_val} for repeated multiplication")
		next_func = partial(carry.mult, val=i_val, center=center)
	else:
		raise ValueError(f"Cannot use {repr(inter)} for animation")

	zero = np.zeros((dims, dims))
	zero[center] = 1
	val = [1]

	fig = plt.gcf()
	plt.title('0')
	image = plt.imshow(zero)
	image.set_clim(0, carry.overflow-1)

	@make_animation(fig, frames, interval=interval)
	def ret(fr):
		next_func(zero)
#		if next_func != xy_2.add:
#			fr = sum(i*(invalid**j) for i,j in zip(zero[0], range(dims))) - 1
		if inter == 'add':
			val[0] += i_val
		else:
			val[0] *= i_val
		plt.title(f"{val[0]}")
		image.set_data(zero)
		fig.tight_layout()

	return ret

def anim_curves(dims=25, invalid=2, frames=None, interval=200, next_func=xy_2.add):
	zero = np.zeros((dims, dims), dtype=np.int32)
	#zero[0,0] = 1
	fig = plt.gcf()
	#plt.colorbar()

	#temp = plt.figure
	#I hate doing it, but there's no other way to get the figure before it's plotted
	plt.figure = bindfig(fig)

	@make_animation(fig, frames, interval=interval)
	def ret(fr):
		next_func(zero, invalid)
		fig.clf()
		plot = plot_implicit(poly_from_array(zero) - fr, backend='matplotlib')
		plt.title(f"{fr+1}")
		print(fr)

	return ret

writer = animation.writers['ffmpeg'](fps=15, metadata={'artist': 'Me'})

def bindfig(fig):
	def ret(**kwargs):
		for i, j in kwargs.items():
			if i == "figsize":
				if j is not None:
					fig.set_figwidth(j[0])
					fig.set_figheight(j[1])
				continue
			fig.__dict__["set_" + i](j)
		return fig
	return ret