%pylab inline

from tempfile import NamedTemporaryFile

VIDEO_TAG = """<video controls>
 <source src="data:video/x-m4v;base64,{0}" type="video/mp4">
 Your browser does not support the video tag.
</video>"""

def anim_to_html(anim):
    if not hasattr(anim, '_encoded_video'):
        with NamedTemporaryFile(suffix='.mp4') as f:
            anim.save(f.name, fps=20, dpi=70,extra_args=['-vcodec', 'libx264', '-pix_fmt', 'yuv420p'])
            video = open(f.name, "rb").read()
        anim._encoded_video = video.encode("base64")
    
    return VIDEO_TAG.format(anim._encoded_video)

from IPython.display import HTML

def display_animation(anim):
    plt.close(anim._fig)
    return HTML(anim_to_html(anim))

class spirograph():
    def __init__(self,a,b,f,noZ=False):
        self._a=a #number of teeth on small gear
        self._b = b #number of teeth on outer wheel
        self._rho = f*a #radius of pen from center of small wheel (in units of tooth spacing) [redundant with below]
        self._f = f  #fraction of the inner wheel's radius for where the pen goes.
        self._noZ = noZ #a switch so that the z-component of the spirograph traces time.
    def inspect(self):
        print self._a, self._b, self._f
    def graph(self,t0):
        #if t0==0: self.inspect()
        a=self._a
        b=self._b
        rho=self._rho
        f=self._f
        lengthscale=5.*2*np.pi/b #scale the spirograph so outer ring is ~5 in graphing coordinates
        timescale=min(a,b)/gcd(a,b) #scale timing so that when t0=2π the spirograph is a closed curve
        return (lengthscale*((b-a)*cos(timescale*t0)+rho*cos(-(1.*b/a -1.)*timescale*t0)),
                lengthscale*((b-a)*sin(timescale*t0)+rho*sin(-(1.*b/a -1.)*timescale*t0)),
                0 if self._noZ else 5+5*t0 )

import numpy as np
from numpy import sin, cos

from matplotlib import pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from matplotlib.colors import cnames
from matplotlib import animation
from fractions import gcd

# Solve for the trajectories (if t->2pi the spirograph will complete)
t = np.linspace(0, 2.1*np.pi, 500)

noZswitch=True
myspiros = [spirograph(a=63,b=96,f=0.6,noZ=True),spirograph(a=63,b=96,f=0.8,noZ=True),
            spirograph(a=51,b=96,f=0.6,noZ=True),spirograph(a=51,b=96,f=0.8,noZ=True)]
N_trajectories = len(myspiros)


x_t = np.asarray([[myspiro.graph(t0) for t0 in t] for myspiro in myspiros])


#use the right sum http://stackoverflow.com/questions/17313853/unexpected-typeerror-with-ipython
#maybe there's a better way?
import __builtin__
sum = __builtin__.sum

#3-d plotting with minimal modifications from Jake's lorentz system example
# Set up figure & 3D axis for animation
fig = plt.figure()
ax = fig.add_axes([0, 0, 1, 1], projection='3d')
ax.axis('off')

# choose a different color for each trajectory
colors = plt.cm.jet(np.linspace(0, 1, N_trajectories))
for c in colors:
    print c

# set up lines and points
lines = sum([ax.plot([], [], [], '-', c=c)
             for c in colors], [])
pts = sum([ax.plot([], [], [], 'o', c=c)
           for c in colors], [])

# prepare the axes limits
ax.set_xlim((-25, 25))
ax.set_ylim((-35, 35))
ax.set_zlim((5, 55))

# set point-of-view: specified by (altitude degrees, azimuth degrees)
ax.view_init(30, 0)

# initialization function: plot the background of each frame
def init():
    for line, pt in zip(lines, pts):
        line.set_data([], [])
        line.set_3d_properties([])

        pt.set_data([], [])
        pt.set_3d_properties([])
    return lines + pts

# animation function.  This will be called sequentially with the frame number
def animate(i):
    # we'll step two time-steps per frame.  This leads to nice results.
    i = (2 * i) % x_t.shape[1]

    for line, pt, xi in zip(lines, pts, x_t):
        x, y, z = xi[:i].T
        line.set_data(x, y)
        line.set_3d_properties(z)

        pt.set_data(x[-1:], y[-1:])
        pt.set_3d_properties(z[-1:])

    ax.view_init(90*cos(np.pi*i/500.), 0.3 * i)
    fig.canvas.draw()
    return lines + pts

anim = animation.FuncAnimation(fig, animate, init_func=init,
                               frames=500, interval=30,blit=True)
display_animation(anim)


#an alternate way to display the animation
#animation.Animation._repr_html_ = anim_to_html
#animation.FuncAnimation(fig, animate, init_func=init,
#                               frames=100, interval=20, blit=True)