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

# In[1]:


# Make output easier to read
from rich import pretty

pretty.install()


# In[2]:


from quantify_scheduler.schemas.examples import utils

# load here to avoid loading every time a fixture is used
transmon_device_cfg = utils.load_json_example_scheme("transmon_test_config.json")
zhinst_hardware_cfg = utils.load_json_example_scheme("zhinst_test_mapping.json")

zhinst_hardware_cfg


# In[3]:


# import statements required to make a schedule

import numpy as np

from quantify_scheduler import Schedule
from quantify_scheduler.operations.acquisition_library import SSBIntegrationComplex
from quantify_scheduler.operations.pulse_library import IdlePulse, SquarePulse
from quantify_scheduler.resources import ClockResource


# In[4]:


pulse_amps = np.linspace(0.05, 0.9, 3)
repetitions = 1024
init_duration = 4000e-6  # 4us should allow for plenty of wait time
mw_port = "q0:mw"
ro_port = "q0:res"
mw_clock = "q0.01"  # chosen to correspond to values in the hardware cfg
ro_clock = "q0.ro"
readout_frequency = (
    6.5e9  # this frequency will be used for both the AWG pulse as well as
)
# for the readout.

pulse_duration = 1e-6
acq_channel = 0
integration_time = 2e-6
acquisition_delay = 0


sched = Schedule(name="AWG staircase", repetitions=repetitions)

sched.add_resource(ClockResource(name=mw_clock, freq=readout_frequency))
sched.add_resource(ClockResource(name=ro_clock, freq=readout_frequency))
pulse_amps = np.asarray(pulse_amps)


for acq_index, pulse_amp in enumerate(pulse_amps):

    sched.add(IdlePulse(duration=init_duration))

    pulse = sched.add(
        SquarePulse(
            duration=pulse_duration,
            amp=pulse_amp,
            port=mw_port,
            clock=mw_clock,
        ),
        label=f"SquarePulse_{acq_index}",
    )

    sched.add(
        SSBIntegrationComplex(
            duration=integration_time,
            port=ro_port,
            clock=ro_clock,
            acq_index=acq_index,
            acq_channel=acq_channel,
        ),
        ref_op=pulse,
        ref_pt="start",
        rel_time=acquisition_delay,
        label=f"Acquisition_{acq_index}",
    )

sched


# In[5]:


from quantify_scheduler.compilation import qcompile

comp_sched = qcompile(
    schedule=sched, device_cfg=transmon_device_cfg, hardware_cfg=zhinst_hardware_cfg
)


# In[6]:


# Pandas dataframes do not render correctly in the sphinx documentation environment. See issue #238.
comp_sched.timing_table


# In[7]:


comp_sched.hardware_timing_table


# In[8]:


comp_sched.hardware_waveform_dict


# In[9]:


comp_sched.compiled_instructions


# In[10]:


# the .as_dict method can be used to generate a "readable" overview of the settings.
hdawg_settings_dict = (
    comp_sched.compiled_instructions["ic_hdawg0"].settings_builder.build().as_dict()
)
# hdawg_settings_dict
hdawg_settings_dict


# In[11]:


awg_index = 0
print(hdawg_settings_dict["compiler/sourcestring"][awg_index])


# In[12]:


# the .as_dict method can be used to generate a "readable" overview of the settings.
uhfqa_settings_dict = (
    comp_sched.compiled_instructions["ic_uhfqa0"].settings_builder.build().as_dict()
)
# uhfqa_settings_dict
uhfqa_settings_dict


# In[13]:


awg_index = 0
print(uhfqa_settings_dict["compiler/sourcestring"][awg_index])


# In[14]:


from quantify_scheduler.schedules.verification import acquisition_staircase_sched

acq_channel = 0
schedule = acquisition_staircase_sched(
    readout_pulse_amps=np.linspace(0, 1, 4),
    readout_pulse_duration=1e-6,
    readout_frequency=6e9,
    acquisition_delay=100e-9,
    integration_time=2e-6,
    port="q0:res",
    clock="q0.ro",
    repetitions=1024,
    acq_channel=acq_channel,
)


comp_sched = qcompile(
    schedule, device_cfg=transmon_device_cfg, hardware_cfg=zhinst_hardware_cfg
)