from qiskit import *
import numpy as np
import math
from matplotlib import pyplot as plt
import qiskit
nshots = 8192
IBMQ.load_account()
provider = qiskit.IBMQ.get_provider(hub='ibm-q', group='open', project='main')
device = provider.get_backend('ibmq_belem')
simulator = Aer.get_backend('qasm_simulator')
from qiskit.tools.monitor import job_monitor
from qiskit.ignis.verification.tomography import state_tomography_circuits, StateTomographyFitter
#from qiskit.ignis.mitigation.measurement import complete_meas_cal, CompleteMeasFitter
from qiskit.visualization import plot_histogram

def qc_dqc1(ph):
    qr = QuantumRegister(3)
    qc = QuantumCircuit(qr, name='DQC1')
    qc.h(0)
    qc.h(1) # cria o estado de Bell dos qubits 1 e 2, que equivale ao estado de 1 sendo maximamente misto
    qc.cx(1,2)
    #qc.barrier()
    qc.cp(ph, 0, 1)
    return qc

ph = math.pi/2
qc_dqc1_ = qc_dqc1(ph)
qc_dqc1_.draw('mpl')

qc_dqc1_.decompose().draw('mpl')

def pTraceL_num(dl, dr, rhoLR): # Retorna traco parcial sobre L de 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): # Retorna traco parcial sobre R de 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

# simulation
phmax = 2*math.pi
dph = phmax/20
ph = np.arange(0,phmax+dph,dph)
d = len(ph);
xm = np.zeros(d)
ym = np.zeros(d)
for j in range(0,d):
    qr = QuantumRegister(3)
    qc = QuantumCircuit(qr)
    qc_dqc1_ = qc_dqc1(ph[j])
    qc.append(qc_dqc1_, [0,1,2])
    qstc = state_tomography_circuits(qc, [1,0])
    job = execute(qstc, backend=simulator, shots=nshots)
    qstf = StateTomographyFitter(job.result(), qstc)
    rho_01 = qstf.fit(method='lstsq')
    rho_0 = pTraceR_num(2, 2, rho_01)
    xm[j] = 2*rho_0[1,0].real
    ym[j] = 2*rho_0[1,0].imag
qc.draw('mpl')

# experiment
phmax = 2*math.pi
dph = phmax/10
ph_exp = np.arange(0,phmax+dph,dph)
d = len(ph_exp);
xm_exp = np.zeros(d)
ym_exp = np.zeros(d)
for j in range(0,d):
    qr = QuantumRegister(3)
    qc = QuantumCircuit(qr)
    qc_dqc1_ = qc_dqc1(ph_exp[j])
    qc.append(qc_dqc1_, [0,1,2])
    qstc = state_tomography_circuits(qc, [1,0])
    job = execute(qstc, backend=device, shots=nshots)
    print(job.job_id())
    job_monitor(job)
    qstf = StateTomographyFitter(job.result(), qstc)
    rho_01 = qstf.fit(method='lstsq')
    rho_0 = pTraceR_num(2, 2, rho_01)
    xm_exp[j] = 2*rho_0[1,0].real
    ym_exp[j] = 2*rho_0[1,0].imag

plt.plot(ph, xm, label = r'$\langle X\rangle_{sim}$')#, marker='*')
plt.plot(ph, ym, label = r'$\langle Y\rangle_{sim}$')#, marker='o')
plt.scatter(ph_exp, xm_exp, label = r'$\langle X\rangle_{exp}$', marker='*', color='r')
plt.scatter(ph_exp, ym_exp, label = r'$\langle Y\rangle_{exp}$', marker='o', color='g')
plt.legend(bbox_to_anchor=(1, 1))#,loc='center right')
#plt.xlim(0,2*math.pi)
#plt.ylim(-1,1)
plt.xlabel(r'$\phi$')
plt.show()

