%matplotlib inline

from qutip import *

from belltests import *

from matplotlib import rcParams
rcParams["font.family"] = "STIXGeneral"
rcParams["mathtext.fontset"] = "stix"

from functools import partial

# number of data points on main axis in plots
Nres = 100
calculate = True

# Hilbert space dimention (signal and idler mode)
N1 = 10
N2 = 10

Delta_0 = 0.0
gamma_1 = gamma_2 = 0.001
gamma_0 = 1.0

n_1_th = n_2_th = 0.0

kappa = 0.15
r = 1.12
E = r**2 * kappa / 2

phi = pi / 4
t_unit = 1/(kappa**2 / gamma_0)
tlist = np.linspace(0, 6*t_unit, 100)
t_idx_vec = [0, 30, 40, 50, 70, 99]

x1_vec = x2_vec = linspace(-5, 5, 50)

# simplified chsh
#bell_chsh_func = partial(bell_quadrature_chsh_simplified, x1_vec, x2_vec, phi)

# full chsh
bell_chsh_func = partial(bell_quadrature_chsh, x1_vec, x2_vec, -2*phi, 3*phi, 0, phi)

a1 = tensor(destroy(N1), qeye(N2))
a2 = tensor(qeye(N1), destroy(N2))

def dynamics_solve(kappa, E, gamma_0, gamma_1, gamma_2, n_th):
    
    n_1_th = n_2_th = n_th    
    
    gamma = kappa**2 * (gamma_0/2) / (abs(gamma_0/2 + 1j * Delta_0)**2)
    chi = - kappa**2 * Delta_0 / (abs(gamma_0/2 + 1j * Delta_0)**2)
    mu = E * kappa / (gamma_0 / 2 + 1j * Delta_0) 
    
    # 2 mode PA Hamiltonian
    H = 1j * (- a1 * a2 * conjugate(mu) + a1.dag() * a2.dag() * mu) + chi * a1.dag() * a2.dag() * a1 * a2

    # start in a thermal state
    psi = tensor([thermal_dm(N1, n_1_th), thermal_dm(N2, n_2_th)])  
    
    # setup collapse operators. only include those with nonzero rates (depends on n_*_th and kappa_*)
    c_ops_raw = [[sqrt(gamma_1 * (1+n_1_th)), a1], 
                 [sqrt(gamma_1 * n_1_th)    , a1.dag()], 
                 [sqrt(gamma_2 * (1+n_2_th)), a2], 
                 [sqrt(gamma_2 * n_2_th)    , a2.dag()],
                 [sqrt(gamma), a1 * a2]]
    
    c_ops = []
    for rate, op in c_ops_raw:
        if rate > 0.0:
            c_ops.append(rate * op)
        
    e_ops = []
        
    result = mesolve(H, psi, tlist, c_ops, e_ops)
    
    return result

rcParams["font.size"] = "18"

def plot_bell_transient(result0, result1):
    
    bell_chsh_vec = [bell_quadrature_chsh(x1_vec, x2_vec, -2*phi, 3*phi, 0, phi, result1.states[t_idx])
                     for t_idx, t in enumerate(tlist)]
    bell_ch_vec   = [bell_quadrature_ch(x1_vec, x2_vec, -2*phi, 3*phi, 0, phi, result1.states[t_idx])
                     for t_idx, t in enumerate(tlist)]


    bell_chsh_vec0 = [bell_quadrature_chsh(x1_vec, x2_vec, -2*phi, 3*phi, 0, phi, result0.states[t_idx])
                      for t_idx, t in enumerate(tlist)]
    bell_ch_vec0   = [bell_quadrature_ch(x1_vec, x2_vec, -2*phi, 3*phi, 0, phi, result0.states[t_idx])
                      for t_idx, t in enumerate(tlist)]
    
    fig, ax = subplots(1, 1, figsize=(8, 5))

    ax.plot(tlist/t_unit, array(bell_ch_vec)   / 1.0, 'b', label=r"$B_{\rm CH}$", linewidth=2)
    ax.plot(tlist/t_unit, array(bell_chsh_vec) / 2.0, 'r', label=r"$B_{\rm CHSH}/2$", linewidth=2)
    ax.plot(tlist/t_unit, array(bell_ch_vec0)   / 1.0, 'b--', linewidth=1)
    ax.plot(tlist/t_unit, array(bell_chsh_vec0) / 2.0, 'r--', linewidth=1)
    ax.plot(tlist/t_unit, ones(shape(tlist)), 'k', linewidth=2)
    
    ax.set_xlabel(r'time $t\kappa^2/\gamma_0$', fontsize=18)
    ax.legend(loc=4)
    
    ax.set_yticks([0.9, 1, 1.1])
    ax.set_ylim(0.9, 1.1)
    fig.tight_layout()
    
    return fig, axes

result0 = dynamics_solve(kappa, E, gamma_0, 0, 0, 0)  # no single-mode dissipation: gamma_1 = gamma_2 = 0

result1 = dynamics_solve(kappa, E, gamma_0, gamma_1, gamma_2, 0)

fig, axes = plot_bell_transient(result0, result1)

fig.savefig("fig-5-a.png")
fig.savefig("fig-5-a.pdf")

def plot_bell_angle(result0, result1):
    
    t_idx = 30  # approx t/t_unit = 1.8
    phi_vec = linspace(0.0, pi, Nres)

    psi = result1.states[t_idx]

    bell_chsh_vec1 = [bell_quadrature_chsh(x1_vec, x2_vec, -2*phi, 3*phi, 0, phi, psi) for phi in phi_vec]
    bell_ch_vec1   = [bell_quadrature_ch(x1_vec, x2_vec, -2*phi, 3*phi, 0, phi, psi)   for phi in phi_vec]

    psi = result0.states[-1]

    bell_chsh_vec0 = [bell_quadrature_chsh(x1_vec, x2_vec, -2*phi, 3*phi, 0, phi, psi) for phi in phi_vec]
    bell_ch_vec0   = [bell_quadrature_ch(x1_vec, x2_vec, -2*phi, 3*phi, 0, phi, psi)   for phi in phi_vec]
    
    fig, ax = subplots(1, 1, figsize=(8, 5))

    ax.plot(phi_vec/(pi), abs(array(bell_ch_vec1))   / 1.0, 'b', linewidth=2, label=r"$B_{\rm CH}$")
    ax.plot(phi_vec/(pi), abs(array(bell_chsh_vec1)) / 2.0, 'r', linewidth=2, label=r"$B_{\rm CHSH}/2$")
    ax.plot(phi_vec/(pi), abs(array(bell_ch_vec0))   / 1.0, 'b--', linewidth=1)
    ax.plot(phi_vec/(pi), abs(array(bell_chsh_vec0)) / 2.0, 'r--', linewidth=1)
    ax.plot(phi_vec/(pi), ones(shape(phi_vec)), 'k', linewidth=2)

    ax.set_xlabel(r"Bell angle $\varphi$", fontsize=18)
    ax.legend(loc=4)    

    ax.set_yticks([0.6, 0.8, 1])
    ax.set_xticks([0, 0.25, 0.5, 0.75, 1])
    ax.set_xticklabels([r'$0$', r'$\pi/4$', r'$\pi/2$', r'$3\pi/4$', r'$\pi$'])

    ax.set_ylim(0.6, 1.05)

    fig.tight_layout()
    
    return fig, ax

fig, axes = plot_bell_angle(result0, result1)

fig.savefig("fig-5-b.png")
fig.savefig("fig-5-b.pdf")

rcParams["font.size"] = "24"

def bell_transient(kappa_vec, gamma_1, gamma_2):
    
    bell_chsh_mat = zeros((len(kappa_vec), len(tlist)))
    for idx, kappa in enumerate(kappa_vec):
        r = dynamics_solve(kappa, E, gamma_0, gamma_1, gamma_2, 0)
        bell_chsh_mat[idx,:] = array([bell_chsh_func(r.states[t_idx]) for t_idx, t in enumerate(r.times)])/2

    return bell_chsh_mat

def plot_bell_transient(kappa_vec, data):
    
    fig, ax = subplots(1, 1, figsize=(8, 6))

    kappa_opt = 2*E / 1.12**2
    
    X, Y = meshgrid(tlist/t_unit, kappa_vec/kappa_opt)
    Z = data
    Z[Z < 1] = 1.0

    norm = mpl.colors.Normalize(1, 1.04) 
    cmap = get_cmap('Reds', lut=250)
    cmap.set_under("white")
    
    p = ax.pcolor(X, Y, Z-0.0001, cmap=cmap, norm=norm)
    ax.plot([0,max(tlist)/t_unit], [1,1], 'k--', lw=2)
    cb = fig.colorbar(p)

    ax.set_ylabel(r"$\kappa/\kappa_{\rm opt}$", fontsize=26, labelpad=0)
    ax.set_xlabel(r"$t{\bar{\kappa}}^2/\bar{\gamma}_0$", fontsize=26, labelpad=-3)
    cb.set_ticks([1.0, 1.01, 1.02, 1.03, 1.04])
    ax.axis('tight')

    ax.set_yticks([0.5, 1.0, 1.5])
    ax.set_yticklabels(["0.5", "1", "1.5"])
    ax.text(6.3, 1.77, r"$B_{\rm CHSH}/2$", fontsize=24)
    
    return fig, axes

kappa_vec = linspace(0.025, 0.25, Nres)

if calculate:
    data = bell_transient(kappa_vec, 0, 0)  # no single-mode dissipation: gamma_1 = gamma_2 = 0
    save("data-fig-6-a.npy", data)
else:
    data = load("data-fig-6-a.npy")

fig, axes = plot_bell_transient(kappa_vec, data);

fig.savefig("fig-6-a.png", dpi=200)
fig.savefig("fig-6-a.pdf")

if calculate:
    data = bell_transient(kappa_vec, gamma_1, gamma_2);
    save("data-fig-6-d.npy", data)
else:
    data = load("data-fig-6-d.npy")

fig, axes = plot_bell_transient(kappa_vec, data);

fig.savefig("fig-6-d.png", dpi=200)
fig.savefig("fig-6-d.pdf")

def bell_transient(E_vec, gamma_1, gamma_2):
    
    bell_chsh_mat = zeros((len(E_vec), len(tlist)))
    for idx, E in enumerate(E_vec):
        r = dynamics_solve(kappa, E, gamma_0, gamma_1, gamma_2, 0)
        bell_chsh_mat[idx,:] = array([bell_chsh_func(r.states[t_idx]) for t_idx, t in enumerate(r.times)])/2
        
    return bell_chsh_mat

def plot_bell_transient(E_vec, data):
    
    fig, ax = subplots(1, 1, figsize=(8, 6))

    E_opt = r**2 * kappa / 2

    X, Y = meshgrid(tlist/t_unit, E_vec/E_opt)
    Z = data
    Z[Z < 1] = 1.0

    norm=mpl.colors.Normalize(1, 1.04) 
    cmap=get_cmap('Reds', lut=250)
    cmap.set_under("white")
    
    p = ax.pcolor(X, Y, Z-0.0001, cmap=cmap, norm=norm)
    ax.plot([0,max(tlist)/t_unit], [1,1], 'k--', lw=2)
    cb = fig.colorbar(p)
    
    ax.set_ylabel(r"$E/E_{\rm opt}$", fontsize=26, labelpad=0)
    ax.set_xlabel(r"$t\bar{\kappa}^2/\bar{\gamma}_0$", fontsize=26, labelpad=-3)
    cb.set_ticks([1.0, 1.01, 1.02, 1.03, 1.04])
    ax.axis('tight')

    ax.set_yticks([1.0, 1.5, 2.0, 2.5])
    ax.set_yticklabels(["1", "1.5", "2", "2.5"])
    ax.text(6.3, 2.78, r"$B_{\rm CHSH}/2$", fontsize=24)
    
    return fig, axes

E_vec = linspace(0.05, 0.25, Nres)

if calculate:
    data = bell_transient(E_vec, 0, 0) # no single-mode dissipation: gamma_1 = gamma_2 = 0
    save("data-fig-6-b.npy", data)
else:
    data = load("data-fig-6-b.npy")

fig, axes = plot_bell_transient(E_vec, data);

fig.savefig("fig-6-b.png", dpi=200)
fig.savefig("fig-6-b.pdf")

if calculate:
    data = bell_transient(E_vec, gamma_1, gamma_2)
    save("data-fig-6-e.npy", data)
else:
    data = load("data-fig-6-e.npy")

fig, axes = plot_bell_transient(E_vec, data);

fig.savefig("fig-6-e.png", dpi=200)
fig.savefig("fig-6-e.pdf")

def bell_transient(gamma_0_vec, gamma_1, gamma_2):
    
    bell_chsh_mat = zeros((len(gamma_0_vec), len(tlist)))
    for idx, gamma_0 in enumerate(gamma_0_vec):
        r = dynamics_solve(kappa, E, gamma_0, gamma_1, gamma_2, 0)
        bell_chsh_mat[idx,:] = array([bell_chsh_func(r.states[t_idx]) for t_idx, t in enumerate(r.times)])/2
    
    return bell_chsh_mat

def plot_bell_transient(gamma_0_vec, data):
    
    fig, ax = subplots(1, 1, figsize=(8, 6))

    X, Y = meshgrid(tlist/t_unit, gamma_0_vec)
    Z = data
    Z[Z < 1] = 1.0

    norm=mpl.colors.Normalize(1, 1.04) 
    cmap=get_cmap('Reds', lut=250)
    cmap.set_under("white")
    
    p = ax.pcolor(X, Y, Z-0.0001, cmap=cmap, norm=norm)
    cb = fig.colorbar(p)

    ax.set_ylabel(r"$\gamma_0$", fontsize=26, labelpad=0)
    ax.set_xlabel(r"$t\bar{\kappa}^2/\bar{\gamma}_0$", fontsize=26, labelpad=-3)
    cb.set_ticks([1.0, 1.01, 1.02, 1.03, 1.04])
    ax.set_yticks([0.5, 1.0, 1.5, 2.0, 2.5])
    ax.set_yticklabels(["0.5", "1", "1.5", "2", "2.5"])    
    ax.text(6.3, 2.64, r"$B_{\rm CHSH}/2$", fontsize=24)
    
    return fig, axes

gamma_0_vec = linspace(0.5, 2.5, Nres)

if calculate:
    data = bell_transient(gamma_0_vec, 0, 0) # no single-mode dissipation: gamma_1 = gamma_2 = 0
    save("data-fig-6-c.npy", data)
else:
    data = load("data-fig-6-c.npy")

fig, axes = plot_bell_transient(gamma_0_vec, data);

fig.savefig("fig-6-c.png", dpi=200)
fig.savefig("fig-6-c.pdf")

if calculate:
    data = bell_transient(gamma_0_vec, gamma_1, gamma_2)
    save("data-fig-6-f.npy", data)
else:
    data = load("data-fig-6-f.npy")

fig, axes = plot_bell_transient(gamma_0_vec, data);

fig.savefig("fig-6-f.png", dpi=200)
fig.savefig("fig-6-f.pdf")

t_unit = 1/(kappa**2 / gamma_0)
tlist = np.linspace(0, 4*t_unit, 250)

n_th_vec = [0.0, 0.025, 0.05, 0.075]

def bell_transient(kappa_vec):
    
    bell_chsh_vec = []
    for n_th in n_th_vec:
        bell_chsh_mat = zeros((len(kappa_vec), len(tlist)))
        for idx, kappa in enumerate(kappa_vec):
            E = 1.12**2 * kappa / 2
            r = dynamics_solve(kappa, E, gamma_0, gamma_1, gamma_2, n_th)
            bell_chsh_mat[idx,:] = array([bell_chsh_func(r.states[t_idx]) for t_idx, t in enumerate(r.times)])/2
        
        bell_chsh_vec.append(bell_chsh_mat)
        
    return array(bell_chsh_vec)

def plot_bell_transient(kappa_vec, data_vec):
    
    fig, ax = subplots(1, 1, figsize=(8, 6))

    for idx, data in enumerate(data_vec):
        
        X, Y = meshgrid(tlist/t_unit, kappa_vec)
        Z = data
        Z[Z < 1] = 1.0
    
        norm = mpl.colors.Normalize(1, 1.04) 
        cmap = get_cmap('Reds', lut=250)
        cmap.set_under("white")
        
        if idx == 0:
            p = ax.pcolor(X, Y, Z-0.0001, cmap=cmap, norm=norm)
        
        CS = ax.contour(X, Y, Z-0.0001, (1.0,), colors='k')
        clabel(CS, array([1.0]), manual=[(2.0,0.6)], inline_spacing=20,
               inline=1, fmt=r'$N = %1.3f$' % (n_th_vec[idx]), fontsize=18)
    
    ax.set_ylim(0, 1)
    ax.set_xlim(0, 4)
    ax.set_ylabel(r"nonlinearity $\kappa/\gamma_0$", fontsize=22)
    ax.set_xlabel(r"time $t\bar{\kappa}^2/\bar{\gamma}_0$", fontsize=22)
    ax.set_xticks([0, 1, 2, 3, 4])

    ax.set_yticks([0, .2, .4, .6, .8, 1])
    ax.set_yticklabels(["0", "0.2", "0.4", "0.6", "0.8", "1"])
    
    fig.tight_layout()
    
    return fig, axes

kappa_vec = linspace(0.1, 1.0, Nres)

if calculate:
    data_vec = bell_transient(kappa_vec)
    save("data-fig-7-c.npy", data_vec)
else:
    data_vec = load("data-fig-7.npy")

fig, axes = plot_bell_transient(kappa_vec, data_vec);

fig.savefig("fig-7.png", dpi=200)
fig.savefig("fig-7.pdf")

%reload_ext version_information
%version_information numpy, scipy, matplotlib, qutip