from sympy import *
init_printing()

from sympsi import *
from sympsi.boson import *
from sympsi.operatorordering import *

omega_r = symbols("omega_r", positive=True)
Hsym = symbols("H")
a = BosonOp("a")

H0 = omega_r * Dagger(a) * a

Eq(Hsym, H0)

omega_d, t = symbols("omega_d, t")
A = symbols("A")

Hdrive = (A * exp(-I * omega_d * t) + conjugate(A) * exp(I * omega_d * t)) * (a + Dagger(a))

Hdrive

H = H0 + Hdrive

Eq(Hsym, H)

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

U

H2 = hamiltonian_transformation(U, H.expand())

H2

Delta = symbols("Delta", positive=True)

H3 = collect(H2, Dagger(a) * a).subs(omega_r - omega_d, Delta)

H3

H3 = drop_terms_containing(H3, [exp( 2 * I * omega_d * t),
                                exp(-2 * I * omega_d * t)])

Eq(Hsym, H3)

alpha = symbols("alpha")
H = Dagger(a) * alpha - conjugate(alpha) * a
U = exp(H)

U

H4 = hamiltonian_transformation(U, H3)

H4 = collect(H4.expand(), [Dagger(a)*a, a, Dagger(a)])

H4

H5 = H4.subs(alpha, A/Delta)

H5 = drop_c_number_terms(H5)

H5

a, b = BosonOp("a"), BosonOp("b")

kappa, omega_p, omega_a, omega_b, t = symbols("kappa, omega_p, omega_a, omega_b, t")
Delta_a, Delta_b = symbols("Delta_a, Delta_b", positive=True)

Hdrive = (A * exp(-I * omega_p * t) + conjugate(A) * exp(I * omega_p * t)) * (b + Dagger(b))
H1 = omega_a * Dagger(a) * a + omega_b * Dagger(b) * b  + kappa * (a ** 2 * Dagger(b) + Dagger(a) ** 2 * b) + Hdrive

Eq(Hsym, H1)

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

U

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

H2

H3 = drop_terms_containing(H2, [exp(2 * I * omega_p * t),
                                exp(-2 * I * omega_p * t)])

H3 = collect(H3, kappa)

H3

H4 = H3.subs(omega_b, Delta_b + omega_p).expand()

H4

beta = symbols("beta")
H = Dagger(b) * beta - conjugate(beta) * b
U = exp(H)

U

H5 = hamiltonian_transformation(U, H4.expand(), independent=True)

H5

H5 = collect(H5.expand(), [Dagger(a) * a, Dagger(b) * b, Dagger(a) ** 2 * b, a ** 2 * Dagger(b), b, Dagger(b)])

H5

H6 = H5.subs(beta, A/Delta_b)

H6 = drop_c_number_terms(H6)

H6 = collect(H6, [exp( I * omega_p * t) * kappa * a ** 2,
                  exp(-I * omega_p * t) * kappa * Dagger(a) ** 2])

H6

H7 = drop_terms_containing(H6.expand(), [b, Dagger(b)])

H7 = collect(H7, kappa / Delta_b)

H7

H = omega_a * Dagger(a) * a  + kappa * (a ** 2 * exp(I * omega_p * t) + Dagger(a) ** 2 * exp(-I * omega_p * t))

Eq(Hsym, H)

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

U

H2 = hamiltonian_transformation(U, H)

H2

H3 = H2.subs(omega_d, omega_p/2)

H3 = collect(H3, Dagger(a) * a)

H3

H4 = H3.subs(omega_a, Delta_a + omega_p/2)

H4

chi = symbols("chi")

U = exp(chi/2 * a **2 - chi/2 * Dagger(a) ** 2)

U

hamiltonian_transformation(U, a)

hamiltonian_transformation(U, Dagger(a))

H5 = hamiltonian_transformation(U, H4)

H5

H6 = normal_ordered_form(H5.expand(), independent=True)

H6 = drop_c_number_terms(H6)

H6 = collect(H6, [Dagger(a) * a, Dagger(a)**2, a**2])

H6

H7 = collect(H6, H6.args[1].args[0])

H7

# Trick to simplify the coefficients for the quantum operators
H8 = Add(*(simplify(arg.args[0]) * Mul(*(arg.args[1:])) for arg in H7.args))

H8

chi_eq = H8.args[0].args[0]

Eq(chi_eq, 0)

chi_sol = simplify(solve(chi_eq, chi)[3])

Eq(chi, chi_sol)

H9 = H8.subs(chi_eq.args[0], -chi_eq.args[1])

H9

Eq(Symbol("omega"), H9.args[0])

%reload_ext version_information

%version_information sympy, sympsi