from sympy import *
init_printing()

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

omega_r, Omega, g, Delta, t, x, Hsym = symbols("omega_r, Omega, g, Delta, t, x, H")

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

H = omega_r * Dagger(a) * a + Omega/2 * sz + g * sx * (a + Dagger(a))

Eq(Hsym, H)

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

U

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

H2

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

U

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

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

H3 = powsimp(H3)

H3

# 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, Delta)

H4

H5 = drop_terms_containing(H4, [exp( I * (Omega + omega_r) * t),
                                exp(-I * (Omega + 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())

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

H8 = simplify_exp(H7).subs(Delta, Omega - omega_r)

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

H8 = drop_c_number_terms(H8)

H = collect(H8, g)

Eq(Hsym, H)

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

U

#H1 = unitary_transformation(U, H, allinone=True, expansion_search=False, N=3).expand()
#H1 = qsimplify(H1)
#H1

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

H1 = qsimplify(H1)

H1

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

H2

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

H3

H4 = drop_c_number_terms(H3)

H4

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

H5

H5.expand()

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

H6 = hamiltonian_transformation(U, H5.expand()); H6

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

H7 = hamiltonian_transformation(U, H6.expand()); H7

H8 = drop_terms_containing(H7, [exp(I * omega_r * t), exp(-I * omega_r * t),
                                exp(I * Omega * t), exp(-I * Omega * 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()); H10

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

H11 = hamiltonian_transformation(U, H10.expand()); H11

H12 = qsimplify(H11)

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

H12 = H12.subs(omega_r, Omega-Delta).expand().collect([Dagger(a)*a, sz]).subs(Omega-Delta,omega_r)

Eq(Hsym, H12)

%reload_ext version_information

%version_information sympy, sympsi