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

from math import *
import matplotlib.pyplot as plt
import numpy as np
from numpy.polynomial.hermite import Hermite

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

# image settings
fname = 'QHO-Fockstate0123'
width, height = 300, 300
ml, mr, mt, mb = 35, 8, 22, 45
x0, x1 = -4, 4
y0, y1 = 0.0, 0.7

def plot():
    for nfock in range(4):
        ax = axi[3-nfock]
        ax.cla()
        ax.axis((x0, x1, y0, y1))
        ax.grid(True)
        plt.yticks([0.0, 0.2, 0.4, 0.6], ['0.0', '0.2', '0.4', ''])
        if nfock != 0:
            ax.set_xticklabels([])
        ax.plot(0, 0, color=(1,1,1,0), label=r'$\vert{}\rangle$'.format(nfock))
        
        psi_fock = np.eye(1, nfock+1, nfock).flatten()
        # Definition of Fock-states in terms of Hermite functions:
        # https://en.wikipedia.org/wiki/Quantum_harmonic_oscillator
        a_hermite = [psi_fock[n] * pi**-0.25 / sqrt(2.**n*factorial(n))
                     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, 2 * width)
        psi_x = np.exp(-x**2 / 2.0) * H(x)
        y = np.abs(psi_x)**2
        
        ax.plot(x, y, lw=2, color='#0000cc')
        leg = ax.legend(handlelength=0, handletextpad=0, borderpad=0.1,
                  borderaxespad=0.45, loc='upper left', fontsize=17)
        leg.get_frame().set_linewidth(0.0)

plt.close('all')
fig, axi = plt.subplots(4, sharey=True, 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')

plot()
plt.savefig(fname + '.png')

# compress with optipng
from os import system
cmd = 'optipng -o9 "' + fname + '.png"'
if system(cmd) != 0:
    print 'warning: optipng not found!'