%matplotlib inline
import matplotlib.pyplot as plt
import numpy as np

from qutip import *

N = 20

def plot_wigner_2d_3d(psi):
    #fig, axes = plt.subplots(1, 2, subplot_kw={'projection': '3d'}, figsize=(12, 6))
    fig = plt.figure(figsize=(17, 8))
    
    ax = fig.add_subplot(1, 2, 1)
    plot_wigner(psi, fig=fig, ax=ax, alpha_max=6);

    ax = fig.add_subplot(1, 2, 2, projection='3d')
    plot_wigner(psi, fig=fig, ax=ax, projection='3d', alpha_max=6);
    
    plt.close(fig)
    return fig

psi = basis(N, 0)
plot_wigner_2d_3d(psi)

psi = thermal_dm(N, 2)
plot_wigner_2d_3d(psi)

psi = coherent(N, 2.0)
plot_wigner_2d_3d(psi)

psi = coherent(N, -1.0)
plot_wigner_2d_3d(psi)

psi = (coherent(N, -2.0) + coherent(N, 2.0)) / np.sqrt(2)
plot_wigner_2d_3d(psi)

psi = (coherent(N, -2.0) - coherent(N, 2.0)) / np.sqrt(2)
plot_wigner_2d_3d(psi)

psi = (coherent(N, -2.0) + coherent(N, -2j) + coherent(N, 2j) + coherent(N, 2.0)).unit()
plot_wigner_2d_3d(psi)

psi = (coherent(N, -2.0) + coherent(N, -1j) + coherent(N, 1j) + coherent(N, 2.0)).unit()
plot_wigner_2d_3d(psi)

NN = 8

fig, axes = plt.subplots(NN, 1, figsize=(5, 5 * NN), sharex=True, sharey=True) 
for n in range(NN):
    psi = sum([coherent(N, 2*np.exp(2j * np.pi * m / (n + 2))) for m in range(n + 2)]).unit()
    plot_wigner(psi, fig=fig, ax=axes[n])

    #if n < NN - 1:
    #    axes[n].set_ylabel("")

psi = (coherent_dm(N, -2.0) + coherent_dm(N, 2.0)) / np.sqrt(2)
plot_wigner_2d_3d(psi)

from IPython.display import display

for n in range(6):
    psi = basis(N, n)
    display(plot_wigner_2d_3d(psi))

NN = MM = 5

fig, axes = plt.subplots(NN, MM, figsize=(18, 18), sharex=True, sharey=True) 
for n in range(NN):
    for m in range(MM):
        psi = (fock(N, n) + fock(N, m)).unit()
        plot_wigner(psi, fig=fig, ax=axes[n, m])
        #axes[n, m].set_title(r"$(\left|%d\right> + \left|%d\right>)/\sqrt{2}$" % (n, m))
        if n < NN - 1:
            axes[n, m].set_xlabel("")
        if m > 0:
            axes[n, m].set_ylabel("")

psi = squeeze(N, 0.5) * basis(N, 0)
display(plot_wigner_2d_3d(psi))

psi = squeeze(N, 0.75j) * basis(N, 0)
display(plot_wigner_2d_3d(psi))

psi = squeeze(N, -1) * basis(N, 0)
display(plot_wigner_2d_3d(psi))

psi = (squeeze(N, 0.75j) * basis(N, 0) - squeeze(N, -0.75j) * basis(N, 0)).unit()
display(plot_wigner_2d_3d(psi))

psi = (ket2dm(squeeze(N, 0.75j) * basis(N, 0)) + ket2dm(squeeze(N, -0.75j) * basis(N, 0))).unit()
display(plot_wigner_2d_3d(psi))

psi = displace(N, 2) * squeeze(N, 0.75) * basis(N, 0)
display(plot_wigner_2d_3d(psi))

psi = (displace(N, -1) * squeeze(N, 0.75) * basis(N, 0) - displace(N, 1) * squeeze(N, -0.75) * basis(N, 0)).unit()
display(plot_wigner_2d_3d(psi))

from qutip.ipynbtools import version_table

version_table()