from qiskit import execute, Aer
from qiskit import QuantumCircuit, QuantumRegister

from qiskit.providers import JobStatus
from qiskit.tools.qi.qi import partial_trace
from qiskit.quantum_info import state_fidelity
from qiskit.compiler import transpile
from qiskit.quantum_info.random.utils import random_state

import time

from pairwise_tomography.pairwise_state_tomography_circuits import pairwise_state_tomography_circuits
from pairwise_tomography.pairwise_fitter import PairwiseStateTomographyFitter
from pairwise_tomography.utils import concurrence

backend = Aer.get_backend('qasm_simulator')

nq = 5

q = QuantumRegister(nq)
qc = QuantumCircuit(q)

qc.ry(3., q[0])
qc.cx(q[0], q[1])
qc.x(q[1])
qc.ry(2.1, q[2])
qc.cx(q[2], q[3])
qc.h(q[4])
qc.cx(q[4], q[1])
qc.cx(q[1], q[4])
qc.cx(q[4], q[3])
qc.cx(q[3], q[4])

qc.draw(output='mpl')

pw_tomo_circs = pairwise_state_tomography_circuits(qc, q)

print(len(pw_tomo_circs))

pw_tomo_circs[10].draw(output='mpl')

job = execute(pw_tomo_circs, Aer.get_backend('qasm_simulator'), shots=8192)

fitter = PairwiseStateTomographyFitter(job.result(), pw_tomo_circs, q)

fit_result = fitter.fit()

print(fit_result[(0,1)])

pairwise_entanglement = {key: concurrence(value) for key, value in fit_result.items()}

print(pairwise_entanglement)

from pairwise_tomography.visualization import draw_entanglement_graph

draw_entanglement_graph(pairwise_entanglement)

