variables.jl

"""
Defines the functions that can be called on parts of an MRI sequence to query or constrain any variables.

In addition this defines:
- [`variables`](@ref): dictionary containing all variables.
- [`VariableType`](@ref): parent type for any variables (whether number or JuMP variable).
- [`get_free_variable`](@ref): helper function to create new JuMP variables.
- [`VariableNotAvailable`](@ref): error raised if variable is not defined for specific [`AbstractBlock`](@ref).
- [`set_simple_constraints`](@ref): call [`apply_simple_constraint`](@ref) for each keyword argument.
- [`apply_simple_constraint`](@ref): set a simple equality constraint.
"""
module Variables
import JuMP: @variable, Model, @objective, objective_function, value, AbstractJuMPScalar
import ..Scanners: gradient_strength, slew_rate
import ..BuildSequences: global_model

"""
Parent type of all components, building block, and sequences that form an MRI sequence.
"""
abstract type AbstractBlock end

all_variables_symbols = [
    :block => [
        :duration => "duration of the building block in ms.",
    ],
    :sequence => [
        :TR => "Time on which an MRI sequence repeats itself in ms.",
    ],
    :pulse => [
        :flip_angle => "The flip angle of the RF pulse in degrees",
        :amplitude => "The maximum amplitude of an RF pulse in kHz",
        :phase => "The angle of the phase of an RF pulse in KHz",
        :frequency => "The off-resonance frequency of an RF pulse (relative to the Larmor frequency of water) in KHz",
        :bandwidth => "Bandwidth of the RF pulse in kHz. To set constraints it is often better to use [`inverse_bandwidth`](@ref).",
        :inverse_bandwidth => "Inverse of bandwidth of the RF pulse in 1/kHz. Also see [`bandwidth`](@ref).",
        :N_left => "The number of zero crossings of the RF pulse before the main peak",
        :N_right => "The number of zero crossings of the RF pulse after the main peak",
        :slice_thickness => "Slice thickness of an RF pulse that is active during a gradient in mm. To set constraints it is often better to use [`inverse_slice_thickness`](@ref).",
        :inverse_slice_thickness => "Inverse of slice thickness of an RF pulse that is active during a gradient in 1/mm. Also, see [`slice_thickness`](@ref).",
    ],
    :gradient => [
        :qvec => "The spatial range and orientation on which the displacements can be detected due to this gradient in rad/um.",
        :qval => "The spatial range on which the displacements can be detected due to this gradient in rad/um (i.e., norm of [`qvec`](@ref)). To set constraints it is often better to use [`qvec`](@ref) or [`qval_square`](@ref).",
        :qval_square => "Square of [`qval`](@ref) in rad^2/um^2.",
        :δ => "Effective duration of a gradient pulse ([`rise_time`](@ref) + [`flat_time`](@ref)) in ms.",
        :rise_time => "Time for gradient pulse to reach its maximum value in ms.",
        :flat_time => "Time of gradient pulse at maximum value in ms.",
        :gradient_strength => "Vector with maximum strength of a gradient along each dimension (kHz/um)",
        :slew_rate => "Vector with maximum slew rate of a gradient along each dimension (kHz/um/ms)",
        :spoiler_scale => "Length-scale on which spins will be dephased by exactly 2π in mm.",
    ],
    :readout => [
        :dwell_time => "Time between two samples in an `ADC` in ms.",
        :nsamples => "Number of samples during a readout. During the optimisation this might produce non-integer values.",
        :fov => "Size of the field of view in mm. To set constraints it is often better to use [`inverse_fov`](@ref).",
        :inverse_fov => "Inverse of size of the field of view in 1/mm. Also see [`fov`](@ref).",
        :voxel_size => "Size of each voxel in mm. To set constraints it is often better to use [`inverse_voxel_size`](@ref).",
        :inverse_voxel_size => "Inverse of voxel size in 1/mm. Also see [`voxel_size`](@ref).",
        :resolution => "Number of voxels in the final readout. During the optimisation this might produce non-integer values, but this will be fixed after optimsation.",
        :oversample => "How much to oversample with ([`nsamples`](@ref) / [`resolution`](@ref))",
        :ramp_overlap => "Fraction of overlap between ADC event and underlying gradient pulse ramp (between 0 and 1)."
    ]
]

"""
Collection of all functions that return variables that can be used to query or constrain their values.
"""
variables = Dict{Symbol, Function}()


for (block_symbol, all_functions) in all_variables_symbols
    for (func_symbol, description) in all_functions
        as_string = "    $func_symbol($block_symbol)\n\n$description\n\nThis represents a variable within the sequence. Variables can be set during the construction of a [`BuildingBlock`](@ref) or used to create constraints after the fact."
        new_func = @eval begin
            function $func_symbol end
            @doc $as_string $func_symbol
            $func_symbol
        end
        variables[func_symbol] = new_func
    end
end


"""
Dictionary with alternative versions of specific function.
    
Setting constraints on these alternative functions can be helpful as it avoids some operations, which the optimiser might struggle with.
"""
alternative_variables = Dict(
    qval => (qval_square, n->n^2, sqrt, false),
    slice_thickness => (inverse_slice_thickness, inv, inv, true),
    spoiler_scale => (qvec, q->1e-3 * 2π/q, l->1e-3 * 2π/l, true),
    bandwidth => (inverse_bandwidth, inv, inv, true),
    fov => (inverse_fov, inv, inv, true),
    voxel_size => (inverse_voxel_size, inv, inv, true),
)


"""
Parent type for any variable in the MRI sequence.

Each variable can be one of:
- a new JuMP variable
- an expression linking this variable to other JuMP variable
- a number

Create these using [`get_free_variable`](@ref).
"""
const VariableType = Union{Number, AbstractJuMPScalar}


"""
    get_free_variable(value; integer=false)

Get a representation of a given `variable` given a user-defined constraint.

The result is guaranteed to be a [`VariableType`](@ref).
"""
get_free_variable(value::Number; integer=false) = integer ? Int(value) : Float64(value)
get_free_variable(value::VariableType; integer=false) = value
get_free_variable(::Nothing; integer=false) = @variable(global_model(), start=0.01, integer=integer)
get_free_variable(value::Symbol; integer=false) = integer ? error("Cannot maximise or minimise an integer variable") : get_free_variable(Val(value))
function get_free_variable(::Val{:min})
    var = get_free_variable(nothing)
    model = global_model()
    @objective model Min objective_function(model) + var
    return var
end
function get_free_variable(::Val{:max})
    var = get_free_variable(nothing)
    model = global_model()
    @objective model Min objective_function(model) - var
    return var
end

"""
    bmat_gradient(gradient::GradientBlock, qstart=(0, 0, 0))

Computes the diffusion-weighting matrix due to a single gradient block in rad^2 ms/um^2.

This should be defined for every `GradientBlock`, but not be called directly.
Instead, the `bmat` and `bval` should be constrained for specific `Pathways`
"""
function bmat_gradient end


"""
    VariableNotAvailable(building_block, variable, alt_variable)

Exception raised when a variable function does not support a specific `BuildingBlock`.
"""
mutable struct VariableNotAvailable <: Exception
    bb :: Type{<:BuildingBlock}
    variable :: Function
    alt_variable :: Union{Nothing, Function}
end
VariableNotAvailable(bb::Type{<:BuildingBlock}, variable::Function) = VariableNotAvailable(bb, variable, nothing)

function Base.showerror(io::IO, e::VariableNotAvailable)
    if isnothing(e.alt_variable)
        print(io, e.variable, " is not available for block of type ", e.bb, ".")
    else
        print(io, e.variable, " is not available for block of type ", e.bb, ". Please use ", e.alt_variable, " instead to set any contsraints or objective functions.")
    end
end


for variable_func in keys(variables)
    if variable_func in [:qval_square, :qval]
        continue
    end
    @eval function Variables.$variable_func(bb::BuildingBlock)
        if Variables.$variable_func in keys(alternative_variables)
            alt_var, forward, backward, _ = alternative_variables[Variables.$variable_func]
            try
                value = alt_var(bb)
                if value isa Number
                    return backward(value)
                elseif value isa AbstractArray{<:Number}
                    return backward.(value)
                end
            catch e
                if e isa VariableNotAvailable
                    throw(VariableNotAvailable(typeof(bb), Variables.$variable_func))
                end
                rethrow()
            end
            throw(VariableNotAvailable(typeof(bb), Variables.$variable_func, alt_var))
        end
        throw(VariableNotAvailable(typeof(bb), Variables.$variable_func))
    end
end

function Variables.qval_square(bb::BuildingBlock, args...; kwargs...)
    vec = Variables.qvec(bb, args...; kwargs...)
    return vec[1]^2 + vec[2]^2 + vec[3]^2
end

Variables.qval(bb::BuildingBlock, args...; kwargs...) = sqrt(Variables.qval_square(bb, args...; kwargs...))


"""
    set_simple_constraints!(block, kwargs)

Add any constraints or objective functions to the variables of a [`BuildingBlock`](@ref).

Each keyword argument has to match one of the functions in [`variables`](@ref)(block).
If set to a numeric value, a constraint will be added to fix the function value to that numeric value.
If set to `:min` or `:max`, minimising or maximising this function will be added to the cost function.
"""
function set_simple_constraints!(block::BuildingBlock, kwargs)
    for (key, value) in kwargs
        if variables[key] in keys(alternative_variables)
            alt_var, forward, backward, to_invert = alternative_variables[variables[key]]
            invert_value(value::VariableType) = forward(value)
            invert_value(value::Symbol) = invert_value(Val(value))
            invert_value(::Val{:min}) = to_invert ? Val(:max) : Val(:min)
            invert_value(::Val{:max}) = to_invert ? Val(:min) : Val(:max)
            invert_value(value::AbstractVector) = invert_value.(value)
            invert_value(value) = value
            try
                apply_simple_constraint!(alt_var(block), invert_value(value))
                return
            catch e
                if !(e isa VariableNotAvailable)
                    rethrow()
                end
            end
        end
        apply_simple_constraint!(variables[key](block), value)
    end
    nothing
end

"""
    apply_simple_constraint!(variable, value)

Add a single constraint or objective to the `variable`.

`value` can be one of:
- `nothing`: do nothing
- `:min`: minimise the variable
- `:max`: maximise the variable
- `number`: fix variable to this value
- `equation`: fix variable to the result of this equation
"""
apply_simple_constraint!(variable, ::Nothing) = nothing
apply_simple_constraint!(variable, value::Symbol) = apply_simple_constraint!(variable, Val(value))
apply_simple_constraint!(variable, ::Val{:min}) = @objective global_model() Min objective_function(global_model()) + variable
apply_simple_constraint!(variable, ::Val{:max}) = @objective global_model() Min objective_function(global_model()) - variable
apply_simple_constraint!(variable, value::VariableType) = @constraint global_model() variable == value
apply_simple_constraint!(variable::AbstractVector, value::AbstractVector) = [apply_simple_constraint!(v1, v2) for (v1, v2) in zip(variable, value)]
apply_simple_constraint!(variable::Number, value::Number) = @assert variable ≈ value "Variable set to multiple incompatible values."


"""
    make_generic(sequence/building_block/component)

Returns a generic version of the `BaseSequence`, `BaseBuildingBlock`, or `BaseComponent`

- Sequences are all flattened and returned as a single `Sequence` containing only `BuildingBlock` objects.
- Any `BaseBuildingBlock` is converted into a `BuildingBlock`.
- Pulses are replaced with `GenericPulse` (except for instant pulses).
- Instant readouts are replaced with `ADC`.
"""
function make_generic end

end