#
# Copyright 2024 Keysight Technologies Inc.
#
# This file generates the asset "calibration.qcs". Run `python calibration.py`
# to generate the asset.
#

import math as math

import keysight.qcs as qcs

n_qubits = n_channels = 4
# getting this to account for even/odd n_qubits
last_qubit_index = 2 * math.floor(n_qubits / 2)
cals = qcs.CalibrationSet()

# define the variables in the calibration set
sq_freqs = [5e9 + 100e6 * n for n in range(n_channels)]
frequency_cals = cals.variables.declare(
    qcs.Array("single_qubit_frequencies", value=sq_freqs, dtype=float)
)
sq_amp_cals = cals.variables.declare(
    qcs.Array("single_qubit_amplitudes", value=[0.1] * n_channels, dtype=float)
)
mq_amp_cals = cals.variables.declare(
    qcs.Array(
        "multi_qubit_amplitudes", value=[0.1] * math.floor(n_qubits / 2), dtype=float
    )
)
sq_dur = cals.variables.declare(
    qcs.Scalar("single_qubit_duration", value=30e-9, dtype=float)
)
mq_dur = cals.variables.declare(
    qcs.Scalar("multi_qubit_duration", value=30e-9, dtype=float)
)

# Measurement calibration values
readout_frequencies = cals.variables.declare(
    qcs.Array(
        "readout_frequencies", value=[5.8e9 + 0.05e9 * i for i in range(n_channels)]
    )
)
readout_frequency_common_detuning = cals.variables.declare(
    qcs.Scalar("readout_frequency_detuning", value=0, dtype=float)
)
readout_amplitudes = cals.variables.declare(
    qcs.Array("readout_amplitudes", value=[0.1] * n_channels)
)
readout_phases = cals.variables.declare(
    qcs.Array("readout_phases", value=[0] * n_channels, dtype=float)
)
readout_plateau_time = cals.variables.declare(
    qcs.Scalar("readout_plateau_time", value=100e-9, dtype=float)
)
readout_rise_time = cals.variables.declare(
    qcs.Scalar("readout_rise_time", value=30e-9, dtype=float)
)

readout_pulse_delay = cals.variables.declare(
    qcs.Scalar("readout_pulse_delay", value=40e-9, dtype=float)
)
acquisition_duration = cals.variables.declare(
    qcs.Scalar("acquisition_duration", value=200e-9, dtype=float)
)

iq_classification_pts = cals.variables.declare(
    qcs.Array("ref_pts", value=[[1 + 1j, 1 - 1j]] * n_channels)
)

qubits = qcs.Qudits(range(n_qubits))
channels = qcs.Channels(range(n_channels), "xy_channels", absolute_phase=True)

# define the single qubit gate linker
layers = [-qcs.PAULIS.sigma_z, qcs.PAULIS.sigma_x, qcs.PAULIS.sigma_z]
zxz = qcs.ParametricGate(layers, ["phi", "theta", "phi"])

thetas = qcs.Array(name="thetas", shape=len(qubits), dtype=float)
phis = qcs.Array(name="phis", shape=len(qubits), dtype=float)
sq_instruction = qcs.ParameterizedGate(zxz, [phis, thetas])

sq_repl_prog = qcs.Program()

envelope = qcs.GaussianEnvelope()
waveforms = qcs.RFWaveform(sq_dur, envelope, sq_amp_cals * thetas, frequency_cals, phis)

sq_repl_prog.add_waveform(waveforms, channels)

sq_linker = qcs.ParameterizedLinker(sq_instruction, qubits, sq_repl_prog)
cals.add_linker("single_qubit_linker", sq_linker)

# define the multi qubit gate linker
zx_ham = qcs.PAULIS.sigma_z & qcs.PAULIS.sigma_x
cr = qcs.ParametricGate([zx_ham], ["beta"])

beta = qcs.Array(name="beta", shape=(math.floor(n_qubits / 2),), dtype=float)
cr_instr = qcs.ParameterizedGate(cr, beta)

alpha = qcs.Scalar("alpha", value=0.33, dtype=float)

envelope = qcs.GaussianEnvelope()
pulse = qcs.RFWaveform(mq_dur, envelope, beta * mq_amp_cals, frequency_cals[1::2])

env_cross = qcs.GaussianEnvelope()
pulse_cross = qcs.RFWaveform(
    mq_dur, env_cross, alpha * beta * mq_amp_cals, frequency_cals[1::2]
)

cr_prog_cross = qcs.Program()
cr_prog_cross.add_waveform(pulse, channels[:last_qubit_index:2])
cr_prog_cross.add_waveform(pulse_cross, channels[1::2])

mq_targets = qcs.MultiQudits.from_qudits((qubits[:last_qubit_index:2], qubits[1::2]))
mq_linker = qcs.ParameterizedLinker(cr_instr, mq_targets, cr_prog_cross)
cals.add_linker("multi_qubit_linker_suppression", mq_linker)

# create Linkers for measurement
replacement_program = qcs.Program()

readout_channels = qcs.Channels(
    range(n_channels), "readout_channels", absolute_phase=True
)
acquire_channels = qcs.Channels(range(n_channels), "acquire_channels")

# Defining Readout pulses
pulse = qcs.RFWaveform(
    2 * readout_rise_time,
    qcs.SineEnvelope(),
    amplitude=readout_amplitudes,
    rf_frequency=readout_frequencies + readout_frequency_common_detuning,
    instantaneous_phase=readout_phases,
)

node_pulse = replacement_program.add_waveform(
    pulse.to_flattop(readout_plateau_time),
    readout_channels,
    pre_delay=readout_pulse_delay,
)

# Acquisition and post-processing
integration_envelope = qcs.ConstantEnvelope()
rf_pulse = qcs.RFWaveform(
    acquisition_duration,
    integration_envelope,
    1,
    readout_frequencies + readout_frequency_common_detuning,
)
integration_filter = qcs.IntegrationFilter(rf_pulse)


classifiers = [
    qcs.MinimumDistanceClassifier(iq_classification_pts[qb]) for qb in range(n_channels)
]

# add the acquisition concurrent with the readout pulses
replacement_program.add_acquisition(
    integration_filter,
    acquire_channels,
    classifier=classifiers,
)


meas_linker = qcs.ParameterizedLinker(qcs.Measure(), qubits, replacement_program)

cals.add_linker("measurement_linker", meas_linker)

qcs.save(cals, "calibration.qcs")