# %% [markdown]
# Characterizing Input Offset
# ===========================

# %% [markdown]
# The Qblox QRM and QRM-RF use analog-to-digital converters (ADCs) to digitize the incoming analog signals. Due to thermal, manufacturing effects or other factors, the digitized signal is not centered around 0V but has a small offset, which we will hereby refer to as input (ADC) offset. This input offset can get demodulated and integrated by the hardware along with the signal. The integrated offset can then show up as oscillations in the result, e.g. during a frequency sweep. In this section we show how to measure and calibrate away this offset to prevent such effects.

# %% [markdown]
# Measuring the Offset
# --------------------
#

# %% [markdown]
# We will use a simple scope acquisition to determine the mean value of this offset. Before proceeding, please **make sure that there is no DC signal going into** to your QRM.

# %% [markdown] tags=[]
# Setup
# -----
#
# First, we are going to import the required packages.

# %% tags=["imports"]
import json
import warnings

import matplotlib.pyplot as plt
import numpy as np

# %% tags=["SUBSTITUTE_HEADER"]
from header import module

# %% [markdown]
# Specify acquisitions
# --------------------
#
# We need to specify the acquisitions so that the instrument can allocate the required memory for its acquisition list. In this case we will create 5 acquisition specifications that each create one or multiple bins.

# %% tags=[]
# Acquisitions
acquisitions = {
    "non_weighed": {"num_bins": 10, "index": 0},
}


# %%
# Create a helper function that creates an acquisition program, runs the sequencer and calculates the ADC offsets based on acquired results.


def acquire_scope_and_calc_offsets() -> tuple[float, float]:
    seq_prog = """
        move    1000, R0        #Loop iterator.

    loop: acquire 0,1,20000   #Acquire bins and store them in "non_weighed" acquisition.
        loop     R0, @loop #Run until number of iterations is done.

        stop                #Stop.
    """

    # Add sequence program, waveforms, weights and acquisitions to single dictionary and write to JSON file.
    sequence = {
        "waveforms": {},
        "weights": {},
        "acquisitions": acquisitions,
        "program": seq_prog,
    }
    with open("sequence.json", "w", encoding="utf-8") as file:
        json.dump(sequence, file, indent=4)
        file.close()

    # Upload sequence.
    module.sequencer0.sequence("sequence.json")

    # Arm and start sequencer.
    module.arm_sequencer(0)
    module.start_sequencer()

    # Wait for the acquisition to stop with one minute timeout
    module.get_acquisition_status(0, timeout=1)

    # Retrieve results
    module.store_scope_acquisition(0, "non_weighed")
    non_weighed_acq = module.get_acquisitions(0)
    I_data = np.array(non_weighed_acq["non_weighed"]["acquisition"]["scope"]["path0"]["data"])
    Q_data = np.array(non_weighed_acq["non_weighed"]["acquisition"]["scope"]["path1"]["data"])

    # Plot results
    fig, ax = plt.subplots(1, 1)
    ax.plot(I_data, label="I")
    ax.plot(Q_data, label="Q")
    ax.set_xlabel("Time (ns)", fontsize=20)
    ax.set_ylabel("Relative amplitude", fontsize=20)
    plt.legend()
    plt.show()

    # Print mean offset values
    I_offset, Q_offset = np.mean(I_data), np.mean(Q_data)
    print(f"I Offset : {I_offset * 1e3:.3f} mV \nQ Offset : {Q_offset * 1e3:.3f} mV")

    if np.isnan(I_offset):
        warnings.warn("Determining offset failed, setting to 0.")
        I_offset = 0
    if np.isnan(Q_offset):
        Q_offset = 0
        warnings.warn("Determining offset failed, setting to 0.")

    return I_offset, Q_offset


# Map sequencer to specific outputs (but first disable all sequencer connections)
module.disconnect_outputs()
module.disconnect_inputs()
# %% tags=["connect"]
# connect sequencers

# %% tags=["adc_offset"]
adc_offset_path0 = lambda _: None  # noqa: E731
adc_offset_path1 = lambda _: None  # noqa: E731

# %%
# Ensure offset is reset to zero
adc_offset_path0(0.0)
adc_offset_path1(0.0)

# Configure scope mode
module.scope_acq_sequencer_select(0)
module.scope_acq_trigger_mode_path0("sequencer")
module.scope_acq_trigger_mode_path1("sequencer")
module.scope_acq_avg_mode_en_path0(True)
module.scope_acq_avg_mode_en_path1(True)

# Configure integration
module.sequencer0.integration_length_acq(16000)
module.sequencer0.thresholded_acq_rotation(0)
module.sequencer0.thresholded_acq_threshold(0)

# %% tags=["configure"]
# Configure the sequencer


# %% tags=[]
I_offset, Q_offset = acquire_scope_and_calc_offsets()

# %% [markdown]
# Correcting the Offsets
# ----------------------
# %% tags=["correct_1"]


# %% tags=[]
adc_offset_path0(-I_offset)  # Negative sign to compensate for the offset
adc_offset_path1(-Q_offset)

# %% [markdown]
# Repeating the offset measurement as before, we get:

# %% tags=[]
I_offset, Q_offset = acquire_scope_and_calc_offsets()


# %% [markdown]
# Advanced ADC Offset Calibration : Curve Fitting Method
# ------------------------------------------------------

# %% [markdown]
# As you may have noticed in the previous section, manually compensating for the offset does not entirely eliminate it, leaving some residual offset. This is because of the non-linear effects of these input offsets. To circumvent this, one can curve-fit the dependence of the set offset value on the actual offset value. In the following section, we do this by fitting 10 setpoints against the 10 measured offset values from the binned acquisition with a function. We then find the roots/zeros of this function to set the actual offset to zero.


# %% tags=[]
# Define helper functions
def get_real_root(coeffs: np.array) -> np.array:
    if all(np.isnan(coeffs)):
        warnings.warn("All NaN array. Returning 0.")
        return 0
    for root in np.roots(coeffs):
        if root.imag == 0:
            output = root
    return np.real(output)


def get_curve(x: np.array, coeffs: np.array) -> np.array:
    y = 0
    for i, coeff in enumerate(coeffs):
        y += coeff * x ** (len(coeffs) - (i + 1))
    return y


# %% tags=[]
# Define the program that we will run
seq_prog = """
      move    0,R0        #Loop iterator.
      nop

loop:
      move    100, R1     #Averages
avg_loop:
      acquire 0,R0,20000   #Acquire bins and store them in "non_weighed" acquisition.
      loop    R1, @avg_loop
      add     R0,1,R0     #Increment iterator
      nop                 #Wait a cycle for R0 to be available.
      jlt     R0,10,@loop #Run until number of iterations is done.

      stop                #Stop.
"""

# %% tags=[]
# Add sequence program, waveforms, weights and acquisitions to single dictionary and write to JSON file.
sequence = {
    "waveforms": {},
    "weights": {},
    "acquisitions": acquisitions,
    "program": seq_prog,
}
with open("sequence.json", "w", encoding="utf-8") as file:
    json.dump(sequence, file, indent=4)
    file.close()

# %% tags=[]
# Upload sequence.
module.sequencer0.sequence("sequence.json")

# %% [markdown]
# Let's start the sequence and retrieve the results.

# %% tags=[]
# Set the domain (X) around these original offset values

I_offset_setpoints = np.linspace(-10e-3 + adc_offset_path0(), +10e-3 + adc_offset_path0(), 12)
Q_offset_setpoints = np.linspace(-10e-3 + adc_offset_path1(), +10e-3 + adc_offset_path1(), 12)

int_len = module.sequencer0.integration_length_acq()
I_offset_sum, Q_offset_sum = [], []

for I_offset_i, Q_offset_i in zip(I_offset_setpoints, Q_offset_setpoints):
    module.delete_acquisition_data(0, all=True)
    adc_offset_path0(I_offset_i)
    adc_offset_path1(Q_offset_i)
    module.arm_sequencer(0)
    module.start_sequencer()
    module.get_acquisition_status(0, timeout=1)
    non_weighed_acq = module.get_acquisitions(0)["non_weighed"]
    I_offset_sum += [
        np.mean(non_weighed_acq["acquisition"]["bins"]["integration"]["path0"]) * 1e3 / int_len
    ]
    Q_offset_sum += [
        np.mean(non_weighed_acq["acquisition"]["bins"]["integration"]["path1"]) * 1e3 / int_len
    ]

output = {
    "offsets_I": I_offset_setpoints,
    "offsets_Q": Q_offset_setpoints,
    "I_m": I_offset_sum,
    "Q_m": Q_offset_sum,
}

# %% tags=["nbsphinx-thumbnail"]
plt.figure()

## Fit I offset and find its root
coeffs = np.polyfit(np.array(output["offsets_I"]) * 1e3, np.array(output["I_m"]), 3)
new_I_offset = get_real_root(coeffs)

plt.plot(
    output["offsets_I"] * 1e3,
    get_curve(np.array(output["offsets_I"]) * 1e3, coeffs),
    c="b",
    label="I",
)
plt.scatter(output["offsets_I"] * 1e3, output["I_m"], c="b")

# Fit Q offset and find its root
coeffs = np.polyfit(np.array(output["offsets_Q"]) * 1e3, np.array(output["Q_m"]), 3)
new_Q_offset = get_real_root(coeffs)

plt.plot(
    output["offsets_Q"] * 1e3,
    get_curve(np.array(output["offsets_Q"]) * 1e3, coeffs),
    c="r",
    label="Q",
)
plt.scatter(output["offsets_Q"] * 1e3, output["Q_m"], c="r")

# Plot zeros on the plot
plt.axvline(x=new_I_offset)
plt.axvline(x=new_Q_offset)
plt.axhline(y=0)

plt.xlabel("Set Offset (mV)")
plt.ylabel("Observed Offset (mV)")
plt.legend()
plt.grid()
print(
    f"Offset setpoints corresponding to observed zero offset: \nI offset: {new_I_offset}\nQ offset: {new_Q_offset}"
)

# %% [markdown]
# Applying The Offset Corrections
# -------------------------------

# %% [markdown]
# Using the zeros obtained from the curve fitting, we attempt to correct for the input offsets again:

# %% tags=[]
adc_offset_path0(new_I_offset * 1e-3)  # Multiplying by 1e-3 to convert to mV
adc_offset_path1(new_Q_offset * 1e-3)

I_offset_calib, Q_offset_calib = acquire_scope_and_calc_offsets()

# %% [markdown]
# As you can see, you have calibrated away the offsets to sub-millivolt range, which is at the limit of the ADC resolution.

# %% [markdown]
# <div class="alert alert-success">
# We advise you to check/calibrate these offset values every few days if they are vital for your experiments, especially when you are dealing with low input signals. </div>

# %% tags=["SUBSTITUTE_FOOTER"]