#!/usr/bin/env python
# coding: utf-8

# In[1]:


import numpy as np
from IPython.core.interactiveshell import InteractiveShell
from qcodes import Instrument, ManualParameter

from quantify_core.data.handling import get_tuids_containing, set_datadir
from quantify_core.measurement import MeasurementControl
from quantify_core.utilities.examples_support import default_datadir
from quantify_core.visualization.pyqt_plotmon import PlotMonitor_pyqt

rng = np.random.default_rng(seed=555555)  # random number generator

# Display any variable or statement on its own line
InteractiveShell.ast_node_interactivity = "all"


# In[2]:


# We recommend to always set the directory at the start of the python kernel
# and stick to a single common data directory for all
# notebooks/experiments within your measurement setup/PC
# This sets a default data directory for tutorial purposes. Change it to your
# desired data directory.


# In[3]:


set_datadir(default_datadir())


# In[4]:


class Device(Instrument):
    """A dummy instrument."""

    def __init__(self, name: str):
        super().__init__(name=name)

        self.add_parameter(name="amp_0", unit="A", parameter_class=ManualParameter)
        self.add_parameter(name="amp_1", unit="A", parameter_class=ManualParameter)
        self.add_parameter(name="offset", unit="A", parameter_class=ManualParameter)

        self.add_parameter(
            name="adc", label="ADC input", unit="V", get_cmd=self._get_dac_value
        )

    def _get_dac_value(self):
        s1 = np.exp(-3 * self.amp_0()) * np.sin(self.amp_0() * 2 * np.pi * 3)
        s2 = np.cos(self.amp_1() * 2 * np.pi * 2)
        return s1 + s2 + rng.uniform(0, 0.2) + self.offset()


# In[5]:


meas_ctrl = MeasurementControl("meas_ctrl")
plotmon = PlotMonitor_pyqt("Plot Monitor")
meas_ctrl.instr_plotmon(plotmon.name)
device = Device("Device")


# In[6]:


plotmon.tuids()
plotmon.tuids_extra()


# In[7]:


# Example of loading datasets onto the plot
# plotmon.tuids(["20201124-184709-137-8a5112", "20201124-184716-237-918bee"])
# plotmon.tuids_extra(["20201124-184722-988-0463d4", "20201124-184729-618-85970f"])


# In[8]:


plotmon.tuids_max_num()


# In[9]:


# set initial values to emulate the instrument state
device.amp_0(0.0)
device.amp_1(0.0)
device.offset(0.0)

n_pnts = 50

meas_ctrl.settables(device.amp_0)
meas_ctrl.setpoints(np.linspace(0, 1, n_pnts))
meas_ctrl.gettables(device.adc)
dset = meas_ctrl.run("ADC scan")

plotmon.main_QtPlot


# In[10]:


n_pnts = 20

meas_ctrl.settables(device.amp_0)
meas_ctrl.setpoints(np.linspace(0, 1, n_pnts))
meas_ctrl.gettables(device.adc)
dset = meas_ctrl.run("ADC scan")

plotmon.main_QtPlot


# In[11]:


n_pnts = 30

meas_ctrl.settables(device.amp_0)
meas_ctrl.setpoints(np.linspace(0, 1, n_pnts))
meas_ctrl.gettables(device.adc)
dset = meas_ctrl.run("ADC scan")

plotmon.main_QtPlot


# In[12]:


# Now the oldest dataset will vanish from the plot

n_pnts = 40

meas_ctrl.settables(device.amp_0)
meas_ctrl.setpoints(np.linspace(0, 1, n_pnts))
meas_ctrl.gettables(device.adc)
dset = meas_ctrl.run("ADC scan")

plotmon.main_QtPlot


# In[13]:


# We can accumulate more datasets on the plot if we want to
plotmon.tuids_max_num(4)
n_pnts = 40

meas_ctrl.settables(device.amp_0)
meas_ctrl.setpoints(np.linspace(0, 1, n_pnts))
meas_ctrl.gettables(device.adc)
dset = meas_ctrl.run("ADC scan")

plotmon.main_QtPlot


# In[14]:


# Or we can disable the accumulation and plot a single dataset
plotmon.tuids_max_num(1)

plotmon.main_QtPlot


# In[15]:


# This can also be reset with
plotmon.tuids([])

plotmon.main_QtPlot  # The plotting window will vanish, it is supposed to


# In[16]:


# For now we will allow two datasets on the plot monitor
plotmon.tuids_max_num(2)


# In[17]:


# Now let's imagine that something strange is happening with our setup
device.offset(1.5)

n_pnts = 40
meas_ctrl.settables(device.amp_0)
meas_ctrl.setpoints(np.linspace(0, 1, n_pnts))
meas_ctrl.gettables(device.adc)
dset = meas_ctrl.run("ADC scan problem")

plotmon.main_QtPlot


# In[18]:


# We would like to compare if the current behavior matches for example
# what we got a few minutes ago

reference_tuids = sorted(get_tuids_containing("ADC"))[-3:-1]

plotmon.tuids_extra(reference_tuids)
plotmon.main_QtPlot


# In[19]:


device.offset(0.0)  # OK... that cable was not connected in the right place...

# Now let's run again our experiments while we compare it to the previous one in realtime

n_pnts = 30
meas_ctrl.settables(device.amp_0)
meas_ctrl.setpoints(np.linspace(0, 1, n_pnts))
meas_ctrl.gettables(device.adc)
dset = meas_ctrl.run("ADC scan")

n_pnts = 40
meas_ctrl.settables(device.amp_0)
meas_ctrl.setpoints(np.linspace(0, 1, n_pnts))
meas_ctrl.gettables(device.adc)
dset = meas_ctrl.run("ADC scan")

print("Yey! We have recovered our setup!")
plotmon.main_QtPlot


# In[20]:


# We do not need the reference datasets anymore
plotmon.tuids_extra([])
plotmon.main_QtPlot


# In[21]:


# Note: both plotmon.tuids_extra and plotmon.tuids can be used
# but keep in mind that meas_ctrl also uses the plotmon.tuids

tuids = get_tuids_containing("problem")[0:1]
tuids
plotmon.tuids(tuids)

n_pnts = 40
meas_ctrl.settables(device.amp_0)
meas_ctrl.setpoints(np.linspace(0, 1, n_pnts))
meas_ctrl.gettables(device.adc)
dset = meas_ctrl.run("ADC scan")

plotmon.main_QtPlot


# In[22]:


# When we have 2D plots only the first dataset from plotmon.tuids or
# plotmon.tuids_extra, in that order of priority, will be plotted in the
# secondary window

meas_ctrl.settables([device.amp_0, device.amp_1])
meas_ctrl.setpoints_grid([np.linspace(0, 1, 20), np.linspace(0, 0.5, 15)])
meas_ctrl.gettables(device.adc)
dset = meas_ctrl.run("ADC scan 2D")
reference_tuid_2D = dset.attrs["tuid"]

plotmon.main_QtPlot
plotmon.secondary_QtPlot



# In[23]:


meas_ctrl.settables([device.amp_0, device.amp_1])
meas_ctrl.setpoints_grid([np.linspace(0, 1, 20), np.linspace(0, 0.5, 15)])
meas_ctrl.gettables(device.adc)
dset = meas_ctrl.run("ADC scan 2D")

plotmon.main_QtPlot
plotmon.secondary_QtPlot


# In[24]:


device.offset(2.03)
plotmon.tuids_extra([reference_tuid_2D])

meas_ctrl.settables([device.amp_0, device.amp_1])
meas_ctrl.setpoints_grid([np.linspace(0, 1, 20), np.linspace(0, 0.5, 15)])
meas_ctrl.gettables(device.adc)
dset = meas_ctrl.run("ADC scan 2D")

plotmon.main_QtPlot
plotmon.secondary_QtPlot


# In[25]:


plotmon.tuids([])
plotmon.tuids_extra([reference_tuid_2D])
plotmon.main_QtPlot
plotmon.secondary_QtPlot


# In[26]:


plotmon.tuids()