#!/usr/bin/env python2
from pylab import *
rcParams['savefig.dpi'] = 100
lightcolor = '#ffc22a'
beamcolor = '#234781'
velocitycolor = '#aaaaaa'
def squigglyarrow(x,y,xp,yp, amplitude, ax, wavelength=1, color='b'):
dx, dy = xp-x, yp-y
L = sqrt(dx ** 2 + dy ** 2) or 1
rot = array([[:en:dx/L, -dy/L],
             [dy/L,  dx/L]])
tailL = L - 3*amplitude
t = r_[arange(0, tailL ,wavelength/50.0), tailL]
xy = c_[t, amplitude*sin(2*pi*t/wavelength)]
rxy = rot.dot(xy.T).T
line = Line2D(rxy[:,0] + x, rxy[:,1] + y, color=color, lw=2, zorder=-1)

arrowpath = rot.dot(array([[:en:tailL, 2*amplitude], 
                           [L,0], 
   [tailL, -2*amplitude]]).T).T
arrowhead = Polygon(arrowpath + array([x,y]), color=color, ec='none')
ax.add_artist(line)
ax.add_patch(arrowhead)
def drawframe(axobs, axlit, t):
axobs.add_patch(Circle((0,12), radius=1, color=lightcolor, ec='none'))
axobs.add_patch(Circle((-8 + 8*t,0), radius=1, color='w', ec='k', lw=1))
squigglyarrow(0,12-t*7,-0,12-4-t*7, 0.2, axobs, wavelength=1.13, color=beamcolor)
axobs.set_xticks([])
axobs.set_yticks([])
axobs.set_xlim(-12,8)
axobs.set_ylim(-5,15)
axobs.set_title("Source's rest frame")
axlit.add_patch(Circle((0-8*t,12), radius=1, color=lightcolor, ec='none'))
axlit.add_patch(Circle((-8,0), radius=1, color='w', ec='k', lw=1))
lx, ly = 4.0*8/sqrt(8*8+12*12), 4*12/sqrt(8*8+12*12)
squigglyarrow(0-(8-lx)*t,12-(11-ly)*t,-lx-(8-lx)*t, 12-ly-(11-ly)*t, 0.2, axlit, wavelength=1.13, color=beamcolor)
axlit.set_xticks([])
axlit.set_yticks([])
axlit.set_xlim(-15,5)
axlit.set_ylim(-5,15)
axlit.set_title("Observer's rest frame")
def drawfigaxes(ax):
ax.set_xlim(0,1)
ax.set_ylim(0,1)
ax.set_xticks([])
ax.set_yticks([])
ax.add_artist(Line2D([0.5,0.5], [0,1], color='#888888', zorder=-1))
def drawframes():
fig = figure(figsize=(6,3*1.1), facecolor='white')
axobs = Axes(fig,[0,0,0.5,0.9],yticks=[],xticks=[],frame_on=False)
axlit = Axes(fig,[0.5,0,0.5,0.9],yticks=[],xticks=[],frame_on=False)
figAxes = Axes(fig, [0,0,1,1], yticks=[], xticks=[], frame_on=False)
fig.add_axes(axobs)
fig.add_axes(axlit)
fig.add_axes(figAxes)
drawfigaxes(figAxes)
drawframe(axobs, axlit, 0)
axobs.arrow(-6,0,3,0, width=1, head_width=2, head_length=1, color=velocitycolor, ec='none')
axlit.arrow(-2,12,-3,0, width=1, head_width=2, head_length=1, color=velocitycolor, ec='none')
savefig('init.png')

figAxes.arrow(0.35,0.49,0.22,0, width=0.25, head_width=0.4, head_length=0.1, color='#c0e8cf', ec='none')
figAxes.text(0.5,0.5, 'Lorentz Transform', horizontalalignment='center')
figAxes.text(0.5,0.43, '(Beaming)', horizontalalignment='center')
savefig('init2.png')

figAxes.clear()
drawfigaxes(figAxes)
for t in range(21):
axobs.clear(); axlit.clear()
drawframe(axobs, axlit, t/20.0)
savefig('t{:02d}.png'.format(t))
axobs.arrow(2,0,3,0, width=1, head_width=2, head_length=1, color=velocitycolor, ec='none')
axlit.arrow(-10,12,-3,0, width=1, head_width=2, head_length=1, color=velocitycolor, ec='none')
figAxes.arrow(0.35,0.49,0.22,0, width=0.25, head_width=0.4, head_length=0.1, color='#c0e8cf', ec='none')
figAxes.text(0.5,0.5, 'Lorentz Transform', horizontalalignment='center')
figAxes.text(0.5,0.43, '(Aberration)', horizontalalignment='center')
savefig('end1.png')

axlit.add_patch(Circle((0,12), radius=1, color=lightcolor, ec='none', alpha=0.5))
axlit.text(0,8, 'Apparent\nposition', alpha=0.5, horizontalalignment='center')
savefig('end2.png')
drawframes()
#convert -delay 200 init.png -delay 300 init2.png -delay 10 t* -delay 100 t20.png -delay 200 end1.png -delay 1000 end2.png -loop 0 out.gif