#!/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-Fockstate3'
width, height = 300, 200
ml, mr, mt, mb = 35, 8, 22, 45
x0, x1 = -4, 4
y0, y1 = 0.0, 0.6

# physics settings
nfock = 3

def plot():
    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.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
    
    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')

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

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