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

import os
import sys
import matplotlib.pyplot as plt
from matplotlib import ticker
from matplotlib import animation, colors, colorbar
import numpy as np
import colorsys
from scipy.interpolate import interp1d
from math import *

plt.rc('path', snap=False)
plt.rc('text', usetex=False)
plt.rc('mathtext', default='regular')
plt.rc('font', size=11)
plt.rcParams['font.sans-serif'] = ['DejaVu Sans']

# image settings
fname = 'Quantum-rigid-rotator_1+2-animation-color'
width, height = 300, 200
ml, mr, mt, mb, mh, mc = 35, 20, 22, 45, 12, 6
x0, x1 = -pi, pi
y0, y1 = 0.0, 0.4
nframes = 80
fps = 20
a = {1:1, 2:1}

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 qrr_wavefunc(a, phi, t):
    # Wavefunction of a quantum rigid rotator
    # https://dx.doi.org/10.1119/1.17340
    ai = [i for i in a.iterkeys()]
    s = sqrt(np.sum(np.abs([v for v in a.itervalues()])**2))
    av = [a[i] / s for i in ai]
    psi = 0j * phi
    for n, aa in zip(ai, av):
        omega = 2*pi * n**2
        psi += aa * np.exp(1j*(n*phi - omega*t)) / sqrt(2*pi)
    return psi

def animate(nframe):
    print str(nframe) + ' ', ; sys.stdout.flush()
    t = float(nframe) / nframes
    
    ax.cla()
    ax.grid(True)
    ax.axis((x0, x1, y0, y1))
    
    x = np.linspace(-pi, pi, int(ceil(1+w_px)))
    x2 = x - px_w/2.
    
    psi = qrr_wavefunc(a, x, t)
    psi2 = qrr_wavefunc(a, x2, t)
    y = np.abs(psi)**2
    phase = np.angle(psi2)
    
    # plot color filling
    for x_, phase_, y_ in zip(x, phase, y):
        ax.plot([x_, x_], [0, y_], color=color(phase_), lw=2*0.72)
    
    ax.plot(x, y, lw=2, color='black')
    ax.yaxis.set_ticks(np.arange(0, 0.4, 0.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, 5./height,
         r'$\phi$', 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 - 12./width, 1.0, '$arg(\psi)$', ha='right', va='top')
plt.sca(ax)

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