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

# In[1]:


import numpy as np
import plotly.graph_objects as go

x = np.linspace(0, 360, 361)
y = np.cos(np.deg2rad(x - 180))
yc = np.minimum(x / 90 - 1, -x / 90 + 3)

fig = go.Figure()
fig.add_trace(go.Scatter(x=x, y=y, name="Quantum"))
fig.add_trace(go.Scatter(x=x, y=yc, name="Classical"))

fig.update_layout(
    title="Bell experiment",
    xaxis_title="Angle between detectors (deg)",
    yaxis_title="Correlation",
)
fig.show()


# In[2]:


# Make output easier to read
from pathlib import Path

from quantify_core.data.handling import set_datadir
from rich import pretty

from quantify_scheduler import Schedule

pretty.install()

set_datadir(Path.home() / "quantify-data")
sched = Schedule("Bell experiment")
sched


# In[3]:


sched.data


# In[4]:


q0, q1 = (
    "q0",
    "q1",
)  # we use strings because qubit resources have not been implemented yet.


# In[5]:


from quantify_scheduler.operations.gate_library import CZ, X90, Measure, Reset, Rxy

# we use a regular for loop as we have to unroll the changing theta variable here
for acq_idx, theta in enumerate(np.linspace(0, 360, 21)):
    sched.add(Reset(q0, q1))
    sched.add(X90(q0))
    sched.add(X90(q1), ref_pt="start")  # this ensures pulses are aligned
    sched.add(CZ(q0, q1))
    sched.add(Rxy(theta=theta, phi=0, qubit=q0))

    sched.add(Measure(q0, acq_index=acq_idx), label="M q0 {:.2f} deg".format(theta))
    sched.add(
        Measure(q1, acq_index=acq_idx),
        label="M q1 {:.2f} deg".format(theta),
        ref_pt="start",
    )


# In[6]:


get_ipython().run_line_magic('matplotlib', 'inline')

f, ax = sched.plot_circuit_diagram()
# all gates are plotted, but it doesn't all fit in a matplotlib figure
ax.set_xlim(-0.5, 9.5)


# In[7]:


sched


# In[8]:


sched.data.keys()


# In[9]:


from itertools import islice

# showing the first 5 elements of the operation dict
dict(islice(sched.data["operation_dict"].items(), 5))


# In[10]:


list(sched.data["schedulables"].values())[:6]


# In[11]:


theta = 314
rxy_theta = Rxy(theta=theta, phi=0, qubit=q0)
rxy_theta.data


# In[12]:


import inspect
import json

from quantify_scheduler.schemas.examples.circuit_to_device_example_cfgs import (
    example_transmon_cfg,
)
from quantify_scheduler.backends.circuit_to_device import DeviceCompilationConfig

transmon_test_config = DeviceCompilationConfig.parse_obj(example_transmon_cfg)


# In[13]:


from quantify_scheduler.compilation import device_compile

sched = device_compile(sched, device_cfg=transmon_test_config)
# add_pulse_information_transmon(sched, device_cfg=transmon_test_config)
# determine_absolute_timing(schedule=sched)


# In[14]:


sched


# In[15]:


from quantify_scheduler.visualization.pulse_scheme import pulse_diagram_plotly

pulse_diagram_plotly(
    sched,
    port_list=["q0:mw", "q0:res", "q0:fl", "q1:mw"],
    modulation_if=10e6,
    sampling_rate=1e9,
)


# In[16]:


sched = Schedule("Bell experiment")
for acq_idx, theta in enumerate(np.linspace(0, 360, 21)):
    sched.add(Reset(q0, q1))
    sched.add(X90(q0))
    sched.add(X90(q1), ref_pt="start")  # this ensures pulses are aligned
    # sched.add(CZ(q0, q1)) # FIXME Commented out because of not implemented error
    sched.add(Rxy(theta=theta, phi=0, qubit=q0))

    sched.add(Measure(q0, acq_index=acq_idx), label=f"M q0 {acq_idx} {theta:.2f} deg")
    sched.add(
        Measure(q1, acq_index=acq_idx),
        label=f"M q1 {acq_idx} {theta:.2f} deg",
        ref_pt="start",
    )


sched = device_compile(sched, device_cfg=transmon_test_config)

# add_pulse_information_transmon(sched, device_cfg=transmon_test_config)
# determine_absolute_timing(schedule=sched)


# In[17]:


import quantify_scheduler.schemas.examples as es

esp = inspect.getfile(es)

cfg_f = Path(esp).parent / "qblox_test_mapping.json"

with open(cfg_f, "r") as f:
    qblox_test_mapping = json.load(f)

qblox_test_mapping


# In[18]:


from qblox_instruments import Pulsar, PulsarType

qcm0 = Pulsar("qcm0", dummy_type=PulsarType.PULSAR_QCM)
qrm0 = Pulsar("qrm0", dummy_type=PulsarType.PULSAR_QRM)


# In[19]:


from qcodes import Instrument

from quantify_scheduler.backends.qblox_backend import hardware_compile

config = hardware_compile(sched, qblox_test_mapping)["compiled_instructions"]


# In[20]:


seq_fn = config["qrm0"]["seq0"]["seq_fn"]
qrm0.sequencer0.sequence(seq_fn)


# In[21]:


qcm0.arm_sequencer()
qrm0.arm_sequencer()

qcm0.start_sequencer()
qrm0.start_sequencer()


# In[22]:


from quantify_scheduler.operations.gate_library import X90, Measure, Reset, X
from quantify_scheduler.operations.pulse_library import SquarePulse
from quantify_scheduler.resources import ClockResource

sched = Schedule("Chevron Experiment")
acq_idx = 0

for duration in np.linspace(
    20e-9, 60e-9, 6
):  # NB multiples of 4 ns need to be used due to limitations of the pulsars
    for amp in np.linspace(0.1, 1.0, 10):
        begin = sched.add(Reset("q0", "q1"))
        sched.add(X("q0"), ref_op=begin, ref_pt="end")
        # NB we specify a clock for tutorial purposes,
        # Chevron experiments do not necessarily use modulated square pulses
        square = sched.add(SquarePulse(amp, duration, "q0:mw", clock="q0.01"))
        sched.add(X90("q0"), ref_op=square)
        sched.add(X90("q1"), ref_op=square)
        sched.add(Measure(q0, acq_index=acq_idx), label=f"M q0 {acq_idx}")
        sched.add(
            Measure(q1, acq_index=acq_idx), label=f"M q1 {acq_idx}", ref_pt="start"
        )

        acq_idx += 1


sched.add_resources([ClockResource("q0.01", 6.02e9)])  # manually add the pulse clock


# In[23]:


from quantify_scheduler.compilation import qcompile

cfg = qcompile(sched, transmon_test_config, qblox_test_mapping)