text:
function select_normal_and_discard!(pairset::Pairset, basis::Basis,
                                matrix::MacaulayMatrix, ht::MonomialHashtable,
                                symbol_ht::MonomialHashtable; maxpairs::Int=0)

    sort_pairset_by_degree!(pairset, 1, pairset.load - 1)
    # sort by degree

    ps = pairset.pairs
    min_deg = ps[1].deg
    min_idx = 0

    while min_idx < pairset.load && ps[min_idx + 1].deg == min_deg
        min_idx += 1
    end

    # number of discarded pairs
    npairs = min_idx

    @debug "Discarder $(npairs) pairs"

    for i in 1:pairset.load - npairs
        ps[i] = ps[i + npairs]
    end
    pairset.load -= npairs
end

function select_normal!(
            pairset::Pairset, basis::Basis, matrix::MacaulayMatrix,
            ht::MonomialHashtable, symbol_ht::MonomialHashtable;
            maxpairs::Int=0)

    sort_pairset_by_degree!(pairset, 1, pairset.load - 1)
    # sort by degree

    ps = pairset.pairs
    min_deg = ps[1].deg
    min_idx = 0

    while min_idx < pairset.load && ps[min_idx + 1].deg == min_deg
        min_idx += 1
    end

    # number of selected pairs
    npairs = min_idx

    @debug "Selected $(npairs) pairs"

    sort_pairset_by_lcm!(pairset, npairs, ht)

    reinitialize_matrix!(matrix, npairs)

    uprows  = matrix.uprows
    lowrows = matrix.lowrows

    # polynomials from pairs in order (p11, p12)(p21, p21)
    # (future rows of the matrix)
    gens   = Vector{Int}(undef, 2*npairs)

    # buffer !
    etmp = ht.exponents[1]
    i = 1
    while i <= npairs
        matrix.ncols += 1
        load = 1
        lcm = ps[i].lcm
        j = i

        # we collect all generators with same lcm into gens
        @inbounds while j <= npairs && ps[j].lcm == lcm
            gens[load] = ps[j].poly1
            load += 1
            gens[load] = ps[j].poly2
            load += 1
            j += 1
        end
        load -= 1

        # sort by number in the basis (by=identity)
        sort_generators_by_position!(gens, load)

        # now we collect reducers, and reduced

        # first generator index in groebner basis
        prev = gens[1]
        # first generator in hash table
        poly = basis.gens[prev]
        # first generator lead monomial index in hash data
        vidx = poly[1]

        # first generator exponent
        eidx = ht.exponents[vidx]
        # exponent of lcm corresponding to first generator
        elcm = ht.exponents[lcm]
        for u in 1:ht.explen
            etmp[u] = elcm[u] - eidx[u]
        end
        # now etmp contents complement to eidx in elcm

        # hash of complement
        htmp = ht.hashdata[lcm].hash - ht.hashdata[vidx].hash

        # add row as a reducer
        matrix.nup += 1
        uprows[matrix.nup] = multiplied_poly_to_matrix_row!(symbol_ht, ht, htmp, etmp, poly)
        # map upper row to index in basis
        matrix.up2coef[matrix.nup] = prev

        # mark lcm column as reducer in symbolic hashtable
        symbol_ht.hashdata[uprows[matrix.nup][1]].idx = 2
        # increase number of rows set
        matrix.nrows += 1

        # over all polys with same lcm,
        # add them to the lower part of matrix
        @inbounds for k in 1:load
            # duplicate generator,
            # we can do so as long as generators are sorted
            if gens[k] == prev
                continue
            end

            # if the table was reallocated
            elcm = ht.exponents[lcm]

            # index in gb
            prev = gens[k]
            # poly of indices of monoms in hash table
            poly = basis.gens[prev]
            vidx = poly[1]
            # leading monom idx
            eidx = ht.exponents[vidx]
            for u in 1:ht.explen
                etmp[u] = elcm[u] - eidx[u]
            end

            htmp = ht.hashdata[lcm].hash - ht.hashdata[vidx].hash

            # add row to be reduced
            matrix.nlow += 1
            lowrows[matrix.nlow] = multiplied_poly_to_matrix_row!(symbol_ht, ht, htmp, etmp, poly)
            # map lower row to index in basis
            matrix.low2coef[matrix.nlow] = prev

            symbol_ht.hashdata[lowrows[matrix.nlow][1]].idx = 2

            matrix.nrows += 1
        end

        i = j
    end

    resize!(matrix.lowrows, matrix.nrows - matrix.ncols)

    # remove selected parirs from pairset
    for i in 1:pairset.load - npairs
        ps[i] = ps[i + npairs]
    end
    pairset.load -= npairs
end

#------------------------------------------------------------------------------

function select_isgroebner!(
        pairset::Pairset, basis::Basis, matrix::MacaulayMatrix,
        ht::MonomialHashtable, symbol_ht::MonomialHashtable)

    ps = pairset.pairs
    npairs = pairset.load

    sort_pairset_by_lcm!(pairset, npairs, ht)

    reinitialize_matrix!(matrix, npairs)

    uprows  = matrix.uprows
    lowrows = matrix.lowrows

    # polynomials from pairs in order (p11, p12)(p21, p21)
    # (future rows of the matrix)
    gens   = Vector{Int}(undef, 2*npairs)

    # buffer !
    etmp = ht.exponents[1]
    i = 1
    while i <= npairs
        matrix.ncols += 1
        load = 1
        lcm = ps[i].lcm
        j = i

        # we collect all generators with same lcm into gens
        @inbounds while j <= npairs && ps[j].lcm == lcm
            gens[load] = ps[j].poly1
            load += 1
            gens[load] = ps[j].poly2
            load += 1
            j += 1
        end
        load -= 1

        # sort by number in the basis (by=identity)
        sort_generators_by_position!(gens, load)

        # now we collect reducers, and reduced

        # first generator index in groebner basis
        prev = gens[1]
        # first generator in hash table
        poly = basis.gens[prev]
        # first generator lead monomial index in hash data
        vidx = poly[1]

        # first generator exponent
        eidx = ht.exponents[vidx]
        # exponent of lcm corresponding to first generator
        elcm = ht.exponents[lcm]
        for u in 1:ht.explen
            etmp[u] = elcm[u] - eidx[u]
        end
        # now etmp contents complement to eidx in elcm

        # hash of complement
        htmp = ht.hashdata[lcm].hash - ht.hashdata[vidx].hash

        # add row as a reducer
        matrix.nup += 1
        uprows[matrix.nup] = multiplied_poly_to_matrix_row!(symbol_ht, ht, htmp, etmp, poly)
        # map upper row to index in basis
        matrix.up2coef[matrix.nup] = prev

        # mark lcm column as reducer in symbolic hashtable
        symbol_ht.hashdata[uprows[matrix.nup][1]].idx = 2
        # increase number of rows set
        matrix.nrows += 1

        # over all polys with same lcm,
        # add them to the lower part of matrix
        @inbounds for k in 1:load
            # duplicate generator,
            # we can do so as long as generators are sorted
            if gens[k] == prev
                continue
            end

            # if the table was reallocated
            elcm = ht.exponents[lcm]

            # index in gb
            prev = gens[k]
            # poly of indices of monoms in hash table
            poly = basis.gens[prev]
            vidx = poly[1]
            # leading monom idx
            eidx = ht.exponents[vidx]
            for u in 1:ht.explen
                etmp[u] = elcm[u] - eidx[u]
            end

            htmp = ht.hashdata[lcm].hash - ht.hashdata[vidx].hash

            # add row to be reduced
            matrix.nlow += 1
            lowrows[matrix.nlow] = multiplied_poly_to_matrix_row!(symbol_ht, ht, htmp, etmp, poly)
            # map lower row to index in basis
            matrix.low2coef[matrix.nlow] = prev

            symbol_ht.hashdata[lowrows[matrix.nlow][1]].idx = 2

            matrix.nrows += 1
        end

        i = j
    end

    resize!(matrix.lowrows, matrix.nrows - matrix.ncols)
end

#------------------------------------------------------------------------------

function find_multiplied_reducer!(
            basis::Basis, matrix::MacaulayMatrix,
            ht::MonomialHashtable, symbol_ht::MonomialHashtable,
            vidx::Int)

    e = symbol_ht.exponents[vidx]
    etmp = ht.exponents[1]
    divmask = symbol_ht.hashdata[vidx].divmask

    blen = basis.ndone     #
    leaddiv = basis.lead  #

    # searching for a poly from state.basis whose leading monom
    # divides the given exponent e
    i = 1
    @label Letsgo
    @inbounds while i <= basis.nlead && (leaddiv[i] & ~divmask) != 0
        i += 1
    end

    #@inbounds while i <= basis.nlead && !divv(e, ht.exponents[basis.gens[basis.nonred[i]][1]])
    #    i += 1
    #end

    # here found polynomial from basis with leading monom
    # dividing symbol_ht.exponents[vidx]
    if i <= basis.nlead
        # reducers index and exponent in hash table
        @inbounds rpoly = basis.gens[basis.nonred[i]]
        @inbounds rexp  = ht.exponents[rpoly[1]]

        @inbounds for j in 1:ht.explen
            # if it actually does not divide and divmask lies
            if e[j] < rexp[j]
                i += 1
                @goto Letsgo
            end
            etmp[j] = e[j] - rexp[j]
        end
        # now etmp = e // rexp in terms of monomias,
        # hash is linear
        h = symbol_ht.hashdata[vidx].hash - ht.hashdata[rpoly[1]].hash

        matrix.uprows[matrix.nup + 1] = multiplied_poly_to_matrix_row!(symbol_ht, ht, h, etmp, rpoly)
        matrix.up2coef[matrix.nup + 1] = basis.nonred[i]

        # upsize matrix
        symbol_ht.hashdata[vidx].idx = 2
        matrix.nup += 1
        i += 1
    end

end

#------------------------------------------------------------------------------

# Given the set of polynomials L and the basis G,
# extends L to contain possible polys for reduction by G,
# and returns it
function symbolic_preprocessing!(
                basis::Basis, matrix::MacaulayMatrix,
                ht::MonomialHashtable, symbol_ht::MonomialHashtable)

    # check matrix sizes (we want to omit this I guess)

    symbol_load = symbol_ht.load

    nrr    = matrix.ncols
    onrr   = matrix.ncols

    #=
        TODO: matrix_enlarge!
    =#
    while matrix.size <= nrr + symbol_load
        matrix.size *= 2
        resize!(matrix.uprows, matrix.size)
        resize!(matrix.up2coef, matrix.size)
    end

    # for each lcm present in symbolic_ht set on select stage
    i = symbol_ht.offset
    #= First round, we add multiplied polynomials which divide  =#
    #= a monomial exponent from selected spairs                 =#
    @inbounds while i <= symbol_load
        # not a reducer already
        if symbol_ht.hashdata[i].idx == 0
            symbol_ht.hashdata[i].idx = 1
            matrix.ncols += 1
            find_multiplied_reducer!(basis, matrix, ht, symbol_ht, i)
        end
        i += 1
    end

    #= Second round, we add multiplied polynomials which divide  =#
    #= lcm added on previous for loop                            =#
    while i <= symbol_ht.load
        if matrix.size == matrix.nup
            matrix.size *= 2
            # TODO:
            resize!(matrix.uprows, matrix.size)
            resize!(matrix.up2coef, matrix.size)
        end

        symbol_ht.hashdata[i].idx = 1
        matrix.ncols += 1
        find_multiplied_reducer!(basis, matrix, ht, symbol_ht, i)
        i += 1
    end

    # shrink matrix sizes, set constants
    resize!(matrix.uprows, matrix.nup)
    # resize!(matrix.up2coef, matrix.nup)

    matrix.nrows += matrix.nup - onrr
    matrix.nlow  = matrix.nrows - matrix.nup
    matrix.size  = matrix.nrows

end

#------------------------------------------------------------------------------

# Given the set of polynomials L and the basis G,
# extends L to contain possible polys for reduction by G,
# and returns it
function symbolic_preprocessing_relaxed!(
                basis::Basis, matrix::MacaulayMatrix,
                ht::MonomialHashtable, symbol_ht::MonomialHashtable)

    # check matrix sizes (we want to omit this I guess)

    symbol_load = symbol_ht.load

    nrr    = matrix.ncols
    onrr   = matrix.ncols

    #=
        TODO: matrix_enlarge!
    =#
    while matrix.size <= nrr + symbol_load
        matrix.size *= 2
        resize!(matrix.uprows, matrix.size)
        resize!(matrix.up2coef, matrix.size)
    end

    # for each lcm present in symbolic_ht set on select stage
    i = symbol_ht.offset
    #= First round, we add multiplied polynomials which divide  =#
    #= a monomial exponent from selected spairs                 =#
    @inbounds while i <= symbol_load
        # not a reducer already
        if symbol_ht.hashdata[i].idx == 0
            symbol_ht.hashdata[i].idx = 1
            matrix.ncols += 1
            find_multiplied_reducer!(basis, matrix, ht, symbol_ht, i)
        end
        i += 1
    end

    #= Second round, we add multiplied polynomials which divide  =#
    #= lcm added on previous for loop                             =#
    #=
    while i <= symbol_ht.load
        if matrix.size == matrix.nup
            matrix.size *= 2
            resize!(matrix.uprows, matrix.size)
            resize!(matrix.up2coef, matrix.size)
        end

        symbol_ht.hashdata[i].idx = 1
        matrix.ncols += 1
        find_multiplied_reducer!(basis, matrix, ht, symbol_ht, i)
        i += 1
    end
    =#

    # shrink matrix sizes, set constants
    resize!(matrix.uprows, matrix.nup)
    # resize!(matrix.up2coef, matrix.nup)

    matrix.nrows += matrix.nup - onrr
    matrix.nlow  = matrix.nrows - matrix.nup
    matrix.size  = matrix.nrows

end

#------------------------------------------------------------------------------

function update_pairset!(
            pairset,
            basis,
            ht,
            update_ht,
            idx,
            plcm)

    pl = pairset.load
    bl = idx
    nl = pl + bl
    ps = pairset.pairs

    new_lead  = basis.gens[idx][1]

    # @error "" bl
    # initialize new critical lcms
    # plcm = Vector{Int}(undef, bl + 1)

    # for each combination (new_Lead, basis.gens[i][1])
    # generate a pair
    @inbounds for i in 1:bl-1
        plcm[i] = get_lcm(basis.gens[i][1], new_lead, ht, update_ht)
        deg = update_ht.hashdata[plcm[i]].deg
        newidx = pl + i
        # TRACE: move isred above
        if basis.isred[i] == 0
            ps[newidx] = SPair(i, idx, plcm[i], deg)
        else
            # lcm == 0 will mark redundancy of spair
            ps[newidx] = SPair(i, idx, 0, deg)
        end

    end

    # @error "" one two

    # traverse existing pairs
    @inbounds for i in 1:pl
        j = ps[i].poly1
        l = ps[i].poly2
        m = max(ps[pl + l].deg, ps[pl + j].deg)

        # if an existing pair is divisible by lead of new poly
        # and has a greater degree than newly generated one
        if is_monom_divisible(ps[i].lcm, new_lead, ht) && ps[i].deg > m
            # mark lcm as 0
            ps[i] = SPair(ps[i].poly1, ps[i].poly2, 0, ps[i].deg)
        end
    end
    # TODO: this can be done faster if we
    # create only non-redundant pairs at first

    # traverse new pairs to check for redundancy
    j = 1
    for i in 1:bl-1
        if basis.isred[i] == 0
            ps[pl + j] = ps[pl + i]
            j += 1
        end
    end

    sort_pairset_by_degree!(pairset, pl + 1, j - 2)

    @inbounds for i in 1:j - 1
        plcm[i] = ps[pl + i].lcm
    end
    plcm[j] = 0
    pc = j
    pc -= 1

    # mark redundancy of some pairs from plcm
    @inbounds for j in 1:pc
        # if is not redundant
        if plcm[j] > 0
            check_monomial_division_in_update(plcm, j + 1, pc, plcm[j], update_ht)
        end
    end

    # remove useless pairs from pairset
    # by moving them to the end
    j = 1
    @inbounds for i in 1:pairset.load
        ps[i].lcm == 0 && continue
        ps[j] = ps[i]
        j += 1
    end

    # assure that basis hashtable can store new lcms
    if ht.size - ht.load <= pc
        enlarge_hash_table!(ht)
    end

    # add new lcms to the basis hashtable,
    # including j and not including pc
    insert_plcms_in_basis_hash_table!(pairset, pl, ht, update_ht, basis, plcm, j, pc+1)

    # mark redundant elements in masis
    nonred = basis.nonred
    lml = basis.nlead
    @inbounds for i in 1:lml
        if basis.isred[nonred[i]] == 0
            if is_monom_divisible(basis.gens[nonred[i]][1], new_lead, ht)
                basis.isred[nonred[i]] = 1
            end
        end
    end

end

#------------------------------------------------------------------------------

function update_basis!(
            basis,
            ht::MonomialHashtable,
            update_ht::MonomialHashtable)

    # here we could check overall redundancy and update basis.lead

    k = 1
    lead   = basis.lead
    nonred = basis.nonred

    for i in 1:basis.nlead
        if basis.isred[nonred[i]] == 0
            basis.lead[k]   = lead[i]
            basis.nonred[k] = nonred[i]
            k += 1
        end
    end
    basis.nlead = k - 1

    for i in basis.ndone+1:basis.ntotal
        if basis.isred[i] == 0
            lead[k]   = ht.hashdata[basis.gens[i][1]].divmask
            nonred[k] = i
            k += 1
        end
    end

    basis.nlead = k - 1
    basis.ndone = basis.ntotal
end

# checks if element of basis at position idx is redundant
function is_redundant!(
            pairset, basis, ht, update_ht, idx)

    # TODO
    if 2*update_ht.load > update_ht.size
        enlarge_hash_table!(update_ht)
    end
    # reinitialize_hash_table!(update_ht, 2*idx)

    ps = pairset.pairs

    # lead of new polynomial
    lead_new = basis.gens[idx][1]
    # degree of lead
    lead_deg = ht.hashdata[lead_new].deg

    for i in idx+1:basis.ntotal
        i == idx && continue
        if basis.isred[i] == 1
            continue
        end

        lead_i = basis.gens[i][1]

        if is_monom_divisible(lead_new, lead_i, ht)
            lcm_new = get_lcm(lead_i, lead_new, ht, ht)

            psidx = pairset.load + 1
            ps[psidx] = SPair(i, idx, lcm_new, ht.hashdata[lcm_new].deg)

            basis.isred[idx] = 1
            pairset.load += 1

            return true
        end
    end

    return false
end

#------------------------------------------------------------------------------

function update!(
        pairset::Pairset,
        basis::Basis,
        ht::MonomialHashtable,
        update_ht::MonomialHashtable,
        plcm)

    #=
        Always check redundancy, for now
    =#

    # number of added elements
    npivs = basis.ntotal

    # number of potential critical pairs to add
    npairs = basis.ndone * npivs
    for i in 1:npivs
        npairs += i
    end

    # make sure pairset and update hashtable have enough
    # space to store new pairs
    # TODO: we create too big array, can be fixed
    # @error "before" pairset.load npairs pairset.load+npairs length(pairset.pairs)
    check_enlarge_pairset!(pairset, npairs)
    pairset_size = length(pairset.pairs)
    # @error "after enlarge" length(pairset.pairs)

    # @error "" basis.ndone+1 basis.ntotal

    if basis.ndone + 1 <= basis.ntotal
        # red = 0
        # for each new element in basis
        # @error pairset.load basis.ndone+1 basis.ntotal
        for i in basis.ndone+1:basis.ntotal
            # check redundancy of new poly
            if is_redundant!(pairset, basis, ht, update_ht, i)
                continue
            end
            if length(plcm) < basis.ntotal + 1
                resize!(plcm, floor(Int, basis.ntotal*1.1)+ 1)
            end
            update_pairset!(pairset, basis, ht, update_ht, i, plcm)
        end
        # @error pairset.load
    end

    # @error "" red basis.ntotal-basis.ndone
    # println("--------------------------")
    # @error "after update" pairset.load length(pairset.pairs)

    update_basis!(basis, ht, update_ht)

    pairset_size
end

function select_tobereduced!(
                basis::Basis, tobereduced::Basis,
                matrix::MacaulayMatrix,
                symbol_ht::MonomialHashtable, ht::MonomialHashtable)

    # prepare to load all elems from tobereduced
    # into low rows
    reinitialize_matrix!(matrix, max(basis.ntotal, tobereduced.ntotal))
    resize!(matrix.lowrows, tobereduced.ntotal)

    # TODO
    etmp = zeros(UInt16, ht.explen)

    for i in 1:tobereduced.ntotal
        matrix.nrows += 1
        gen = tobereduced.gens[i]
        h = UInt32(0)
        matrix.lowrows[matrix.nrows] = multiplied_poly_to_matrix_row!(symbol_ht, ht, h, etmp, gen)
        matrix.low2coef[matrix.nrows] = i
    end

    basis.ntotal
    basis.nlead = basis.ndone = basis.ntotal
    basis.isred .= 0
    for i in 1:basis.nlead
        basis.nonred[i] = i
        basis.lead[i] = ht.hashdata[basis.gens[i][1]].divmask
    end
end


text:
@testset "Rocket Landing SOCP" begin
## Rocket landing problem
D,N = 3,101
n,m = 2D,D
dt = 0.1
tf = (N-1)*dt
x0 = SA_F64[10,10,100, 0,0,-10]
xf = @SVector zeros(n)
Q = Diagonal(@SVector fill(1.0, n)) * dt
R = Diagonal(@SVector fill(1e-1, m)) * dt
Qf = (N-1)*Q * 100

##
model = RobotZoo.DoubleIntegrator(3, gravity=SA[0,0,-9.81])
obj = LQRObjective(Q,R,Qf,xf,N)
cons = ConstraintList(n,m,N)
u_bnd = 100.0
connorm = NormConstraint(n, m, u_bnd, TO.SecondOrderCone(), :control)
add_constraint!(cons, GoalConstraint(xf), N)
add_constraint!(cons, connorm, 1:N-1)

prob = Problem(model, obj, x0, tf, xf=xf, constraints=cons)
solver = ALTROSolver(prob, projected_newton=false, show_summary=false)
solve!(solver)
@test iterations(solver) == 13 # 12

@test abs(maximum(norm.(controls(solver))) - u_bnd) < 1e-6
@test norm(states(solver)[end] - xf) < 1e-6
end

ext: jl


text:
"""
   update_perceiv_publ_opinion!(state, agent_idx)

Update an agent's perceived public opinion depending on its neighbors and current feed

# Arguments
- `state`: a tuple of the current graph and agent_list
- `agent_idx`: agent index

See also: [`Config`](@ref), [`update_opinion!`](@ref)
"""
function update_perceiv_publ_opinion!(
    state::Tuple{AbstractGraph, AbstractArray},
    agent_idx::Integer
)

    graph, agent_list = state
    this_agent = agent_list[agent_idx]

    input = inneighbors(graph, agent_idx)
    if length(input) != 0
        input_opinion_mean = mean(
            [agent_list[input_agent].opinion for input_agent in input]
        )
    else
        input_opinion_mean = this_agent.opinion
    end

    feed_opinions = [post.opinion for post in this_agent.feed]
    feed_weights = [post.weight for post in this_agent.feed]

    if length(feed_opinions) > 0
        feed_opinion_mean = (
            sum(
                [opinion * weight for (opinion, weight)
                in zip(feed_opinions, feed_weights)]
            )
            / sum(feed_weights)
        )
    else
        feed_opinion_mean = this_agent.opinion
    end

    this_agent.perceiv_publ_opinion = mean(
        [input_opinion_mean, feed_opinion_mean]
    )

    return state
end

# suppress output of include()
;


text:
"""
    function calc_new_potential_by_neighbors_3D(
        volume_weight::T,
        weights::NTuple{6,T},
        neighbor_potentials::NTuple{6,T},
    ) where {T}

Calculates and returns the new potential value of a grid point given
the potential values of the six neighbouring grid points (2 in each dimension) `neighbor_potentials`,
the weights corresponding to those six potentials of the neighboring grid points,
the weight due to the volume of the grid point itself (the voxel around it).

For more detailed information see Chapter 5.2.2 "Calculation of the Electric Potential"
Eqs. 5.31 to 5.37 in https://mediatum.ub.tum.de/doc/1620954/1620954.pdf.
"""
@inline function calc_new_potential_by_neighbors_3D(
    volume_weight::T,
    weights::NTuple{6,T},
    neighbor_potentials::NTuple{6,T},
) where {T}
    new_potential = weights[1] * neighbor_potentials[1]
    new_potential = muladd(weights[2], neighbor_potentials[2], new_potential)
    new_potential = muladd(weights[3], neighbor_potentials[3], new_potential)
    new_potential = muladd(weights[4], neighbor_potentials[4], new_potential)
    new_potential = muladd(weights[5], neighbor_potentials[5], new_potential)
    new_potential = muladd(weights[6], neighbor_potentials[6], new_potential)
    return new_potential * volume_weight
end

"""
    function handle_depletion(
        new_potential::T,
        imp_scale::T,
        neighbor_potentials::NTuple{6,T},
        q_eff_imp::T,
        volume_weight::T,
    )::Tuple{T, PointType} where {T}

This function handles the grid points with volumes which are not fully depleted.
The decision depends on the proposal for the new potential value, `new_potential_proposal`,
for the respective grid point in that iteration and the potential values of the neighboring grid points, `neighbor_potentials`.

If `vmin = minimum(neighbor_potentials) < new_potential_proposal < vmax = maximum(neighbor_potentials)` => fully depleted => imp_scale = 1

If it undershoots or overshoots, the impurity density in the grid point is scaled down via reducing `imp_scale ∈ [0, 1]` of this grid point.

This decision is based on the fact that the potential inside a solid-state
detector increases monotonically from a `p+`-contact towards an `n+`-contact.
Thus, there cannot be local extrema.

!!! note
    If a fixed charge impurity is present, e.g. due to a charged passivated surface,
    this handling is probably not valid anymore as the potential between a
    `p+`-contact towards an `n+`-contact is not required to change monotonically anymore.

"""
@inline @fastmath function handle_depletion(
    new_potential::T,
    imp_scale::T,
    neighbor_potentials::NTuple{6,T},
    q_eff_imp::T,
    volume_weight::T,
)::Tuple{T, T} where {T}
    vmin::T = min(neighbor_potentials...)
    vmax::T = max(neighbor_potentials...)

    is_ptype = q_eff_imp < 0
    neighbor_relevant_extremum = is_ptype ? vmin : vmax
    imp_contribution = q_eff_imp * volume_weight

    if !iszero(imp_contribution)
        imp_scale = (neighbor_relevant_extremum - new_potential) / imp_contribution
        imp_scale = min(max(imp_scale, T(0)), T(1))
        if vmax == vmin imp_scale = T(0) end
        new_potential += imp_scale * imp_contribution
        Δ = new_potential - neighbor_relevant_extremum
        if (is_ptype && Δ < 0) || (!is_ptype && Δ > 0)
            new_potential = neighbor_relevant_extremum
        end
    end

    new_potential, imp_scale
end

@inline function apply_over_relaxation(new_potential::T, old_potential::T, sor_const::T)::T where {T}
    new_potential -= old_potential
    new_potential = muladd(new_potential, sor_const, old_potential)
end

include("CPU_outerloop.jl")
include("CPU_middleloop.jl")
include("CPU_innerloop.jl")
include("CPU_update.jl")

include("GPU_kernel_Cylindrical.jl")
include("GPU_kernel_Cartesian3D.jl")
include("GPU_kernel.jl")
include("GPU_update.jl")

include("convergence.jl")

function mark_undep_bits!(point_types::Array{PointType, 3}, imp_scale::Array{T, 3}) where {T}
    @inbounds for i in eachindex(imp_scale)
        point_types[i] += undepleted_bit * (is_pn_junction_point_type(point_types[i]) && imp_scale[i] < 1)
    end
    nothing
end


text:
# Create a dictionary that bins a real-valued vector
# Used to compute fit-diff statistics
#  Example:
#  julia> w = [-2.,-1,0.,1.]
#
#  4-element Array{Float64,1}:
#   -2.0
#   -1.0
#    0.0
#    1.0
#
#  julia> wbins = create_bins(w,1.0)
#  DataStructures.Accumulator{Int64,Int64} with 4 entries:
#    0  => 1
#    -2 => 1
#    -1 => 1
#    1  => 1
#
#  julia> summarize_bins(wbins)
#  (1,1,1,1)

import DataStructures.Accumulator
import DataStructures.counter
export create_bins, increment_bins, summarize_bins

@doc """ function create_bins()
 Starting with a list of floats, count the number that are in the intervals [(i-1)/coutff,i/cutoff] for all integers i.
"""
function create_bins( vect::Vector{Float64}, cutoff::Float64 )
  bins = counter(Int64)
  for v in vect
    index = Int(floor(v/cutoff))
    push!(bins, index)
  end
  bins
end

@doc """ function increment_bins()
  Add value  x  to bins.
  The cutoff value should be the value used to create the bins.
"""
function increment_bins( bins::Accumulator{Int64,Int64}, x::Float64, cutoff::Float64 )
  index = Int(floor(x/cutoff))
  #println("inc bins x: ",x,"  index: ",index)
  push!(bins, index)
end

@doc """ function create_bins_by_filter()
 Create bins by using the Julia filter functon.
 Used to check the correctness of create_bins()
"""
function create_bins_by_filter( vect::Vector{Float64}, cutoff::Float64 )
  bins = Dict{Int64,Int64}()
  Min = minimum(vect)
  Max = maximum(vect)
  for i = collect(Int(floor(Min/cutoff)):(Int(floor(Max/cutoff))))
    len = length(filter(x->(((i*cutoff))<=x && x <(i+1)*cutoff),vect))
    println("i: ",i,"  lb: ",i*cutoff,"  ub: ",(i+1)*cutoff,"  len: ",len)
    #println("  len: ", length(filter(x->(((i*cutoff))<=x && x <(i+1)/cutoff),vect)) )
    if len > 0
      bins[i] = len
    end
  end
  bins
end

@doc """ check_bins()
  Creates bins dictionary, and then uses create_bins_by_filter to check correctness
"""
function check_bins( vect::Vector{Float64}, cutoff::Float64 )
  bins = create_bins( vect, cutoff )
  fbins =  create_bins_by_filter( vect, cutoff )
  for k in keys(bins)
    @assert bins[k] == fbins[k]
  end
  for k in keys(fbins)
    @assert bins[k] == fbins[k]
  end
end

# Checks that the lower_bounds_vect produce by bins_to_vector() is correct
function checkbins( vect::Vector{Float64}, cutoff::Float64 )
  bins = counter(Int64)
  for v in vect
    increment_bins( bins, v, cutoff )
  end
  println("bins: ",bins)
  (bin_vect,lower_bounds_vect) = bins_to_vector( bins, cutoff )
  println("lower_bounds_vect: ",lower_bounds_vect)
  counts = zeros(Int64,length(bin_vect))
  for i = 1:length(bin_vect)
    if i < length(bin_vect)
      len = length(filter(x->lower_bounds_vect[i]<=x && x<lower_bounds_vect[i+1],vect))
    else
      len = length(filter(x->lower_bounds_vect[i]<=x && x<lower_bounds_vect[i]+cutoff,vect))
    end
    counts[i] = len
  end
  if bin_vect == counts
    println("passed")
  else
    println("failed")
    println("bin_vect: ",bin_vect)
    println("counts:   ",counts )
  end
  counts
end

@doc """ function summarize_bins()
  Returns a 4-tuple of vect values in the following intervals:
  (-infinity,-cutoff)
  [-cutoff,0.0)
  [0.0,cutoff)
  [cutoff,infinity)
"""
function summarize_bins( bins::Accumulator{Int64,Int64} )
  min = minimum(keys(bins))
  max = maximum(keys(bins))
  neg_neutral = bins[-1]
  pos_neutral = bins[0]
  neg_count = min <= -2 ? sum( bins[k] for k = min:-2 ) : 0
  pos_count = 1 <= max ? sum( bins[k] for k = 1:max ) : 0
  return (neg_count,neg_neutral,pos_neutral,pos_count)
end

function bins_to_vector( bins::Accumulator{Int64,Int64}, cutoff::Float64 )
  min = minimum(keys(bins))
  max = maximum(keys(bins))
  println("bins cutoff: ",cutoff,"  min: ",min,"  max: ",max)
  bin_vect = zeros(Int64,(max-min+1))
  lower_bounds_vect = [cutoff*(min+i) for i = 0:(max-min)]
  for k=1:(max-min+1)
    bin_vect[k] = bins[k+min-1]
    lower_bounds_vec=k+min
    #println("k: ",k,"  k+min-1: ",k+min-1,"  bins[k+min-1]: ", bins[k+min-1])
  end
  (bin_vect,lower_bounds_vect)
end


text:
# This file is a part of SimilaritySearch.jl

using Dates
export LocalSearchAlgorithm, SearchGraph, SearchGraphOptions, VisitedVertices, NeighborhoodReduction, SatNeighborhood, IdentityNeighborhood, find_neighborhood, push_neighborhood!
export Callback, RandomHintsCallback, DisjointNeighborhoodHints

"""
    abstract type NeighborhoodReduction end

Overrides `Base.reduce(::NeighborhoodReduction, res::KnnResult, index::SearchGraph)` to postprocess `res` using some criteria.
Called from `find_neighborhood`, and returns a new KnnResult struct (perhaps a copy of res) since `push_neighborhood` captures
the reference of its output.
"""
abstract type NeighborhoodReduction end
abstract type LocalSearchAlgorithm end

"""
    abstract type Callback end

Abstract type to trigger callbacks after some number of insertions.
SearchGraph stores the callbacks in `callbacks` (a dictionary that associates symbols and callback objects);
A SearchGraph object controls when callbacks are fired using `callback_logbase` and `callback_starting`

"""
abstract type Callback end

### Basic operations on the index

@with_kw mutable struct OptimizeParametersCallback <: Callback
    error = :distance # :recall, :distance, :distance_and_searchtime
    ksearch::Int32 = 10
    numqueries::Int32 = 32
    initialpopulation::Int32 = 4
    maxpopulation::Int32 = 4
    tol::Float32 = 0.01
    maxiters::Int32 = 4
end

"""
    @with_kw mutable struct Neighborhood

Determines the size of the neighborhood, \$k\$ is adjusted as a callback, and it is intended to affect previously inserted vertices.
The neighborhood is designed to consider two components \$k=in+out\$, i.e. _in_coming and _out_going edges for each vertex.
- The \$out\$ size is computed as \$minsize + \\log(logbase, n)\$ where \$n\$ is the current number of indexed elements; this is computed searching
for \$out\$  elements in the current index.
- The \$in\$ size is computed as \$\\Delta in\$, i.e., this is not searched in the current index yet for accepting future edges.
- reduce is intended to postprocess neighbors (after search process, i.e., once out edges are computed); do not change \$k\$ but always must return a copy of the reduced result set.

Note: The underlying graph is undirected, in and out edges are fused in the same priority queue; old edges can be discarded when closer elements are found.
Note: Set \$logbase=Inf\$ to obtain a fixed number of \$in\$ nodes; and set \$minsize=0\$ to obtain a pure logarithmic growing neighborhood.

"""
@with_kw mutable struct Neighborhood
    k::Int32 = 2 # actual neighborhood
    ksearch::Int32 = 2
    logbase::Float32 = 2
    minsize::Int32 = 2
    Δ::Float32 = 1
    reduce::NeighborhoodReduction = IdentityNeighborhood()
end

Base.copy(N::Neighborhood; k=N.k, ksearch=N.ksearch, logbase=N.logbase, minsize=N.minsize, Δ=N.Δ, reduce=copy(N.reduce)) =
    Neighborhood(; k, ksearch, logbase, minsize, Δ, reduce)

struct NeighborhoodCallback <: Callback end

"""
    struct SearchGraph <: AbstractSearchContext

SearchGraph index. It stores a set of points that can be compared through a distance function `dist`.
The performance is determined by the search algorithm `search_algo` and the neighborhood policy.
It supports callbacks to adjust parameters as insertions are made.

- `hints`: Initial points for exploration (empty hints imply using random points)

Note: Parallel insertions should be made through `append!` function with `parallel_block > 1`

"""
@with_kw struct SearchGraph{DistType<:PreMetric, DataType<:AbstractVector, SType<:LocalSearchAlgorithm}<:AbstractSearchContext
    dist::DistType = SqL2Distance()
    db::DataType = Vector{Float32}[]
    links::Vector{KnnResult{Int32,Float32}} = KnnResult{Int32,Float32}[]
    locks::Vector{Threads.SpinLock} = Threads.SpinLock[]
    hints::Vector{Int32} = Int32[]
    search_algo::SType = BeamSearch()
    neighborhood::Neighborhood = Neighborhood()
    callbacks::Dict{Symbol,Callback} = Dict(
        #:parameters => OptimizeParametersCallback(),
        :hints => RandomHintsCallback(),
        :neighborhood => NeighborhoodCallback()
    )
    callback_logbase::Float32 = 1.5
    callback_starting::Int32 = 8
    verbose::Bool = true
end

@inline Base.length(g::SearchGraph) = length(g.locks)
include("visitedvertices.jl")

Base.copy(g::SearchGraph;
        dist=g.dist,
        db=g.db,
        links=g.links,
        locks=g.locks,
        hints=g.hints,
        search_algo=copy(g.search_algo),
        neighborhood=copy(g.neighborhood),
        callbacks=g.callbacks,
        callback_logbase=g.callback_logbase,
        callback_starting=g.callback_starting,
        verbose=true
) = SearchGraph(; dist, db, links, locks, hints, search_algo, neighborhood, callbacks, callback_logbase, callback_starting, verbose)

## search algorithms

include("ihc.jl")
include("beamsearch.jl")
## parameter optimization and neighborhood definitions
include("opt.jl")
include("neighborhood.jl")
include("hints.jl")

"""
    append!(index::SearchGraph, db; parallel_block=1, parallel_minimum_first_block=parallel_block, apply_callbacks=true)

Appends all items in db to the index. It can be made in parallel or sequentially.
In case of a parallel appending, then:
- `parallel_block` must be bigger than 1 and describes the batch size to append in parallel (i.e., in the order of thousands,
  depending on the size of the `db` and the number of available threads).
- `parallel_minimum_first_block` indicates the minimum number of items inserted sequentially before going parallel, it can be 0 if the index is already
  populated, defaults to `parallel_block`.

Note: Parallel doesn't trigger callbacks inside blocks.
"""
function Base.append!(index::SearchGraph, db;
    parallel_block=1, parallel_minimum_first_block=parallel_block, apply_callbacks=true)

    if parallel_block == 1
        for item in db
            push!(index, item)
        end

        return index
    end

    m = 0
    n = length(index) + length(db)
    while length(index) < parallel_minimum_first_block
        m += 1
        push!(index, db[m])
    end

    sp = length(index) + 1
    sp > n && return index

    resize!(index.db, n)
    resize!(index.links, n)
    for i in sp:n
        m += 1
        index.db[i] = db[m]
    end

    while sp < n
        ep = min(n, sp + parallel_block)
        index.verbose && println(stderr, "appending chunk ", (sp=sp, ep=ep, n=n), " ", Dates.now())

        # searching neighbors
        # @show length(index.links), length(index.db), length(db), length(index.locks), length(index), sp, ep
        Threads.@threads for i in sp:ep
           @inbounds index.links[i] = find_neighborhood(index, index.db[i], getknnresult(), getvisitedvertices(index))
        end

        # connecting neighbors
        k = index.neighborhood.k
        Threads.@threads for i in sp:ep
            @inbounds for (id, dist) in index.links[i]
                lock(index.locks[id])
                try
                    vertex = index.links[id]
                    vertex.k = max(maxlength(vertex), k)
                    push!(vertex, i => dist)
                finally
                    unlock(index.locks[id])
                end
            end
        end

        # increasing locks => new items are enabled for searching (and reported by length so they can also be hints)
        resize!(index.locks, ep)
        for i in sp:ep
            @inbounds index.locks[i] = Threads.SpinLock()
        end

        # apply callbacks
        apply_callbacks && callbacks(index, sp, ep)
        sp = ep + 1
    end

    index
end

"""
    push!(index::SearchGraph, item)

Appends `item` into the index.
"""
function push!(index::SearchGraph, item)
    neighbors = find_neighborhood(index, item, getknnresult(), getvisitedvertices(index))
    push_neighborhood!(index, item, neighbors)
    neighbors
end

"""
    search(index::SearchGraph, q, res::KnnResult; hints=index.hints, vstate=nothing)

Solves the specified query `res` for the query object `q`.
"""
function search(index::SearchGraph, q, res::KnnResult; hints=index.hints, vstate=nothing)
    if vstate === nothing
        vstate = getvisitedvertices(index)
    end

    length(index) > 0 && search(index.search_algo, index, q, res, hints, vstate)
    res
end

"""
    callback(opt::NeighborhoodCallback, index)

SearchGraph's callback for adjusting neighborhood strategy
"""
function callback(opt::NeighborhoodCallback, index)
    N = index.neighborhood
    N.ksearch = ceil(Int, N.minsize + log(N.logbase, length(index)))
    N.k = N.ksearch + ceil(Int, N.ksearch * N.Δ)
end

"""
    callbacks(index::SearchGraph, n=length(index), m=n+1)

Process all registered callbacks in `index`
"""
function callbacks(index::SearchGraph, n=length(index), m=n+1)
    if n >= index.callback_starting && ceil(Int, log(index.callback_logbase, n)) != ceil(Int, log(index.callback_logbase, m))
        for (name, callback_object) in index.callbacks
            index.verbose && println(stderr, "calling callback ", name, "; n=$(length(index)), type=", typeof(callback_object))
            callback(callback_object, index)
        end
    end
end


text:
# p23a.jl - eigenvalues of UNperturbed Laplacian on [-1,1]x[-1,1]
#         (compare p16.jl)

# Set up tensor product Laplacian and compute 4 eigenmodes:
N = 16; (D,x) = cheb(N); y = x;
xx = x[2:N]; yy = y[2:N];
D2 = D^2; D2 = D2[2:N,2:N]; I = eye(N-1);
L = -kron(I,D2) - kron(D2,I);                #  Laplacian
#f = @. exp(20*(yy-xx'-1));      # perturbation
#L = L + diagm(f[:]);
(D,V) = eig(L); ii = sortperm(D)[1:4];
D = D[ii]; V = V[:,ii];

# Reshape them to 2D grid, interpolate to finer grid, and plot:
xx = yy = x[end:-1:1];
xxx = yyy = -1:.02:1;
uu = zeros(N+1,N+1);
(ay,ax) = (repmat([.56 .04],2,1),repmat([.1,.5],1,2)); clf();
for i = 1:4
    uu[2:N,2:N] = reshape(V[:,i],N-1,N-1);
    uu = uu/norm(uu[:],Inf);
    s = Spline2D(xx,yy,reduce(flipdim,uu,1:2));
    uuu = evalgrid(s,xxx,yyy);
    axes([ax[i],ay[i],.38,.38]);
    contour(xxx,yyy,uuu,levels=-.9:.2:.9);
    axis("square");
    title("eig = $(signif(D[i]/(pi^2/4),13)) π^2/4");
end


text:
# addition of positive definite matrices

using PDMats

for T in [Float64,Float32]

  printstyled("Testing addition with eltype = $T\n", color=:blue)
  M = convert(Array{T,2},[4. -2. -1.; -2. 5. -1.; -1. -1. 6.])
  V = convert(Array{T,1},[1.5, 2.5, 2.0])
  local X = convert(T,2.0)

  pm1 = PDMat(M)
  pm2 = PDiagMat(V)
  pm3 = ScalMat(3,X)
  pm4 = X*I
  pm5 = PDSparseMat(sparse(M))

  pmats = Any[pm1, pm2, pm3] #, pm5]

  for p1 in pmats, p2 in pmats
      pr = p1 + p2
      @test size(pr) == size(p1)
      @test Matrix(pr) ≈ Matrix(p1) + Matrix(p2)

      pr = pdadd(p1, p2, convert(T,1.5))
      @test size(pr) == size(p1)
      @test Matrix(pr) ≈ Matrix(p1) + Matrix(p2) * convert(T,1.5)
  end

  for p1 in pmats
        pr = p1 + pm4
        @test size(pr) == size(p1)
        @test Matrix(pr) ≈ Matrix(p1) + pm4
  end
end


text:
export duration, nframes, framerate, nchannels, sampletype
"""

    duration(x)

Return the duration of the signal in seconds, if known. May return `missing`
(e.g. for a stream).

!!! note

    A fallback implementation of `duration` uses `nframes(x) / framerate(x)`.
    However, if one or both of these is `missing` and you want `duartion` to
    return a non-missing value, you can define a custom method of `duration`.

"""
duration(x) = nframes(x) / framerate(x)

"""

    nframes(x)

Returns the number of frames in the signal, if known. May return `missing`
 (e.g. for a file stream).

"""
function nframes
end

"""

    framerate(x)

Returns the frame rate of the signal (in Hertz). May return `missing` if the
frame rate is unknown.

"""
function framerate
end

"""

    nchannels(x)

Returns the number of channels in the signal.

"""
function nchannels
end

"""

    sampletype(x)

Returns the element type of an individual channel of a signal (e.g. `Float64`).

!!! note

    The result of `sampletype` and `eltype` (when defined) are often the
    same. They are distinct so that these two can diverge when appropriate.

"""
function sampletype
end


text:
module SWT

export
    # stationary wavelet transform
    sdwt,
    sdwt!,
    swpt,
    swpt!,
    swpd,
    swpd!,
    # inverse stationary wavelet transform
    isdwt,
    isdwt!,
    iswpt,
    iswpt!,
    iswpd,
    iswpd!,
    # Transforms on all signals
    sdwtall,
    swptall,
    swpdall,
    isdwtall,
    iswptall,
    iswpdall

using Wavelets

using ..Utils

# ========== Stationary DWT ==========
# ----- SDWT with allocation -----
@doc raw"""
    sdwt(x, wt[, L])

Computes the stationary discrete wavelet transform (SDWT) for `L` levels.

# Arguments
- `x::AbstractVector{T}` or `x::AbstractMatrix{T} where T<:Number`: Original signal,
    preferably of size 2ᴷ where ``K \in \mathbb{N}``.
- `wt::OrthoFilter`: Orthogonal wavelet filter.
- `L::Integer`: (Default: `maxtransformlevels(x)`) Number of levels of decomposition.

# Returns
`::Matrix{T}` or `::Array{T,3}`: Output from SDWT on `x`.

# Examples
```julia
using Wavelets, WaveletsExt

# Setup
x = generatesignals(:heavysine)
wt = wavelet(WT.haar)

# SDWT
xw = sdwt(x, wt)
```

**See also:** [`swpd`](@ref), [`swpt`](@ref), [`isdwt`](@ref)
"""
function sdwt(x::AbstractArray{T},
              wt::OrthoFilter,
              L::Integer = maxtransformlevels(x)) where T<:Number
    # Sanity check
    @assert 1 ≤ ndims(x) ≤ 2
    @assert L ≤ maxtransformlevels(x) || throw(ArgumentError("Too many transform levels (length(x) < 2^L"))
    @assert L ≥ 1 || throw(ArgumentError("L must be ≥ 1"))
    # Setup
    sz = size(x)
    N = ndims(x)+1
    k = N==2 ? L+1 : 3*L+1
    xw = Array{T,N}(undef, (sz...,k))
    # Compute transforms
    sdwt!(xw, x, wt, L)
    return xw
end

# ----- 1D SDWT without allocation -----
@doc raw"""
    sdwt!(xw, x, wt[, L])

Same as `sdwt` but without array allocation.

# Arguments
- `xw::AbstractArray{T,2}` or `xw::AbstractArray{T,3} where T<:Number`: An allocated array
  of dimension `(n,L+1)` to write the outputs of `x` onto.
- `x::AbstractVector{T}` or `x::AbstractMatrix{T} where T<:Number`: Original signal, preferably of size 2ᴷ where ``K
  \in \mathbb{N}``.
- `wt::OrthoFilter`: Orthogonal wavelet filter.
- `L::Integer`: (Default: `maxtransformlevels(x)`) Number of levels of decomposition.

# Returns
`xw::Array{T,2}` or `xw::Array{T,3}`: Output from SDWT on `x`.

# Examples
```julia
using Wavelets, WaveletsExt

# Setup
x = generatesignals(:heavysine, 7)
wt = wavelet(WT.haar)

# SDWT
xw = Matrix{Float64}(undef, (128,5))
sdwt!(xw, x, wt, 4)
```

**See also:** [`sdwt`](@ref)
"""
function sdwt!(xw::AbstractArray{T,2},
               x::AbstractVector{T},
               wt::OrthoFilter,
               L::Integer = maxtransformlevels(x)) where T<:Number
    # Sanity check
    @assert L ≤ maxtransformlevels(x) ||
        throw(ArgumentError("Too many transform levels (length(x) < 2^L"))
    @assert L ≥ 1 || throw(ArgumentError("L must be ≥ 1"))

    # Setup
    g, h = WT.makereverseqmfpair(wt, true)
    xw[:,end] = x

    # SDWT
    for d in 0:(L-1)
        @inbounds v = xw[:, L-d+1]       # Parent node
        w₁ = @view xw[:, L-d]            # Scaling coefficients
        w₂ = @view xw[:, L-d+1]          # Detail coefficients
        @inbounds sdwt_step!(w₁, w₂, v, d, h, g)
    end
    return xw
end
# ----- 2D SDWT with no allocation -----
function sdwt!(xw::AbstractArray{T,3},
               x::AbstractMatrix{T},
               wt::OrthoFilter,
               L::Integer = maxtransformlevels(x)) where T<:Number
    # Sanity check
    @assert L ≤ maxtransformlevels(x) ||
        throw(ArgumentError("Too many transform levels"))
    @assert L ≥ 1 || throw(ArgumentError("L must be ≥ 1"))
    @assert size(xw,3) == 3*L+1

    # Setup
    n, m = size(x)
    g, h = WT.makereverseqmfpair(wt, true)
    xw[:,:,end] = x
    temp = Array{T,3}(undef, (n,m,2))

    # SDWT
    for d in 0:(L-1)
        @inbounds v = xw[:,:, 3*(L-d)+1]       # Parent node
        @inbounds w₁ = @view xw[:,:, 3*(L-d)-2]          # Scaling + Scaling coefficients
        @inbounds w₂ = @view xw[:,:, 3*(L-d)-1]          # Detail + Scaling coefficients
        @inbounds w₃ = @view xw[:,:, 3*(L-d)]            # Scaling + Detail coefficients
        @inbounds w₄ = @view xw[:,:, 3*(L-d)+1]          # Detail + Detail coefficients
        @inbounds sdwt_step!(w₁, w₂, w₃, w₄, v, d, h, g, temp)
    end
    return xw
end

# ----- ISDWT (Shift based) with allocation -----
"""
    isdwt(xw, wt[, sm])

Computes the inverse stationary discrete wavelet transform (iSDWT) on `xw`.

# Arguments
- `xw::AbstractArray{T,2}` or `xw::AbstractArray{T,3} where T<:Number`: SDWT-transformed
  array.
- `wt::OrthoFilter`: Orthogonal wavelet filter.
- `sm::Integer`: If `sm` is included as an argument, the `sm`-shifted inverse transform will
  be computed. This results in significantly faster computation, but fails to fully utilize
  the strength of redundant wavelet transforms.

# Returns
`::Vector{T}` or `::Matrix{T}`: Inverse transformed signal.

# Examples
```julia
using Wavelets, WaveletsExt

# Setup
x = generatesignals(:heavysine)
wt = wavelet(WT.haar)

# SDWT
xw = sdwt(x, wt)

#  Shift-based iSDWT
x̂ = isdwt(xw, wt, 5)

# Average-based iSDWT
x̃ = isdwt(xw, wt)
```

**See also:** [`isdwt_step`](@ref), [`iswpt`](@ref), [`sdwt`](@ref)
"""
function isdwt(xw::AbstractArray{T}, wt::OrthoFilter, sm::Integer) where T<:Number
    @assert 2 ≤ ndims(xw) ≤ 3
    # Setup
    sz = size(xw)[1:(end-1)]
    N = ndims(xw)-1
    x = Array{T,N}(undef, sz)
    # Transform
    isdwt!(x, xw, wt, sm)
    return x
end
# ----- ISDWT (Average based) with allocation -----
function isdwt(xw::AbstractArray{T}, wt::OrthoFilter) where T<:Number
    @assert 2 ≤ ndims(xw) ≤ 3
    # Setup
    sz = size(xw)[1:(end-1)]
    N = ndims(xw)-1
    x = Array{T,N}(undef, sz)
    # Transform
    isdwt!(x, xw, wt)
    return x
end

# ----- 1D ISDWT (Shift based) without allocation -----
"""
    isdwt!(x, xw, wt[, sm])

Same as `isdwt` but with no array allocation.

# Arguments
- `x::AbstractVector{T}` or `x::AbstractMatrix{T} where T<:Number`: Allocation for
  reconstructed signal.
- `xw::AbstractArray{T,2}` or `xw::AbstractArray{T,3} where T<:Number`: SDWT-transformed
  array.
- `wt::OrthoFilter`: Orthogonal wavelet filter.
- `sm::Integer`: If `sm` is included as an argument, the `sm`-shifted inverse transform will
  be computed. This results in significantly faster computation, but fails to fully utilize
  the strength of redundant wavelet transforms.

# Returns
`x::Vector{T}` or `x::Matrix{T}`: Inverse transformed signal.

# Examples
```julia
using Wavelets, WaveletsExt

# Setup
x = generatesignals(:heavysine)
wt = wavelet(WT.haar)

# SDWT
xw = sdwt(x, wt)

#  Shift-based iSDWT
x̂ = similar(x)
isdwt!(x̂, xw, wt, 5)

# Average-based iSDWT
isdwt(x̂, xw, wt)
```

**See also:** [`isdwt`](@ref)
"""
function isdwt!(x::AbstractVector{T},
                xw::AbstractArray{T,2},
                wt::OrthoFilter,
                sm::Integer) where T<:Number
    # Sanity check
    _, k = size(xw)
    L = k-1
    @assert 0 ≤ log2(sm) < L

    # Setup
    g, h = WT.makereverseqmfpair(wt, true)
    sd = Utils.main2depthshift(sm, L)
    for i in eachindex(x)
        @inbounds x[i] = xw[i,1]
    end

    # ISDWT
    for d in reverse(0:(L-1))
        sv = sd[d+1]
        sw = sd[d+2]
        w₁ = copy(x)
        @inbounds w₂ = @view xw[:,L-d+1]
        @inbounds isdwt_step!(x, w₁, w₂, d, sv, sw, h, g)
    end
    return x
end
# ----- 2D ISDWT (Shift based) without allocation -----
function isdwt!(x::AbstractMatrix{T},
                xw::AbstractArray{T,3},
                wt::OrthoFilter,
                sm::Integer) where T<:Number
    # Sanity check
    n, m, k = size(xw)
    L = (k-1) ÷ 3
    @assert 0 ≤ log2(sm) ≤ L

    # Setup
    g, h = WT.makereverseqmfpair(wt, true)
    sd = Utils.main2depthshift(sm, L)
    temp = Array{T,3}(undef, (n,m,2))
    for i in eachindex(x)
        @inbounds x[i] = xw[i]
    end

    # ISDWT
    for d in reverse(0:(L-1))
        sv = sd[d+1]
        sw = sd[d+2]
        w₁ = copy(x)
        @inbounds w₂ = @view xw[:,:,3*(L-d)-1]
        @inbounds w₃ = @view xw[:,:,3*(L-d)]
        @inbounds w₄ = @view xw[:,:,3*(L-d)+1]
        @inbounds isdwt_step!(x, w₁, w₂, w₃, w₄, d, sv, sw, h, g, temp)
    end
    return x
end

# ----- 1D ISDWT (Average based) without allocation -----
function isdwt!(x::AbstractVector{T}, xw::AbstractArray{T,2}, wt::OrthoFilter) where
                T<:Number
    # Setup
    g, h = WT.makereverseqmfpair(wt, true)
    _, k = size(xw)
    L = k-1
    for i in eachindex(x)
        @inbounds x[i] = xw[i,1]
    end

    # ISDWT
    for d in reverse(0:(L-1))
        w₁ = copy(x)
        @inbounds w₂ = @view xw[:,L-d+1]
        @inbounds isdwt_step!(x, w₁, w₂, d, h, g)
    end
    return x
end
# ----- 2D ISDWT (Average based) without allocation -----
function isdwt!(x::AbstractMatrix{T}, xw::AbstractArray{T,3}, wt::OrthoFilter) where
                T<:Number
    # Sanity check
    n, m, k = size(xw)
    L = (k-1) ÷ 3

    # Setup
    g, h = WT.makereverseqmfpair(wt, true)
    temp = Array{T,3}(undef, (n,m,2))
    for i in eachindex(x)
        @inbounds x[i] = xw[i]
    end

    # ISDWT
    for d in reverse(0:(L-1))
        w₁ = copy(x)
        @inbounds w₂ = @view xw[:,:,3*(L-d)-1]
        @inbounds w₃ = @view xw[:,:,3*(L-d)]
        @inbounds w₄ = @view xw[:,:,3*(L-d)+1]
        @inbounds isdwt_step!(x, w₁, w₂, w₃, w₄, d, h, g, temp)
    end
    return x
end

# ========== Stationary WPT ==========
# ----- SWPT with allocation -----
@doc raw"""
    swpt(x, wt[, L])

Computes `L` levels of stationary wavelet packet transform (SWPT) on `x`.

# Arguments
- `x::AbstractVector{T}` or `x::AbstractMatrix{T} where T<:Number`: Original signal, preferably of size 2ᴷ where ``K
  \in \mathbb{N}``.
- `wt::OrthoFilter`: Orthogonal wavelet filter.
- `L::Integer`: (Default: `maxtransformlevels(x)`) Number of levels of decomposition.

# Returns
`::Matrix{T}` or `::Array{T,3}`: Output from SWPT on `x`.

# Examples
```julia
using Wavelets, WaveletsExt

# Setup
x = generatesignals(:heavysine)
wt = wavelet(WT.haar)

# SWPT
xw = swpt(x, wt)
```

**See also:** [`sdwt`](@ref), [`swpd`](@ref)
"""
function swpt(x::AbstractArray{T},
              wt::OrthoFilter,
              L::Integer = maxtransformlevels(x)) where T<:Number
    # Sanity check
    @assert 1 ≤ ndims(x) ≤ 2
    @assert L ≤ maxtransformlevels(x) ||
            throw(ArgumentError("Too many transform levels (length(x) < 2ᴸ)"))
    @assert L ≥ 1 || throw(ArgumentError("L must be ≥ 1"))
    # Setup
    sz = size(x)
    N = ndims(x)+1
    k = N==2 ? 1<<L : 1<<(2*L)
    xw = Array{T,N}(undef, (sz..., k))
    # Transform
    swpt!(xw, x, wt, L)
    return xw
end

# ----- 1D SWPT without allocation -----
@doc raw"""
    swpt!(xw, x, wt[, L])

Same as `swpt` but without array allocation.

# Arguments
- `xw::AbstractArray{T,2}` or `xw::AbstractArray{T,3} where T<:Number`: Allocation for transformed signal.
- `x::AbstractVector{T}` or `x::AbstractMatrix{T} where T<:Number`: Original signal, preferably of size 2ᴷ where ``K
  \in \mathbb{N}``.
- `wt::OrthoFilter`: Orthogonal wavelet filter.
- `L::Integer`: (Default: `maxtransformlevels(x)`) Number of levels of decomposition.

# Returns
`xw::Matrix{T}` or `xw::Array{T,3}`: Output from SWPT on `x`.

# Examples
```julia
using Wavelets, WaveletsExt

# Setup
x = generatesignals(:heavysine)
wt = wavelet(WT.haar)

# SWPT
xw = Array{Float64,2}(undef, (128,128))
swpt!(xw, x, wt)
```

**See also:** [`swpt`](@ref)
"""
function swpt!(xw::AbstractArray{T,2},
               x::AbstractVector{T},
               wt::OrthoFilter,
               L::Integer = maxtransformlevels(x)) where T<:Number
    # Sanity check
    @assert L ≤ maxtransformlevels(x) ||
            throw(ArgumentError("Too many transform levels (length(x) < 2ᴸ)"))
    @assert L ≥ 1 || throw(ArgumentError("L must be ≥ 1"))
    @assert size(xw,2) == 1<<L
    @assert size(xw,1) == size(x,1)

    # Setup
    g, h = WT.makereverseqmfpair(wt, true)
    for i in eachindex(x)
        @inbounds xw[i,1] = x[i]
    end

    # SWPT for L levels
    for d in 0:(L-1)
        nn = 1<<d               # Number of nodes at current level
        for b in 0:(nn-1)
            np = (1<<L)÷nn      # Parent node length
            nc = np÷2           # Child node length
            j₁ = (2*b)*nc + 1   # Parent and (scaling) child index
            j₂ = (2*b+1)*nc + 1 # Detail child index
            @inbounds v = xw[:,j₁]
            @inbounds w₁ = @view xw[:,j₁]
            @inbounds w₂ = @view xw[:,j₂]
            # Overwrite output of SDWT directly onto xw
            @inbounds sdwt_step!(w₁, w₂, v, d, h, g)
        end
    end
    return xw
end
# ----- 2D SWPT without allocation -----
function swpt!(xw::AbstractArray{T,3},
               x::AbstractMatrix{T},
               wt::OrthoFilter,
               L::Integer = maxtransformlevels(x)) where T<:Number
    @assert L ≤ maxtransformlevels(x) ||
        throw(ArgumentError("Too many transform levels."))
    @assert L ≥ 1 || throw(ArgumentError("L must be ≥ 1."))
    @assert size(xw,3) == 4^L
    @assert size(xw,1) == size(x,1)
    @assert size(xw,2) == size(x,2)

    # Setup
    n, m = size(x)
    g, h = WT.makereverseqmfpair(wt, true)
    temp = Array{T,3}(undef, (n,m,2))
    for i in eachindex(x)
        @inbounds xw[i] = x[i]
    end

    # SWPT for L levels
    for d in 0:(L-1)
        nn = 4^d            # Number of nodes at current level
        for b in 0:(nn-1)
            np = (4^L)÷nn       # Number of slices occupied by parent
            nc = np÷4           # Number of slices occupied by children
            j₁ = (4*b)*nc + 1   # Parent and scaling+scaling child index
            j₂ = (4*b+1)*nc + 1 # Scaling+detail child index
            j₃ = (4*b+2)*nc + 1 # Detail+scaling child index
            j₄ = (4*b+3)*nc + 1 # Detail+detail child index
            v = xw[:,:,j₁]
            w₁ = @view xw[:,:,j₁]
            w₂ = @view xw[:,:,j₂]
            w₃ = @view xw[:,:,j₃]
            w₄ = @view xw[:,:,j₄]
            # Overwrite output of SDWT directly onto xw
            @inbounds sdwt_step!(w₁, w₂, w₃, w₄, v, d, h, g, temp)
        end
    end
    return xw
end

# ----- ISWPT (Shift based) with allocation -----
"""
    iswpt(xw, wt[, sm])

Computes the inverse stationary wavelet packet transform (iSWPT) on `xw`.

# Arguments
- `xw::AbstractArray{T,2}` or `xw::AbstractArray{T,3} where T<:Number`: SWPT-transformed array.
- `wt::OrthoFilter`: Orthogonal wavelet filter.
- `sm::Integer`: If `sm` is included as an argument, the `sm`-shifted inverse transform will
  be computed. This results in significantly faster computation, but fails to fully utilize
  the strength of redundant wavelet transforms.

# Returns
`::Vector{T}` or `::Matrix{T}`: Inverse transformed signal.

# Examples
```julia
using Wavelets, WaveletsExt

# Setup
x = generatesignals(:heavysine)
wt = wavelet(WT.haar)

# SWPT
xw = swpt(x, wt)

# Shift-based iSWPT
x̂ = iswpt(xw, wt, 5)

# Average-based iSWPT
x̃ = iswpt(xw, wt)
```

**See also:** [`isdwt_step`](@ref), [`isdwt`](@ref), [`swpt`](@ref)
"""
function iswpt(xw::AbstractArray{T}, wt::OrthoFilter, sm::Integer) where T<:Number
    @assert 2 ≤ ndims(xw) ≤ 3
    # Setup
    sz = size(xw)[1:(end-1)]
    N = ndims(xw)-1
    x = Array{T,N}(undef, sz)
    # Transform
    iswpt!(x, xw, wt, sm)
    return x
end
# ----- ISWPT (Average based) with allocation -----
function iswpt(xw::AbstractArray{T}, wt::OrthoFilter) where T<:Number
    @assert 2 ≤ ndims(xw) ≤ 3
    # Setup
    sz = size(xw)[1:(end-1)]
    N = ndims(xw)-1
    x = Array{T,N}(undef, sz)
    # Transform
    iswpt!(x, xw, wt)
    return x
end

# ----- 1D ISWPT (Shift based) without allocation -----
"""
    iswpt!(x, xw, wt[, sm])

Same as `iswpt` but with no array allocation.

# Arguments
- `x::AbstractVector{T}` or `x::AbstractMatrix{T} where T<:Number`: Allocation for inverse
  transform.
- `xw::AbstractArray{T,2}` or `xw::AbstractArray{T,3} where T<:Number`: SWPT-transformed
  array.
- `wt::OrthoFilter`: Orthogonal wavelet filter.
- `sm::Integer`: If `sm` is included as an argument, the `sm`-shifted inverse transform will
  be computed. This results in significantly faster computation, but fails to fully utilize
  the strength of redundant wavelet transforms.

# Returns
`x::Vector{T}` or `x::Matrix{T}`: Inverse transformed signal.

# Examples
```julia
using Wavelets, WaveletsExt

# Setup
x = generatesignals(:heavysine)
wt = wavelet(WT.haar)

# SWPT
xw = swpt(x, wt)

# Shift-based iSWPT
x̂ = similar(x)
iswpt!(x̂, xw, wt, 5)

# Average-based iSWPT
iswpt!(x̂, xw, wt)
```

**See also:** [`iswpt`](@ref)
"""
function iswpt!(x::AbstractVector{T},
                xw::AbstractArray{T,2},
                wt::OrthoFilter,
                sm::Integer) where T<:Number
    # Sanity check
    n, m = size(xw)
    @assert isdyadic(m) || throw(ArgumentError("Number of columns of xw is not dyadic."))
    @assert ndyadicscales(m) ≤ maxtransformlevels(n) ||
            throw(ArgumentError("Number of nodes in `xw` is more than possible number of nodes at any depth for signal of length `n`"))

    # Setup
    g, h = WT.makereverseqmfpair(wt, true)
    temp = copy(xw)                         # Temp. array to store intermediate outputs
    L = ndyadicscales(m)                    # Number of decompositions from xw
    sd = Utils.main2depthshift(sm, L)       # Shifts at each depth

    # ISWPT
    for d in reverse(0:(L-1))
        nn = 1<<d
        for b in 0:(nn-1)
            sv = sd[d+1]                    # Parent shift
            sw = sd[d+2]                    # Child shift
            np = (1<<L)÷nn                  # Parent node length
            nc = np÷2                       # Child node length
            j₁ = (2*b)*nc + 1               # Parent and (scaling) child index
            j₂ = (2*b+1)*nc + 1             # Detail child index
            @inbounds v = d==0 ? x : @view temp[:,j₁]
            @inbounds w₁ = temp[:,j₁]
            @inbounds w₂ = @view temp[:,j₂]
            # Overwrite output of iSWPT directly onto temp
            @inbounds isdwt_step!(v, w₁, w₂, d, sv, sw, h, g)
        end
    end
    return x
end
# ----- 2D ISWPT (Shift based) without allocation
function iswpt!(x::AbstractMatrix{T},
                xw::AbstractArray{T,3},
                wt::OrthoFilter,
                sm::Integer) where T<:Number
    # Sanity check
    n, m, k = size(xw)
    @assert isinteger(log(4,k)) || throw(ArgumentError("Size of dimension 3 is not a power of 4."))
    @assert size(x,1) == n
    @assert size(x,2) == m
    @assert log(4,k) ≤ maxtransformlevels(x)

    # Setup
    g, h = WT.makereverseqmfpair(wt, true)
    xwₜ = copy(xw)
    temp = Array{T,3}(undef, (n,m,2))       # Temporary array
    L = log(4,k) |> Int                     # Number of decomposition levels of xw
    sd = Utils.main2depthshift(sm, L)       # Shifts at each depth

    # ISWPT
    for d in reverse(0:(L-1))
        nn = 4^d
        for b in 0:(nn-1)
            sv = sd[d+1]                    # Parent shift
            sw = sd[d+2]                    # Child shift
            np = (4^L)÷nn                   # Number of slices occupied by parent
            nc = np÷4                       # Number of slices occupied by children
            j₁ = (4*b)*nc + 1               # Parent and scaling+scaling child index
            j₂ = (4*b+1)*nc + 1             # Scaling+detail child index
            j₃ = (4*b+2)*nc + 1             # Detail+scaling child index
            j₄ = (4*b+3)*nc + 1             # Detail+detail child index
            v = d==0 ? x : @view xwₜ[:,:,j₁]
            w₁ = xwₜ[:,:,j₁]
            w₂ = @view xwₜ[:,:,j₂]
            w₃ = @view xwₜ[:,:,j₃]
            w₄ = @view xwₜ[:,:,j₄]
            # Overwrite output of iSWPT directly onto temp
            @inbounds isdwt_step!(v, w₁, w₂, w₃, w₄, d, sv, sw, h, g, temp)
        end
    end
    return x
end

# ----- 1D ISWPT (Average based) withou allocation -----
function iswpt!(x::AbstractVector{T}, xw::AbstractArray{T,2}, wt::OrthoFilter) where
                T<:Number
    # Sanity check
    n, m = size(xw)
    @assert isdyadic(m) || throw(ArgumentError("Number of columns of xw is not dyadic."))
    @assert ndyadicscales(m) ≤ maxtransformlevels(n) ||
            throw(ArgumentError("Number of nodes in `xw` is more than possible number of nodes at any depth for signal of length `n`"))

    # Setup
    g, h = WT.makereverseqmfpair(wt, true)
    temp = copy(xw)                         # Temp. array to store intermediate outputs
    L = ndyadicscales(m)                    # Number of decompositions from xw

    # ISWPT
    for d in reverse(0:(L-1))
        nn = 1<<d
        for b in 0:(nn-1)
            np = (1<<L)÷nn                              # Parent node length
            nc = np÷2                                   # Child node length
            j₁ = (2*b)*nc + 1                           # Parent and (scaling) child index
            j₂ = (2*b+1)*nc + 1                         # Detail child index
            @inbounds v = d==0 ? x : @view temp[:,j₁]   # Parent node
            @inbounds w₁ = temp[:,j₁]                   # Scaling node
            @inbounds w₂ = @view temp[:,j₂]             # Detail node
            # Overwrite output of iSWPT directly onto temp
            @inbounds isdwt_step!(v, w₁, w₂, d, h, g)
        end
    end
    return x
end
# ----- 2D ISWPT (Average based) without allocation -----
function iswpt!(x::AbstractMatrix{T}, xw::AbstractArray{T,3}, wt::OrthoFilter) where
                T<:Number
    # Sanity check
    n, m, k = size(xw)
    @assert isinteger(log(4,k)) || throw(ArgumentError("Size of dimension 3 is not a power of 4."))
    @assert size(x,1) == n
    @assert size(x,2) == m
    @assert log(4,k) ≤ maxtransformlevels(x)

    # Setup
    g, h = WT.makereverseqmfpair(wt, true)
    xwₜ = copy(xw)
    temp = Array{T,3}(undef, (n,m,2))       # Temporary array
    L = log(4,k) |> Int                     # Number of decomposition levels of xw

    # ISWPT
    for d in reverse(0:(L-1))
        nn = 4^d
        for b in 0:(nn-1)
            np = (4^L)÷nn                   # Number of slices occupied by parent
            nc = np÷4                       # Number of slices occupied by children
            j₁ = (4*b)*nc + 1               # Parent and scaling+scaling child index
            j₂ = (4*b+1)*nc + 1             # Scaling+detail child index
            j₃ = (4*b+2)*nc + 1             # Detail+scaling child index
            j₄ = (4*b+3)*nc + 1             # Detail+detail child index
            v = d==0 ? x : @view xwₜ[:,:,j₁]
            w₁ = xwₜ[:,:,j₁]
            w₂ = @view xwₜ[:,:,j₂]
            w₃ = @view xwₜ[:,:,j₃]
            w₄ = @view xwₜ[:,:,j₄]
            # Overwrite output of iSWPT directly onto temp
            @inbounds isdwt_step!(v, w₁, w₂, w₃, w₄, d, h, g, temp)
        end
    end
    return x
end

# ========== Stationary WPD ==========
# ----- SWPD with allocation -----
@doc raw"""
    swpd(x, wt[, L])

Computes `L` levels of stationary wavelet packet decomposition (SWPD) on `x`.

# Arguments
- `x::AbstractVector{T}` or `x::AbstractMatrix{T} where T<:Number`: Original signal,
  preferably of size 2ᴷ where ``K \in \mathbb{N}``.
- `wt::OrthoFilter`: Orthogonal wavelet filter.
- `L::Integer`: (Default: `maxtransformlevels(x)`) Number of levels of decomposition.

# Returns
`::Matrix{T}` or `::Array{T,3}`: Output from SWPD on `x`.

# Examples
```julia
using Wavelets, WaveletsExt

# Setup
x = generatesignals(:heavysine)
wt = wavelet(WT.haar)

# SWPD
xw = swpd(x, wt)
```

**See also:** [`iswpd`](@ref), [`swpt`](@ref)
"""
function swpd(x::AbstractArray{T},
              wt::OrthoFilter,
              L::Integer = maxtransformlevels(x)) where T<:Number
    # Sanity check
    @assert 1 ≤ ndims(x) ≤ 2
    @assert L ≤ maxtransformlevels(x) ||
        throw(ArgumentError("Too many transform levels (length(x) < 2^L"))
    @assert L ≥ 1 || throw(ArgumentError("L must be ≥ 1"))
    # Setup
    sz = size(x)                            # Signal size
    N = ndims(x)+1
    k = N==2 ? 1<<(L+1)-1 : sum(4 .^(0:L))  # Total number of nodes
    xw = Array{T}(undef, (sz...,k))
    # Transform
    swpd!(xw, x, wt, L)
    return xw
end

# ----- 1D SWPD without allocation -----
@doc raw"""
    swpd!(xw, x, wt[, L])

Same as `swpd` but without array allocation.

# Arguments
- `xw::AbstractArray{T,2}` or `xw::AbstractArray{T,3} where T<:Number`: Allocation for
  transformed signal.
- `x::AbstractVector{T}` or `x::AbstractMatrix{T} where T<:Number`: Original signal,
  preferably of size 2ᴷ where ``K \in \mathbb{N}``.
- `wt::OrthoFilter`: Orthogonal wavelet filter.
- `L::Integer`: (Default: `maxtransformlevels(x)`) Number of levels of decomposition.

# Returns
`xw::Matrix{T}` or `xw::Array{T,3}`: Output from SWPD on `x`.

# Examples
```julia
using Wavelets, WaveletsExt

# Setup
x = generatesignals(:heavysine)
wt = wavelet(WT.haar)

# SWPD
xw = Matrix{T}(undef, (128, 255))
swpd!(xw, x, wt)
```

**See also:** [`swpd`](@ref)
"""
function swpd!(xw::AbstractArray{T,2},
               x::AbstractVector{T},
               wt::OrthoFilter,
               L::Integer = maxtransformlevels(x)) where T<:Number
    # Sanity check
    @assert L <= maxtransformlevels(x) ||
        throw(ArgumentError("Too many transform levels (length(x) < 2^L"))
    @assert L >= 1 || throw(ArgumentError("L must be >= 1"))

    # Setup
    g, h = WT.makereverseqmfpair(wt, true)
    n₀ = 1<<(L+1)-1                     # Total number of nodes
    n₁ = n₀ - (1<<L)                    # Total number of nodes excluding leaf nodes
    for i in eachindex(x)
        @inbounds xw[i,1] = x[i]
    end

    # SWPD
    for i in 1:n₁
        d = getdepth(i,:binary)
        j₁ = getchildindex(i,:left)
        j₂ = getchildindex(i,:right)
        @inbounds v = @view xw[:,i]
        @inbounds w₁ = @view xw[:,j₁]
        @inbounds w₂ = @view xw[:,j₂]
        @inbounds sdwt_step!(w₁, w₂, v, d, h, g)
    end
    return xw
end
# ----- 2D SWPD without allocation -----
function swpd!(xw::AbstractArray{T,3},
               x::AbstractMatrix{T},
               wt::OrthoFilter,
               L::Integer = maxtransformlevels(x)) where T<:Number
    # Sanity check
    n, m, k = size(xw)
    @assert k == sum(4 .^(0:L))
    @assert n == size(x,1)
    @assert m == size(x,2)

    # Setup
    g, h = WT.makereverseqmfpair(wt, true)
    temp = Array{T,3}(undef, (n,m,2))
    n₁ = k - (4^L)                    # Total number of nodes excluding leaf nodes
    for i in eachindex(x)
        @inbounds xw[i] = x[i]
    end

    # SWPD
    for i in 1:n₁
        d = getdepth(i,:quad)
        j₁ = getchildindex(i,:topleft)
        j₂ = getchildindex(i,:topright)
        j₃ = getchildindex(i,:bottomleft)
        j₄ = getchildindex(i,:bottomright)
        v = @view xw[:,:,i]
        w₁ = @view xw[:,:,j₁]
        w₂ = @view xw[:,:,j₂]
        w₃ = @view xw[:,:,j₃]
        w₄ = @view xw[:,:,j₄]
        sdwt_step!(w₁, w₂, w₃, w₄, v, d, h, g, temp)
    end
end

# ----- ISWPD (Shift based by level) with allocation -----
"""
    iswpd(xw, wt, L, sm)
    iswpd(xw, wt[, L])
    iswpd(xw, wt, tree[, sm])

Computes the inverse stationary wavelet packet transform (iSWPT) on `xw`.

!!! note
    This function might not be very useful if one is looking to reconstruct a raw
    decomposed signal. The purpose of this function would be better utilized in applications
    such as denoising, where a signal is decomposed (`swpd`) and thresholded
    (`denoise`/`denoiseall`) before being reconstructed.

# Arguments
- `xw::AbstractArray{T,2}` or `x::AbstractArray{T,3} where T<:Number`: SWPD-transformed
  array.
- `wt::OrthoFilter`: Orthogonal wavelet filter.
- `L::Integer`: (Default: `maxtransformlevels(x)` or `minimum(size(xw)[1:end-1]) |>
  maxtransformlevels`) Number of levels of decomposition used for reconstruction.
- `tree::BitVector`: Binary/Quad tree for inverse transform to be computed accordingly.
- `sm::Integer`: If `sm` is included as an argument, the `sm`-shifted inverse transform will
  be computed. This results in significantly faster computation, but fails to fully utilize
  the strength of redundant wavelet transforms.

# Returns
`::Vector{T}` or `::Matrix{T}`: Inverse transformed signal.

# Examples
```julia
using Wavelets, WaveletsExt

# Setup
x = generatesignals(:heavysine)
wt = wavelet(WT.haar)

# SWPD
xw = swpt(x, wt)

# Shift-based iSWPD
x̂ = iswpd(xw, wt, maxtransformlevels(xw,1), 5)

# Average-based iSWPD
x̃ = iswpd(xw, wt)
```

**See also:** [`isdwt_step`](@ref), [`iswpt`](@ref), [`swpd`](@ref), [`iswpd!`](@ref)
"""
function iswpd(xw::AbstractArray{T}, wt::OrthoFilter, L::Integer, sm::Integer) where
               T<:Number
    @assert 2 ≤ ndims(xw) ≤ 3
    sz = size(xw)[1:(end-1)]
    return iswpd(xw, wt, maketree(sz..., L, :full), sm)
end
# ----- ISWPD (Average based by level) with allocation -----
function iswpd(xw::AbstractArray{T},
               wt::OrthoFilter,
               L::Integer = minimum(size(xw)[1:end-1]) |> maxtransformlevels) where
               T<:Number
    @assert 2 ≤ ndims(xw) ≤ 3
    sz = size(xw)[1:(end-1)]
    return iswpd(xw, wt, maketree(sz..., L, :full))
end
# ----- ISWPD (Shift based by tree) with allocation -----
function iswpd(xw::AbstractArray{T},
               wt::OrthoFilter,
               tree::BitVector,
               sm::Integer) where T<:Number
    @assert 2 ≤ ndims(xw) ≤ 3
    sz = size(xw)[1:(end-1)]
    x = Array{T}(undef, sz)
    iswpd!(x, xw, wt, tree, sm)
    return x
end
# ----- ISWPD (Average based by tree) with allocation -----
function iswpd(xw::AbstractArray{T},
               wt::OrthoFilter,
               tree::BitVector) where T<:Number
    @assert 2 ≤ ndims(xw) ≤ 3
    sz = size(xw)[1:(end-1)]
    x = Array{T}(undef, sz)
    iswpd!(x, xw, wt, tree)
    return x
end

# ----- ISWPD (Shift based by level) without allocation -----
"""
    iswpd!(x, xw, wt, L, sm)
    iswpd!(x, xw, wt[, L])
    iswpd!(x, xw, wt, tree, sm)
    iswpd!(x, xw, wt, tree)

Same as `iswpd` but with no array allocation.

# Arguments
- `x::AbstractVector{T}` or `x::AbstractMatrix{T} where T<:Number`: Allocated array for
  output.
- `xw::AbstractArray{T,2}` or `xw::AbstractArray{T,3} where T<:Number`: SWPD-transformed
  array.
- `wt::OrthoFilter`: Orthogonal wavelet filter.
- `L::Integer`: (Default: `maxtransformlevels(x)` or `minimum(size(xw)[1:end-1]) |>
  maxtransformlevels`) Number of levels of decomposition used for reconstruction.
- `tree::BitVector`: Binary/Quad tree for inverse transform to be computed accordingly.
- `sm::Integer`: If `sm` is included as an argument, the `sm`-shifted inverse transform will
  be computed. This results in significantly faster computation, but fails to fully utilize
  the strength of redundant wavelet transforms.

# Returns
`x::Vector{T}` or `x::Matrix{T}`: Inverse transformed signal.

# Examples
```julia
using Wavelets, WaveletsExt

# Setup
x = generatesignals(:heavysine)
wt = wavelet(WT.haar)

# SWPD
xw = swpd(x, wt)

# ISWPD
x̂ = similar(x)
iswpd!(x̂, xw, wt, 4, 5)
iswpd!(x̂, xw, wt, maketree(x), 5)
iswpd!(x̂, xw, wt, 4)
iswpd!(x̂, xw, wt, maketree(x))
```

**See also:** [`iswpd`](@ref)
"""
function iswpd!(x::AbstractArray{T},
                xw::AbstractArray{T},
                wt::OrthoFilter,
                L::Integer,
                sm::Integer) where T<:Number
    # Sanity check
    @assert ndims(x) == ndims(xw)-1
    @assert size(x) == size(xw)[1:(end-1)]
    @assert L ≤ maxtransformlevels(x) ||
            throw(ArgumentError("Too many transform levels."))
    @assert L ≥ 1 || throw(ArgumentError("L must be ≥ 1."))
    return iswpd!(x, xw, wt, maketree(size(x)..., L, :full), sm)
end
# ----- ISWPD (Average based by level) without allocation -----
function iswpd!(x::AbstractArray{T},
                xw::AbstractArray{T},
                wt::OrthoFilter,
                L::Integer = minimum(size(xw)[1:end-1]) |> maxtransformlevels) where
                T<:Number
    # Sanity check
    @assert ndims(x) == ndims(xw)-1
    @assert size(x) == size(xw)[1:(end-1)]
    @assert L ≤ maxtransformlevels(x) ||
            throw(ArgumentError("Too many transform levels."))
    @assert L ≥ 1 || throw(ArgumentError("L must be ≥ 1"))
    return iswpd!(x, xw, wt, maketree(size(xw,1), L, :full))
end
# ----- 1D ISWPD (Shift based by tree) without allocation -----
function iswpd!(x::AbstractVector{T},
                xw::AbstractArray{T,2},
                wt::OrthoFilter,
                tree::BitVector,
                sm::Integer) where T<:Number
    # Sanity check
    @assert isvalidtree(xw[:,1], tree)

    # Setup
    _, m = size(xw)
    g, h = WT.makereverseqmfpair(wt, true)
    tmp = copy(xw)                          # Temp. array to store intermediate outputs
    L = getdepth(m,:binary)                 # Number of decompositions from xw
    sd = Utils.main2depthshift(sm, L)       # Shifts at each depth

    # iSWPD
    for (i, haschild) in Iterators.reverse(enumerate(tree))
        # Current node has child => Compute one step of inverse transform
        if haschild
            d = getdepth(i,:binary)                     # Parent depth
            sv = sd[d+1]                                # Parent shift
            sw = sd[d+2]                                # Child shift
            j₁ = getchildindex(i,:left)                 # Scaling child index
            j₂ = getchildindex(i,:right)                # Detail child index
            @inbounds v = i==1 ? x : @view tmp[:,i]     # Parent node
            @inbounds w₁ = @view tmp[:,j₁]              # Scaling child node
            @inbounds w₂ = @view tmp[:,j₂]              # Detail child node
            # Inverse transform
            @inbounds isdwt_step!(v, w₁, w₂, d, sv, sw, h, g)
        end
    end
    return x
end
# ----- 2D ISWPD (Shift based by tree) without allocation -----
function iswpd!(x::AbstractMatrix{T},
                xw::AbstractArray{T,3},
                wt::OrthoFilter,
                tree::BitVector,
                sm::Integer) where T<:Number
    # Sanity check
    @assert size(x) == size(xw)[1:(end-1)]
    @assert isvalidtree(x, tree)

    # Setup
    n, m, k = size(xw)
    g, h = WT.makereverseqmfpair(wt, true)
    xwₜ = copy(xw)                           # Temp. array to store intermediate outputs
    temp = Array{T,3}(undef, (n,m,2))
    L = getdepth(k,:quad)                   # Number of decompositions from xw
    sd = Utils.main2depthshift(sm, L)       # Shifts at each depth

    # ISWPD
    for (i, haschild) in Iterators.reverse(enumerate(tree))
        # Current node has child => Compute one step of inverse transform
        if haschild
            d = getdepth(i,:quad)                       # Parent depth
            sv = sd[d+1]                                # Parent shift
            sw = sd[d+2]                                # Child shift
            j₁ = getchildindex(i,:topleft)              # Scaling child index
            j₂ = getchildindex(i,:topright)             # Detail child index
            j₃ = getchildindex(i,:bottomleft)           # Detail child index
            j₄ = getchildindex(i,:bottomright)          # Detail child index
            @inbounds v = i==1 ? x : @view xwₜ[:,:,i]    # Parent node
            @inbounds w₁ = @view xwₜ[:,:,j₁]             # Scaling child node
            @inbounds w₂ = @view xwₜ[:,:,j₂]             # Detail child node
            @inbounds w₃ = @view xwₜ[:,:,j₃]             # Detail child node
            @inbounds w₄ = @view xwₜ[:,:,j₄]             # Detail child node
            # Inverse transform
            @inbounds isdwt_step!(v, w₁, w₂, w₃, w₄, d, sv, sw, h, g, temp)
        end
    end
    return x
end

# ----- 1D ISWPD (Average based by tree) without allocation -----
function iswpd!(x::AbstractVector{T},
                xw::AbstractArray{T,2},
                wt::OrthoFilter,
                tree::BitVector) where T<:Number
    # Sanity check
    @assert isvalidtree(xw[:,1], tree)

    # Setup
    g, h = WT.makereverseqmfpair(wt, true)
    tmp = copy(xw)

    # iSWPD
    for (i, haschild) in Iterators.reverse(enumerate(tree))
        # Current node has child => Compute one step of inverse transform
        if haschild
            d = getdepth(i,:binary)                     # Parent depth
            j₁ = getchildindex(i,:left)                 # Scaling child index
            j₂ = getchildindex(i,:right)                # Detail child index
            @inbounds v = i==1 ? x : @view tmp[:,i]     # Parent node
            @inbounds w₁ = @view tmp[:,j₁]              # Scaling child node
            @inbounds w₂ = @view tmp[:,j₂]              # Detail child node
            # Inverse transform
            @inbounds isdwt_step!(v, w₁, w₂, d, h, g)
        end
    end
    return x
end
# ----- 2D ISWPD (Average based by tree) without allocation -----
function iswpd!(x::AbstractMatrix{T},
                xw::AbstractArray{T,3},
                wt::OrthoFilter,
                tree::BitVector) where T<:Number
    # Sanity check
    @assert size(x) == size(xw)[1:(end-1)]
    @assert isvalidtree(x, tree)

    # Setup
    n, m, k = size(xw)
    g, h = WT.makereverseqmfpair(wt, true)
    xwₜ = copy(xw)                           # Temp. array to store intermediate outputs
    temp = Array{T,3}(undef, (n,m,2))

    # ISWPD
    for (i, haschild) in Iterators.reverse(enumerate(tree))
        # Current node has child => Compute one step of inverse transform
        if haschild
            d = getdepth(i,:quad)                       # Parent depth
            j₁ = getchildindex(i,:topleft)              # Scaling child index
            j₂ = getchildindex(i,:topright)             # Detail child index
            j₃ = getchildindex(i,:bottomleft)           # Detail child index
            j₄ = getchildindex(i,:bottomright)          # Detail child index
            @inbounds v = i==1 ? x : @view xwₜ[:,:,i]    # Parent node
            @inbounds w₁ = @view xwₜ[:,:,j₁]             # Scaling child node
            @inbounds w₂ = @view xwₜ[:,:,j₂]             # Detail child node
            @inbounds w₃ = @view xwₜ[:,:,j₃]             # Detail child node
            @inbounds w₄ = @view xwₜ[:,:,j₄]             # Detail child node
            # Inverse transform
            @inbounds isdwt_step!(v, w₁, w₂, w₃, w₄, d, h, g, temp)
        end
    end
    return x
end

include("swt_one_level.jl")
include("swt_all.jl")

end # end module


text:
export rule

@rule Categorical(:out, Marginalisation) (q_p::Dirichlet, ) = begin
    rho = clamp.(exp.(logmean(q_p)), tiny, Inf) # Softens the parameter
    return Categorical(rho ./ sum(rho))
end

@rule Categorical(:out, Marginalisation) (m_p::PointMass, ) = begin
    return Categorical(mean(m_p))
end

@rule Categorical(:out, Marginalisation) (q_p::PointMass, ) = begin
    return Categorical(mean(q_p))
end