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

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

# settings
plt.rc('path', snap=False)
plt.rc('mathtext', default='regular')
fname = 'Cat-state-even-loss-animation'
width, height = 220, 128
dpi = 65.
nframes = 40
fps = 10
alpha0 = 2.
xlim = -4, 4
ylim = -2, 2
res = 40.

x_range = np.linspace(xlim[0], xlim[1], res*4+1)
y_range = np.linspace(ylim[0], ylim[1], res*2+1)
x_pixels = (x_range[1:] + x_range[:-1]) / 2.
y_pixels = (y_range[1:] + y_range[:-1]) / 2.
X, Y = np.meshgrid(x_pixels, y_pixels)

delays = [100 // fps for i in range(nframes)]
delays[0] = delays[20] = 60

plt.close('all')

def cat_state(x, y, alpha, theta, L):
    # Eq. (64), Phys. Rev. A 80, 032318 (2009)
    gaussm = np.exp(-(x + sqrt(2 * (1-L)) * alpha)**2 - y**2)
    gaussp = np.exp(-(x - sqrt(2 * (1-L)) * alpha)**2 - y**2)
    gauss0 = np.exp(-x**2 - y**2)
    fringe = np.cos(sqrt(8 * (1-L)) * alpha * y + theta)
    damp = exp(-2 * L * alpha**2)
    norm = 2 * pi * (1. + exp(-2 * alpha**2) * cos(theta))
    return (gaussm + gaussp + 2. * gauss0 * fringe * damp) / norm

def animate(nframe):
    print 'frame', nframe+1, '/', nframes
    t = nframe / float(nframes)
    plt.clf()
    
    if t < 0.5:
        alpha = alpha0
        L = t / 0.5
    else:
        alpha = alpha0 * (t - 0.5) / 0.5
        L = 0.
    
    img = cat_state(X, Y, alpha, 0., L)
    plt.imshow(img, cmap='RdBu', vmin=-1/pi, vmax=1/pi, origin='lower',
        extent=(xlim[0], xlim[1], ylim[0], ylim[1]))
    
    plt.title(r'$\alpha={:.1f}$   loss$={:.2f}$'.format(alpha, L))
    plt.gca().set_aspect('equal', adjustable='box')
    plt.gca().yaxis.set_major_locator(ticker.MultipleLocator(1))
    plt.subplots_adjust(left=0.1, right=1-0.06, bottom=0.1)

fig = plt.figure(figsize=(width/float(dpi), height/float(dpi)))

# start animation
if 0 != os.system('convert -version > ' +  os.devnull):
    print 'imagemagick not installed!'
    # warning: imagemagick produces dithered and therefore large gifs
    anim = animation.FuncAnimation(fig, animate, frames=nframes)
    anim.save(fname + '.gif', writer='imagemagick', fps=fps, dpi=dpi)
else:
    # unfortunately the matplotlib imagemagick backend does not support
    # options which are necessary to generate high quality output without
    # framewise color palettes. Therefore save all frames and convert then.
    if not os.path.isdir(fname):
        os.mkdir(fname)
    fnames = []
    
    for frame in range(nframes):
        animate(frame)
        imgname = os.path.join(fname, fname + '{:03d}'.format(frame) + '.png')
        fig.savefig(imgname, dpi=dpi)
        fnames.append(imgname)
    
    # compile optimized animation with ImageMagick
    cmd = 'convert -loop 0 '
    nameslist = ['-delay ' + str(delays[i]) + ' ' + n for i,n in enumerate(fnames)]
    cmd += ' '.join(nameslist)
    cmd += ' +dither -colors 255 '
    cmd += fname + '.gif'
    os.system(cmd)
    
    for fnamei in fnames:
        os.remove(fnamei)
    os.rmdir(fname)