from qiskit import QuantumRegister, QuantumCircuit
# Quantum register
q = QuantumRegister(2, name='q')

# Quantum circuit
bell = QuantumCircuit(q)

# Create a Bell state
bell.x(q[0])
bell.x(q[1])
bell.h(q[0])
bell.cx(q[0], q[1])

# Draw circuit
bell.draw(output = 'mpl')

# Quantum register
q = QuantumRegister(2, name='q')

# Quantum circuit
entangled = QuantumCircuit(q)

# Entangled state
entangled.h(q[0])
entangled.s(q[0])
entangled.cx(q[0], q[1])
entangled.h(q[0])

# Draw circuit
entangled.draw(output='mpl')

# Quantum register
q = QuantumRegister(2, name='q')

# Quantum circuit
channel = QuantumCircuit(q)

# CNOT gate
channel.cx(q[0], q[1])

# Draw circuit
channel.draw(output='mpl')

# Quantum register
q = QuantumRegister(2, name='q')

# Quantum circuit
channel = QuantumCircuit(q)

# CNOT gate
channel.h(q[1])
channel.cz(q[0], q[1])
channel.cx(q[1], q[0])

# Draw circuit
channel.draw(output='mpl')

from qiskit import IBMQ, execute
from qiskit import QuantumCircuit, QuantumRegister, ClassicalRegister
from qiskit.tools.visualization import plot_histogram, plot_state_city

IBMQ.load_account()

backend = IBMQ.get_provider().get_backend('ibmqx2')

# qiskit-ignis provides tools for noise mitigation
from qiskit.ignis.mitigation.measurement import complete_meas_cal, CompleteMeasFitter, MeasurementFilter

# For visualizing the circuits
import matplotlib.pyplot as plt

# Suppose we have a 5-qubit quantum computer, and we want to perform measurements on qubits 0 and 2

# We first initialize a 5-qubit quantum register
q = QuantumRegister(5)

# Next, we generate the calibration circuits
meas_calibs, state_labels = complete_meas_cal(qubit_list=[0,2], qr=q)

# We can plot them 
fig, axes = plt.subplots(1, 4)
for i, circuit in enumerate(meas_calibs):
    circuit.draw(output='mpl', ax=axes[i], scale=2)
    
print("State labels:",state_labels)

# We now perform the calibration on the real device
job_calibration = execute(meas_calibs, backend=backend)
cal_results = job_calibration.result()

# We now calculate the calibration matrix from the outcomes
meas_fitter = CompleteMeasFitter(cal_results, state_labels)
meas_fitter.plot_calibration()

# What is the measurement fidelity?
print("Average Measurement Fidelity: %f" % meas_fitter.readout_fidelity())

# Make a 2Q Bell state between Q0 and Q2
c = ClassicalRegister(2)
bell = QuantumCircuit(q)

bell.x(q[0])
bell.x(q[2])
bell.h(q[0])
bell.cx(q[0], q[2])

meas = QuantumCircuit(q, c)
meas.measure(q[0],c[0])
meas.measure(q[2],c[1])
qc = bell + meas

qc.draw(output='mpl')

# We execute the circuit
job = execute(qc, backend=backend, shots=5000)
results = job.result()

# We also simulate the circuit in a noiseless case
from qiskit import Aer
simjob = execute(qc, backend=Aer.get_backend('qasm_simulator'), shots=10000)
simresults = simjob.result()

# We get the raw counts
raw_counts = results.get_counts()
sim_counts = simresults.get_counts()

# ... and the mitigated counts
# First we need to get the measurement filter object
meas_filter = meas_fitter.filter
# and then we apply it to the results
mitigated_results = meas_filter.apply(results)
mitigated_counts = mitigated_results.get_counts(0)

plot_histogram([sim_counts, raw_counts, mitigated_counts], legend=['simulated', 'raw', 'mitigated'])

# We import the required classes from ignis.
from qiskit.ignis.verification.tomography import state_tomography_circuits, StateTomographyFitter

from qiskit.tools.qi.qi import partial_trace

job_state_vec = execute(bell, Aer.get_backend('statevector_simulator'))
state_vec = partial_trace(job_state_vec.result().get_statevector(), [1, 3, 4])
plot_state_city(state_vec)

# We construct all the tomographic circuits
qst = state_tomography_circuits(bell, [q[0], q[2]])

# We can plot them
fig, axes = plt.subplots(3, 3)
for i, circuit in enumerate(qst):
    row = i // 3
    col = i % 3
    circuit.draw(output='mpl', ax=axes[row, col], scale=2)

# We execute the tomography circuits on the real device
job = execute(qst, backend=backend)

# We create a StateTomographyFitter from the job result and the tomography circuit
tomo_fitter = StateTomographyFitter(job.result(), qst)

# To obtain the 
rho_tomo = tomo_fitter.fit()

plot_state_city(rho_tomo)

plot_state_city(state_vec)

# Import the state_fidelity function
from qiskit.quantum_info import state_fidelity

# Exercise code:
# ...
# ...

