from __future__ import annotations

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


def deer(
    tau1: Quantity,
    tau2: Quantity,
    pi_duration_nv1: Quantity,
    pi_duration_nv2: Quantity,
    frequency_nv1: Quantity,
    frequency_nv2: Quantity,
    amplitude_nv1: Quantity,
    amplitude_nv2: Quantity,
) -> None:
    """Add a DEER-style sequence with an NV1 echo and an NV2 pump pulse.

    tau1 and tau2 are specified center-to-center between the pulses, matching deer_snippet.py.
    """
    assert tau1.is_compatible_with(unit.s), "tau1 must have time units"
    assert tau2.is_compatible_with(unit.s), "tau2 must have time units"
    assert pi_duration_nv1.is_compatible_with(unit.s), "pi_duration_nv1 must have time units"
    assert pi_duration_nv2.is_compatible_with(unit.s), "pi_duration_nv2 must have time units"
    assert frequency_nv1.is_compatible_with(unit.Hz), "frequency_nv1 must have frequency units"
    assert frequency_nv2.is_compatible_with(unit.Hz), "frequency_nv2 must have frequency units"
    assert amplitude_nv1.is_compatible_with(unit.T), "amplitude_nv1 must have magnetic field units"
    assert amplitude_nv2.is_compatible_with(unit.T), "amplitude_nv2 must have magnetic field units"

    pi_half_nv1 = pi_duration_nv1 / 2
    nv1_delay = tau1 - (pi_half_nv1 + pi_duration_nv1) / 2
    nv2_delay = tau1 + tau2 - (pi_half_nv1 + pi_duration_nv2) / 2
    assert nv1_delay >= 0 * unit.s, "tau1 must be at least (pi/2 + pi)/2 to allow non-negative NV1 delays"
    assert nv2_delay >= 0 * unit.s, "tau1 + tau2 must allow a non-negative NV2 delay"

    with sequence():
        drive(duration=pi_half_nv1, frequency=frequency_nv1, amplitude=amplitude_nv1)
        with parallel():
            with sequence():
                delay(duration=nv1_delay)
                drive(duration=pi_duration_nv1, frequency=frequency_nv1, amplitude=amplitude_nv1)
                delay(duration=nv1_delay)
            with sequence():
                delay(duration=nv2_delay)
                drive(duration=pi_duration_nv2, frequency=frequency_nv2, amplitude=amplitude_nv2)
        drive(duration=pi_half_nv1, frequency=frequency_nv1, amplitude=amplitude_nv1)


def deer_sequence(
    start: Quantity,
    stop: Quantity,
    step: Quantity,
    tau1: Quantity,
    pi_duration_nv1: Quantity,
    pi_duration_nv2: Quantity,
    frequency_nv1: Quantity,
    frequency_nv2: Quantity,
    amplitude_nv1: Quantity,
    amplitude_nv2: Quantity,
) -> PulseProgram:
    """Return a DEER program that sweeps tau2 relative to the NV1 echo."""
    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 tau1.is_compatible_with(unit.s), "tau1 must have time units"
    assert pi_duration_nv1.is_compatible_with(unit.s), "pi_duration_nv1 must have time units"
    assert pi_duration_nv2.is_compatible_with(unit.s), "pi_duration_nv2 must have time units"
    assert frequency_nv1.is_compatible_with(unit.Hz), "frequency_nv1 must have frequency units"
    assert frequency_nv2.is_compatible_with(unit.Hz), "frequency_nv2 must have frequency units"
    assert amplitude_nv1.is_compatible_with(unit.T), "amplitude_nv1 must have magnetic field units"
    assert amplitude_nv2.is_compatible_with(unit.T), "amplitude_nv2 must have magnetic field units"

    with pulse_program() as program:
        with parameter_sweep(start=start, stop=stop, step=step) as (_, tau2):
            deer(
                tau1,
                tau2,
                pi_duration_nv1,
                pi_duration_nv2,
                frequency_nv1,
                frequency_nv2,
                amplitude_nv1,
                amplitude_nv2,
            )

    return program
