# import SymPy and define symbols
import sympy as sp
sp.init_printing(use_unicode=True)
wr = sp.Symbol('\omega_r') # resonator frequency
wq = sp.Symbol('\omega_q') # qubit frequency
g = sp.Symbol('g', real=True) # vacuum Rabi coupling
Delta = sp.Symbol('Delta', real=True) # wr - wq; defined later

# import operator relations and define them
from sympy.physics.quantum.boson import BosonOp
a = BosonOp('a') # resonator photon annihilation operator
from sympy.physics.quantum import pauli, Dagger, Commutator
from sympy.physics.quantum.operatorordering import normal_ordered_form

# Pauli matrices
sx = pauli.SigmaX()
sy = pauli.SigmaY()
sz = pauli.SigmaZ()

# qubit raising and lowering operators
splus = pauli.SigmaPlus() 
sminus = pauli.SigmaMinus()

# define J-C Hamiltonian in terms of diagonal and non-block diagonal terms
H0 = wr*Dagger(a)*a - (1/2)*wq*sz; 
H1 = 0
H2 = g*(Dagger(a)*sminus + a*splus); 
HJC = H0 + H1 + H2; HJC # print

# using the above method for finding the ansatz
eta = Commutator(H0, H2); eta

pauli.qsimplify_pauli(normal_ordered_form(eta.doit().expand()))

A = sp.Symbol('A')
B = sp.Symbol('B')
eta = A * Dagger(a) * sminus - B * a * splus;
pauli.qsimplify_pauli(normal_ordered_form(Commutator(H0, eta).doit().expand()))

H2

S1 = eta.subs(A, g/Delta)
S1 = S1.subs(B, g/Delta); S1.factor()

Heff = H0 + H1 + 0.5*pauli.qsimplify_pauli(normal_ordered_form(Commutator(H2, S1).doit().expand())).simplify(); Heff

from qiskit.tools.jupyter import *
from qiskit import IBMQ
IBMQ.load_account()
provider = IBMQ.get_provider(hub='ibm-q', group='open', project='main')
backend = provider.get_backend('ibmq_armonk') 

backend_config = backend.configuration()
assert backend_config.open_pulse, "Backend doesn't support Pulse"

dt = backend_config.dt
print(f"Sampling time: {dt*1e9} ns") 
backend_defaults = backend.defaults()

import numpy as np

# unit conversion factors -> all backend properties returned in SI (Hz, sec, etc)
GHz = 1.0e9 # Gigahertz
MHz = 1.0e6 # Megahertz
kHz = 1.0e3
us = 1.0e-6 # Microseconds
ns = 1.0e-9 # Nanoseconds

# We will find the qubit frequency for the following qubit.
qubit = 0

# The sweep will be centered around the estimated qubit frequency.
center_frequency_Hz = backend_defaults.qubit_freq_est[qubit]        # The default frequency is given in Hz
                                                                    # warning: this will change in a future release
print(f"Qubit {qubit} has an estimated frequency of {center_frequency_Hz / GHz} GHz.")

# scale factor to remove factors of 10 from the data
scale_factor = 1e-14
# We will sweep 40 MHz around the estimated frequency
frequency_span_Hz = 40 * MHz
# in steps of 1 MHz.
frequency_step_Hz = 1 * MHz
# We will sweep 20 MHz above and 20 MHz below the estimated frequency
frequency_min = center_frequency_Hz - frequency_span_Hz / 2
frequency_max = center_frequency_Hz + frequency_span_Hz / 2
# Construct an np array of the frequencies for our experiment
frequencies_GHz = np.arange(frequency_min / GHz, 
                            frequency_max / GHz, 
                            frequency_step_Hz / GHz)

print(f"The sweep will go from {frequency_min / GHz} GHz to {frequency_max / GHz} GHz \
in steps of {frequency_step_Hz / MHz} MHz.")

from qiskit import pulse            # This is where we access all of our Pulse features!

inst_sched_map = backend_defaults.instruction_schedule_map
measure = inst_sched_map.get('measure', qubits=[qubit])
x_pulse = inst_sched_map.get('x', qubits=[qubit])

### Collect the necessary channels
drive_chan = pulse.DriveChannel(qubit)
meas_chan = pulse.MeasureChannel(qubit)
acq_chan = pulse.AcquireChannel(qubit)

# Create the base schedule
# Start with drive pulse acting on the drive channel
schedule = pulse.Schedule(name='Frequency sweep')
schedule += x_pulse
# The left shift `<<` is special syntax meaning to shift the start time of the schedule by some duration
schedule += measure << schedule.duration

# Create the frequency settings for the sweep (MUST BE IN HZ)
frequencies_Hz = frequencies_GHz*GHz
schedule_frequencies = [{drive_chan: freq} for freq in frequencies_Hz]

schedule.draw(label=True, scaling=0.8)

from qiskit import assemble
frequency_sweep_program = assemble(schedule,
                                   backend=backend, 
                                   meas_level=1,
                                   meas_return='avg',
                                   shots=1024,
                                   schedule_los=schedule_frequencies)

# RUN the job on a real device
job = backend.run(frequency_sweep_program)
print(job.job_id())
from qiskit.tools.monitor import job_monitor
job_monitor(job)

# OR retreive result from previous run
# job = backend.retrieve_job('5ef3b081fbc24b001275b03b')

frequency_sweep_results = job.result() 

import matplotlib.pyplot as plt
plt.style.use('dark_background')

sweep_values = []
for i in range(len(frequency_sweep_results.results)):
    # Get the results from the ith experiment
    res = frequency_sweep_results.get_memory(i)*scale_factor
    # Get the results for `qubit` from this experiment
    sweep_values.append(res[qubit])

plt.scatter(frequencies_GHz, np.real(sweep_values), color='white') # plot real part of sweep values
plt.xlim([min(frequencies_GHz), max(frequencies_GHz)])
plt.xlabel("Frequency [GHz]")
plt.ylabel("Measured signal [a.u.]")
plt.show()

from scipy.optimize import curve_fit

def fit_function(x_values, y_values, function, init_params):
    fitparams, conv = curve_fit(function, x_values, y_values, init_params)
    y_fit = function(x_values, *fitparams)
    
    return fitparams, y_fit

fit_params, y_fit = fit_function(frequencies_GHz,
                                 np.real(sweep_values), 
                                 lambda x, A, q_freq, B, C: (A / np.pi) * (B / ((x - q_freq)**2 + B**2)) + C,
                                 [5, 4.975, 1, 3] # initial parameters for curve_fit
                                )

plt.scatter(frequencies_GHz, np.real(sweep_values), color='white')
plt.plot(frequencies_GHz, y_fit, color='red')
plt.xlim([min(frequencies_GHz), max(frequencies_GHz)])

plt.xlabel("Frequency [GHz]")
plt.ylabel("Measured Signal [a.u.]")
plt.show()

 # Create the schedules for 0 and 1
schedule_0 = pulse.Schedule(name='0')
schedule_0 += measure

schedule_1 = pulse.Schedule(name='1')
schedule_1 += x_pulse
schedule_1 += measure << schedule_1.duration

schedule_0.draw()

schedule_1.draw()

frequency_span_Hz = 320 * kHz
frequency_step_Hz = 8 * kHz

center_frequency_Hz = backend_defaults.meas_freq_est[qubit]
print(f"Qubit {qubit} has an estimated readout frequency of {center_frequency_Hz / GHz} GHz.")
frequency_min = center_frequency_Hz - frequency_span_Hz / 2
frequency_max = center_frequency_Hz + frequency_span_Hz / 2
frequencies_GHz = np.arange(frequency_min / GHz, 
                            frequency_max / GHz, 
                            frequency_step_Hz / GHz)

print(f"The sweep will go from {frequency_min / GHz} GHz to {frequency_max / GHz} GHz\
      in steps of {frequency_step_Hz / MHz} MHz.")

num_shots_per_frequency = 2048

frequencies_Hz = frequencies_GHz*GHz
schedule_los = [{meas_chan: freq} for freq in frequencies_Hz]

cavity_sweep_0 = assemble(schedule_0,
                               backend=backend, 
                               meas_level=1,
                               meas_return='avg',
                               shots=num_shots_per_frequency,
                               schedule_los=schedule_los)

cavity_sweep_1 = assemble(schedule_1,
                               backend=backend, 
                               meas_level=1,
                               meas_return='avg',
                               shots=num_shots_per_frequency,
                               schedule_los=schedule_los)

# RUN the job on a real device
#job_0 = backend.run(cavity_sweep_0)
#job_monitor(job_0)
#job_0.error_message()

#job_1 = backend.run(cavity_sweep_1)
#job_monitor(job_1)
#job_1.error_message()

# OR retreive result from previous run
job_0 = backend.retrieve_job('5efa5b447c0d6800137fff1c')
job_1 = backend.retrieve_job('5efa6b2720eee10013be46b4')

cavity_sweep_0_results = job_0.result()
cavity_sweep_1_results = job_1.result()

scale_factor = 1e-14

sweep_values_0 = []
for i in range(len(cavity_sweep_0_results.results)):
    res_0 = cavity_sweep_0_results.get_memory(i)*scale_factor
    sweep_values_0.append(res_0[qubit])

sweep_values_1 = []
for i in range(len(cavity_sweep_1_results.results)):
    res_1 = cavity_sweep_1_results.get_memory(i)*scale_factor
    sweep_values_1.append(res_1[qubit])
    
plotx = frequencies_Hz/kHz
ploty_0 = np.abs(sweep_values_0)
ploty_1 = np.abs(sweep_values_1)

plt.plot(plotx, ploty_0, color='blue', marker='.') # plot real part of sweep values
plt.plot(plotx, ploty_1, color='red', marker='.') # plot real part of sweep values
plt.legend([r'$\vert0\rangle$', r'$\vert1\rangle$'])

plt.grid()
plt.xlabel("Frequency [kHz]")
plt.ylabel("Measured signal [a.u.]")
plt.yscale('log')
plt.show()