%run init.ipynb

def gell_mann(d, g, j, k):
    gm = np.zeros((d,d), dtype= complex)
    if g == 'i':  # identity
        for j in range(0,d):
            gm[j,j] = 1
    elif g == 'd':  # diagonal
        nt = math.sqrt(2/(j*(j+1)))
        for m in range(0,j):
            gm[m,m] = nt
        gm[j,j] = -j*nt
    elif g == 's':  # symmetric
        gm[j-1,k-1] = 1;  gm[k-1,j-1] = 1
    elif g == 'a':  # anti-symmetric
        gm[j-1,k-1] = -1j;  gm[k-1,j-1] = 1j
    return gm

gell_mann(2, 'i', 0, 0), gell_mann(2, 'd', 1, 1), gell_mann(2, 's', 1, 2), gell_mann(2, 'a', 1, 2)

def bloch_vector(A):
    d = A.shape[0]; bv = np.zeros((d**2-1), dtype = complex)
    #bv[0] = trace(d,A)
    for j in range(1,d):   # diagonal
        bv[j-1] = 0
        for k in range(1,j+1):
            bv[j-1] += A[k-1,k-1]
        bv[j-1] -= j*A[j,j]
        bv[j-1] *= math.sqrt(2/(j*(j+1)))
    for k in range(1,d):  # symmetric
        for l in range(k+1,d+1):
            j += 1
            bv[j-1] = A[l-1,k-1] + A[k-1,l-1]
    for k in range(1,d):  # anti-symmetric
        for l in range(k+1,d+1):
            j += 1
            bv[j-1] = -1j*(A[l-1,k-1] - A[k-1,l-1])
    return bv.real

s0 = np.array([[1,0],[0,1]]); s1 = np.array([[0,1],[1,0]])
s2 = np.array([[0,-1j],[1j,0]]); s3 = np.array([[1,0],[0,-1]])
A = 0.5*(s0 + 0.1*s1 + 0.2*s2 + 0.3*s3); bloch_vector(A)

def corr_mat_dd(da, db, M):
    cmdd = np.zeros((da-1,db-1), dtype = complex)
    for j in range(1,da):
        for k in range(1,db):
            for m in range(1,j+1):
                for n in range(1,k+1):
                    cmdd[j-1,k-1] += M[(m-1)*db+(n-1),(m-1)*db+(n-1)]
            m = j+1
            for n in range(1,k+1):
                cmdd[j-1,k-1] -= j*M[(m-1)*db+(n-1),(m-1)*db+(n-1)]
            n = k+1
            for m in range(1,j+1):
                cmdd[j-1,k-1] -= k*M[(m-1)*db+(n-1),(m-1)*db+(n-1)]
            m = j+1
            n = k+1
            cmdd[j-1,k-1] += j*k*M[(m-1)*db+(n-1),(m-1)*db+(n-1)]
            cmdd[j-1,k-1] *= 2/sqrt(j*(j+1)*k*(k+1))
    return cmdd.real

def Pauli(j):
    if j == 0:
        return np.array([[1,0],[0,1]])
    elif j == 1:
        return np.array([[0,1],[1,0]])
    elif j == 2:
        return np.array([[0,-1j],[1j,0]])
    elif j == 3:
        return np.array([[1,0],[0,-1]])

def rho2qb(a1, a2, a3, b1, b2, b3, c11, c12, c13, c21, c22, c23, c31, c32, c33):
    return (1/4)*(tp(Pauli(0),Pauli(0)) + b1*tp(Pauli(0),Pauli(1)) + b2*tp(Pauli(0),Pauli(2)) + b3*tp(Pauli(0),Pauli(3))
           + a1*tp(Pauli(1),Pauli(0)) + c11*tp(Pauli(1),Pauli(1)) + c12*tp(Pauli(1),Pauli(2)) + c13*tp(Pauli(1),Pauli(3))
           + a2*tp(Pauli(2),Pauli(0)) + c21*tp(Pauli(2),Pauli(1)) + c22*tp(Pauli(2),Pauli(2)) + c23*tp(Pauli(2),Pauli(3))
           + a3*tp(Pauli(3),Pauli(0)) + c31*tp(Pauli(3),Pauli(1)) + c32*tp(Pauli(3),Pauli(2)) + c33*tp(Pauli(3),Pauli(3)))

a1, a2, a3, b1, b2, b3 = 0.11, 0.12, 0.13, 0.25, 0.26, 0.27
c11, c12, c13, c21, c22, c23, c31, c32, c33 = 0.11, 0.12, 0.13, 0.21, 0.22, 0.23, 0.31, 0.32, 0.33
corr_mat_dd(2,2,rho2qb(a1, a2, a3, b1, b2, b3, c11, c12, c13, c21, c22, c23, c31, c32, c33))

def corr_mat_ds(da, db, M):
    cmds = np.zeros((da-1,db*(db-1)//2), dtype = complex)
    for j in range(1,da):
        n = 0
        for k in range(1,db):
            for l in range(k+1,db+1):
                n += 1
                for m in range(1,j+1):
                    cmds[j-1,n-1] += (M[(m-1)*db+(l-1),(m-1)*db+(k-1)] + M[(m-1)*db+(k-1),(m-1)*db+(l-1)])
                m = j+1
                cmds[j-1,n-1] -= j*(M[(m-1)*db+(l-1),(m-1)*db+(k-1)] + M[(m-1)*db+(k-1),(m-1)*db+(l-1)])
                cmds[j-1,n-1] *= sqrt(2/(j*(j+1)))
    return cmds.real

corr_mat_ds(2,2,rho2qb(a1, a2, a3, b1, b2, b3, c11, c12, c13, c21, c22, c23, c31, c32, c33))

def corr_mat_da(da, db, M):
    cmda = np.zeros((da-1,db*(db-1)//2), dtype = complex)
    for j in range(1,da):
        n = 0
        for k in range(1,db):
            for l in range(k+1,db+1):
                n += 1
                for m in range(1,j+1):
                    cmda[j-1,n-1] += (M[(m-1)*db+(l-1),(m-1)*db+(k-1)] - M[(m-1)*db+(k-1),(m-1)*db+(l-1)])
                m = j+1
                cmda[j-1,n-1] -= j*(M[(m-1)*db+(l-1),(m-1)*db+(k-1)] - M[(m-1)*db+(k-1),(m-1)*db+(l-1)])
                cmda[j-1,n-1] *= -1j*sqrt(2/(j*(j+1)))
    return cmda.real

corr_mat_da(2,2,rho2qb(a1, a2, a3, b1, b2, b3, c11, c12, c13, c21, c22, c23, c31, c32, c33))

def corr_mat_sd(da, db, M):
    cmsd = np.zeros((da*(da-1)//2,db-1), dtype = complex)
    n = 0
    for k in range(1,da):
        for l in range(k+1,da+1):
            n += 1
            for j in range(1,db):    
                for m in range(1,j+1):
                    cmsd[n-1,j-1] += (M[(l-1)*db+(m-1),(k-1)*db+(m-1)] + M[(k-1)*db+(m-1),(l-1)*db+(m-1)])
                m = j+1
                cmsd[n-1,j-1] -= j*(M[(l-1)*db+(m-1),(k-1)*db+(m-1)] + M[(k-1)*db+(m-1),(l-1)*db+(m-1)])
                cmsd[n-1,j-1] *= sqrt(2/(j*(j+1)))
    return cmsd.real

corr_mat_sd(2,2,rho2qb(a1, a2, a3, b1, b2, b3, c11, c12, c13, c21, c22, c23, c31, c32, c33))

def corr_mat_ad(da, db, M):
    cmad = np.zeros((da*(da-1)//2,db-1), dtype = complex)
    n = 0
    for k in range(1,da):
        for l in range(k+1,da+1):
            n += 1
            for j in range(1,db):    
                for m in range(1,j+1):
                    cmad[n-1,j-1] += (M[(l-1)*db+m-1,(k-1)*db+m-1] - M[(k-1)*db+m-1,(l-1)*db+m-1])
                m = j+1
                cmad[n-1,j-1] -= j*(M[(l-1)*db+m-1,(k-1)*db+m-1] - M[(k-1)*db+m-1,(l-1)*db+m-1])
                cmad[n-1,j-1] *= -1j*sqrt(2/(j*(j+1)))
    return cmad.real

corr_mat_ad(2,2,rho2qb(a1, a2, a3, b1, b2, b3, c11, c12, c13, c21, c22, c23, c31, c32, c33))

def corr_mat_ss(da, db, M):
    cmss = np.zeros((da*(da-1)//2,db*(db-1)//2), dtype = complex)
    p = 0
    for k in range(1,da):
        for l in range(k+1,da+1):
            p += 1
            q = 0
            for m in range(1,db):
                for n in range(m+1,db+1):
                    q += 1
                    cmss[p-1,q-1] += (M[(l-1)*db+n-1,(k-1)*db+m-1] + M[(k-1)*db+m-1,(l-1)*db+n-1])
                    cmss[p-1,q-1] += (M[(l-1)*db+m-1,(k-1)*db+n-1] + M[(k-1)*db+n-1,(l-1)*db+m-1])
    return cmss.real

corr_mat_ss(2,2,rho2qb(a1, a2, a3, b1, b2, b3, c11, c12, c13, c21, c22, c23, c31, c32, c33))

def corr_mat_sa(da, db, M):
    cmsa = np.zeros((da*(da-1)//2,db*(db-1)//2), dtype = complex)
    p = 0
    for k in range(1,da):
        for l in range(k+1,da+1):
            p += 1
            q = 0
            for m in range(1,db):
                for n in range(m+1,db+1):
                    q += 1
                    cmsa[p-1,q-1] -= 1j*(M[(l-1)*db+n-1,(k-1)*db+m-1] - M[(k-1)*db+m-1,(l-1)*db+n-1])
                    cmsa[p-1,q-1] -= 1j*(M[(k-1)*db+n-1,(l-1)*db+m-1] - M[(l-1)*db+m-1,(k-1)*db+n-1])
    return cmsa.real

simplify(corr_mat_sa(2,2,rho2qb(a1, a2, a3, b1, b2, b3, c11, c12, c13, c21, c22, c23, c31, c32, c33)))

def corr_mat_as(da, db, M):
    cmas = np.zeros((da*(da-1)//2,db*(db-1)//2), dtype = complex)
    p = 0
    for k in range(1,da):
        for l in range(k+1,da+1):
            p += 1
            q = 0
            for m in range(1,db):
                for n in range(m+1,db+1):
                    q += 1
                    cmas[p-1,q-1] -= 1j*(M[(l-1)*db+n-1,(k-1)*db+m-1] - M[(k-1)*db+m-1,(l-1)*db+n-1])
                    cmas[p-1,q-1] -= 1j*(M[(l-1)*db+m-1,(k-1)*db+n-1] - M[(k-1)*db+n-1,(l-1)*db+m-1])
    return cmas.real

simplify(corr_mat_as(2,2,rho2qb(a1, a2, a3, b1, b2, b3, c11, c12, c13, c21, c22, c23, c31, c32, c33)))

def corr_mat_aa(da, db, M):
    cmaa = np.zeros((da*(da-1)//2,db*(db-1)//2), dtype = complex)
    p = 0
    for k in range(1,da):
        for l in range(k+1,da+1):
            p += 1
            q = 0
            for m in range(1,db):
                for n in range(m+1,db+1):
                    q += 1
                    cmaa[p-1,q-1] += (M[(l-1)*db+m-1,(k-1)*db+n-1] + M[(k-1)*db+n-1,(l-1)*db+m-1])
                    cmaa[p-1,q-1] -= (M[(l-1)*db+n-1,(k-1)*db+m-1] + M[(k-1)*db+m-1,(l-1)*db+n-1])
    return cmaa.real

corr_mat_aa(2,2,rho2qb(a1, a2, a3, b1, b2, b3, c11, c12, c13, c21, c22, c23, c31, c32, c33))

def corr_mat(da, db, M):
    dda = int((da*(da-1))/2); ddb = int((db*(db-1))/2); cm = np.zeros((da**2-1,db**2-1))
    k = -1; l = -1
    cmdd = corr_mat_dd(da, db, M); cmds = corr_mat_ds(da, db, M); cmda = corr_mat_da(da, db, M)
    for m in range(0,da-1): 
        k += 1
        for n in range(0,db-1):  # diagonal-diagonal
            l += 1; cm[k,l] = cmdd[m,n]
        for n in range(0,ddb):  # diagonal-symmetric
            l += 1; cm[k,l] = cmds[m,n]
        for n in range(0,ddb):  # diagonal-antisymmetric
            l += 1; cm[k,l] = cmda[m,n]
    cmsd = corr_mat_sd(da, db, M); cmss = corr_mat_ss(da, db, M); cmsa = corr_mat_sa(da, db, M)
    l = -1
    for m in range(0,dda): 
        k += 1
        for n in range(0,db-1):  # symmetric-diagonal
            l += 1; cm[k,l] = cmsd[m,n]
        for n in range(0,ddb):  # symmetric-symmetric
            l += 1; cm[k,l] = cmss[m,n]
        for n in range(0,ddb):  # symmetric-antisymmetric
            l += 1; cm[k,l] = cmsa[m,n]
    cmad = corr_mat_ad(da, db, M); cmas = corr_mat_as(da, db, M); cmaa = corr_mat_aa(da, db, M)
    l = -1
    for m in range(0,dda): 
        k += 1
        for n in range(0,db-1):  # antisymmetric-diagonal
            l += 1; cm[k,l] = cmad[m,n]
        for n in range(0,ddb):  # antisymmetric-symmetric
            l += 1; cm[k,l] = cmas[m,n]
        for n in range(0,ddb):  # antisymmetric-antisymmetric
            l += 1; cm[k,l] = cmaa[m,n]
    return cm.real

simplify(corr_mat(2,2,rho2qb(a1, a2, a3, b1, b2, b3, c11, c12, c13, c21, c22, c23, c31, c32, c33)))

