%run init.ipynb
%run 08_gellmann.ipynb

def coh_l1(rho):
    d = rho.shape[0]; C = 0
    for j in range(0,d-1):
        for k in range(j+1,d):
            C += np.abs(rho[j,k])
    return 2*C

rho = np.array([[1,0],[0,0]]); rho, coh_l1(rho)

rho = np.array([[1,1],[1,1]])/2; rho, coh_l1(rho)

def mutual_info(da, db, rhoAB):
    rhoA = pTraceR_num(da, db, rhoAB); Sa = von_neumann_num(rhoA)
    rhoB = pTraceL_num(da, db, rhoAB); Sb = von_neumann_num(rhoB)
    Sab = von_neumann_num(rhoAB) 
    return Sa + Sb - Sab

def shannon_num(pv):
    d = pv.shape[0]; SE = 0.0; j = -1
    while (j < d-1):
        j = j + 1
        if pv[j] > 10**-15 and pv[j] < (1.0-10**-15):
            SE -= pv[j]*math.log(pv[j], 2)
    return SE

import scipy.linalg.lapack as lapak
def von_neumann_num(rho):
    d = rho.shape[0]; b = lapak.zheevd(rho)
    return shannon_num(b[0])

def pTraceL_num(dl, dr, rhoLR):
    rhoR = np.zeros((dr, dr), dtype=complex)
    for j in range(0, dr):
        for k in range(j, dr):
            for l in range(0, dl):
                rhoR[j,k] += rhoLR[l*dr+j,l*dr+k]
            if j != k:
                rhoR[k,j] = np.conj(rhoR[j,k])
    return rhoR
def pTraceR_num(dl, dr, rhoLR):
    rhoL = np.zeros((dl, dl), dtype=complex)
    for j in range(0, dl):
        for k in range(j, dl):
            for l in range(0, dr):
                rhoL[j,k] += rhoLR[j*dr+l,k*dr+l]
        if j != k:
            rhoL[k,j] = np.conj(rhoL[j,k])
    return rhoL

rho = proj(tp(cb(2,0),cb(2,0))); rho, mutual_info(2, 2, rho)

rho = 0.5*proj(tp(cb(2,0),cb(2,0))) + 0.5*proj(tp(cb(2,1),cb(2,1))); rho, mutual_info(2, 2, rho)

psi = (1/math.sqrt(2))*(tp(cb(2,0),cb(2,0))+tp(cb(2,1),cb(2,1)))
rho = proj(psi); rho, mutual_info(2, 2, rho)

import scipy
def Dhellinger(da, db, rhoAB):
    M = scipy.linalg.sqrtm(rhoAB)
    A = pTraceR_num(da, db, M); bva = bloch_vector(A)/math.sqrt(2*db)
    B = pTraceL_num(da, db, M); bvb = bloch_vector(B)/2
    cm = corr_mat(da, db, M)/2; ev = lapak.zheevd(bva*bva.T + cm*cm.T)
    D_he = 2 - 2*math.sqrt((np.trace(A)/math.sqrt(2*db))**2 + (scipy.linalg.norm(bvb))**2 + max(ev[0].real))
    D_he = D_he/(2 - 2/math.sqrt(da))
    return D_he

da = 2; db = 2
M = scipy.linalg.sqrtm(np.array([[1,0,0,0],[0,1,0,0],[0,0,1,0],[0,0,0,1]])/4)
A = pTraceR_num(2, 2, M); bva = bloch_vector(A)/math.sqrt(2*2); bva
B = pTraceL_num(da, db, M); bvb = bloch_vector(B)/2; bvb
cm = corr_mat(da, db, M)/2; ev = lapak.zheevd(bva*bva.T + cm*cm.T); cm
D_he = 2 - 2*math.sqrt((np.trace(A)/math.sqrt(2*db))**2 + (scipy.linalg.norm(bvb))**2 + max(ev[0].real))
D_he = D_he/(2 - 2/math.sqrt(da)); D_he

# teste para a raiz quadrada de matriz calculada usando scipy
A = np.array([[1,0],[0,-1]]); Asr = scipy.linalg.sqrtm(A); Asr # raiz quadrada (testes)
#Asr*Asr

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]])

ev = lapak.zheevd(Pauli(1)); ev, ev[0], ev[0][0], ev[1][0], ev[1][0][0]

def Econcurrence(rho):
    R = rho@tp(Pauli(2),Pauli(2))@np.matrix.conjugate(rho)@tp(Pauli(2),Pauli(2))
    ev = lapak.zheevd(R)
    evm = max(abs(ev[0][0]), abs(ev[0][1]), abs(ev[0][2]), abs(ev[0][3]))
    Ec = 2*math.sqrt(abs(evm)) - math.sqrt(abs(ev[0][0])) - math.sqrt(abs(ev[0][1]))\
    - math.sqrt(abs(ev[0][2])) - math.sqrt(abs(ev[0][3]))
    if Ec < 0.0:
        Ec = 0.0
    return Ec

# Emaranhamento de formação
def Eof(rho):
    Ec = Econcurrence(rho)
    pv = np.zeros(2); pv[0] = (1+math.sqrt(1-Ec**2))/2; pv[1] = 1 - pv[0]
    return shannon_num(pv)

pv=np.array([1,2]); pv.shape

x = np.arange(10**-3,1,0.01); p0 = (1 + np.sqrt(1-x**2))/2; p1 = 1 - p0
ef = np.zeros(x.shape[0]); ef = -p0*np.log2(p0) - p1*np.log2(p1)
plt.plot(x , ef); plt.xlabel(r'$E_{c}$'); plt.ylabel(r'$E_{f}$'); plt.show()

def rho_werner(w):
    rho = np.zeros((4,4), dtype = complex)
    rho =+ ((1-w)/4)*np.identity(4) + (w/2)*np.array([[1,0,0,1],[0,0,0,0],[0,0,0,0],[1,0,0,1]])
    return rho
rho_werner(0)

w = np.arange(0, 1.05, 0.05); C = np.zeros(w.shape[0]); MI = np.zeros(w.shape[0]); D_he = np.zeros(w.shape[0])
Ec = np.zeros(w.shape[0]); Ef = np.zeros(w.shape[0])
for j in range(0, w.shape[0]):
    rho = rho_werner(w[j]); C[j] = coh_l1(rho); MI[j] = mutual_info(2, 2, rho)/2
    D_he[j] = Dhellinger(2, 2, rho); Ec[j] = Econcurrence(rho); Ef[j] = Eof(rho)

matplotlib.rcParams.update({'font.size':12})
plt.figure(figsize = (6,4), dpi = 100)
plt.plot(w, C, label = r'$C$'); plt.plot(w, MI, label = r'$I/2$'); 
plt.plot(w, D_he, label = r'$D_{A}^{he}$')
plt.plot(w, Ec, '*', label = r'$E_{c}$'); plt.plot(w, Ef, '-.', label = r'$E_{f}$')
plt.xlabel(r'$w$'); plt.legend(); plt.show()

