%run init.ipynb

rp = np.arange(0,1.1,0.1); vth = 2*np.arccos(np.sqrt(rp)); rp, vth

def ket_rp(th,phi):
    return Matrix([[cos(th/2)],[(cos(ph)+1j*sin(ph))*sin(th/2)]])
def ket_rm(th,phi):
    return Matrix([[sin(th/2)],[-(cos(ph)+1j*sin(ph))*cos(th/2)]])

rp,th,ph = symbols('r_{+} theta phi', Real=True); ket_rp(th,ph), ket_rm(th,ph)

def rho(rp,rm,th,ph):
    return rp*proj(ket_rp(th,ph)) + rm*proj(ket_rm(th,ph))
rp,rm,th,ph = symbols('r_{+} r_{-} theta phi', Real=True); simplify(rho(rp,rm,th,ph))

rp = 1/2; vth = 2*math.acos(sqrt(rp)); th = 0*math.pi/4; ph = 0*math.pi/3
simplify(rho(rp,1-rp,th,ph))

print(rp, vth, th, ph)

import qiskit
qiskit.IBMQ.load_account()
provider = qiskit.IBMQ.get_provider(hub='ibm-q', group='open', project='main')
simulator = provider.get_backend('ibmq_qasm_simulator')
device = provider.get_backend('ibmq_belem')
nshots = 8192
from qiskit.ignis.verification.tomography import state_tomography_circuits, StateTomographyFitter
from qiskit.tools.monitor import job_monitor

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

qr = qiskit.QuantumRegister(2); qc = qiskit.QuantumCircuit(qr)
qc.u(vth,0,0,qr[0]); qc.cx(qr[0],qr[1]); qc.u(th,ph,math.pi,qr[0]);
qc.draw()

qstc = state_tomography_circuits(qc, [qr[0],qr[1]]); # circuito para TEQ
job = qiskit.execute(qstc, backend = simulator, shots = nshots) # executa no simulador
qstf = StateTomographyFitter(job.result(), qstc) # ajusta os dados
rhoBA = qstf.fit(method='lstsq'); rhoA = pTraceL_num(2, 2, rhoBA); rhoA

qc.draw()

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

a,b,c,d = symbols('a,b,c,d', Real=True, Positive=True);
z,w = symbols('z,w', Real=True)
rhoX = Matrix([[a,0,0,w],[0,b,z,0],[0,z,c,0],[w,0,0,d]]); rhoX

rhoX.eigenvals()

rhoX.eigenvects()

a,b,g = symbols('alpha beta gamma', Real = True, Positive = True)
p00 = (cos(a)*cos(b)*cos(g) + sin(a)*sin(b)*sin(g))**2
p01 = (cos(a)*cos(b)*sin(g) - sin(a)*sin(b)*cos(g))**2
p10 = (cos(a)*sin(b)*cos(g) - sin(a)*cos(b)*sin(g))**2
p11 = (cos(a)*sin(b)*sin(g) + sin(a)*cos(b)*cos(g))**2
p00,p01,p10,p11

simplify(p00+p01+p10+p11) # ok

th,ph = symbols('theta phi', Real = True, Positive = True)
psi00 = cos(th)*tp(cb(2,0),cb(2,0)) + sin(th)*tp(cb(2,1),cb(2,1))
psi10 = -sin(th)*tp(cb(2,0),cb(2,0)) + cos(th)*tp(cb(2,1),cb(2,1))
psi01 = sin(ph)*tp(cb(2,0),cb(2,1)) + cos(ph)*tp(cb(2,1),cb(2,0))
psi11 = cos(ph)*tp(cb(2,0),cb(2,1)) - sin(ph)*tp(cb(2,1),cb(2,0))
psi00, psi10, psi01, psi11

rho23 = p00*proj(psi00) + p01*proj(psi01) + p10*proj(psi10) + p11*proj(psi11);
simplify(rho23)

p00,p01,p10,p11 = symbols('p_{00},p_{01},p_{10},p_{11}')
rho23 = p00*proj(psi00) + p01*proj(psi01) + p10*proj(psi10) + p11*proj(psi11);
rho23

nonlinsolve([((cos(a)*cos(b)*cos(g) + sin(a)*sin(b)*sin(g))**2)*cos(th)**2 + (cos(a)*sin(b)*cos(g) - sin(a)*cos(b)*sin(g))**2*sin(th)**2 - a, 
       ((cos(a)*cos(b)*sin(g) - sin(a)*sin(b)*cos(g))**2)*sin(ph)**2 + ((cos(a)*sin(b)*sin(g) + sin(a)*cos(b)*cos(g))**2)*cos(ph)**2 - b, 
       ((cos(a)*cos(b)*sin(g) - sin(a)*sin(b)*cos(g))**2)*cos(ph)**2 + ((cos(a)*sin(b)*sin(g) + sin(a)*cos(b)*cos(g))**2)*sin(ph)**2 - c, 
       ((cos(a)*cos(b)*cos(g) + sin(a)*sin(b)*sin(g))**2)*sin(ph)**2 + (cos(a)*sin(b)*cos(g) - sin(a)*cos(b)*sin(g))**2*cos(th)**2 - d, 
       ((cos(a)*cos(b)*sin(g) - sin(a)*sin(b)*cos(g))**2-(cos(a)*sin(b)*sin(g) + sin(a)*cos(b)*cos(g))**2)*sin(ph)*cos(ph) - z,
       ((cos(a)*cos(b)*cos(g) + sin(a)*sin(b)*sin(g))**2-(cos(a)*sin(b)*cos(g) - sin(a)*cos(b)*sin(g))**2)*sin(th)*cos(th) - w],
      [th,ph,a,b,g])

