%matplotlib inline

from numpy import *

import matplotlib.pyplot as plt

from qutip import *

# number of fock states in the cavity and mirror modes
N = 15
M = 30
r = 1.0

def solve_dynamics_undamped(r, k, alpha, beta, t, N, M):
    
    wm = 1
    w0 = r * wm
    g  = k * wm
    
    a = tensor(destroy(N), identity(M))
    b = tensor(identity(N), destroy(M))
    
    H = w0 * a.dag() * a + wm * b.dag() * b - g * a.dag() * a * (b + b.dag())
    
    psi0 = tensor(coherent(N, alpha), coherent(M, beta))
    
    result = mesolve(H, psi0, t, [], [], options=Odeoptions(nsteps=10000))
    
    return result, expect(commutator(a, a.dag()), result.states), expect(commutator(b, b.dag()), result.states)
    #return result, expect(a.dag() * a, result.states), expect(b.dag() * b, result.states)

t = linspace(0, 2 * pi, 50)

alpha = beta = 2

result, C1_A, C1_B = solve_dynamics_undamped(r, 0.1, alpha, beta, t, N, M)
S1 = [entropy_linear(ptrace(rho, 1)) for rho in result.states]

result, C2_A, C2_B = solve_dynamics_undamped(r, 0.5, alpha, beta, t, N, 2*M)
S2 = [entropy_linear(ptrace(rho, 1)) for rho in result.states]

result, C3_A, C3_B = solve_dynamics_undamped(r, 1.0, alpha, beta, t, N, 10*M)
S3 = [entropy_linear(ptrace(rho, 1)) for rho in result.states]

fig, ax = plt.subplots(figsize=(10,5))
ax.plot(t, S1, label=r'$k = 0.1$')
ax.plot(t, S2, label=r'$k = 0.5$')
ax.plot(t, S3, label=r'$k = 1.0$')
ax.legend()
ax.set_title('Linear entropy of the mirror')
ax.set_xlabel(r'$t$', fontsize=18)
ax.set_ylabel(r'$S$', fontsize=18)
ax.set_xlim(0, t.max());

fig, axes = plt.subplots(1, 2, figsize=(16,5))
axes[0].plot(t, C1_A, label=r'$k = 0.1$')
axes[0].plot(t, C2_A, label=r'$k = 0.5$')
axes[0].plot(t, C3_A, label=r'$k = 1.0$')
axes[0].legend()
axes[0].set_title('[a, a.dag()]')
axes[0].set_xlabel(r'$t$', fontsize=18)
axes[0].set_ylabel(r'$S$', fontsize=18);
axes[0].set_ylim(0, 1.1);

axes[1].plot(t, C1_B, label=r'$k = 0.1$')
axes[1].plot(t, C2_B, label=r'$k = 0.5$')
axes[1].plot(t, C3_B, label=r'$k = 1.0$')
axes[1].legend()
axes[1].set_title('[b, b.dag()]]')
axes[1].set_xlabel(r'$t$', fontsize=18)
axes[1].set_ylabel(r'$S$', fontsize=18)
axes[1].set_ylim(0, 1.1);

t = [0, pi/2, pi, 3*pi/2, 2*pi]

result, C_A, C_B = solve_dynamics_undamped(r, 0.3, alpha, beta, t, N, M)

fig, axes = plt.subplots(len(t), 2, figsize=(8,16))

for idx, rho in enumerate(result.states):
    rho_mirror = ptrace(rho, 1)
    plot_wigner_fock_distribution(rho_mirror, fig=fig, axes=axes[idx])
    
fig.tight_layout()

t = [0, 2*pi]

k_vec = [0.5, 0.4082, 0.3536]

states = [solve_dynamics_undamped(r, k, alpha, beta, t, 3*N, 3*M)[0].states[-1] for k in k_vec]

fig, axes = plt.subplots(len(k_vec), 2, figsize=(8,12))

for idx, rho in enumerate(states):
    rho_cavity = ptrace(rho, 0)
    plot_wigner_fock_distribution(rho_cavity, fig=fig, axes=axes[idx])
    
fig.tight_layout()

# number of fock states in the cavity and mirror modes
N = 20
M = 30

def solve_dynamics_damped(r, k, gamma, alpha, beta, t, N, M):
    
    wm = 1
    w0 = r * wm
    g  = k * wm
    
    a = tensor(destroy(N), identity(M))
    b = tensor(identity(N), destroy(M))
    
    H = w0 * a.dag() * a + wm * b.dag() * b - g * a.dag() * a * (b + b.dag())

    psi0 = tensor(coherent(N, alpha), coherent(M, beta))
    
    c_ops = [sqrt(gamma) * b]
    
    return mesolve(H, psi0, t, c_ops, [], options=Odeoptions(nsteps=20000))

gamma = 1.0
alpha = 2
beta = 0

t = linspace(0, 2 * pi, 25)

S1 = [entropy_linear(ptrace(rho, 1)) for rho in solve_dynamics_damped(r, 0.1, gamma, alpha, beta, t, N, M).states]

S2 = [entropy_linear(ptrace(rho, 1)) for rho in solve_dynamics_damped(r, 0.5, gamma, alpha, beta, t, N, M).states]

S3 = [entropy_linear(ptrace(rho, 1)) for rho in solve_dynamics_damped(r, 1.0, gamma, alpha, beta, t, N, M).states]

fig, ax = plt.subplots(figsize=(10,5))
ax.plot(t, S1, label=r'$k = 0.1$')
ax.plot(t, S2, label=r'$k = 0.5$')
ax.plot(t, S3, label=r'$k = 1.0$')
ax.legend()
ax.set_title('Linear entropy of the mirror')
ax.set_xlabel(r'$t$', fontsize=18)
ax.set_ylabel(r'$S$', fontsize=18)
ax.set_xlim(0, t.max());

t = [0, 2*pi]
k = 0.5
alpha = beta = 2

gamma_vec = [0.001, 0.01, 0.1, 1.0]

states = [solve_dynamics_damped(r, k, gamma, alpha, beta, t, 3*N, 3*M).states[-1] for gamma in gamma_vec]

fig, axes = plt.subplots(len(gamma_vec), 2, figsize=(8,16))

for idx, rho in enumerate(states):
    rho_cavity = ptrace(rho, 0)
    plot_wigner_fock_distribution(rho_cavity, fig=fig, axes=axes[idx])
    
fig.tight_layout()

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