instant_gradients.jl

module InstantGradients
import StaticArrays: SVector
import JuMP: @constraint, @variable, AbstractJuMPScalar
import ...Variables: qvec, bmat_gradient, duration, variables, get_free_variable, VariableType
import ...BuildingBlocks: GradientBlock, fixed
import ...BuildSequences: global_model

"""
    InstantGradientBlock(; orientation=nothing, qval=nothing, qvec=nothing, rotate=nothing, scale=nothing)

Defines an instantaneous gradient.

## Parameters
- `orientation` sets the gradient orienation (ignored if `qvec` is set). Can be set to a vector for a fixed orientation. Otherwise the orientation will be aligned with the `rotate` (if set) or fully free (if `rotate` is nothing). Set to :flip to point in the inverse of the user-provided `rotate`.
- `rotate`: with which user-set parameter will this gradient be rotated (e.g., :bvec). Default is no rotation.
- `scale`: with which user-set parameter will this gradient be scaled (e.g., :bval). Default is no scaling.

## Variables
- [`qvec`](@ref): Spatial scale and direction on which spins will be dephased due to this pulsed gradient in rad/um.
- [`qval`](@ref): Spatial scale on which spins will be dephased due to this pulsed gradient in rad/um.
"""
struct InstantGradientBlock <: GradientBlock
    qvec :: SVector{3, VariableType}
    rotate :: Union{Nothing, Symbol}
    scale :: Union{Nothing, Symbol}
end

function interpret_orientation(orientation, qval, qvec, will_rotate)
    if !isnothing(qvec)
        return (false, qvec)
    end
    if orientation == :flip
        @assert will_rotate "setting `orientation=:flip` only makes sense if the `rotate` is set as well."
        return (true, SVector{3}(
            -get_free_variable(qval),
            0.,
            0.,
        ))
    end
    if isnothing(orientation)
        if !will_rotate
            return (false, SVector{3}(
                get_free_variable(nothing),
                get_free_variable(nothing),
                get_free_variable(nothing),
            ))
        else
            return (true, SVector{3}(
                get_free_variable(qval),
                0.,
                0.,
            ))
        end
    else
        qval = get_free_variable(qval)
        return (true, (orientation ./ norm(orientation)) * qval)
    end
end

function InstantGradientBlock(; orientation=nothing, qval=nothing, qvec=nothing, rotate=nothing, scale=nothing)
    model = global_model()
    (used_qval, qvec) = interpret_orientation(orientation, qval, qvec, !isnothing(rotate))
    res = InstantGradientBlock(
        qvec,
        rotate,
        scale
    )
    if !used_qval && !isnothing(qval)
        @constraint model qval_sqr(res) == qval^2
    end
    if qval isa AbstractJuMPScalar
        @constraint model qval >= 0
    end
    return res
end


qvec(instant::InstantGradientBlock) = instant.qvec
bmat_gradient(::InstantGradientBlock, qstart=nothing) = zeros(3, 3)
duration(instant::InstantGradientBlock) = 0.
variables(::Type{<:InstantGradientBlock}) = [qvec, qval]


end