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

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

plt.rc('path', snap=False)
plt.rc('mathtext', default='regular')

# image settings
fname = 'QHO-catstate-even3-animation'
width, height = 300, 200
ml, mr, mt, mb = 35, 8, 22, 45
x0, x1 = -6.5, 6.5
y0, y1 = 0, 1.2
nframes = 150
fps = 20

# physics settings
alpha0 = 3.0
omega = 2*pi

def coherent(alpha, x, omega, t, l=1.0):
    # Definition of coherent states
    # https://en.wikipedia.org/wiki/Coherent_states
    psi = (pi*l**2)**-0.25 * np.exp(
                -0.5/l**2 * (x - sqrt(2)*l * alpha.real)**2
                + 1j*sqrt(2)/l * alpha.imag * x
                + 0.5j * (alpha0**2*sin(2*omega*t) - omega*t))
    return psi

def animate(nframe):
    print str(nframe) + ' ',
    t = float(nframe) / nframes * 0.5 # animation repeats after t=0.5
    alpha = e ** (-1j * omega * t) * alpha0
    
    ax.cla()
    ax.axis((x0, x1, y0, y1))
    ax.grid(True)
    
    x = np.linspace(x0, x1, 2 * width)
    
    # Definition of cat states in terms of coherent states:
    # https://en.wikipedia.org/wiki/Cat_state
    psi = coherent(alpha, x, omega, t) + coherent(-alpha, x, omega, t)
    psi /= sqrt(2 * (1 + exp(-2*alpha0**2)))
    y = np.abs(psi)**2
    
    plt.plot(x, y, lw=2, color='#0000cc')
    ax.set_yticks(ax.get_yticks()[:-1])

# create figure and axes
plt.close('all')
fig, ax = plt.subplots(1, figsize=(width/100., height/100.))
bounds = [float(ml)/width, float(mb)/height,
          1.0 - float(mr)/width, 1.0 - float(mt)/height]
fig.subplots_adjust(left=bounds[0], bottom=bounds[1],
                    right=bounds[2], top=bounds[3], hspace=0)

# axes labels
fig.text(0.5 + 0.5 * float(ml-mr)/width, 4./height,
         r'$x\ \ [(\hbar/(m\omega))^{1/2}]$', ha='center')
fig.text(5./width, 1.0, '$|\psi|^2$', va='top')

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

import os
# compress with gifsicle
commons = 'https://commons.wikimedia.org/wiki/File:'
cmd = 'gifsicle -O3 -k64 --careful --comment="' + commons + fname + '.gif"'
cmd += ' < ' + fname + '_.gif > ' + fname + '.gif'
if os.system(cmd) == 0:
    os.remove(fname + '_.gif')
else:
    print 'warning: gifsicle not found!'
    os.remove(fname + '.gif')
    os.rename(fname + '_.gif', fname + '.gif')