import numpy as np
from scipy.linalg import sqrtm

#Computation basis definition
LO = np.array([[1.0],[0.0]],dtype=complex)
HI = np.array([[0.0],[1.0]],dtype=complex)

#hermite conjugation
dagger = lambda x: x.transpose().conjugate()
#inner product of two kets (kets->number)
braket = lambda x, y: np.dot(dagger(x),y).item(0,0)
#outer product of two kets (kets->matrix)
ketbra = lambda x, y: np.dot(x,dagger(y))
#computation base vector (tuple of 0,1 -> ket)
binket = lambda x: reduce(np.kron, map(lambda y: ((y==1)*HI + (y==0)*LO), x))

#Bell states ket vectors
Bell1 = binket((0,1))+binket((1,0))
Bell2 = binket((0,1))-binket((1,0))
Bell3 = binket((0,0))+binket((1,1))
Bell4 = binket((0,0))-binket((1,1))

#density matrices of Bell states
rBells = map(lambda x: ketbra(x,x)*0.5, [Bell1,Bell2,Bell3,Bell4])

print("*"*10)
print("SciPy square roots of Matrix")
print("*"*10+"\n")

for Bell in rBells:
    print("Bell state:")
    print(Bell)
    print("and it square root >>>")
    print(sqrtm(Bell))
    print("*"*5)

def sqrtm2(M):
    #Eigen value decomposition based matrix square root computation
    Di, Rot = np.linalg.eig(M)
    Di = np.diag(Di)
    Di = np.sqrt(Di)
    try:
        N = np.dot(np.dot(Rot,Di),np.linalg.inv(Rot))
    except:
        raise
    return N

print("\n"*2+"*"*10)
print("Eigenvalue-based square roots of Matrix")
print("*"*10+"\n")
for Bell in rBells:
    print("Bell state:")
    print(Bell)
    print("and it square root >>>")
    print(sqrtm2(Bell))
    print("*"*5)