# Needed for functions
import numpy as np
import time

# Import Qiskit classes
import qiskit 
from qiskit import QuantumRegister, QuantumCircuit, ClassicalRegister, Aer
from qiskit.quantum_info import state_fidelity
from qiskit.providers.aer import noise

# Tomography functions
from qiskit.ignis.verification.tomography import state_tomography_circuits, StateTomographyFitter
import qiskit.ignis.mitigation.measurement as mc

# Create the expected density matrix
q2 = QuantumRegister(2)
bell = QuantumCircuit(q2)
bell.h(q2[0])
bell.cx(q2[0], q2[1])
print(bell)

job = qiskit.execute(bell, Aer.get_backend('statevector_simulator'))
psi_bell = job.result().get_statevector(bell)
print(psi_bell)

# Create the actual circuit 
q2 = QuantumRegister(6)
bell = QuantumCircuit(q2)
bell.h(q2[3])
bell.cx(q2[3], q2[5])
print(bell)

# Generate circuits and run on simulator
t = time.time()
# Generate the state tomography circuits. Only pass in the 
# registers we want to measure (in this case 3 and 5)
qst_bell = state_tomography_circuits(bell, [q2[3],q2[5]])
job = qiskit.execute(qst_bell, Aer.get_backend('qasm_simulator'), shots=5000)
print('Time taken:', time.time() - t)

# Generate the state tomography circuits using the default settings for 
# basis

tomo_bell = StateTomographyFitter(job.result(), qst_bell)

# Perform the tomography fit
# which outputs a density matrix
rho_bell = tomo_bell.fit()
F_bell = state_fidelity(psi_bell, rho_bell)
print('Fit Fidelity =', F_bell)

#Add measurement noise
noise_model = noise.NoiseModel()
for qi in range(6):
    read_err = noise.errors.readout_error.ReadoutError([[0.75, 0.25],[0.1,0.9]])
    noise_model.add_readout_error(read_err,[qi])
    
#generate the calibration circuits
meas_calibs, state_labels = mc.complete_meas_cal(qubit_list=[3,5])

backend = Aer.get_backend('qasm_simulator')
qobj_cal = qiskit.compile(meas_calibs, backend=backend, shots=15000)
qobj_tomo = qiskit.compile(qst_bell, backend=backend, shots=15000)

job_cal = backend.run(qobj_cal, noise_model=noise_model)
job_tomo = backend.run(qobj_tomo, noise_model=noise_model)

meas_fitter = mc.CompleteMeasFitter(job_cal.result(),state_labels)

tomo_bell = StateTomographyFitter(job_tomo.result(), qst_bell)

#no correction
rho_bell = tomo_bell.fit()
F_bell = state_fidelity(psi_bell, rho_bell)
print('Fit Fidelity (no correction) =', F_bell)

#correct data
correct_tomo_results = meas_fitter.filter.apply(job_tomo.result(), method='least_squares')
tomo_bell = StateTomographyFitter(correct_tomo_results, qst_bell)
rho_bell = tomo_bell.fit()
F_bell = state_fidelity(psi_bell, rho_bell)
print('Fit Fidelity (w/ correction) =', F_bell)

def random_u3_tomo(nq, shots):
    
    def rand_angles():
        return tuple(2 * np.pi * np.random.random(3) - np.pi)
    q = QuantumRegister(nq)
    circ = QuantumCircuit(q)
    for j in range(nq):
        circ.u3(*rand_angles(), q[j])
    job = qiskit.execute(circ, Aer.get_backend('statevector_simulator'))
    psi_rand = job.result().get_statevector(circ)
    
    qst_circs = state_tomography_circuits(circ, q)
    job = qiskit.execute(qst_circs, Aer.get_backend('qasm_simulator'),
                         shots=shots)
    tomo_data = StateTomographyFitter(job.result(), qst_circs)
    rho_cvx = tomo_data.fit(method='cvx')
    rho_lstsq = tomo_data.fit(method='lstsq')
    
    print('F fit (CVX) =', state_fidelity(psi_rand, rho_cvx))
    print('F fit (LSTSQ) =', state_fidelity(psi_rand, rho_lstsq))

for j in range(5):
    print('Random single-qubit unitaries: set {}'.format(j))
    random_u3_tomo(3, 5000)

# Create a state preparation circuit
q5 = QuantumRegister(5)
bell5 = QuantumCircuit(q5)
bell5.h(q5[0])
for j in range(4):
    bell5.cx(q5[j], q5[j + 1])

# Get ideal output state
job = qiskit.execute(bell5, Aer.get_backend('statevector_simulator'))
psi_bell5 = job.result().get_statevector(bell5)

# Generate circuits and run on simulator
t = time.time()
qst_bell5 = state_tomography_circuits(bell5, q5)
job = qiskit.execute(qst_bell5, Aer.get_backend('qasm_simulator'), shots=5000)

# Extract tomography data so that counts are indexed by measurement configuration
tomo_bell5 = StateTomographyFitter(job.result(), qst_bell5)
print('Time taken:', time.time() - t)

t = time.time()
rho_lstsq_bell5 = tomo_bell5.fit(method='lstsq')
print('Least-Sq Reconstruction')
print('Time taken:', time.time() - t)
print('Fit Fidelity:', state_fidelity(psi_bell5, rho_lstsq_bell5))

t = time.time()
rho_cvx_bell5 = tomo_bell5.fit(method='cvx')
print('CVX Reconstruction')
print('Time taken:', time.time() - t)
print('Fidelity:', state_fidelity(psi_bell5, rho_cvx_bell5))