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

from qutip import *

chi = 0.2
N1 = 75
N2 = 75

a  = tensor(destroy(N1), qeye(N2))
na = tensor(num(N1),     qeye(N2)) 
b  = tensor(qeye(N1),    destroy(N2))
nb = tensor(qeye(N1),    num(N2)) 

H = - chi * (a * b + a.dag() * b.dag())

# start in the ground (vacuum) state
psi0 = tensor(basis(N1,0), basis(N2,0))

tlist = np.linspace(0, 10, 100)

c_ops = []

e_ops = []

output = mesolve(H, psi0, tlist, c_ops, e_ops)
output

na_e = zeros(shape(tlist))
na_s = zeros(shape(tlist))
nb_e = zeros(shape(tlist))
nb_s = zeros(shape(tlist))

for idx, psi in enumerate(output.states):
    na_e[idx] = expect(na, psi)
    na_s[idx] = expect(na*na, psi)
    nb_e[idx] = expect(nb, psi)
    nb_s[idx] = expect(nb*nb, psi)

# substract the average squared to obtain variances
na_s = na_s - na_e ** 2
nb_s = nb_s - nb_e ** 2

fig, axes = plt.subplots(2, 2, sharex=True, figsize=(8,5))

line1 = axes[0,0].plot(tlist, na_e, 'r', linewidth=2)
axes[0,0].set_ylabel(r'$\langle a^\dagger a \rangle$', fontsize=18)

line2 = axes[0,1].plot(tlist, nb_e, 'b', linewidth=2)

line3 = axes[1,0].plot(tlist, na_s, 'r', linewidth=2)
axes[1,0].set_xlabel('$t$', fontsize=18)
axes[1,0].set_ylabel(r'$Std[a^\dagger a]$, $Std[b^\dagger b]$', fontsize=18)

line4 = axes[1,1].plot(tlist, nb_s, 'b', linewidth=2)
axes[1,1].set_xlabel('$t$', fontsize=18)

fig.tight_layout()


# pick an arbitrary time and calculate the wigner functions for each mode
xvec = np.linspace(-5,5,200)
t_idx_vec = [0, 10, 20, 30]

fig, axes = plt.subplots(len(t_idx_vec), 2, sharex=True, sharey=True, figsize=(8,4*len(t_idx_vec)))

for idx, t_idx in enumerate(t_idx_vec):
    psi_a = ptrace(output.states[t_idx], 0)
    psi_b = ptrace(output.states[t_idx], 1)
    W_a = wigner(psi_a, xvec, xvec)
    W_b = wigner(psi_b, xvec, xvec)
    
    cont1 = axes[idx,0].contourf(xvec, xvec, W_a, 100)
    cont2 = axes[idx,1].contourf(xvec, xvec, W_b, 100)


# pick arbitrary times and plot the photon distributions at those times
#t_idx_vec = [0, 10, 20, 30]
t_idx_vec = range(0,len(tlist),25) 

fig, axes = plt.subplots(len(t_idx_vec), 2, sharex=True, sharey=True, figsize=(8,2*len(t_idx_vec)))

for idx, t_idx in enumerate(t_idx_vec):
    psi_a = ptrace(output.states[t_idx], 0)
    psi_b = ptrace(output.states[t_idx], 1)
    
    cont1 = axes[idx,0].bar(range(0, N1), real(psi_a.diag()))
    cont2 = axes[idx,1].bar(range(0, N2), real(psi_b.diag()))
    
fig.tight_layout()

# second-order photon correlations
g2_1  = zeros(shape(tlist))
g2_2  = zeros(shape(tlist))
g2_12 = zeros(shape(tlist))

ad_ad_a_a = a.dag() * a.dag() * a * a
bd_bd_b_b = b.dag() * b.dag() * b * b
ad_a_bd_b = a.dag() * a * b.dag() * b

cs_rhs = zeros(shape(tlist))
cs_lhs = zeros(shape(tlist))

for idx, psi in enumerate(output.states):
    # g2 correlations
    g2_1[idx]  = expect(ad_ad_a_a, psi)
    g2_2[idx]  = expect(bd_bd_b_b, psi)
    g2_12[idx] = expect(ad_a_bd_b, psi)
    
    # cauchy-schwarz 
    cs_lhs[idx] = expect(ad_a_bd_b, psi)
    cs_rhs[idx] = expect(ad_ad_a_a, psi)
    
# normalize the correlation functions
g2_1  = g2_1  / (na_e ** 2)
g2_2  = g2_2  / (nb_e ** 2)
g2_12 = g2_12 / (na_e * nb_e)

fig, axes = plt.subplots(2, 2, figsize=(8,5))

line1 = axes[0,0].plot(tlist, g2_1, 'r', linewidth=2)
axes[0,0].set_xlabel("$t$", fontsize=18)
axes[0,0].set_ylabel(r'$g_1^{(2)}(t)$', fontsize=18)
axes[0,0].set_ylim(0,3)

line2 = axes[0,1].plot(tlist, g2_2, 'b', linewidth=2)
axes[0,1].set_xlabel("$t$", fontsize=18)
axes[0,1].set_ylabel(r'$g_2^{(2)}(t)$', fontsize=18)
axes[0,1].set_ylim(0,3)

line3 = axes[1,0].plot(tlist[10:], g2_12[10:], 'b', linewidth=2)
axes[1,0].set_xlabel("$t$", fontsize=18)
axes[1,0].set_ylabel(r'$g_{12}^{(2)}(t)$', fontsize=18)


line4 = axes[1,1].plot(tlist[20:], abs(g2_12[20:])**2, 'b', linewidth=2)
line5 = axes[1,1].plot(tlist, g2_1 * g2_2, 'r', linewidth=2)
axes[1,1].set_xlabel("$t$", fontsize=18)
axes[1,1].set_ylabel(r'$|g_{12}^{(2)}(t)|^2$', fontsize=18)

fig.tight_layout()

fig, axes = plt.subplots(1, 2, figsize=(10,4))

line1 = axes[0].plot(tlist, cs_lhs, 'b', tlist, cs_rhs, 'r', linewidth=2)
axes[0].set_xlabel("$t$", fontsize=18)
axes[0].set_title(r'Cauchy-Schwarz inequality', fontsize=18)

line1 = axes[1].plot(tlist, cs_lhs / (cs_rhs), 'k', linewidth=2)
axes[1].set_xlabel("$t$", fontsize=18)
axes[1].set_title(r'Cauchy-Schwarz ratio inequality', fontsize=18)

fig.tight_layout()

# pre-compute operators outside the loop
op_a_b = a * b
op_ad_bd = a.dag() * b.dag()
op_ad_a_p_a_ad  = a.dag() * a + a * a.dag()
op_bd_b_p_b_bd  = b.dag() * b + b * b.dag()

e_a_b   = np.zeros(shape(tlist), dtype=complex)
e_ad_bd = np.zeros(shape(tlist), dtype=complex)
e_ad_a_p_a_ad = np.zeros(shape(tlist), dtype=complex)
e_bd_b_p_b_bd = np.zeros(shape(tlist), dtype=complex)

for idx, psi in enumerate(output.states):
    
    e_a_b[idx]         = expect(op_a_b, psi)
    e_ad_bd[idx]       = expect(op_ad_bd, psi)
    e_ad_a_p_a_ad[idx] = expect(op_ad_a_p_a_ad, psi)
    e_bd_b_p_b_bd[idx] = expect(op_bd_b_p_b_bd, psi)
    
# calculate the sigma_2
theta = 3*pi/4
w_a = w_b = 1
sigma2 = 2 * sqrt(w_a * w_b) * (e_a_b * exp(2j * theta) + e_ad_bd * exp(-2j * theta)) / (w_a * e_ad_a_p_a_ad + w_b * e_bd_b_p_b_bd)

fig, axes = plt.subplots(1, 1, figsize=(8,4))

line1 = axes.plot(tlist, real(sigma2), 'b', tlist, imag(sigma2), 'r--', linewidth=2)
axes.set_xlabel("$t$", fontsize=18)
axes.set_ylabel(r'$\sigma_2(t)$', fontsize=18)
axes.set_ylim(-1, 1)
fig.tight_layout()

# pre-compute operators outside the loop
op_ad_a  = a.dag() * a
op_bd_b  = b.dag() * b
op_a_b   = a * b
op_ad_bd = a.dag() * b.dag()

e_ad_a  = np.zeros(shape(tlist), dtype=complex)
e_bd_b  = np.zeros(shape(tlist), dtype=complex)
e_a_b   = np.zeros(shape(tlist), dtype=complex)
e_ad_bd = np.zeros(shape(tlist), dtype=complex)

for idx, psi in enumerate(output.states):
    
    e_ad_a[idx]  = expect(op_ad_a, psi)
    e_bd_b[idx]  = expect(op_bd_b, psi)
    e_a_b[idx]   = expect(op_a_b, psi)
    e_ad_bd[idx] = expect(op_ad_bd, psi)
    
# calculate the sigma_2, function of the angle parameter theta
def F_theta(theta):
    return 2 * e_ad_a + 2 * e_bd_b + 2j * (exp(2j * theta) * e_a_b - exp(-2j * theta) * e_ad_bd)

fig, axes = plt.subplots(1, 1, figsize=(8,3))

for theta in np.linspace(0.0, 2*pi, 100):
    line1 = axes.plot(tlist, real(F_theta(theta)), 'b', tlist, imag(F_theta(theta)), 'g--', linewidth=2)

line = axes.plot(tlist, real(F_theta(0)), 'r', linewidth=4)
    
axes.set_xlabel("$t$", fontsize=18)
axes.set_ylabel(r'$\langle:f_\theta^\dagger f_\theta:\rangle$', fontsize=18)
axes.set_ylim(-2, 5)
fig.tight_layout()

R_op = correlation_matrix_quadrature(a, b)

def plot_covariance_matrix(V, ax):
    """
    Plot a matrix-histogram representation of the supplied Wigner covariance matrix.
    """
    num_elem = 16
    xpos,ypos = meshgrid(range(4),range(4))
    xpos = xpos.T.flatten()-0.5 
    ypos = ypos.T.flatten()-0.5 
    zpos = zeros(num_elem)   
    dx = 0.75 * np.ones(num_elem) 
    dy = dx.copy()           
    dz = V.flatten()

    nrm = mpl.colors.Normalize(-0.5,0.5) 
    colors = cm.jet(nrm((np.sign(dz)*abs(dz)**0.75))) 
    
    ax.view_init(azim=-40,elev=60)
    ax.bar3d(xpos, ypos, zpos, dx, dy, dz, color=colors)
    ax.axes.w_xaxis.set_major_locator(plt.IndexLocator(1,-0.5))
    ax.axes.w_yaxis.set_major_locator(plt.IndexLocator(1,-0.5))
    ax.axes.w_xaxis.set_ticklabels(("$q_-$", "$p_-$", "$q_+$", "$p_+$"), fontsize=20)
    ax.axes.w_yaxis.set_ticklabels(("$q_-$", "$p_-$", "$q_+$", "$p_+$"), fontsize=20)

# pick arbitrary times and plot the photon distributions at those times
t_idx_vec = [0, 20, 40]

fig, axes = plt.subplots(len(t_idx_vec), 1, subplot_kw={'projection':'3d'}, figsize=(6,3*len(t_idx_vec)))

for idx, t_idx in enumerate(t_idx_vec):
    
    # calculate the wigner covariance matrix
    V = wigner_covariance_matrix(R=R_op, rho=output.states[idx])

    plot_covariance_matrix(V, axes[idx])
    
fig.tight_layout()

"""
Calculate the wigner covariance matrix logarithmic negativity for each time step
"""
logneg = np.zeros(shape(tlist))

for idx, t_idx in enumerate(tlist):
    
    V = wigner_covariance_matrix(R=R_op, rho=output.states[idx])
    
    logneg[idx] = logarithmic_negativity(V)
    
fig, axes = plt.subplots(1, 1, figsize=(8,4))
axes.plot(tlist, logneg, 'r')
axes.set_xlabel("$t$", fontsize=18)
axes.set_ylabel("Logarithmic negativity", fontsize=18)
fig.tight_layout()

from qutip.ipynbtools import version_table

version_table()