%run init.ipynb

from qiskit import *

X = Matrix([[0,1],[1,0]]); Y = Matrix([[0,-1j],[1j,0]]); Z = Matrix([[1,0],[0,-1]]); I = Matrix([[1,0],[0,1]])
xa,ya,za,xb,yb,zb = symbols('x_a y_a z_a x_b y_b z_b')
xx,xy,xz,yx,yy,yz,zx,zy,zz = symbols('xx xy xz yx yy yz zx zy zz')
rho_ab = (1/4)*(tp(I,I) + xa*tp(X,I) + ya*tp(Y,I) + za*tp(Z,I) + xb*tp(I,X) + yb*tp(I,Y) + zb*tp(I,Z))
rho_ab += (1/4)*(xx*tp(X,X) + xy*tp(X,Y) + xz*tp(X,Z) + yx*tp(Y,X) + yy*tp(Y,Y) + yz*tp(Y,Z))
rho_ab += (1/4)*(zx*tp(Z,X) + xy*tp(Z,Y) + zz*tp(Z,Z))
rho_ab

trace(rho_ab)

def ptraceA(da, db, rho):
    rhoB = zeros(db,db)
    for j in range(0, db):
        for k in range(0, db):
            for l in range(0, da):
                rhoB[j,k] += rho[l*db+j,l*db+k]
    return rhoB
def ptraceB(da, db, rho):
    rhoA = zeros(da,da)
    for j in range(0, da):
        for k in range(0, da):
            for l in range(0, db):
                rhoA[j,k] += rho[j*db+l,k*db+l]
    return rhoA

rho_a = ptraceB(2, 2, rho_ab); rho_a

rho_b = ptraceA(2, 2, rho_ab); rho_b

# alguns resultados prontos que usare de init.ipynb
cb(2,0), cb(2,1), id(2), pauli(1), proj(cb(2,0)), tp(cb(2,0),cb(2,0))

CX_ab = tp(proj(cb(2,0)),id(2)) + tp(proj(cb(2,1)),pauli(1)); CX_ab

CX_ba = tp(id(2),proj(cb(2,0))) + tp(pauli(1),proj(cb(2,1))); CX_ba

SWAP = CX_ab*CX_ba*CX_ab; SWAP

SWAP = CX_ba*CX_ab*CX_ba; SWAP

# aplica a troca ao estado de 2 qubits
rho_ab_swap = SWAP*rho_ab*SWAP # pois a SWAP é hermitiana

# verificação que a troca ocorreu
rho_a_swap = ptraceB(2, 2, rho_ab_swap); rho_a_swap

rho_b_swap = ptraceA(2, 2, rho_ab_swap); rho_b_swap

def pTraceL_num(dl, dr, rhoLR):
    # Returns the left partial trace over the 'left' subsystem of 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):
    # Returns the right partial trace over the 'right' subsystem of 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

qr = QuantumRegister(2); qc = QuantumCircuit(qr)

qc.h(qr[1]); qc.s(qr[1]); qc.draw()

from qiskit.ignis.verification.tomography import state_tomography_circuits, StateTomographyFitter
nshots = 8192
qstc = state_tomography_circuits(qc, [qr[0],qr[1]]); # circuito para TEQ
job = qiskit.execute(qstc, Aer.get_backend('qasm_simulator'), shots=nshots) # executa no simulador
qstf = StateTomographyFitter(job.result(), qstc) # ajusta os dados
rhoBA = qstf.fit(method='lstsq'); rhoBA # extrai o operador densidade

rhoA = pTraceL_num(2, 2, rhoBA); rhoA

rhoB = pTraceR_num(2, 2, rhoBA); rhoB

qc.cx(qr[0],qr[1]); qc.cx(qr[0],qr[1]); qc.draw()

# data contém todas as portas lógicas do cicuito quântico em uma lista (ou dicionário)
data = qc.data; data

# imprime uma componente de qc.data
print( data[3] )

# deleta uma componente de data
data.pop(3)

data.pop(4)

data.pop(3)

# agora o circuit está como queremos
qc.draw()

qc.cx(qr[1],qr[0]); qc.draw()

data = qc.data; data.pop(3)

qc.cx(qr[0],qr[1]); qc.draw()

qstc = state_tomography_circuits(qc, [qr[0],qr[1]]); # circuito para TEQ
job = qiskit.execute(qstc, Aer.get_backend('qasm_simulator'), shots=nshots) # executa no simulador
qstf = StateTomographyFitter(job.result(), qstc) # ajusta os dados
rhoBA = qstf.fit(method='lstsq'); # extrai o operador densidade

rhoA = pTraceL_num(2, 2, rhoBA); rhoA

rhoB = pTraceR_num(2, 2, rhoBA); rhoB

qr = QuantumRegister(2); qc = QuantumCircuit(qr)
qc.h(qr[1]); qc.s(qr[1]); qc.barrier(); qc.draw()

qiskit.IBMQ.load_account();
provider = IBMQ.get_provider(hub = 'ibm-q-research-2', group = 'federal-uni-sant-1', project = 'main')
device = provider.get_backend('ibmq_belem')
from qiskit.tools.monitor import job_monitor

qstc = state_tomography_circuits(qc, [qr[0],qr[1]]); # circuito para TEQ
job = qiskit.execute(qstc, backend = device, shots = nshots)
print(job.job_id()); job_monitor(job)

qstf = StateTomographyFitter(job.result(), qstc) # ajusta os dados
rhoBA = qstf.fit(method='lstsq'); # extrai o operador densidade
rhoA = pTraceL_num(2, 2, rhoBA); rhoA

rhoB = pTraceR_num(2, 2, rhoBA); rhoB

qc.cx(qr[0],qr[1]); qc.cx(qr[1],qr[0]); qc.cx(qr[0],qr[1]); 

data = qc.data; data.pop(5)

qc.draw()

qstc = state_tomography_circuits(qc, [qr[0],qr[1]]); # circuito para TEQ
job = qiskit.execute(qstc, backend = device, shots = nshots)
print(job.job_id()); job_monitor(job)

qstf = StateTomographyFitter(job.result(), qstc) # ajusta os dados
rhoBA = qstf.fit(method='lstsq'); # extrai o operador densidade
rhoA = pTraceL_num(2, 2, rhoBA); rhoA

rhoB = pTraceR_num(2, 2, rhoBA); rhoB

def psi(th):
    return cos(th/2)*cb(2,0) + sin(th/2)*cb(2,1)
th, ph = symbols('theta phi'); psi(th)

def Psi(th,ph):
    return CX_ab*tp(psi(th),psi(ph))
th, ph = symbols('theta phi'); Psi(th,ph)

th, ph = math.pi/3, math.pi/4; Psi(th,ph)

def rhoA(th,ph):
    return ptraceB(2, 2, proj(Psi(th,ph)))
th, ph = symbols('theta phi'); simplify(rhoA(th,ph))

th, ph = math.pi/3, math.pi/4; simplify(rhoA(th,ph))

rhoa = rhoA(th,ph); rhoa.eigenvals()

def rhoB(th,ph):
    return ptraceA(2, 2, proj(Psi(th,ph)))
th, ph = symbols('theta phi'); simplify(rhoB(th,ph))

th, ph = math.pi/3, math.pi/4; simplify(rhoB(th,ph))

rhob = rhoB(th,ph); rhob.eigenvals()

qr = QuantumRegister(2); qc = QuantumCircuit(qr)

th = math.pi/3; ph = math.pi/4
qc.u(th,0,0,qr[0]); qc.u(ph,0,0,qr[1]); qc.cx(qr[0],qr[1]); qc.barrier();

qc.draw()

data = qc.data; data.pop(1);

qstc = state_tomography_circuits(qc, [qr[0],qr[1]]); # circuito para TEQ
job = qiskit.execute(qstc, Aer.get_backend('qasm_simulator'), shots=nshots) # executa no simulador
qstf = StateTomographyFitter(job.result(), qstc) # ajusta os dados
rhoBA = qstf.fit(method='lstsq'); # extrai o operador densidade

rhoA = pTraceL_num(2, 2, rhoBA); rhoA

rhoB = pTraceR_num(2, 2, rhoBA); rhoB

qstc = state_tomography_circuits(qc, [qr[0],qr[1]]); # circuito para TEQ
job = qiskit.execute(qstc, backend = device, shots = nshots)
print(job.job_id()); job_monitor(job)

rhoA = pTraceL_num(2, 2, rhoBA); rhoA

rhoB = pTraceR_num(2, 2, rhoBA); rhoB

qc.cx(qr[1],qr[0]); qc.cx(qr[0],qr[1]); qc.cx(qr[1],qr[0])
qc.draw()

qstc = state_tomography_circuits(qc, [qr[0],qr[1]]); # circuito para TEQ
job = qiskit.execute(qstc, Aer.get_backend('qasm_simulator'), shots=nshots) # executa no simulador
qstf = StateTomographyFitter(job.result(), qstc) # ajusta os dados
rhoBA = qstf.fit(method='lstsq'); # extrai o operador densidade

rhoA = pTraceL_num(2, 2, rhoBA); rhoA

rhoB = pTraceR_num(2, 2, rhoBA); rhoB

qstc = state_tomography_circuits(qc, [qr[0],qr[1]]); # circuito para TEQ
job = qiskit.execute(qstc, backend = device, shots = nshots)
print(job.job_id()); job_monitor(job)

rhoA = pTraceL_num(2, 2, rhoBA); rhoA

rhoB = pTraceR_num(2, 2, rhoBA); rhoB

