from sympy import *
init_printing()

from sympsi import *
from sympsi.boson import *
from sympsi.pauli import *

"""
Explanation: Symbolic derivation qubit-cavity Hamiltonian
<style>
p {
    font-family: "Liberation Serif", serif;
    font-size: 12pt;
}
</style>

Based on: J. R. Johansson (robert@riken.jp), http://jrjohansson.github.io, and Eunjong Kim.
Setup modules
End of explanation
"""


# CPW, qubit and NR energies
omega_r, omega_q, omega_nr = symbols("omega_r, omega_q, omega_{NR}") 
# Coupling CPW-qubit, NR_qubit
g, L, chi, eps = symbols("g, lambda, chi, epsilon")
# Detuning 

# Drives and detunnings

Delta_d, Delta_s, Delta_p  = symbols(" Delta_d, Delta_s, Delta_p ")

A, B, C = symbols("A,B,C") # Electric field amplitude

omega_d, omega_s, omega_p = symbols("omega_d, omega_s, omega_p") # drive frequencies

# Detunning CPW-qubit, NR-qubit

Delta_CPW, Delta_NR = symbols("Delta_{CPW},Delta_{NR}")

# auxilary variables

x, y, t, Hsym = symbols("x, y, t, H ") 

Delta_SP, Delta_SD = symbols("Delta_{SP},Delta_{SD}")

# omega_r, omega_q, g, Delta_d, Delta_s, t, x, chi, Hsym = symbols("omega_r, omega_q, g, Delta_d, Delta_s, t, x, chi, H")

# A, B, C = symbols("A,B,C") # Electric field amplitude

# omega_d, omega_s = symbols("omega_d, omega_s") # 

# omega_nr, L = symbols("omega_{NR},lambda")

# Delta,Delta_t = symbols("Delta,Delta_t")

# y, omega_t = symbols("y, omega_t")


sx, sy, sz, sm, sp = SigmaX(), SigmaY(), SigmaZ(), SigmaMinus(), SigmaPlus()
a = BosonOp("a")
b = BosonOp("b")

H = omega_r * Dagger(a) * a + omega_q/2 * sz + omega_nr * Dagger(b) * b 

H_int = sx * (g  * (a + Dagger(a)))
H_int_2 = sx *( L * (b + Dagger(b)))


H_drive_r = A * (exp(I*omega_d*t)*a + exp(-I*omega_d*t)*Dagger(a))
H_drive_q = B * (exp(I*omega_s*t)*sm + exp(-I*omega_s*t)*sp)
H_drive_NR = C * (exp(I*omega_p*t)*b + exp(-I*omega_p*t)*Dagger(b))
H_total  =  H +   H_drive_r + H_drive_q + H_drive_NR #+ H_int + H_int_2 +
Eq(Hsym,H_total)


"""
Explanation: The Jaynes-Cummings model
The Jaynes-Cummings model is one of the most elementary quantum mechanical models light-matter interaction. It describes a single two-level atom that interacts with a single harmonic-oscillator mode of a electromagnetic cavity.
The Hamiltonian for a two-level system in its eigenbasis can be written as
$$
H = \frac{1}{2}\omega_q \sigma_z
$$
and the Hamiltonian of a quantum harmonic oscillator (cavity or nanomechanical resonator, NR) is
$$
H = \hbar\omega_r (a^\dagger a + 1/2)
$$
$$
H = \hbar\omega_{NR} (b^\dagger b + 1/2)
$$
The atom interacts with the electromagnetic field produced by the cavity (NR) mode $a + a^\dagger$ ($b + b^\dagger$) through its dipole moment. The dipole-transition operators is $\sigma_x$ (which cause a transition from the two dipole states of the atom). The combined atom-cavity Hamiltonian can therefore be written in the form
$$
H = 
\hbar \omega_r (a^\dagger a + 1/2)
+
\hbar \omega_{NR} (b^\dagger b + 1/2)
+ \frac{1}{2}\hbar\Omega\sigma_z 
+
\hbar
\sigma_x \left( g(a + a^\dagger)
+
\lambda(b + b^\dagger)\right)
$$
Although the Jaynes-Cumming Hamiltonian allow us to evolve the given initial state according to the Schrödinger Equation, in an experiment we would like to predicte the response of the coupled cavity-NR-qubit system under the influence of driving fields for the cavity and qubit, and also account for the effects of dissipation and dephasing (not treated here)
The external coherent-state input may be incorporated in the Jaynes-Cummings Hamiltonian by addition of terms involving the amplitude of the driving field $\vec{E_d} \left(\vec{E_s}\right)$ and it's frequency $\omega_d\left(\omega_s\right)$
$$
H_{cavity} = 
E_d \left(e^{i\omega_dt}a +e^{-i\omega_dt}a^\dagger\right)
$$
$$
H_{qubit} = 
E_s \left(e^{i\omega_st}\sigma_- +e^{-i\omega_st}\sigma_+\right)
$$
$$
H_{NR} = 
E_p \left(e^{i\omega_st}b +e^{-i\omega_st}b^\dagger\right)
$$
To obtain the Jaynes-Cumming Hamiltonian 
$$
?
H = 
\hbar\omega_r (a^\dagger a + 1/2)
%-\frac{1}{2}\Delta\sigma_x 
+ \frac{1}{2}\hbar\Omega\sigma_z 
+
\hbar
\sigma_x \left( g(a + a^\dagger)
+
\lambda(b + b^\dagger)\right)
$$
we also need to perform a rotating-wave approximation which simplifies the interaction part of the Hamiltonian. In the following we will begin with looking at how these two Hamiltonians are related.
To represent the atom-cavity Hamiltonian in SymPy we creates an instances of the operator classes BosonOp and SigmaX, SigmaY, and SigmaZ, and use these to construct the Hamiltonian (we work in units where $\hbar = 1$).
End of explanation
"""


U = exp(I * omega_r * t * Dagger(a) * a)

U


H2 = hamiltonian_transformation(U, H_total.expand(),independent = True)

H2

U = exp(I * omega_q * t * sp * sm)

U

H3 = hamiltonian_transformation(U, H2.expand(), independent = True)

H3 = H3.subs(sx, sm + sp).expand()

H3 = powsimp(H3)

H3

"""
Explanation: First move to the interaction picture:
End of explanation
"""


# trick to simplify exponents
def simplify_exp(e):
    if isinstance(e, exp):
        return exp(simplify(e.exp.expand()))

    if isinstance(e, (Add, Mul)):
        return type(e)(*(simplify_exp(arg) for arg in e.args)) 

    return e

H4 = simplify_exp(H3).subs(-omega_r + omega_q, Delta_CPW)

H4

"""
Explanation: We introduce the detuning parameter $\Delta = \omega_q - \omega_r$ and substitute into this expression
End of explanation
"""


H5 = drop_terms_containing(H4, [exp( I * (omega_q + omega_r) * t),
                                exp(-I * (omega_q + omega_r) * t)])

H5 = drop_c_number_terms(H5.expand())

Eq(Hsym, H5)

"""
Explanation: Now, in the rotating-wave approximation we can drop the fast oscillating terms containing the factors $ e^{\pm i( (\omega_q + \omega_r)t}$
End of explanation
"""


U = exp(-I * omega_r * t * Dagger(a) * a)
H6 = hamiltonian_transformation(U, H5.expand(), independent =True)

U = exp(-I * omega_q * t * sp * sm)
H7 = hamiltonian_transformation(U, H6.expand(),independent = True)

H8 = simplify_exp(H7).subs(Delta_CPW, omega_q - omega_r)

H8 = simplify_exp(powsimp(H8)).expand()

H8 = drop_c_number_terms(H8)

H = collect(H8, [g])

Eq(Hsym, H)

H = collect(H,[L*sm])
H = collect(H,[L*sp])
H = collect(H,[L*(Dagger(b)+b), A, B ,C])
H = H.subs(sm+sp,sx)
H

"""
Explanation: This is the interaction term of in the Jaynes-Cumming model in the interaction picture. If we transform back to the Schrödinger picture we have:
End of explanation
"""


U = exp(I * Dagger(a) * a * omega_d * t)


H1 = hamiltonian_transformation(U, H_total, independent=True)

H1

U = exp(I * Dagger(b) * b * omega_p * t)

H2 = hamiltonian_transformation(U, H1, independent=True)

H2

U = exp(I * Dagger(sm) * sm * omega_s * t)

H3 = hamiltonian_transformation(U, H2, independent=True)

H3

# H4 = simplify_exp(H3).subs(omega_d - omega_s, -Delta_SD)

H4 = H3.collect([A,B,C,Dagger(a)*a,Dagger(b)*b, g, L])

H4

H5 = H4.expand().subs(-sp*sm , -sz/2 ).collect([A,B,C,Dagger(a)*a,Dagger(b)*b, g, L, sz])
H5

H5 = H5.subs(-omega_d + omega_r, -Delta_d).subs(-omega_p + omega_nr, - Delta_p).subs(omega_q/2 - omega_s/2, Delta_s); H5

H_total_01 = H5
Eq(Hsym,H_total_01)

H_total_01 = H_total_01 + H_int + H_int_2;
Eq(Hsym,H_total_01)

U = exp(I * omega_r * t * Dagger(a) * a)

U



H2 = hamiltonian_transformation(U, H_total_01.expand(),independent = True)

H2

U = exp(I * omega_q * t * sp * sm)

U

H3 = hamiltonian_transformation(U, H2.expand(), independent = True)

H3 = H3.subs(sx, sm + sp).expand()

H3 = powsimp(H3)

H3

H4 = simplify_exp(H3).subs(-omega_r + omega_q, Delta_CPW)

H4

H5 = drop_terms_containing(H4, [exp( I * (omega_q + omega_r) * t),
                                exp(-I * (omega_q + omega_r) * t)])

H5 = drop_c_number_terms(H5.expand())

Eq(Hsym, H5)

U = exp(-I * omega_r * t * Dagger(a) * a)
H6 = hamiltonian_transformation(U, H5.expand(), independent =True)

U = exp(-I * omega_q * t * sp * sm)
H7 = hamiltonian_transformation(U, H6.expand(),independent = True)

H8 = simplify_exp(H7).subs(Delta_CPW, omega_q - omega_r)

H8 = simplify_exp(powsimp(H8)).expand()

H8 = drop_c_number_terms(H8)

H = collect(H8, [g])



H = collect(H,[L*sm])
H = collect(H,[L*sp])
H = collect(H,[L*(Dagger(b)+b), A, B ,C])
# H = H.subs(sm+sp,sx)
H

"""
Explanation: Linearized interaction
End of explanation
"""


U = exp((x * (  -Dagger(a) * sm + a * sp )).expand())

U

H1 = hamiltonian_transformation(U, H, expansion_search=False, N=4).expand()

H1 = qsimplify(H1)

# H1

H2 = drop_terms_containing(H1.expand(), [ x**4,x**5,x**6])

# H2

H3 = H2.subs(x, g/Delta_CPW)

# H3

H4 = drop_c_number_terms(H3)

# H4

H5 = collect(H4, [Dagger(a) * a, sz])

# H5


U = exp(I * omega_r * t * Dagger(a) * a)

H6 = hamiltonian_transformation(U, H5.expand(),independent = True);
# H6

U = exp(I * omega_q * t * Dagger(sm) * sm)

H7 = hamiltonian_transformation(U, H6.expand(),independent=True); 
# H7


H8 = drop_terms_containing(H7, [exp(I * omega_r * t), exp(-I * omega_r * t),
                                exp(I * omega_q * t), exp(-I * omega_q * t)])

# H8

H9 = qsimplify(H8)

H9 = collect(H9, [Dagger(a) * a, sz])

# H9

U = exp(-I * omega_r * t * Dagger(a) * a)

H10 = hamiltonian_transformation(U, H9.expand(),independent = True);
# H10

U = exp(-I * omega_q * t * Dagger(sm) * sm)

H11 = hamiltonian_transformation(U, H10.expand(),independent=True); 
# H11

# H11  = drop_c_number_terms(H11)
H12 = qsimplify(H11)

H12 = collect(H12, [Dagger(a) * a, sz])

H12 = H12.subs(omega_r, omega_q-Delta_CPW).expand().collect([Dagger(a)*a, Dagger(b)*b, sz, L*g/2/Delta_CPW]).subs(omega_q-Delta_CPW,omega_r)

Eq(Hsym, H12)

"""
Explanation: Dispersive Regime
End of explanation
"""


