from __future__ import annotations

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


def x_pulse(frequency: Quantity, amplitude: Quantity, rabi_period: Quantity) -> None:
    drive(duration=rabi_period / 2, frequency=frequency, amplitude=amplitude)


def y_pulse(frequency: Quantity, amplitude: Quantity, rabi_period: Quantity) -> None:
    drive(
        duration=rabi_period / 2,
        frequency=frequency,
        amplitude=amplitude,
        phase=90 * unit.deg,
    )


def x90_pulse(frequency: Quantity, amplitude: Quantity, rabi_period: Quantity) -> None:
    drive(duration=rabi_period / 4, frequency=frequency, amplitude=amplitude)


def xy8_block(
    tau: Quantity,
    rabi_period: Quantity,
    frequency: Quantity,
    amplitude: Quantity,
) -> None:
    assert tau.is_compatible_with(unit.s), "tau must have time units"
    assert rabi_period.is_compatible_with(unit.s), "rabi_period must have time units"
    assert tau >= rabi_period / 2, "tau must be at least rabi_period / 2"

    with sequence():
        x_pulse(frequency, amplitude, rabi_period)
        delay(duration=tau - rabi_period / 2)
        y_pulse(frequency, amplitude, rabi_period)
        delay(duration=tau - rabi_period / 2)
        x_pulse(frequency, amplitude, rabi_period)
        delay(duration=tau - rabi_period / 2)
        y_pulse(frequency, amplitude, rabi_period)
        delay(duration=tau - rabi_period / 2)

        y_pulse(frequency, amplitude, rabi_period)
        delay(duration=tau - rabi_period / 2)
        x_pulse(frequency, amplitude, rabi_period)
        delay(duration=tau - rabi_period / 2)
        y_pulse(frequency, amplitude, rabi_period)
        delay(duration=tau - rabi_period / 2)
        x_pulse(frequency, amplitude, rabi_period)


def xy8(
    tau: Quantity,
    rabi_period: Quantity,
    frequency: Quantity,
    amplitude: Quantity,
    order: int,
) -> None:
    assert tau.is_compatible_with(unit.s), "tau must have time units"
    assert rabi_period.is_compatible_with(unit.s), "rabi_period must have time units"
    assert order >= 1, "order must be at least 1"

    edge_delay = tau / 2 - 3 * rabi_period / 8
    assert edge_delay >= 0 * unit.s, "tau must be at least 3 * rabi_period / 4"

    with sequence():
        x90_pulse(frequency, amplitude, rabi_period)
        delay(duration=edge_delay)
        with for_idx(1, order, 1) as (_, _loop_param):
            xy8_block(tau, rabi_period, frequency, amplitude)
        delay(duration=edge_delay)
        x90_pulse(frequency, amplitude, rabi_period)


def xy8_sequence(
    start: Quantity,
    stop: Quantity,
    step: Quantity,
    rabi_period: Quantity,
    frequency: Quantity,
    amplitude: Quantity,
    order: int,
) -> PulseProgram:
    """Return an XY8 pulse program that sweeps tau."""
    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 rabi_period.is_compatible_with(unit.s), "rabi_period 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"
    assert order >= 1, "order must be at least 1"

    with pulse_program() as program:
        with parameter_sweep(start=start, stop=stop, step=step) as (_, tau):
            xy8(tau, rabi_period, frequency, amplitude, order)

    return program
