from sympy import *
init_printing(use_latex='png')

from sympsi.operator import Operator
from sympsi.boson import BosonOp
from sympsi.dagger import Dagger

# Constants
d = symbols("Delta", real=True)  # detuning
hbar, g, O, wp, wm = symbols("hbar, g, Omega, omega_+, omega_-", positive=True)

M = Matrix([[-d, 0, -g, 0],
            [0, -d, 0, 0],
            [-g, 0, O, 0],
            [0, 0, 0, O]])
M

# Operators
Xc, Xm = Operator('X_c'), Operator('X_m')  # Position Quadratures
Pc, Pm = Operator('P_c'), Operator('P_m')  # Momentum Quadratures
ac, am = BosonOp('a_c'), BosonOp('a_m')    # Boson Operators

R = Matrix([Xc, Pc, Xm, Pm])
a = Matrix([ac, Dagger(ac), am, Dagger(am)])
Q = Matrix([[1, 1, 0, 0],
            [-I, I, 0, 0],
            [0, 0, 1, 1],
            [0, 0, -I, I]])/sqrt(2)
Eq(R, MatMul(Q, a))

# Constraint representative of the canonical commutation relation
J = Matrix([[0, 1, 0, 0],
            [-1, 0, 0, 0],
            [0, 0, 0, 1],
            [0, 0, -1, 0]])
J

# Operators in the normal mode coordinates
X_p, X_m = Operator('X_+'), Operator('X_-')  # Position Quadratures
P_p, P_m = Operator('P_+'), Operator('P_-')  # Momentum Quadratures
a_p, a_m = BosonOp('a_+'), BosonOp('a_-')    # Boson Operators

Rnm = Matrix([X_p, P_p, X_m, P_m])
anm = Matrix([a_p, Dagger(a_p), a_m, Dagger(a_m)])
Eq(Rnm, MatMul(Q, anm))

D = diag(wp, wp, wm, wm)
D

# extracting eigenvectors and eigenvalues for J*M
[(e1,_,[v1]), (e2,_,[v2]), (e3,_,[v3]), (e4,_,[v4])] = (J * M).eigenvects()

e1, e2, e3, e4

# extracting eigenvectors and eigenvalues for J*D
[(f1,_,[w1]), (f2,_,[w2]), (f3,_,[w3]), (f4,_,[w4])] = (J * D).eigenvects()

f1, f2, f3, f4

wpsol = solve(e2-f2, wp)[0]
Eq(wp, wpsol)

wmsol = solve(e4-f4, wm)[0]
Eq(wm, wmsol)

# simpler expression without an imaginary number
wpsol = sqrt((d**2 + O**2 + sqrt((d**2 -O**2)**2 -4*d*O*g**2))/2)
wmsol = sqrt((d**2 + O**2 - sqrt((d**2 -O**2)**2 -4*d*O*g**2))/2)

Eq(wp, wpsol), Eq(wm, wmsol)

# arbitrary constants
aa, bb, cc, dd = symbols("alpha, beta, gamma, delta")
x, y, z, w = symbols("x, y, z, w")

ee2, ff2 = simplify(sqrt(2)*e2), sqrt(2)*f2
ee4, ff4 = simplify(sqrt(2)*e4), sqrt(2)*f4
subs_dic = {ee2:ff2, ee2**2+2*d**2:ff2**2+2*d**2, ee4:ff4, ee4**2+2*d**2: ff4**2+2*d**2}
u1, u2, u3, u4 = (v1.subs(subs_dic), v2.subs(subs_dic),
                  v3.subs(subs_dic), v4.subs(subs_dic))

A = (x*u1).row_join(y*u2).row_join(z*u3).row_join(w*u4)
A = simplify(A)
A

B = (aa*w1).row_join(bb*w2).row_join(cc*w3).row_join(dd*w4)
B

Ainv = A.inv()
Ainv = simplify(Ainv)
S = simplify(B * Ainv)

S = Matrix([[simplify(S[i, j].factor()) for j in range(4)] for i in range(4)])
S

SJST = simplify(S * J * S.T)
SJST = Matrix([[simplify(SJST[i,j].subs({wp:wpsol, wm:wmsol}).factor()) for j in range(4)] for i in range(4)])

Eq(1, simplify(SJST[0, 1].subs(subs_dic).factor()))

Eq(1, simplify(SJST[2, 3].subs(subs_dic).factor()))

S_ = S.subs({x: sqrt(wp * (wp**2 - d**2)/(O*(wp**2-wm**2))),
        y: sqrt(wp * (wp**2 - d**2)/(O*(wp**2-wm**2))),
        z: sqrt(wm * (d**2 - wm**2)/(O*(wp**2-wm**2))),
        w: sqrt(wm * (d**2 - wm**2)/(O*(wp**2-wm**2))),
        aa:1, bb:1, cc:1, dd:1})
S_ = Matrix([[simplify(S_[i,j].factor()).factor() for j in range(4)] for i in range(4)])
S_

#check1
STDS_ = simplify(S_.T * D * S_)
STDS_ = Matrix([[simplify(STDS_[i,j].subs({wp:wpsol, wm:wmsol}).factor()) for j in range(4)] for i in range(4)])
STDS_

#check2
SJST_ = simplify(S_ * J * S_.T)
SJST_ = Matrix([[simplify(SJST_[i,j].subs({wp:wpsol, wm:wmsol}).factor()) for j in range(4)] for i in range(4)])
SJST_

T = simplify((Q.inv() * S_ * Q))

lhs = anm[0]
rhs = simplify((T * a)[0])
Eq(lhs, rhs)

lhs = anm[1]
rhs = simplify((T * a)[1])
Eq(lhs, rhs)

lhs = anm[2]
rhs = simplify((T * a)[2])
Eq(lhs, rhs)

lhs = anm[3]
rhs = simplify((T * a)[3])
Eq(lhs, rhs)

S__ = simplify(S_.subs({d: -O, wp: wpsol.subs(d, -O), wm: wmsol.subs(d, -O)}))
lhs = Rnm
rhs = MatMul(S__, R)
Eq(lhs, rhs)

H = hbar * wp/2 *(X_p**2 + P_p**2) + hbar * wm/2 * (X_m**2 + P_m**2)
H

H_ = H.subs({wp: wpsol.subs(d, -O), wm: wmsol.subs(d, -O),
             X_p: (rhs.doit())[0],
             P_p: (rhs.doit())[1],
             X_m: (rhs.doit())[2],
             P_m: (rhs.doit())[3],
             })
simplify(H_.expand().factor()).subs({sqrt(((O+g)**2).expand()): O+g,
                                     Xm * Xc: Xc * Xm,
                                     Pm * Pc: Pc * Pm}).expand()

%reload_ext version_information

%version_information sympy, sympsi