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

'''
Copyright (C) 2015 Geek3, Wikimedia Foundation

This program is free software; you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation; either version 2 of the License, or
(at your option) any later version.

This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU General Public License for more details.

You should have received a copy of the GNU General Public License
along with this program; if not, write to the Free Software
Foundation, Inc., 59 Temple Place - Suite 330, Boston, MA 02111-1307, USA.
'''

from math import *
import scipy
import matplotlib.pyplot as plt
from matplotlib import animation
from PIL import Image

# settings
fname = 'gaussianbeam_travelling_PRGn'
width, height = 240, 140
nframes = 16
fps=20

lambda_px = 16.
w0 = 16.

def E(r, z, w0):
    z0 = pi * w0**2 / lambda_px
    w = w0 * sqrt(1. + (z / z0)**2)
    k = 2. * pi / lambda_px
    R = z * (1. + (z0 / z)**2)
    gouy = atan(z / z0)
    E = w0 / w
    E *= e ** (-(r/w)**2)
    E *= e ** (-1j * k * r**2 / (2 * R))
    E *= e ** (1j * (gouy - k * z))
    return E

def draw(nframe):
    # draw an image
    image = scipy.zeros((height, width))
    
    phase = nframe * 2 * pi / nframes
    for y in range(height):
        for x in range(width):
            z = (x + 0.5) - width / 2.0
            r = (y + 0.5) - height / 2.0
            EE = E(r, z, w0)
            image[y, x] = (EE * e**(1j*phase)).real
    return image

def animate(nframe):
    # prepare a clean and image-filling canvas for each frame
    plt.clf()
    fig.gca().set_position((0, 0, 1, 1))
    plt.xlim(-0.5, width + 0.5)
    plt.ylim(-0.5, height + 0.5)
    plt.axis('off')
    
    print 'frame', nframe
    image = draw(nframe)
    
    imgplot = plt.imshow(image, cmap='PRGn')
    imgplot.set_clim(-1.0, 1.0)
    imgplot.set_interpolation('nearest')

fig = plt.figure(figsize=(width/100., height/100.))
anim = animation.FuncAnimation(fig, animate, frames=nframes, repeat=False)
anim.save(fname + '.gif', writer='imagemagick', fps=fps)