# Needed for functions
import numpy as np
import time

# Import QISKit classes
import qiskit
from qiskit import QuantumRegister, QuantumCircuit, Aer
from qiskit.quantum_info import state_fidelity
from qiskit.tools.qi.qi import outer

# Tomography functions
from qiskit.ignis.verification.tomography import process_tomography_circuits, ProcessTomographyFitter

# Process tomography of a Hadamard gate
q = QuantumRegister(1)
circ = QuantumCircuit(q)
circ.h(q[0])

# Run circuit on unitary simulator to find ideal unitary
job = qiskit.execute(circ, Aer.get_backend('unitary_simulator'))
ideal_unitary = job.result().get_unitary(circ)
# convert to Choi-matrix in column-major convention
choi_ideal = outer(ideal_unitary.ravel(order='F'))

# Generate process tomography circuits and run on qasm simulator
qpt_circs = process_tomography_circuits(circ, q)
job = qiskit.execute(qpt_circs, Aer.get_backend('qasm_simulator'), shots=4000)

# Extract tomography data so that counts are indexed by measurement configuration
qpt_tomo = ProcessTomographyFitter(job.result(), qpt_circs)
qpt_tomo.data

# MLE Least-Squares tomographic reconstruction
t = time.time()
choi_lstsq = qpt_tomo.fit(method='lstsq')
print('Least-Sq Fitter')
print('fit time:', time.time() - t)
print('fit fidelity:', state_fidelity(choi_ideal / 2, choi_lstsq / 2))

# CVXOPT Semidefinite-Program tomographic reconstruction
t = time.time()
choi_cvx = qpt_tomo.fit(method='cvx')
print('\nCVXOPT Fitter')
print('fit time:', time.time() - t)
print('fit fidelity:', state_fidelity(choi_ideal / 2, choi_cvx / 2))

# Process tomography of a Hadamard gate
q = QuantumRegister(2)
circ = QuantumCircuit(q)
circ.swap(q[0], q[1])

# Ideal channel is a unitary
ideal_unitary = np.eye(2)
choi_ideal = outer(ideal_unitary.ravel(order='F'))

# Generate process tomography circuits and run on qasm simulator
# We use the optional prepared_qubits kwarg to specify that the prepared qubit was different to measured qubit
qpt_circs = process_tomography_circuits(circ, q[1], prepared_qubits=q[0])
job = qiskit.execute(qpt_circs, Aer.get_backend('qasm_simulator'), shots=2000)

# Extract tomography data so that counts are indexed by measurement configuration
qpt_tomo = ProcessTomographyFitter(job.result(), qpt_circs)
qpt_tomo.data

# Least-Squares tomographic reconstruction
t = time.time()
choi_lstsq = qpt_tomo.fit(method='lstsq')
print('Least-Sq Fitter')
print('fit time:', time.time() - t)
print('fit fidelity:', state_fidelity(choi_ideal / 2, choi_lstsq / 2))

# CVXOPT Semidefinite-Program tomographic reconstruction
t = time.time()
choi_cvx = qpt_tomo.fit(method='cvx')
print('\nCVXOPT Fitter')
print('fit time:', time.time() - t)
print('fit fidelity:', state_fidelity(choi_ideal / 2, choi_cvx / 2))

# Bell-state entangling circuit
q = QuantumRegister(2)
circ = QuantumCircuit(q)
circ.h(q[0])
circ.cx(q[0], q[1])

# Run circuit on unitary simulator to find ideal unitary
job = qiskit.execute(circ, Aer.get_backend('unitary_simulator'))
ideal_unitary = job.result().get_unitary(circ)
# convert to Choi-matrix in column-major convention
choi_ideal = outer(ideal_unitary.ravel(order='F'))

# Generate process tomography circuits and run on qasm simulator
qpt_circs = process_tomography_circuits(circ, q)
job = qiskit.execute(qpt_circs, Aer.get_backend('qasm_simulator'), shots=2000)

# Extract tomography data so that counts are indexed by measurement configuration
qpt_tomo = ProcessTomographyFitter(job.result(), qpt_circs)

t = time.time()
choi_lstsq = qpt_tomo.fit(method='lstsq')
print('Least-Sq Fitter')
print('fit time:', time.time() - t)
print('fit fidelity:', state_fidelity(choi_ideal / 4, choi_lstsq / 4))

t = time.time()
choi_cvx = qpt_tomo.fit(method='cvx')
print('\nCVXOPT Fitter')
print('fit time:', time.time() - t)
print('fit fidelity:', state_fidelity(choi_ideal / 4, choi_cvx / 4))

# Process tomography of a Hadamard gate
q = QuantumRegister(1)
circ = QuantumCircuit(q)
circ.h(q[0])

# Run circuit on unitary simulator to find ideal unitary
job = qiskit.execute(circ, Aer.get_backend('unitary_simulator'))
ideal_unitary = job.result().get_unitary(circ)
# convert to Choi-matrix in column-major convention
choi_ideal = outer(ideal_unitary.ravel(order='F'))

# Generate process tomography circuits and run on qasm simulator
qpt_circs = process_tomography_circuits(circ, q, prep_labels='SIC', prep_basis='SIC')
job = qiskit.execute(qpt_circs, Aer.get_backend('qasm_simulator'), shots=2000)

# Extract tomography data so that counts are indexed by measurement configuration
qpt_tomo = ProcessTomographyFitter(job.result(), qpt_circs, prep_basis='SIC')
qpt_tomo.data

# MLE Least-Squares tomographic reconstruction
t = time.time()
choi_lstsq = qpt_tomo.fit(method='lstsq')
print('Least-Sq Fitter')
print('fit time:', time.time() - t)
print('fit fidelity:', state_fidelity(choi_ideal / 2, choi_lstsq / 2))

# CVXOPT Semidefinite-Program tomographic reconstruction
t = time.time()
choi_cvx = qpt_tomo.fit(method='cvx')

print('\nCVXOPT Fitter')
print('fit time:', time.time() - t)
print('fit fidelity:', state_fidelity(choi_ideal / 2, choi_cvx / 2))