import numpy as np

import qiskit.pulse as pulse
from qiskit.circuit import Parameter

from qiskit_experiments.calibration_management.backend_calibrations import BackendCalibrations

from qiskit import IBMQ, schedule

API_KEY = "5bd4ecfdc74e6680da7c79998259781431661e5326ae2f88eea95dee8f74b87530ba63fbca8105404de4ffd36e4b484631907acff73c805580928218a5ccf0b3"

import qiskit as q
from qiskit import IBMQ,schedule
import numpy as np

import qiskit.pulse as pulse
from qiskit.circuit import Parameter

from qiskit_experiments.calibration_management.backend_calibrations import BackendCalibrations


%matplotlib inline
import sys
sys.path.insert(0, '..') # Add qiskit_runtime directory to the path

IBMQ.save_account(API_KEY)
# Details in: https://qiskit.org/documentation/install.html 

IBMQ.save_account(API_KEY, overwrite=True)
IBMQ.load_account()

provider = IBMQ.get_provider("ibm-q")

for backend in provider.backends():
    try:
        qubit_count = len(backend.properties().quibits)
    except:
        qubit_count = "simulated"
        
    print(f"{backend.name()} has {backend.status().pending_jobs} queued and { qubit_count} qubits")

qubit = 0  # The qubit we will work with

def setup_cals(backend) -> BackendCalibrations:
    """A function to instantiate calibrations and add a couple of template schedules."""
    cals = BackendCalibrations(backend)

    dur = Parameter("dur")
    amp = Parameter("amp")
    sigma = Parameter("σ")
    beta = Parameter("β")
    drive = pulse.DriveChannel(Parameter("ch0"))

    # Define and add template schedules.
    with pulse.build(name="xp") as xp:
        pulse.play(pulse.Drag(dur, amp, sigma, beta), drive)

    with pulse.build(name="xm") as xm:
        pulse.play(pulse.Drag(dur, -amp, sigma, beta), drive)
        
    with pulse.build(name="x90p") as x90p:
        pulse.play(pulse.Drag(dur, Parameter("amp"), sigma, Parameter("β")), drive)

    cals.add_schedule(xp)
    cals.add_schedule(xm)
    cals.add_schedule(x90p)
    
    return cals

def add_parameter_guesses(cals: BackendCalibrations):
    """Add guesses for the parameter values to the calibrations."""
    for sched in ["xp", "x90p"]:
        cals.add_parameter_value(80, "σ", schedule=sched)
        cals.add_parameter_value(0.5, "β", schedule=sched)
        cals.add_parameter_value(320, "dur", schedule=sched)
        cals.add_parameter_value(0.5, "amp", schedule=sched)

cals = setup_cals(backend)
add_parameter_guesses(cals)

from qiskit_experiments.calibration_management.basis_gate_library import FixedFrequencyTransmon

library = FixedFrequencyTransmon(default_values={"duration": 320})
cals = BackendCalibrations(backend, library)

from qiskit_experiments.library.characterization.qubit_spectroscopy import QubitSpectroscopy

import pandas as pd

pd.DataFrame(cals.parameters_table(qubit_list=[qubit, ()]))

freq01_estimate = backend.defaults().qubit_freq_est[qubit]
frequencies = np.linspace(freq01_estimate -15e6, freq01_estimate + 15e6, 51)
spec = QubitSpectroscopy(qubit, frequencies)
spec.set_experiment_options(amp=0.1)

circuit = spec.circuits(backend)[0]
circuit.draw(output="mpl")

schedule(circuit, backend).draw()

spec_data = spec.run(backend)

spec_data.block_for_results()

spec_data.figure(0)

print(spec_data.analysis_results("f01"))

from qiskit_experiments.calibration_management.update_library import Frequency

Frequency.update(cals, spec_data)

pd.DataFrame(cals.parameters_table(qubit_list=[qubit]))

from qiskit_experiments.library.calibration import Rabi
from qiskit_experiments.calibration_management.update_library import Amplitude

rabi = Rabi(qubit)
rabi.set_experiment_options(
    amplitudes=np.linspace(-0.95, 0.95, 51), 
    schedule=cals.get_schedule("x", (qubit,), assign_params={"amp": Parameter("amp")}),
)

rabi_data = rabi.run(backend)
rabi_data.block_for_results()

rabi_data.figure(0)

print(rabi_data.analysis_results("rabi_rate"))

Amplitude.update(cals, rabi_data, angles_schedules=[(np.pi, "amp", "x"), (np.pi/2, "amp", "sx")])

pd.DataFrame(cals.parameters_table(qubit_list=[qubit, ()], parameters="amp"))

cals.get_schedule("sx", qubit)

cals.get_schedule("x", qubit)

cals.get_schedule("y", qubit)

cals.save(file_type="csv", overwrite=True, file_prefix="Armonk")

cals = BackendCalibrations(backend, library)
cals.load_parameter_values(file_name="Armonkparameter_values.csv")

pd.DataFrame(cals.parameters_table(qubit_list=[qubit, ()], parameters="amp"))

from qiskit_experiments.library.calibration.drag import DragCal
from qiskit_experiments.calibration_management.update_library import Drag

cal_drag = DragCal(qubit)

cal_drag.set_experiment_options(
    rp=cals.get_schedule("x", qubit, assign_params={"β": Parameter("β")}),
    betas=np.linspace(-20, 20, 25),
    reps=[3, 5, 7]
)

cal_drag.circuits(backend)[1].draw(output='mpl')

drag_data = cal_drag.run(backend)
drag_data.block_for_results()

drag_data.figure(0)

print(drag_data.analysis_results("beta"))

Drag.update(cals, drag_data, parameter="β", schedule="x")

pd.DataFrame(cals.parameters_table(qubit_list=[qubit, ()], parameters="β"))

from qiskit_experiments.library.calibration.fine_amplitude import FineXAmplitude
from qiskit_experiments.calibration_management.update_library import Amplitude

amp_x_cal = FineXAmplitude(qubit)
amp_x_cal.set_experiment_options(schedule=cals.get_schedule("x", qubit))

amp_x_cal.circuits(backend)[5].draw(output="mpl")

data_fine = amp_x_cal.run(backend)
data_fine.block_for_results()

data_fine.figure(0)

print(data_fine.analysis_results("d_theta"))

dtheta = data_fine.analysis_results("d_theta").value.value
target_angle = np.pi
scale = target_angle / (target_angle + dtheta)
pulse_amp = cals.get_parameter_value("amp", qubit, "x")
print(f"The ideal angle is {target_angle:.2f} rad. We measured a deviation of {dtheta:.3f} rad.")
print(f"Thus, scale the {pulse_amp:.4f} pulse amplitude by {scale:.3f} to obtain {pulse_amp*scale:.5f}.")

Amplitude.update(cals, data_fine, angles_schedules=[(target_angle, "amp", "x")])

pd.DataFrame(cals.parameters_table(qubit_list=[qubit, ()], parameters="amp"))

amp_x_cal.set_experiment_options(schedule=cals.get_schedule("x", qubit))

data_fine2 = amp_x_cal.run(backend)
data_fine2.block_for_results()

data_fine2.figure(0)

print(data_fine2.analysis_results("d_theta"))

from qiskit_experiments.library.calibration.fine_amplitude import FineSXAmplitude

amp_sx_cal = FineSXAmplitude(qubit)
amp_sx_cal.set_experiment_options(schedule=cals.get_schedule("sx", qubit))

amp_sx_cal.circuits(backend)[5].draw(output="mpl")

data_fine_sx = amp_sx_cal.run(backend)
data_fine_sx.block_for_results()

data_fine_sx.figure(0)

print(data_fine_sx.analysis_results("d_theta"))

Amplitude.update(cals, data_fine_sx, angles_schedules=[(np.pi/2, "amp", "sx")])

pd.DataFrame(cals.parameters_table(qubit_list=[qubit, ()], parameters="amp"))

cals.get_schedule("sx", qubit)

cals.get_schedule("x", qubit)

cals.get_schedule("y", qubit)

import qiskit.tools.jupyter
%qiskit_copyright

