from sympy import symbols, I, im, re, sympify, plot_parametric, sin, cos
import numpy as np
import matplotlib.pyplot as plt
from matplotlib import animation

class LazyList:
	def __init__(self, generator):
		self._list = []
		self._generator = generator

	def __getitem__(self, idx):
		final = idx
		if isinstance(idx, slice):
			final = idx.stop
		while final >= len(self._list):
			self._list.append(next(self._generator))
		return self._list[idx]

	def __repr__(self):
		if len(self._list):
			return repr(self._list)[:-1] + ", ...]"
		return "[...]"

t = symbols('t', real=True)

e = (1 + I*t) / (1 - I*t)

def generate_es():
	ret = sympify(1)
	while True:
		yield ret
		ret = (ret * e).expand().cancel().simplify()

es = LazyList(generate_es())
cs = LazyList(map(lambda x: re(x).simplify(), es))
ss = LazyList(map(lambda x: im(x).simplify(), es))


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 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

def anim_curves(stuff=[], interval=200):
	#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, len(stuff), interval=interval)
	def ret(fr):
		fig.clf()
		idx1, idx2 = stuff[fr]

		plot_parametric(idx1, idx2, xlim=(-2,2), ylim=(-2,2), backend='matplotlib')
		#plt.title(f"$x = c_{{{idx1}}}c_{{{idx2}}}; y = c_{{{idx1}}}s_{{{idx2}}}$")
		plt.tight_layout()

	return ret

plot_roses = lambda roses: anim_curves([(cs[i]*cs[j], cs[i]*ss[j]) for i,j in roses], interval=500)
#plot_roses([(1,1),(1,2),(2,1),(3,1),(3,2),(2,3),(1,3),(1,4),(3,4),(4,3),(4,1)]).save()

def archimedes_frames():
	writer = animation.writers['ffmpeg'](fps=15, metadata={'artist': 'Me'})
	#fig = plt.gcf()

	for i in range(10):
	#for i in range(1):
		i += 1
		p = plot_parametric(t*cos(t), t*sin(t), xlim=(-4,4), ylim=(-4,4), label="Archimedean Spiral", show=False)
		p.save(f"frame{i}.png")
		p.extend( 
			plot_parametric(2*i*t*cs[i], 2*i*t*ss[i], line_color="black", label=f"$R_{{{i}}}(t) = {2*i}t$", 
			show=False) 
		)
		p.legend = True
		p.save(f"frame%02d.png" % i)