import numpy as np
import matplotlib.pyplot as plt

import qiskit
from qiskit.providers.aer.noise.errors.standard_errors import coherent_unitary_error
from qiskit.providers.aer.noise import NoiseModel

from qiskit.ignis.characterization.hamiltonian import ZZFitter, zz_circuits

from qiskit.ignis.characterization.gates import (AmpCalFitter, ampcal_1Q_circuits, 
                                                 AngleCalFitter, anglecal_1Q_circuits, 
                                                 AmpCalCXFitter, ampcal_cx_circuits, 
                                                 AngleCalCXFitter, anglecal_cx_circuits)

# ZZ rates are typically ~ 100kHz so we want Ramsey oscillations around 1MHz
# 12 numbers ranging from 10 to 1000, logarithmically spaced
# extra point at 1500
num_of_gates = np.arange(0,150,5)
gate_time = 0.1

# Select the qubits whose ZZ will be measured
qubits = [0]
spectators = [1]

# Generate experiments
circs, xdata, osc_freq = zz_circuits(num_of_gates, gate_time, qubits, spectators, nosc=2)

# Set the simulator with ZZ
zz_unitary = np.eye(4,dtype=complex)
zz_unitary[3,3] = np.exp(1j*2*np.pi*0.02*gate_time)
error = coherent_unitary_error(zz_unitary)
noise_model = NoiseModel()
noise_model.add_nonlocal_quantum_error(error, 'id', [0], [0,1])

# Run the simulator
backend = qiskit.Aer.get_backend('qasm_simulator')
shots = 500
# For demonstration purposes split the execution into two jobs
print("Running the first 20 circuits")
backend_result1 = qiskit.execute(circs[0:20], backend,
                                shots=shots, noise_model=noise_model).result()
print("Running the rest of the circuits")
backend_result2 = qiskit.execute(circs[20:], backend,
                                shots=shots, noise_model=noise_model).result()

%matplotlib inline
# Fit the data to an oscillation

plt.figure(figsize=(10, 6))
initial_a = 1
initial_c = 0
initial_f = osc_freq 
initial_phi = -np.pi/20

# Instantiate the fitter
# pass the 2 results in as a list of results
fit = ZZFitter([backend_result1, backend_result2], xdata, qubits, spectators,
                   fit_p0=[initial_a, initial_f, initial_phi, initial_c],
                   fit_bounds=([-0.5, 0, -np.pi, -0.5],
                               [1.5, 2*osc_freq, np.pi, 1.5]))

fit.plot_ZZ(0, ax=plt.gca())
    
print("ZZ Rate: %f kHz"%(fit.ZZ_rate()[0]*1e3))
plt.show()

qubits = [4,2]
circs, xdata = ampcal_1Q_circuits(10, qubits)

print(circs[2])

# Set the simulator 
# Add a rotation error
err_unitary = np.zeros([2,2],dtype=complex)
angle_err = 0.1
for i in range(2):
    err_unitary[i,i] = np.cos(angle_err)
    err_unitary[i,(i+1) % 2] = np.sin(angle_err)
err_unitary[0,1] *= -1.0

error = coherent_unitary_error(err_unitary)
noise_model = NoiseModel()
noise_model.add_all_qubit_quantum_error(error, 'u2')

# Run the simulator
backend = qiskit.Aer.get_backend('qasm_simulator')
shots = 500
backend_result1 = qiskit.execute(circs, backend,
                                shots=shots, noise_model=noise_model).result()

%matplotlib inline
# Fit the data to an oscillation

plt.figure(figsize=(10, 6))
initial_theta = 0.02
initial_c = 0.5
initial_phi = 0.1


fit = AmpCalFitter(backend_result1, xdata, qubits,
                   fit_p0=[initial_theta, initial_c],
                   fit_bounds=([-np.pi, -1],
                               [np.pi, 1]))

# plot the result for the number 1 indexed qubit.
# In this case that refers to Q2 since we passed in as [4, 2])
fit.plot(1, ax=plt.gca())
    
print("Rotation Error on U2: %f rads"%(fit.angle_err()[0]))
plt.show()

qubits = [0,1]
circs, xdata = anglecal_1Q_circuits(10, qubits, angleerr=0.1)

#The U1 gates are added errors to test the procedure
print(circs[2])

# Set the simulator 
# Run the simulator
backend = qiskit.Aer.get_backend('qasm_simulator')
shots = 1000
backend_result1 = qiskit.execute(circs, backend,
                                shots=shots).result()

%matplotlib inline
# Fit the data to an oscillation

plt.figure(figsize=(10, 6))
initial_theta = 0.02
initial_c = 0.5
initial_phi = 0.01


fit = AngleCalFitter(backend_result1, xdata, qubits,
                   fit_p0=[initial_theta, initial_c],
                   fit_bounds=([-np.pi, -1],
                               [np.pi, 1]))

fit.plot(0, ax=plt.gca())
    
print("Angle error between X and Y: %f rads"%(fit.angle_err()[0]))
plt.show()

# We can specify more than one CX gate to calibrate in parallel
# but these lists must be the same length and not contain
# any duplicate elements
qubits = [0,2]
controls = [1,3]
circs, xdata = ampcal_cx_circuits(15, qubits, controls)

print(circs[2])

# Set the simulator 
# Add a rotation error on CX
# only if the control is in the excited state
err_unitary = np.eye(4,dtype=complex)
angle_err = 0.15
for i in range(2):
    err_unitary[2+i,2+i] = np.cos(angle_err)
    err_unitary[2+i,2+(i+1) % 2] = -1j*np.sin(angle_err)

error = coherent_unitary_error(err_unitary)
noise_model = NoiseModel()
noise_model.add_nonlocal_quantum_error(error, 'cx', [1,0], [0,1])

# Run the simulator
backend = qiskit.Aer.get_backend('qasm_simulator')
shots = 1500
backend_result1 = qiskit.execute(circs, backend,
                                shots=shots, noise_model=noise_model).result()

%matplotlib inline
# Fit the data to an oscillation

plt.figure(figsize=(10, 6))
initial_theta = 0.02
initial_c = 0.5
initial_phi = 0.01


fit = AmpCalCXFitter(backend_result1, xdata, qubits,
                   fit_p0=[initial_theta, initial_c],
                   fit_bounds=([-np.pi, -1],
                               [np.pi, 1]))

fit.plot(0, ax=plt.gca())
    
print("Rotation Error on CX: %f rads"%(fit.angle_err()[0]))
plt.show()

qubits = [0,2]
controls = [1,3]
circs, xdata = anglecal_cx_circuits(15, qubits, controls, angleerr=0.1)

print(circs[2])

# Set the simulator 
# Run the simulator
backend = qiskit.Aer.get_backend('qasm_simulator')
shots = 1000
backend_result1 = qiskit.execute(circs, backend,
                                shots=shots).result()

%matplotlib inline
# Fit the data to an oscillation

plt.figure(figsize=(10, 6))
initial_theta = 0.02
initial_c = 0.5
initial_phi = 0.01


fit = AngleCalCXFitter(backend_result1, xdata, qubits,
                   fit_p0=[initial_theta, initial_c],
                   fit_bounds=([-np.pi, -1],
                               [np.pi, 1]))

fit.plot(0, ax=plt.gca())
    
print("Rotation Error on CX: %f rads"%(fit.angle_err()[0]))
plt.show()