%matplotlib inline
import matplotlib as mpl
import matplotlib.pyplot as plt
from numpy import *

from IPython.display import HTML

from matplotlib import animation

from qutip import *

N = 15
chi = 1 * 2 * pi              # Kerr-nonlinearity
tlist = linspace(0, 1.0, 101) # time

# operators: the annihilation operator of the field
a = destroy(N)

# and we'll also need the following operators in calculation of
# expectation values when visualizing the dynamics 
n = num(N)
x = a + a.dag()
p = -1j * (a - a.dag())

# the Kerr Hamiltonian
H = 0.5 * chi * a.dag() * a.dag() * a * a

def plot_expect_with_variance(N, op_list, op_title, states):
    """
    Plot the expectation value of an operator (list of operators)
    with an envelope that describes the operators variance.
    """
    
    fig, axes = plt.subplots(1, len(op_list), figsize=(14,3))

    for idx, op in enumerate(op_list):
        
        e_op = expect(op, states)
        v_op = variance(op, states)

        axes[idx].fill_between(tlist, e_op - sqrt(v_op), e_op + sqrt(v_op), color="green", alpha=0.5);
        axes[idx].plot(tlist, e_op)
        axes[idx].set_xlabel('Time')
        axes[idx].set_title(op_title[idx])
        axes[idx].set_xlim(0, max(tlist))

    return fig, axes

def plot_wigner(rho, fig=None, ax=None):
    """
    Plot the Wigner function and the Fock state distribution given a density matrix for
    a harmonic oscillator mode.
    """
    
    if fig is None or ax is None:
        fig, ax = plt.subplots(1, 1, figsize=(8,8))

    if isket(rho):
        rho = ket2dm(rho)
    
    xvec = linspace(-7.5,7.5,200)

    W = wigner(rho, xvec, xvec)
    wlim = abs(W).max()

    ax.contourf(xvec, xvec, W, 100, norm=mpl.colors.Normalize(-wlim,wlim), cmap=mpl.cm.get_cmap('RdBu'))
    ax.set_xlabel(r'$x_1$', fontsize=16)
    ax.set_ylabel(r'$x_2$', fontsize=16)

    return fig, ax

def plot_fock_distribution_vs_time(tlist, states, fig=None, ax=None):
    
    Z = zeros((len(tlist), states[0].shape[0]))
    
    for state_idx, state in enumerate(states):
        Z[state_idx,:] = real(ket2dm(state).diag())
        
    if fig is None or axes is None:
        fig, ax = plt.subplots(1, 1, figsize=(8,6))

    Y, X = meshgrid(tlist, range(states[0].shape[0]))
    p = ax.pcolor(X, Y, Z.T, norm=mpl.colors.Normalize(0, 0.5), cmap=mpl.cm.get_cmap('Reds'), edgecolors='k')
    ax.set_xlabel(r'$N$', fontsize=16)
    ax.set_ylabel(r'$t$', fontsize=16)    
    
    cb = fig.colorbar(p)
    cb.set_label('Probability')
    
    return fig, ax

from base64 import b64encode

def display_embedded_video(filename):
    video = open(filename, "rb").read()
    video_encoded = b64encode(video).decode("ascii")
    video_tag = '<video controls alt="test" src="data:video/x-m4v;base64,{0}">'.format(video_encoded)
    return HTML(video_tag)

# we start with a coherent state with alpha=2.0
psi0 = coherent(N, 2.0)

# and evolve the state under the influence of the hamiltonian. 
# by passing an empty list as expecation value operators argument, 
# we get the full state of the system in result.states
result = mesolve(H, psi0, tlist, [], [])

plot_expect_with_variance(N, [n, x, p], [r'n', r'x', r'p'], result.states);

plot_fock_distribution_vs_time(tlist, result.states);

fig, ax = plt.subplots(1, 1, figsize=(8,8))

def update(n): 
    plot_wigner(result.states[n], fig=fig, ax=ax)

anim = animation.FuncAnimation(fig, update, frames=len(result.states), blit=True)

anim.save('animation-kerr-coherent-state.mp4', fps=10, writer="avconv", codec="libx264")
#anim.save('animation-kerr-coherent-state.gif', writer='imagemagick', fps=10)

plt.close(fig)

display_embedded_video("animation-kerr-coherent-state.mp4")

fig, ax = plt.subplots(1, 1, figsize=(8, 8))

def update(n): 
    plot_wigner(result.states[n], fig=fig, ax=ax)

anim = animation.FuncAnimation(fig, update, frames=int(len(result.states)/2+1), blit=True)

#anim.save('animation-kerr-coherent-state-half-period.mp4', fps=10, extra_args=['-vcodec', 'libx264'])
anim.save('animation-kerr-coherent-state-half-period.mp4', fps=10, writer="avconv", codec="libx264")

plt.close(fig)

display_embedded_video('animation-kerr-coherent-state-half-period.mp4')

psi = (exp(1j*pi/4) * coherent(N, -2.0j) + exp(-1j*pi/4) * coherent(N, 2.0j)).unit()

plot_wigner(psi);

from qutip.ipynbtools import version_table

version_table()