using SIMDPirates
using LinearAlgebra: Transpose, Adjoint
using Base.Cartesian: @nexprs, @nif
using ..LoopInfo
using LoopVectorization: @avx
using SIMDPirates: vfmadd231
using VectorizationBase: VectorizationBase
using TimerOutputs: TimerOutputs, @timeit_debug, TimerOutput
using UnPack: @unpack

# General direction: we want to avoid pointer arithmetic as much as possible,
# also, don't strict type the container type

# I'll try to implement everything without defining a type, and all variables
# should stay local if possible. We can always clean up the code later when we
# find a good direction/structure.

###
### Prelude
###

const N_THREADS_BUFFERS = map(_->Vector{UInt8}(undef, 0), 1:Threads.nthreads())
const PAGE_SIZE = ccall(:jl_getpagesize, Int, ())

const BLAS_TIMER = TimerOutput()
reset_timer!() = TimerOutputs.reset_timer!(BLAS_TIMER)
get_timer() = BLAS_TIMER
enable_timer() = TimerOutputs.enable_debug_timings(@__MODULE__)
disable_timer() = TimerOutputs.disable_debug_timings(@__MODULE__)

# alias scopes from https://github.com/JuliaLang/julia/pull/31018
struct Const{T<:Array}
    a::T
end

@eval Base.getindex(A::Const, i1::Int) = Core.const_arrayref($(Expr(:boundscheck)), A.a, i1)
@eval Base.getindex(A::Const, i1::Int, i2::Int, I::Int...) = (Base.@_inline_meta; Core.const_arrayref($(Expr(:boundscheck)), A.a, i1, i2, I...))
constarray(a::Array) = Const(a)
constarray(a) = a

macro aliasscope(body)
    sym = gensym()
    esc(quote
        $(Expr(:aliasscope))
        $sym = $body
        $(Expr(:popaliasscope))
        $sym
    end)
end

###
### User-level API
###

function mul!(C, A, B, α=true, β=false; cache_params=(cache_m=72, cache_k=256, cache_n=4080), kernel_params=(Val(8), Val(6)), packing=nothing)
    m, k, n = checkmulsize(C, A, B)
    α = convert(eltype(C), α)
    β = convert(eltype(C), β)
    iszeroα = iszero(α)
    if iszero(β) && iszeroα
        return fill!(C, zero(eltype(C)))
    elseif isone(β) && iszeroα
        return C
    elseif iszeroα
        # TODO: optimize this
        for ii in eachindex(C)
            Cii = @inbounds C[ii]
            Cii = β * Cii
            @inbounds C[ii] = Cii
        end
        return C
    end
    # TODO: packing strategy
    packing === nothing && (packing = (false, false))
    packa, packb = packing
    if packa isa Bool && packb isa Bool
        if packa # manual union split
            if packb
                _mul!(C, A, B, α, β, (Val(true),Val(true)), cache_params, kernel_params)
            else
                _mul!(C, A, B, α, β, (Val(true),Val(false)), cache_params, kernel_params)
            end
        else
            if packb
                _mul!(C, A, B, α, β, (Val(false),Val(true)), cache_params, kernel_params)
            else
                _mul!(C, A, B, α, β, (Val(false),Val(false)), cache_params, kernel_params)
            end
        end
    else
        _mul!(C, A, B, α, β, packing, cache_params, kernel_params)
    end
    return C
end

partition_k(k, cache_k) = cld(k, cld(k, cache_k))

###
### Lower-level `_mul!`
###

# use lowercase here because these are not actually types
struct GemmContext{packa,packb,micro_m,micro_n,C1,C2,A1,A2,B1,B2}
    cache_m::Int
    cache_k::Int
    cache_n::Int
    cs1::C1
    cs2::C2
    as1::A1
    as2::A2
    bs1::B1
    bs2::B2
end
function GemmContext{packa,packb,micro_m,micro_n}(cache_m, cache_k, cache_n, cs1, cs2, as1, as2, bs1, bs2) where {packa,packb,micro_m,micro_n}
    return GemmContext{packa,packb,micro_m,micro_n,
                       typeof(cs1),typeof(cs2),
                       typeof(as1),typeof(as2),
                       typeof(bs1),typeof(bs2)}(
                                                cache_m, cache_k, cache_n,
                                                cs1, cs2,
                                                as1, as2,
                                                bs1, bs2
                                               )
end

function _mul!(C, A, B, α, β, packing::Tuple{Val{packa},Val{packb}}, (cache_m, cache_k, cache_n), kernel_params::Tuple{Val{micro_m},Val{micro_n}}) where {packa,packb,micro_m,micro_n}
    cs1, cs2 = _strides(C)
    as1, as2 = _strides(A)
    bs1, bs2 = _strides(B)
    cs1 == 1 || throw(ArgumentError("Packing kernel doesn't support nonunit leading stride C matrix. Got stride(C, 1) = $cs1."))
    if !packa
        as1 == 1 || throw(ArgumentError("Kernel with packed A doesn't support nonunit leading stride C and A matrices. Got stride(C, 1) = $cs1, stride(A, 1) = $as1, and stride(B, 1) = $bs1."))
    end
    ctx = GemmContext{packa,packb,micro_m,micro_n}(cache_m, cache_k, cache_n,
                                                   cs1, cs2, as1, as2, bs1, bs2)
    m, k, n = checkmulsize(C, A, B)
    A′, B′ = canonicalize(A), canonicalize(B)
    if !packa && !packb
        macrokernel!(C, A′, B′, α, β, (1, m), (1, k), (1, n), ctx)
        return C
    end

    cache_k = partition_k(k, cache_k) # More evenly divide K

    # get buffer
    Abuffersize = cache_m * cache_k
    Bbuffersize = cache_k * cache_n
    T = eltype(C)
    buffer = N_THREADS_BUFFERS[Threads.threadid()]
    resize!(buffer, (Abuffersize + Bbuffersize) * sizeof(T) + 2PAGE_SIZE)
    ptrbuffer = align(convert(Ptr{T}, pointer(buffer)), PAGE_SIZE)
    Abuffer = packa ? unsafe_wrap(Array, ptrbuffer, Abuffersize) : A′
    ptrbuffer = align(ptrbuffer + Abuffersize * sizeof(T), PAGE_SIZE)
    Bbuffer = packb ? unsafe_wrap(Array, ptrbuffer, Bbuffersize) : B′
    ptrbuffer = align(ptrbuffer + Bbuffersize * sizeof(T), PAGE_SIZE)

    for cachej₁ in 1:cache_n:n; cachej₂ = min(cachej₁ + cache_n - 1, n)
        for cachep₁ in 1:cache_k:k; cachep₂ = min(cachep₁ + cache_k - 1, k) # TODO: add adaptive packing?
            β′ = cachep₁ == 1 ? β : one(β)
            packb && @timeit_debug BLAS_TIMER "Pack B" packBbuffer!(Bbuffer, B, cachep₁, cachep₂, cachej₁, cachej₂, Val(micro_n))
            for cachei₁ in 1:cache_m:m; cachei₂ = min(cachei₁ + cache_m - 1, m)
                packa && @timeit_debug BLAS_TIMER "Pack A" packAbuffer!(Abuffer, A, cachei₁, cachei₂, cachep₁, cachep₂, Val(micro_m))
                # macrokernel
                macrokernel!(C, Abuffer, Bbuffer, α, β′, (cachei₁, cachei₂), (cachep₁, cachep₂), (cachej₁, cachej₂), ctx)
            end # cachei
        end # cachep
    end # cachej
    return C
end

###
### Macro kernel
###

@generated function macrokernel!(C, A, B, α, β, cacheis, cacheps, cachejs, ctx::GemmContext{packa,packb,micro_m,micro_n})::Nothing where {packa,packb,micro_m,micro_n}
    T = eltype(C)
    N = VectorizationBase.pick_vector_width(T)
    micro_m % N == 0 || error("`micro_m` must be an integer multiple of vector register width for efficient microkernel generation, got width = $N, micro_m = $micro_m.")
    mregister = max(1, micro_m ÷ N)

    mloopexpr = quote
        ii = cachei₁
        m′ = $micro_m
        while (m = cachei₂ - ii + 1) >= $micro_m
            Coffset = (ii - 1) * cs1 + (jj - 1) * cs2
            Aoffset′ = Aoffset + (packa ? (ii - cachei₁) * ps : (ii - 1) * as1)
            Boffset′ = Boffset + (packb ? (jjfull - cachej₁) * ps + micro_n_offset : (jj - 1) * bs2)
            microkernel!(C, A, B, α, β, Coffset, Aoffset′, Boffset′, ps, (Val(m′), Val(n′)), Val(4), nothing, ctx) # no mask
            ii += m′
        end
        m = cachei₂ - ii + 1
        mask = VectorizationBase.mask($T, m)
        Coffset = (ii - 1) * cs1 + (jj - 1) * cs2
        Aoffset′ = Aoffset + (packa ? (ii - cachei₁) * ps : (ii - 1) * as1)
        Boffset′ = Boffset + (packb ? (jjfull - cachej₁) * ps + micro_n_offset : (jj - 1) * bs2)
        @nif $(mregister+1) midx -> begin
            m′ = $micro_m - (midx - 1) * $N
            m > m′ - $N
        end midx -> begin
            microkernel!(C, A, B, α, β, Coffset, Aoffset′, Boffset′, ps, (Val(m′), Val(n′)), Val(1), mask, ctx) # mask
        end midx -> nothing
        jj += n′
    end

    nloopexpr = quote
        cachei₁, cachei₂ = cacheis
        cachep₁, cachep₂ = cacheps
        cachej₁, cachej₂ = cachejs
        ps = cachep₂ - cachep₁ + 1
        @unpack cs1, cs2, as1, as2, bs1, bs2 = ctx
        jj = jjfull = cachej₁
        micro_n_offset = 0
        Aoffset = packa ? 0 : (cachep₁ - 1) * as2
        Boffset = packb ? 0 : (cachep₁ - 1) * bs1
        n′ = $micro_n
        while (n = cachej₂ - jj + 1) >= n′
            $mloopexpr
            jjfull = jj
        end # full n
        jjfull = jj # set `jjfull` after completing a full `micro_n`

        @nif $micro_n nidx -> begin
            n′ = $micro_n - nidx
            n = cachej₂ - jj + 1
            n >= n′
        end nidx -> begin
            $mloopexpr
            micro_n_offset += n′
        end nidx -> nothing
        return nothing
    end
    # macroexpand(@__MODULE__, nloopexpr)
    return nloopexpr
end

###
### Packing
###

function packAbuffer!(Abuffer, A, cachei₁, cachei₂, cachep₁, cachep₂, ::Val{micro_m}) where micro_m
    # terminology:
    # a `micro_m × 1` vector is a panel
    # a `micro_m × cache_k` matrix is a block
    l = 0 # write position in packing buffer
    @inbounds @aliasscope for i₁ in cachei₁:micro_m:cachei₂ # iterate over blocks
        i₂ = i₁ + micro_m - 1
        if i₂ <= cachei₂ # full panel
            # copy a panel to the buffer contiguously
            @simd for p in cachep₁:cachep₂
                @unroll for i in i₁:i₂ # has constant loop count
                    Abuffer[l += 1] = constarray(A)[i, p]
                end
            end
        else # a panel is not full
            for p in cachep₁:cachep₂ # iterate through panel columns
                for i in i₁:cachei₂  # iterate through live panel rows
                    Abuffer[l += 1] = constarray(A)[i, p]
                end
                for i in (cachei₂ + 1):i₂ # pad the rest of the panel with zero
                    Abuffer[l += 1] = zero(eltype(Abuffer))
                end
            end
        end
    end
    return nothing
end

function packBbuffer!(Bbuffer, B, cachep₁, cachep₂, cachej₁, cachej₂, ::Val{micro_n}) where micro_n
    # terminology:
    # a `1 × micro_n` vector is a panel
    # a `cache_k × micro_n` matrix is a block
    l = 0 # write position in packing buffer
    @inbounds @aliasscope for j₁ in cachej₁:micro_n:cachej₂ # iterate over blocks
        j₂ = j₁ + micro_n - 1
        if j₂ <= cachej₂ # full panel
            # copy a panel to the buffer contiguously
            @simd for p in cachep₁:cachep₂
                @unroll for j in j₁:j₂ # has constant loop count
                    Bbuffer[l += 1] = constarray(B)[p, j]
                end
            end
        else # a panel is not full
            for p in cachep₁:cachep₂ # iterate through panel columns
                for j in j₁:cachej₂  # iterate through live panel rows
                    Bbuffer[l += 1] = constarray(B)[p, j]
                end
                for j in (cachej₂ + 1):j₂ # pad the rest of the panel with zero
                    Bbuffer[l += 1] = zero(eltype(Bbuffer))
                end
            end
        end
    end
    return nothing
end

###
### Micro kernel
###

"""
    microkernel!(C, Abuffer, Bbuffer, α, β, Coffset, Aoffset, Boffset, m, ps,
        kernel_params::Tuple{Val{micro_m},Val{micro_n}}) where {micro_m,micro_n}

A kernel that computes
```julia
Ĉ[1:m, 1:n] = α * Â[1:m, 1:ps] * B̂[1:ps, 1:n] + β * Ĉ[1:m, 1:n]
```
where ``{⋅̂}`` denotes the matrix with offset.
"""
@generated function microkernel!(C, A, B, α, β, Coffset, Aoffset, Boffset, ps, ::Tuple{Val{micro_m},Val{micro_n}},::Val{unroll}, mask,
                                 ctx::GemmContext{packa,packb,fullmicro_m,fullmicro_n})::Nothing where {micro_m,micro_n,unroll,packa,packb,fullmicro_m,fullmicro_n}
    T = eltype(C)
    TA = eltype(A)
    TB = eltype(B)
    T === TA === TB || throw(ArgumentError("eltype(C) === eltype(A) === eltype(B) must hold in the microkernel! Got eltype(C) = $T, eltype(A) = $TA, eltype(B) = $TB."))
    N = VectorizationBase.pick_vector_width(T)
    mregister = max(1, micro_m ÷ N)
    V = SVec{N,T}
    dounroll = unroll > 1
    matA = !packa
    matB = !packb
    usemask = mask != Nothing

    nounroll_loopexpr = quote
        prefetcht0(ptrÂ + 64 * $fullmicro_m + 2 * $fullmicro_n)
        @nexprs $mregister m̂ -> begin
            # assumption: A has unit leading stride
            A_m̂ = vload($V, ptrÂ + (m̂ - 1) * $N * st, mask_m̂)
        end
        @nexprs $micro_n n̂ -> begin
            if $matB
                B_n̂ = $V(vload(ptrB̂ + (n̂ - 1) * bs2))
            else
                B_n̂ = $V(vload(ptrB̂ + (n̂ - 1) * st))
            end
            @nexprs $mregister m̂ -> begin
                AB_m̂_n̂ = vfmadd231(A_m̂, B_n̂, AB_m̂_n̂)
            end
        end
        ptrÂ += $matA ? as2 : $fullmicro_m * st
        ptrB̂ += $matB ? bs1 : $fullmicro_n * st
    end

    unroll_loopexpr = quote
        @nexprs $unroll u -> begin
            #TODO
            if !$matA
                u == 1 && prefetcht0(ptrÂ + 64 * $micro_m + 2 * $fullmicro_n)
                u == 3 && prefetcht0(ptrÂ + 76 * $micro_m + 2 * $fullmicro_n)
            end
            ## unroll variable: u
            @nexprs $mregister m̂ -> begin
                # assumption: A has unit leading stride, as1 == 1
                ptrA′ = ptrÂ + (u - 1) * ($matA ? as2 : $fullmicro_m * st)
                A_m̂ = vload($V, ptrA′ + (m̂ - 1) * $N * st, mask_m̂)
            end
            @nexprs $micro_n n̂ -> begin
                ptrB′ = ptrB̂ + (n̂ - 1) * ($matB ? bs2 : st) + (u - 1) * ($matB ? bs1 : $fullmicro_n * st)
                B_n̂ = $V(unsafe_load(ptrB′))
                @nexprs $mregister m̂ -> begin
                    AB_m̂_n̂ = vfmadd231(A_m̂, B_n̂, AB_m̂_n̂)
                end
            end
        end
        ptrÂ += $matA ? $unroll * as2 : $unroll * $fullmicro_m * st
        ptrB̂ += $matB ? $unroll * bs1 : $unroll * $fullmicro_n * st
    end

    kernel_code = quote
        # tell the compiler that the iteration is nonempty
        SIMDPirates.assume(ps > 0)

        st = sizeof($T)
        cs1 = 1 * st # ctx.cs1 * st
        cs2 = ctx.cs2 * st
        if $matA # A is not packed
            as1 = 1 * st # ctx.as1 * st
            as2 = ctx.as2 * st
        end
        if $matB # B is not packed
            bs1 = ctx.bs1 * st
            bs2 = ctx.bs2 * st
        end
        @nexprs $mregister m̂ -> begin
            mask_m̂ = ($usemask && m̂ == $mregister) ? mask : # nontrival mask
                                                     VectorizationBase.max_mask($T) # trival mask that should be optimized away
        end

        ptrĈ = pointer(C) + Coffset * st
        ptrÂ = pointer(A) + Aoffset * st
        ptrB̂ = pointer(B) + Boffset * st

        # prefetch C matrix
        @nexprs $micro_n n̂ -> begin
            prefetcht0(ptrĈ + (n̂ - 1) * cs2 + 7 * 8)
        end

        # intializing AB registers
        @nexprs $micro_n n̂ -> @nexprs $mregister m̂ -> AB_m̂_n̂ = zero($V)

        # rank-1 updates
        if $dounroll
            punroll, pleft = divrem(ps, $unroll)
            # tell the compiler that the iteration is nonempty
            SIMDPirates.assume(punroll > 0 || pleft > 0)
            for _ in 1:punroll
                $unroll_loopexpr
            end

            for _ in 1:pleft
                $nounroll_loopexpr
            end
        else
            for _ in 1:ps
                $nounroll_loopexpr
            end
        end

        _α = $V(α)
        @nexprs $micro_n n̂ -> @nexprs $mregister m̂ -> begin
            C_m̂_n̂ = _α * AB_m̂_n̂
        end

        if iszero(β)
            @nexprs $micro_n n̂ -> @nexprs $mregister m̂ -> begin
                addr = ptrĈ + (m̂ - 1) * $N * cs1 + (n̂ - 1) * cs2
                vstore!(addr, C_m̂_n̂, mask_m̂)
            end
        else
            _β = $V(β)
            @nexprs $micro_n n̂ -> @nexprs $mregister m̂ -> begin
                addr = ptrĈ + (m̂ - 1) * $N * cs1 + (n̂ - 1) * cs2
                C_m̂_n̂ = vfmadd231(_β, vload($V, addr, mask_m̂), C_m̂_n̂)
                vstore!(addr, C_m̂_n̂, mask_m̂)
            end
        end
    end

    expr = quote
        begin
            $(Expr(:meta, usemask ? :noinline : :inline))
            $kernel_code
            return nothing
        end
    end
    return expr
end

###
### Fast stride handling
###

struct One
end
@inline Base.:*(::One, x::Int) = x
@inline Base.:*(x::Int, ::One) = x
@inline Base.:+(::One, x::Int) = x + 1
@inline Base.:+(x::Int, ::One) = x + 1
@inline Base.:(==)(x::Int, ::One) = isone(x)
@inline Base.:(==)(::One, x::Int) = isone(x)
@inline _strides(A) = strides(A)
@inline _strides(A::DenseArray) = (One(), stride(A, 2))
@inline _strides(A::SubArray{<:Any,2,<:DenseArray,<:Tuple{Union{Base.Slice, AbstractUnitRange}, Vararg{Any}}}) = (One(), stride(A, 2))
@inline _strides(A::Union{Adjoint,Transpose}) = reverse(_strides(parent(A)))
@inline function _strides(A::PermutedDimsArray{T,N,perm}) where {T,N,perm}
    s = _strides(parent(A))
    return ntuple(d->s[perm[d]], Val(N))
end
@inline canonicalize(A) = A
@inline canonicalize(A::Union{Adjoint,Transpose,PermutedDimsArray}) = canonicalize(parent(A))

###
### Utilities
###

function checkmulsize(C, A, B)
    cm, cn = size(C)
    am, ak = size(A)
    bk, bn = size(B)
    (cm == am && ak == bk && cn == bn) || throw(DimensionMismatch("C has dimensions ($cm, $cn), A has dimensions ($am, $ak), but B has dimensions ($bk, $bn)"))
    return cm, ak, bn
end

prefetcht0(ptr::Ptr) = __prefetch(ptr, Val(:read), Val(3), Val(:data))
# From https://github.com/vchuravy/GPUifyLoops.jl/pull/5
@generated function __prefetch(ptr::T, ::Val{RW}, ::Val{Locality}, ::Val{Cache}) where {T, RW, Locality, Cache}
    decls = """
    declare void @llvm.prefetch(i8*, i32, i32, i32)
    """

    if RW == :read
        f_rw = 0
    elseif RW == :write
        f_rw = 1
    end

    f_locality = Locality

    if Cache == :data
        f_cache = 1
    elseif Cache == :instruction
        f_cache = 0
    end

    ir = """
        %ptr = inttoptr i64 %0 to i8*
        call void @llvm.prefetch(i8* %ptr, i32 $f_rw, i32 $f_locality, i32 $f_cache)
        ret void
    """

    quote
        Base.@_inline_meta
        Base.llvmcall(($decls, $ir), Nothing, Tuple{T}, ptr)
    end
end

@inline align(x::Ptr{T}, n) where {T} = reinterpret(Ptr{T}, align(reinterpret(UInt, x), n))
@inline align(x, n) = (nm1 = n - 1; (x + nm1) & -n)


@with_kw mutable struct DLevelSetData{T <: Real}
    Q::T = 1e10
    Q̃::T = 1e10
    θ::T = -1e10
    timestamp::Int = 1
    iterations::Int = 0
end

@with_kw mutable struct DLevelSetParameters{T <: Real}
    κ::T = 0.3
    τ::T = 1e-6
    λ::T = 0.5
    log::Bool = true
end

struct DLevelSet{T <: Real, A <: AbstractVector, M <: LQSolver, S <: LQSolver} <: AbstractLShapedSolver{T,A,M,S}
    structuredmodel::JuMP.Model
    solverdata::DLevelSetData{T}

    # Master
    mastersolver::M
    projectionsolver::M
    mastervector::A
    c::A
    x::A
    Q_history::A

    # Subproblems
    nscenarios::Int
    subobjectives::Vector{A}
    finished::Vector{Int}

    # Workers
    subworkers::Vector{SubWorker{T,A,S}}
    work::Vector{Work}
    decisions::Decisions{A}
    cutqueue::CutQueue{T}

    # Regularizer
    ξ::A
    Q̃_history::A
    levels::A

    # Cuts
    θs::A
    cuts::Vector{SparseHyperPlane{T}}
    θ_history::A

    # Params
    parameters::DLevelSetParameters{T}
    progress::ProgressThresh{T}

    @implement_trait DLevelSet HasLevels
    @implement_trait DLevelSet IsParallel

    function (::Type{DLevelSet})(model::JuMP.Model,ξ₀::AbstractVector,mastersolver::AbstractMathProgSolver,subsolver::AbstractMathProgSolver; kw...)
        if nworkers() == 1
            warn("There are no worker processes, defaulting to serial version of algorithm")
            return LevelSet(model,ξ₀,mastersolver,subsolver; kw...)
        end
        length(ξ₀) != model.numCols && error("Incorrect length of starting guess, has ",length(ξ₀)," should be ",model.numCols)
        !haskey(model.ext,:SP) && error("The provided model is not structured")

        T = promote_type(eltype(ξ₀),Float32)
        c_ = convert(AbstractVector{T},JuMP.prepAffObjective(model))
        c_ *= model.objSense == :Min ? 1 : -1
        mastervector = convert(AbstractVector{T},copy(ξ₀))
        x₀_ = convert(AbstractVector{T},copy(ξ₀))
        ξ₀_ = convert(AbstractVector{T},copy(ξ₀))
        A = typeof(x₀_)

        msolver = LQSolver(model,mastersolver)
        psolver = LQSolver(model,mastersolver)
        M = typeof(msolver)
        S = LQSolver{typeof(LinearQuadraticModel(subsolver)),typeof(subsolver)}
        n = StochasticPrograms.nscenarios(model)

        lshaped = new{T,A,M,S}(model,
                               DLevelSetData{T}(),
                               msolver,
                               psolver,
                               mastervector,
                               c_,
                               x₀_,
                               A(),
                               n,
                               Vector{A}(),
                               Vector{Int}(),
                               Vector{SubWorker{T,A,S}}(nworkers()),
                               Vector{Work}(nworkers()),
                               RemoteChannel(() -> DecisionChannel(Dict{Int,A}())),
                               RemoteChannel(() -> Channel{QCut{T}}(4*nworkers()*n)),
                               ξ₀_,
                               A(),
                               A(),
                               A(fill(-Inf,n)),
                               Vector{SparseHyperPlane{T}}(),
                               A(),
                               DLevelSetParameters{T}(;kw...),
                               ProgressThresh(1.0, "Distributed Leveled L-Shaped Gap "))
        lshaped.progress.thresh = lshaped.parameters.τ
        push!(lshaped.subobjectives,zeros(n))
        push!(lshaped.finished,0)
        push!(lshaped.Q_history,Inf)
        push!(lshaped.Q̃_history,Inf)
        push!(lshaped.θ_history,-Inf)

        init!(lshaped,subsolver)

        return lshaped
    end
end
DLevelSet(model::JuMP.Model,mastersolver::AbstractMathProgSolver,subsolver::AbstractMathProgSolver; kw...) = DLevelSet(model,rand(model.numCols),mastersolver,subsolver; kw...)

function (lshaped::DLevelSet)()
    # Reset timer
    lshaped.progress.tfirst = lshaped.progress.tlast = time()
    # Start workers
    active_workers = init_workers!(lshaped)
    # Start procedure
    while true
        status = iterate!(lshaped)
        if status != :Valid
            close_workers!(lshaped,active_workers)
            return status
        end
    end
end


using Revise
using FlippingModel
using DataFrames, Optim
# using BrowseTables
using Test
import Distributions: pdf

##
pokes
names(pokes)
pokes[!,:Leave] = Vector{Union{Missing,Bool}}(missing,length(pokes.Side))
pokes[!,:Bout] .= 0
bouts = by(pokes, :Session) do dd
    process_bouts(dd)
end

b = filter(bouts) do row
    !ismissing(row.Leave) &&
    row.Correct_start &&
    row.Start_Reward &&
    row.Streak > 10 &&
    3 < row.Streak_within_Block < 23 &&
    row.ReverseStreak > 20 &&
    (ismissing(row.Pre_Interpoke) || row.Pre_Interpoke < 1) &&
    !ismissing(row.Leave) &&
    row.Protocol_Day > 1 &&
    row.Day > 20191215
end
disallowmissing!(b, :Leave)
##
using CSV
file = "/Volumes/GoogleDrive/My Drive/Flipping/Datasets/Stimulations/WT_new_protocols/Results/bouts.csv"
CSV.write(file,b)
# open_html_table(b[end-100:end,:])


@with_kw mutable struct CrosswalkDriver <: DriverModel{LonAccelDirection}
    a::LonAccelDirection = LonAccelDirection(0., 1)
    navigator::RouteFollowingIDM = RouteFollowingIDM()
    crosswalk::Lane = Lane()
    conflict_lanes::Vector{Lane{Float64}} = Lane{Float64}[]
    intersection_entrances::Vector{Lane{Float64}} = Lane{Float64}[]
    ped_model::ConstantPedestrian = ConstantPedestrian()
    ped_start::Float64 = 4.0 # assume a 5meter buffer before a pedestrian reaches the road
    stop_delta::Float64 = 0.7
    accel_tol::Float64 = 0.1
    d_tol::Float64 = 0.5

    #states
    yield::Bool = true
    priority::Bool = false
    stop::Bool = false
    wait_list::Vector{Int64} = Int64[]

    # transition
    clear::Bool = false

    debug::Bool = false
end

function AutomotiveDrivingModels.reset_hidden_state!(model::CrosswalkDriver)
    model.priority = false
    model.stop = false
    model.clear = false
    model
end

AutomotiveDrivingModels.get_name(::CrosswalkDriver) = "CrosswalkDriver"
Base.rand(model::CrosswalkDriver) = model.a

function AutomotiveDrivingModels.observe!(model::CrosswalkDriver, scene::Scene, roadway::Roadway, egoid::Int)
    ego = scene[findfirst(egoid, scene)]
    AutomotiveDrivingModels.observe!(model.navigator, scene, roadway, egoid) # set the direction
    dir = model.navigator.dir
    a_lon =0.
    a_lon_idm = model.navigator.a
    dist_to_cw = get_distance_to_crosswalk(model, ego, roadway, -model.stop_delta)
    # if !model.yield
    #     a_lon = a_lon_idm
    #     # println("veh ", egoid, " not yielding to crosswalk ", model.crosswalk.tag)
    # else
    #     if !model.priority && !model.stop # reach stop line
    #         a_lon = min(a_lon_idm, stop_at_dist(model, ego, dist_to_cw))
    #     elseif !model.priority && model.stop # wait
    #         a_lon = -model.navigator.d_max # negative to make sure v stays 0
    #     else
    #         a_lon = a_lon_idm # just idm
    #     end
    # end
    # if !model.stop
    #     update_stop!(model, ego, roadway, dist_to_cw)
    # end
    # grow_wait_list!(model, scene, roadway, egoid)
    # ungrow_wait_list!(model, scene, roadway, egoid)
    # update_priority!(model, scene, roadway, egoid)

    if !model.yield
        a_lon = a_lon_idm
    else
        grow_wait_list!(model, scene, roadway, egoid)
        ungrow_wait_list!(model, scene, roadway, egoid)
        update_priority!(model, scene, roadway, egoid)
        if !model.priority
            a_lon = min(a_lon_idm, AutomotivePOMDPs.stop_at_dist(model, ego, dist_to_cw))
            # println("dist to cw $dist_to_cw, crosswalk $(model.crosswalk.tag), $(get_lane(roadway, car) ∈ model.intersection_entrances)")
            # println("priority $(model.priority), crossing $(is_crossing(ped, model.crosswalk, model.conflict_lanes, roadway)), a_lon $a_lon, v $(car.v)")
        else
            a_lon = a_lon_idm # just idm
        end
    end

    if model.debug
        println("ID ", egoid, " priority ", model.priority, " stop ", model.stop, " a ", a_lon, " wait list ", model.wait_list, " dist_to_cw ", dist_to_cw)
    end
    model.a = LonAccelDirection(a_lon, dir)
end

"""
Check if all the pedestrian have crossed
"""
function update_priority!(model::CrosswalkDriver, scene::Scene, roadway::Roadway, egoid::Int)
    ego = scene[findfirst(egoid, scene)]
    model.priority = has_passed(model, ego, roadway) || isempty(model.wait_list)
end

function has_passed(model::CrosswalkDriver, ego::Vehicle, roadway::Roadway)
    cw_length = get_end(model.crosswalk)
    # cw_width = model.crosswalk.width
    cw_center = posg(Frenet(model.crosswalk, cw_length/2, 0.), roadway)
    cw_to_car = get_front(ego) - cw_center
    car_vec = get_front(ego) - ego.state.posG
    has_passed = dot(cw_to_car, car_vec) > 0.
    return has_passed
end

function has_passed(model::CrosswalkDriver, ego::VehicleState, roadway::Roadway)
    cw_length = get_end(model.crosswalk)
    cw_center = posg(Frenet(model.crosswalk, cw_length/2), roadway)
    cw_to_car = ego.posG - cw_center
    car_vec = polar(1., ego.posG.θ) - ego.posG + ego.posG
    has_passed = dot(cw_to_car, car_vec) > 0.
    return has_passed
end

function grow_wait_list!(model::CrosswalkDriver, scene::Scene, roadway::Roadway, egoid::Int)
    for veh in scene
        if veh.def.class == AgentClass.PEDESTRIAN &&
           is_crossing(veh, model.crosswalk, model.conflict_lanes, scene, roadway) &&
           !(veh.id ∈ model.wait_list)
            push!(model.wait_list, veh.id)
        end
    end
end

function ungrow_wait_list!(model::CrosswalkDriver, scene::Scene, roadway::Roadway, egoid::Int)
    n = length(model.wait_list)
    to_remove = Int64[]
    for i=1:n
        pedid = model.wait_list[i]
        ped_ind = findfirst(pedid, scene)
        if ped_ind == nothing
            push!(to_remove, i)
            continue
        end
        ped = scene[ped_ind]
        if !is_crossing(ped, model.crosswalk, model.conflict_lanes, scene, roadway)
            push!(to_remove, i)
        end
    end
    deleteat!(model.wait_list, to_remove)
end

function is_crossing(ped::Vehicle, crosswalk::Lane{T}, conflict_lanes::Vector{Lane{T}}, scene::Scene, roadway::Roadway) where T
    # check if the pedestrian is in the conflict zone
    ped_lane = get_lane(roadway, ped)
    if ped_lane.tag != crosswalk.tag
        return false
    end
    for lane in conflict_lanes
        ped_f = Frenet(proj(ped.state.posG, lane, roadway, move_along_curves=false), roadway)
        if abs(ped_f.t) < lane.width/2 && get_lane(roadway, ped).tag == crosswalk.tag
            return true
        end
    end
    # at this point, the pedestrian is not on the road
    # check if the pedestrian is going to cross or not
    if AutomotivePOMDPs.direction_from_center(ped, crosswalk) > 0. && get_lane(roadway, ped).tag == crosswalk.tag && is_free(crosswalk, scene, roadway) #&& !isapprox(ped.state.v, 0.)
        return true
    end
    return false
end

function is_free(crosswalk::Lane, scene::Scene, roadway::Roadway)
    #check that no vehicle is stopped on the crosswalk
    for veh in scene
        if veh.def.class == AgentClass.PEDESTRIAN
            continue
        end
        posG = veh.state.posG
        cw_posF = Frenet(proj(posG, crosswalk, roadway,move_along_curves=false), roadway)
        lane = get_lane(roadway, veh)
        t_car = cw_posF.t
        if abs(cw_posF.t) < crosswalk.width/2 + veh.def.length/2*sin(cw_posF.ϕ)
            return false
        end
    end
    return true
end

function get_distance_to_crosswalk(model::CrosswalkDriver, veh::VehicleState, roadway::Roadway, delta::Float64 = 0.)
    lane = get_lane(roadway, veh)
    Δs = 0.
    # if lane ∈ model.conflict_lanes
    #     cw_length = get_end(model.crosswalk)
    #     cw_center = posg(Frenet(model.crosswalk, cw_length/2), roadway)
    #     collision_point = VecSE2(cw_center.x, veh.posG.y)
    #     cw_proj = Frenet(collision_point, lane, roadway)
    #     Δs = cw_proj.s - veh.posF.s + delta - model.crosswalk.width
    # elseif lane ∈ model.intersection_entrances # stop before the intersection to not block traffic
    #     Δs = get_end(lane) - veh.posF.s + delta - model.crosswalk.width
    # end
    if lane ∈ model.intersection_entrances # stop before the intersection to not block traffic
        Δs = get_end(lane) - veh.posF.s + delta - model.crosswalk.width
    elseif lane ∈ model.conflict_lanes
        cw_length = get_end(model.crosswalk)
        cw_start = model.crosswalk.curve[1].pos
        cw_center = posg(Frenet(model.crosswalk, cw_length/2), roadway)
        cw_to_car = veh.posG - cw_center
        car_vec = polar(1., veh.posG.θ) - veh.posG + veh.posG
        d = dot(car_vec, cw_to_car)
        Δs = d + delta - model.crosswalk.width/2
    end
    return Δs
end

function get_distance_to_crosswalk(model::CrosswalkDriver, veh::Vehicle, roadway::Roadway, delta::Float64 = 0.)
    return get_distance_to_crosswalk(model, veh.state, roadway, delta)
end

function update_stop!(model::CrosswalkDriver, veh::Vehicle, roadway::Roadway, dist_to_cw::Float64)
    update_stop!(model, veh.state, roadway, dist_to_cw)
end

function update_stop!(model::CrosswalkDriver, veh::VehicleState, roadway::Roadway, dist_to_cw::Float64)
    if veh.v ≈ 0. && isapprox(dist_to_cw, 0, atol=1.)
        model.stop = true
    end
end

# # croswalk driver
# @with_kw mutable struct CrosswalkDriver <: DriverModel{LonAccelDirection}
#     a::LonAccelDirection = LonAccelDirection(0., 1)
#     navigator::RouteFollowingIDM = RouteFollowingIDM()
#     crosswalk::Lane = Lane()
#     intersect::Lane = Lane()
#     ped_model::ConstantPedestrian = ConstantPedestrian()
#     ped_start::Float64 = 4.0 # assume a 5meter buffer before a pedestrian reaches the road
#     stop_delta::Float64 = 0.0
#     accel_tol::Float64 = 0.1
#     d_tol::Float64 = 0.5
#
#     #states
#     go::Bool = false
#     stop::Bool = false
#
#     # transition
#     clear::Bool = false
#
#     debug::Bool = false
# end
#
# AutomotiveDrivingModels.get_name(::CrosswalkDriver) = "CrosswalkDriver"
# Base.rand(model::CrosswalkDriver) = model.a
#
# function AutomotiveDrivingModels.reset_hidden_state!(model::CrosswalkDriver)
#     model.go = false
#     model.stop = false
#     model
# end
#
# function AutomotiveDrivingModels.observe!(model::CrosswalkDriver, scene::Scene, roadway::Roadway, egoid::Int)
#     ego = scene[findfirst(egoid, scene)]
#     AutomotiveDrivingModels.observe!(model.navigator, scene, roadway, egoid) # set the direction
#     dir = model.navigator.dir
#     a_lon =0.
#     a_lon_idm = model.navigator.a
#     if model.go
#         a_lon = a_lon_idm
#     else
#         update_clear!(model, ego, scene, roadway)
#         if model.clear
#             model.go = true
#         else # brake!
#             model.clear = false
#             d = dist_to_stop(model, ego, roadway)
#             if model.debug
#                 println("braking! d = ", d, "a = ", stop_at_dist(model, ego, d))
#             end
#             a_lon = min(a_lon_idm, stop_at_dist(model, ego, d))
#             if ego.state.v ≈ 0. && d < model.d_tol
#                 a_lon = 0.
#             end
#         end
#     end
#     if model.debug
#         println("veh ", egoid, " go ", model.go, " stop ", model.stop, " clear ", model.clear)
#     end
#     model.a = LonAccelDirection(a_lon, dir)
# end
#
#
# function update_clear!(model::CrosswalkDriver, ego::Vehicle, scene::Scene, roadway::Roadway)
#     # check how many pedestrians are present
#     peds = Vehicle[]
#     for veh in scene
#         if veh.def.class == AgentClass.PEDESTRIAN && get_lane(roadway, veh) == model.crosswalk
#             push!(peds, veh)
#         end
#     end
#     cw_length = get_end(model.crosswalk)
#     cw_center = posg(Frenet(model.crosswalk, cw_length/2), roadway)
#     collision_point = VecSE2(cw_center.x, ego.state.posG.y)
#     # set to the maximum value
#     buffer_ped = VehicleState(Frenet(model.crosswalk, model.ped_start), roadway,
#                               model.ped_model.v_desired)
#     # println("compute time ped")
#     min_ped_ttc = compute_time_of_approach(model, collision_point, buffer_ped, roadway, -model.intersect.width/2)
#     if model.debug
#         println("buffer ped ttc ", min_ped_ttc)
#     end
#     ego_ttc = compute_time_of_approach(model, collision_point, ego.state, roadway,
#                                        +model.crosswalk.width/2 + ego.def.length/2)
#     if ego.state.v ≈ 0.
#         ego_ttc = sqrt((model.crosswalk.width + ego.def.length)/model.navigator.a_max/2) # it is stopped
#     end
#     if model.debug
#         println("ego ttc ", ego_ttc)
#     end
#     has_crossed = -0.5*model.intersect.width/model.ped_model.v_desired
#     if model.debug
#         println("has_crossed ttc", has_crossed)
#     end
#     for ped in peds
#         ttc = compute_time_of_approach(model, collision_point, ped.state, roadway, -model.intersect.width/2)
#         if model.debug
#             println("ped $(ped.id) ttc ", ttc)
#         end
#         if has_crossed < ttc < min_ped_ttc
#             min_ped_ttc = ttc
#         end
#     end
#     if model.debug
#         println("min ped ttc ", min_ped_ttc)
#     end
#     if has_crossed < min_ped_ttc < ego_ttc
#         model.clear = false
#     else
#         model.clear = true
#     end
# end
#
# """
# compute the time to reach the projection of cw on the ego lane
# """
# function compute_time_of_approach(model::CrosswalkDriver, cw::VecSE2, ego::VehicleState, roadway::Roadway, delta::Float64=0.)
#     lane = get_lane(roadway, ego)
#     cw_proj = Frenet(cw, lane, roadway)
#     Δs = cw_proj.s - ego.posF.s + delta
#     # println("lane ", lane)
#     # println("cw_proj ", cw_proj.s)
#     # println("delta ", delta)
#     # println("ego posF", ego.posF.s)
#     # println("ds ", Δs)
#     v = max(ego.v, model.ped_model.v_desired)
#     time_of_approach = Δs/v
# end
#
# function dist_to_stop(model::CrosswalkDriver, ego::Vehicle, roadway::Roadway)
#     cw_length = get_end(model.crosswalk)
#     cw_center = posg(Frenet(model.crosswalk, cw_length/2), roadway)
#     collision_point = VecSE2(cw_center.x, ego.state.posG.y)
#     intersect_pt = Frenet(collision_point, model.intersect, roadway)
#     return intersect_pt.s - model.crosswalk.width/2 - ego.state.posF.s - ego.def.length/2 - model.stop_delta
# end
#
# function is_crossing(ped::Vehicle, crosswalk::Lane, roadway::Roadway)
# end


# Default implmenentation for CuArrays
using CUDA: @workspace, @argout

memsize(a::CuArray) = sizeof(a)

scalar(C::CuArray) = ndims(C)==0 ? collect(C)[] : throw(DimensionMismatch())

function add!(α, A::CuArray{<:Any, N}, CA::Symbol,
                β, C::CuArray{<:Any, N}, indCinA) where {N}

    T = eltype(C)
    N == length(indCinA) || throw(IndexError("Invalid permutation of length $N: $indCinA"))
    CA == :N || CA == :C ||
        throw(ArgumentError("Value of conjA should be :N or :C instead of $CA"))
    opA = (T <: Real || CA == :N) ? CUTENSOR_OP_IDENTITY : CUTENSOR_OP_CONJ
    opAC = CUTENSOR_OP_ADD

    descA = CuTensorDescriptor(A; op = opA)
    descC = CuTensorDescriptor(C; op = CUTENSOR_OP_IDENTITY)
    descD = descC
    typeCompute = convert(cudaDataType, T)
    modeA = collect(Cint, 1:N)
    modeC = collect(Cint, indCinA)
    stream = CuDefaultStream()
    if β == zero(β)
        cutensorPermutation(handle(), T[α], A, descA, modeA, C, descC, modeC,
                            typeCompute, stream)
    else
        cutensorElementwiseBinary(handle(), T[α], A, descA, modeA, T[β], C, descC,
                                  modeC, C, descC, modeC, opAC, typeCompute, stream)
    end

    return C
end

# TODO: the following seems broken due to bug in CUTENSOR?#
# function add!(α, A::CuArray{T, N}, CA::Symbol,
#                 β, C::CuArray{Complex{T}, N}, indCinA) where {T<:CublasReal,N}
#
#     N == length(indCinA) || throw(IndexError("Invalid permutation of length $N: $indCinA"))
#     CA == :N || CA == :C ||
#         throw(ArgumentError("Value of conjA should be :N or :C instead of $CA"))
#     opid = CUTENSOR_OP_IDENTITY
#     opAC = CUTENSOR_OP_ADD
#
#     if β == zero(β)
#         fill!(C, zero(Complex{T}))
#     elseif β != one(β)
#         rmul!(C, β)
#     end
#     sizeC = size(C)
#     stridesC = 2 .* strides(C)
#     Cr = reinterpret(T, C, (2*length(C),))
#
#     @show stridesC, sizeC
#
#     descA = CuTensorDescriptor(A; op = opid)
#     descCr = CuTensorDescriptor(Cr; size = sizeC, strides = stridesC, op = opid)
#     typeCompute = cudaDataType(T)
#     modeA = collect(Cint, 1:N)
#     modeC = collect(Cint, indCinA)
#     stream = CuDefaultStream()
#     cutensorElementwiseBinary(handle(), T[real(α)], A, descA, modeA, T[1], Cr, descCr,
#                               modeC, Cr, descCr, modeC, opAC, typeCompute, stream)
#     if imag(α) != 0
#         Ci = view(Cr, 2:length(Cr))
#         descCi = CuTensorDescriptor(Ci; size = sizeC, strides = stridesC, op = opid)
#         cutensorElementwiseBinary(handle(), T[imag(α)], A, descA, modeA, T[1], Ci, descCi,
#                                   modeC, Ci, descCi, modeC, opAC, typeCompute, stream)
#     end
#     return C
# end

function trace!(α, A::CuArray, CA::Symbol, β, C::CuArray,
                indCinA, cindA1, cindA2)

    T = eltype(C)
    NA, NC = ndims(A), ndims(C)
    NC == length(indCinA) ||
        throw(IndexError("Invalid selection of $NC out of $NA: $indCinA"))
    NA-NC == 2*length(cindA1) == 2*length(cindA2) ||
        throw(IndexError("invalid number of trace dimension"))

    opA = (T <: Real || CA == :N) ? CUTENSOR_OP_IDENTITY : CUTENSOR_OP_CONJ
    opReduce = CUTENSOR_OP_ADD

    sizeA = i->size(A, i)
    strideA = i->stride(A, i)
    tracesize = sizeA.(cindA1)
    tracesize == sizeA.(cindA2) || throw(DimensionMismatch("non-matching trace sizes"))
    size(C) == sizeA.(indCinA) || throw(DimensionMismatch("non-matching sizes"))

    newstrides = (strideA.(indCinA)..., (strideA.(cindA1) .+ strideA.(cindA2))...)
    newsize = (size(C)..., tracesize...)
    descA = CuTensorDescriptor(A; op = opA, size = newsize, strides = newstrides)
    descC = CuTensorDescriptor(C; op = CUTENSOR_OP_IDENTITY)
    descD = descC
    typeCompute = cutensorComputeType(T)
    modeA = collect(Cint, 1:NA)
    modeC = collect(Cint, 1:NC)
    stream = CuDefaultStream()
    @workspace fallback=1<<13 size=@argout(
            cutensorReductionGetWorkspace(handle(),
                A, descA, modeA,
                C, descC, modeC,
                C, descC, modeC,
                opReduce, typeCompute,
                out(Ref{UInt64}(C_NULL)))
        )[] workspace->begin
            cutensorReduction(handle(),
                T[α], A, descA, modeA,
                T[β], C, descC, modeC,
                        C, descC, modeC,
                opReduce, typeCompute,
                workspace, sizeof(workspace), stream)
        end
    return C
end

function contract!(α, A::CuArray, CA::Symbol,
                    B::CuArray, CB::Symbol,
                    β, C::CuArray,
                    oindA::IndexTuple, cindA::IndexTuple,
                    oindB::IndexTuple, cindB::IndexTuple,
                    indCinoAB::IndexTuple, syms::Union{Nothing, NTuple{3,Symbol}} = nothing)

    pA = (oindA...,cindA...)
    (length(pA) == ndims(A) && isperm(pA)) ||
        throw(IndexError("invalid permutation of length $(ndims(A)): $pA"))
    pB = (oindB...,cindB...)
    (length(pB) == ndims(B) && isperm(pB)) ||
        throw(IndexError("invalid permutation of length $(ndims(B)): $pB"))
    (length(oindA) + length(oindB) == ndims(C)) ||
        throw(IndexError("non-matching output indices in contraction"))
    (ndims(C) == length(indCinoAB) && isperm(indCinoAB)) ||
        throw(IndexError("invalid permutation of length $(ndims(C)): $indCinoAB"))

    sizeA = i->size(A, i)
    sizeB = i->size(B, i)
    sizeC = i->size(C, i)

    csizeA = sizeA.(cindA)
    csizeB = sizeB.(cindB)
    osizeA = sizeA.(oindA)
    osizeB = sizeB.(oindB)

    csizeA == csizeB ||
        throw(DimensionMismatch("non-matching sizes in contracted dimensions"))
    sizeAB = let osize = (osizeA..., osizeB...)
        i->osize[i]
    end
    sizeAB.(indCinoAB) == size(C) ||
        throw(DimensionMismatch("non-matching sizes in uncontracted dimensions"))

    TC = eltype(C)

    CA == :N || CA == :C ||
        throw(ArgumentError("Value of conjA should be :N or :C instead of $CA"))
    CB == :N || CB == :C ||
        throw(ArgumentError("Value of conjB should be :N or :C instead of $CB"))

    opA = (TC <: Real || CA == :N) ? CUTENSOR_OP_IDENTITY : CUTENSOR_OP_CONJ
    opB = (TC <: Real || CB == :N) ? CUTENSOR_OP_IDENTITY : CUTENSOR_OP_CONJ
    opC = CUTENSOR_OP_IDENTITY

    strideA = i->stride(A, i)
    strideB = i->stride(B, i)

    cstrideA = strideA.(cindA)
    cstrideB = strideB.(cindB)
    ostrideA = strideA.(oindA)
    ostrideB = strideB.(oindB)

    descA = CuTensorDescriptor(A; op = opA, size = (osizeA..., csizeA...),
                                    strides = (ostrideA..., cstrideA...))
    descB = CuTensorDescriptor(B; op = opB, size = (osizeB..., csizeB...),
                                    strides = (ostrideB..., cstrideB...))
    descC = CuTensorDescriptor(C)
    T = eltype(C)
    typeCompute = cutensorComputeType(T)
    opOut = CUTENSOR_OP_IDENTITY

    NoA = length(osizeA)
    NoB = length(osizeB)
    Nc = length(csizeA)
    modeoA = ntuple(n->n, NoA)
    modeoB = ntuple(n->NoA+n, NoB)
    modec = ntuple(n->NoA+NoB+n, Nc)

    modeA = collect(Cint, (modeoA..., modec...))
    modeB = collect(Cint, (modeoB..., modec...))
    modeC = collect(Cint, indCinoAB)

    algo = CUTENSOR_ALGO_DEFAULT
    stream = CuDefaultStream()
    pref = CUTENSOR_WORKSPACE_RECOMMENDED

    alignmentRequirementA = Ref{UInt32}(C_NULL)
    cutensorGetAlignmentRequirement(handle(), A, descA, alignmentRequirementA)
    alignmentRequirementB = Ref{UInt32}(C_NULL)
    cutensorGetAlignmentRequirement(handle(), B, descB, alignmentRequirementB)
    alignmentRequirementC = Ref{UInt32}(C_NULL)
    cutensorGetAlignmentRequirement(handle(), C, descC, alignmentRequirementC)
    desc = Ref{cutensorContractionDescriptor_t}()
    cutensorInitContractionDescriptor(handle(),
                                      desc,
                   descA, modeA, alignmentRequirementA[],
                   descB, modeB, alignmentRequirementB[],
                   descC, modeC, alignmentRequirementC[],
                   descC, modeC, alignmentRequirementC[],
                   typeCompute)

    find = Ref{cutensorContractionFind_t}()
    cutensorInitContractionFind(handle(), find, algo)

    @workspace fallback=1<<27 size=@argout(
            cutensorContractionGetWorkspace(handle(), desc, find, pref,
                                            out(Ref{UInt64}(C_NULL)))
        )[] workspace->begin
            plan_ref = Ref{cutensorContractionPlan_t}()
            cutensorInitContractionPlan(handle(), plan_ref, desc, find, sizeof(workspace))

            cutensorContraction(handle(), plan_ref, T[α], A, B, T[β], C, C,
                                workspace, sizeof(workspace), stream)
        end

    return C
end

# overwrite similar_from_indices to return zero initialized arrays
function similar_from_indices(T::Type, ind::IndexTuple, A::CuArray, CA::Symbol)
    sz = similarstructure_from_indices(T, ind, A, CA)
    return fill!(similar(A, T, sz), zero(T))
end
function similar_from_indices(T::Type, poA::IndexTuple, poB::IndexTuple,
                                p1::IndexTuple, p2::IndexTuple,
                                A::CuArray, B::CuArray, CA::Symbol, CB::Symbol)
    sz = similarstructure_from_indices(T, poA, poB, p1, p2, A, B, CA, CB)
    return fill!(similar(A, T, sz), zero(T))
end

# overwrite cached_similar_from_indices in order not to use cache for CuArray objects
function cached_similar_from_indices(sym::Symbol, T::Type, p1::IndexTuple, p2::IndexTuple, A::CuArray, CA::Symbol)
    # also zero fill to avoid problems with cutensorElementwiseBinary
    return similar_from_indices(T, p1, p2, A, CA)
end

function cached_similar_from_indices(sym::Symbol, T::Type, poA::IndexTuple, poB::IndexTuple,
    p1::IndexTuple, p2::IndexTuple, A::CuArray, B::CuArray, CA::Symbol, CB::Symbol)
    # also zero fill to avoid problems with cutensorElementwiseBinary
    return similar_from_indices(T, poA, poB, p1, p2, A, B, CA, CB)
end


function PSI.construct_device!(
    optimization_container::PSI.OptimizationContainer,
    sys::PSY.System,
    model::PSI.DeviceModel{R, D},
    ::Type{S},
) where {R <: PSY.RenewableGen, D <: RenewableOutageDispatch, S <: PM.AbstractPowerModel}
    devices = PSI.get_available_components(R, sys)

    if !PSI.validate_available_devices(R, devices)
        return
    end

    # Variables
    PSI.add_variables!(optimization_container, PSI.ActivePowerVariable, devices, D())
    PSI.add_variables!(optimization_container, PSI.ReactivePowerVariable, devices, D())

    # Constraints
    PSI.add_constraints!(
        optimization_container,
        PSI.RangeConstraint,
        PSI.ActivePowerVariable,
        devices,
        model,
        S,
        PSI.get_feedforward(model),
    )
    PSI.add_constraints!(
        optimization_container,
        PSI.RangeConstraint,
        PSI.ReactivePowerVariable,
        devices,
        model,
        S,
        PSI.get_feedforward(model),
    )
    outage_constraints!(
        optimization_container,
        devices,
        model,
        S,
        PSI.get_feedforward(model),
    )
    PSI.eedforward!(optimization_container, devices, model, PSI.get_feedforward(model))

    # Cost Function
    PSI.cost_function!(optimization_container, devices, model, S)

    return
end

function PSI.construct_device!(
    optimization_container::PSI.OptimizationContainer,
    sys::PSY.System,
    model::PSI.DeviceModel{R, D},
    ::Type{S},
) where {
    R <: PSY.RenewableGen,
    D <: RenewableOutageDispatch,
    S <: PM.AbstractActivePowerModel,
}
    devices = PSI.get_available_components(R, sys)

    if !PSI.validate_available_devices(R, devices)
        return
    end

    # Variables
    PSI.add_variables!(optimization_container, PSI.ActivePowerVariable, devices, D())

    # Constraints
    PSI.add_constraints!(
        optimization_container,
        PSI.RangeConstraint,
        PSI.ActivePowerVariable,
        devices,
        model,
        S,
        PSI.get_feedforward(model),
    )
    outage_constraints!(
        optimization_container,
        devices,
        model,
        S,
        PSI.get_feedforward(model),
    )
    PSI.feedforward!(optimization_container, devices, model, PSI.get_feedforward(model))

    # Cost Function
    PSI.cost_function!(optimization_container, devices, model, S)

    return
end


"""
Bridges `CP.GlobalCardinality` to `CP.Count`.
"""
struct GlobalCardinalityFixedOpen2CountBridge{T} <: MOIBC.AbstractBridge
    cons_count::Vector{MOI.ConstraintIndex{MOI.VectorAffineFunction{T}, CP.Count{MOI.EqualTo{T}}}}
end

function MOIBC.bridge_constraint(
    ::Type{GlobalCardinalityFixedOpen2CountBridge{T}},
    model,
    f::MOI.VectorOfVariables,
    s::CP.GlobalCardinality{CP.FIXED_COUNTED_VALUES, CP.OPEN_COUNTED_VALUES, T},
) where {T}
    return MOIBC.bridge_constraint(
        GlobalCardinalityFixedOpen2CountBridge{T},
        model,
        MOI.VectorAffineFunction{T}(f),
        s,
    )
end

function MOIBC.bridge_constraint(
    ::Type{GlobalCardinalityFixedOpen2CountBridge{T}},
    model,
    f::MOI.VectorAffineFunction{T},
    s::CP.GlobalCardinality{CP.FIXED_COUNTED_VALUES, CP.OPEN_COUNTED_VALUES, T},
) where {T}
    f_scalars = MOIU.scalarize(f)
    f_array = f_scalars[1:s.dimension]

    cons_count = MOI.ConstraintIndex{MOI.VectorAffineFunction{T}, CP.Count{MOI.EqualTo{T}}}[
        MOI.add_constraint(
            model,
            MOIU.vectorize(
                MOI.ScalarAffineFunction{T}[
                    f_scalars[s.dimension + i],
                    f_array...,
                ]
            ),
            CP.Count(s.dimension, MOI.EqualTo(s.values[i]))
        )
        for i in 1:length(s.values)
    ]

    return GlobalCardinalityFixedOpen2CountBridge(cons_count)
end

function MOI.supports_constraint(
    ::Type{GlobalCardinalityFixedOpen2CountBridge{T}},
    ::Union{Type{MOI.VectorOfVariables}, Type{MOI.VectorAffineFunction{T}}},
    ::Type{CP.GlobalCardinality{CP.FIXED_COUNTED_VALUES, CP.OPEN_COUNTED_VALUES, T}},
) where {T}
    return true
end

function MOIB.added_constrained_variable_types(::Type{GlobalCardinalityFixedOpen2CountBridge{T}}) where {T}
    return Tuple{Type}[]
end

function MOIB.added_constraint_types(::Type{GlobalCardinalityFixedOpen2CountBridge{T}}) where {T}
    return [
        (MOI.VectorAffineFunction{T}, CP.Count{MOI.EqualTo{T}}),
    ]
end

function MOI.get(
    b::GlobalCardinalityFixedOpen2CountBridge{T},
    ::MOI.NumberOfConstraints{
        MOI.VectorAffineFunction{T}, CP.Count{MOI.EqualTo{T}},
    },
) where {T}
    return length(b.cons_count)
end

function MOI.get(
    b::GlobalCardinalityFixedOpen2CountBridge{T},
    ::MOI.ListOfConstraintIndices{
        MOI.VectorAffineFunction{T}, CP.Count{MOI.EqualTo{T}},
    },
) where {T}
    return copy(b.cons_count)
end


__precompile__()
module AppleAccelerate

const libacc = "/System/Library/Frameworks/Accelerate.framework/Accelerate"

if !isfile(libacc)
    error("Accelerate framework not found at $(libacc)")
end

include("Array.jl")
include("DSP.jl")
include("Util.jl")

end # module


using FreeTypeAbstraction: height_insensitive_boundingbox

"""
Calculates the exact boundingbox of a Scene/Plot, without considering any transformation
"""
raw_boundingbox(x::Atomic) = data_limits(x)

rootparent(x) = rootparent(parent(x))
rootparent(x::Scene) = x

# function raw_boundingbox(x::Annotations)
#     bb = raw_boundingbox(x.plots)
#     inv(modelmatrix(rootparent(x))) * bb
# end

raw_boundingbox(x::Combined) = raw_boundingbox(x.plots)
boundingbox(x) = raw_boundingbox(x)

function combined_modelmatrix(x)
    m = Mat4f0(I)
    while true
        m = modelmatrix(x) * m
        if parent(x) !== nothing && parent(x) isa Combined
            x = parent(x)
        else
            break
        end
    end
    return m
end

function modelmatrix(x)
    t = transformation(x)
    transformationmatrix(t.translation[], t.scale[], t.rotation[])
end

function boundingbox(x::Atomic)
    bb = raw_boundingbox(x)
    return combined_modelmatrix(x) * bb
end

boundingbox(scene::Scene) = raw_boundingbox(scene)
function raw_boundingbox(scene::Scene)
    if scene[OldAxis] !== nothing
        return raw_boundingbox(scene[OldAxis])
    elseif scene.limits[] !== automatic
        return scene_limits(scene)
    elseif cameracontrols(scene) == EmptyCamera()
        # Empty camera means this is a parent scene that itself doesn't display anything
        return raw_boundingbox(scene.children)
    else
        plots = plots_from_camera(scene)
        children = filter(scene.children) do child
            child.camera == scene.camera
        end
        return raw_boundingbox([plots; children])
    end
end

function raw_boundingbox(plots::Vector)
    isempty(plots) && return FRect3D()
    plot_idx = iterate(plots)
    bb = FRect3D()
    while plot_idx !== nothing
        plot, idx = plot_idx
        plot_idx = iterate(plots, idx)
        # isvisible(plot) || continue
        bb2 = boundingbox(plot)
        isfinite_rect(bb) || (bb = bb2)
        isfinite_rect(bb2) || continue
        bb = union(bb, bb2)
    end
    return bb
end

function project_widths(matrix, vec)
    pr = project(matrix, vec)
    zero = project(matrix, zeros(typeof(vec)))
    return pr - zero
end

function rotate_bbox(bb::FRect3D, rot)
    points = decompose(Point3f0, bb)
    FRect3D(Ref(rot) .* points)
end

function boundingbox(x::Text, text::String, position::VecTypes)

    glyphlayout = x._glyphlayout[]

    pos = to_ndim(Point3f0, x.position[], 0)
    rot = convert_attribute(x.rotation[], key"rotation"())

    if x.space[] == :data
        data_text_boundingbox(x[1][], glyphlayout, rot, pos)
    elseif x.space[] == :screen
        data_limits(x)
    end
end

function data_text_boundingbox(string::String, glyphlayout::Glyphlayout, rotation::Quaternion, position::Point3f0)
    bb = FRect3D()
    glyphorigins, glyphbbs = glyphlayout.origins, glyphlayout.bboxes

    for (char, charo, glyphbb) in zip(string, glyphorigins, glyphbbs)
        # ignore line breaks
        # char in ('\r', '\n') && continue

        charbb = rotate_bbox(FRect3D(glyphbb), rotation) + charo + position
        if !isfinite_rect(bb)
            bb = charbb
        else
            bb = union(bb, charbb)
        end
    end

    bb
end

function boundingbox(x::Text, texts::AbstractArray, positions::AbstractArray)

    layouts = x._glyphlayout[]
    rotations = convert_attribute(x.rotation[], key"rotation"())
    positions_3d = to_ndim.(Point3f0, positions, 0)

    if x.space[] == :data
        bb = FRect3D()
        broadcast_foreach(texts, positions_3d, layouts, rotations) do text, pos, layout, rot
            if !isfinite_rect(bb)
                bb = data_text_boundingbox(text, layout, rot, pos)
            else
                bb = union(bb, data_text_boundingbox(text, layout, rot, pos))
            end
        end
        bb
    elseif x.space[] == :screen
        data_limits(x)
    end
end

function boundingbox(x::Text)

    # TODO: this is only necessary because of the "text in text" recipe weirdness
    if !isempty(x.plots)
        # the "real" text is inside the outer text
        x = x.plots[1]
    end

    boundingbox(x, to_value(x[1]), to_value(x[:position]))
end

#     if x.space[] == :data
#         if x[1][] isa AbstractArray
#             bb = FRect3D()
#             for
#         else
#             boundingbox(x, to_value(x[1]))
#         end
#     elseif x.space[] == :screen
#         data_limits(x)
#     end
# end

# boundingbox(x::Text) = boundingbox(x, to_value(x[1]))

# function boundingbox(
#         text::String, position, textsize;
#         font = "default", align = (:left, :bottom), rotation = 0.0
#     )
#     return boundingbox(
#         text, position, textsize,
#         to_font(font), to_align(align), to_rotation(rotation)
#     )
# end

"""
Calculate an approximation of a tight rectangle around a 2D rectangle rotated by `angle` radians.
This is not perfect but works well enough. Check an A vs X to see the difference.
"""
function rotatedrect(rect::Rect{2}, angle)
    ox, oy = rect.origin
    wx, wy = rect.widths
    points = @SMatrix([
        ox oy;
        ox oy+wy;
        ox+wx oy;
        ox+wx oy+wy;
    ])
    mrot = @SMatrix([
        cos(angle) -sin(angle);
        sin(angle) cos(angle);
    ])
    rotated = mrot * points'

    rmins = minimum(rotated, dims = 2)
    rmaxs = maximum(rotated, dims = 2)

    return Rect2D(rmins..., (rmaxs .- rmins)...)
end

function quaternion_to_2d_angle(quat)
    # this assumes that the quaternion was calculated from a simple 2d rotation as well
    return 2acos(quat[4]) * (signbit(quat[1]) ? -1 : 1)
end

function boundingbox(
        text::String, position, textsize, font,
        align, rotation, model, justification, lineheight;
        # use the font's ascenders and descenders for the bounding box
        # this means that a string's boundingbox doesn't change in the vertical
        # dimension when characters change (for example numbers during an animation)
        # this is not wanted in most cases because of the jitter it creates when
        # the boundingbox slightly changes size in each frame (in MakieLayout mostly)
        use_vertical_dimensions_from_font = true
    )
    atlas = get_texture_atlas()
    N = length(text)
    ctext_state1 = iterate(text)
    ctext_state1 === nothing && return FRect3D()

    # call the layouting algorithm to find out where all the glyphs end up
    # this is kind of a doubling, maybe it could be avoided if at creation all
    # positions would be populated in the text object, but that seems convoluted
    if position isa VecTypes
        glyph_offsets = layout_text(text, textsize, font, align,
            rotation, model, justification, lineheight)
        position = [go .+ to_ndim(Point3f0, position, 0) for go in glyph_offsets]
    end

    bbox = Ref(FRect3D())            # use ref to avoid julia#15276
    ctext_state = Ref{Union{Nothing,typeof(ctext_state1)}}(ctext_state1)

    let position = position
        broadcast_foreach(1:N, rotation, font, textsize) do i, rotation, font, scale
            c, text_state = ctext_state[]::typeof(ctext_state1)
            ctext_state[] = iterate(text, text_state)

            if !(c in ('\r', '\n'))
                bb_unitspace = if use_vertical_dimensions_from_font
                    height_insensitive_boundingbox(
                        FreeTypeAbstraction.get_extent(font, c), font)
                else
                    inkboundingbox(FreeTypeAbstraction.get_extent(font, c))
                end

                scaled_bb = bb_unitspace * scale

                # TODO this only works in 2d
                rot_2d_radians = quaternion_to_2d_angle(rotation)
                rotated_bb = rotatedrect(scaled_bb, rot_2d_radians)

                # bb = rectdiv(bb, 1.5)
                shifted_bb = FRect3D(rotated_bb) + position[i]
                bboxc = bbox[]
                if !isfinite_rect(bboxc)
                    bbox[] = shifted_bb
                else
                    bbox[] = union(bboxc, shifted_bb)
                end
            end
        end
    end
    return bbox[]
end


@testset "DataSubset constructor" begin
    @test_throws DimensionMismatch DataSubset((rand(2,10),rand(9)))
    @test_throws DimensionMismatch DataSubset((rand(2,10),rand(9)),1:2)
    @test_throws DimensionMismatch DataSubset((rand(2,10),rand(4,9,10),rand(9)))

    @testset "bounds check" begin
        for var in (vars..., tuples..., CustomType())
            @test_throws BoundsError DataSubset(var, -1:100)
            @test_throws BoundsError DataSubset(var, 1:151)
            @test_throws BoundsError DataSubset(var, [1, 10, 0, 3])
            @test_throws BoundsError DataSubset(var, [1, 10, -10, 3])
            @test_throws BoundsError DataSubset(var, [1, 10, 180, 3])
        end
    end

    @testset "Tuple unrolling" begin
        @test_throws DimensionMismatch DataSubset((X,X), 1:150, (ObsDim.Last(), ObsDim.Last(), ObsDim.Last()))
        @test_throws DimensionMismatch DataSubset((X,X), 1:150, (ObsDim.Last(),))
        @test_throws DimensionMismatch DataSubset((X,X), (ObsDim.Last(), ObsDim.Last(), ObsDim.Last()))
        @test_throws DimensionMismatch DataSubset((X,X), (ObsDim.Last(),))
        @test typeof(@inferred(DataSubset((X,X)))) <: Tuple
        @test eltype(@inferred(DataSubset((X,X)))) <: DataSubset
        @test typeof(@inferred(DataSubset((X,X), ObsDim.Last()))) <: Tuple
        @test eltype(@inferred(DataSubset((X,X), ObsDim.Last()))) <: DataSubset
        @test typeof(@inferred(DataSubset((X,X), (ObsDim.Last(), ObsDim.Last())))) <: Tuple
        @test eltype(@inferred(DataSubset((X,X), (ObsDim.Last(), ObsDim.Last())))) <: DataSubset
        @test typeof(@inferred(DataSubset((X,X), 1:150, (ObsDim.Last(), ObsDim.Last())))) <: Tuple
        @test eltype(@inferred(DataSubset((X,X), 1:150, (ObsDim.Last(), ObsDim.Last())))) <: DataSubset
        D1 = @inferred(DataSubset((X',X), (ObsDim.First(),ObsDim.Last())))
        D2 = @inferred(DataSubset((X',X), 1:150, (ObsDim.First(),ObsDim.Last())))
        D3 = DataSubset((X',X), obsdim = (1,:last))
        D4 = DataSubset((X',X), 1:150, obsdim = (:first,:last))
        for (s1,s2) in (D1,D2,D3,D4)
            @test typeof(datasubset(s1,2:10)) <: DataSubset
            @test @inferred(datasubset(s1,2:10)) == @inferred(s1[2:10])
            @test @inferred(datasubset(s1,2:10)) == @inferred(DataSubset(s1,2:10))
            @test s1.obsdim == ObsDim.First()
            @test s2.obsdim == ObsDim.Last()
            @test getobs(s1,2) == getobs(s2,2)
            @test getobs(s1,9:10) == getobs(s2,9:10)'
            @test getobs((s1,s2),9:10) == (getobs(s1,9:10),getobs(s2,9:10))
            @test nobs(s1) == nobs(s2) == 150
        end
    end

    @testset "Array, SubArray, SparseArray" begin
        @test nobs(DataSubset(X, obsdim = 1)) == 4
        @test nobs(DataSubset(X, 1:3, obsdim = 1)) == 3
        @test_reference "references/DataSubset1.txt" DataSubset(X, Int64(1):Int64(nobs(X)))
        @test_reference "references/DataSubset2.txt" @io2str showcompact(::IO, DataSubset(X))
        for var in (Xs, ys, vars...)
            subset = @inferred(DataSubset(var))
            @test subset.data === var
            @test subset.indices === 1:150
            @test typeof(subset) <: DataSubset
            @test @inferred(nobs(subset)) === nobs(var)
            @test @inferred(getobs(subset)) == getobs(var)
            @test @inferred(DataSubset(subset)) === subset
            @test @inferred(DataSubset(subset, 1:150)) === subset
            @test subset[end] == DataSubset(var, 150)
            @test @inferred(subset[150]) == DataSubset(var, 150)
            @test @inferred(subset[20:25]) == DataSubset(var, 20:25)
            for idx in (1:100, [1,10,150,3], [2])
                @test DataSubset(var)[idx] == DataSubset(var, idx)
                @test DataSubset(var)[idx] == DataSubset(var, collect(idx))
                subset = @inferred(DataSubset(var, idx))
                @test typeof(subset) <: DataSubset{typeof(var), typeof(idx)}
                @test subset.data === var
                @test subset.indices === idx
                @test subset.obsdim === ObsDim.Last()
                @test @inferred(nobs(subset)) === length(idx)
                @test @inferred(getobs(subset)) == getobs(var, idx)
                @test @inferred(DataSubset(subset)) === subset
                @test @inferred(subset[1]) == DataSubset(var, idx[1])
                if typeof(idx) <: AbstractRange
                    @test typeof(@inferred(subset[1:1])) == typeof(DataSubset(var, idx[1:1]))
                    @test nobs(subset[1:1]) == nobs(DataSubset(var, idx[1:1]))
                else
                    @test typeof(@inferred(subset[1:1])) == typeof(DataSubset(var, view(idx, 1:1)))
                    @test nobs(subset[1:1]) == nobs(DataSubset(var, view(idx, 1:1)))
                end
            end
        end
    end

    @testset "custom types" begin
        @test_throws MethodError DataSubset(EmptyType())
        @test_throws MethodError DataSubset(EmptyType(), 1:10)
        @test_throws MethodError DataSubset(EmptyType(), 1:10, ObsDim.First())
        @test_throws MethodError DataSubset(CustomType(), obsdim=1)
        @test_throws MethodError DataSubset(CustomType(), obsdim=:last)
        @test_throws MethodError DataSubset(CustomType(), 2:10, obsdim=1)
        @test_throws MethodError DataSubset(CustomType(), 2:10, obsdim=:last)
        @test_throws BoundsError getobs(DataSubset(CustomType(), 11:20), 11)
        @test typeof(@inferred(DataSubset(CustomType()))) <: DataSubset
        @test nobs(DataSubset(CustomType())) === 100
        @test nobs(DataSubset(CustomType(), 11:20)) === 10
        @test getobs(DataSubset(CustomType())) == collect(1:100)
        @test getobs(DataSubset(CustomType(),11:20),10) == 20
        @test getobs(DataSubset(CustomType(),11:20),[3,5]) == [13,15]
    end
end

@testset "DataSubset getindex and getobs" begin
    @testset "Matrix and SubArray{T,2}" begin
        for var in (X, Xv)
            subset = @inferred(DataSubset(var, 101:150))
            @test typeof(@inferred(getobs(subset))) <: Array{Float64,2}
            @test @inferred(nobs(subset)) == length(subset) == 50
            @test @inferred(subset[10:20]) == DataSubset(X, 110:120)
            @test @inferred(subset[11:21]) != DataSubset(X, 110:120)
            @test @inferred(getobs(subset, 10:20)) == X[:, 110:120]
            @test @inferred(getobs(subset, [11,10,14])) == X[:, [111,110,114]]
            @test typeof(subset[10:20]) <: DataSubset
            @test @inferred(subset[collect(10:20)]) == DataSubset(X, collect(110:120))
            @test typeof(subset[collect(10:20)]) <: DataSubset
            @test @inferred(getobs(subset)) == getobs(subset[1:end]) == X[:, 101:150]
        end
    end

    @testset "Vector and SubArray{T,1}" begin
        for var in (y, yv)
            subset = @inferred(DataSubset(var, 101:150))
            @test typeof(getobs(subset)) <: Array{String,1}
            @test @inferred(nobs(subset)) == length(subset) == 50
            @test @inferred(subset[10:20]) == DataSubset(y, 110:120)
            @test @inferred(getobs(subset, 10:20)) == y[110:120]
            @test @inferred(getobs(subset, [11,10,14])) == y[[111,110,114]]
            @test typeof(subset[10:20]) <: DataSubset
            @test @inferred(subset[collect(10:20)]) == DataSubset(y, collect(110:120))
            @test typeof(subset[collect(10:20)]) <: DataSubset
            @test @inferred(getobs(subset)) == getobs(subset[1:end]) == y[101:150]
        end
    end

    @testset "2-Tuple of Matrix, Vector, or SubArray"  begin
        for v1 in (X, Xv), v2 in (y, yv)
            subset = @inferred(DataSubset((v1,v2), 101:150))
            @test typeof(getobs(subset)) <: Tuple{Array{Float64,2},Array{String,1}}
            @test @inferred(nobs(subset)) == nobs(subset[1]) == nobs(subset[2]) == 50
            @test @inferred(subset[1][10:20]) == DataSubset(X, 110:120)
            @test @inferred(subset[2][10:20]) == DataSubset(y, 110:120)
            @test @inferred(getobs(subset)) == (X[:, 101:150], y[101:150])
        end
    end

    @testset "2-Tuple of SparseArray"  begin
        subset = @inferred(DataSubset((Xs,ys), 101:150))
        @test typeof(subset) <: Tuple
        @test typeof(subset[1]) <: DataSubset
        @test typeof(subset[2]) <: DataSubset
        @test typeof(@inferred(getobs(subset))) <: Tuple
        @test typeof(getobs(subset)[1]) <: SparseMatrixCSC
        @test typeof(getobs(subset)[2]) <: SparseVector
        @test @inferred(nobs(subset)) == nobs(subset[1]) == nobs(subset[2]) == 50
        @test @inferred(getobs(subset[1][10:20])) == getindex(Xs, :, 110:120)
        @test @inferred(getobs(subset[2][10:20])) == getindex(ys, 110:120)
        @test @inferred(getobs(subset)) == (getindex(Xs, :, 101:150), getindex(ys, 101:150))
    end
end

println("<HEARTBEAT>")

@testset "datasubset" begin
    @testset "Array and SubArray" begin
        @test @inferred(datasubset(X)) == Xv
        @test @inferred(datasubset(X, ObsDim.Last())) == Xv
        @test @inferred(datasubset(X, ObsDim.Last())) == Xv
        @test typeof(datasubset(X)) <: SubArray
        @test @inferred(datasubset(Xv)) === Xv
        @test @inferred(datasubset(XX)) == XX
        @test @inferred(datasubset(XXX)) == XXX
        @test typeof(datasubset(XXX)) <: SubArray
        @test @inferred(datasubset(y)) == y
        @test typeof(datasubset(y)) <: SubArray
        @test @inferred(datasubset(yv)) === yv
        for i in (2, 1:150, 2:10, [2,5,7], [2,1])
            @test @inferred(datasubset(X,i))   === view(X,:,i)
            @test @inferred(datasubset(Xv,i))  === view(X,:,i)
            @test @inferred(datasubset(Xv,i))  === view(Xv,:,i)
            @test @inferred(datasubset(XX,i))  === view(XX,:,:,i)
            @test @inferred(datasubset(XXX,i)) === view(XXX,:,:,:,i)
            @test @inferred(datasubset(y,i))   === view(y,i)
            @test @inferred(datasubset(yv,i))  === view(y,i)
            @test @inferred(datasubset(yv,i))  === view(yv,i)
            @test @inferred(datasubset(Y,i))   === view(Y,:,i)
        end
    end

    @testset "Tuple of Array and Subarray" begin
        @test_throws DimensionMismatch datasubset((X,X), 1:150, (ObsDim.Last(), ObsDim.Last(), ObsDim.Last()))
        @test_throws DimensionMismatch datasubset((X,X), 1:150, (ObsDim.Last(),))
        @test_throws DimensionMismatch datasubset((X,X), (ObsDim.Last(), ObsDim.Last(), ObsDim.Last()))
        @test_throws DimensionMismatch datasubset((X,X), (ObsDim.Last(),))
        @test @inferred(datasubset((X,y),ObsDim.Last())) == (X,y)
        @test @inferred(datasubset((X,y),(ObsDim.Last(),ObsDim.Last()))) == (X,y)
        @test @inferred(datasubset((X,y)))   == (X,y)
        @test @inferred(datasubset((X,yv)))  == (X,yv)
        @test @inferred(datasubset((X,yv)))  === (view(X,:,1:150),yv)
        @test @inferred(datasubset((Xv,y)))  == (Xv,y)
        @test @inferred(datasubset((Xv,y)))  === (Xv,view(y,1:150))
        @test @inferred(datasubset((Xv,yv))) === (Xv,yv)
        @test @inferred(datasubset((X,Y)))   == (X,Y)
        @test @inferred(datasubset((XX,X,y))) == (XX,X,y)
        for i in (1:150, 2:10, [2,5,7], [2,1])
            @test @inferred(datasubset((X,y),i))   === (view(X,:,i), view(y,i))
            @test @inferred(datasubset((Xv,y),i))  === (view(X,:,i), view(y,i))
            @test @inferred(datasubset((X,yv),i))  === (view(X,:,i), view(y,i))
            @test @inferred(datasubset((Xv,yv),i)) === (view(X,:,i), view(y,i))
            @test @inferred(datasubset((XX,X,y),i)) === (view(XX,:,:,i), view(X,:,i),view(y,i))
            # compare if obs match in tuple
            x1, y1 = getobs(datasubset((X1,Y1), i))
            @test all(x1' .== y1)
            x1, y1, z1 = getobs(datasubset((X1,Y1,X1), i))
            @test all(x1' .== y1)
            @test all(x1 .== z1)
        end
    end

    println("<HEARTBEAT>")

    @testset "SparseArray" begin
        @test @inferred(datasubset(Xs)) === DataSubset(Xs)
        @test @inferred(datasubset(ys)) === DataSubset(ys)
        for i in (1:150, 2:10, [2,5,7], [2,1])
            @test @inferred(datasubset(Xs,i)) === DataSubset(Xs,i)
            @test @inferred(datasubset(ys,i)) === DataSubset(ys,i)
        end
    end

    @testset "Tuple of SparseArray" begin
        @test @inferred(datasubset((Xv,ys))) === (Xv,DataSubset(ys))
        @test @inferred(datasubset((X,ys)))  === (datasubset(X),DataSubset(ys))
        @test @inferred(datasubset((Xs,y)))  === (DataSubset(Xs),datasubset(y))
        @test @inferred(datasubset((Xs,ys))) === (DataSubset(Xs),DataSubset(ys))
        @test @inferred(datasubset((Xs,Xs))) === (DataSubset(Xs),DataSubset(Xs))
        @test @inferred(datasubset((ys,Xs))) === (DataSubset(ys),DataSubset(Xs))
        @test @inferred(datasubset((XX,Xs,yv))) === (datasubset(XX),DataSubset(Xs),yv)
        for i in (1:150, 2:10, [2,5,7], [2,1])
            @test @inferred(datasubset((X,ys),i))  === (view(X,:,i), DataSubset(ys,i))
            @test @inferred(datasubset((Xs,y),i))  === (DataSubset(Xs,i), view(y,i))
            @test @inferred(datasubset((Xs,ys),i)) === (DataSubset(Xs,i), DataSubset(ys,i))
            @test @inferred(datasubset((Xs,Xs),i)) === (DataSubset(Xs,i), DataSubset(Xs,i))
            @test @inferred(datasubset((ys,Xs),i)) === (DataSubset(ys,i), DataSubset(Xs,i))
            @test @inferred(datasubset((XX,Xs,y),i)) === (view(XX,:,:,i),DataSubset(Xs,i),view(y,i))
            # compare if obs match in tuple
            x1, y1 = getobs(datasubset((X1,sparse(Y1)), i))
            @test all(x1' .== y1)
            x1, y1, z1 = getobs(datasubset((X1,Y1,sparse(X1)), i))
            @test all(x1' .== y1)
            @test all(x1 .== z1)
        end
    end

    @testset "custom types" begin
        @test_throws MethodError datasubset(EmptyType())
        @test_throws MethodError datasubset(EmptyType(), 1:10)
        @test_throws MethodError datasubset(EmptyType(), 1:10, ObsDim.First())
        @test_throws MethodError datasubset(CustomType(), obsdim=1)
        @test_throws MethodError datasubset(CustomType(), obsdim=:last)
        @test_throws MethodError datasubset(CustomType(), 2:10, obsdim=1)
        @test_throws MethodError datasubset(CustomType(), 2:10, obsdim=:last)
        @test_throws BoundsError getobs(datasubset(CustomType(), 11:20), 11)
        @test typeof(@inferred(datasubset(CustomType()))) <: DataSubset
        @test datasubset(CustomType()) == DataSubset(CustomType())
        @test datasubset(CustomType(), 11:20) == DataSubset(CustomType(), 11:20)
        @test nobs(datasubset(CustomType())) === 100
        @test nobs(datasubset(CustomType(), 11:20)) === 10
        @test getobs(datasubset(CustomType())) == collect(1:100)
        @test getobs(datasubset(CustomType(), 11:20), 10) == 20
        @test getobs(datasubset(CustomType(), 11:20), [3,5]) == [13,15]
    end
end

@testset "getobs!" begin
    @test getobs!(nothing, datasubset(y, 1)) == "setosa"

    @testset "DataSubset" begin
        xbuf1 = zeros(4,8)
        s1 = DataSubset(X, 2:9)
        @test @inferred(getobs!(xbuf1,s1)) === xbuf1
        @test xbuf1 == getobs(s1)
        xbuf1 = zeros(4,5)
        s1 = DataSubset(X, 10:17)
        @test @inferred(getobs!(xbuf1,s1,2:6)) === xbuf1
        @test xbuf1 == getobs(s1,2:6) == getobs(X,11:15)

        xbuf2 = zeros(8,4)
        s2 = DataSubset(X', 2:9, obsdim=1)
        @test @inferred(getobs!(xbuf2,s2)) === xbuf2
        @test xbuf2 == getobs(s2)
        xbuf2 = zeros(5,4)
        s2 = DataSubset(X', 10:17, obsdim=1)
        @test @inferred(getobs!(xbuf2,s2,2:6)) === xbuf2
        @test xbuf2 == getobs(s2,2:6) == getobs(X',11:15,obsdim=1)

        s3 = DataSubset(Xs, 11:15)
        @test @inferred(getobs!(nothing,s3)) == getobs(Xs,11:15)

        s4 = DataSubset(CustomType(), 6:10)
        @test @inferred(getobs!(nothing,s4)) == getobs(s4)
        s5 = DataSubset(CustomType(), 9:20)
        @test @inferred(getobs!(nothing,s5,2:6)) == getobs(s5,2:6)
    end

    @testset "Tuple with DataSubset" begin
        xbuf = zeros(4,2)
        ybuf = ["foo", "bar"]
        s1 = DataSubset(Xs, 5:9)
        s2 = DataSubset(X, 5:9)
        @test_throws AssertionError getobs!((nothing,xbuf),(s1,s2),2:3,ObsDim.First())
        @test_throws AssertionError getobs!((nothing,xbuf),(s1,s2),2:3,(ObsDim.First(),ObsDim.Last()))
        @test getobs!((nothing,xbuf),(s1,s2), 2:3) == (getobs(Xs,6:7),xbuf)
        @test xbuf == getobs(X,6:7)
        @test getobs!((nothing,xbuf),(s1,s2), 2:3, ObsDim.Last()) == (getobs(Xs,6:7),xbuf)
        @test getobs!((nothing,xbuf),(s1,s2), 2:3, (ObsDim.Last(),ObsDim.Last())) == (getobs(Xs,6:7),xbuf)
    end
end


#__precompile__(true)

module ImagineAnalyses

using AxisArrays, Interpolations, Unitful, IntervalSets
using Reexport
@reexport using CachedCalls, ImagineInterface, UnitAliases

import UnitAliases.HasTimeUnits
import CachedCalls.CachedCall
import Unitful: μm, s, Hz
using DSP
using Statistics

export dumb_bench,
        get_cycles,
        largest_cycle_diff,
        largest_cycle_std,
        mean_cycle,

        find_circular,
        get_circular,
        mon_delay,
        pulse_timings,
        flash_cam_cycs,
        max_exp,
        framerate_efficiency

#This will be a bit of a hodgepodge of analyses one can do on ImagineSignals and acquired images.  May be worth organizing differently if it gets too big.

#this may do something like measure the lag between sig1 and sig2
dumb_bench(sig1::ImagineSignal, sig2::ImagineSignal) = 0.0

#Define new constructors for CachedCalls on a per-function basis when the return type is known (warning: risky when function isn't type stable)
CachedCall(f::typeof(dumb_bench), args...) = CachedCall("dumb benchmark", f, (args...,), Float64)

include("consistency.jl")
include("exposure_placement.jl")

#To implement:

#lag(sig1::ImagineSignal, sig2::ImagineSignal)
    #lag(sig1::AxisArray, sig2::AxisArray)
#offset(sig1::ImagineSignal, sig2::ImagineSignal)
#gain(sig_out::ImagineSignal, sig_in::ImagineSignal)
#corr(sig1::ImagineSignal, sig2::ImagineSignal)
#pulse_count(sig::ImagineSignal)
#range(sig::ImagineSignal)
#signal_consistency(sig::ImagineSignal, cycle_dur::HasTimeUnits) #consistency of a periodic 1d timeseries. may change the interface

#The below use images
#frame_count(camsig::ImagineSignal, img::AbstractImage{T,N}) where {T,N}
#frame_order(camsig::ImagineSignal, img::AbstractImage{T,N}) where {T,N}  #returns an integer sequence of frame numbers
#bidi_image_consistency(img::AbstractImage{T,4}) where {T} #returns each pair of ssds in a bidi sequence
#image_consistency(img::AbstractImage{T,N}, imgtemplate::AbstractImage{T,N-1}) where {T, N} #returns ssd between imgtemplate and each image in the img sequence

end # module


module Stats_module

push!(LOAD_PATH, "../src")

using JLD, Aircraft_module, Constants_module
include("check_status.jl")

export Stats, run_stats

mutable struct Stats
	# airspace_controller::Airspace_Controller
	# cas_ctrl_type::DataType
	# sensor_ctrl::Sensor_Property_Controller
	# dyn_update_rule::Update_Rule

	generated_time::Float64

	# Metrics within simulation wrt time steps:
	t::Array{Int64,1}
	num_ac::Array{Int64,1}
	num_NMAC_t::Array{Int64,1}
	num_NMAC_cum::Array{Int64,1}
	average_flight_time::Array{Float64,1}
	num_ac_arrived::Array{Int64,1}
	num_ac_created::Array{Int64,1}
	traffic_density::Array{Float64,1}
	avg_ac_encountered::Array{Float64,1}
	ac_encountered_per_step_arr::Array{Float64,1}

	# Simulation-wise metrics:
	# Safety:
	total_NMACs::Int64
	total_flight_hours::Int64
	final_traffic_density::Float64

	# Efficiency:
	landed_over_created::Float64
	final_avg_flight_hours::Float64
	total_route_len::Float64
	total_nominal_route_len::Float64
	encounter_distribution::Vector{Float64}

	function Stats(
		generated_time::Float64)
		return new(generated_time,
			[], [], [], [], [], [], [], [], [], [],
			0, 0, 0., 0., 0., 0., 0., zeros(6))
	end
end

function run_stats(All_AC::Array{Aircraft,1}, stats::Stats, t::Int64)
	push!(stats.t, t)
	push!(stats.num_ac, length(All_AC))
	push!(stats.num_NMAC_t, check_NMAC(All_AC))
	push!(stats.num_NMAC_cum, sum(stats.num_NMAC_t))
	avg_flight_time = 0
	if isempty(stats.num_ac_arrived)
		num_ac_arrived = 0
	else
		num_ac_arrived = stats.num_ac_arrived[t-1]
	end
	avg_ac_encountered = 0.
	temp_num_intruders_per_step = 0.
	for ac in All_AC
		avg_flight_time += (t - ac.created_time)
		if check_desti(ac) == true
			num_ac_arrived += 1
		end
		ac.t_to_nominal_t = (t - ac.created_time) / (ac.nominal_route_len / ac.dyn.v)
		ac.d_to_nominal_d = sqrt((ac.dyn.x - ac.x_dest)^2 + (ac.dyn.y - ac.y_dest)^2) / ac.nominal_route_len
		avg_ac_encountered += ac.cumulative_num_ac_encountered / (t - ac.created_time)
		temp_num_intruders_per_step += ac.num_ac_encountered_arr[end]
		if ac.num_ac_encountered_arr[end] > 5
			stats.encounter_distribution[end] += 1
		elseif ac.num_ac_encountered_arr[end] >= 1 && ac.num_ac_encountered_arr[end] <= 5
			stats.encounter_distribution[ac.num_ac_encountered_arr[end]] += 1
		end
	end
	avg_flight_time /= length(All_AC)
	push!(stats.average_flight_time, avg_flight_time)
	push!(stats.num_ac_arrived, num_ac_arrived)
	push!(stats.num_ac_created, num_ac_arrived + length(All_AC))
	push!(stats.traffic_density, length(All_AC) / airspace_dim^2)
	stats.total_flight_hours += length(All_AC) * dt
	stats.total_NMACs = stats.num_NMAC_cum[end]
	stats.final_traffic_density = stats.traffic_density[end]
	stats.landed_over_created = stats.num_ac_arrived[end] / stats.num_ac_created[end]
	stats.final_avg_flight_hours = stats.total_flight_hours / stats.num_ac_created[end]
	avg_ac_encountered /= length(All_AC)
	push!(stats.avg_ac_encountered, avg_ac_encountered)
	temp_num_intruders_per_step /= length(All_AC)
	push!(stats.ac_encountered_per_step_arr, temp_num_intruders_per_step)
end

end # module


using KernelAbstractions
using Oceananigans.Utils
using Oceananigans.Architectures: device, @hascuda, CPU, GPU, array_type
using Oceananigans.Fields: datatuple
using Oceananigans.BoundaryConditions: apply_x_bcs!, apply_y_bcs!, apply_z_bcs!

@kernel function update_total_density!(total_density, grid, gases, tracers)
    i, j, k = @index(Global, NTuple)

    @inbounds total_density[i, j, k] = diagnose_density(i, j, k, grid, gases, tracers)
end

# This is for users. Do not use in time_stepping.jl as it doesn't make use of intermediate fields.
function update_total_density!(model)

    density_update_event =
        launch!(model.architecture, model.grid, :xyz, update_total_density!,
                datatuple(model.total_density), model.grid, model.gases, datatuple(model.tracers),
                dependencies=Event(device(model.architecture)))

    wait(device(model.architecture), density_update_event)

    return nothing
end

#####
##### Computing slow source terms (viscous dissipation, diffusion, and Coriolis terms).
#####

function compute_slow_source_terms!(slow_source_terms, arch, grid, thermodynamic_variable, gases, gravity, coriolis, closure, total_density, momenta, tracers, diffusivities, forcing, clock)

    slow_source_terms, total_density, momenta, tracers, diffusivities =
        datatuples(slow_source_terms, total_density, momenta, tracers, diffusivities)

    workgroup, worksize = work_layout(grid, :xyz)
    barrier = Event(device(arch))

    momentum_kernel! = compute_slow_momentum_source_terms!(device(arch), workgroup, worksize)
    tracer_kernel! = compute_slow_tracer_source_terms!(device(arch), workgroup, worksize)
    thermodynamic_variable_kernel! = compute_slow_thermodynamic_variable_source_terms!(device(arch), workgroup, worksize)

    momentum_event = momentum_kernel!(slow_source_terms, grid, coriolis, closure, total_density, momenta, tracers, diffusivities, forcing, clock, dependencies=barrier)

    events = [momentum_event]

    for (tracer_index, ρc_name) in enumerate(propertynames(tracers))
        ρc   = getproperty(tracers, ρc_name)
        S_ρc = getproperty(slow_source_terms.tracers, ρc_name)
        forcing_ρc = getproperty(forcing, ρc_name)

        tracer_event = tracer_kernel!(S_ρc, grid, closure, tracer_index, total_density, ρc, momenta, tracers, diffusivities, forcing_ρc, clock, dependencies=barrier)
        push!(events, tracer_event)
    end

    thermodynamic_variable_event = thermodynamic_variable_kernel!(slow_source_terms.tracers[1], grid, thermodynamic_variable, gases, gravity, closure, total_density, momenta, tracers, diffusivities, dependencies=barrier)
    push!(events, thermodynamic_variable_event)

    wait(device(arch), MultiEvent(Tuple(events)))

    return nothing
end

@kernel function compute_slow_momentum_source_terms!(slow_source_terms, grid, coriolis, closure, total_density, momenta, tracers, diffusivities, forcing, clock)
    i, j, k = @index(Global, NTuple)

    @inbounds slow_source_terms.ρu[i, j, k] = ρu_slow_source_term(i, j, k, grid, coriolis, closure, total_density, momenta, diffusivities) + forcing.ρu(i, j, k, grid, clock, merge(momenta, tracers))
    @inbounds slow_source_terms.ρv[i, j, k] = ρv_slow_source_term(i, j, k, grid, coriolis, closure, total_density, momenta, diffusivities) + forcing.ρv(i, j, k, grid, clock, merge(momenta, tracers))
    @inbounds slow_source_terms.ρw[i, j, k] = ρw_slow_source_term(i, j, k, grid, coriolis, closure, total_density, momenta, diffusivities) + forcing.ρw(i, j, k, grid, clock, merge(momenta, tracers))
end

@kernel function compute_slow_tracer_source_terms!(S_ρc, grid, closure, tracer_index, total_density, ρc, momenta, tracers, diffusivities, forcing, clock)
    i, j, k = @index(Global, NTuple)

    @inbounds S_ρc[i, j, k] = ρc_slow_source_term(i, j, k, grid, closure, tracer_index, total_density, ρc, diffusivities) + forcing(i, j, k, grid, clock, merge(momenta, tracers))
end

@kernel function compute_slow_thermodynamic_variable_source_terms!(S_ρt, grid, thermodynamic_variable, gases, gravity, closure, total_density, momenta, tracers, diffusivities)
    i, j, k = @index(Global, NTuple)

    @inbounds S_ρt[i, j, k] += ρt_slow_source_term(i, j, k, grid, closure, thermodynamic_variable, gases, gravity, total_density, momenta, tracers, diffusivities)
end

#####
##### Computing fast source terms (advection, pressure gradient, and buoyancy terms).
#####

function compute_fast_source_terms!(fast_source_terms, arch, grid, thermodynamic_variable, gases, gravity, advection_scheme, total_density, momenta, tracers, slow_source_terms)

    fast_source_terms, total_density, momenta, tracers, slow_source_terms =
        datatuples(fast_source_terms, total_density, momenta, tracers, slow_source_terms)

    workgroup, worksize = work_layout(grid, :xyz)
    barrier = Event(device(arch))

    momentum_kernel! = compute_fast_momentum_source_terms!(device(arch), workgroup, worksize)
    tracer_kernel! = compute_fast_tracer_source_terms!(device(arch), workgroup, worksize)
    thermodynamic_variable_kernel! = compute_fast_thermodynamic_variable_source_terms!(device(arch), workgroup, worksize)

    momentum_event = momentum_kernel!(fast_source_terms, grid, thermodynamic_variable, gases, gravity, advection_scheme, total_density, momenta, tracers, slow_source_terms, dependencies=barrier)

    events = [momentum_event]

    for ρc_name in propertynames(tracers)
        ρc   = getproperty(tracers, ρc_name)
        F_ρc = getproperty(fast_source_terms.tracers, ρc_name)
        S_ρc = getproperty(slow_source_terms.tracers, ρc_name)

        tracer_event = tracer_kernel!(F_ρc, grid, advection_scheme, total_density, momenta, ρc, S_ρc, dependencies=barrier)
        push!(events, tracer_event)
    end

    thermodynamic_variable_event = thermodynamic_variable_kernel!(fast_source_terms.tracers[1], grid, thermodynamic_variable, gases, gravity, total_density, momenta, tracers, dependencies=barrier)
    push!(events, thermodynamic_variable_event)

    wait(device(arch), MultiEvent(Tuple(events)))

    return nothing
end

@kernel function compute_fast_momentum_source_terms!(fast_source_terms, grid, thermodynamic_variable, gases, gravity, advection_scheme, total_density, momenta, tracers, slow_source_terms)
    i, j, k = @index(Global, NTuple)

    @inbounds fast_source_terms.ρu[i, j, k] = ρu_fast_source_term(i, j, k, grid, thermodynamic_variable, gases, gravity, advection_scheme, total_density, momenta, tracers, slow_source_terms.ρu)
    @inbounds fast_source_terms.ρv[i, j, k] = ρv_fast_source_term(i, j, k, grid, thermodynamic_variable, gases, gravity, advection_scheme, total_density, momenta, tracers, slow_source_terms.ρv)
    @inbounds fast_source_terms.ρw[i, j, k] = ρw_fast_source_term(i, j, k, grid, thermodynamic_variable, gases, gravity, advection_scheme, total_density, momenta, tracers, slow_source_terms.ρw)
end

@kernel function compute_fast_tracer_source_terms!(F_ρc, grid, advection_scheme, total_density, momenta, ρc, S_ρc)
    i, j, k = @index(Global, NTuple)

    @inbounds F_ρc[i, j, k] = ρc_fast_source_term(i, j, k, grid, advection_scheme, total_density, momenta, ρc, S_ρc)
end

@kernel function compute_fast_thermodynamic_variable_source_terms!(F_ρt, grid, thermodynamic_variable, gases, gravity, total_density, momenta, tracers)
    i, j, k = @index(Global, NTuple)

    @inbounds F_ρt[i, j, k] += ρt_fast_source_term(i, j, k, grid, thermodynamic_variable, gases, gravity, total_density, momenta, tracers)
end

#####
##### Calculating boundary tendency contributions
#####

function calculate_boundary_tendency_contributions!(source_terms, arch, momenta, tracers, clock, model_fields)

    barrier = Event(device(arch))

    events = []

    # Momentum fields
    momentum_source_terms = (source_terms.ρu, source_terms.ρv, source_terms.ρw)

    for (ρϕ_source_term, ρϕ) in zip(momentum_source_terms, momenta)
        x_bcs_event = apply_x_bcs!(ρϕ_source_term, ρϕ, arch, barrier, clock, model_fields)
        y_bcs_event = apply_y_bcs!(ρϕ_source_term, ρϕ, arch, barrier, clock, model_fields)
        z_bcs_event = apply_z_bcs!(ρϕ_source_term, ρϕ, arch, barrier, clock, model_fields)

        push!(events, x_bcs_event, y_bcs_event, z_bcs_event)
    end

    # Tracer fields
    for (ρϕ_source_term, ρϕ) in zip(source_terms.tracers, tracers)
        x_bcs_event = apply_x_bcs!(ρϕ_source_term, ρϕ, arch, barrier, clock, model_fields)
        y_bcs_event = apply_y_bcs!(ρϕ_source_term, ρϕ, arch, barrier, clock, model_fields)
        z_bcs_event = apply_z_bcs!(ρϕ_source_term, ρϕ, arch, barrier, clock, model_fields)

        push!(events, x_bcs_event, y_bcs_event, z_bcs_event)
    end

    events = filter(e -> typeof(e) <: Event, events)

    wait(device(arch), MultiEvent(Tuple(events)))

    return nothing
end

#####
##### Advancing state variables
#####

function advance_state_variables!(state_variables, arch, grid, momenta, tracers, fast_source_terms; Δt)

    state_variables, momenta, tracers, fast_source_terms =
        datatuples(state_variables, momenta, tracers, fast_source_terms )

    workgroup, worksize = work_layout(grid, :xyz)
    barrier = Event(device(arch))

    momentum_kernel! = advance_momentum!(device(arch), workgroup, worksize)
    tracer_kernel! = advance_tracer!(device(arch), workgroup, worksize)

    momentum_event = momentum_kernel!(state_variables, grid, momenta, fast_source_terms, Δt, dependencies=barrier)

    events = [momentum_event]

    for ρc_name in propertynames(tracers)
        ρc   = getproperty(tracers, ρc_name)
        ρc⁺  = getproperty(state_variables.tracers, ρc_name)
        F_ρc = getproperty(fast_source_terms.tracers, ρc_name)

        tracer_event = tracer_kernel!(ρc⁺, grid, ρc, F_ρc, Δt, dependencies=barrier)
        push!(events, tracer_event)
    end

    wait(device(arch), MultiEvent(Tuple(events)))

    return nothing
end

@kernel function advance_momentum!(momenta⁺, grid, momenta, fast_source_terms, Δt)
    i, j, k = @index(Global, NTuple)

    @inbounds momenta⁺.ρu[i, j, k] = momenta.ρu[i, j, k] + Δt * fast_source_terms.ρu[i, j, k]
    @inbounds momenta⁺.ρv[i, j, k] = momenta.ρv[i, j, k] + Δt * fast_source_terms.ρv[i, j, k]
    @inbounds momenta⁺.ρw[i, j, k] = momenta.ρw[i, j, k] + Δt * fast_source_terms.ρw[i, j, k]
end

@kernel function advance_tracer!(ρc⁺, grid, ρc, F_ρc, Δt)
    i, j, k = @index(Global, NTuple)

    @inbounds ρc⁺[i, j, k] = ρc[i, j, k] + Δt * F_ρc[i, j, k]
end