import auspex.config as config
config.auspex_dummy_mode = True

from QGL import *
from auspex.qubit import *

"""
Explanation: Example Q5: Experiment Sweeps
This example notebook shows how to add sweeps to Auspex qubit experiments
© Raytheon BBN Technologies 2018
End of explanation
"""


cl = ChannelLibrary(":memory:")
pl = PipelineManager()

q1 = cl.new_qubit("q1")
aps2_1 = cl.new_APS2("BBNAPSa", address="192.168.2.4", trigger_interval=200e-6) 
aps2_2 = cl.new_APS2("BBNAPSb", address="192.168.2.2")
dig_1  = cl.new_Alazar("Alazar_1", address="1", sampling_rate=500e6, record_length=1024)
h1 = cl.new_source("Holz_1", "HolzworthHS9000", "HS9004A-009-1", reference='10MHz', power=-30)
h2 = cl.new_source("Holz_2", "HolzworthHS9000", "HS9004A-009-2", reference='10MHz', power=-30)

cl.set_measure(q1, aps2_1, dig_1.ch("1"), trig_channel=aps2_1.ch("m2"), gate=False, generator=h1)
cl.set_control(q1, aps2_2, generator=h2)
cl.set_master(aps2_1, aps2_1.ch("m1"))
cl["q1"].measure_chan.frequency = 0e6
cl["q1"].measure_chan.autodyne_freq = 10e6

"""
Explanation: Channel library setup
End of explanation
"""


pl.create_default_pipeline(buffers=True)
pl["q1"].add(Display(label="blub"))
pl["q1"]["Demodulate"]["Integrate"].add(Display(label="int", plot_dims=1))
pl.show_pipeline()

"""
Explanation: Pipeline setup: Take Note: we use the buffers keyword argument to automatically generate buffers instead of writers. This is sometimes convenient if you don't require data to be written to file. It becomes immediately available in the notebook after running!
End of explanation
"""


demod = pl["q1"]["Demodulate"]
demod.frequency = cl["q1"].measure_chan.frequency
demod.decimation_factor = 16

integ = pl["q1"]["Demodulate"]["Integrate"]
integ.box_car_start = 0.2e-6
integ.box_car_stop= 1.9e-6

"""
Explanation: Initialize software demodulation parameters. If these are not properly configured than the Channelizer filter will report 'insufficient decimation' or other errors. The integration boxcar parameters are then defined.
End of explanation
"""


lengths = np.linspace(20e-9, 2020e-9, 31)
exp = QubitExperiment(RabiWidth(q1,lengths),averages=50)
exp.set_fake_data(dig_1, np.exp(-lengths/1e-6)*np.cos(1e7*lengths))
# exp.add_qubit_sweep(q1,"measure", "frequency", np.linspace(6.512e9, 6.522e9, 11))
exp.run_sweeps()

"""
Explanation: Adding experiment sweeps
Once a QubitExperiment has been created, we can programmatically add sweeps as shown here.
End of explanation
"""


data, descriptor = exp.outputs_by_qubit["q1"][0].get_data()

descriptor.axes

data.shape

"""
Explanation: We fetch the data and data descriptor directly from the buffer. The data is automatically reshaped to match the experiment axes, and the descriptor enumerates all of the values of these axes for convenience plotting, etc..
End of explanation
"""


import matplotlib.pyplot as plt
%matplotlib inline

plt.imshow(np.real(data), aspect='auto', extent=descriptor.extent())
plt.xlabel("Delay (µs)")
plt.ylabel("Frequency (GHz)");

"""
Explanation: We even include a convenience extent function conforming to the infinitely forgettable matplotlib format.
End of explanation
"""


exp = QubitExperiment(RabiWidth(q1,lengths),averages=50)
exp.add_qubit_sweep(q1,"measure", "frequency", np.linspace(6.512e9, 6.522e9, 5))
exp.add_qubit_sweep(q1,"measure", "amplitude", np.linspace(0.0, 1.0, 21))

"""
Explanation: Adding Multiple Sweeps
An arbitrary number of sweeps can be added. For example:
End of explanation
"""


exp.output_connectors["q1"].descriptor.axes

"""
Explanation: If we inspect the internal representation of the "output connector" into which the instrument driver will dump its data, we can see all of the axes it contains.
End of explanation
"""


