import numpy as np
import matplotlib.pyplot as plt

import qiskit
from qiskit.providers.aer.noise.errors.standard_errors import amplitude_damping_error
from qiskit.providers.aer.noise.errors.standard_errors import phase_damping_error
from qiskit.providers.aer.noise import NoiseModel

from qiskit.ignis.characterization.coherence import T1Fitter, T2StarFitter, T2Fitter
from qiskit.ignis.characterization.coherence import t1_circuits, t2_circuits, t2star_circuits

# 12 numbers ranging from 10 to 1000, logarithmically spaced
# extra point at 1500
num_of_gates = np.append((np.logspace(1, 3, 12)).astype(int), np.array([1500]))
gate_time = 0.1

# Select the qubits whose T1 are to be measured
qubits = [0]

# Generate experiments
circs, xdata = t1_circuits(num_of_gates, gate_time, qubits)

# Set the simulator with amplitude damping noise
t1 = 25.0
gamma = 1 - np.exp(-gate_time/t1)
error = amplitude_damping_error(gamma)
noise_model = NoiseModel()
noise_model.add_quantum_error(error, 'id', [0])

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

%matplotlib inline
# Fit the data to an exponential
# The correct answers are a=1, and c=0, and t1=25/15 for qubit 0/2
# The user does not know the correct answer exactly,
# so starts the fit from a different but close location

initial_t1 = t1*1.2
initial_a = 1.0
initial_c = 0.0

fit = T1Fitter(backend_result, xdata, qubits,
               fit_p0=[initial_a, initial_t1, initial_c],
               fit_bounds=([0, 0, -1], [2, initial_t1*2, 1]))

fit.plot(0)

# 50 points linearly spaced in two regions (fine and coarse)
# 30 from 10->150, 20 from 160->450
num_of_gates = np.append((np.linspace(10, 150, 30)).astype(int), (np.linspace(160,450,20)).astype(int))
gate_time = 0.1

# Select the qubits whose T2* are to be measured
qubits = [0]

# Generate experiments
circs, xdata, osc_freq = t2star_circuits(num_of_gates, gate_time, qubits, nosc=5)

backend = qiskit.Aer.get_backend('qasm_simulator')

# Set the simulator with phase damping noise
t2 = 10
p = 1 - np.exp(-2*gate_time/t2)
error = phase_damping_error(p)
noise_model = NoiseModel()
noise_model.add_quantum_error(error, 'id', [0])

# Run the simulator
shots = 300
backend_result = qiskit.execute(circs, backend,
                                shots=shots, noise_model=noise_model).result()

%matplotlib inline
# Fit the data to an oscillator
# The correct answers are a=0.5, f=osc_freq, phi=0, c=0.5, and t2=10/5 for qubit 0/2
# The user does not know the correct answer exactly,
# so starts the fit from a different but close location

initial_t2 = t2*1.1
initial_a = 0.5
initial_c = 0.5
initial_f = osc_freq 
initial_phi = -np.pi/20

fit = T2StarFitter(backend_result, xdata, qubits,
                   fit_p0=[initial_a, initial_t2, initial_f, initial_phi, initial_c],
                   fit_bounds=([-0.5, 0, 0, -np.pi, -0.5],
                               [1.5, 2*t2, 2*osc_freq, np.pi, 1.5]))

fit.plot(0)

# 50 points linearly spaced to 300
num_of_gates = (np.linspace(10, 300, 50)).astype(int)
gate_time = 0.1

# Select the qubits whose T2 are to be measured
qubits = [0]

# Generate experiments
circs, xdata = t2_circuits(num_of_gates, gate_time, qubits)

backend = qiskit.Aer.get_backend('qasm_simulator')

# Set the simulator with phase damping noise
t2 = 10
p = 1 - np.exp(-2*gate_time/t2)
error = phase_damping_error(p)
noise_model = NoiseModel()
noise_model.add_quantum_error(error, 'id', [0])

# Run the simulator
shots = 300
backend_result = qiskit.execute(circs, backend,
                                shots=shots, noise_model=noise_model).result()

%matplotlib inline
# Fit the data to an exponent
# The correct answers are a=1, c=0, and t2=10/5 for qubit 0/2
# The user does not know the correct answer exactly,
# so starts the fit from a different but close location

initial_t2 = t2*1.1
initial_a = 0.5
initial_c = 0.5

fit = T2Fitter(backend_result, xdata, qubits,
               fit_p0=[initial_a, initial_t2, initial_c],
               fit_bounds=([-0.5, 0, -0.5],
                           [1.5, 2*t2, 1.5]))

fit.plot(0)

num_of_gates = (np.linspace(1, 30, 30)).astype(int)
gate_time = 0.1

# Select the qubits whose T2 are to be measured
qubits = [0]

# Echo parameters
n_echos = 5
alt_phase_echo = True

# Generate experiments
circs, xdata = t2_circuits(num_of_gates, gate_time, qubits, n_echos, alt_phase_echo)

backend = qiskit.Aer.get_backend('qasm_simulator')

# Set the simulator with phase damping noise
t2 = 10
p = 1 - np.exp(-2*gate_time/t2)
error = phase_damping_error(p)
noise_model = NoiseModel()
noise_model.add_quantum_error(error, 'id', [0])

# Run the simulator
shots = 300
backend_result = qiskit.execute(circs, backend,
                                shots=shots, noise_model=noise_model).result()

%matplotlib inline
# Fit the data to an exponent
# The correct answers are a=1, c=0, and t2=10/5 for qubit 0/2
# The user does not know the correct answer exactly,
# so starts the fit from a different but close location

initial_t2 = t2*1.1
initial_a = 0.5
initial_c = 0.5

fit = T2Fitter(backend_result, xdata, qubits,
               fit_p0=[initial_a, initial_t2, initial_c],
               fit_bounds=([-0.5, 0, -0.5],
                           [1.5, 2*t2, 1.5]))

fit.plot(0)