import numpy as np
import matplotlib.pyplot as plt
from qutip import *

qutip.about()

# Let's define a quantum object. We can start from a simple 
# bosonic destruction operator
n_max = 2
a = destroy(n_max)

a

# in QuTiP, all number (Fock) states are defined on the "basis". 
psi0 = basis(n_max)

psi0 # it's a ket!

a*psi0 # destroing the ground state gives the vacuum

a.dag()*psi0

psi0 = basis(n_max)
for n in range(1,n_max):
    print("n = ",n)
    psi = a.dag()*psi0/np.sqrt(n)
    print(psi)
    psi0 = psi
psi0 = basis(n_max)

rho0 = psi0*psi0.dag()
rho0

rho1 = fock(n_max,1)*fock(n_max,1).dag()
rho1

mixed = (rho0 + rho1)/2
mixed

# let's define the Schrodinger cat state
cat = fock(n_max,0)+fock(n_max,1)
cat = cat 
cat_rho = ket2dm(cat)/2
cat_rho

# trace Tr[rho]
print(cat_rho.tr())
print(mixed.tr())

# purity Tr[rho*rho]
print((cat_rho*cat_rho).tr())
print((mixed*mixed).tr())

# We can now build a Hamiltonian, 
# that of the Quantum Harmonic Oscillator 
# exploring its features with QuTiP
n_max = 6
a = destroy(n_max)

H = a.dag()*a

H

E_n , Psi_n = H.eigenstates()

E_n

Psi_n

g = .6
H_driven = a.dag()*a + g*(a + a.dag())
En_driven , Psin_driven = H_driven.eigenstates()

plt.figure()
plt.plot(E_n,"o-",markersize = 10, label="QHO")
plt.plot(En_driven,"s--",markersize = 10, label="Driven QHO")

plt.ylabel("energy")
plt.xlabel("eigenvalue")
plt.legend()
plt.show()
plt.close()

rho0 = ket2dm(basis(n_max,n_max-1))
t = np.linspace(0,30,300)
kappa = 0.1
results = mesolve(H,rho0,t,[kappa*a])
rho_t = results.states
rho_t[0]

rho_t[-1]

results = mesolve(H_driven,rho0,t,c_ops=[kappa*a], e_ops=[a,a.dag()*a])
at = results.expect[0]
nt = results.expect[1]

plt.figure()
plt.plot(t,nt,"-",markersize = 10, label="$\\langle a^\dagger a \\rangle$")
plt.ylabel("photons in cavity")
plt.xlabel("time")
plt.legend()
plt.show()
plt.close()

rho0qubit = ket2dm(basis(2,1))
rho0_tot = tensor(rho0,rho0qubit)
Om = 0.5
Hint_tot = Om*tensor(a+a.dag(),sigmax())
H0_tot = tensor(H,qeye(2))+tensor(qeye(n_max),sigmaz())
H_tot=H0_tot+Hint_tot
t = np.linspace(0,50,300)
kappa = 0.1
a_tot = tensor(a,qeye(2))
sz_tot = tensor(qeye(n_max),sigmaz())
sm_tot = tensor(qeye(n_max),sigmam())
results_tot = mesolve(H_tot,rho0_tot,t,
                      c_ops=[kappa*a_tot, kappa*sm_tot], 
                      e_ops= [a_tot.dag()*a_tot,sz_tot])
nphot_tot_t = results_tot.expect[0]
sz_tot_t = results_tot.expect[1]

plt.figure()
plt.plot(t,sz_tot_t,"-",markersize = 10, 
         label="$\\langle \sigma_z \\rangle (t)$")
plt.plot(t,nphot_tot_t,"--",markersize = 10, 
         label="$\\langle a^\dagger a \\rangle (t)$")
plt.ylabel("System excitation")
plt.xlabel("time")
plt.legend()
plt.show()
plt.close()

b = Bloch()

b.show()

b2 = Bloch()
vec = [0,1,0]
b2.add_vectors(vec)
b2.render()

x = (basis(2,0)+(1+0j)*basis(2,1)).unit() 
y = (basis(2,0)+(0+1j)*basis(2,1)).unit()
z = (basis(2,0)+(0+0j)*basis(2,1)).unit() 
b3 = Bloch()      
b3.add_states([x,y,z])
b3.show()

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

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

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

psi = squeeze(N, 0.5) * basis(N, 0)
display(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])

q = QubitCircuit(3, reverse_states=False)
q.add_gate("TOFFOLI", controls=[0, 2], targets=[1])
q.add_gate("CNOT", targets=[1], controls=[0])
display("q",q.png)

q.add_gate("TOFFOLI", controls=[0, 2], targets=[1])
display("q  updated",q.png)

q2 = QubitCircuit(3, reverse_states=False)
q2.add_gate("TOFFOLI", controls=[0, 2], targets=[1])
q2.add_gate("TOFFOLI", controls=[0, 2], targets=[1])
display("q2",q2.png)

U = gate_sequence_product(q.propagators())
U

qutip.about()