generic_pulses.jl

module GenericPulses

import Polynomials: fit
import ...AbstractTypes: RFPulseComponent, split_timestep
import ....Variables: make_generic, variables, @defvar, adjust_internal


"""
    GenericPulse(time, amplitude, phase, effective_time=<halfway>)
    GenericPulse(time, amplitude; phase=0., frequency=0., effective_time=<halfway>)

Create a Pulse profile that has been fully defined by N control point.

All arguments should be arrays of the same length N defining these control points.

This pulse has no free variables.

- `time`: time since the start of this [`BuildingBlock`](@ref) in ms.
- `amplitude`: amplitude of the RF pulse at every timepoint in kHz.
- `phase`: phase of the RF pulse at every timpoint in degrees. If not set explicitly it will be determined by the provided starting `phase` (degrees) and the `frequency` (kHz).
- `effective_time`: the time that the RF pulse should be considered to have taken place when computing a `Pathway` (defaults: whenever half of the final flip angle has been achieved for on-resonance spins).
"""
struct GenericPulse <: RFPulseComponent
    time :: Vector{Float64}
    amplitude :: Vector{Float64}
    phase :: Vector{Float64}
    effective_time :: Float64
    function GenericPulse(time::AbstractVector{<:Number}, amplitude::AbstractVector{<:Number}, phase::AbstractVector{<:Number}, effective_time=nothing)
        @assert length(time) == length(amplitude)
        @assert length(time) == length(phase)
        @assert all(time .>= 0)
        if isnothing(effective_time)
            test_res = new(Float64.(time), Float64.(amplitude), Float64.(phase), 0.)
            effective_time = variables.time_halfway_flip(test_res)
        end
        return new(Float64.(time), Float64.(amplitude), Float64.(phase), Float64(effective_time))
    end
end

function GenericPulse(time::AbstractVector{<:Number}, amplitude::AbstractVector{<:Number}; phase=0., frequency=0., effective_time=nothing)
    return GenericPulse(time, amplitude, (time .- time[1]) .* (frequency * 360) .+ phase, effective_time)
end

"""
    GenericPulse(pulse, t1, t2)

Creates a new [`GenericPulse`](@ref) by slicing another pulse between `t1` and `t2`
"""
function GenericPulse(pulse::GenericPulse, t1::Number, t2::Number)
    if t1 < pulse.time[1] || t2 > pulse.time[end]
        error("Cannot extrapolate GenericPulse")
    end
    use = t1 .<= pulse.time .<= t2
    tnew = pulse.time[use]
    anew = pulse.amplitude[use]
    pnew = pulse.phase[use]
    if !(t1 ≈ tnew[1])
        pushfirst!(tnew, t1)
        pushfirst!(anew, variables.amplitude(pulse, t1))
        pushfirst!(pnew, variables.phase(pulse, t1))
    elseif !(t2 ≈ tnew[end])
        push!(tnew, t2)
        push!(anew, variables.amplitude(pulse, t2))
        push!(pnew, variables.phase(pulse, t2))
    end
    return GenericPulse(tnew .- t1, anew, pnew, pulse.effective_time - t1)
end

GenericPulse(pulse::RFPulseComponent, t1::Number, t2::Number) = GenericPulse(make_generic(pulse), t1, t2)

@defvar duration(fp::GenericPulse) = maximum(fp.time)
@defvar effective_time(pulse::GenericPulse) = pulse.time[findmax(abs.(pulse.amplitude))]

@defvar pulse begin
    amplitude(fp::GenericPulse) = maximum(abs.(fp.amplitude))
    phase(pulse::GenericPulse) = pulse.phase[findmax(abs.(pulse.amplitude))[2]]
    flip_angle(pulse::GenericPulse) = sum(get_weights(pulse) .* pulse.amplitude) * 360
    function time_halfway_flip(pulse::GenericPulse)
        w = get_weights(pulse)
        flip_so_far = cumsum(w .* pulse.amplitude)
        return pulse.time[findfirst(f -> f >= flip_so_far[end] / 2, flip_so_far)]
    end
end

for fn in (:amplitude, :phase)
    @eval function variables.$fn.f(fp::GenericPulse, time::Number)
        i2 = findfirst(t->t > time, fp.time)
        if isnothing(i2)
            @assert time ≈ fp.time[end]
            return fp.$(fn)[end]
        end
        i1 = i2 - 1
        if time ≈ fp.time[i1]
            return fp.$(fn)[i1]
        end
        w2 = (time - fp.time[i1]) / (fp.time[i2] - fp.time[i1])
        w1 = (fp.time[i2] - time) / (fp.time[i2] - fp.time[i1])
        return w1 * fp.$(fn)[i1] + w2 * fp.$(fn)[i2]
    end
end


@defvar pulse function frequency(gp::GenericPulse, time::Number)
    i2 = findfirst(t -> t > time, gp.time)
    if isnothing(i2)
        @assert time ≈ gp.time[end]
        i2 = length(gp.time)
    end
    if !isone(i2) && time ≈ gp.time[i2 - 1]
        i2 -= 1
    end
    if time ≈ gp.time[i2]
        if i2 == 1
            return (gp.phase[2] - gp.phase[1]) / (gp.time[2] - gp.time[1]) / 360
        elseif i2 == length(gp.time)
            return (gp.phase[end] - gp.phase[end-1]) / (gp.time[end] - gp.time[end-1]) / 360
        end
        return (gp.phase[i2 + 1] - gp.phase[i2 - 1]) / (gp.time[i2 + 1] - gp.time[i2 - 1]) / 360
    end
    @assert !isone(i2)
    return (gp.phase[i2] - gp.phase[i2 - 1]) / (gp.time[i2] - gp.time[i2 - 1]) / 360
end


function get_weights(pulse::GenericPulse)
    weights_double = pulse.time[2:end] - pulse.time[1:end-1]
    return [weights_double[1] / 2, ((weights_double[1:end-1] + weights_double[2:end]) / 2)..., weights_double[end]/2]
end



make_generic(gp::GenericPulse) = gp

function split_timestep(gp::GenericPulse, precision)
    function second_derivative(arr)
        max_second_der = 0.
        for index in 2:length(arr)-1
            poly = fit(gp.time[index-1:index+1], arr[index-1:index+1])
            if iszero(length(poly.coeffs))
                continue
            end
            second_der = abs(poly.coeffs[end])
            if second_der > max_second_der
                max_second_der = second_der
            end
        end
        return max_second_der
    end
    max_second_der = max(
        second_derivative(gp.amplitude ./ maximum(gp.amplitude)),
        second_derivative(gp.phase ./ 360),
    )
    if iszero(max_second_der)
        return Inf
    end
    return sqrt(2 * precision / max_second_der)
end

function adjust_internal(block::GenericPulse; scale=1., frequency=0., stretch=1.)
    GenericPulse(
        block.time .* stretch,
        block.amplitude .* scale,
        block.phase .+ (360. * frequency) .* (block.time .- effective_time(block)),
        block.effective_time * stretch,
    )
end

end