text:
# struct TileDescription
    # vectorized::Symbol
    # u₁loop::Symbol
    # u₂loop::Symbol
    # U::Int
    # T::Int
# end

function lower!(
    q::Expr, op::Operation, ls::LoopSet, unrollargs::UnrollArgs, mask::Union{Nothing,Symbol,Unsigned}, ::Nothing
)
    if isconstant(op)
        zerotyp = zerotype(ls, op)
        if zerotyp == INVALID
            lower_constant!(q, op, ls, unrollargs)
        else
            lower_zero!(q, op, ls, unrollargs, zerotyp)
        end
    elseif isload(op)
        lower_load!(q, op, ls, unrollargs, mask)
    elseif iscompute(op)
        lower_compute!(q, op, ls, unrollargs, mask)
    elseif isstore(op)
        lower_store!(q, ls, op, unrollargs, mask)
    # elseif isloopvalue(op)
    end
end
function lower!(
    q::Expr, op::Operation, ls::LoopSet, unrollargs::UnrollArgs, mask::Union{Nothing,Symbol,Unsigned}, filterstore::Bool
)
    if filterstore
        if isstore(op)
            lower_store!(q, ls, op, unrollargs, mask)
        end
    else
        if isconstant(op)
            zerotyp = zerotype(ls, op)
            if zerotyp == INVALID
                lower_constant!(q, op, ls, unrollargs)
            else
                lower_zero!(q, op, ls, unrollargs, zerotyp)
            end
        elseif isload(op)
            lower_load!(q, op, ls, unrollargs, mask)
        elseif iscompute(op)
            lower_compute!(q, op, ls, unrollargs, mask)
        end
    end
end
function lower!(
    q::Expr, ops::AbstractVector{Operation}, ls::LoopSet, unrollsyms::UnrollSymbols, u₁::Int, u₂::Int,
    suffix::Union{Nothing,Int}, mask::Union{Nothing,Symbol,Unsigned}, filterstore = nothing
)
    foreach(op -> lower!(q, op, ls, UnrollArgs(u₁, unrollsyms, u₂, suffix), mask, filterstore), ops)
end

function lower_block(ls::LoopSet, us::UnrollSpecification, n::Int, inclmask::Bool, UF)
    if inclmask
        lower_block(ls, us, n, Symbol("##mask##"), UF)
    else
        lower_block(ls, us, n, nothing, UF)
    end
end

function lower_block(
    ls::LoopSet, us::UnrollSpecification, n::Int, mask::Union{Nothing,Symbol}, UF::Int
)
    @unpack u₁loopnum, u₂loopnum, vectorizedloopnum, u₁, u₂ = us
    ops = oporder(ls)
    order = names(ls)
    u₁loopsym = order[u₁loopnum]
    u₂loopsym = order[u₂loopnum]
    vectorized = order[vectorizedloopnum]
    unrollsyms = UnrollSymbols(u₁loopsym, u₂loopsym, vectorized)
    u₁ = n == u₁loopnum ? UF : u₁
    dontmaskfirsttiles = !isnothing(mask) && vectorizedloopnum == u₂loopnum
    blockq = Expr(:block)
    delay_u₁ = true
    # delay_u₁ = false
    dontdelayat = nothing#4
    for prepost ∈ 1:2
        # !u₁ && !u₂
        lower!(blockq, ops[1,1,prepost,n], ls, unrollsyms, u₁, u₂, nothing, mask)
        if !delay_u₁ || u₁ == dontdelayat
            lower!(blockq, ops[2,1,prepost,n], ls, unrollsyms, u₁, u₂, nothing, mask)
        end
        opsv1 = ops[1,2,prepost,n]
        opsv2 = ops[2,2,prepost,n]
        if length(opsv1) + length(opsv2) > 0
            # if u₁ == 3
                # lower!(blockq, ops[2,1,prepost,n], vectorized, ls, u₁loop, u₂loop, u₁, u₂, nothing, mask)
            # end
            for store ∈ (false,true)
                # let store = nothing
                nstores = 0
                iszero(length(opsv1)) || (nstores += sum(isstore, opsv1))
                iszero(length(opsv2)) || (nstores += sum(isstore, opsv2))
                if delay_u₁ && !store && length(opsv1) + length(opsv2) == nstores
                    u₁ != dontdelayat && lower!(blockq, ops[2,1,prepost,n], ls, unrollsyms, u₁, u₂, nothing, mask) # for u ∈ 0:u₁-1
                    continue
                end
                for t ∈ 0:u₂-1
                    # if t == 0
                    #     push!(blockq.args, Expr(:(=), u₂loop, tiledsym(u₂loop)))
                    # elseif u₂loopnum == vectorizedloopnum
                    #     push!(blockq.args, Expr(:(=), u₂loop, Expr(:call, lv(:vadd_fast), VECTORWIDTHSYMBOL, staticexpr(u₂loop))))
                    # else
                    #     push!(blockq.args, Expr(:+=, u₂loop, 1))
                    # end
                    if dontmaskfirsttiles && t < u₂ - 1 # !u₁ &&  u₂
                        lower!(blockq, opsv1, ls, unrollsyms, u₁, u₂, t, nothing, store)
                    else # !u₁ &&  u₂
                        lower!(blockq, opsv1, ls, unrollsyms, u₁, u₂, t, mask, store)
                    end
                    if delay_u₁ && iszero(t) && !store && u₁ != dontdelayat #  u₁ && !u₂
                        # for u ∈ 0:u₁-1
                        lower!(blockq, ops[2,1,prepost,n], ls, unrollsyms, u₁, u₂, nothing, mask)
                        # end
                    end
                    if dontmaskfirsttiles && t < u₂ - 1 #  u₁ && u₂
                        # for u ∈ 0:u₁-1
                        lower!(blockq, opsv2, ls, unrollsyms, u₁, u₂, t, nothing, store)
                        # end
                    else #  u₁ && u₂
                        # for u ∈ 0:u₁-1
                        lower!(blockq, opsv2, ls, unrollsyms, u₁, u₂, t, mask, store)
                        # end
                    end
                end
                nstores == 0 && break
            end
        elseif delay_u₁ && u₁ != dontdelayat
            # for u ∈ 0:u₁-1     #  u₁ && !u₂
            lower!(blockq, ops[2,1,prepost,n], ls, unrollsyms, u₁, u₂, nothing, mask)
            # end
        end
        if n > 1 && prepost == 1
            push!(blockq.args, lower_unrolled_dynamic(ls, us, n-1, !isnothing(mask)))
        end
    end
    # loopsym = mangletiledsym(order[n], us, n)
    loopsym = order[n]
    # push!(blockq.args, incrementloopcounter(us, n, loopsym, UF))
    # if n > 1 || iszero(ls.align_loops[])
    push!(blockq.args, incrementloopcounter(ls, us, n, UF))
    # else
    #     loopsym = names(ls)[n]
    #     push!(blockq.args, Expr(:(=), loopsym, Expr(:call, lv(:vadd_fast), loopsym, Symbol("##ALIGNMENT#STEP##"))))
    # end
    blockq
end

# function lower_llvm_unroll(ls::LoopSet, us::UnrollSpecification, n::Int, loop::Loop)
#     loopsym = names(ls)[n]
#     loop = getloop(ls, loopsym)
#     # loopsym = mangletiledsym(loopsym, us, n)
#     nisvectorized = false#isvectorized(us, n)
#     sl = startloop(loop, nisvectorized, loopsym)
#     # tc = terminatecondition(loop, us, n, loopsym, inclmask, 1)
#     looprange = if loop.startexact
#         if loop.stopexact
#             Expr(:(=), loopsym, Expr(:call, :(:), loop.starthint-1, loop.stophint-1))
#         else
#             # expectation = expect(Expr(:call, :(>), loop.stopsym, loop.starthint-5))
#             Expr(:(=), loopsym, Expr(:call, :(:), loop.starthint-1, Expr(:call, lv(:staticm1), loop.stopsym)))
#         end
#     elseif loop.stopexact
#         # expectation = expect(Expr(:call, :(>), loop.stophint+5, loop.startsym))
#         Expr(:(=), loopsym, Expr(:call, :(:), Expr(:call, lv(:staticm1), loop.startsym), loop.stophint - 1))
#     else
#         # expectation = expect(Expr(:call, :(>), loop.stopsym, Expr(:call, lv(:vsub_fast), loop.startsym, 5)))
#         Expr(:(=), loopsym, Expr(:call, :(:), Expr(:call, lv(:staticm1), loop.startsym), Expr(:call, lv(:staticm1), loop.stopsym)))
#     end
#     body = lower_block(ls, us, n, false, 1)
#     push!(body.args, Expr(:loopinfo, (Symbol("llvm.loop.unroll.count"), 4)))
#     q = Expr(:for, looprange, body)
#     # if nisvectorized
#     #     tc = terminatecondition(loop, us, n, loopsym, true, 1)
#     #     body = lower_block(ls, us, n, true, 1)
#     #     push!(q.args, Expr(:if, tc, body))
#     # end
#     if loop.startexact && loop.stopexact
#         q
#     else
#         # Expr(:block, assumption, expectation, q)
#         q
#         Expr(:block, loopiteratesatleastonce(loop), q)
#     end
# end
# function lower_unroll_for_throughput(ls::LoopSet, us::UnrollSpecification, loop::Loop, loopsym::Symbol)
#     UF = 4
#     sl = startloop(ls, us, 1, UF)
#     tcc = terminatecondition(ls, us, 1, false, 1)
#     tcu = terminatecondition(ls, us, 1, false, UF)
#     body = lower_block(ls, us, 1, false, 1)
#     loopisstatic = isstaticloop(loop)
#     tcu = loopisstatic ? tcu : expect(tcu)
#     termcondu = gensym(:maybetermu)
#     unrolledbody = Expr(:block)
#     foreach(_ -> push!(unrolledbody.args, body), 1:UF)
#     # q = Expr(
#     #     :block,
#     #     Expr(:while, tcu, unrolledbody),
#     #     Expr(:while, tcc, body)
#     # )
#     # return Expr(:let, sl, q)
#     push!(unrolledbody.args, Expr(:(=), termcondu, tcu))
#     unrolledloop = Expr(
#         :block,
#         Expr(:while, termcondu, unrolledbody),
#         Expr(:while, tcc, body)
#     )
#     termcond = gensym(:maybeterm)
#     singleloop = Expr(
#         :block,
#         Expr(:(=), termcond, true),
#         Expr(:while, termcond, Expr(:block, body, Expr(:(=), termcond, tcc)))
#     )
#     q = Expr(
#         :block,
#         assume(tcc),
#         Expr(:(=), termcondu, tcu),
#         Expr(:if, termcondu, unrolledloop, singleloop)
#     )
#     Expr(:let, sl, q)
# end

function assume(ex)
    Expr(:call, Expr(:(.), Expr(:(.), :LoopVectorization, QuoteNode(:VectorizationBase)), QuoteNode(:assume)), ex)
end
function expect(ex)
    # use_expect() || return ex
    # Expr(:call, Expr(:(.), Expr(:(.), :LoopVectorization, QuoteNode(:VectorizationBase)), QuoteNode(:expect)), ex)
    ex
end
function loopiteratesatleastonce(loop::Loop, as::Bool = true)
    comp = if loop.startexact # requires !loop.stopexact
        Expr(:call, :>, loop.stopsym, loop.starthint - 1)
    elseif loop.stopexact # requires !loop.startexact
        Expr(:call, :>, loop.stopexact + 1, loop.startsym)
    else
        Expr(:call, :>, loop.stopsym, Expr(:call, lv(:vsub_fast), loop.startsym, 1))
    end
    # as ? assume(comp) : expect(comp)
    assume(comp)
end
# @inline step_to_align(x, ::Val{W}) where {W} = step_to_align(pointer(x), Val{W}())
# @inline step_to_align(x::Ptr{T}, ::Val{W}) where {W,T} = vsub_fast(W, reinterpret(Int, x) & (W - 1))
# function align_inner_loop_expr(ls::LoopSet, us::UnrollSpecification, loop::Loop)
#     alignincr = Symbol("##ALIGNMENT#STEP##")
#     looplength = gensym(:inner_loop_length)
#     pushpreamble!(ls, Expr(:(=), looplength, looplengthexpr(loop)))
#     vp = vptr(operations(ls)[ls.align_loops[]])
#     align_step = Expr(:call, :min, Expr(:call, lv(:step_to_align), vp, VECTORWIDTHSYMBOL), looplength)
#     Expr(
#         :block,
#         Expr(:(=), alignincr, align_step),
#         maskexpr(alignincr),
#         lower_block(ls, us, 1, true, 1)
#     )
# end

function check_full_conv_kernel(ls, us, N)
    loop₁ = getloop(ls, names(ls)[us.u₁loopnum])
    (isstaticloop(loop₁) && length(loop₁) == us.u₁) && return true
    loop₂ = getloop(ls, names(ls)[us.u₂loopnum])
    (isstaticloop(loop₂) && length(loop₂) == us.u₂) && return true
    false
end
function allinteriorunrolled(ls::LoopSet, us::UnrollSpecification, N)
    if ls.loadelimination[]
        check_full_conv_kernel(ls, us, N) || return false
    end
    unroll_total = 1
    for n ∈ 1:N-1
        loop = getloop(ls, names(ls)[n])
        nisvectorized = isvectorized(us, n)
        W = nisvectorized ? ls.vector_width[] : 1
        ((length(loop) ≤ 8W) && (isstaticloop(loop) & (!iszero(W)))) || return false
        unroll_total *= cld(length(loop),W)
    end
    if us.u₁loopnum > N
        unroll_total *= us.u₁
    end
    if us.u₂loopnum > N
        unroll_total *= us.u₂
    end
    unroll_total ≤ 16
end

function lower_no_unroll(ls::LoopSet, us::UnrollSpecification, n::Int, inclmask::Bool, initialize::Bool = true, maxiters::Int=-1)
    usorig = ls.unrollspecification[]
    nisvectorized = isvectorized(us, n)
    loopsym = names(ls)[n]
    loop = getloop(ls, loopsym)
    # if !nisvectorized && !inclmask && isone(n) && !iszero(ls.lssm[].terminators[1]) && !ls.loadelimination[] && (us.u₁ > 1) && (usorig.u₁ == us.u₁) && (usorig.u₂ == us.u₂) && length(loop) > 15
    #     return lower_unroll_for_throughput(ls, us, loop, loopsym)
    #     # return lower_llvm_unroll(ls, us, n, loop)
    # end
    # sl = startloop(loop, nisvectorized, loopsym)

    tc = terminatecondition(ls, us, n, inclmask, 1)
    body = lower_block(ls, us, n, inclmask, 1)
    # align_loop = isone(n) & (ls.align_loops[] > 0)
    loopisstatic = isstaticloop(loop)
    if !loopisstatic && (usorig.u₁ == us.u₁) && (usorig.u₂ == us.u₂) && !inclmask
        tc = expect(tc)
    end
    W = nisvectorized ? ls.vector_width[] : 1
    loopisstatic &= (!iszero(W))
    # q = if align_loop
    #     Expr(:block, align_inner_loop_expr(ls, us, loop), Expr(:while, tc, body))
    # elseif nisvectorized
    if loopisstatic && (isone(length(loop) ÷ W) || (n ≤ 3 && length(loop) ≤ 8W && allinteriorunrolled(ls, us, n)))
        q = Expr(:block)
        foreach(_ -> push!(q.args, body), 1:(length(loop) ÷ W))
    elseif nisvectorized
            # Expr(:block, loopiteratesatleastonce(loop, true), Expr(:while, expect(tc), body))
        q = Expr(:block, Expr(maxiters == 1 ? :if : :while, tc, body))
    else
        termcond = gensym(:maybeterm)
        push!(body.args, Expr(:(=), termcond, tc))
        q = Expr(:block, Expr(:(=), termcond, true), Expr(maxiters == 1 ? :if : :while, termcond, body))
        # Expr(:block, Expr(:while, expect(tc), body))
        # Expr(:block, assume(tc), Expr(:while, tc, body))
        # push!(body.args, Expr(:&&, expect(Expr(:call, :!, tc)), Expr(:break)))
        # Expr(:block, assume(tc), Expr(:while, true, body))
        # push!(body.args, Expr(:||, expect(tc), Expr(:break)))
        # Expr(:block, Expr(:while, true, body))
    end
    if nisvectorized && !(loopisstatic && iszero(length(loop) & (W - 1)))
        # tc = terminatecondition(loop, us, n, loopsym, true, 1)
        body = lower_block(ls, us, n, true, 1)
        if isone(num_loops(ls))
            pushfirst!(body.args, definemask(loop))
        # elseif align_loop
        #     pushfirst!(body.args, definemask_for_alignment_cleanup(loop))
        end
        if loopisstatic
            push!(q.args, body)
        else
            tc = terminatecondition(ls, us, n, true, 1)
            push!(q.args, Expr(:if, tc, body))
        end
    end
    if initialize
        Expr(:let, startloop(ls, us, n), q)
    else
        q
    end
end
function lower_unrolled_dynamic(ls::LoopSet, us::UnrollSpecification, n::Int, inclmask::Bool)
    UF = unrollfactor(us, n)
    isone(UF) && return lower_no_unroll(ls, us, n, inclmask)
    @unpack u₁loopnum, vectorizedloopnum, u₁, u₂ = us
    order = names(ls)
    loopsym = order[n]
    loop = getloop(ls, loopsym)
    vectorized = order[vectorizedloopnum]
    nisunrolled = isunrolled1(us, n)
    nisvectorized = isvectorized(us, n)
    W = nisvectorized ? ls.vector_width[] : 1
    loopisstatic = isstaticloop(loop) & (!iszero(W))
    UFW = UF * W
    looplength = length(loop)
    if loopisstatic & (UFW > looplength)
        UFWnew = cld(looplength, cld(looplength, UFW))
        UF = cld(UFWnew, W)
        UFW = UF * W
        us = nisunrolled ? UnrollSpecification(us, UF, u₂) : UnrollSpecification(us, u₁, UF)
    end
    remmask = inclmask | nisvectorized
    Ureduct = (n == num_loops(ls) && (u₂ == -1)) ? calc_Ureduct(ls, us) : -1
    # sl = startloop(loop, nisvectorized, loopsym)
    sl = startloop(ls, us, n)
    UFt = loopisstatic ? cld(looplength % UFW, W) : 1
    # Don't place remainder first if we're going to have to mask this loop (i.e., if this loop is vectorized)
    remfirst = loopisstatic & (!nisvectorized) & (UFt > 0) & !(unsigned(Ureduct) < unsigned(UF))
    tc = terminatecondition(ls, us, n, inclmask, remfirst ? 1 : UF)
    # usorig = ls.unrollspecification[]
    # tc = (usorig.u₁ == us.u₁) && (usorig.u₂ == us.u₂) && !loopisstatic && !inclmask && !ls.loadelimination[] ? expect(tc) : tc
    body = lower_block(ls, us, n, inclmask, UF)
    if loopisstatic
        iters = length(loop) ÷ UFW
        if (iters ≤ 1) || (iters*UF ≤ 16 && allinteriorunrolled(ls, us, n))# Let's set a limit on total unrolling
            q = Expr(:block)
            foreach(_ -> push!(q.args, body), 1:iters)
        else
            q = Expr(:while, tc, body)
        end
        remblock = Expr(:block)
        (nisvectorized && (UFt > 0) && isone(num_loops(ls))) && push!(remblock.args, definemask(loop))
        unroll_cleanup = true
    else
        remblock = init_remblock(loop, ls.lssm[], n)#loopsym)
        # unroll_cleanup = Ureduct > 0 || (nisunrolled ? (u₂ > 1) : (u₁ > 1))
        # remblock = unroll_cleanup ? init_remblock(loop, ls.lssm[], n)#loopsym) : Expr(:block)
        q = Expr(:while, tc, body)
    end
    q = if unsigned(Ureduct) < unsigned(UF) # unsigned(-1) == typemax(UInt); is logic relying on twos-complement bad?
        add_cleanup = true
        if isone(Ureduct)
            UF_cleanup = 1
            if nisvectorized
                blockhead = :while
            else
                blockhead = if UF == 2
                    if loopisstatic
                        add_cleanup = UFt == 1
                        :block
                    else
                        :if
                    end
                else
                    :while
                end
                UFt = 0
            end
        elseif 2Ureduct < UF
            UF_cleanup = 2
            blockhead = :while
        else
            UF_cleanup = UF - Ureduct
            blockhead = :if
        end
        _q = Expr(:block, add_upper_outer_reductions(ls, q, Ureduct, UF, loop, vectorized, nisvectorized))
        if add_cleanup
            cleanup_expr = Expr(blockhead)
            blockhead === :block || push!(cleanup_expr.args, terminatecondition(ls, us, n, inclmask, UF_cleanup))
            us_cleanup = nisunrolled ? UnrollSpecification(us, UF_cleanup, u₂) : UnrollSpecification(us, u₁, UF_cleanup)
            push!(cleanup_expr.args, lower_block(ls, us_cleanup, n, inclmask, UF_cleanup))
            push!(_q.args, cleanup_expr)
        end
        UFt > 0 && push!(_q.args, remblock)
        _q
    elseif remfirst
        numiters = length(loop) ÷ UF
        if numiters > 2
            Expr( :block, remblock, q )
        else
            q = Expr(:block, remblock)
            for i ∈ 1:numiters
                push!(q.args, body)
            end
            q
        end
    elseif iszero(UFt)
        Expr( :block, q )
    elseif !nisvectorized && !loopisstatic && UF ≥ 10
        rem_uf = UF - 1
        UF = rem_uf >> 1
        UFt = rem_uf - UF
        ust = nisunrolled ? UnrollSpecification(us, UFt, u₂) : UnrollSpecification(us, u₁, UFt)
        newblock = lower_block(ls, ust, n, remmask, UFt)
        # comparison = unrollremcomparison(ls, loop, UFt, n, nisvectorized, remfirst)
        comparison = terminatecondition(ls, us, n, inclmask, UFt)
        UFt = 1
        UF += 1 - iseven(rem_uf)
        Expr( :block, q, Expr(iseven(rem_uf) ? :while : :if, comparison, newblock), remblock )
    else
        # if (usorig.u₁ == us.u₁) && (usorig.u₂ == us.u₂) && !isstaticloop(loop) && !inclmask# && !ls.loadelimination[]
        #     # Expr(:block, sl, assumeloopiteratesatleastonce(loop), Expr(:while, tc, body))
        #     Expr(:block, sl, expect(tc), q, remblock)
        # else
        #     Expr(:block, sl, q, remblock)
        # end
        Expr( :block, q, remblock )
    end
    if !iszero(UFt)
        # if unroll_cleanup
        while true
            ust = nisunrolled ? UnrollSpecification(us, UFt, u₂) : UnrollSpecification(us, u₁, UFt)
            newblock = lower_block(ls, ust, n, remmask, UFt)
            if (UFt ≥ UF - 1 + nisvectorized) || UFt == Ureduct || loopisstatic
                if isone(num_loops(ls)) && isone(UFt) && isone(Ureduct)
                    newblock = Expr(:block, definemask(loop), newblock)
                end
                push!(remblock.args, newblock)
                break
            end
            comparison = unrollremcomparison(ls, loop, UFt, n, nisvectorized, remfirst)
            if isone(num_loops(ls)) && isone(UFt)
                remblocknew = Expr(:if, comparison, newblock)
                push!(remblock.args, Expr(:block, definemask(loop), remblocknew))
                remblock = remblocknew
            else
                remblocknew = Expr(:elseif, comparison, newblock)
                push!(remblock.args, remblocknew)
                remblock = remblocknew
            end
            UFt += 1
        end
        # else
        #     ust = nisunrolled ? UnrollSpecification(us, 1, u₂) : UnrollSpecification(us, u₁, 1)
        #     # newblock = lower_block(ls, ust, n, remmask, 1)
        #     push!(remblock.args, lower_no_unroll(ls, ust, n, inclmask, false, UF-1))
        # end
    end
    Expr(:block, Expr(:let, sl, q))
end
function unrollremcomparison(ls::LoopSet, loop::Loop, UFt::Int, n::Int, nisvectorized::Bool, remfirst::Bool)
    termind = ls.lssm[].terminators[n]
    if iszero(termind)
        loopvarremcomparison(loop, UFt, nisvectorized, remfirst)
    else
        pointerremcomparison(ls, termind, UFt, n, nisvectorized, remfirst, loop)
    end
end
function loopvarremcomparison(loop::Loop, UFt::Int, nisvectorized::Bool, remfirst::Bool)
    loopsym = loop.itersymbol
    if nisvectorized
        itercount = if loop.stopexact
            Expr(:call, lv(:vsub_fast), loop.stophint - 1, Expr(:call, lv(:vmul_fast), VECTORWIDTHSYMBOL, UFt))
        else
            Expr(:call, lv(:vsub_fast), loop.stopsym, Expr(:call, lv(:vadd_fast), Expr(:call, lv(:vmul_fast), VECTORWIDTHSYMBOL, UFt), staticexpr(1)))
        end
        Expr(:call, :>, loopsym, itercount)
    elseif remfirst
        Expr(:call, :<, loopsym, loop.starthint + UFt - 1)
    elseif loop.stopexact
        Expr(:call, :>, loopsym, loop.stophint - UFt - 1)
    else
        Expr(:call, :>, loopsym, Expr(:call, lv(:vsub_fast), loop.stopsym, UFt + 1))
    end
end
function pointerremcomparison(ls::LoopSet, termind::Int, UFt::Int, n::Int, nisvectorized::Bool, remfirst::Bool, loop::Loop)
    lssm = ls.lssm[]
    termar = lssm.incrementedptrs[n][termind]
    ptrdef = lssm.incrementedptrs[n][termind]
    ptr = vptr(termar)
    ptrex = callpointerforcomparison(ptr)
    if remfirst
        Expr(:call, :<, ptrex, pointermax(ls, ptrdef, n, 1 - UFt, nisvectorized, loop.startexact))
    else
        # Expr(:call, :≥, ptrex, pointermax(ls, ptrdef, n, UFt, nisvectorized, loop))
        Expr(:call, :≥, ptrex, maxsym(ptr, UFt))
    end
end

function initialize_outer_reductions!(
    q::Expr, op::Operation, Umin::Int, Umax::Int, vectorized::Symbol, suffix::Union{Symbol,Nothing} = nothing
)
    reduct_zero = reduction_zero(op.instruction)
    isvectorized = vectorized ∈ reduceddependencies(op)
    # typeTr = Symbol("##TYPEOF##", name(op))
    typeTr = Expr(:call, :typeof, mangledvar(op))
    z = if isvectorized
        if reduct_zero === :zero
            Expr(:call, lv(:vzero), VECTORWIDTHSYMBOL, typeTr)
        else
            Expr(:call, lv(:vbroadcast), VECTORWIDTHSYMBOL, Expr(:call, reduct_zero, typeTr))
        end
    else
        Expr(:call, reduct_zero, typeTr)
    end
    mvar = variable_name(op, suffix)
    for u ∈ Umin:Umax-1
        push!(q.args, Expr(:(=), Symbol(mvar, u), z))
    end
    nothing
end
function initialize_outer_reductions!(q::Expr, ls::LoopSet, Umin::Int, Umax::Int, vectorized::Symbol, suffix::Union{Symbol,Nothing} = nothing)
    foreach(or -> initialize_outer_reductions!(q, ls.operations[or], Umin, Umax, vectorized, suffix), ls.outer_reductions)
end
function initialize_outer_reductions!(ls::LoopSet, Umin::Int, Umax::Int, vectorized::Symbol, suffix::Union{Symbol,Nothing} = nothing)
    initialize_outer_reductions!(ls.preamble, ls, Umin, Umax, vectorized, suffix)
end
function add_upper_comp_check(unrolledloop, loopbuffer)
    if isstaticloop(unrolledloop)
        Expr(:call, lv(:scalar_greaterequal), length(unrolledloop), loopbuffer)
    elseif unrolledloop.startexact
        if isone(unrolledloop.starthint)
            Expr(:call, lv(:scalar_greaterequal), unrolledloop.stopsym, loopbuffer)
        else
            Expr(:call, lv(:scalar_greaterequal), Expr(:call, lv(:vsub_fast), unrolledloop.stopsym, unrolledloop.starthint-1), loopbuffer)
        end
    elseif unrolledloop.stopexact
        Expr(:call, lv(:scalar_greaterequal), Expr(:call, lv(:vsub_fast), unrolledloop.stophint+1, unrolledloop.startsym), loopbuffer)
    else# both are given by symbols
        Expr(:call, lv(:scalar_greaterequal), Expr(:call, lv(:vsub_fast), unrolledloop.stopsym, Expr(:call,lv(:vsub_fast),unrolledloop.startsym, staticexpr(1))), loopbuffer)
    end
end
function add_upper_outer_reductions(ls::LoopSet, loopq::Expr, Ulow::Int, Uhigh::Int, unrolledloop::Loop, vectorized::Symbol, reductisvectorized::Bool)
    ifq = Expr(:block)
    initialize_outer_reductions!(ifq, ls, Ulow, Uhigh, vectorized)
    push!(ifq.args, loopq)
    _Ulow = Uhigh >>> 1; _Uhigh = Uhigh
    while _Ulow > Ulow
        reduce_range!(ifq, ls, _Ulow, _Uhigh)
        _Uhigh = _Ulow
        _Ulow >>>= 1
    end
    reduce_range!(ifq, ls, Ulow, _Uhigh)
    ncomparison = if reductisvectorized
        add_upper_comp_check(unrolledloop, Expr(:call, lv(:vmul_fast), VECTORWIDTHSYMBOL, Uhigh))
    else
        add_upper_comp_check(unrolledloop, Uhigh)
    end
    Expr(:if, ncomparison, ifq)
end
function reduce_expr!(q::Expr, ls::LoopSet, U::Int)
    us = ls.unrollspecification[]
    # u₁loop, u₂loop = getunrolled(ls)
    for or ∈ ls.outer_reductions
        op = ls.operations[or]
        var = name(op)
        mvar = mangledvar(op)
        instr = instruction(op)
        reduce_expr!(q, mvar, instr, U)
        if !iszero(length(ls.opdict))
            if (isu₁unrolled(op) | isu₂unrolled(op))
                push!(q.args, Expr(:(=), var, Expr(:call, lv(reduction_scalar_combine(instr)), Symbol(mvar, 0), var)))
            else
                push!(q.args, Expr(:(=), var, mvar))
            end
        end
    end
end

"""
For structs wrapping arrays, using `GC.@preserve` can trigger heap allocations.
`preserve_buffer` attempts to extract the heap-allocated part. Isolating it by itself
will often allow the heap allocations to be elided. For example:

```julia
julia> using StaticArrays, BenchmarkTools

julia> # Needed until a release is made featuring https://github.com/JuliaArrays/StaticArrays.jl/commit/a0179213b741c0feebd2fc6a1101a7358a90caed
       Base.elsize(::Type{<:MArray{S,T}}) where {S,T} = sizeof(T)

julia> @noinline foo(A) = unsafe_load(A,1)
foo (generic function with 1 method)

julia> function alloc_test_1()
           A = view(MMatrix{8,8,Float64}(undef), 2:5, 3:7)
           A[begin] = 4
           GC.@preserve A foo(pointer(A))
       end
alloc_test_1 (generic function with 1 method)

julia> function alloc_test_2()
           A = view(MMatrix{8,8,Float64}(undef), 2:5, 3:7)
           A[begin] = 4
           pb = parent(A) # or `LoopVectorization.preserve_buffer(A)`; `perserve_buffer(::SubArray)` calls `parent`
           GC.@preserve pb foo(pointer(A))
       end
alloc_test_2 (generic function with 1 method)

julia> @benchmark alloc_test_1()
BenchmarkTools.Trial:
  memory estimate:  544 bytes
  allocs estimate:  1
  --------------
  minimum time:     17.227 ns (0.00% GC)
  median time:      21.352 ns (0.00% GC)
  mean time:        26.151 ns (13.33% GC)
  maximum time:     571.130 ns (78.53% GC)
  --------------
  samples:          10000
  evals/sample:     998

julia> @benchmark alloc_test_2()
BenchmarkTools.Trial:
  memory estimate:  0 bytes
  allocs estimate:  0
  --------------
  minimum time:     3.275 ns (0.00% GC)
  median time:      3.493 ns (0.00% GC)
  mean time:        3.491 ns (0.00% GC)
  maximum time:     4.998 ns (0.00% GC)
  --------------
  samples:          10000
  evals/sample:     1000
```
"""
@inline preserve_buffer(A::AbstractArray) = A
@inline preserve_buffer(A::SubArray) = preserve_buffer(parent(A))
@inline preserve_buffer(A::PermutedDimsArray) = preserve_buffer(parent(A))
@inline preserve_buffer(A::Union{Transpose,Adjoint}) = preserve_buffer(parent(A))
@inline preserve_buffer(x) = x

function gc_preserve(ls::LoopSet, q::Expr)
    length(ls.opdict) == 0 && return q
    q2 = Expr(:block)
    gcp = Expr(:gc_preserve, q)
    # gcp = Expr(:macrocall, Expr(:(.), :GC, QuoteNode(Symbol("@preserve"))), LineNumberNode(@__LINE__, Symbol(@__FILE__)))
    for array ∈ ls.includedactualarrays
        pb = gensym(array);
        push!(q2.args, Expr(:(=), pb, Expr(:call, lv(:preserve_buffer), array)))
        push!(gcp.args, pb)
    end
    q.head === :block && push!(q.args, nothing)
    # push!(gcp.args, q)
    push!(q2.args, gcp)
    q2
    # Expr(:block, gcp)
end

function typeof_outer_reduction_init(ls::LoopSet, op::Operation)
    opid = identifier(op)
    for (id, sym) ∈ ls.preamble_symsym
        opid == id && return Expr(:call, :typeof, sym)
    end
    for (id,intval) ∈ ls.preamble_symint
        opid == id && return :Int
    end
    for (id,floatval) ∈ ls.preamble_symfloat
        opid == id && return :Float64
    end
    for (id,typ) ∈ ls.preamble_zeros
        instruction(ops[id]) === LOOPCONSTANT || continue
        opid == id || continue
        if typ == IntOrFloat
            return :Float64
        elseif typ == HardInt
            return :Int
        else#if typ == HardFloat
            return :Float64
        end
    end
    throw("Could not find initializing constant.")
end
function typeof_outer_reduction(ls::LoopSet, op::Operation)
    for opp ∈ operations(ls)
        opp === op && continue
        name(op) === name(opp) && return typeof_outer_reduction_init(ls, opp)
    end
    throw("Could not find initialization op.")
end

function determine_eltype(ls::LoopSet)::Union{Symbol,Expr}
    if length(ls.includedactualarrays) == 0
        if length(ls.outer_reductions) == 0
            return Expr(:call, lv(:typeof), 0)
        elseif length(ls.outer_reductions) == 1
            op = ls.operations[only(ls.outer_reductions)]
            return typeof_outer_reduction(ls, op)
        else
            pt = Expr(:call, lv(:promote_type))
            for j ∈ ls.outer_reductions
                push!(pt.args, typeof_outer_reduction(ls, ls.operations[j]))
            end
            return pt
        end
    elseif length(ls.includedactualarrays) == 1
        return Expr(:call, lv(:eltype), first(ls.includedactualarrays))
    end
    promote_q = Expr(:call, lv(:promote_type))
    for array ∈ ls.includedactualarrays
        push!(promote_q.args, Expr(:call, lv(:eltype), array))
    end
    promote_q
end
@inline _eltype(x) = eltype(x)
@inline _eltype(::BitArray) = VectorizationBase.Bit
function determine_width(
    ls::LoopSet, vectorized::Union{Symbol,Nothing}
)
    vwidth_q = Expr(:call, lv(:pick_vector_width))
    if !isnothing(vectorized)
        vloop = getloop(ls, vectorized)
        if isstaticloop(vloop)
            push!(vwidth_q.args, Expr(:call, Expr(:curly, :Val, length(vloop))))
        end
    end
    # push!(vwidth_q.args, ls.T)
    if length(ls.includedactualarrays) < 2
        push!(vwidth_q.args, ELTYPESYMBOL)
    else
        for array ∈ ls.includedactualarrays
            push!(vwidth_q.args, Expr(:call, lv(:_eltype), array))
        end
    end
    vwidth_q
end
function init_remblock(unrolledloop::Loop, lssm::LoopStartStopManager, n::Int)#u₁loop::Symbol = unrolledloop.itersymbol)
    termind = lssm.terminators[n]
    condition = if iszero(termind)
        loopsym = unrolledloop.itersymbol
        if unrolledloop.stopexact
            Expr(:call, :<, loopsym, unrolledloop.stophint)
        else
            Expr(:call, :<, loopsym, unrolledloop.stopsym)
        end
    else
        termar = lssm.incrementedptrs[n][termind]
        ptr = vptr(termar)
        Expr(:call, :<, callpointerforcomparison(ptr), maxsym(ptr, 0))
    end
    Expr(:if, condition)
end

function maskexpr(looplimit)
    Expr(:(=), Symbol("##mask##"), Expr(:call, lv(:mask), VECTORWIDTHSYMBOL, looplimit))
    # rem = Expr(:call, lv(:valrem), W, looplimit)
    # Expr(:(=), Symbol("##mask##"), Expr(:call, lv(:masktable), W, rem))
end
function definemask(loop::Loop)
    if isstaticloop(loop)
        maskexpr(length(loop))
    elseif loop.startexact && isone(loop.starthint)
        maskexpr(loop.stopsym)
    else
        lexpr = if loop.startexact
            Expr(:call, lv(:vsub_fast), loop.stopsym, loop.starthint - 1)
        elseif loop.stopexact
            Expr(:call, lv(:vsub_fast), loop.stophint + 1, loop.startsym)
        else
            Expr(:call, lv(:vsub_fast), Expr(:call, lv(:vadd_fast), loop.stopsym, 1), loop.startsym)
        end
        maskexpr(lexpr)
    end
end
# function definemask_for_alignment_cleanup(loop::Loop)
#     lexpr = if loop.stopexact
#         Expr(:call, lv(:vsub_fast), loop.stophint + 1, loop.itersym)
#     else
#         Expr(:call, lv(:vsub_fast), Expr(:call, lv(:vadd_fast), loop.stopsym, 1), loop.itersymbol)
#     end
#     maskexpr(lexpr)
# end
function define_eltype_vec_width!(q::Expr, ls::LoopSet, vectorized)
    push!(q.args, Expr(:(=), ELTYPESYMBOL, determine_eltype(ls)))
    push!(q.args, Expr(:(=), VECTORWIDTHSYMBOL, determine_width(ls, vectorized)))
    nothing
end
function setup_preamble!(ls::LoopSet, us::UnrollSpecification, Ureduct::Int)
    @unpack u₁loopnum, u₂loopnum, vectorizedloopnum, u₁, u₂ = us
    order = names(ls)
    u₁loopsym = order[u₁loopnum]
    u₂loopsym = order[u₂loopnum]
    vectorized = order[vectorizedloopnum]
    set_vector_width!(ls, vectorized)
    iszero(length(ls.includedactualarrays) + length(ls.outer_reductions)) || define_eltype_vec_width!(ls.preamble, ls, vectorized)
    lower_licm_constants!(ls)
    isone(num_loops(ls)) || pushpreamble!(ls, definemask(getloop(ls, vectorized)))#, u₁ > 1 && u₁loopnum == vectorizedloopnum))
    initialize_outer_reductions!(ls, 0, Ureduct, vectorized)
    for op ∈ operations(ls)
        (iszero(length(loopdependencies(op))) && iscompute(op)) && lower_compute!(ls.preamble, op, ls, UnrollArgs(u₁, u₁loopsym, u₂loopsym, vectorized, u₂, nothing), nothing)
    end
end
function lsexpr(ls::LoopSet, q)
    Expr(:block, ls.preamble, q)
end
function calc_Ureduct(ls::LoopSet, us::UnrollSpecification)
    @unpack u₁loopnum, u₁, u₂, vectorizedloopnum = us
    if iszero(length(ls.outer_reductions))
        -1
    elseif u₂ == -1
        loopisstatic = isstaticloop(getloop(ls, names(ls)[u₁loopnum]))
        loopisstatic &= ((vectorizedloopnum != u₁loopnum) | (!iszero(ls.vector_width[])))
        # loopisstatic ? u₁ : min(u₁, 4) # much worse than the other two options, don't use this one
        if Sys.CPU_NAME === "znver1"
            loopisstatic ? u₁ : 1
        else
            loopisstatic ? u₁ : (u₁ ≥ 4 ? 2 : 1)
        end
    else
        8#u₂#u₁
    # elseif num_loops(ls) == u₁loopnum
    #     min(u₁, 4)
    # else
    #     # u₂ == -1 ? u₁ : 4
    #     u₁
    end
end
function lower_unrollspec(ls::LoopSet)
    us = ls.unrollspecification[]
    @unpack vectorizedloopnum, u₁, u₂ = us
    # @show u₁, u₂
    order = names(ls)
    vectorized = order[vectorizedloopnum]
    Ureduct = calc_Ureduct(ls, us)
    setup_preamble!(ls, us, Ureduct)
    initgesps = add_loop_start_stop_manager!(ls)
    q = Expr(:let, initgesps, lower_unrolled_dynamic(ls, us, num_loops(ls), false))
    q = gc_preserve( ls, Expr(:block, q) )
    reduce_expr!(q, ls, Ureduct)
    lsexpr(ls, q)
end

function lower(ls::LoopSet, order, u₁loop, u₂loop, vectorized, u₁, u₂, inline::Bool)
    cacheunrolled!(ls, u₁loop, u₂loop, vectorized)
    fillorder!(ls, order, u₁loop, u₂loop, u₂ != -1, vectorized)
    ls.unrollspecification[] = UnrollSpecification(ls, u₁loop, u₂loop, vectorized, u₁, u₂)
    q = lower_unrollspec(ls)
    inline && pushfirst!(q.args, Expr(:meta, :inline))
    q
end

function lower(ls::LoopSet, inline::Int = -1)
    order, u₁loop, u₂loop, vectorized, u₁, u₂, c, shouldinline = choose_order_cost(ls)
    lower(ls, order, u₁loop, u₂loop, vectorized, u₁, u₂, inlinedecision(inline, shouldinline))
end
function lower(ls::LoopSet, u₁::Int, u₂::Int, inline::Int)
    if u₂ > 1
        @assert num_loops(ls) > 1 "There is only $(num_loops(ls)) loop, but specified blocking parameter u₂ is $u₂."
        order, u₁loop, u₂loop, vectorized, _u₁, _u₂, c, shouldinline = choose_tile(ls)
        copyto!(ls.loop_order.bestorder, order)
    else
        u₂ = -1
        order, vectorized, c = choose_unroll_order(ls, Inf)
        u₁loop = first(order); u₂loop = Symbol("##undefined##"); shouldinline = true
        copyto!(ls.loop_order.bestorder, order)
    end
    doinline = inlinedecision(inline, shouldinline)
    lower(ls, order, u₁loop, u₂loop, vectorized, u₁, u₂, doinline)
end

# Base.convert(::Type{Expr}, ls::LoopSet) = lower(ls)
Base.show(io::IO, ls::LoopSet) = println(io, lower(ls))

"""
This function is normally called
isunrolled_sym(op, u₁loop)
isunrolled_sym(op, u₁loop, u₂loop)

It returns `true`/`false` for each loop, indicating whether they're unrolled.

If there is a third argument, it will avoid unrolling that symbol along reductions if said symbol is part of the reduction chain.

"""
function isunrolled_sym(op::Operation, u₁loop::Symbol, u₂loop::Symbol)
    u₁ild = u₁loop ∈ loopdependencies(op)
    u₂ild = u₂loop ∈ loopdependencies(op)
    (accesses_memory(op) | isloopvalue(op)) && return (u₁ild, u₂ild)
    if isconstant(op)
        if !u₁ild
            u₁ild = u₁loop ∈ reducedchildren(op)
        end
        if !u₂ild
            u₂ild = u₂loop ∈ reducedchildren(op)
        end
    end
    # @show op u₁ild, u₂ild
    (u₁ild & u₂ild) || return u₁ild, u₂ild
    reductops = isconstant(op) ? reducedchildren(op) : reduceddependencies(op)
    # @show op reductops
    iszero(length(reductops)) && return true, true
    u₁reduced = u₁loop ∈ reductops
    u₂reduced = u₂loop ∈ reductops
    # We want to only unroll one of them.
    # We prefer not to unroll a reduced loop
    # @show u₁reduced, u₂reduced
    if u₂reduced # if both are reduced, we unroll u₁
        true, false
    elseif u₁reduced
        false, true
    else
        true, true
    end
end

function isunrolled_sym(op::Operation, u₁loop::Symbol)
    u₁loop ∈ loopdependencies(op) || (isconstant(op) && (u₁loop ∈ reducedchildren(op)))
end

isunrolled_sym(op::Operation, u₁loop::Symbol, u₂loop::Symbol, ::Nothing) = (isunrolled_sym(op, u₁loop), false)
isunrolled_sym(op::Operation, u₁loop::Symbol, u₂loop::Symbol, ::Int) = isunrolled_sym(op, u₁loop, u₂loop)

variable_name(op::Operation, ::Nothing) = mangledvar(op)
variable_name(op::Operation, suffix) = Symbol(mangledvar(op), suffix, :_)

function variable_name_and_unrolled(op::Operation, u₁loop::Symbol, u₂loop::Symbol, ::Nothing)
    mangledvar(op), isunrolled_sym(op, u₁loop), false
end
function variable_name_and_unrolled(op::Operation, u₁loop::Symbol, u₂loop::Symbol, u₂iter::Int)
    u₁op, u₂op = isunrolled_sym(op, u₁loop, u₂loop)
    mvar = u₂op ? variable_name(op, u₂iter) : mangledvar(op)
    mvar, u₁op, u₂op
end


text:
# 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

ext: jl


text:
################################################################################
#
# Project root module.
#
################################################################################

module FluxHELayers

    include("LinearTransformations.jl")
    include("Polynomials.jl")

    using Flux
    using SEAL

    using .LinearTransformations: naive_matrix_multiply
    using .Polynomials

    export Square

    """
        (a::Dense)(x::AbstractVecOrMat{T}) where {T<:Ciphertext}

    Performs a fully connected layer forward pass over a vector of individually
    encrypted ciphertexts.

    In what can be described as an ugly hack, this function uses a global variable
    called "SEALParams" which is an object that contains:
        - the SEAL CKKS evaluator, called "evaluator"
        - the SEAL CKKS encoder, called "encoder"
        - the SEAL CKKS relinearization keys, called "relinearization_keys"
    """
    function (a::Dense)(x::AbstractVecOrMat{T}) where {T<:Ciphertext}
        if !isdefined(Main, :SEALParams)
            throw("SEALParams needs to be defined to run this function.")
        else
            SEALParams = Main.SEALParams
        end

        @debug "Performing forward-pass of Dense layer"

        W, b = a.W, a.b

        # Compute the layer weight multiplication.
        result = naive_matrix_multiply(W,
                                       x,
                                       SEALParams.evaluator,
                                       SEALParams.encoder,
                                       SEALParams.relinearization_keys)

        # For each intermediate result, add the bias term.
        for i in 1:length(b)
            # Encode the bias term with the same initial scale as the
            # intermediate result ciphertext and mod switch to the same level.
            bias_plaintext = Plaintext()
            encode!(bias_plaintext, Float64(b[i]), scale(result[i]), SEALParams.encoder)
            mod_switch_to_inplace!(bias_plaintext, parms_id(result[i]), SEALParams.evaluator)

            # Add the bias plaintext to the intermediate result.
            add_plain_inplace!(result[i], bias_plaintext, SEALParams.evaluator)
        end

        result
    end

    """
        (BN::BatchNorm)(x::AbstractVecOrMat{T}) where {T<:Ciphertext}

    Performs a Batch Norm forward pass over a vector of individually encrypted
    ciphertexts.

    In what can be described as an ugly hack, this function uses a global variable
    called "SEALParams" which is an object that contains:
        - the SEAL CKKS evaluator, called "evaluator"
        - the SEAL CKKS encoder, called "encoder"
        - the SEAL CKKS relinearization keys, called "relinearization_keys"
    """
    function (bn::BatchNorm)(x::AbstractVecOrMat{T}) where {T<:Ciphertext}
        if !isdefined(Main, :SEALParams)
            throw("SEALParams needs to be defined to run this function.")
        else
            SEALParams = Main.SEALParams
        end

        @debug "Performing forward-pass of BatchNorm layer"

        β, γ, μ, σ², ϵ = bn.β, bn.γ, bn.μ, bn.σ², bn.ϵ

        output = Array{Ciphertext}(undef, length(x))

        for i in 1:length(x)
            # Negate, encode and mod switch μ to calculate the numerator.
            μ_neg_plaintext = Plaintext()
            numerator_ciphertext = Ciphertext()
            encode!(μ_neg_plaintext, Float64(-μ[i]), scale(x[i]), SEALParams.encoder)
            mod_switch_to_inplace!(μ_neg_plaintext, parms_id(x[i]), SEALParams.evaluator)

            # Calculate the numerator.
            add_plain!(numerator_ciphertext, x[i], μ_neg_plaintext, SEALParams.evaluator)

            # Calculate and encode the denominator.
            denominator_plaintext = Plaintext()
            encode!(denominator_plaintext, Float64(1 / sqrt(σ²[i] + ϵ)), scale(x[i]), SEALParams.encoder)
            mod_switch_to_inplace!(denominator_plaintext, parms_id(x[i]), SEALParams.evaluator)

            # Calculate x.
            xi_ciphertext = Ciphertext()
            multiply_plain!(xi_ciphertext, numerator_ciphertext, denominator_plaintext, SEALParams.evaluator)
            relinearize_inplace!(xi_ciphertext, SEALParams.relinearization_keys, SEALParams.evaluator)
            rescale_to_next_inplace!(xi_ciphertext, SEALParams.evaluator)

            # Calculate the linear transformation w.r.t β.
            β_plaintext = Plaintext()
            encode!(β_plaintext, Float64(β[i]), scale(xi_ciphertext), SEALParams.encoder)
            mod_switch_to_inplace!(β_plaintext, parms_id(xi_ciphertext), SEALParams.evaluator)

            result_i_ciphertext = Ciphertext()
            add_plain!(result_i_ciphertext, xi_ciphertext, β_plaintext, SEALParams.evaluator)

            # Calculate the linear transforamtion w.r.t γ.
            γ_plaintext = Plaintext()
            encode!(γ_plaintext, Float64(γ[i]), scale(result_i_ciphertext), SEALParams.encoder)
            mod_switch_to_inplace!(γ_plaintext, parms_id(result_i_ciphertext), SEALParams.evaluator)

            multiply_plain_inplace!(result_i_ciphertext, γ_plaintext, SEALParams.evaluator)
            relinearize_inplace!(result_i_ciphertext, SEALParams.relinearization_keys, SEALParams.evaluator)
            rescale_to_next_inplace!(result_i_ciphertext, SEALParams.evaluator)

            output[i] = result_i_ciphertext
        end

        output
    end

    """
        Square(x::AbstractVecOrMat)

    Defines a square activation function that can be used for the forwards and
    backwards pass over normal values.
    """
    function Square(x::AbstractVecOrMat)
        x .^ 2
    end

    """
        Square(x::AbstractVecOrMat{T}) where {T<:Ciphertext}

    Performs a square activation function forward pass over a vector of
    individually encrypted ciphertexts.

    In what can be described as an ugly hack, this function uses a global variable
    called "SEALParams" which is an object that contains:
        - the SEAL CKKS evaluator, called "evaluator"
        - the SEAL CKKS encoder, called "encoder"
        - the SEAL CKKS relinearization keys, called "relinearization_keys"
    """
    function Square(x::AbstractVecOrMat{T}) where {T<:Ciphertext}
        if !isdefined(Main, :SEALParams)
            throw("SEALParams needs to be defined to run this function.")
        else
            SEALParams = Main.SEALParams
        end

        @debug "Performing forward-pass of Square layer"

        for i in 1:length(x)
            square_inplace!(x[i], SEALParams.evaluator)
            relinearize_inplace!(x[i], SEALParams.relinearization_keys, SEALParams.evaluator)
            rescale_to_next_inplace!(x[i], SEALParams.evaluator)
        end

        x
    end

    """
        ReLU(x::AbstractVecOrMat)

    Defines a ReLU activation function that can be used for the forwards and
    backwards pass over plaintext values.
    """
    function ReLU(x::AbstractVecOrMat)
        relu.(x)
    end

    """
        ReLU(x::AbstractVecOrMat{T}) where {T<:Ciphertext}

    Throw a friendly reminder error, since we can't use the ReLU function
    on homomorphically encrypted data.
    """
    function ReLU(x::AbstractVecOrMat{T}) where {T<:Ciphertext}
        throw(ArgumentError("ReLU on encryted data not available."))
    end
end


text:
# Most functions here are copied and adapted from `constraints.jl`
"""
    supports_objective( model::ModelLike,
		::Type{F})::Bool where {F<MOI.AbstractFunction}

Return a `Bool` indicating whether `model` generally supports adding an objective
of type `F`.
That is, `copy_to(model, src)` does not throw [`UnsupportedGoal`](@ref) when
`src` contains objectives of type `F`.
"""
function supports_objective(
		::ModelLike,
		::Type{F},
		::Type{S}
	) where{F<:AbstractFunction, S<:OptimSense}
	if !(is_multi_model(model))
		return supports_single_objective(model,F,S)
	end
    return false
end

# this is a helper function to be used if `!(is_multi_model(model))`
# to check whether model supports F-S objectives via the "old" interface
function supports_single_objective(
	model :: ModelLike,
	::Type{F},
	::Type{S} ) where{F<:AbstractFunction, S<:OptimSense}
	return(
		supports(model, ObjectiveFunction{F}()) &&
		supports(model, ObjectiveSense()) #, _optimization_sense(S))
	)
end

"""
    struct UnsupportedGoal{F<:MOI.AbstractFunction}
        message::String # Human-friendly explanation why the attribute cannot be set
    end

An error indicating that objectives of type `F` are not supported by
the model, i.e. that [`supports_constraint`](@ref) returns `false`.
"""
struct UnsupportedGoal{F<:AbstractFunction,S<:OptimSense} <: UnsupportedError
    message::String # Human-friendly explanation why the attribute cannot be set
end
UnsupportedGoal{F,S}() where {F,S} = UnsupportedGoal{F,S}("")

function element_name(::UnsupportedGoal{F,S}) where {F,S}
    return "`$F-$S` objective"
end

"""
    struct AddGoalNotAllowed{F<:AbstractFunction} <: NotAllowedError
        message::String # Human-friendly explanation why the attribute cannot be set
    end
An error indicating that objectives of type `F` are supported (see
[`supports_objective`](@ref)) but cannot be added.
"""
struct AddGoalNotAllowed{F<:AbstractFunction, S<:OptimSense} <:
       NotAllowedError
    message::String # Human-friendly explanation why the attribute cannot be set
end
AddGoalNotAllowed{F,S}() where {F,S} = AddConstraintNotAllowed{F,S}("")

function operation_name(::AddGoalNotAllowed{F,S}) where {F,S}
    return "Adding `$F-$S` objectives."
end

"""
    add_objective(model::ModelLike, func::F, sense::S)::GoalIndex{F,S} where {F,S}

Add the objective ``f: ℝⁿ → ℝᵏ`` to `model`, where ``f`` is defined by `func` and
`sense` is the optimization sense.

* An [`UnsupportedGoal`](@ref) error is thrown if `model` does not support
  `F`-objectives.
* a [`AddGoalNotAllowed`](@ref) error is thrown if it supports `F`-objectives
  but it cannot add the objectives(s) in its current state.
"""

function add_objective(
		model::ModelLike,
		func::F,
		sense::S
	) where {F<:AbstractFunction,S<:OptimSense}

	if !(is_multi_model(model))
		if supports_single_objective(model, F, S)
			set(model, ObjectiveSense(), _optimization_sense(sense) )
			set(model, ObjectiveFunction{F}(), func )
		end
		return GoalIndex{F,S}(-1)
	end

	return throw_add_objective_error_fallback(model, func, sense)
end

# if the model generally supports the objective type,
# then throw `error_if_supported` to indicate that currently
# it is not possible to add it.
# else throw `UnsupportedObjective(F)`.
function throw_add_objective_error_fallback(
	model::ModelLike,
	func::F,
	sense::S;
	error_if_supported = AddObjectiveNotAllowed{F,S}(),
) where {F<:AbstractFunction, S<:OptimSense}
    if supports_objective(model, F, S)
        throw(error_if_supported)
    else
        throw(UnsupportedObjective{F,S}())
    end
end

"""
    add_objectives(model::ModelLike, funcs::Vector{F}, senses::Vector{S})::Vector{GoalIndex{F,S}} where {F,S}

Add the set of objectives specified by each function-sense pair in `funcs` and `sets`.
`F` and `S` should be concrete types.
This call is equivalent to `add_constraint.(model, funcs, sets)` but may be more efficient.
"""
function add_objectives end

# default fallback
function add_objectives(model::ModelLike, funcs, senses)
    return add_objective.(model, funcs, senses)
end

# default fallback
function transform(model::ModelLike, o::GoalIndex, new_sense)
    f = get(model, GoalFunction(), o)
    delete_fallback(model, o)
    return add_objective(model, f, new_sense)
end

function delete_fallback(model::ModelLike, o::GoalIndex)
	!is_multi_model(model) && return nothing
	return delete(model, o)
end


text:
struct NILSASSensitivityFunction{iip,NILSS,ASF,Mtype} <: DiffEqBase.AbstractODEFunction{iip}
  nilss::NILSS
  S::ASF # Adjoint sensitivity function
  M::Mtype
  discrete::Bool
end

struct QuadratureCache{A1,A2,A3,A4,A5}
  dwv::A1
  dwf::A1
  dwfs::A2
  dvf::A3
  dvfs::A4
  dJs::A4
  C::A5
  R::A5
  b::A1
end

function QuadratureCache(u0, M, nseg, numparams)
  dwv = Array{eltype(u0)}(undef, M, nseg)
  dwf = Array{eltype(u0)}(undef, M, nseg)
  dwfs = Array{eltype(u0)}(undef, numparams*M, nseg)
  dvf = Array{eltype(u0)}(undef, 1, nseg)
  dvfs = Array{eltype(u0)}(undef, numparams, nseg)
  dJs = Array{eltype(u0)}(undef, numparams, nseg)
  C = Array{eltype(u0)}(undef, M, M, nseg)
  R = Array{eltype(u0)}(undef, M, M, nseg)
  b = Array{eltype(u0)}(undef, M, nseg)

  QuadratureCache{typeof(dwv),typeof(dwfs),typeof(dvf),typeof(dvfs),typeof(C)}(dwv,dwf,dwfs,dvf,dvfs,dJs,C,R,b)
end

struct NILSASProblem{A,NILSS,Aprob,Qcache,solType,z0Type,G,DG,T}
  sensealg::A
  nilss::NILSS # diffcache
  adjoint_prob::Aprob
  quadcache::Qcache
  sol::solType
  z0::z0Type
  g::G
  dg::DG
  T_seg::T
  dtsave::T
end

function NILSASProblem(sol, sensealg::NILSAS, g, t=nothing, dg = nothing; kwargs...)

  @unpack f, p, u0, tspan = sol.prob
  @unpack nseg, nstep, rng, adjoint_sensealg, M = sensealg  #number of segments on time interval, number of steps saved on each segment

  numindvar = length(u0)
  numparams = length(p)

  p === nothing && error("You must have parameters to use parameter sensitivity calculations!")
  !(u0 isa AbstractVector) && error("`u` has to be an AbstractVector.")

  nstep <= 1 && error("At least the start and the end point of each segment must be stored. Please use `nstep >=2`.")

  !(u0 isa AbstractVector) && error("`u` has to be an AbstractVector.")

  # segmentation: determine length of segmentation and spacing between saved points
  T_seg = (tspan[2]-tspan[1])/nseg # length of each segment
  dtsave = T_seg/(nstep-1)

  # homogenous + inhomogenous adjoint sensitivity problem
  # assign initial values to y, vstar, w
  y = copy(sol.u[end])
  z0 = terminate_conditions(adjoint_sensealg,rng,f,y,p,tspan[2],numindvar,numparams,M)
  nilss = NILSSSensitivityFunction(sensealg,f,u0,p,tspan,g,dg)
  tspan = (tspan[2] - T_seg, tspan[2])
  checkpoints = tspan[1]:dtsave:tspan[2]

  adjoint_prob = ODEAdjointProblem(sol,adjoint_sensealg,g,t,dg; checkpoints=checkpoints, z0=z0, M=M, nilss=nilss, tspan=tspan, kwargs...)

  # pre-allocate variables for integration Eq.(23) in NILSAS paper.
  quadcache = QuadratureCache(u0, M, nseg, numparams)

  NILSASProblem{typeof(sensealg),typeof(nilss),typeof(adjoint_prob),typeof(quadcache),
    typeof(sol),typeof(z0),typeof(g),typeof(dg),typeof(T_seg)}(sensealg,nilss,
    adjoint_prob,quadcache,sol,deepcopy(z0),g,dg,T_seg,dtsave)
end

function terminate_conditions(alg::BacksolveAdjoint,rng,f,y,p,t,numindvar,numparams,M)
  if isinplace(f)
    f_unit = zero(y)
    f(f_unit,y,p,t)
    normalize!(f_unit)
  else
    f_unit = f(y,p,t)
    normalize!(f_unit)
  end

  if M>1
    W = rand(rng,numindvar,M-1)
    W .-= (f_unit'*W) .* f_unit
    w, _ = qr(W)
    _w = @view w[:,1:M-1]
    W = hcat(_w, f_unit)
  else
    W = f_unit
  end
  vst = zeros(numindvar)

  # quadratures
  C = zeros(M,M)
  dwv = zeros(M)
  dwf = zeros(M)
  dvf = zeros(1)
  dJs = zeros(numparams)

  return ArrayPartition([vst;vec(W)],zeros(numparams*(M+1)),y,C,dwv,dwf,dvf,dJs)
end

function split_states(du,u,t,NS::NILSASSensitivityFunction,j;update=true)
  @unpack nilss, S = NS
  @unpack numindvar,numparams = nilss

  indx1 = (j-1)*(numindvar) + 1
  indx2 = indx1 + (numindvar-1)
  indx3 = (j-1)*(numparams) + 1
  indx4 = indx3 + (numparams-1)

  λ     = @view u.x[1][indx1:indx2]
  grad  = @view u.x[2][indx3:indx4]
  _y    = u.x[3]

  # like ODE/Drift term and scalar noise
  dλ    = @view du.x[1][indx1:indx2]
  dgrad = @view du.x[2][indx3:indx4]
  dy    = du.x[3]

  λ,grad,_y,dλ,dgrad,dy
end

function split_quadratures(du,u,t,NS::NILSASSensitivityFunction;update=true)
  @unpack nilss, S = NS
  @unpack numindvar,numparams = nilss

  C     = u.x[4]
  dwv   = u.x[5]
  dwf   = u.x[6]
  dvf   = u.x[7]
  dJs   = u.x[8]

  dC     = du.x[4]
  ddwv   = du.x[5]
  ddwf   = du.x[6]
  ddvf   = du.x[7]
  ddJs   = du.x[8]

  dC,ddwv,ddwf,ddvf,ddJs, C,dwv,dwf,dvf,dJs
end

function (NS::NILSASSensitivityFunction)(du,u,p,t)
  @unpack nilss, S, M = NS
  @unpack f, dg, dg_val, pgpu, pgpu_config, pgpp, pgpp_config, numparams, numindvar, alg = nilss
  @unpack y, discrete = S

  λ,_,_y,dλ,dgrad,dy = split_states(du,u,t,NS,1)
  copyto!(vec(y), _y)

  #  compute gradient of objective wrt. state
  if !discrete
    accumulate_cost!(dg, y, p, t, nilss)
  end

  # loop over all adjoint states
  for j=1:M+1
    λ,_,_,dλ,dgrad,dy = split_states(du,u,t,NS,j)
    vecjacobian!(dλ, y, λ, p, t, S, dgrad=dgrad, dy=dy)
    dλ .*= -1

    if j==1
      # j = 1 is the inhomogenous adjoint solution
      if !discrete
        if dg_val isa Tuple
          dλ .-= vec(dg_val[1])
        else
          dλ .-= vec(dg_val)
        end
      end
    end
  end

  # quadratures
  dC,ddwv,ddwf,ddvf,ddJs, _,_,_,_,_ = split_quadratures(du,u,t,NS)
  # j = 1 is the inhomogenous adjoint solution
  λv,_,_,_,_,dy = split_states(du,u,t,NS,1)
  ddvf .= -dot(λv,dy)
  for j=1:M
    λ,_,_,_,_,_ = split_states(du,u,t,NS,j+1)
    ddwf[j] = -dot(λ,dy)
    ddwv[j] = -dot(λ,λv)
    for i=j+1:M
      λ2,_,_,_,_,_ = split_states(du,u,t,NS,i+1)
      dC[j,i] = -dot(λ,λ2)
      dC[i,j] = dC[j,i]
    end
    dC[j,j] = -dot(λ,λ)
  end

  if dg_val isa Tuple && !discrete
    ddJs = -vec(dg_val[2])
  end

  return nothing
end

function accumulate_cost!(dg, y, p, t, nilss::NILSSSensitivityFunction)
  @unpack dg_val, pgpu, pgpu_config, pgpp, pgpp_config, alg = nilss

  if dg===nothing
    if dg_val isa Tuple
      DiffEqSensitivity.gradient!(dg_val[1],pgpu,y,alg,pgpu_config)
      DiffEqSensitivity.gradient!(dg_val[2],pgpp,y,alg,pgpp_config)
      dg_val[1] .*= -1
      dg_val[2] .*= -1
    else
      DiffEqSensitivity.gradient!(dg_val,pgpu,y,alg,pgpu_config)
      dg_val .*= -1
    end
  else
    if dg_val isa Tuple
      dg[1](dg_val[1],y,p,t)
      dg[2](dg_val[2],y,p,t)
    else
      dg(dg_val,y,p,t)
    end
  end

  return nothing
end

function adjoint_sense(prob::NILSASProblem,nilsas::NILSAS,alg; kwargs...)

  @unpack M, nseg, nstep, adjoint_sensealg = nilsas
  @unpack sol, nilss, z0, g, dg, T_seg, dtsave, adjoint_prob = prob
  @unpack u0, tspan = adjoint_prob

  copyto!(z0,u0)

  @assert haskey(adjoint_prob.kwargs, :callback)
  # get loss callback
  cb = adjoint_prob.kwargs[:callback]

  # adjoint sensitivities on segments
  for iseg=nseg:-1:1
    t1 = tspan[1]-(nseg-iseg+1)*T_seg
    t2 = tspan[1]-(nseg-iseg)*T_seg
    checkpoints=t1:dtsave:t2
    _prob = ODEAdjointProblem(sol,adjoint_sensealg,g,nothing,dg;
              checkpoints=checkpoints,z0=z0,M=M,nilss=nilss,tspan=(t1,t2),kwargs...)
    _sol = solve(_prob,alg; save_everystep=false,save_start=false,saveat=eltype(sol[1])[],
                  dt = dtsave,
                  tstops=checkpoints,
                  callback = cb,
                  kwargs...)

    # renormalize at interfaces and store quadratures
    # update sense problem
    renormalize!(prob, _sol, z0, iseg)
  end
  return nothing
end

function renormalize!(prob::NILSASProblem, sol, z0, iseg)
  @unpack quadcache, nilss, sensealg = prob
  @unpack M = sensealg
  @unpack numparams, numindvar = nilss
  @unpack R,b = quadcache

  x = sol.u[end].x
  # vstar_right (inhomogenous adjoint on the rhs of the interface)
  vstar = @view x[1][1:numindvar]
  # homogenous adjoint on the rhs of the interface
  W = @view x[1][numindvar+1:end]
  W = reshape(W, numindvar, M)

  Q_, R_ = qr(W)
  Q = @view Q_[:,1:M]
  b_ = (Q'*vstar)

  # store R and b to solve NILSAS problem
  copyto!( (@view R[:,:,iseg]), R_)
  copyto!( (@view b[:,iseg]), b_)

  # store quadrature values
  store_quad(quadcache, x, numparams, iseg)

  # reset z0
  reset!(z0, numindvar, vstar, b_, Q)

  return nothing
end

function store_quad(quadcache, x, numparams, iseg)
  @unpack dwv,dwf,dwfs,dvf,dvfs,dJs,C = quadcache

  grad_vfs  = @view x[2][1:numparams]
  copyto!( (@view dvfs[:,iseg]), grad_vfs)

  grad_wfs  = @view x[2][numparams+1:end]
  copyto!( (@view dwfs[:,iseg]), grad_wfs)

  # C_i = x[4]
  copyto!( (@view C[:,:,iseg]), x[4])
  # dwv_i = x[5]
  copyto!( (@view dwv[:,iseg]), x[5])
  # dwf_i = x[6]
  copyto!( (@view dwf[:,iseg]), x[6])
  # dvf_i = x[7]
  copyto!( (@view dvf[:,iseg]), x[7])
  # dJs_i = x[8]
  copyto!( (@view dJs[:,iseg]), x[8])
  return nothing
end

function reset!(z0, numindvar, vstar, b, Q)
  # modify z0
  x0 = z0.x

  # vstar_left
  v = @view x0[1][1:numindvar]
  v .= vstar - vec(b'*Q')

  # W_left (homogenous adjoint on lhs of the interface)
  w = @view x0[1][numindvar+1:end]
  w .= vec(Q)

  # reset all other values t0 zero
  x0[2] .*= false
  x0[4] .*= false
  x0[5] .*= false
  x0[6] .*= false
  x0[7] .*= false
  x0[8] .*= false
  return nothing
end

function nilsas_min(cache::QuadratureCache)
  @unpack dwv,dwf,dvf,C,R,b = cache

  # Construct Schur complement of the Lagrange multiplier method of the NILSAS problem.

  # see description in Appendix A of Nilsas paper.
  # M= # unstable CLVs, K = # segments
  M, K = size(dwv)

  # construct Cinv
  # Cinv is a block diagonal matrix
  Cinv = zeros(eltype(C), M*K, M*K)

  for i=1:K
    Ci = @view C[:, :, i]
    _Cinv = @view Cinv[(i-1)*M+1:i*M, (i-1)*M+1:i*M]
    Ciinv = inv(Ci)
    copyto!(_Cinv,Ciinv)
  end

  # construct B, also very sparse if K >> M
  B = zeros(eltype(C), M*K-M+1, M*K)

  for i=1:K
    if i<K
      # off diagonal Rs
      _B = @view B[(i-1)*M+1:i*M, i*M+1:(i+1)*M]
      _R = @view R[:,:,i+1]
      copyto!(_B, _R)
      _B .*= -1

      # diagonal ones
      for j=1:M
        B[(i-1)*M+j, (i-1)*M+j] = one(eltype(R))
      end
    end
    # last row
    dwfi = dwf[:,i]
    _B = @view B[end, (i-1)*M+1:i*M]
    copyto!(_B, dwfi)
  end

  # construct d
  d = vec(dwv)

  # construct b
  _b = [b[M+1:end]; -sum(dvf)]

  # compute Lagrange multiplier
  λ = (-B*Cinv*B') \ (B*Cinv*d + _b)

  # return a
  return reshape(-Cinv*(B'*λ + d), M, K)
end

function shadow_adjoint(prob::NILSASProblem,alg; kwargs...)
  shadow_adjoint(prob,prob.sensealg,alg; kwargs...)
end

function shadow_adjoint(prob::NILSASProblem,sensealg::NILSAS,alg; kwargs...)

  # compute adjoint sensitivities
  adjoint_sense(prob,sensealg,alg; kwargs...)

  # compute NILSAS problem on multiple segments
  a = nilsas_min(prob.quadcache)

  # compute gradient, Eq. (28) -- second part to avoid explicit construction of vbar
  @unpack M, nseg = sensealg
  @unpack dvfs, dJs, dwfs = prob.quadcache

  res = vec(sum(dvfs,dims=2)) + vec(sum(dJs,dims=2))
  NP = length(res) # number of parameters

  # loop over segments
  for (i,ai) in enumerate(eachcol(a))
    dwfsi = @view dwfs[:,i]
    dwfsi = reshape(dwfsi,NP,M)
    res .+= dwfsi*ai
  end

  return res/(nseg*prob.T_seg)
end

check_for_g(sensealg::NILSAS,g) = (g===nothing && error("To use NILSAS, g must be passed as a kwarg."))


text:
using MultipleViewGeometry, Test, LinearAlgebra

@test rotx(0.0) == Matrix(1.0I, 3, 3)
@test rotx(2.0*pi) ≈  Matrix(1.0I, 3, 3)

@test roty(0.0) == Matrix(1.0I, 3, 3)
@test roty(2.0*pi) ≈  Matrix(1.0I, 3, 3)

@test rotz(0.0) == Matrix(1.0I, 3, 3)
@test rotz(2.0*pi) ≈  Matrix(1.0I, 3, 3)

@test rotxyz(0.0, 0.0, 0.0) == Matrix(1.0I, 3, 3)
@test rotxyz(2.0*pi, 2.0*pi, 2.0*pi) ≈  Matrix(1.0I, 3, 3)

@test rodrigues2matrix(0.0, 0.0, 0.0) == Matrix(1.0I, 3, 3)