from __future__ import annotations

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


def hahn_echo(
    tau: Quantity,
    pi_duration: Quantity,
    frequency: Quantity,
    amplitude: Quantity,
) -> None:
    """Add a Hahn-echo sequence to the current pulse context.

    The delays are applied as explicit free-precession intervals between pulses.
    """
    assert tau.is_compatible_with(unit.s), "tau must have time units"
    assert pi_duration.is_compatible_with(unit.s), "pi_duration 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"

    pi_half = pi_duration / 2
    delay_between_centers = tau - (pi_half + pi_duration) / 2
    assert delay_between_centers >= 0 * unit.s, "tau must be at least (pi/2 + pi)/2 to allow non-negative delays"

    with sequence():
        drive(duration=pi_half, frequency=frequency, amplitude=amplitude)
        delay(duration=delay_between_centers)
        drive(duration=pi_duration, frequency=frequency, amplitude=amplitude)
        delay(duration=delay_between_centers)
        drive(duration=pi_half, frequency=frequency, amplitude=amplitude)


def hahn_echo_sequence(
    start: Quantity,
    stop: Quantity,
    step: Quantity,
    pi_duration: Quantity,
    frequency: Quantity,
    amplitude: Quantity,
) -> PulseProgram:
    """Return a Hahn-echo program that sweeps the free-precession delay."""
    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 pi_duration.is_compatible_with(unit.s), "pi_duration 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 (_, tau):
            hahn_echo(tau, pi_duration, frequency, amplitude)

    return program
