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

from math import *
import numpy as np
import matplotlib.pyplot as plt
from matplotlib import animation

# settings
fname = 'gaussian-beam-animation'
size = 320, 200
nframes = 2 * 40
lambda_px = 30
angle0 = radians(14)
angle1 = radians(32)

angles = angle0+(angle1-angle0)*(1-np.abs(1-np.linspace(0, 2, nframes+1)[:-1]))

def animate(nframe):
    angle = angles[nframe]
    w0 = lambda_px / (pi * tan(angle))
    z0 = pi * w0**2 / lambda_px
    w = lambda z: w0 * sqrt(1 + (z / z0)**2)
    
    plt.clf()
    fig.gca().set_position((0, 0, 1, 1))
    plt.xlim(-size[0]/2., size[0]/2.)
    plt.ylim(-size[1]/2., size[1]/2.)
    plt.axis('off')
    
    # asymptotes
    edges = np.array([-size[0]/2, size[0]/2.])
    plt.plot(edges, w0/z0*edges, 'k', lw=0.5)
    plt.plot(edges, -w0/z0*edges, 'k', lw=0.5)
    
    # parameters
    plt.vlines(0, 0, w0, lw=1)
    plt.hlines(0, 0, z0, lw=1)
    plt.vlines(z0, 0, w0, lw=0.5)
    plt.hlines(w0, 0, z0, lw=0.5)
    plt.text(-1, 6, r'$w_0$', fontsize=20, horizontalalignment='right')
    plt.text(11, 3, r'$z_0$', fontsize=20, verticalalignment='top')
    
    # laser beam edges
    x = np.linspace(-size[0]/2., size[0]/2., int(size[0]+1))
    plt.plot(x, [w(z) for z in x], 'k', lw=4)
    plt.plot(x, [-w(z) for z in x], 'k', lw=4)
    
    # wavelength legend
    px, py = 56, -88
    plt.text(px, py + 4, r'$\lambda$', fontsize=20, horizontalalignment='center')
    plt.hlines(py, px-lambda_px/2., px+lambda_px/2.)
    plt.vlines(px-lambda_px/2., py-5, py+5)
    plt.vlines(px+lambda_px/2., py-5, py+5)

fig = plt.figure(figsize=(size[0]/100., size[1]/100.))
anim = animation.FuncAnimation(fig, animate, frames=nframes)
anim.save(fname + '.gif', writer='imagemagick', fps=20)