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

from math import *
import matplotlib.pyplot as plt
from matplotlib import animation, colors, colorbar
import numpy as np
from numpy.polynomial.hermite import Hermite
import colorsys
from scipy.interpolate import interp1d

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

# image settings
fname = 'QHO-groundstate-animation-color'
width, height = 300, 200
ml, mr, mt, mb, mh, mc = 35, 19, 22, 45, 12, 6
x0, x1 = -4, 4
y0, y1 = 0, 0.6
nframes = 3 * 5 * 7
fps = 20

# physics settings
nfock = 0
omega = 2*pi * 2

def color(phase):
    phase1 = ((phase / (2*pi)) % 1 + 1) % 1
    hue = (interp1d([0, 1./3, 1.2/3, 0.5, 1], # spread yellow a bit
                    [0, 1./3, 1.3/3, 0.5, 1])(phase1) + 2./3.) % 1
    light = interp1d([0, 1, 2, 3, 4, 5, 6], # adjust lightness
                     [0.64, 0.5, 0.56, 0.48, 0.75, 0.57, 0.64])(6 * hue)
    hls = (hue, light, 1.0) # maximum saturation
    rgb = colorsys.hls_to_rgb(*hls)
    return rgb

def animate(nframe):
    print str(nframe) + ' ',
    t = float(nframe) / nframes
    
    ax.cla()
    ax.axis((x0, x1, y0, y1))
    ax.grid(True)
    
    psi_fock = np.eye(1, nfock+1, nfock).flatten()
    # Definition of Fock-states in terms of Hermite functions:
    # https://en.wiki.x.io/wiki/Quantum_harmonic_oscillator
    a_hermite = [psi_fock[n] * pi**-0.25 / sqrt(2.**n*factorial(n))
                    * e**(-1j * omega * (n+0.5) * t) for n in range(1+nfock)]
    # doc: http://docs.scipy.org/doc/numpy/reference/generated/numpy.polynomial.hermite.Hermite.html
    H = Hermite(a_hermite)
    
    x = np.linspace(x0, x1, int(ceil(1+w_px)))
    x2 = x - px_w/2.
    psi_x = np.exp(-x**2 / 2.0) * H(x)
    phi_x = np.angle(np.exp(-(x2)**2 / 2.0) * H(x2))
    y = np.abs(psi_x)**2
    
    # plot color filling
    for x_, phi_, y_ in zip(x, phi_x, y):
        ax.plot([x_, x_], [0, y_], color=color(phi_), lw=2*0.72)
    
    ax.plot(x, y, lw=2, color='black')
    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+mc+mh)/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+mh+mc) # plot width in pixels
px_w = float(x1 - x0) / w_px # width of one pixel in plot units

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

# colorbar for phase
cax = fig.add_axes([1.0 - float(mr+mc)/width, float(mb)/height,
                    float(mc)/width, 1.0 - float(mb+mt)/height])
cax.yaxis.set_tick_params(length=2)
cmap = colors.ListedColormap([color(phase) for phase in
                              np.linspace(0, 2*pi, 384, endpoint=False)])
norm = colors.Normalize(0, 2*pi)
cbar = colorbar.ColorbarBase(cax, cmap=cmap, norm=norm,
                    orientation='vertical', ticks=np.linspace(0, 2*pi, 3))
cax.set_yticklabels(['$0$', r'$\pi$', r'$2\pi$'], rotation=90)
fig.text(1.0 - 10./width, 1.0, '$arg(\psi)$', ha='right', va='top')
plt.sca(ax)

# start animation
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 -k256 --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')