from __future__ import annotations

from xq_pulse.pulse.dsl import drive, parameter_sweep, pulse_program
from xq_pulse.pulse.program import PulseProgram
from xq_pulse.util import Quantity, unit


def rabi(
    start: Quantity,
    stop: Quantity,
    step: Quantity,
    frequency: Quantity,
    amplitude: Quantity,
) -> PulseProgram:
    """Return a Rabi pulse program that sweeps the drive duration."""
    assert start.is_compatible_with(unit.s), "start must have time units"
    assert stop.is_compatible_with(unit.s), "stop must have time units"
    assert step.is_compatible_with(unit.s), "step must have time units"
    assert frequency.is_compatible_with(unit.Hz), "frequency must have frequency units"
    assert amplitude.is_compatible_with(unit.T), "amplitude must have magnetic field units"

    with pulse_program() as program:
        with parameter_sweep(start=start, stop=stop, step=step) as (_, duration):
            drive(duration=duration, frequency=frequency, amplitude=amplitude)

    return program
