# General tools
import numpy as np
import matplotlib.pyplot as plt

# Qiskit Circuit Functions
from qiskit import  execute,QuantumCircuit, QuantumRegister, ClassicalRegister, Aer, transpile
import qiskit.quantum_info as qi

# Tomography functions
from qiskit.ignis.verification.tomography import process_tomography_circuits, ProcessTomographyFitter
from qiskit.ignis.mitigation.measurement import complete_meas_cal, CompleteMeasFitter

import warnings 
warnings.filterwarnings('ignore')

target = QuantumCircuit(2)
#target = # YOUR CODE HERE
target.h(0)
target.h(1)
target.rx(np.pi/2,0)
target.rx(np.pi/2,1)
target.cx(0,1)
target.p(np.pi,1)
target.cx(0,1)
target_unitary = qi.Operator(target)
target.draw('mpl')

from qc_grader import grade_lab5_ex1

# Note that the grading function is expecting a quantum circuit with no measurements
grade_lab5_ex1(target)

simulator = Aer.get_backend('qasm_simulator')
qpt_circs = process_tomography_circuits(target,measured_qubits=[0,1])# YOUR CODE HERE
qpt_job = execute(qpt_circs,simulator,seed_simulator=3145,seed_transpiler=3145,shots=8192)
qpt_result = qpt_job.result()





# YOUR CODE HERE
qpt_tomo= ProcessTomographyFitter(qpt_result,qpt_circs)
Choi_qpt= qpt_tomo.fit(method='lstsq')
fidelity = qi.average_gate_fidelity(Choi_qpt,target_unitary)

from qc_grader import grade_lab5_ex2

# Note that the grading function is expecting a floating point number
grade_lab5_ex2(fidelity)

# T1 and T2 values for qubits 0-3
T1s = [15000, 19000, 22000, 14000]
T2s = [30000, 25000, 18000, 28000]

# Instruction times (in nanoseconds)
time_u1 = 0   # virtual gate
time_u2 = 50  # (single X90 pulse)
time_u3 = 100 # (two X90 pulses)
time_cx = 300
time_reset = 1000  # 1 microsecond
time_measure = 1000 # 1 microsecond

from qiskit.providers.aer.noise import thermal_relaxation_error
from qiskit.providers.aer.noise import NoiseModel

# QuantumError objects
errors_reset = [thermal_relaxation_error(t1, t2, time_reset)
                for t1, t2 in zip(T1s, T2s)]
errors_measure = [thermal_relaxation_error(t1, t2, time_measure)
                  for t1, t2 in zip(T1s, T2s)]
errors_u1  = [thermal_relaxation_error(t1, t2, time_u1)
              for t1, t2 in zip(T1s, T2s)]
errors_u2  = [thermal_relaxation_error(t1, t2, time_u2)
              for t1, t2 in zip(T1s, T2s)]
errors_u3  = [thermal_relaxation_error(t1, t2, time_u3)
              for t1, t2 in zip(T1s, T2s)]
errors_cx = [[thermal_relaxation_error(t1a, t2a, time_cx).expand(
             thermal_relaxation_error(t1b, t2b, time_cx))
              for t1a, t2a in zip(T1s, T2s)]
               for t1b, t2b in zip(T1s, T2s)]

# Add errors to noise model
noise_thermal = NoiseModel()

for j in range(4):
    
    noise_thermal.add_quantum_error(errors_measure[j],'measure',[j])
    
    noise_thermal.add_quantum_error(errors_reset[j],'reset',[j])
    
    noise_thermal.add_quantum_error(errors_u3[j],'u3',[j])
    
    noise_thermal.add_quantum_error(errors_u2[j],'u2',[j])
    
    noise_thermal.add_quantum_error(errors_u1[j],'u1',[j])
    
    for k in range(4):
        
        noise_thermal.add_quantum_error(errors_cx[j][k],'cx',[j,k])

# YOUR CODE HERE

from qc_grader import grade_lab5_ex3

# Note that the grading function is expecting a NoiseModel
grade_lab5_ex3(noise_thermal)

np.random.seed(0)

noise_qpt_circs = process_tomography_circuits(target,measured_qubits=[0,1])# YOUR CODE HERE
noise_qpt_job = execute(noise_qpt_circs,simulator,seed_simulator=3145,seed_transpiler=3145,shots=8192,noise_model=noise_thermal)
noise_qpt_result = noise_qpt_job.result()
# YOUR CODE HERE
noise_qpt_tomo= ProcessTomographyFitter(noise_qpt_result,noise_qpt_circs)
noise_Choi_qpt= noise_qpt_tomo.fit(method='lstsq')
fidelity = qi.average_gate_fidelity(noise_Choi_qpt,target_unitary)

from qc_grader import grade_lab5_ex4

# Note that the grading function is expecting a floating point number
grade_lab5_ex4(fidelity)

np.random.seed(0)

# YOUR CODE HERE
#qr = QuantumRegister(2)

qubit_list = [0,1]

meas_calibs, state_labels = complete_meas_cal(qubit_list=qubit_list)

meas_calibs_job = execute(meas_calibs,simulator,seed_simulator=3145,seed_transpiler=3145,shots=8192,noise_model=noise_thermal)

cal_results = meas_calibs_job.result()

meas_fitter = CompleteMeasFitter(cal_results, state_labels, qubit_list=qubit_list)

Cal_result = meas_fitter.filter.apply(noise_qpt_result)

# YOUR CODE HERE
Cal_qpt_tomo= ProcessTomographyFitter(Cal_result,noise_qpt_circs)
Cal_Choi_qpt= Cal_qpt_tomo.fit(method='lstsq')
fidelity = qi.average_gate_fidelity(Cal_Choi_qpt,target_unitary)

from qc_grader import grade_lab5_ex5

# Note that the grading function is expecting a floating point number
grade_lab5_ex5(fidelity)

