#!/usr/bin/python
# -*- coding: utf8 -*-

import os
import inspect
from math import *
import matplotlib as mpl
import matplotlib.pyplot as plt
from matplotlib import animation

# settings
mpl.rcParams['path.snap'] = False
fname = 'driven-pendulums-resonance-animation'
size = 220, 136
nframes = 100

susp = 16
f0 = 3.0
pendulums = [
{'x0':29, 'y0':117.5, 'f':2.0, 'A':4, 'D':0.05, 'l':89, 'phi0':0.},
{'x0':110, 'y0':117.5, 'f':3.0, 'A':4, 'D':0.05, 'l':89, 'phi0':0.},
{'x0':191, 'y0':117.5, 'f':4.0, 'A':4, 'D':0.05, 'l':89, 'phi0':0.}]

for p in pendulums:
    D = p['D']
    eta = p['f'] / f0
    # use harmonic approximation
    p['Arel'] = 1.0 / sqrt((1. - eta**2)**2 + (2.*eta*D)**2)
    p['Phaserel'] = atan2(2 * eta * D, 1.0 - eta**2)

def animate(nframe, saveframes=False):
    print nframe, nframes
    t = float(nframe) / nframes
    
    plt.clf()
    fig.gca().set_position((0, 0, 1, 1))
    plt.xlim(0, size[0])
    plt.ylim(0, size[1])
    plt.axis('off')
    
    for p in pendulums:
        dxSusp = p['A'] * sin(p['phi0'] + 2*pi*p['f']*t)
        dsPend = p['A'] * sin(p['phi0'] + 2*pi*p['f']*t - p['Phaserel']) * p['Arel']
        PhiPend = (dsPend - dxSusp) / p['l']
        dxPend = dxSusp + p['l'] * sin(PhiPend)
        dyPend = -p['l'] * cos(PhiPend)
        
        xSusp = p['x0'] + dxSusp
        xPend = p['x0'] + dxPend
        yPend = p['y0'] + dyPend
        
        # draw pendulum
        s = 0.72
        plt.plot([p['x0'] - susp, p['x0'] + susp],
                 [p['y0'], p['y0']], '-k', lw=3)
        plt.plot([xSusp, xPend], [p['y0'], yPend], '-k', lw=1.5)
        plt.plot([xSusp], [p['y0']], '.k', markersize=12)
        plt.plot([xPend], [yPend], 'o', color='#aaaaaa',
                 markersize=14, markeredgewidth=1.5)
    
    if saveframes:
        # export frame
        dig = int(ceil(log10(nframes)))
        fsavename = ('frame{:0' + str(dig) + '}.svg').format(nframe)
        fig.savefig(fsavename)
        with open(fsavename) as f: content = f.read()
        content = content.replace('pt"', 'px"').replace('pt"', 'px"')
        with open(fsavename, 'w') as f: f.write(content)

fig = plt.figure(figsize=(size[0]/72., size[1]/72.))

#anim = animation.FuncAnimation(fig, animate, frames=nframes)
#anim.save(fname + '.gif', writer='imagemagick', fps=25)

# use imagemagick on svgs rather than builtin imagemagick support
# this avoids snapping and offers some custom settings
os.chdir(os.path.dirname(os.path.abspath(inspect.getfile(inspect.currentframe()))))
for i in range(nframes):
    animate(i, True)
os.system('convert -loop 0 -delay 4 frame*.svg +dither -posterize 16 ' + fname + '.gif')
for i in os.listdir('.'):
    if i.startswith('frame') and i.endswith('.svg'):
        os.remove(i)