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

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

# image settings
fname = 'QHO-coherent3-squeezed2dB-animation'
plt.rc('path', snap=False)
plt.rc('mathtext', default='regular')
width, height = 300, 200
ml, mr, mt, mb = 35, 8, 22, 45
x0, x1 = -8,8
y0, y1 = 0.0, 1.0
nframes = 120
fps = 20

# physics settings
alpha0 = 3.0
xi0 = 0.2 * log(10) * e**(-0.5j*pi) # 2dB of squeezing

def squeezed_coherent(alpha0, xi0, x, omega_t):
    # Definition of coherent states
    # https://en.wikipedia.org/wiki/Coherent_states
    alpha = alpha0 * e**(-1j * omega_t)
    xi = xi0 * e**(-2j * omega_t)
    r = np.abs(xi)
    tr = tanh(r)
    kk = (r - tr * xi) / (r + tr * xi)
    psi = (kk.real/pi)**0.25 * np.exp(-0.5j * omega_t
                - 0.5 * ((x - sqrt(2) * alpha.real))**2 * kk
                - 1j * alpha.imag * (alpha.real - sqrt(2) * x))
    return psi

def animate(nframe):
    print str(nframe) + ' ',; sys.stdout.flush()
    t = float(nframe) / nframes * 1.0 # animation repeats after t=1.0
    omega = 2 * pi
    
    ax.cla()
    ax.grid(True)
    ax.axis((x0, x1, y0, y1))
    
    x = np.linspace(x0, x1, int(ceil(1+w_px)))
    psi = squeezed_coherent(alpha0, xi0, x, omega*t)
    y = np.abs(psi)**2
    
    plt.plot(x, y, lw=2, color='#0000cc')
    ax.set_yticks(ax.get_yticks()[:-1])
    ax.set_yticklabels([l for l in ax.get_yticks() if l < y0+0.9*(y1-y0)])
    

# 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)
w_px = width - (ml+mr) # plot width in pixels

# 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')

# start animation
if 0 != os.system('convert -version > ' +  os.devnull):
    print 'imagemagick not installed!'
    # warning: imagemagick produces somewhat jagged and therefore large gifs
    anim = animation.FuncAnimation(fig, animate, frames=nframes)
    anim.save(fname + '.gif', writer='imagemagick', fps=fps)
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)
        fnames.append(imgname)
    
    # compile optimized animation with ImageMagick
    cmd = 'convert -loop 0 -delay ' + str(100 / fps) + ' '
    cmd += ' '.join(fnames) # now create optimized palette from all frames
    cmd += r' \( -clone 0--1 \( -clone 0--1 -fill black -colorize 100% \) '
    cmd += '-append +dither -colors 63 -unique-colors '
    cmd += '-write mpr:colormap +delete \) +dither -map mpr:colormap '
    cmd += '-alpha activate -layers OptimizeTransparency '
    cmd += fname + '.gif'
    os.system(cmd)
    
    for fnamei in fnames:
        os.remove(fnamei)
    os.rmdir(fname)