"""
Implementation of the Universally Robust dynamical decoupling sequence from
"Arbitrarily Accurate Pulse Sequences for Robust Dynamical Decoupling" (http://arxiv.org/abs/1609.09416).
"""

from __future__ import annotations

import numpy as np

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


def x(frequency, rabi_period, amplitude):
    drive(duration=rabi_period / 2, frequency=frequency, amplitude=amplitude)


def x_phi(frequency, rabi_period, amplitude, phase):
    drive(duration=rabi_period / 2, frequency=frequency, amplitude=amplitude, phase=phase)


def y(frequency, rabi_period, amplitude):
    drive(duration=rabi_period / 2, frequency=frequency, amplitude=amplitude, phase=90 * unit.deg)


def x90(frequency, rabi_period, amplitude):
    drive(duration=rabi_period / 4, frequency=frequency, amplitude=amplitude)


def y90(frequency, rabi_period, amplitude):
    drive(duration=rabi_period / 4, frequency=frequency, amplitude=amplitude, phase=90 * unit.deg)


def ur8_block(tau, rabi_period, frequency, amplitude):
    with sequence():
        x_phi(frequency, rabi_period, amplitude, phase=0)
        delay(duration=tau - rabi_period / 2)
        x_phi(frequency, rabi_period, amplitude, phase=1 * np.pi / 2)
        delay(duration=tau - rabi_period / 2)
        x_phi(frequency, rabi_period, amplitude, phase=3 * np.pi / 2)
        delay(duration=tau - rabi_period / 2)
        x_phi(frequency, rabi_period, amplitude, phase=2 * np.pi / 2)
        delay(duration=tau - rabi_period / 2)

        x_phi(frequency, rabi_period, amplitude, phase=2 * np.pi / 2)
        delay(duration=tau - rabi_period / 2)
        x_phi(frequency, rabi_period, amplitude, phase=3 * np.pi / 2)
        delay(duration=tau - rabi_period / 2)
        x_phi(frequency, rabi_period, amplitude, phase=1 * np.pi / 2)
        delay(duration=tau - rabi_period / 2)
        x_phi(frequency, rabi_period, amplitude, phase=0)


def ur8(tau, rabi_period, frequency, amplitude, order: int):
    with sequence():
        x90(frequency, rabi_period, amplitude)
        with for_idx(1, order, 1) as (i, _):
            delay(duration=tau / 2 - rabi_period / 4)  # - rabi_period_nv1/8)
            ur8_block(tau, rabi_period, frequency, amplitude)
            delay(duration=tau / 2 - rabi_period / 4)  # - rabi_period_nv1/8)
        x90(frequency, rabi_period, amplitude)


def ur8_sequence(
    start: Quantity,
    stop: Quantity,
    step: Quantity,
    amplitude: Quantity,
    rabi_period: Quantity,
    frequency: Quantity,
    order: int,
):
    with pulse_program() as program:
        with parameter_sweep(start=start, stop=stop, step=step) as (i, tau):
            ur8(tau, rabi_period, frequency, amplitude, order)
    return program


def ur16_block(tau, rabi_period, frequency, amplitude):
    with sequence():
        x_phi(frequency, rabi_period, amplitude, phase=0)
        delay(duration=tau - rabi_period / 2)
        x_phi(frequency, rabi_period, amplitude, phase=1 * np.pi / 4)
        delay(duration=tau - rabi_period / 2)
        x_phi(frequency, rabi_period, amplitude, phase=3 * np.pi / 4)
        delay(duration=tau - rabi_period / 2)
        x_phi(frequency, rabi_period, amplitude, phase=6 * np.pi / 4)
        delay(duration=tau - rabi_period / 2)
        x_phi(frequency, rabi_period, amplitude, phase=2 * np.pi / 4)
        delay(duration=tau - rabi_period / 2)
        x_phi(frequency, rabi_period, amplitude, phase=7 * np.pi / 4)
        delay(duration=tau - rabi_period / 2)
        x_phi(frequency, rabi_period, amplitude, phase=5 * np.pi / 4)
        delay(duration=tau - rabi_period / 2)
        x_phi(frequency, rabi_period, amplitude, phase=4 * np.pi / 4)
        delay(duration=tau - rabi_period / 2)

        x_phi(frequency, rabi_period, amplitude, phase=4 * np.pi / 4)
        delay(duration=tau - rabi_period / 2)
        x_phi(frequency, rabi_period, amplitude, phase=5 * np.pi / 4)
        delay(duration=tau - rabi_period / 2)
        x_phi(frequency, rabi_period, amplitude, phase=7 * np.pi / 4)
        delay(duration=tau - rabi_period / 2)
        x_phi(frequency, rabi_period, amplitude, phase=2 * np.pi / 4)
        delay(duration=tau - rabi_period / 2)
        x_phi(frequency, rabi_period, amplitude, phase=6 * np.pi / 4)
        delay(duration=tau - rabi_period / 2)
        x_phi(frequency, rabi_period, amplitude, phase=3 * np.pi / 4)
        delay(duration=tau - rabi_period / 2)
        x_phi(frequency, rabi_period, amplitude, phase=1 * np.pi / 4)
        delay(duration=tau - rabi_period / 2)
        x_phi(frequency, rabi_period, amplitude, phase=0)


def ur16(tau, rabi_period, frequency, amplitude, order: int):
    with sequence():
        x90(frequency, rabi_period, amplitude)
        with for_idx(1, order, 1) as (i, _):
            delay(duration=tau / 2 - rabi_period / 4)
            ur16_block(tau, rabi_period, frequency, amplitude)
            delay(duration=tau / 2 - rabi_period / 4)
        x90(frequency, rabi_period, amplitude)


def ur16_sequence(
    start: Quantity,
    stop: Quantity,
    step: Quantity,
    amplitude: Quantity,
    rabi_period: Quantity,
    frequency: Quantity,
    order: int,
):
    with pulse_program() as program:
        with parameter_sweep(start=start, stop=stop, step=step) as (i, tau):
            ur16(tau, rabi_period, frequency, amplitude, order)
    return program
