from sympy import *
from sympy.physics.quantum import *
from sympy.physics.quantum.density import *
from sympy.physics.quantum.spin import (
    Jx, Jy, Jz, Jplus, Jminus, J2,
    JxBra, JyBra, JzBra,
    JxKet, JyKet, JzKet,
)
from IPython.core.display import display_pretty

%load_ext sympyprinting

psi = Ket('psi')
phi = Ket('phi')

d = Density((psi,0.5),(phi,0.5)); d

display_pretty(d)

d.states()

d.probs()

d.doit()

Dagger(d)

A = Operator('A')

d.apply_op(A)

up = JzKet(S(1)/2,S(1)/2)
down = JzKet(S(1)/2,-S(1)/2)

d2 = Density((up,0.5),(down,0.5)); d2

represent(d2)

d2.apply_op(Jz)

qapply(_)

qapply((Jy*d2).doit())

entropy(d2)

entropy(represent(d2))

entropy(represent(d2,format="numpy"))

entropy(represent(d2,format="scipy.sparse"))

from sympy.core.trace import Tr

A, B, C, D = symbols('A B C D',commutative=False)

t1 = TensorProduct(A,B,C)

d = Density([t1, 1.0])
d.doit()

t2 = TensorProduct(A,B)
t3 = TensorProduct(C,D)

d = Density([t2, 0.5], [t3, 0.5])
d.doit() 

#mixed states
d = Density([t2+t3, 1.0])
d.doit() 

from sympy.physics.quantum.density import Density
from sympy.physics.quantum.spin import (
    Jx, Jy, Jz, Jplus, Jminus, J2,
    JxBra, JyBra, JzBra,
    JxKet, JyKet, JzKet,
)
from sympy.core.trace import Tr

d = Density([JzKet(1,1),0.5],[JzKet(1,-1),0.5]);
t = Tr(d); 

display_pretty(t)
print t.doit()

#Partial trace on mixed state
from sympy.core.trace import Tr

A, B, C, D = symbols('A B C D',commutative=False)

t1 = TensorProduct(A,B,C)

d = Density([t1, 1.0])
d.doit()

t2 = TensorProduct(A,B)
t3 = TensorProduct(C,D)

d = Density([t2, 0.5], [t3, 0.5])


tr = Tr(d,[1])
tr.doit()

from sympy.physics.quantum.density import Density
from sympy.physics.quantum.spin import (
    Jx, Jy, Jz, Jplus, Jminus, J2,
    JxBra, JyBra, JzBra,
    JxKet, JyKet, JzKet,
)
from sympy.core.trace import Tr

tp1 = TensorProduct(JzKet(1,1), JzKet(1,-1))

#trace out 0 index
d = Density([tp1,1]);
t = Tr(d,[0]); 

display_pretty(t)
display_pretty(t.doit())

#trace out 1 index
t = Tr(d,[1])
display_pretty(t)
t.doit()

psi = Ket('psi')
phi = Ket('phi')

u = UnitaryOperator()
d = Density((psi,0.5),(phi,0.5)); d

display_pretty(qapply(u*d))
 
# another example
up = JzKet(S(1)/2, S(1)/2)
down = JzKet(S(1)/2, -S(1)/2)
d = Density((up,0.5),(down,0.5))

uMat = Matrix([[0,1],[1,0]])
display_pretty(qapply(uMat*d))




from sympy.physics.quantum.gate import UGate
from sympy.physics.quantum.qubit import Qubit

uMat = UGate((0,), Matrix([[0,1],[1,0]]))
d = Density([Qubit('0'),0.5],[Qubit('1'), 0.5])

display_pretty(d)

#after applying Not gate
display_pretty(qapply(uMat*d) )