from sympy import *
init_printing()

from sympy.physics.quantum import *
from sympy.physics.quantum.boson import *
from sympy.physics.quantum.fermion import *
from sympy.physics.quantum.operatorordering import *

from sympy_quantum_utils import *

omega_0, t, alpha = symbols("omega_0, t, alpha")
n = 8

a = BosonOp("a")

A = I * omega_0 * Dagger(a) * a * t
U = exp(A)

U

e = bch_expansion(A, Dagger(a) * a, n).doit()

normal_ordered_form(e)

unitary_transformation_auto(U, Dagger(a) * a)

unitary_transformation_auto(U, a)

e = bch_expansion(A, a, n).doit()

e = normal_ordered_form(expand(e))

e = collect(e, a)

e

e == (exp(-I * omega_0 * t).series(t, 0, n=n).removeO()) * a

unitary_transformation_auto(U, Dagger(a))

e = bch_expansion(A, Dagger(a), n).doit()

e = normal_ordered_form(expand(e))

e = collect(e, Dagger(a))

e

e == (exp(I * omega_0 * t).series(t, 0, n=n).removeO()) * Dagger(a)

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

U

unitary_transformation_auto(U, a)

unitary_transformation_auto(U, Dagger(a))

unitary_transformation_auto(U, Dagger(a) * a)

n = 6
omega_a, omega_b, g, A, Delta = symbols("omega_a, omega_b, g, A, Delta", positive=True)
Hsym, omega_d = symbols("H, omega_d")

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

H = omega_a * Dagger(a) * a + omega_b * Dagger(b) * b + g * Dagger(a) * a * (b + Dagger(b)) \
    + (A * exp(-I * omega_d * t) + conjugate(A) * exp(I * omega_d * t)) * (a + Dagger(a))

Eq(Hsym, H)

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

U

unitary_transformation_auto(U, H, independent=True)

H1 = hamiltonian_transformation_auto(U, H, independent=True)

H1

H2 = drop_terms_containing(H1.expand(), [exp(-2*I*omega_d*t), exp(2*I*omega_d*t)])

Eq(Symbol("H_{rwa}"), H2)

H3 = H2.subs(omega_a, Delta + omega_d).expand()

H3

alpha = symbols("alpha")

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

U

H4 = hamiltonian_transformation_auto(U, H3, independent=True)

H4

H5 = H4.expand().subs({A: alpha * Delta, conjugate(alpha): alpha}) 

H5 = collect(H5, [g * Dagger(a) * a, - alpha * g])

H5

H6 = drop_c_number_terms(H5)

H6

H7 = drop_terms_containing(H6, [g * Dagger(a) * a])  # XXX not working

Eq(Hsym, H7)

%reload_ext version_information

%version_information sympy