# scripts to analyse behavioral data in various ways

using SpecialFunctions, Distributions, Optim, Calculus, MAT, DataFrames, Gadfly, Colors

# constants and common functions
const sqrttwo = √(2)
const sqrttwoπ = √(2 * π)
Φ(x) = (1 + erf(x / sqrttwo)) / 2

###############################################################################
# loading experimental data
#
# Most of the dataset is self-explanatory.
# Stimulus ranges from 1 to 8 (Stimulus = 1..4 -> right, 5..8 -> left). OSD is
# odor sampling duration [s]. Outcome is 1 (rewarded), 0 (not rewarded due to
# wrong choice), or NaN (not rewarded due to too fast choice).
#
# For the interleaved :itlvd dataset, the format is slightly different, in
# that Stimulus ranges from 1 to 32, with the following ordering:
#  1.. 4 (right easy->hard)  8.. 5 (left hard->easy) concentration 0.1
#  9..12 (right easy->hard) 16..13 (left hard->easy) concentration 0.01
# 17..20                    24..21                   concentration 0.001
# 25..28                    32..29                   concentration 0.0001
# In the below, the Id's for left will be inverted to recover the 1-8, 9-16, ...
# ordering.
#
# On top of the stimulus id, the method adds the following stimulus identifier
# mixw 1  mixmag 4  for 100/  0 mix | concentr 4  for 0.1     concentration
# mixw 2  mixmag 3  for  80/ 20 mix | concentr 3  for 0.01    concentration
# mixw 3  mixmag 2  for  68/ 32 mix | concentr 2  for 0.001   concentration
# mixw 4  mixmag 1  for  56/ 44 mix | concentr 1  for 0.0001  concentration
# mixw 5  mixmag 1  for  44/ 56 mix
# mixw 6  mixmag 2  for  32/ 68 mix
# mixw 7  mixmag 3  for  20/ 80 mix
# mixw 8  mixmag 4  for   0/100 mix

# dataset is :categ, :ident, :itlvd
# a ratid of 0 causes the function to return all rats
function load_data(dataset::Symbol, ratid::Int = 0)
    @assert dataset == :categ || dataset == :ident || dataset == :itlvd
    @assert ratid ∈ [0, 1, 2, 3, 4]  # only 4 rats
    # load MAT file
    d = dataset == :categ ?
        matread("../data/DATA CAT.mat")["DATA_cat"] :
        (dataset == :ident ?
         matread("../data/DATA ID.mat")["DATA_id"] :
         matread("../data/DATA INT.mat")["DATA_int"])
    stimid = round.(Int, vec(d["Stimulus"]))
    resp = round.(Int, vec(d["ChoiceDir"]))
    trials = length(stimid)
    if dataset == :itlvd
        # :itlvd dataset
        @assert minimum(stimid) == 1 && maximum(stimid) == 32
        corrresp = zeros(Int, trials)
        stimmag = zeros(Int, trials)
        corr = zeros(Bool, trials)
        mixw = zeros(Int, trials)
        mixmag = zeros(Int, trials)
        concentr = zeros(Int, trials)
        # process stimulus properties separately per concentration (1-8, 9-16, ...)
        for i = 1:4
            # reverse stimulus ordering for 5 -> 8, 13 -> 16, etc.
            righttrials = (stimid .>= 8i-3) .& (stimid .<= 8i)
            lefttrials = (stimid .>= 8i-7) .& (stimid .<= 8i-4)
            alltrials = righttrials .| lefttrials
            stimid[righttrials] = (16i - 3) .- stimid[righttrials]
            corrresp[righttrials] .= 1 # left = 0, right = 1
            stimmag[righttrials] .= stimid[righttrials]
            stimmag[lefttrials] .= (16i - 7) .- stimid[lefttrials]
            corr[lefttrials] = resp[lefttrials] .== 0
            corr[righttrials] = resp[righttrials] .== 1
            mixw[alltrials] = stimid[alltrials] .- 8(i-1)       # to 1-8 range
            mixmag[alltrials] = stimmag[alltrials] .- (8i - 4)  # to 1-4 range
            concentr[alltrials] .= 5-i
        end
    else
        # :categ or :ident dataset - those are simpler
        @assert minimum(stimid) == 1 && maximum(stimid) == 8
        corrresp = Int[i < 5 ? 0 : 1 for i in stimid]
        stimmag = Int[i < 5 ? 9 - i : i for i in stimid]
        corr = (stimid .>= 5) .& (resp .== 1) .| (stimid .<= 4) .& (resp .== 0)
        if dataset == :categ
            mixw = copy(stimid)
            mixmag = stimmag .- 4
            concentr = 4 * ones(Int, trials)
        else
            mixw = 7corrresp .+ 1  # either 1 or 8
            mixmag = 4 * ones(Int, trials)
            concentr = stimmag .- 4
        end
    end
    rew = [!isnan(i) && round(Int, i) == 1 for i in vec(d["Outcome"])]
    # turn into DataFrame
    df = DataFrame(rat = round.(Int, vec(d["RAT"])),
                   session = round.(Int, vec(d["Session"])),
                   trial = round.(Int, vec(d["Trial"])),
                   stimid = stimid,
                   stimmag = stimmag,
                   mixw = mixw,
                   mixmag = mixmag,
                   concentr = concentr,
                   corrresp = corrresp,
                   previd = [stimid[end]; stimid[1:end-1]],    # ign. boundaries
                   prevmag = [stimmag[end]; stimmag[1:end-1]], # "
                   resp = resp,
                   t = vec(d["OSD"]),
                   corr = corr,
                   rew = rew,
                   prevcorr = [corr[end]; corr[1:end-1]],      # ign. boundaries
                   prevrew = [rew[end]; rew[1:end-1]])
    if ratid != 0
        df = df[df[:rat] .== ratid, :]
    end
    return df
end

###############################################################################
# fitting psychometric function, assuming the probabilistic model
#
# p(y = 1 | x) = (1-pl) Phi((x - μ) / σ) + pl pb
#
# with mean μ, standard deviation σ, lapse rate pl and lapse bias bp.
# THe log-likelihood is defined over J xj's, nj's and pj's, where nj is the
# number of observations for the corresponding xj, and pj = sum yj / nj is the
# fraction of y's that are 1 at this xj.
#
# The parameter vector is θ = [μ, σ, pl, pb]

# p(y = 1 | x) for given parameters
psychpy1(θ::AbstractVector, x) = (1 - θ[3]) .* Φ.((x .- θ[1]) ./ θ[2]) .+ θ[3] * θ[4]

# log-likelihood
function psychllh(θ::AbstractVector,
                  x::AbstractVector, n::AbstractVector, p::AbstractVector)
    @assert length(x) == length(n) == length(p)
    @assert length(θ) == 4

    # parameter bound checking
    μ, σ, pl, pb = θ
    σ > 0.0 && 0.0 <= pl <= 1.0 && 0.0 <= pb <= 1.0 || return -Inf
    # compute llh
    llh = 0.0
    for j = 1:length(x)
        res = (x[j] - μ) / σ
        llh += n[j] * (p[j] * log((1-pl) * Φ(res) + pl * pb) +
                       (1-p[j]) * log((1-pl) * Φ(-res) + pl * (1-pb)))
    end
    return llh
end

# gradient g
function psychgrad(θ::AbstractVector,
                   x::AbstractVector, n::AbstractVector, p::AbstractVector)
    @assert length(x) == length(n) == length(p)
    @assert length(θ) == 4
    g = Vector{Float64}(undef, 4)

    # parameter bound checking
    μ, σ, pl, pb = θ
    if !(σ > 0.0 && 0.0 <= pl <= 1.0 && 0.0 <= pb <= 1.0)
        g[1] = 0.0
        g[2] = σ > 0.0 ? 0.0 : Inf
        g[3] = 0.0 <= pl <= 1.0 ? 0.0 : (pl < 0.0 ? Inf : -Inf)
        g[4] = 0.0 <= pb <= 1.0 ? 0.0 : (pb < 0.0 ? Inf : -Inf)
        return g
    end

    # compute gradient
    g[:] .= 0.0
    for j = 1:length(x)
        res = (x[j] - μ) / σ
        Nres = exp(- 0.5 * abs2(res)) / (sqrttwoπ * σ)
        Φres, Φnres = Φ(res), Φ(-res)
        plΦres, plΦnres = (1-pl) * Φres + pl * pb, (1-pl) * Φnres + pl * (1-pb)
        plfrac = (1-p[j]) / plΦnres - p[j] / plΦres

        g[1] += n[j] * (1-pl) * Nres * plfrac
        g[2] += n[j] * (1-pl) * Nres * res * plfrac
        g[3] += n[j] * (p[j] * (pb - Φres) / plΦres +
                        (1 - p[j]) * (1 - Φnres - pb) / plΦnres)
        g[4] -= n[j] * pl * plfrac
    end
    return g
end
# returing llh, modifying gradient, to be used by NLopt.jl
function psychllhgrad!(θ::AbstractVector, g::AbstractVector,
                       x::AbstractVector, n::AbstractVector, p::AbstractVector)
    if !isempty(g)
        g[:] = psychgrad(θ, x, n, p)
    end
    return psychllh(θ, x, n, p)
end
# modifying negative gradient, to be used by Optim.jl
function negpsychgrad!(θ::AbstractVector, g::AbstractVector,
                       x::AbstractVector, n::AbstractVector, p::AbstractVector)
    g[:] = -psychgrad(θ, x, n, p)
    return g
end

# Hessian h, by finite differences
psychhess(θ, x, n, p) = hessian(θt -> psychllh(θt, x, n, p), θ)
# Hessian diagonal by taking (finite difference) derivative of gradient
function psychhessdiag(θ, x, n, p)
    @assert length(θ) == 4
    psychgradi(i, θi) = (θtmp = copy(θ); θtmp[i] = θi; psychgrad(θtmp, x, n, p)[i])
    return Float64[derivative(θi -> psychgradi(i, θi), θ[i]) for i in 1:4]
end

###############################################################################
# functions analyzing behavioral data
#
# For psycometric/chronometric functions, objective correctness of the choices
# is used. For sequential dependencies, the measure of 'correctness' is if
# reward has been given in the previous trial.

# statistics that also work for empty lists
_0mean(x) = isempty(x) ? zero(eltype(x)) : mean(x)
_0var(x) = isempty(x) ? Inf : var(x)
_0sem(x, n) = n == 0 ? Inf : √(_0var(x) / n)
_0sem(x) = _0sem(x, length(x))

# returns various RT / choice statistics over the whole set
_tcstats(df, datasets::Int=1) = DataFrame(
    n = size(df, 1) / datasets,
    nr = sum(df[:resp] .== 1) / datasets,
    ncorr = sum(df[:corr]) / datasets,
    pcorrμ = _0mean(df[:corr]),
    pcorrsem = √(datasets * _0mean(df[:corr]) .* (1 - _0mean(df[:corr])) ./ size(df, 1)),
    prμ = _0mean(df[:resp] .== 1),
    prsem = √(datasets * _0mean(df[:resp] .== 1) .* (1 - _0mean(df[:resp] .== 1)) ./ size(df, 1)),
    tμ = _0mean(df[:t]),
    tσ = √(_0var(df[:t])),
    tsem = _0sem(df[:t], size(df, 1) / datasets),
    tcorrμ = _0mean(df[df[:corr], :t]),
    tcorrσ = √(_0var(df[df[:corr], :t])),
    tcorrsem = _0sem(df[df[:corr], :t], sum(df[:corr]) / datasets),
    tincorrμ = _0mean(df[.!df[:corr], :t]),
    tincorrσ = √(_0var(df[.!df[:corr], :t])),
    tincorrsem = _0sem(df[.!df[:corr], :t], sum(.!df[:corr]) / datasets)
    )

# returns psychometric/chronometric statistics, once ordered by stimid, once by
# stimmag.
# If datasets > 1, the returned sem's are average sem's per dataset, assuming
# that the provided data is from the given number of datasets.
getpsychchronstats(df, datasets::Int=1) = (
    by(df, :stimid, df -> _tcstats(df, datasets)),
    by(df, :stimmag, df -> _tcstats(df, datasets)))

# returns sequential statistics, seq_df, fit_df, biasid_df, biasmag_df
# seq_df returns sequential choice statistics, conditional on previous reward,
# previous stimulus id, and current stimulus id.
# fit_df fits these sequential psychometric curve, with one fit per
# (pre. rew, prev. stim) combination. The function first performs a global
# fit (stimid/stimmag = 0) to determine the slope of the psychometric curve,
# and fixes this slope in all subsequent (conditional) fits.
# biasid_df and biasmag_df use the fits to compute the biases, either per
# id or magnitude of the previous stimulus.
function getseqstats(df, datasets::Int=1)
    seq_df = by(df, [:stimid, :previd, :prevrew], df -> _tcstats(df, datasets))

    # find psychometric curve fit parameters
    fit_df = DataFrame(previd = Int[], prevrew = Bool[],
                       μ = Float64[], σ = Float64[], pl = Float64[], pb = Float64[],
                       μsd = Float64[], σsd = Float64[], plsd = Float64[], pbsd = Float64[])
    # fit overall psychometric curve
    psych_df = by(df, :stimid,
                  df -> DataFrame(n = size(df, 1), pr = _0mean(df[:resp])))
    x, n, p = psych_df[:stimid], psych_df[:n] ./ datasets, psych_df[:pr]
    θini = [mean(x), 1.0, 0.1, 0.5]
    θmin = [-Inf, 0.0, 0.0, 0.0]
    θmax = [Inf, Inf, 1.0, 1.0]
    res = optimize(
        OnceDifferentiable(
            θ -> -psychllh(θ, x, n, p),
            (g, θ) -> negpsychgrad!(θ, g, x, n, p), θini),
        θmin, θmax, θini, Fminbox(LBFGS()))
    globalθ = Optim.minimizer(res)
    globalθsd = sqrt.(max.(0.0, -1 ./ psychhessdiag(globalθ, x, n, p)))
    push!(fit_df, [0, true,
                   globalθ[1], globalθ[2], globalθ[3], globalθ[4],
                   globalθsd[1], globalθsd[2], globalθsd[3], globalθsd[4]])

    # fit per-condition psychometric curve, fixing slope to above slope
    stimids = sort(unique(psych_df[:stimid]))
    for prevrew in [true, false], sid in stimids
        psych_df = by(df[(df[:prevrew] .== prevrew) .& (df[:previd] .== sid), :], :stimid,
                      df -> DataFrame(n = size(df, 1), pr = _0mean(df[:resp])))
        x, n, p = psych_df[:stimid], psych_df[:n] ./ datasets, psych_df[:pr]
        # optimize over [μ, pl, pb], while using σ = globalθ[2]
        function psychllhgradσ!(subθ, subg)
            θ = [subθ[1], globalθ[2], subθ[2], subθ[3]]
            if !isempty(subg)
                subg[:] = -psychgrad(θ, x, n, p)[[1, 3, 4]] # ignore σ gradient
            end
        end
        res = optimize(
            OnceDifferentiable(
                subθ -> -psychllh([subθ[1], globalθ[2], subθ[2], subθ[3]], x, n, p),
                (subg, subθ) -> psychllhgradσ!(subθ, subg), θini[[1, 3, 4]]),
            θmin[[1, 3, 4]], θmax[[1, 3, 4]], θini[[1, 3, 4]], Fminbox(LBFGS()))
        θ = Optim.minimizer(res)
        θ = Float64[θ[1], globalθ[2], θ[2], θ[3]]
        # approximation, as off-diagonal terms might influence inv. Hessian
        θsd = sqrt.(max.(0.0, -1 ./ psychhessdiag(θ, x, n, p)))
        # ensure σ^2 < μ (1 - μ)
        θsd[3] = min(θsd[3], 0.999999 * sqrt.(θ[3] * (1 - θ[3])))
        θsd[4] = min(θsd[4], 0.999999 * sqrt.(θ[4] * (1 - θ[4])))
        push!(fit_df, [sid, prevrew,
                       θ[1], θ[2], θ[3], θ[4], θsd[1], θsd[2], θsd[3], θsd[4]])
    end
    fit_df[:prevmag] = [id == 0 ? 0 : (id < 5 ? 9 - id : id) for id in fit_df[:previd]]

    # compute biases
    bootstrapn = 1000
    # draw n samples from Beta distribution with mean μ, standard deviation σ
    samplebeta(μ, σ, n) = σ <= 1e-30 ? fill(μ, n) :
        (x = (μ - abs2(μ)) / abs2(σ) - 1; rand(Beta(μ * x, (1 - μ) * x), n))
    biasid_df = DataFrame(previd = Int[], prevrew = Bool[],
                          biasμ = Float64[], biasσ = Float64[])
    for prevrew in [true, false], sid in stimids
        μ, σ, pl, pb, μsd, σsd, plsd, pbsd = Vector{Float64}(
            fit_df[(fit_df[:previd] .== sid) .& (fit_df[:prevrew] .== prevrew),
                   [:μ, :σ, :pl, :pb, :μsd, :σsd, :plsd, :pbsd]][1,:])
        # bias = intersection of psychometric curve with 4.5
        biases = Float64[psychpy1([μ + μsd * randn(),
                                   σ + σsd * randn(),
                                   samplebeta(pl, plsd, 1)[1],
                                   samplebeta(pb, pbsd, 1)[1]], 4.5) - 0.5
                         for n in 1:bootstrapn]
        push!(biasid_df, [sid, prevrew, mean(biases), √(var(biases))])
    end
    # biases over magnitudes -> flip for id < 5, and combine by variance
    stimmags = sort(unique(fit_df[fit_df[:prevmag] .!= 0, :prevmag]))
    biasmag_df = DataFrame(prevmag = Int[], prevrew = Bool[],
                           biasμ = Float64[], biasσ = Float64[])
    for prevrew in [true, false], smag in stimmags
        μs = [-biasid_df[(biasid_df[:prevrew] .== prevrew) .&
                         (biasid_df[:previd] .== 9 - smag), :biasμ][1],
              biasid_df[(biasid_df[:prevrew] .== prevrew) .&
                         (biasid_df[:previd] .== smag), :biasμ][1]]
        σs = [biasid_df[(biasid_df[:prevrew] .== prevrew) .&
                        (biasid_df[:previd] .== 9 - smag), :biasσ][1],
              biasid_df[(biasid_df[:prevrew] .== prevrew) .&
                         (biasid_df[:previd] .== smag), :biasσ][1]]
        ws = 1 ./ abs2.(σs)
        push!(biasmag_df, [smag, prevrew, sum(μs .* ws) / sum(ws), √(1 / sum(ws))])
    end

    return seq_df, fit_df, biasid_df, biasmag_df
end

###############################################################################
# functions plotting behavioral data
#
# These functions use the above analysis functions to perform the plotting.

# returns chron/psych plot over stimulus id
function plotchronpsychid(d_bystim, stimx, xlabel, ttitle, t=Theme())
    d_bystim[:stimx] = stimx[d_bystim[:stimid]]
    xmin, xmax = minimum(stimx), maximum(stimx)
    tp = copy(t)
    tp.default_color = colorant"black"
    # chronometric curve
    d_bystim[:tmin] = d_bystim[:tμ] - d_bystim[:tsem]
    d_bystim[:tmax] = d_bystim[:tμ] + d_bystim[:tsem]
    pchronid = plot(d_bystim, x="stimx", y="tμ", ymin="tmin", ymax="tmax",
        Geom.point, Geom.errorbar,
        Guide.xlabel(xlabel, orientation=:horizontal),
        Guide.ylabel("reaction time [s]", orientation=:vertical),
        Guide.title("Chronometric curve, $ttitle"),
        Scale.x_continuous(minvalue=xmin, maxvalue=xmax),
        Scale.y_continuous(minvalue=0.26, maxvalue=0.44),
        tp)
    # psychometric curve
    d_bystim[:prmin] = d_bystim[:prμ] - d_bystim[:prsem]
    d_bystim[:prmax] = d_bystim[:prμ] + d_bystim[:prsem]
    ppsychid = plot(d_bystim, x="stimx", y="prμ", ymin="prmin", ymax="prmax",
        Geom.point, Geom.errorbar,
        Guide.xlabel(xlabel, orientation=:horizontal),
        Guide.ylabel("left choices (fraction)", orientation=:vertical),
        Guide.title("Psychometric curve, $ttitle"),
        Scale.x_continuous(minvalue=xmin, maxvalue=xmax),
        Scale.y_continuous(minvalue=0.0, maxvalue=1.0),
        tp)
    return pchronid, ppsychid
end
function plotchronpsychidm(d_bystim, m_bystim, stimx, xlabel, ttitle, t=Theme())
    d_bystim[:stimx] = stimx[d_bystim[:stimid]]
    m_bystim[:stimx] = stimx[m_bystim[:stimid]]
    xmin, xmax = minimum(stimx), maximum(stimx)
    tp = copy(t)
    tp.default_color = colorant"black"
    # chronometric curve
    d_bystim[:tmin] = d_bystim[:tμ] - d_bystim[:tsem]
    d_bystim[:tmax] = d_bystim[:tμ] + d_bystim[:tsem]
    pchronid = plot(
        layer(d_bystim, x="stimx", y="tμ", ymin="tmin", ymax="tmax",
            Geom.point, Geom.errorbar, order=1),
        layer(m_bystim, x="stimx", y="tμ", Geom.line, order=0),
        Guide.xlabel(xlabel, orientation=:horizontal),
        Guide.ylabel("reaction time [s]", orientation=:vertical),
        Guide.title("Chronometric curve, $ttitle"),
        Scale.x_continuous(minvalue=xmin, maxvalue=xmax),
        Scale.y_continuous(minvalue=0.26, maxvalue=0.44),
        tp)
    # psychometric curve
    d_bystim[:prmin] = d_bystim[:prμ] - d_bystim[:prsem]
    d_bystim[:prmax] = d_bystim[:prμ] + d_bystim[:prsem]
    ppsychid = plot(
        layer(d_bystim, x="stimx", y="prμ", ymin="prmin", ymax="prmax",
            Geom.point, Geom.errorbar, order=1),
        layer(m_bystim, x="stimx", y="prμ", Geom.line, order=0),
        Guide.xlabel(xlabel, orientation=:horizontal),
        Guide.ylabel("left choices (fraction)", orientation=:vertical),
        Guide.title("Psychometric curve, $ttitle"),
        Scale.x_continuous(minvalue=xmin, maxvalue=xmax),
        Scale.y_continuous(minvalue=0.0, maxvalue=1.0),
        tp)
    return pchronid, ppsychid
end

# returns psych/chron over stimulus magnitude
function plotchronpsychmag(d_bymag, magctrst, xlabel, ttitle, t=Theme())
    d_bymag[:stimctrst] = magctrst[d_bymag[:stimmag]]
    xmin, xmax = minimum(magctrst), maximum(magctrst)
    tp = copy(t)
    tp.default_color = colorant"black"
    # chronometric curve
    d_bymag[:tmin] = d_bymag[:tμ] - d_bymag[:tsem]
    d_bymag[:tmax] = d_bymag[:tμ] + d_bymag[:tsem]
    pchronmag = plot(d_bymag, x="stimctrst", y="tμ", ymin="tmin", ymax="tmax",
        Geom.point, Geom.errorbar,
        Guide.xlabel(xlabel, orientation=:horizontal),
        Guide.ylabel("reaction time [s]", orientation=:vertical),
        Guide.title("Chronometric curve, $ttitle"),
        Scale.x_continuous(minvalue=xmin, maxvalue=xmax),
        Scale.y_continuous(minvalue=0.26, maxvalue=0.44),
        tp)
    # psychometric curve
    d_bymag[:pcorrmin] = d_bymag[:pcorrμ] - d_bymag[:pcorrsem]
    d_bymag[:pcorrmax] = d_bymag[:pcorrμ] + d_bymag[:pcorrsem]
    ppsychmag = plot(d_bymag, x="stimctrst", y="pcorrμ", ymin="pcorrmin", ymax="pcorrmax",
        Geom.point, Geom.errorbar,
        Guide.xlabel(xlabel, orientation=:horizontal),
        Guide.ylabel("correct choices (fraction)", orientation=:vertical),
        Guide.title("Psychometric curve, $ttitle"),
        Scale.x_continuous(minvalue=xmin, maxvalue=xmax),
        Scale.y_continuous(minvalue=0.5, maxvalue=1.0),
        tp)
    return pchronmag, ppsychmag
end
function plotchronpsychmagm(d_bymag, m_bymag, magctrst, xlabel, ttitle, t=Theme())
    d_bymag[:stimctrst] = magctrst[d_bymag[:stimmag]]
    m_bymag[:stimctrst] = magctrst[m_bymag[:stimmag]]
    xmin, xmax = minimum(magctrst), maximum(magctrst)
    tp = copy(t)
    tp.default_color = colorant"black"
    # chronometric curve
    d_bymag[:tmin] = d_bymag[:tμ] - d_bymag[:tsem]
    d_bymag[:tmax] = d_bymag[:tμ] + d_bymag[:tsem]
    pchronmag = plot(
        layer(d_bymag, x="stimctrst", y="tμ", ymin="tmin", ymax="tmax",
            Geom.point, Geom.errorbar, order=1),
        layer(m_bymag, x="stimctrst", y="tμ", Geom.line, order=0),
        Guide.xlabel(xlabel, orientation=:horizontal),
        Guide.ylabel("reaction time [s]", orientation=:vertical),
        Guide.title("Chronometric curve, $ttitle"),
        Scale.x_continuous(minvalue=xmin, maxvalue=xmax),
        Scale.y_continuous(minvalue=0.26, maxvalue=0.44),
        tp)
    # psychometric curve
    d_bymag[:pcorrmin] = d_bymag[:pcorrμ] - d_bymag[:pcorrsem]
    d_bymag[:pcorrmax] = d_bymag[:pcorrμ] + d_bymag[:pcorrsem]
    ppsychmag = plot(
        layer(d_bymag, x="stimctrst", y="pcorrμ", ymin="pcorrmin", ymax="pcorrmax",
            Geom.point, Geom.errorbar, order=1),
        layer(m_bymag, x="stimctrst", y="pcorrμ", Geom.line, order=0),
        Guide.xlabel(xlabel, orientation=:horizontal),
        Guide.ylabel("correct choices (fraction)", orientation=:vertical),
        Guide.title("Psychometric curve, $ttitle"),
        Scale.x_continuous(minvalue=xmin, maxvalue=xmax),
        Scale.y_continuous(minvalue=0.5, maxvalue=1.0),
        tp)
    return pchronmag, ppsychmag
end

# returns array of conditional psychometric curves/fits
function plotcondpsych(d_seq, d_seqfit, stimx, xlabel, ttitle, t=Theme())
    # ignore the actual x stimulus distribution, plot on [0..1] scale
    stimx = range(0, stop=1, length=8)
    d_seq[:stimx] = stimx[d_seq[:stimid]]
    stimidcol(stimid) = RGB(0.0, 1 - 0.125stimid, 0.125stimid)
    pcondpsych = Vector{Any}(undef, 4)
    xcont = range(0.0, stop=1.0, length=100)
    globalθ = Vector{Float64}(
        d_seqfit[(d_seqfit[:prevrew] .== true) .& (d_seqfit[:previd] .== 0),
                 [:μ, :σ, :pl, :pb]][1,:])
    for stimmag = 1:4
        id1, id2 = stimmag, 9-stimmag
        pr1 = [d_seq[(d_seq[:prevrew] .== true) .& (d_seq[:previd] .== id1) .&
                     (d_seq[:stimid] .== id), :prμ][1] for id = 1:8]
        pr1sem = [d_seq[(d_seq[:prevrew] .== true) .& (d_seq[:previd] .== id1) .&
                        (d_seq[:stimid] .== id), :prsem][1] for id = 1:8]
        θ1 = Vector{Float64}(
            d_seqfit[(d_seqfit[:prevrew] .== true) .& (d_seqfit[:previd] .== id1),
                     [:μ, :σ, :pl, :pb]][1,:])
        pr2 = [d_seq[(d_seq[:prevrew] .== true) .& (d_seq[:previd] .== id2) .&
                     (d_seq[:stimid] .== id), :prμ][1] for id = 1:8]
        pr2sem = [d_seq[(d_seq[:prevrew] .== true) .& (d_seq[:previd] .== id2) .&
                        (d_seq[:stimid] .== id), :prsem][1] for id = 1:8]
        θ2 = Vector{Float64}(
            d_seqfit[(d_seqfit[:prevrew] .== true) .& (d_seqfit[:previd] .== id2),
                     [:μ, :σ, :pl, :pb]][1,:])

        pcondpsych[stimmag] = plot(
            Scale.x_continuous(minvalue=0.0, maxvalue=1.0),
            Scale.y_continuous(minvalue=0.0, maxvalue=1.0),
            layer(x = xcont, y = psychpy1(globalθ, 1 .+ 7xcont),
                  Geom.line, Theme(default_color=colorant"black"), order=0),
            layer(x = xcont, y = psychpy1(θ1, 1 .+ 7xcont),
                  Geom.line, Theme(default_color=stimidcol(id1)), order=1),
            layer(x = xcont, y = psychpy1(θ2, 1 .+ 7xcont),
                  Geom.line, Theme(default_color=stimidcol(id2)), order=2),
            layer(x = stimx, y = pr1, ymin = max.(0.0, pr1-pr1sem), ymax = min.(1.0, pr1+pr1sem),
                  Geom.point, Geom.errorbar, Theme(default_color=stimidcol(id1)), order=3),
            layer(x = stimx, y = pr2, ymin = max.(0.0, pr2-pr2sem), ymax = min.(1.0, pr2+pr2sem),
                  Geom.point, Geom.errorbar, Theme(default_color=stimidcol(id2)), order=4),
            Guide.xlabel(xlabel, orientation=:horizontal),
            Guide.ylabel("left choices (fraction)", orientation=:vertical),
            t)
    end

    return pcondpsych
end
# returns array of conditional psychometric curves/fits
function plotcondpsychm(d_seq, m_seq, m_seqfit, stimx, xlabel, ttitle, t=Theme())
    # ignore the actual x stimulus distribution, plot on [0..1] scale
    stimx = range(0, stop=1, length=8)
    d_seq[:stimx] = stimx[d_seq[:stimid]]
    stimidcol(stimid) = RGB(0.0, 1 - 0.125stimid, 0.125stimid)
    fitidcol(stimid) = RGB(0.5, 1 - 0.0625stimid, 0.5 + 0.0625stimid)
    pcondpsych = Vector{Any}(undef, 4)
    xcont = range(0.0, stop=1.0, length=100)
    for stimmag = 1:4
        id1, id2 = stimmag, 9-stimmag
        pr1 = [d_seq[(d_seq[:prevrew] .== true) .& (d_seq[:previd] .== id1) .&
                     (d_seq[:stimid] .== id), :prμ][1] for id = 1:8]
        pr1sem = [d_seq[(d_seq[:prevrew] .== true) .& (d_seq[:previd] .== id1) .&
                        (d_seq[:stimid] .== id), :prsem][1] for id = 1:8]
        mpr1 = [m_seq[(m_seq[:prevrew] .== true) .& (m_seq[:previd] .== id1) .&
                      (m_seq[:stimid] .== id), :prμ][1] for id = 1:8]
        θ1 = Vector{Float64}(
            m_seqfit[(m_seqfit[:prevrew] .== true) .& (m_seqfit[:previd] .== id1),
                     [:μ, :σ, :pl, :pb]][1,:])
        pr2 = [d_seq[(d_seq[:prevrew] .== true) .& (d_seq[:previd] .== id2) .&
                     (d_seq[:stimid] .== id), :prμ][1] for id = 1:8]
        pr2sem = [d_seq[(d_seq[:prevrew] .== true) .& (d_seq[:previd] .== id2) .&
                        (d_seq[:stimid] .== id), :prsem][1] for id = 1:8]
        mpr2 = [m_seq[(m_seq[:prevrew] .== true) .& (m_seq[:previd] .== id2) .&
                      (m_seq[:stimid] .== id), :prμ][1] for id = 1:8]
        θ2 = Vector{Float64}(
            m_seqfit[(m_seqfit[:prevrew] .== true) .& (m_seqfit[:previd] .== id2),
                     [:μ, :σ, :pl, :pb]][1,:])

        pcondpsych[stimmag] = plot(
            Scale.x_continuous(minvalue=0.0, maxvalue=1.0),
            Scale.y_continuous(minvalue=0.0, maxvalue=1.0),
            layer(x = xcont, y = psychpy1(θ1, 1 .+ 7xcont),
                  Geom.line, Theme(default_color=fitidcol(id1)), order=0),
            layer(x = xcont, y = psychpy1(θ2, 1 .+ 7xcont),
                  Geom.line, Theme(default_color=fitidcol(id2)), order=1),
            layer(x = stimx, y = mpr1, Geom.line,
                  Theme(default_color=stimidcol(id1)), order=2),
            layer(x = stimx, y = mpr2, Geom.line,
                  Theme(default_color=stimidcol(id2)), order=3),
            layer(x = stimx, y = pr1, ymin = max.(0.0, pr1-pr1sem), ymax = min.(1.0, pr1+pr1sem),
                  Geom.point, Geom.errorbar, Theme(default_color=stimidcol(id1)), order=4),
            layer(x = stimx, y = pr2, ymin = max.(0.0, pr2-pr2sem), ymax = min.(1.0, pr2+pr2sem),
                  Geom.point, Geom.errorbar, Theme(default_color=stimidcol(id2)), order=5),
            Guide.xlabel(xlabel, orientation=:horizontal),
            Guide.ylabel("left choices (fraction)", orientation=:vertical),
            t)
    end

    return pcondpsych
end

# returns bias plot over simulus id
function plotbiasid(d_biasid, stimx, xlabel, ttitle, t=Theme())
    d_biasid[:prevx] = stimx[d_biasid[:previd]]
    d_biasid[:biasmin] = d_biasid[:biasμ] - d_biasid[:biasσ]
    d_biasid[:biasmax] = d_biasid[:biasμ] + d_biasid[:biasσ]
    tp = copy(t)
    tp.default_color = colorant"red"
    pbiasid = plot(d_biasid[d_biasid[:prevrew],:],
        x="prevx", y="biasμ", ymin="biasmin", ymax="biasmax",
        Geom.point, Geom.errorbar,
        Guide.xlabel(xlabel, orientation=:horizontal),
        Guide.ylabel("choice bias", orientation=:vertical),
        Guide.title("Choice biases, $ttitle"),
        Scale.x_continuous(minvalue=minimum(stimx), maxvalue=maximum(stimx)),
        Scale.y_continuous(minvalue=-0.5, maxvalue=0.5),
        tp)
    return pbiasid
end
function plotbiasidm(d_biasid, m_biasid, stimx, xlabel, ttitle, t=Theme())
    d_biasid[:prevx] = stimx[d_biasid[:previd]]
    m_biasid[:prevx] = stimx[m_biasid[:previd]]
    d_biasid[:biasmin] = d_biasid[:biasμ] - d_biasid[:biasσ]
    d_biasid[:biasmax] = d_biasid[:biasμ] + d_biasid[:biasσ]
    tp = copy(t)
    tp.default_color = colorant"red"
    pbiasid = plot(
        layer(d_biasid[d_biasid[:prevrew],:],
            x="prevx", y="biasμ", ymin="biasmin", ymax="biasmax",
            Geom.point, Geom.errorbar),
        layer(m_biasid[m_biasid[:prevrew],:],
            x="prevx", y="biasμ", Geom.line),
        Guide.xlabel(xlabel, orientation=:horizontal),
        Guide.ylabel("choice bias", orientation=:vertical),
        Guide.title("Choice biases, $ttitle"),
        Scale.x_continuous(minvalue=minimum(stimx), maxvalue=maximum(stimx)),
        Scale.y_continuous(minvalue=-0.5, maxvalue=0.5),
        tp)
    return pbiasid
end

# returns bias plot over stimulus magnitude
function plotbiasmag(d_biasmag, magctrst, xlabel, ttitle, t=Theme())
    d_biasmag[:prevctrst] = magctrst[d_biasmag[:prevmag]]
    d_biasmag[:biasmin] = d_biasmag[:biasμ] - d_biasmag[:biasσ]
    d_biasmag[:biasmax] = d_biasmag[:biasμ] + d_biasmag[:biasσ]
    tp = copy(t)
    tp.default_color = colorant"red"
    pbiasmag = plot(d_biasmag[d_biasmag[:prevrew],:],
        x="prevctrst", y="biasμ", ymin="biasmin", ymax="biasmax",
        Geom.point, Geom.errorbar,
        Guide.xlabel(xlabel, orientation=:horizontal),
        Guide.ylabel("choice bias", orientation=:vertical),
        Guide.title("Choice biases, $ttitle"),
        Scale.x_continuous(minvalue=minimum(magctrst), maxvalue=maximum(magctrst)),
        Scale.y_continuous(minvalue=-0.05, maxvalue=0.25),
        tp)
    return pbiasmag
end
function plotbiasmagm(d_biasmag, m_biasmag, magctrst, xlabel, ttitle, t=Theme())
    d_biasmag[:prevctrst] = magctrst[d_biasmag[:prevmag]]
    m_biasmag[:prevctrst] = magctrst[m_biasmag[:prevmag]]
    d_biasmag[:biasmin] = d_biasmag[:biasμ] - d_biasmag[:biasσ]
    d_biasmag[:biasmax] = d_biasmag[:biasμ] + d_biasmag[:biasσ]
    tp = copy(t)
    tp.default_color = colorant"red"
    pbiasmag = plot(
        layer(d_biasmag[d_biasmag[:prevrew],:],
            x="prevctrst", y="biasμ", ymin="biasmin", ymax="biasmax",
            Geom.point, Geom.errorbar),
        layer(m_biasmag[m_biasmag[:prevrew],:],
            x="prevctrst", y="biasμ", Geom.line),
        Guide.xlabel(xlabel, orientation=:horizontal),
        Guide.ylabel("choice bias", orientation=:vertical),
        Guide.title("Choice biases, $ttitle"),
        Scale.x_continuous(minvalue=minimum(magctrst), maxvalue=maximum(magctrst)),
        Scale.y_continuous(minvalue=-0.05, maxvalue=0.25),
        tp)
    return pbiasmag
end


nelements(node::Leaf) = length(node)
nelements(node::Branch) = mapreduce(nelements, +, children(node), init=0)

# check the R-tree node (branch or leaf):
# * the parent
# * the level
# * the number of attached elements
# * element id uniqueness
# * the MBR (should fit the MBR of children)
function check(node::Node{T,N,V}, tree::RTree{T,N,V},
               nnodes_perlevel::AbstractVector{Int}, ids::Union{Set, Nothing}) where {T,N,V}
    nnodes_perlevel[level(node)] += 1
    for (i, child) in enumerate(children(node))
        if node isa Branch
            child.parent === node ||
                throw(SpatialIndexException("Node (lev=$(level(node)) len=$(length(node))): child #$i references different parent"))
            level(child) + 1 == level(node) ||
                throw(SpatialIndexException("Node (lev=$(level(node)) len=$(length(node))): child #$i level ($(level(child))) doesn't match parent's ($(level(node)))"))
            check(child, tree, nnodes_perlevel, ids)
        elseif idtrait(typeof(child)) !== HasNoID
            cid = id(child)
            (cid ∈ ids) && throw(SpatialIndexException("Duplicate data id=$cid"))
            push!(ids, cid)
        end
    end
    br = !isempty(node) ? mapreduce(mbr, combine, node.children) : empty(Rect{T,N})
    if isempty(node)
        isequal(br, mbr(node)) ||
            throw(SpatialIndexException("Node (lev=$(level(node)) len=$(length(node))) MBR ($(mbr(node))) should be empty"))
    elseif tree.tight_mbrs && br != mbr(node) || !in(br, mbr(node))
        throw(SpatialIndexException("Node (lev=$(level(node)) len=$(length(node))) MBR ($(mbr(node))) doesn't match its children MBR ($(br))"))
    end
    return true
end

# check the tree:
# * the root
# * the height
# * the number of elements
# * the number of nodes per level
# * element id uniqueness
# * MBRs of the nodes
function check(tree::RTree)
    tree.root.parent === nothing ||
        throw(SpatialIndexException("RTree: root has a parent"))
    level(tree.root) == height(tree) ||
        throw(SpatialIndexException("RTree: root level ($(level(tree.root))) doesn't match tree height ($(height(tree)))"))
    length(tree.nnodes_perlevel) == height(tree) ||
        throw(SpatialIndexException("RTree: nnodes_perlevel length ($(length(tree.nnodes_perlevel))) doesn't match tree height ($(height(tree)))"))
    ids = idtype(tree) !== Union{} ? Set{idtype(tree)}() : nothing
    nnodes_perlevel = zeros(Int, height(tree))
    check(tree.root, tree, nnodes_perlevel, ids)
    nelms = nelements(tree.root)
    length(tree) == nelms ||
        throw(SpatialIndexException("RTree: actual ($(nelms)) and reported ($(length(tree))) number of elements do not match"))
    ids === nothing || length(tree) == length(ids) ||
        throw(SpatialIndexException("RTree: the number of ids ($(length(ids))) doesn't match the reported number of elements ($(length(tree)))"))
    tree.nnodes_perlevel == nnodes_perlevel ||
        throw(SpatialIndexException("RTree: actual ($(nnodes_perlevel)) and reported ($(tree.nnodes_perlevel)) number of nodes per level do not match"))
    return true
end


# PART 5 – ASSESSMENT OF LOCAL METAL LOSS

# Determine Asessment Applicability
#Determine the assessment applicability
# @doc DesignCodeCriteria
# @doc MaterialToughness
# @doc CyclicService
# @doc Part5ComponentType
print("Begin -- Assessment Applicability and Component Type Checks\n")
creep_range = CreepRangeTemperature("Carbon Steel (UTS ≤ 414MPa (60 ksi))"; design_temperature=100.0, units="lbs-in-psi")
design = DesignCodeCriteria("ASME B31.3 Piping Code")
toughness = MaterialToughness("Certain")
cyclic = CyclicService(100, "Meets Part 14")
x = Part5ComponentType("Straight Section of Piping, Elbow or Bend - No Structural Attachments", vessel_orientation="horizontal", material="Carbon and Low Alloy Steels", D=24.75,Lss=120.0,H=0.0, NPS=10.0, design_temperature=100.0, units="lbs-in-psi")
part5_applicability = Part5AsessmentApplicability(x,design,toughness,cyclic,creep_range)

# For all assessments - determine the inspection data grid
M1 = [0.300, 0.300, 0.300, 0.300, 0.300, 0.300, 0.300, 0.300, 0.300, 0.300, 0.300, 0.300, 0.300, 0.300, 0.300, 0.300, 0.300, 0.300, 0.300, 0.300, 0.300, 0.300]
M2 = [0.300, 0.300, 0.300, 0.300, 0.300, 0.300, 0.300, 0.100, 0.220, 0.280, 0.250, 0.240, 0.300, 0.300, 0.300, 0.300, 0.300, 0.300, 0.300, 0.300, 0.300, 0.300]
M3 = [0.300, 0.300, 0.300, 0.300, 0.300, 0.215, 0.255, 0.215, 0.145, 0.275, 0.170, 0.240, 0.250, 0.250, 0.280, 0.290, 0.300, 0.300, 0.300, 0.300, 0.300, 0.300]
M4 = [0.300, 0.300, 0.300, 0.300, 0.300, 0.170, 0.270, 0.190, 0.190, 0.285, 0.250, 0.225, 0.275, 0.300, 0.300, 0.300, 0.300, 0.300, 0.300, 0.300, 0.300, 0.300]
M5 = [0.300, 0.300, 0.300, 0.300, 0.300, 0.300, 0.300, 0.300, 0.300, 0.300, 0.300, 0.300, 0.300, 0.300, 0.300, 0.300, 0.300, 0.300, 0.300, 0.300, 0.300, 0.300]
M6 = [0.300, 0.300, 0.300, 0.300, 0.300, 0.300, 0.300, 0.300, 0.300, 0.300, 0.300, 0.300, 0.300, 0.300, 0.300, 0.300, 0.300, 0.300, 0.300, 0.300, 0.300, 0.300]
CTPGrid = hcat(M6,M5,M4,M3,M2,M1) # build in descending order
CTPGrid = rotl90(CTPGrid) # rotate to correct orientation

# Input variables
    annex2c_tmin_category = "Straight Pipes Subject To Internal Pressure" # ["Cylindrical Shell","Spherical Shell","Hemispherical Head","Elliptical Head","Torispherical Head","Conical Shell","Toriconical Head","Conical Transition","Nozzles Connections in Shells",
    # "Junction Reinforcement Requirements at Conical Transitions","Tubesheets","Flat head to cylinder connections","Bolted Flanges","Straight Pipes Subject To Internal Pressure","Boiler Tubes","Pipe Bends Subject To Internal Pressure",
    # "MAWP for External Pressure","Branch Connections","API 650 Storage Tanks"]
    equipment_group = "piping" # "vessel", "tank"
    flaw_location = "internal" # "External","Internal"
    pipe_code = "B31.8" # "B31.3", "B31.8"
    FCA_string = "internal"
    metal_loss_categorization = "LTA" # "LTA" or "Groove-Like Flaw"
    units = "lbs-in-psi" # "lbs-in-psi" or "nmm-mm-mpa"
    remaining_life = true
    tmm_forcing = true
    tmm = .949
    tnom = 1.0 # nominal or furnished thickness of the component adjusted for mill undertolerance as applicable.
    trd = 1.0 # uniform thickness away from the local metal loss location established by thickness measurements at the time of the assessment.
    FCAml = 0.1 # Future Corrosion Allowance applied to the region of metal loss.
    FCA = 0.0 # Future Corrosion Allowance applied to the region away from the metal loss (see Annex 2C, paragraph 2C.2.8).
    LOSS = 0.0 #the amount of uniform metal loss away from the local metal loss location at the time of the assessment.
    Do = 8.625 # Outside Diameter
    D = Do - 2*(tnom) # Inside Dia.
    P = 8139.13 # internal design pressure.
    S = 65000.0 # code allowable stress note if using B31.8 its SMYS - check design code criteria
    E = 1.0 # weld joint efficiency or quality factor from the original construction code, if unknown use 0.7.
    F = 0.6 # B31.8 design factor
    T = 1.0 # temperature factor B31.8
    MA = 0.0 # mechanical allowances (thread or groove depth); for threaded components, the nominal thread depth (dimension h of ASME B.1.20.1) shall apply.
    Yb31 = 0.4 # coefficient from ASME B31 Piping codes used for determining the pipe wall thickness, the coefficient can be determined from the following table that is valid for tmin < Do / 6 Annex 2C .
    t = trd - LOSS - FCA # thickness of the shell or pipe adjusted for mill tolerance, LOSS and FCA , or cylinder thickness at a conical transition for a junction reinforcement calculation adjusted for mill tolerance, LOSS and FCA , as applicable.
    t = 1.0
    tsl = 0.0 # supplemental thickness for mechanical loads other than pressure that result in longitudinal stress; this thickness is usually obtained from the results of a weight case in a stress analysis of the piping system (see paragraph 2C.2.7).
    spacings = 0.5 # spacings determine by visual inspection to adequately ccategorizse the corrosion -----------+ may add to CTP_Grid function for plotting purposes
    # Flaw dimensions
    s = 20.0 # longitudinal extent or length of the region of local metal loss based on future corroded thickness,
    c = 6.7  # circumferential extent or length of the region of local metal loss (see Figure 5.2 and Figure 5.10), based on future corroded thickness, tc .
    Ec = 1.0 # circumferential weld joint efficiency. note if damage on weld see # 2C.2.5 Treatment of Weld and Riveted Joint Efficiency, and Ligament Efficiency
    El = 1.0 # longitudinal weld joint efficiency. note if damage on weld see # 2C.2.5 Treatment of Weld and Riveted Joint Efficiency, and Ligament Efficiency
    RSFa = 0.9 # remaining strength factor - consult API 579 is go lower than 0.9

    # For all assessments determine far enough from structural discontinuity
    # Flaw-To-Major Structural Discontinuity Spacing
    L1msd = [12.0] # distance to the nearest major structural discontinuity.
    L2msd = [12.0] # distance to the nearest major structural discontinuity.
    L3msd = [12.0] # distance to the nearest major structural discontinuity.
    L4msd = [12.0] # distance to the nearest major structural discontinuity.
    L5msd = [12.0] # distance to the nearest major structural discontinuity.
    Lmsd = minimum([L1msd,L2msd,L3msd,L4msd,L5msd])[1]
    if (Lmsd >= (1.8*(sqrt(D*(trd - LOSS - FCA)))))
        print("Satisfied - Flaw is located far enough from structural discontinuity\n")
        lmsd_satisfied = 1
    else
        print("Not satisfied - Flaw is too close to a structural discontinuity - Conduct a level 3 assessment\n")
        lmsd_satisfied = 0
    end

    # Groove Like Flaw dimensions
    gl = .05 # length of the Groove-Like Flaw based on future corroded condition.
    gw = .4 # width of the Groove-Like Flaw based on future corroded condition.
    gr = 0.1 # radius at the base of a Groove-Like Flaw based on future corroded condition.
    β = 40.0 # see (Figure 5.4) :: orientation of the groove-like flaw with respect to the longitudinal axis or a parameter to compute an effective fracture toughness for a groove being evaluated as a crack-like flaw, as applicable.

# Perform level 1 assessment
if (part5_applicability[1] == 1 && lmsd_satisfied == 1) # begin level 1 assessment
    #let part_5_lta_output = Array{Any,2},
    part_5_lta_output = Part5LTALevel1(CTPGrid; remaining_life=remaining_life,tmm_forcing=tmm_forcing, tmm=tmm, annex2c_tmin_category=annex2c_tmin_category, pipe_code=pipe_code,equipment_group=equipment_group, flaw_location=flaw_location, FCA_string=FCA_string, metal_loss_categorization=metal_loss_categorization, units=units, Lmsd=Lmsd,tnom=tnom,
        trd=trd, FCA=FCA, FCAml=FCAml, LOSS=LOSS, Do=Do, D=D, P=P, S=S, E=E, MA=MA, Yb31=Yb31, t=t,tsl=tsl, F=F, T=T, spacings=spacings, s=s, c=c, El=El, Ec=Ec, RSFa=RSFa, gl=gl, gw=gw, gr=gr,β=β)
    #end # let end
    part_5_lta_output

elseif (part5_applicability[1] == 0 && lmsd_satisfied == 0)
    print("Level 1 Criteria Not Met - Perform Level 2 or 3 as applicable")
elseif (part5_applicability[1] == 1 && lmsd_satisfied == 0)
    print("Level 1 Criteria Not Met - Perform Level 2 or 3 as applicable")
elseif (part5_applicability[1] == 0 && lmsd_satisfied == 1)
    print("Level 1 Criteria Not Met - Perform Level 2 or 3 as applicable")
end

# Perform level 2 assessment
if (part5_applicability[2] == 1 && lmsd_satisfied == 1) # begin level 2 assessment
    #let part_5_lta_output = Array{Any,2},
    part_5_lta_output = Part5LTALevel2(CTPGrid; annex2c_tmin_category=annex2c_tmin_category, pipe_code=pipe_code, equipment_group=equipment_group, flaw_location=flaw_location, FCA_string=FCA_string, metal_loss_categorization=metal_loss_categorization, units=units, Lmsd=Lmsd, tnom=tnom,
        trd=trd, FCA=FCA, FCAml=FCAml, LOSS=LOSS, Do=Do, D=D, P=P, S=S, E=E, MA=MA, Yb31=Yb31, tsl=tsl, t=t,F=F, T=T, spacings=spacings, s=s, c=c, El=El, Ec=Ec, RSFa=RSFa, gl=gl, gw=gw, gr=gr,β=β)
    #end # let end
elseif (part5_applicability[1] == 0 && lmsd_satisfied == 0)
    print("Level 1 Criteria Not Met - Perform Level 2 or 3 as applicable")
elseif (part5_applicability[1] == 1 && lmsd_satisfied == 0)
    print("Level 1 Criteria Not Met - Perform Level 2 or 3 as applicable")
elseif (part5_applicability[1] == 0 && lmsd_satisfied == 1)
    print("Level 1 Criteria Not Met - Perform Level 2 or 3 as applicable")
end


isautodifferentiable(alg::Union{OrdinaryDiffEqAlgorithm,DAEAlgorithm}) = true

DiffEqBase.isdiscrete(alg::FunctionMap) = true

isfsal(alg::Union{OrdinaryDiffEqAlgorithm,DAEAlgorithm}) = true
isfsal(tab::DiffEqBase.ExplicitRKTableau{MType,VType,fsal}) where {MType,VType,fsal} = fsal
# isfsal(alg::CompositeAlgorithm) = isfsal(alg.algs[alg.current])
isfsal(alg::FunctionMap) = false
isfsal(alg::Rodas4) = false
isfsal(alg::Rodas42) = false
isfsal(alg::Rodas4P) = false
isfsal(alg::Vern7) = false
isfsal(alg::Vern8) = false
isfsal(alg::Vern9) = false
# Pseudo Non-FSAL
isfsal(alg::ORK256) = false
isfsal(alg::CarpenterKennedy2N54) = false
isfsal(alg::SHLDDRK64) = false
isfsal(alg::DGLDDRK73_C) = false
isfsal(alg::DGLDDRK84_C) = false
isfsal(alg::DGLDDRK84_F) = false
isfsal(alg::NDBLSRK124) = false
isfsal(alg::NDBLSRK134) = false
isfsal(alg::NDBLSRK144) = false
isfsal(alg::PDIRK44) = false
isfsal(alg::DImplicitEuler) = false
isfsal(alg::RKO65) = false
isfsal(alg::FRK65) = true
#isfsal(alg::RKM) = false

isfsal(alg::SSPRK22) = false
isfsal(alg::SSPRK33) = false
isfsal(alg::SSPRK53) = false
isfsal(alg::SSPRK53_2N1) = false
isfsal(alg::SSPRK53_2N2) = false
isfsal(alg::SSPRK53_H) = false
isfsal(alg::SSPRK63) = false
isfsal(alg::SSPRK73) = false
isfsal(alg::SSPRK83) = false
isfsal(alg::SSPRK43) = false
isfsal(alg::SSPRK432) = false
isfsal(alg::SSPRK932) = false
isfsal(alg::SSPRK54) = false
isfsal(alg::SSPRK104) = false

get_current_isfsal(alg, cache) = isfsal(alg)
get_current_isfsal(alg::CompositeAlgorithm, cache) = isfsal(alg.algs[cache.current])

issplit(alg::Union{OrdinaryDiffEqAlgorithm,DAEAlgorithm}) = false
issplit(alg::SplitAlgorithms) = true

@generated function _composite_beta1_default(algs::T, current, ::Type{QT}, beta2) where {T <: Tuple, QT}
  expr = Expr(:block)
  for i in 1:length(T.types)
    push!(expr.args, quote
      if current == $i
        return QT(beta1_default(algs[$i], beta2))
      end
    end)
  end
  return expr
end
@generated function _composite_beta2_default(algs::T, current, ::Type{QT}) where {T <: Tuple, QT}
  expr = Expr(:block)
  for i in 1:length(T.types)
    push!(expr.args, quote
      if current == $i
        return QT(beta2_default(algs[$i]))
      end
    end)
  end
  return expr
end

fsal_typeof(alg::Union{OrdinaryDiffEqAlgorithm,DAEAlgorithm},rate_prototype) = typeof(rate_prototype)
fsal_typeof(alg::ETD2,rate_prototype) = ETD2Fsal{typeof(rate_prototype)}
function fsal_typeof(alg::CompositeAlgorithm,rate_prototype)
  fsal = map(x->fsal_typeof(x,rate_prototype), alg.algs)
  @assert length(unique(fsal)) == 1 "`fsal_typeof` must be consistent"
  return fsal[1]
end

isimplicit(alg::Union{OrdinaryDiffEqAlgorithm,DAEAlgorithm}) = false
isimplicit(alg::OrdinaryDiffEqAdaptiveImplicitAlgorithm) = true
isimplicit(alg::OrdinaryDiffEqImplicitAlgorithm) = true
isimplicit(alg::CompositeAlgorithm) = any(isimplicit.(alg.algs))

isdtchangeable(alg::Union{OrdinaryDiffEqAlgorithm,DAEAlgorithm}) = true
isdtchangeable(alg::CompositeAlgorithm) = all(isdtchangeable.(alg.algs))
isdtchangeable(alg::Union{LawsonEuler,NorsettEuler,LieEuler,MagnusGauss4,CayleyEuler,ETDRK2,ETDRK3,ETDRK4,HochOst4,ETD2}) = false # due to caching

ismultistep(alg::Union{OrdinaryDiffEqAlgorithm,DAEAlgorithm}) = false
ismultistep(alg::CompositeAlgorithm) = any(ismultistep.(alg.algs))
ismultistep(alg::ETD2) = true

isadaptive(alg::Union{OrdinaryDiffEqAlgorithm,DAEAlgorithm}) = false
isadaptive(alg::OrdinaryDiffEqAdaptiveAlgorithm) = true
isadaptive(alg::OrdinaryDiffEqCompositeAlgorithm) = all(isadaptive.(alg.algs))
isadaptive(alg::DImplicitEuler) = true
isadaptive(alg::DABDF2) = true

isautoswitch(alg) = false
isautoswitch(alg::CompositeAlgorithm) = alg.choice_function isa AutoSwitch

qmin_default(alg::Union{OrdinaryDiffEqAlgorithm,DAEAlgorithm}) = isadaptive(alg) ? 1//5 : 0
qmin_default(alg::CompositeAlgorithm) = maximum(qmin_default.(alg.algs))
qmin_default(alg::DP8) = 1//3

qmax_default(alg::Union{OrdinaryDiffEqAlgorithm,DAEAlgorithm}) = 10
qmax_default(alg::CompositeAlgorithm) = minimum(qmax_default.(alg.algs))
qmax_default(alg::DP8) = 6
qmax_default(alg::Union{RadauIIA3,RadauIIA5}) = 8

get_chunksize(alg::OrdinaryDiffEqAlgorithm) = error("This algorithm does not have a chunk size defined.")
get_chunksize(alg::OrdinaryDiffEqAdaptiveImplicitAlgorithm{CS,AD}) where {CS,AD} = Val(CS)
get_chunksize(alg::OrdinaryDiffEqImplicitAlgorithm{CS,AD}) where {CS,AD} = Val(CS)
get_chunksize(alg::DAEAlgorithm{CS,AD}) where {CS,AD} = Val(CS)
get_chunksize(alg::ExponentialAlgorithm) = Val(alg.chunksize)
# get_chunksize(alg::CompositeAlgorithm) = get_chunksize(alg.algs[alg.current_alg])

alg_autodiff(alg::OrdinaryDiffEqAlgorithm) = error("This algorithm does not have an autodifferentiation option defined.")
alg_autodiff(alg::OrdinaryDiffEqAdaptiveImplicitAlgorithm{CS,AD}) where {CS,AD} = AD
alg_autodiff(alg::DAEAlgorithm{CS,AD}) where {CS,AD} = AD
alg_autodiff(alg::OrdinaryDiffEqImplicitAlgorithm{CS,AD}) where {CS,AD} = AD
alg_autodiff(alg::ExponentialAlgorithm) = alg.autodiff
# alg_autodiff(alg::CompositeAlgorithm) = alg_autodiff(alg.algs[alg.current_alg])
get_current_alg_autodiff(alg, cache) = alg_autodiff(alg)
get_current_alg_autodiff(alg::CompositeAlgorithm, cache) = alg_autodiff(alg.algs[cache.current])

alg_extrapolates(alg::Union{OrdinaryDiffEqAlgorithm,DAEAlgorithm}) = false
alg_extrapolates(alg::CompositeAlgorithm) = any(alg_extrapolates.(alg.algs))
alg_extrapolates(alg::ImplicitEuler) = true
alg_extrapolates(alg::DImplicitEuler) = true
alg_extrapolates(alg::DABDF2) = true
alg_extrapolates(alg::Trapezoid) = true
alg_extrapolates(alg::ImplicitMidpoint) = true
alg_extrapolates(alg::TRBDF2) = true
alg_extrapolates(alg::SSPSDIRK2) = true
alg_extrapolates(alg::SDIRK2) = true
alg_extrapolates(alg::SDIRK22) = true
alg_extrapolates(alg::Kvaerno3) = true
alg_extrapolates(alg::Kvaerno4) = true
alg_extrapolates(alg::Kvaerno5) = true
alg_extrapolates(alg::ESDIRK54I8L2SA) = true
alg_extrapolates(alg::KenCarp3) = true
alg_extrapolates(alg::KenCarp4) = true
alg_extrapolates(alg::KenCarp47) = true
alg_extrapolates(alg::KenCarp5) = true
alg_extrapolates(alg::KenCarp58) = true
alg_extrapolates(alg::Cash4) = true
alg_extrapolates(alg::Hairer4) = true
alg_extrapolates(alg::Hairer42) = true
alg_extrapolates(alg::IRKN4) = true
alg_extrapolates(alg::IRKN3) = true
alg_extrapolates(alg::ABDF2) = true
alg_extrapolates(alg::SBDF) = true
alg_extrapolates(alg::MEBDF2) = true
alg_extrapolates(alg::IRKC) = true
alg_extrapolates(alg::MagnusLeapfrog) = true

alg_order(alg::Union{OrdinaryDiffEqAlgorithm,DAEAlgorithm}) = error("Order is not defined for this algorithm")
get_current_alg_order(alg::Union{OrdinaryDiffEqAlgorithm,DAEAlgorithm},cache) = alg_order(alg)
get_current_alg_order(alg::CompositeAlgorithm,cache) = alg_order(alg.algs[cache.current])

get_current_alg_order(alg::OrdinaryDiffEqAdamsVarOrderVarStepAlgorithm,cache) = cache.order
get_current_adaptive_order(alg::OrdinaryDiffEqAdamsVarOrderVarStepAlgorithm,cache) = cache.order
get_current_alg_order(alg::JVODE,cache) = get_current_adaptive_order(alg,cache)
get_current_alg_order(alg::QNDF,cache) = cache.order
get_current_adaptive_order(alg::QNDF,cache) = cache.order
get_current_adaptive_order(alg::OrdinaryDiffEqExtrapolationVarOrderVarStepAlgorithm,cache) = cache.cur_order
get_current_alg_order(alg::OrdinaryDiffEqExtrapolationVarOrderVarStepAlgorithm,cache) = cache.cur_order
get_current_alg_order(alg::ExtrapolationMidpointDeuflhard,cache) = 2(cache.n_curr + 1)
get_current_alg_order(alg::ImplicitDeuflhardExtrapolation,cache) = 2(cache.n_curr + 1)
get_current_adaptive_order(alg::ExtrapolationMidpointDeuflhard,cache) = 2cache.n_curr
get_current_adaptive_order(alg::ImplicitDeuflhardExtrapolation,cache) = 2cache.n_curr
get_current_alg_order(alg::ExtrapolationMidpointHairerWanner,cache) = 2(cache.n_curr + 1)
get_current_alg_order(alg::ImplicitHairerWannerExtrapolation,cache) = 2(cache.n_curr + 1)
get_current_alg_order(alg::ImplicitEulerBarycentricExtrapolation,cache) = cache.n_curr
get_current_alg_order(alg::ImplicitEulerExtrapolation,cache) = cache.n_curr + 1
get_current_adaptive_order(alg::ExtrapolationMidpointHairerWanner,cache) = 2cache.n_curr
get_current_adaptive_order(alg::ImplicitHairerWannerExtrapolation,cache) = 2cache.n_curr
get_current_adaptive_order(alg::ImplicitEulerExtrapolation,cache) = cache.n_curr - 1
get_current_adaptive_order(alg::ImplicitEulerBarycentricExtrapolation,cache) = cache.n_curr - 2

#alg_adaptive_order(alg::OrdinaryDiffEqAdaptiveAlgorithm) = error("Algorithm is adaptive with no order")
get_current_adaptive_order(alg::Union{OrdinaryDiffEqAlgorithm,DAEAlgorithm},cache) = alg_adaptive_order(alg)
get_current_adaptive_order(alg::CompositeAlgorithm,cache) = alg_adaptive_order(alg.algs[cache.current])

alg_order(alg::FunctionMap) = 0
alg_order(alg::Euler) = 1
alg_order(alg::AitkenNeville) = alg.init_order
alg_order(alg::Heun) = 2
alg_order(alg::Ralston) = 2
alg_order(alg::LawsonEuler) = 1
alg_order(alg::NorsettEuler) = 1
alg_order(alg::LieEuler) = 1
alg_order(alg::CayleyEuler) = 2
alg_order(alg::ETDRK2) = 2
alg_order(alg::ETDRK3) = 3
alg_order(alg::ETDRK4) = 4
alg_order(alg::HochOst4) = 4
alg_order(alg::Exp4) = 4
alg_order(alg::EPIRK4s3A) = 4
alg_order(alg::EPIRK4s3B) = 4
alg_order(alg::EPIRK5s3) = 5
alg_order(alg::EPIRK5P1) = 5
alg_order(alg::EPIRK5P2) = 5
alg_order(alg::EXPRB53s3) = 5
alg_order(alg::SplitEuler) = 1
alg_order(alg::ETD2) = 2
alg_order(alg::Exprb32) = 3
alg_order(alg::Exprb43) = 4
alg_order(alg::Anas5) = 5
alg_order(alg::RK46NL) = 4
alg_order(alg::KuttaPRK2p5) = 5
alg_order(alg::RKO65) = 5
alg_order(alg::FRK65) = 6

alg_order(alg::SymplecticEuler) = 1
alg_order(alg::VelocityVerlet) = 2
alg_order(alg::VerletLeapfrog) = 2
alg_order(alg::PseudoVerletLeapfrog) = 2
alg_order(alg::McAte2) = 2
alg_order(alg::Ruth3) = 3
alg_order(alg::McAte3) = 3
alg_order(alg::McAte4) = 4
alg_order(alg::CandyRoz4) = 4
alg_order(alg::CalvoSanz4) = 4
alg_order(alg::McAte42) = 4
alg_order(alg::McAte5) = 5
alg_order(alg::Yoshida6) = 6
alg_order(alg::KahanLi6) = 6
alg_order(alg::McAte8) = 8
alg_order(alg::KahanLi8) = 8
alg_order(alg::SofSpa10) = 10

alg_order(alg::IRKN3) = 3
alg_order(alg::Nystrom4) = 4
alg_order(alg::Nystrom4VelocityIndependent) = 4
alg_order(alg::IRKN4) = 4
alg_order(alg::Nystrom5VelocityIndependent) = 5
alg_order(alg::DPRKN6) = 6
alg_order(alg::DPRKN8) = 8
alg_order(alg::DPRKN12) = 12
alg_order(alg::ERKN4) = 4
alg_order(alg::ERKN5) = 5

alg_order(alg::Midpoint) = 2

alg_order(alg::ORK256) = 2
alg_order(alg::CarpenterKennedy2N54) = 4
alg_order(alg::SHLDDRK52) = 2
alg_order(alg::SHLDDRK_2N) = 4
alg_order(alg::SHLDDRK64) = 4
alg_order(alg::DGLDDRK73_C) = 3
alg_order(alg::DGLDDRK84_C) = 4
alg_order(alg::DGLDDRK84_F) = 4
alg_order(alg::NDBLSRK124) = 4
alg_order(alg::NDBLSRK134) = 4
alg_order(alg::NDBLSRK144) = 4
alg_order(alg::CFRLDDRK64) = 4
alg_order(alg::TSLDDRK74) = 4
alg_order(alg::CKLLSRK43_2) = 3
alg_order(alg::CKLLSRK54_3C) = 4
alg_order(alg::CKLLSRK95_4S) = 5
alg_order(alg::CKLLSRK95_4C) = 5
alg_order(alg::CKLLSRK95_4M) = 5
alg_order(alg::CKLLSRK54_3C_3R) = 4
alg_order(alg::CKLLSRK54_3M_3R) = 4
alg_order(alg::CKLLSRK54_3N_3R) = 4
alg_order(alg::CKLLSRK85_4C_3R) = 5
alg_order(alg::CKLLSRK85_4M_3R) = 5
alg_order(alg::CKLLSRK85_4P_3R) = 5
alg_order(alg::CKLLSRK54_3N_4R) = 4
alg_order(alg::CKLLSRK54_3M_4R) = 4
alg_order(alg::CKLLSRK65_4M_4R) = 5
alg_order(alg::CKLLSRK85_4FM_4R) = 5
alg_order(alg::CKLLSRK75_4M_5R) = 5
alg_order(alg::ParsaniKetchesonDeconinck3S32) = 2
alg_order(alg::ParsaniKetchesonDeconinck3S82) = 2
alg_order(alg::ParsaniKetchesonDeconinck3S53) = 3
alg_order(alg::ParsaniKetchesonDeconinck3S173) = 3
alg_order(alg::ParsaniKetchesonDeconinck3S94) = 4
alg_order(alg::ParsaniKetchesonDeconinck3S184) = 4
alg_order(alg::ParsaniKetchesonDeconinck3S105) = 5
alg_order(alg::ParsaniKetchesonDeconinck3S205) = 5
alg_order(alg::KYK2014DGSSPRK_3S2) = 2

alg_order(alg::SSPRK22) = 2
alg_order(alg::SSPRKMSVS32) = 2
alg_order(alg::SSPRK33) = 3
alg_order(alg::KYKSSPRK42) = 2
alg_order(alg::SSPRK53) = 3
alg_order(alg::SSPRK53_2N1) = 3
alg_order(alg::SSPRK53_2N2) = 3
alg_order(alg::SSPRK53_H) = 3
alg_order(alg::SSPRK63) = 3
alg_order(alg::SSPRK73) = 3
alg_order(alg::SSPRK83) = 3
alg_order(alg::SSPRK43) = 3
alg_order(alg::SSPRK432) = 3
alg_order(alg::SSPRKMSVS43) = 3
alg_order(alg::SSPRK932) = 3
alg_order(alg::SSPRK54) = 4
alg_order(alg::SSPRK104) = 4

alg_order(alg::RK4) = 4
alg_order(alg::RKM) = 4
alg_order(alg::ExplicitRK) = alg.tableau.order

alg_order(alg::BS3) = 3
alg_order(alg::BS5) = 5
alg_order(alg::OwrenZen3) = 3
alg_order(alg::OwrenZen4) = 4
alg_order(alg::OwrenZen5) = 5
alg_order(alg::DP5) = 5
alg_order(alg::Tsit5) = 5
alg_order(alg::DP8) = 8
alg_order(alg::Vern6) = 6
alg_order(alg::Vern7) = 7
alg_order(alg::Vern8) = 8
alg_order(alg::Vern9) = 9
alg_order(alg::TanYam7) = 7
alg_order(alg::TsitPap8) = 8
alg_order(alg::RadauIIA3) = 3
alg_order(alg::RadauIIA5) = 5
alg_order(alg::ImplicitEuler) = 1
alg_order(alg::RKMK2) = 2
alg_order(alg::RKMK4) = 4
alg_order(alg::LieRK4) = 4
alg_order(alg::CG2) = 2
alg_order(alg::CG3) = 3
alg_order(alg::MagnusMidpoint) = 2
alg_order(alg::MagnusGauss4) = 4
alg_order(alg::MagnusNC6) = 6
alg_order(alg::MagnusGL6) = 6
alg_order(alg::MagnusGL8) = 8
alg_order(alg::MagnusNC8) = 8
alg_order(alg::MagnusGL4) = 4
alg_order(alg::MagnusAdapt4) = 4
alg_order(alg::LinearExponential) = 1
alg_order(alg::MagnusLeapfrog) = 2
alg_order(alg::Trapezoid) = 2
alg_order(alg::ImplicitMidpoint) = 2
alg_order(alg::TRBDF2) = 2
alg_order(alg::SSPSDIRK2) = 2
alg_order(alg::SDIRK2) = 2
alg_order(alg::SDIRK22) = 2
alg_order(alg::Kvaerno3) = 3
alg_order(alg::Kvaerno4) = 4
alg_order(alg::Kvaerno5) = 5
alg_order(alg::ESDIRK54I8L2SA) = 5
alg_order(alg::KenCarp3) = 3
alg_order(alg::CFNLIRK3) = 3
alg_order(alg::KenCarp4) = 4
alg_order(alg::KenCarp47) = 4
alg_order(alg::KenCarp5) = 5
alg_order(alg::KenCarp58) = 5
alg_order(alg::Cash4) = 4
alg_order(alg::SFSDIRK4) = 4
alg_order(alg::SFSDIRK5) = 4
alg_order(alg::SFSDIRK6) = 4
alg_order(alg::SFSDIRK7) = 4
alg_order(alg::SFSDIRK8) = 4
alg_order(alg::Hairer4) = 4
alg_order(alg::Hairer42) = 4
alg_order(alg::Feagin10) = 10
alg_order(alg::Feagin12) = 12
alg_order(alg::Feagin14) = 14
alg_order(alg::PFRK87) = 8

alg_order(alg::Rosenbrock23) = 2
alg_order(alg::Rosenbrock32) = 3
alg_order(alg::ROS3P) = 3
alg_order(alg::Rodas3) = 3
alg_order(alg::ROS34PW1a) = 3
alg_order(alg::ROS34PW1b) = 3
alg_order(alg::ROS34PW2) = 3
alg_order(alg::ROS34PW3) = 4
alg_order(alg::RosShamp4) = 4
alg_order(alg::Veldd4) = 4
alg_order(alg::Velds4) = 4
alg_order(alg::GRK4T) = 4
alg_order(alg::GRK4A) = 4
alg_order(alg::Ros4LStab) = 4
alg_order(alg::RosenbrockW6S4OS) = 4
alg_order(alg::Rodas4) = 4
alg_order(alg::Rodas42) = 4
alg_order(alg::Rodas4P) = 4
alg_order(alg::Rodas5) = 5

alg_order(alg::AB3) = 3
alg_order(alg::AB4) = 4
alg_order(alg::AB5) = 5
alg_order(alg::ABM32) = 3
alg_order(alg::ABM43) = 4
alg_order(alg::ABM54) = 5

alg_order(alg::VCAB3) = 3
alg_order(alg::VCAB4) = 4
alg_order(alg::VCAB5) = 5
alg_order(alg::VCABM3) = 3
alg_order(alg::VCABM4) = 4
alg_order(alg::VCABM5) = 5

alg_order(alg::VCABM) = 1  #dummy value

alg_order(alg::CNAB2) = 2
alg_order(alg::CNLF2) = 2

alg_order(alg::AN5) = 5
alg_order(alg::JVODE) = 1  #dummy value

alg_order(alg::ABDF2) = 2
alg_order(alg::QNDF1) = 1
alg_order(alg::QNDF2) = 2

alg_order(alg::QNDF) = 1 #dummy value

alg_order(alg::SBDF) = alg.order

alg_order(alg::ROCK2) = 2
alg_order(alg::ROCK4) = 4

alg_order(alg::ESERK4) = 4
alg_order(alg::ESERK5) = 5
alg_order(alg::SERK2) = 2

alg_order(alg::RKC) = 2
alg_order(alg::IRKC) = 2

alg_order(alg::MEBDF2) = 2
alg_order(alg::PDIRK44) = 4

alg_order(alg::DImplicitEuler) = 1
alg_order(alg::DABDF2) = 2

alg_maximum_order(alg) = alg_order(alg)
alg_maximum_order(alg::CompositeAlgorithm) = maximum(alg_order(x) for x in alg.algs)
alg_maximum_order(alg::ExtrapolationMidpointDeuflhard) = 2(alg.n_max+1)
alg_maximum_order(alg::ImplicitDeuflhardExtrapolation) = 2(alg.n_max+1)
alg_maximum_order(alg::ExtrapolationMidpointHairerWanner) = 2(alg.n_max+1)
alg_maximum_order(alg::ImplicitHairerWannerExtrapolation) = 2(alg.n_max+1)
alg_maximum_order(alg::ImplicitEulerExtrapolation) = 2(alg.n_max+1)

alg_maximum_order(alg::ImplicitEulerBarycentricExtrapolation) = alg.n_max
alg_adaptive_order(alg::ExplicitRK) = alg.tableau.adaptiveorder
alg_adaptive_order(alg::Union{OrdinaryDiffEqAlgorithm,DAEAlgorithm}) = alg_order(alg)-1
alg_adaptive_order(alg::Feagin10) = 8
alg_adaptive_order(alg::Feagin14) = 12

alg_adaptive_order(alg::Rosenbrock23) = 3
alg_adaptive_order(alg::Rosenbrock32) = 2

alg_adaptive_order(alg::RadauIIA3) = 1
alg_adaptive_order(alg::RadauIIA5) = 3
alg_adaptive_order(alg::RKC) = 2
alg_adaptive_order(alg::IRKC) = 1

alg_adaptive_order(alg::ImplicitEuler) = 0
alg_adaptive_order(alg::Trapezoid) = 1
# this is actually incorrect and is purposefully decreased as this tends
# to track the real error much better
alg_adaptive_order(alg::ImplicitMidpoint) = 1
# this is actually incorrect and is purposefully decreased as this tends
# to track the real error much better

alg_adaptive_order(alg::Exprb32) = 2
alg_adaptive_order(alg::Exprb43) = 4
alg_adaptive_order(alg::AN5) = 5

beta2_default(alg::Union{OrdinaryDiffEqAlgorithm,DAEAlgorithm}) = isadaptive(alg) ? 2//(5alg_order(alg)) : 0
beta2_default(alg::FunctionMap) = 0
beta2_default(alg::DP8) = 0//1
beta2_default(alg::DP5) = 4//100
beta2_default(alg::ExtrapolationMidpointDeuflhard) = 0//1
beta2_default(alg::ImplicitDeuflhardExtrapolation) = 0//1
beta2_default(alg::ExtrapolationMidpointHairerWanner) = 0//1
beta2_default(alg::ImplicitHairerWannerExtrapolation) = 0//1
beta2_default(alg::ImplicitEulerExtrapolation) = 0//1
beta2_default(alg::ImplicitEulerBarycentricExtrapolation) = 0//1

beta1_default(alg::Union{OrdinaryDiffEqAlgorithm,DAEAlgorithm},beta2) = isadaptive(alg) ? 7//(10alg_order(alg)) : 0
beta1_default(alg::FunctionMap,beta2) = 0
beta1_default(alg::DP8,beta2) = typeof(beta2)(1//alg_order(alg)) - beta2/5
beta1_default(alg::DP5,beta2) = typeof(beta2)(1//alg_order(alg)) - 3beta2/4
beta1_default(alg::ExtrapolationMidpointDeuflhard,beta2) =  1//(2alg.n_init+1)
beta1_default(alg::ImplicitDeuflhardExtrapolation,beta2) =  1//(2alg.n_init+1)
beta1_default(alg::ExtrapolationMidpointHairerWanner,beta2) =  1//(2alg.n_init+1)
beta1_default(alg::ImplicitHairerWannerExtrapolation,beta2) =  1//(2alg.n_init+1)
beta1_default(alg::ImplicitEulerExtrapolation,beta2) =  1//(alg.n_init+1)
beta1_default(alg::ImplicitEulerBarycentricExtrapolation,beta2) =  1//(alg.n_init - 1)

gamma_default(alg::Union{OrdinaryDiffEqAlgorithm,DAEAlgorithm}) = isadaptive(alg) ? 9//10 : 0
gamma_default(alg::RKC) = 8//10
gamma_default(alg::IRKC) = 8//10
gamma_default(alg::ExtrapolationMidpointDeuflhard) = (1//4)^beta1_default(alg,beta2_default(alg))
gamma_default(alg::ImplicitDeuflhardExtrapolation) = (1//4)^beta1_default(alg,beta2_default(alg))
gamma_default(alg::ExtrapolationMidpointHairerWanner) = (65//100)^beta1_default(alg,beta2_default(alg))
gamma_default(alg::ImplicitHairerWannerExtrapolation) = (65//100)^beta1_default(alg,beta2_default(alg))
gamma_default(alg::ImplicitEulerExtrapolation) = (65//100)^beta1_default(alg,beta2_default(alg))

gamma_default(alg::ImplicitEulerBarycentricExtrapolation) = (80//100)^beta1_default(alg,beta2_default(alg))

qsteady_min_default(alg::Union{OrdinaryDiffEqAlgorithm,DAEAlgorithm}) = 1
qsteady_max_default(alg::Union{OrdinaryDiffEqAlgorithm,DAEAlgorithm}) = 1
qsteady_max_default(alg::OrdinaryDiffEqAdaptiveImplicitAlgorithm) = 6//5
# But don't re-use Jacobian if not adaptive: too risky and cannot pull back
qsteady_max_default(alg::OrdinaryDiffEqImplicitAlgorithm) = isadaptive(alg) ? 1//1 : 0
qsteady_max_default(alg::AN5) = 3//2
qsteady_max_default(alg::JVODE) = 3//2
qsteady_max_default(alg::QNDF1) = 2//1
qsteady_max_default(alg::QNDF2) = 2//1
qsteady_max_default(alg::QNDF) = 2//1

FunctionMap_scale_by_time(alg::FunctionMap{scale_by_time}) where {scale_by_time} = scale_by_time

# SSP coefficients
"""
    ssp_coefficient(alg)

Return the SSP coefficient of the ODE algorithm `alg`. If one time step of size
`dt` with `alg` can be written as a convex combination of explicit Euler steps
with step sizes `cᵢ * dt`, the SSP coefficient is the minimal value of `1/cᵢ`.

# Examples
```julia-repl
julia> ssp_coefficient(SSPRK104())
6
```
"""
ssp_coefficient(alg) = error("$alg is not a strong stability preserving method.")
ssp_coefficient(alg::Euler) = 1
ssp_coefficient(alg::SSPRK22) = 1
ssp_coefficient(alg::SSPRK33) = 1
ssp_coefficient(alg::KYKSSPRK42) = 2.459
ssp_coefficient(alg::SSPRK53) = 2.65
ssp_coefficient(alg::SSPRK53_2N1) = 2.18
ssp_coefficient(alg::SSPRK53_2N2) = 2.148
ssp_coefficient(alg::SSPRK53_H) = 2.65
ssp_coefficient(alg::SSPRK63) = 3.518
ssp_coefficient(alg::SSPRK73) = 4.2879
ssp_coefficient(alg::SSPRK83) = 5.107
ssp_coefficient(alg::SSPRK43) = 2
ssp_coefficient(alg::SSPRK432) = 2
ssp_coefficient(alg::SSPRKMSVS32) = 0.5
ssp_coefficient(alg::SSPRKMSVS43) = 0.33
ssp_coefficient(alg::SSPRK932) = 6
ssp_coefficient(alg::SSPRK54) = 1.508
ssp_coefficient(alg::SSPRK104) = 6
ssp_coefficient(alg::KYK2014DGSSPRK_3S2) =  0.8417

# We shouldn't do this probably.
#ssp_coefficient(alg::ImplicitEuler) = Inf
ssp_coefficient(alg::SSPSDIRK2) = 4

# stability regions
alg_stability_size(alg::DP5) = 3.3066
alg_stability_size(alg::Tsit5) = 3.5068
alg_stability_size(alg::Vern6) = 4.8553
alg_stability_size(alg::Vern7) = 4.6400
alg_stability_size(alg::Vern8) = 5.8641
alg_stability_size(alg::Vern9) = 4.4762

alg_can_repeat_jac(alg::Union{OrdinaryDiffEqAlgorithm,DAEAlgorithm}) = false
alg_can_repeat_jac(alg::OrdinaryDiffEqNewtonAdaptiveAlgorithm) = true
alg_can_repeat_jac(alg::IRKC) = false

function unwrap_alg(integrator, is_stiff)
  alg = integrator.alg
  iscomp = typeof(alg) <: CompositeAlgorithm
  if !iscomp
    return alg
  elseif typeof(alg.choice_function) <: AutoSwitchCache
    num = is_stiff ? 2 : 1
    if num == 1
      return alg.algs[1]
    else
      return alg.algs[2]
    end
  else
    if integrator.cache.current == 1
      return alg.algs[1]
    elseif integrator.cache.current == 2
      return alg.algs[2]
    else
      return alg.algs[integrator.cache.current]
    end
  end
end

# Whether `uprev` is used in the algorithm directly.
uses_uprev(alg::Union{OrdinaryDiffEqAlgorithm,DAEAlgorithm}, adaptive::Bool) = true
uses_uprev(alg::ORK256, adaptive::Bool) = false
uses_uprev(alg::CarpenterKennedy2N54, adaptive::Bool) = false
uses_uprev(alg::SHLDDRK64, adaptive::Bool) = false
uses_uprev(alg::DGLDDRK73_C, adaptive::Bool) = false
uses_uprev(alg::DGLDDRK84_C, adaptive::Bool) = false
uses_uprev(alg::DGLDDRK84_F, adaptive::Bool) = false
uses_uprev(alg::NDBLSRK124, adaptive::Bool) = false
uses_uprev(alg::NDBLSRK134, adaptive::Bool) = false
uses_uprev(alg::NDBLSRK144, adaptive::Bool) = false
uses_uprev(alg::CFRLDDRK64, adaptive::Bool) = false
uses_uprev(alg::TSLDDRK74, adaptive::Bool) = false
uses_uprev(alg::OrdinaryDiffEqAdaptiveAlgorithm, adaptive::Bool) = true
uses_uprev(alg::CKLLSRK43_2, adaptive::Bool) = adaptive
uses_uprev(alg::CKLLSRK54_3C, adaptive::Bool) = adaptive
uses_uprev(alg::CKLLSRK95_4S, adaptive::Bool) = adaptive
uses_uprev(alg::CKLLSRK95_4C, adaptive::Bool) = adaptive
uses_uprev(alg::CKLLSRK95_4M, adaptive::Bool) = adaptive
uses_uprev(alg::CKLLSRK54_3C_3R, adaptive::Bool) = adaptive
uses_uprev(alg::CKLLSRK54_3M_3R, adaptive::Bool) = adaptive
uses_uprev(alg::CKLLSRK54_3N_3R, adaptive::Bool) = adaptive
uses_uprev(alg::CKLLSRK85_4C_3R, adaptive::Bool) = adaptive
uses_uprev(alg::CKLLSRK85_4M_3R, adaptive::Bool) = adaptive
uses_uprev(alg::CKLLSRK85_4P_3R, adaptive::Bool) = adaptive
uses_uprev(alg::CKLLSRK54_3N_4R, adaptive::Bool) = adaptive
uses_uprev(alg::CKLLSRK54_3M_4R, adaptive::Bool) = adaptive
uses_uprev(alg::CKLLSRK65_4M_4R, adaptive::Bool) = adaptive
uses_uprev(alg::CKLLSRK85_4FM_4R, adaptive::Bool) = adaptive
uses_uprev(alg::CKLLSRK75_4M_5R, adaptive::Bool) = adaptive

ispredictive(alg::Union{OrdinaryDiffEqAlgorithm,DAEAlgorithm}) = false
ispredictive(alg::Union{RKC}) = true
ispredictive(alg::Union{SERK2}) = alg.controller === :Predictive
ispredictive(alg::OrdinaryDiffEqNewtonAdaptiveAlgorithm) = alg.controller === :Predictive
isstandard(alg::OrdinaryDiffEqNewtonAdaptiveAlgorithm) = alg.controller === :Standard
isstandard(alg::VCABM) = true
isstandard(alg::Union{OrdinaryDiffEqAlgorithm,DAEAlgorithm}) = false
ispi(alg::Union{OrdinaryDiffEqAlgorithm,DAEAlgorithm}) = !(ispredictive(alg) || isstandard(alg))

isWmethod(alg::Union{OrdinaryDiffEqAlgorithm,DAEAlgorithm}) = false
isWmethod(alg::Rosenbrock23) = true
isWmethod(alg::Rosenbrock32) = true
isWmethod(alg::ROS34PW1a) = true
isWmethod(alg::ROS34PW1b) = true
isWmethod(alg::ROS34PW2) = true
isWmethod(alg::ROS34PW3) = true
isWmethod(alg::RosenbrockW6S4OS) = true

isesdirk(alg::TRBDF2) = true
isesdirk(alg::Union{KenCarp3, KenCarp4, KenCarp5, KenCarp58,
                    Kvaerno3, Kvaerno4, Kvaerno5,
                    ESDIRK54I8L2SA,CFNLIRK3}) = true
isesdirk(alg::Union{OrdinaryDiffEqAlgorithm,DAEAlgorithm}) = false

is_mass_matrix_alg(alg::Union{OrdinaryDiffEqAlgorithm,DAEAlgorithm}) = false
is_mass_matrix_alg(alg::RosenbrockAlgorithm) = true
is_mass_matrix_alg(alg::NewtonAlgorithm) = !isesdirk(alg)


function uniquearrayrefs(ls::LoopSet)
    uniquerefs = ArrayReferenceMeta[]
    # for arrayref ∈ ls.refs_aliasing_syms
    for op ∈ operations(ls)
        arrayref = op.ref
        arrayref === NOTAREFERENCE && continue
        any(ref -> sameref(arrayref, ref), uniquerefs) || push!(uniquerefs, arrayref)
    end
    uniquerefs
end

otherindexunrolled(loopsym::Symbol, ind::Symbol, loopdeps::Vector{Symbol}) = (loopsym !== ind) && (loopsym ∈ loopdeps)
function otherindexunrolled(ls::LoopSet, ind::Symbol, ref::ArrayReferenceMeta)
    us = ls.unrollspecification[]
    u₁sym = names(ls)[us.u₁loopnum]
    u₂sym = us.u₂loopnum > 0 ? names(ls)[us.u₂loopnum] : Symbol("##undefined##")
    otherindexunrolled(u₁sym, ind, loopdependencies(ref)) || otherindexunrolled(u₂sym, ind, loopdependencies(ref))
end
multiple_with_name(n::Symbol, v::Vector{ArrayReferenceMeta}) = sum(ref -> n === vptr(ref), v) > 1
# TODO: DRY between indices_calculated_by_pointer_offsets and use_loop_induct_var
function indices_calculated_by_pointer_offsets(ls::LoopSet, ar::ArrayReferenceMeta)
    indices = getindices(ar)
    ls.isbroadcast[] && return fill(false, length(indices))
    looporder = names(ls)
    offset = isdiscontiguous(ar)
    gespinds = Expr(:tuple)
    out = Vector{Bool}(undef, length(indices))
    li = ar.loopedindex
    # @show ls.vector_width[]
    for i ∈ eachindex(li)
        ii = i + offset
        ind = indices[ii]
        # j = findfirst(isequal(ind), view(indices, offset+1:ii-1))
        # if !isnothing(j)
        #     out[i] = out[j - offset]
        #     continue
        # end
        if (!li[i]) || multiple_with_name(vptr(ar), ls.lssm[].uniquearrayrefs) ||
            (iszero(ls.vector_width[]) && isstaticloop(getloop(ls, ind)))# ||
            # ((ls.align_loops[] > 0) && (first(names(ls)) == ind))
            out[i] = false
        elseif (isone(ii) && (first(looporder) === ind))
            out[i] = otherindexunrolled(ls, ind, ar)
        else
            out[i] = true
        end
    end
    out
end

"""
Returns a vector of length equal to the number of indices.
A value > 0 indicates which loop number that index corresponds to when incrementing the pointer.
A value < 0 indicates that abs(value) is the corresponding loop, and a `loopvalue` will be used.
"""
function use_loop_induct_var!(ls::LoopSet, q::Expr, ar::ArrayReferenceMeta, allarrayrefs::Vector{ArrayReferenceMeta})
    us = ls.unrollspecification[]
    li = ar.loopedindex
    looporder = reversenames(ls)
    uliv = Vector{Int}(undef, length(li))
    indices = getindices(ar)
    offset = first(indices) === Symbol("##DISCONTIGUOUSSUBARRAY##")
    if length(indices) != offset + length(li)
        println(ar)
        throw("Length of indices and length of offset do not match!")
    end
    isbroadcast = ls.isbroadcast[]
    gespinds = Expr(:tuple)
    offsetprecalc_descript = Expr(:tuple)
    use_offsetprecalc = false
    for i ∈ eachindex(li)
        ii = i + offset
        ind = indices[ii]
        # j = findfirst(isequal(ind), view(indices, offset+1:ii-1))
        # if !isnothing(j)
        #     j -= offset
        #     push!(gespinds.args, gespinds.args[j])
        #     uliv[i] = uliv[j]
        # else
        if (!li[i])
            uliv[i] = 0
            push!(gespinds.args, Expr(:call, lv(:Zero)))
            push!(offsetprecalc_descript.args, 0)
        elseif isbroadcast ||
            ((isone(ii) && (last(looporder) === ind)) && !(otherindexunrolled(ls, ind, ar)) ||
             multiple_with_name(vptr(ar), allarrayrefs)) ||
             (iszero(ls.vector_width[]) && isstaticloop(getloop(ls, ind)))# ||
             # ((ls.align_loops[] > 0) && (first(names(ls)) == ind))

            # Not doing normal offset indexing
            uliv[i] = -findfirst(isequal(ind), looporder)::Int
            push!(gespinds.args, Expr(:call, lv(:Zero)))
            push!(offsetprecalc_descript.args, 0) # not doing offset indexing, so push 0
        else
            uliv[i] = findfirst(isequal(ind), looporder)::Int
            loop = getloop(ls, ind)
            if loop.startexact
                push!(gespinds.args, Expr(:call, Expr(:curly, lv(:Static), loop.starthint - 1)))
            else
                push!(gespinds.args, Expr(:call, lv(:staticm1), loop.startsym))
            end
            if ind === names(ls)[us.vectorizedloopnum]
                push!(offsetprecalc_descript.args, 0)
            elseif (ind === names(ls)[us.u₁loopnum]) & (us.u₁ > 3)
                use_offsetprecalc = true
                push!(offsetprecalc_descript.args, us.u₁)
            elseif (ind === names(ls)[us.u₂loopnum]) & (us.u₂ > 3)
                use_offsetprecalc = true
                push!(offsetprecalc_descript.args, us.u₂)
            else
                push!(offsetprecalc_descript.args, 0)
            end
        end
    end
    if use_offsetprecalc
        push!(q.args, Expr(:(=), vptr(ar), Expr(:call, lv(:offsetprecalc), Expr(:call, lv(:gesp), vptr(ar), gespinds), Expr(:call, Expr(:curly, :Val, offsetprecalc_descript)))))
    else
        push!(q.args, Expr(:(=), vptr(ar), Expr(:call, lv(:gesp), vptr(ar), gespinds)))
    end
    uliv
end

# Plan here is that we increment every unique array
function add_loop_start_stop_manager!(ls::LoopSet)
    q = Expr(:block)
    us = ls.unrollspecification[]
    # Presence of an explicit use of a loopinducation var means we should use that, so we look for one
    # TODO: replace first with only once you add Compat as a dep or drop support for older Julia versions
    loopinductvars = map(op -> first(loopdependencies(op)), filter(isloopvalue, operations(ls)))
    # Filtered ArrayReferenceMetas, we must increment each
    arrayrefs = uniquearrayrefs(ls)
    use_livs = map(ar -> use_loop_induct_var!(ls, q, ar, arrayrefs), arrayrefs)
    # loops, sorted from outer-most to inner-most
    looporder = reversenames(ls)
    # For each loop, we need to choose an induction variable
    nloops = length(looporder)
    # loopstarts = Vector{Vector{ArrayReferenceMeta}}(undef, nloops)
    loopstarts = fill(ArrayReferenceMeta[], nloops)
    terminators = Vector{Int}(undef, nloops) # zero is standard loop induct var
    # loopincrements = Vector{Vector{ArrayReferenceMeta}}(undef, nloops) # Not needed; copy of loopstarts
    # The purpose of the reminds thing here is to pick which of these to use for the terminator
    # We want to pick the one with the fewest outer loops with respect to this one, to minimize
    # the number of redefinitions of max-pointer used for the termination comparison.
    reached_indices = zeros(Int, length(arrayrefs))
    for (i,loopsym) ∈ enumerate(looporder) # iterates from outer to inner
        loopstartᵢ = ArrayReferenceMeta[]; arⱼ = 0; minrem = typemax(Int);
        ric = Tuple{Int,Int}[]
        for j ∈ eachindex(use_livs) # j is array ref number
            for (l,k) ∈ enumerate(use_livs[j])# l is index number, k is loop number
                if k == i
                    push!(loopstartᵢ, arrayrefs[j])
                    push!(ric, ((reached_indices[j] += 1), length(loopstartᵢ)))
                end
            end
        end
        loopstarts[i] = loopstartᵢ
        terminators[i] = if (loopsym ∈ loopinductvars) || (any(r -> any(isequal(-i), r), use_livs)) || iszero(length(loopstartᵢ))
            0
        else
            # @show i, loopsym loopdependencies.(operations(ls)) operations(ls)
            # @assert !iszero(length(loopstartᵢ))
            last(ric[argmin(first.(ric))]) # index corresponds to array ref's position in loopstart
        end
    end
    ls.lssm[] = LoopStartStopManager(
        reverse!(terminators), reverse!(loopstarts), arrayrefs
    )
    q
end
maxsym(ptr, sub) = Symbol(ptr, "##MAX##", sub, "##")
function pointermax(ls::LoopSet, ar::ArrayReferenceMeta, n::Int, sub::Int, isvectorized::Bool)::Expr
    pointermax(ls, ar, n, sub, isvectorized, getloop(ls, names(ls)[n]))
end
function pointermax(ls::LoopSet, ar::ArrayReferenceMeta, n::Int, sub::Int, isvectorized::Bool, loop::Loop)::Expr
    pointermax(ls, ar, n, sub, isvectorized, looplengthexpr(loop, n))::Expr
end
function pointermax(ls::LoopSet, ar::ArrayReferenceMeta, n::Int, sub::Int, isvectorized::Bool, stophint::Int)::Expr
    # @unpack u₁loopnum, u₂loopnum, vectorizedloopnum, u₁, u₂ = us
    loopsym = names(ls)[n]
    index = Expr(:tuple)
    for i ∈ getindicesonly(ar)
        if i === loopsym
            if iszero(sub)
                push!(index.args, stophint)
            elseif isvectorized
                if isone(sub)
                    push!(index.args, Expr(:call, lv(:valsub), stophint, VECTORWIDTHSYMBOL))
                else
                    push!(index.args, Expr(:call, lv(:vsub), stophint, Expr(:call, lv(:valmul), VECTORWIDTHSYMBOL, sub)))
                end
            else
                push!(index.args, stophint - sub)
            end
            ptr = vptr(ar)
            return Expr(:call, lv(:pointerforcomparison), ptr, index)
        else
            push!(index.args, Expr(:call, lv(:Zero)))
        end
    end
    @show ar, loopsym
end
function pointermax(ls::LoopSet, ar::ArrayReferenceMeta, n::Int, sub::Int, isvectorized::Bool, stopsym)::Expr
    # @unpack u₁loopnum, u₂loopnum, vectorizedloopnum, u₁, u₂ = us
    loopsym = names(ls)[n]
    index = Expr(:tuple)
    for i ∈ getindicesonly(ar)
        if i === loopsym
            if iszero(sub)
                push!(index.args, stopsym)
            elseif isvectorized
                if isone(sub)
                    push!(index.args, Expr(:call, lv(:valsub), stopsym, VECTORWIDTHSYMBOL))
                else
                    push!(index.args, Expr(:call, lv(:vsub), stopsym, Expr(:call, lv(:valmul), VECTORWIDTHSYMBOL, sub)))
                end
            else
                push!(index.args, Expr(:call, lv(:vsub), stopsym, sub))
            end
            return Expr(:call, lv(:pointerforcomparison), vptr(ar), index)
        else
            push!(index.args, Expr(:call, lv(:Zero)))
        end
    end
    @show ar, loopsym
end

function defpointermax(ls::LoopSet, ar::ArrayReferenceMeta, n::Int, sub::Int, isvectorized::Bool)::Expr
    Expr(:(=), maxsym(vptr(ar), sub), pointermax(ls, ar, n, sub, isvectorized))
end

function maxunroll(us::UnrollSpecification, n)
    @unpack u₁loopnum, u₂loopnum, u₁, u₂ = us
    if n == u₁loopnum# && u₁ > 1
        u₁
    elseif n == u₂loopnum# && u₂ > 1
        u₂
    else
        1
    end
end

function startloop(ls::LoopSet, us::UnrollSpecification, n::Int, submax = maxunroll(us, n))
    @unpack u₁loopnum, u₂loopnum, vectorizedloopnum, u₁, u₂ = us
    lssm = ls.lssm[]
    termind = lssm.terminators[n]
    ptrdefs = lssm.incrementedptrs[n]
    loopstart = Expr(:block)
    firstloop = n == num_loops(ls)
    for ar ∈ ptrdefs
        ptr = vptr(ar)
        push!(loopstart.args, Expr(:(=), ptr, ptr))
    end
    if iszero(termind)
        loopsym = names(ls)[n]
        push!(loopstart.args, startloop(getloop(ls, loopsym), loopsym))
    else
        isvectorized = n == vectorizedloopnum
        for sub ∈ 0:submax
            push!(loopstart.args, defpointermax(ls, ptrdefs[termind], n, sub, isvectorized))
        end
    end
    loopstart
end
function offset_ptr(ar::ArrayReferenceMeta, us::UnrollSpecification, loopsym::Symbol, n::Int, UF::Int, offsetinds::Vector{Bool})
    indices = getindices(ar)
    offset = first(indices) === DISCONTIGUOUS
    gespinds = Expr(:tuple)
    li = ar.loopedindex
    for i ∈ eachindex(li)
        ii = i + offset
        ind = indices[ii]
        if !offsetinds[i]
            push!(gespinds.args, Expr(:call, lv(:Zero)))
        elseif ind == loopsym
            incrementloopcounter!(gespinds, us, n, UF)
        else
            push!(gespinds.args, Expr(:call, lv(:Zero)))
        end
        ind == loopsym && break
    end
    Expr(:(=), vptr(ar), Expr(:call, lv(:gesp), vptr(ar), gespinds))
end
function incrementloopcounter(ls::LoopSet, us::UnrollSpecification, n::Int, UF::Int)
    @unpack u₁loopnum, u₂loopnum, vectorizedloopnum, u₁, u₂ = us
    lssm = ls.lssm[]
    ptrdefs = lssm.incrementedptrs[n]
    looporder = names(ls)
    loopsym = looporder[n]
    q = Expr(:block)
    termind = lssm.terminators[n]
    if iszero(termind) # increment liv
        push!(q.args, incrementloopcounter(us, n, loopsym, UF))
    end
    for (j,ar) ∈ enumerate(ptrdefs)
        offsetinds = indices_calculated_by_pointer_offsets(ls, ar)
        push!(q.args, offset_ptr(ar, us, loopsym, n, UF, offsetinds))
    end
    q
end
function terminatecondition(ls::LoopSet, us::UnrollSpecification, n::Int, inclmask::Bool, UF::Int)
    lssm = ls.lssm[]
    termind = lssm.terminators[n]
    if iszero(termind)
        loop = getloop(ls, names(ls)[n])
        return terminatecondition(loop, us, n, loop.itersymbol, inclmask, UF)
    end

    termar = lssm.incrementedptrs[n][termind]
    ptr = vptr(termar)
    # @show UF, isvectorized(us, n)
    if inclmask && isvectorized(us, n)
        Expr(:call, :<, callpointerforcomparison(ptr), maxsym(ptr, 0))
    else
        Expr(:call, :≤, callpointerforcomparison(ptr), maxsym(ptr, UF))
    end
end


export AtmosFilterPerturbations
export AtmosSpecificFilterPerturbations

struct AtmosFilterPerturbations{M} <: AbstractFilterTarget
    atmos::M
end

vars_state_filtered(target::AtmosFilterPerturbations, FT) =
    vars_state(target.atmos, Prognostic(), FT)

function compute_filter_argument!(
    target::AtmosFilterPerturbations,
    filter_state::Vars,
    state::Vars,
    aux::Vars,
)
    # copy the whole state
    parent(filter_state) .= parent(state)
    # remove reference state
    filter_state.ρ -= aux.ref_state.ρ
    filter_state.energy.ρe -= aux.ref_state.ρe
    if !(moisture_model(target.atmos) isa DryModel)
        filter_state.moisture.ρq_tot -= aux.ref_state.ρq_tot
    end
    if (moisture_model(target.atmos) isa NonEquilMoist)
        filter_state.moisture.ρq_liq -= aux.ref_state.ρq_liq
        filter_state.moisture.ρq_ice -= aux.ref_state.ρq_ice
    end
end
function compute_filter_result!(
    target::AtmosFilterPerturbations,
    state::Vars,
    filter_state::Vars,
    aux::Vars,
)
    # copy the whole filter state
    parent(state) .= parent(filter_state)
    # add reference state
    state.ρ += aux.ref_state.ρ
    state.energy.ρe += aux.ref_state.ρe
    if !(moisture_model(target.atmos) isa DryModel)
        state.moisture.ρq_tot += aux.ref_state.ρq_tot
    end
    if (moisture_model(target.atmos) isa NonEquilMoist)
        filter_state.moisture.ρq_liq += aux.ref_state.ρq_liq
        filter_state.moisture.ρq_ice += aux.ref_state.ρq_ice
    end
end

struct AtmosSpecificFilterPerturbations{M} <: AbstractFilterTarget
    atmos::M
end

vars_state_filtered(target::AtmosSpecificFilterPerturbations, FT) =
    vars_state_filtered(target.atmos, FT)

function compute_filter_argument!(
    target::AtmosSpecificFilterPerturbations,
    filter_state::Vars,
    state::Vars,
    aux::Vars,
)

    ρ_inv = 1 / state.ρ
    ρ_ref_inv = 1 / aux.ref_state.ρ

    # copy the whole state
    parent(filter_state) .= parent(state) * ρ_inv

    # remove reference state
    filter_state.energy.e -= aux.ref_state.ρe * ρ_ref_inv
    if !(moisture_model(target.atmos) isa DryModel)
        filter_state.moisture.q_tot -= aux.ref_state.ρq_tot * ρ_ref_inv
    end
    if (moisture_model(target.atmos) isa NonEquilMoist)
        filter_state.moisture.q_liq -= aux.ref_state.ρq_liq * ρ_ref_inv
        filter_state.moisture.q_ice -= aux.ref_state.ρq_ice * ρ_ref_inv
    end

end

function compute_filter_result!(
    target::AtmosSpecificFilterPerturbations,
    state::Vars,
    filter_state::Vars,
    aux::Vars,
)

    ρ = state.ρ
    ρ_ρ_ref_ratio = ρ / aux.ref_state.ρ

    # copy the whole filter state
    parent(state) .= parent(filter_state) * ρ

    # add reference state
    state.energy.ρe += aux.ref_state.ρe * ρ_ρ_ref_ratio
    if !(moisture_model(target.atmos) isa DryModel)
        state.moisture.ρq_tot += aux.ref_state.ρq_tot * ρ_ρ_ref_ratio
    end
    if (moisture_model(target.atmos) isa NonEquilMoist)
        state.moisture.ρq_liq += aux.ref_state.ρq_liq * ρ_ρ_ref_ratio
        state.moisture.ρq_ice += aux.ref_state.ρq_ice * ρ_ρ_ref_ratio
    end
end


"""
    solve(problem::WaveguideProblem; [ρ₀, Nouts], kwargs...)

Solves the `WaveguideProblem`, using `QuantumOptics.timeevolution.master_nh[_dynamic]`.

# Arguments
- `problem`: the waveguide problem
- `Nouts::Array{Int}` (optional): specifies the dimension of the output Fock spaces
- `ρ₀` (optional): Alternative initial state for the simulation.
	Overrides the initial state provided by the `WavePacket` and the quantum system.
- Additional keyword arguments are passed to `QuantumOptics.jl`'s master equation solver,
  see [it's documentation](https://docs.qojulia.org/api/#QuantumOptics.timeevolution.master)

# Returns
- ts:    The times on which the output is determined.
- ρ(t):  The density matrix at each point in time. The Hilbert space depends on the specific
         problem; if unsure check out `basis(ρ[1])`.
"""
function solve(problem::_DrivenProblem{O, Coherent, ContinuousWave};
		ρ₀ = nothing, kwargs...) where {O}
	ψ₀, Ls, H_nh = _generateLindbladian(problem)
	ρ₀ = isnothing(ρ₀) ? ψ₀ : ρ₀

	return timeevolution.master_nh(problem.ts, ρ₀, H_nh, Ls; kwargs...)
end

function solve(problem::WaveguideProblem; Nouts::Union{Nothing,Array} = nothing,
		ρ₀ = nothing,  kwargs...)
	if problem isa _ScatterProblem
		if !isnothing(ρ₀)
			!isnothing(Nouts) && @warn "`Nouts` and `ρ₀` simultaneously set. Ignoring `Nouts`"
			Nouts = [basis.N for basis ∈ basis(ρ₀).bases[3:end]]
		end
		isnothing(Nouts) && ( Nouts = _expectedOutputPhotons(problem) )
		@assert length(Nouts) == length(problem.ψₒ)
	end

	ψ₀, Ls, H_nh = _generateLindbladian(problem, Nouts)
	ρ₀ = isnothing(ρ₀) ? ψ₀ : ρ₀

	function Operators(t, ψ)
		Ls_t = [Ls[1](t), Ls[2:end]...]
		return H_nh(t), H_nh(t)', Ls_t, dagger.(Ls_t)
	end

	return timeevolution.master_nh_dynamic(problem.ts, ρ₀, Operators; kwargs...)
end

function _expectedOutputPhotons(problem::_ScatterProblem{O, Coherent, ContinuousWave}) where {O}
	if problem.displace_output
		return fill(ceil(Int, 3 * sqrt(problem.ts[end] - problem.ts[1]) ), length(problem.ψₒ) )
	end
	fill(coherent_cutoff( sqrt( (problem.ts[end] - problem.ts[1]) * abs2(problem.ψᵢ.α) ) ),
		length(problem.ψₒ)
	)
end

function _expectedOutputPhotons(problem::_ScatterProblem{O, Coherent, WavePacket{Coherent}}) where {O}
	if problem.displace_output
		return fill(ceil(Int, 3 * sqrt(problem.ts[end] - problem.ts[1]) ), length(problem.ψₒ) )
	end
	N = problem.ψᵢ.state.N_cutoff
	fill(N, length(problem.ψₒ))
end

function _expectedOutputPhotons(problem::_ScatterProblem{O, S, WavePacket{S}}) where {O, S <: Displaced}
	if problem.displace_output
		return fill(
			roblem.ψᵢ.state.N_cutoff + ceil(Int, 3 * sqrt(problem.ts[end] - problem.ts[1]) ),
			length(problem.ψₒ)
		)
	end
	N = problem.ψᵢ.state.N_cutoff + coherent_cutoff(problem.ψᵢ.state.α)
	fill(N, length(problem.ψₒ))
end

function _expectedOutputPhotons(problem::_ScatterProblem{O, S, WavePacket{S}}) where {O, S}
	N = problem.ψᵢ.state.N_cutoff
	fill(N, length(problem.ψₒ))
end

# ------------------------------------------------------------
# _DrivenProblems

function _generateLindbladian(problem::_DrivenProblem{O, Coherent, ContinuousWave}) where {O}
	# ----------------------------------------
	# Scatter System
	H_sys, σ⁻, σ⁺ = problem.H, problem.σ, problem.σ'
	σ⁺σ⁻ = σ⁺ * σ⁻

	# ----------------------------------------
	# Constructing Disspator and Hamiltonian
	α  = problem.ψᵢ.α
	atomic_dissipation = isempty(problem.Ls) ?
		-0.5im * σ⁺σ⁻ :
		-0.5im * σ⁺σ⁻ - 0.5im * sum([L'*L for L ∈ problem.Ls])

	H_nh = atomic_dissipation - 1.0im * (σ⁺ * α - σ⁻ * α) + H_sys

	ψ₀ = problem.system_state
	return ψ₀, [σ⁻, problem.Ls...], H_nh
end

function _generateLindbladian(problem::_DrivenProblem{O, Coherent, WavePacket{Coherent}}, _) where {O}
	# ----------------------------------------
	# Scatter System
	H_sys, σ⁻, σ⁺ = problem.H, problem.σ, problem.σ'
	σ⁺σ⁻ = σ⁺ * σ⁻

	# ----------------------------------------
	# Constructing Disspator and Hamiltonian
	α  = problem.ψᵢ.state.α
	mf = problem.ψᵢ.mode.modeFunction
	atomic_dissipation = isempty(problem.Ls) ?
		-0.5im * σ⁺σ⁻ :
		-0.5im * σ⁺σ⁻ - 0.5im * sum([L'*L for L ∈ problem.Ls])

	H_nh(t) = atomic_dissipation - 1.0im * (σ⁺ * α * mf(t) - σ⁻ * α * mf(t)) + H_sys

	ψ₀ = problem.system_state
	return ψ₀, [t -> σ⁻, problem.Ls...], H_nh
end

function _generateLindbladian(problem::_DrivenProblem{O, T, WavePacket{T}}, _) where {O, T}
	# ----------------------------------------
	# Scatter System
	H_sys, σ⁻, σ⁺ = problem.H, problem.σ, problem.σ'
	σ⁺σ⁻ = σ⁺ * σ⁻
	idₛ = one(basis(H_sys))

	# ----------------------------------------
	# Input System
	inBasis = FockBasis(problem.ψᵢ.state.N_cutoff)
	gᵢ      = problem.ψᵢ.mode.gᵢ
	idᵢ     = identityoperator(inBasis)
	aᵢ, aᵢᵀ = destroy(inBasis) ⊗ idₛ, create(inBasis) ⊗ idₛ
	aᵢᵀaᵢ   = aᵢᵀ * aᵢ

	# ----------------------------------------
	# Embed Operators and abbreviations
	H_sys, σ⁻, σ⁺, σ⁺σ⁻ = idᵢ ⊗ H_sys, idᵢ ⊗ σ⁻, idᵢ ⊗ σ⁺, idᵢ ⊗ σ⁺σ⁻
	aᵢᵀσ⁻,  aᵢσ⁺ = aᵢᵀ * σ⁻, aᵢ * σ⁺

	# ----------------------------------------
	# Constructing Disspator and Hamiltonian
	L  = t -> σ⁻ + gᵢ(t)aᵢ
	Lᵀ = t -> σ⁺ + gᵢ(t)aᵢᵀ

	atomic_dissipation = isempty(problem.Ls) ?
		-0.5im * σ⁺σ⁻ :
		-0.5im * σ⁺σ⁻ - (0.5im * idᵢ) ⊗ sum([L'*L for L ∈ problem.Ls])

	H_nh = _createHamiltonian(problem, atomic_dissipation, aᵢᵀaᵢ, aᵢσ⁺, σ⁺, σ⁻, H_sys)

	ψ₀ = Ket(inBasis, createState(problem.ψᵢ.state)) ⊗ problem.system_state
	return ψ₀, [L, [idᵢ ⊗ L for L ∈ problem.Ls]...], H_nh
end

function _createHamiltonian(problem::_DrivenProblem{O, S, WavePacket{S}},
		atomic_dissipation, aᵢᵀaᵢ, aᵢσ⁺, σ⁺, σ⁻, H_sys) where {O, S <: Displaced}
	gᵢ, mf, α = problem.ψᵢ.mode.gᵢ, problem.ψᵢ.mode.modeFunction, problem.ψᵢ.state.α

	return t -> atomic_dissipation - 0.5im * gᵢ(t)^2 * aᵢᵀaᵢ -
		1.0im * (σ⁺ * α * mf(t) - σ⁻ * α' * mf(t)') - 1.0im * gᵢ(t)aᵢσ⁺ + H_sys
end

function _createHamiltonian(problem::_DrivenProblem{O, S, WavePacket{S}},
		atomic_dissipation, aᵢᵀaᵢ, aᵢσ⁺, σ⁺, σ⁻, H_sys) where {O, S <: NonDisplaced}
	gᵢ = problem.ψᵢ.mode.gᵢ

	return t -> atomic_dissipation - 0.5im * gᵢ(t)^2 * aᵢᵀaᵢ - 1.0im * gᵢ(t)aᵢσ⁺ + H_sys
end

# ------------------------------------------------------------
# _ScatterProblems

function _generateLindbladian(problem::_ScatterProblem{O, Coherent, ContinuousWave}, Nouts) where {O}
	_generateLindbladian(
		WaveguideProblem(
			[problem.H, problem.Ls, problem.σ, problem.system_state],
			ContinuousWave(t -> problem.ψᵢ.α),
			problem.ψₒ,
			problem.ts,
			problem.displace_output
		),
		Nouts
	)
end

function _generateLindbladian(problem::_ScatterProblem{O, Coherent, WavePacket{Coherent}}, Nouts) where {O}
	# ----------------------------------------
	# Output System
	outBasis, groundstateOutput, aₒs, aₒᵀs = _generateOutputCavities(Nouts)
	idₒ = one(outBasis)

	# ----------------------------------------
	# Scatter System
	H_sys, σ⁻, σ⁺ = problem.H ⊗ idₒ, problem.σ ⊗ idₒ, problem.σ' ⊗ idₒ
	σ⁺σ⁻ = σ⁺ * σ⁻
	idₛ = one(basis(problem.H))

	# ----------------------------------------
	# Embed Operators and abbreviations
	aₒs, aₒᵀs = [idₛ ⊗ a for a ∈ aₒs], [idₛ ⊗ aᵀ for aᵀ ∈ aₒᵀs]

	gₒs = [mode.gₒ for mode in problem.ψₒ]

	gₒaₒs(t)  = sum([gₒ(t) for gₒ ∈ gₒs] .* aₒs)
	gₒaₒᵀs(t) = sum([gₒ(t) for gₒ ∈ gₒs] .* aₒᵀs)

	aₒᵀσ⁻s = [aₒᵀ * σ⁻ for aₒᵀ in aₒᵀs]
	aₒᵀaₒs = aₒᵀs .* aₒs
	# TODO: maybe it is faster to do `gₒaₒs(t)  = idₛ ⊗ sum(gₒs(t) .* aₒs)` and leave aₒs unchanged
	# due to the sum

	# ----------------------------------------
	# Constructing Disspator and Hamiltonian
	L  = t -> (σ⁻ + gₒaₒs(t) )
	Lᵀ = t -> (σ⁺ + gₒaₒᵀs(t) )

	α  = problem.ψᵢ.state.α
	mf = problem.ψᵢ.mode.modeFunction
	atomic_dissipation = isempty(problem.Ls) ?
		-0.5im * σ⁺σ⁻ :
		-0.5im * σ⁺σ⁻ - 0.5im * sum([L'*L for L ∈ problem.Ls]) ⊗ idₒ

	if problem.displace_output
		H_nh = function(t)
			return atomic_dissipation - 0.5im * sum([gₒ(t)^2 for gₒ ∈ gₒs] .* aₒᵀaₒs) -
				1.0im * sum([gₒ(t) for gₒ ∈ gₒs] .* aₒᵀσ⁻s) -
				1.0im * (σ⁺ * α * mf(t) - σ⁻ * α' * mf(t)) +
				H_sys
		end
	else
		H_nh = function(t)
			return atomic_dissipation - 0.5im * sum([gₒ(t)^2 for gₒ ∈ gₒs] .* aₒᵀaₒs) -
				1.0im * sum([gₒ(t) for gₒ ∈ gₒs] .* aₒᵀσ⁻s) -
				1.0im * (Lᵀ(t) * α * mf(t) - L(t) * α' * mf(t)) +
				H_sys
		end
	end

	ψ₀ = problem.system_state ⊗ groundstateOutput
	return ψ₀, [t -> L(t), [L ⊗ idₒ for L ∈ problem.Ls]...], H_nh
end

function _generateLindbladian(problem::_ScatterProblem{O, T, WavePacket{T}}, Nouts) where {O, T}
	# ----------------------------------------
	# Output System
	outBasis, groundstateOutput, aₒs, aₒᵀs = _generateOutputCavities(Nouts)
	idₒ = one(outBasis)

	# ----------------------------------------
	# Scatter System
	H_sys, σ⁻, σ⁺  = problem.H, problem.σ, problem.σ'
	σ⁺σ⁻ = σ⁺ * σ⁻
	idₛ = one(basis(problem.H))

	# ----------------------------------------
	# Input System
	inBasis = FockBasis(problem.ψᵢ.state.N_cutoff)
	gᵢ      = problem.ψᵢ.mode.gᵢ
	idᵢ     = identityoperator(inBasis)
	aᵢ, aᵢᵀ = destroy(inBasis) ⊗ idₛ ⊗ idₒ, create(inBasis) ⊗ idₛ ⊗ idₒ
	aᵢᵀaᵢ = aᵢᵀ * aᵢ

	# ----------------------------------------
	# Embed Operators and abbreviations
	H_sys, σ⁻, σ⁺, σ⁺σ⁻ = idᵢ ⊗ H_sys ⊗ idₒ, idᵢ ⊗ σ⁻ ⊗ idₒ, idᵢ ⊗ σ⁺ ⊗ idₒ, idᵢ ⊗ σ⁺σ⁻ ⊗ idₒ
	aₒs, aₒᵀs = [idᵢ ⊗ idₛ ⊗ a for a ∈ aₒs], [idᵢ ⊗ idₛ ⊗ aᵀ for aᵀ ∈ aₒᵀs]

	gₒs = [mode.gₒ for mode in problem.ψₒ]

	gₒaₒs(t)  = sum([gₒ(t)' for gₒ ∈ gₒs] .* aₒs)
	gₒaₒᵀs(t) = sum([gₒ(t)  for gₒ ∈ gₒs] .* aₒᵀs)

	aᵢσ⁺   = aᵢ * σ⁺
	aₒᵀaᵢs = [aₒᵀ * aᵢ for aₒᵀ in aₒᵀs]
	aₒᵀσ⁻s = [aₒᵀ * σ⁻ for aₒᵀ in aₒᵀs]
	aₒᵀaₒs = aₒᵀs .* aₒs
	# TODO: maybe it is faster to do `gₒaₒs(t)  = idₛ ⊗ sum(gₒs(t) .* aₒs)` and leave aₒs unchanged
	# due to the sum

	# ----------------------------------------
	# Constructing Disspator and Hamiltonian
	L  = t -> σ⁻ + gᵢ(t)aᵢ  + gₒaₒs(t)

	atomic_dissipation = isempty(problem.Ls) ?
		-0.5im * σ⁺σ⁻ :
		-0.5im * σ⁺σ⁻ - (0.5im * idᵢ) ⊗ sum([L'*L for L ∈ problem.Ls]) ⊗ idₒ

	H_nh = _createHamiltonian(problem, atomic_dissipation,
		aₒᵀaₒs, aᵢᵀaᵢ, aᵢσ⁺, aₒᵀaᵢs, aₒᵀσ⁻s, σ⁺, σ⁻, gₒaₒᵀs, gₒaₒs, H_sys
	)

	ψ₀ = Ket(inBasis, createState(problem.ψᵢ.state)) ⊗ problem.system_state ⊗ groundstateOutput
	return ψ₀, [t -> L(t), [idᵢ ⊗ L ⊗ idₒ for L ∈ problem.Ls]...], H_nh
end

function _generateOutputCavities(Nouts)
	# Preconverting into `CompositeBasis` allows us to `embed` into it, even for a single system
	outBasis  = reduce(⊗, FockBasis.(Nouts) .|> CompositeBasis)

	groundstateOutput = reduce(⊗, [fockstate(FockBasis(N), 0) for N in Nouts])

	aₒs  = [embed(outBasis, i, destroy(FockBasis(N))) for (i, N) in enumerate(Nouts)]
	aₒᵀs = [embed(outBasis, i, create(FockBasis(N)))  for (i, N) in enumerate(Nouts)]

	return outBasis, groundstateOutput, aₒs, aₒᵀs
end

function _createHamiltonian(problem::_ScatterProblem{O, S, WavePacket{S}},
		atomic_dissipation, aₒᵀaₒs, aᵢᵀaᵢ, aᵢσ⁺, aₒᵀaᵢs, aₒᵀσ⁻s, σ⁺, σ⁻, gₒaₒᵀs, gₒaₒs,
		H_sys) where {O, S<: Displaced}
	gᵢ, mf, α = problem.ψᵢ.mode.gᵢ, problem.ψᵢ.mode.modeFunction, problem.ψᵢ.state.α
	gₒs   = [mode.gₒ for mode in problem.ψₒ]

	return t -> atomic_dissipation -
		0.5im * (gᵢ(t)^2 * aᵢᵀaᵢ + sum([gₒ(t)^2 for gₒ ∈ gₒs] .* aₒᵀaₒs) ) -
		1.0im * (gᵢ(t)aᵢσ⁺ + sum([gₒ(t) for gₒ ∈ gₒs] .* aₒᵀσ⁻s) +
			sum(gᵢ(t) .* [gₒ(t) for gₒ ∈ gₒs] .* aₒᵀaᵢs) ) -
		1.0im * (( σ⁺ + gₒaₒᵀs(t) ) * α * mf(t) - ( σ⁻ + gₒaₒs(t) )  * α' * mf(t)') + H_sys
end

function _createHamiltonian(problem::_ScatterProblem{O, S, WavePacket{S}},
		atomic_dissipation, aₒᵀaₒs, aᵢᵀaᵢ, aᵢσ⁺, aₒᵀaᵢs, aₒᵀσ⁻s, σ⁺, σ⁻, gₒaₒᵀs, gₒaₒs,
		H_sys) where {O, S<: NonDisplaced}
	gᵢ    = problem.ψᵢ.mode.gᵢ
	gₒs   = [mode.gₒ for mode in problem.ψₒ]

	return t -> atomic_dissipation -
		0.5im * (gᵢ(t)^2 * aᵢᵀaᵢ + sum([gₒ(t)^2 for gₒ ∈ gₒs] .* aₒᵀaₒs) ) -
		1.0im * (gᵢ(t)aᵢσ⁺ + sum([gₒ(t) for gₒ ∈ gₒs] .* aₒᵀσ⁻s) +
			sum(gᵢ(t) .* [gₒ(t) for gₒ ∈ gₒs] .* aₒᵀaᵢs) ) +
		H_sys
end