# Grab some modules
%matplotlib inline
import numpy as np
import matplotlib.pyplot as plt
from matplotlib import animation
import io
import base64
from IPython.display import HTML
from tempfile import NamedTemporaryFile

# Clever bits that we'll need for embedding animations
# http://jakevdp.github.io/blog/2013/05/12/embedding-matplotlib-animations/

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, fps):
    if not hasattr(anim, '_encoded_video'):
        with NamedTemporaryFile(suffix='.mp4') as f:
            anim.save(f.name, fps=fps, writer='avconv', extra_args=['-vcodec', 'libx264'])
            video = open(f.name, "rb").read()
        anim._encoded_video = video.encode("base64")
    
    return VIDEO_TAG.format(anim._encoded_video)

def display_html_video(video_tag_string):
    return HTML(video_tag_string)

class Disk:
    """
    A rotating wheel connected to a parent, part of the spirograph stack
    """
    
    def __init__(self, index, parent, radius, angvel, inner, store=None, pencil=None):
        """
        Create a disk and initialise state
        """
        
        # Properties
        self.index  = index
        self.inner  = inner
        self.radius = radius
        self.angvel = angvel
        self.pencil = pencil
        
        # Relative position
        self.body_angle = 0    # Angle of the disk relative to the parent rotor
        self.rotor_angle = 0   # Angle of the rotor relative to the disk
        
        # Attach to parent
        self.child = None
        self.parent = None
        if parent is not None:
            self.parent = parent
            self.parent.attach(self)
        
        # Attach the store object
        self.store = store
            
    
    def attach(self, child):
        """
        Attach child disk
        """
        
        self.child = child
    
    
    def increment(self):
        """
        Increment time
        """
        
        # Update rotor angle
        self.rotor_angle += self.angvel
        
        # Modify child's body angle accordingly and recurse
        if self.child is not None:
            if self.inner:
                self.child.body_angle -= self.angvel * (self.radius/self.child.radius)
            else:
                self.child.body_angle += self.angvel * (self.radius/self.child.radius)
            self.child.increment()
    
    
    def emit_position(self):
        """
        Calculate absolute position and store it
        """
        
        if self.parent is None:
            x = 0;
            y = 0;
            body = self.body_angle;
            rotor = self.rotor_angle;
        else:
            parent_x, parent_y, parent_body, parent_rotor = self.parent.emit_position()
            if self.inner:
                x = parent_x + (self.parent.radius-self.radius) * np.sin(parent_rotor)
                y = parent_y + (self.parent.radius-self.radius) * np.cos(parent_rotor)
            else:
                x = parent_x + (self.parent.radius+self.radius) * np.sin(parent_rotor)
                y = parent_y + (self.parent.radius+self.radius) * np.cos(parent_rotor)
            body = parent_rotor + self.body_angle
            rotor = body + self.rotor_angle
        
        if self.pencil is not None:
            pencilx = x + self.pencil * self.radius * np.sin(rotor)
            pencily = y + self.pencil * self.radius * np.cos(rotor)
        
        self.store[self.index]['x'].append(x)
        self.store[self.index]['y'].append(y)
        self.store[self.index]['body'].append(body)
        self.store[self.index]['rotor'].append(rotor)
        self.store[self.index]['radius'].append(self.radius)
        if self.pencil is not None:
            self.store[self.index]['pencilx'].append(pencilx)
            self.store[self.index]['pencily'].append(pencily)
        
        
        return x, y, body, rotor

def simulate(T, R, radius, angvel, inner, pencil=None):
    """
    Simulate the motion of the disks and record their trajectories over time.
    """
    
    # R is the number of time steps to do 1 rotation at a frequency of 1
    # T is the number of rotations at frequency 1 to simulate
    
    N = len(radius)
    assert len(angvel)==N, "Lists of radius and angular velocity must be of the same length."
    assert len(inner)==N, "Lists of radius and inner-ness must be of the same length."
    if pencil is not None:
        assert len(pencil)==N, "Lists of radius and pencil must be of the same length."
    else:
        pencil = [None for ii in range(N)];
    
    # Make the disks
    disk_list = []
    parent = None
    store = dict()
    for ii in range(N):
        disk = Disk(ii, parent, radius[ii], 2*np.pi*angvel[ii]/R, inner[ii], pencil=pencil[ii], store=store)
        disk_list.append(disk)
        store[ii] = dict()
        store[ii]['x'] = []
        store[ii]['y'] = []
        store[ii]['body'] = []
        store[ii]['rotor'] = []
        store[ii]['radius'] = []
        store[ii]['pencilx'] = []
        store[ii]['pencily'] = []
        parent = disk
    
    # Go!
    for tt in range(int(T*R)):
        disk_list[0].increment()
        disk_list[-1].emit_position()
    
    return disk_list, store

def drawit(store, centre_idxs, pencil_idxs):
    """
    Draw the trajectory of the disks, centres or pencils
    """
    
    fig = plt.figure(figsize=(10,10))
    ax = fig.add_subplot(1,1,1)
    ax.set_aspect('equal')
    
    for idx in centre_idxs:
        ax.plot(store[idx]['x'], store[idx]['y'])
    
    for idx in pencil_idxs:
        ax.plot(store[idx]['pencilx'], store[idx]['pencily'])
    
    plt.show()
    return fig

def disk_drawing_elements(x, y, r, ba, ra):
    disk_elements = []
    th = np.linspace(0, 2*np.pi, 100)
    circumx = x + r*np.sin(th)
    circumy = y + r*np.cos(th)
    basex = np.array([x, x+r*np.sin(ba)])
    basey = np.array([y, y+r*np.cos(ba)])
    rotorx = np.array([x, x+r*np.sin(ra)])
    rotory = np.array([y, y+r*np.cos(ra)])
    return circumx, circumy, basex, basey, rotorx, rotory

def disk_animation(store, centre_idxs, pencil_idxs, fps, ds, lim):
    
    # Definte a palate
    colours = ['r', 'b', 'g', 'c', 'm']
    
    # Some useful constants
    T = len(store[0]['x'])  # Number of time steps
    N = len(store)                 # Number of disks
    
    # Create the real figure
    fig = plt.figure(figsize=(8,8), dpi=300)
    ax = plt.axes(xlim=[-lim,lim], ylim=[-lim,lim])
    ax.set_aspect('equal')
    ax.set_axis_off()

    # Create the plot elements
    centre_traces = []
    pencil_traces = []
    outlines = []
    baselines = []
    rotors = []
    for ii in range(N):
        centre_traces.append( ax.plot([],[], color=colours[ii])[0] )
        pencil_traces.append( ax.plot([],[], color=colours[ii])[0] )
        outlines.append( ax.plot([],[], color=colours[ii])[0])
        baselines.append( ax.plot([],[], color='k')[0] )
        rotors.append( ax.plot([],[], color='k')[0] )
    lines = tuple(centre_traces + pencil_traces + outlines + baselines + rotors)
        
    def init():
        for line in lines:
            line.set_data([], [])
        return lines
    
    def animate(t):
        tt = min(len(store[0]['x'])-1, (t-1)*ds)
        for ii in range(N):
            idx = ii
            hist = store[idx]
            circumx, circumy, basex, basey, rotorx, rotory = disk_drawing_elements(
                hist['x'][tt], hist['y'][tt], hist['radius'][tt], hist['body'][tt], hist['rotor'][tt])
            if idx in centre_idxs:
                centre_traces[ii].set_data(hist['x'][:tt+1], hist['y'][:tt+1])
            if idx in pencil_idxs:
                pencil_traces[ii].set_data(hist['pencilx'][:tt+1], hist['pencily'][:tt+1])
            outlines[ii].set_data(circumx, circumy)
            baselines[ii].set_data(basex, basey)
            rotors[ii].set_data(rotorx, rotory)
            
        return lines
    
    anim = animation.FuncAnimation(fig, animate, init_func=init,
                               frames=int(T/ds)+3, interval=1000.0/fps, blit=True)
    encoded_anim = anim_to_html(anim, fps)
    plt.close(anim._fig)
    
    return encoded_anim

R = 100
T = 16+1.0/R
radius = [1.0, 1.0/2.0, 1.0/4.0, 1.0/8.0]
angvel = [1.0, 1.0/2.0, 1.0/16.0, 1.0/4.0]
inner = [True, True, True, True]

disk_list, store = simulate(T, R, radius, angvel, inner)
fig = drawit(store, range(4), [])

html_anim = disk_animation(store, range(4), [], 20, 2, 1.1)
display_html_video(html_anim)

R = 10000
T = 1+1.0/R
radius = [1.0, 1.0/4.0, 1.0/16.0, 1.0/64.0]
angvel = [1.0, 10.0, 10.0**2, 10.0**3]
inner = [False, False, False, False]
pencil = [0.8, 0.8, 0.8, 0.8]

disk_list, store = simulate(T, R, radius, angvel, inner, pencil=pencil)
fig = drawit(store, [], range(4))

html_anim = disk_animation(store, [], range(4), 20, 10, 2.0)
display_html_video(html_anim)

N = 5
R = 50000
T = 1+1.0/R
radius = [1.0, 1.0/3.0, 1.0/9.0, 1.0/27.0, 1/81.0]
angvel = [1.0, -10, 100, -1000, 10000]
inner = [False]*N
pencil = [0.9]*N

disk_list, store = simulate(T, R, radius, angvel, inner, pencil=pencil)
fig = drawit(store, [], range(N))

html_anim = disk_animation(store, [], range(N), 20, 50, 2.0)
display_html_video(html_anim)