import numpy as np

# Importing standard Qiskit libraries
from qiskit import QuantumCircuit, transpile, IBMQ, schedule

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

from qiskit_experiments.calibration_management import Calibrations

import pandas as pd

# Loading your IBM Quantum account(s)


qubit = 1

phase = np.pi

# Generating an instance of Calibrations class
def setup_cals( backend ) -> Calibrations:
    """ A function to instantiate calibrations and add a couple of template schedules. """
    # Instantiating Calibrations as cals
    cals = Calibrations.from_backend( backend )
    
    # Parameters to sweep
    dur = Parameter( "dur" )
    amp = Parameter( "amp" )
    amp_2 = Parameter( "amp" )
    sigma = Parameter( "σ" )
    beta = Parameter( "β" )
    beta_2 = Parameter( "β" )
    
    
    drive = pulse.DriveChannel( Parameter( "ch0" ) )
    
    # Define and add template schedules
    with pulse.build( name="xp" ) as xp:  # X positive
        pulse.play( pulse.Drag( dur, amp, sigma, beta ),
                    drive )
        
    with pulse.build( name="x" ) as x:  # X positive
        pulse.play( pulse.Drag( dur, amp, sigma, beta ),
                    drive )
    
    with pulse.build( name="y" ) as y:  # Y positive
        pulse.shift_phase( phase, drive )
        pulse.play( pulse.Drag( dur, amp, sigma, beta ),
                    drive )

    with pulse.build( name="xm" ) as xm:  # X minus
        pulse.play( pulse.Drag( dur, -amp, sigma, beta ),
                    drive )
        
    with pulse.build( name="x90p" ) as x90p:  # X/2 positive
        # Different pulse amplitude and correction pulse amplitude for X/2
        pulse.play( pulse.Drag( dur, amp_2, sigma, beta_2 ),
                    drive )

    with pulse.build( name="sx" ) as sx:  # X/2 positive
        # Different pulse amplitude and correction pulse amplitude for X/2
        pulse.play( pulse.Drag( dur, amp_2, sigma, beta_2 ),
                    drive )
    
    cals.add_schedule( xp, num_qubits=1 )
    cals.add_schedule( xm, num_qubits=1 )
    cals.add_schedule( x90p, num_qubits=1 )
    cals.add_schedule( x, num_qubits=1 )
    cals.add_schedule( y, num_qubits=1 )
    cals.add_schedule( sx, num_qubits=1 )
    
    return cals


def add_parameter_guesses( cals: Calibrations ):
    """ Add guesses for the parameter values to the calibrations. """
    for sched in [ "x", "xp", "y", "sx", "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 )

# Calling calibration functions
cals = setup_cals( backend )
add_parameter_guesses( cals )

# Using a default frequency

import pandas as pd

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

from qiskit_experiments.library.calibration import RoughXSXAmplitudeCal

# Define Rabi amplitude sweep sequence
rabi = RoughXSXAmplitudeCal( qubit, cals, amplitudes=np.linspace( -0.2, 0.2, 51 ), backend=backend )

# Run the experiment
rabi_data = rabi.run().block_for_results()

rabi_data.figure( 0 )

# Save pulse amplitude
pi_amp = np.pi / ( 2*np.pi*rabi_data.analysis_results( "rabi_rate" ).value.value )
print( pi_amp )

defaults = backend.defaults()
x_schedule = defaults.instruction_schedule_map.get( 'x', 1 )
print( x_schedule )

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

cals.get_schedule("sx", qubit)

cals.get_schedule("x", qubit)

from qiskit_experiments.library import RoughDragCal

# Define DRAG coefficient sweep sequence
cal_drag = RoughDragCal( qubit, cals, backend=backend, betas=np.linspace( -20, 20, 51 ) )
cal_drag.set_experiment_options( reps=[ 3, 5, 7 ] )

# Run the DRAG measurement
drag_data = cal_drag.run().block_for_results()

drag_data.figure( 0 )

print( drag_data.analysis_results( "beta" ) )

# Save DRAG coefficient
beta = drag_data.analysis_results( "beta" ).value.value

print( beta )

# Define DRAG coefficient sweep sequence
cal_drag = RoughDragCal( qubit, cals, backend=backend, schedule_name="sx", betas=np.linspace( -20, 20, 51 ) )
cal_drag.set_experiment_options( reps=[ 9, 11, 17 ] )

# Run the DRAG measurement
drag_data = cal_drag.run().block_for_results()

drag_data.figure( 0 )

# Save DRAG coefficient
beta_2 = drag_data.analysis_results( "beta" ).value.value

print( beta_2 )

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

#print( "Qubit frequency: ", qubit_freq / 1e9, " GHz" )
print( "π-pulse amplitude: ", pi_amp )
print( "DRAG coefficient β fox X: ", beta )
print( "DRAG coefficient β for SX: ", beta_2 )

from qiskit_experiments.library.calibration.fine_amplitude import FineXAmplitudeCal

amp_x_cal = FineXAmplitudeCal(qubit, cals, backend=backend, schedule_name="x")

amp_x_fine = amp_x_cal.run().block_for_results()

amp_x_fine.figure(0)

print( amp_x_fine.analysis_results( "d_theta" ) )

dtheta = amp_x_fine.analysis_results( "d_theta" ).value.value
target_angle = np.pi
scale = target_angle / ( target_angle + dtheta )

print( "Deviation of", dtheta, "is detected for", target_angle, "rotation." )
print( "Hence, switch the π-amplitude from", pi_amp, "to", pi_amp*scale )

fine_x_amp = pi_amp * scale

# Check if we have better π-gate.
# Note that we don't have to update cals as it is updated automatically.
sanity_check = amp_x_cal.run().block_for_results()

print( "d_theta has been decreased from", dtheta, "to", sanity_check.analysis_results( "d_theta" ).value.value )

from qiskit_experiments.library.calibration.fine_amplitude import FineSXAmplitudeCal

amp_sx_cal = FineSXAmplitudeCal(qubit, cals, backend=backend, schedule_name="sx")

amp_sx_fine = amp_sx_cal.run().block_for_results()

print( amp_sx_fine.analysis_results( "d_theta" ) )

amp_sx_fine.figure(0)

dtheta = amp_sx_fine.analysis_results( "d_theta" ).value.value
target_angle = np.pi/2
scale = target_angle / ( target_angle + dtheta )

print( "Deviation of", dtheta, "is detected for", target_angle, "rotation." )
print( "Hence, switch the π/2-amplitude from", pi_amp/2, "to", pi_amp/2*scale )

fine_sx_amp = pi_amp/2 * scale

from qiskit_experiments.library.calibration import FineXDragCal, FineSXDragCal

drag_x_cal = FineXDragCal( qubit, cals, backend=backend )

drag_x_fine = drag_x_cal.run().block_for_results()

drag_x_fine.figure(0)

fine_beta = cals.get_parameter_value( "β", qubit, "x" )

print( fine_beta )

drag_sx_cal = FineSXDragCal( qubit, cals, backend=backend )

drag_sx_fine = drag_sx_cal.run().block_for_results()

fine_beta_2 = cals.get_parameter_value( "β", qubit, "sx" )

print( fine_beta_2 )

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

from qiskit_experiments.library.calibration.half_angle_cal import HalfAngleCal

half_angle_cal = HalfAngleCal( qubit, cals, backend=backend )

half_angle_fine = half_angle_cal.run().block_for_results()

half_angle_fine.figure(0)

finer_sx_amp = cals.get_parameter_value( "amp", qubit, "sx" )
print( finer_sx_amp )

# print( "Qubit frequency:", qubit_freq/1e9, "GHz to Fine Qubit frequency:", fine_qubit_freq/1e9, "GHz." )
print( "π-pulse amplitude:", pi_amp, "to Fine π-pulse amplitude:", fine_x_amp)
print( "π/2-pulse amplitude:", pi_amp/2, "to Fine π/2-pulse amplitude:", finer_sx_amp)
print( "DRAG coefficient β fox X:", beta, "to Fine DRAG coefficient β for X:", fine_beta )
print( "DRAG coefficient β for SX: ", beta_2, "to Fine DRAG coefficient β for SX:", fine_beta_2  )

# Importing standard Qiskit libraries
from qiskit import QuantumCircuit, transpile, pulse, IBMQ

# Importing RB-related libraries
from qiskit_experiments.library import StandardRB, InterleavedRB
from qiskit_experiments.framework import ParallelExperiment
from qiskit_experiments.library.randomized_benchmarking import RBUtils
import qiskit.circuit.library as circuits

x_amp = fine_x_amp
sx_amp = fine_sx_amp
x_drag = fine_beta
sx_drag = fine_beta_2

with pulse.build(backend) as X_pulse:
    drive_duration=320
    drive_sigma=80
    drive_chan=pulse.drive_channel(qubit)
    pulse.play(pulse.library.Drag(duration=drive_duration,
                              amp=x_amp,
                              sigma=drive_sigma,
                              beta=x_drag,
                              name='X pulse'), drive_chan)

with pulse.build(backend) as SX_pulse:
    drive_duration=320
    drive_sigma=80
    drive_chan=pulse.drive_channel(qubit)
    pulse.play(pulse.library.Drag(duration=drive_duration,
                              amp=sx_amp,
                              sigma=drive_sigma,
                              beta=sx_drag,
                              name='SX pulse'), drive_chan)

inst_map = backend.defaults().instruction_schedule_map
print( inst_map.get('x',qubit) )
print( inst_map.get('sx',qubit) )
inst_map.add( 'x', qubit, X_pulse )
inst_map.add( 'sx', qubit, SX_pulse )
print( inst_map.get('x',qubit) )
print( inst_map.get('sx',qubit) )

lengths = np.arange( 1, 2000, 200 )
num_samples = 16
seed = 1010
qubits = [qubit]

# Run an RB experiment on qubit 1
expDyn = StandardRB(qubits, lengths, num_samples=num_samples, seed=seed)
expdataDyn = expDyn.run(backend=backend)

# View result data
display(expdataDyn.figure(0))