# preliminary definitions: constants, macros
# and functions used throughout the code
const _msk64 = ~uint64(0)
macro _mskr(l) :(_msk64 >>> (63 & (64-$(esc(l))))) end
macro _div64(l) :($(esc(l)) >>> 6) end
macro _mod64(l) :($(esc(l)) & 63) end
macro _msk_end(l) :(@_mskr @_mod64 $(esc(l))) end
num_bit_chunks(n::Int) = @_div64 (n+63)

## BitArray

# notes: bits are stored in contiguous chunks
#        unused bits must always be set to 0
type BitArray{N} <: DenseArray{Bool, N}
    chunks::Vector{Uint64}
    len::Int
    dims::NTuple{N,Int}
    function BitArray(dims::Int...)
        if length(dims) != N
            error("number of dimensions must be $N (got $(length(dims)))")
        end
        n = 1
        for d in dims
            if d < 0
                error("dimension size must be nonnegative (got $d)")
            end
            n *= d
        end
        nc = num_bit_chunks(n)
        chunks = Array(Uint64, nc)
        if nc > 0
            chunks[end] = uint64(0)
        end
        b = new(chunks, n)
        if N != 1
            b.dims = dims
        end
        return b
    end
end

BitArray{N}(dims::NTuple{N,Int}) = BitArray{N}(dims...)
BitArray(dims::Int...) = BitArray(dims)

typealias BitVector BitArray{1}
typealias BitMatrix BitArray{2}

## utility functions ##

length(B::BitArray) = B.len
size(B::BitVector) = (B.len,)
size(B::BitArray) = B.dims

size(B::BitVector, d) = (d==1 ? B.len : d>1 ? 1 : error("dimensions should be positive (got $d)"))
size{N}(B::BitArray{N}, d) = (d>N ? 1 : B.dims[d])

isassigned{N}(B::BitArray{N}, i::Int) = 1 <= i <= length(B)

## Aux functions ##

get_chunks_id(i::Integer) = @_div64(int(i)-1)+1, @_mod64(int(i)-1)

function glue_src_bitchunks(src::Vector{Uint64}, k::Int, ks1::Int, msk_s0::Uint64, ls0::Int)
    @inbounds begin
        chunk = ((src[k] & msk_s0) >>> ls0)
        if ks1 > k && ls0 > 0
            chunk_n = (src[k + 1] & ~msk_s0)
            chunk |= (chunk_n << (64 - ls0))
        end
    end
    return chunk
end

function copy_chunks!(dest::Vector{Uint64}, pos_d::Integer, src::Vector{Uint64}, pos_s::Integer, numbits::Integer)
    numbits == 0 && return
    if dest === src && pos_d > pos_s
        return copy_chunks_rtol!(dest, pos_d, pos_s, numbits)
    end

    kd0, ld0 = get_chunks_id(pos_d)
    kd1, ld1 = get_chunks_id(pos_d + numbits - 1)
    ks0, ls0 = get_chunks_id(pos_s)
    ks1, ls1 = get_chunks_id(pos_s + numbits - 1)

    delta_kd = kd1 - kd0
    delta_ks = ks1 - ks0

    u = _msk64
    if delta_kd == 0
        msk_d0 = ~(u << ld0) | (u << (ld1+1))
    else
        msk_d0 = ~(u << ld0)
        msk_d1 = (u << (ld1+1))
    end
    if delta_ks == 0
        msk_s0 = (u << ls0) & ~(u << (ls1+1))
    else
        msk_s0 = (u << ls0)
    end

    chunk_s0 = glue_src_bitchunks(src, ks0, ks1, msk_s0, ls0)

    dest[kd0] = (dest[kd0] & msk_d0) | ((chunk_s0 << ld0) & ~msk_d0)

    if delta_kd == 0
        return
    end

    for i = 1 : kd1 - kd0 - 1
        chunk_s1 = glue_src_bitchunks(src, ks0 + i, ks1, msk_s0, ls0)

        chunk_s = (chunk_s0 >>> (64 - ld0)) | (chunk_s1 << ld0)

        dest[kd0 + i] = chunk_s

        chunk_s0 = chunk_s1
    end

    if ks1 >= ks0 + delta_kd
        chunk_s1 = glue_src_bitchunks(src, ks0 + delta_kd, ks1, msk_s0, ls0)
    else
        chunk_s1 = uint64(0)
    end

    chunk_s = (chunk_s0 >>> (64 - ld0)) | (chunk_s1 << ld0)

    dest[kd1] = (dest[kd1] & msk_d1) | (chunk_s & ~msk_d1)

    return
end

function copy_chunks_rtol!(chunks::Vector{Uint64}, pos_d::Integer, pos_s::Integer, numbits::Integer)
    if pos_d == pos_s
        return
    elseif pos_d < pos_s
        return copy_chunks!(chunks, pos_d, chunks, pos_s, numbits)
    end

    left = numbits
    s = min(left, 64)
    b = left - s
    ps = pos_s + b
    pd = pos_d + b
    u = _msk64
    while left > 0
        kd0, ld0 = get_chunks_id(pd)
        kd1, ld1 = get_chunks_id(pd + s - 1)
        ks0, ls0 = get_chunks_id(ps)
        ks1, ls1 = get_chunks_id(ps + s - 1)

        delta_kd = kd1 - kd0
        delta_ks = ks1 - ks0

        if delta_kd == 0
            msk_d0 = ~(u << ld0) | (u << (ld1+1))
        else
            msk_d0 = ~(u << ld0)
            msk_d1 = (u << (ld1+1))
        end
        if delta_ks == 0
            msk_s0 = (u << ls0) & ~(u << (ls1+1))
        else
            msk_s0 = (u << ls0)
        end

        chunk_s0 = glue_src_bitchunks(chunks, ks0, ks1, msk_s0, ls0) & ~(u << s)
        chunks[kd0] = (chunks[kd0] & msk_d0) | ((chunk_s0 << ld0) & ~msk_d0)

        if delta_kd != 0
            chunk_s = (chunk_s0 >>> (64 - ld0))

            chunks[kd1] = (chunks[kd1] & msk_d1) | (chunk_s & ~msk_d1)
        end

        left -= s
        s = min(left, 64)
        b = left - s
        ps = pos_s + b
        pd = pos_d + b
    end
end

function fill_chunks!(Bc::Array{Uint64}, x::Bool, pos::Integer, numbits::Integer)
    numbits <= 0 && return
    k0, l0 = get_chunks_id(pos)
    k1, l1 = get_chunks_id(pos+numbits-1)

    u = _msk64
    if k1 == k0
        msk0 = (u << l0) & ~(u << (l1+1))
    else
        msk0 = (u << l0)
        msk1 = ~(u << (l1+1))
    end
    @inbounds if x
        Bc[k0] |= msk0
        for k = k0+1:k1-1
            Bc[k] = u
        end
        k1 > k0 && (Bc[k1] |= msk1)
    else
        Bc[k0] &= ~msk0
        for k = k0+1:k1-1
            Bc[k] = 0
        end
        k1 > k0 && (Bc[k1] &= ~msk1)
    end
end

## custom iterator ##
start(B::BitArray) = 0
next(B::BitArray, i::Int) = (B.chunks[@_div64(i)+1] & (uint64(1)<<@_mod64(i)) != 0, i+1)
done(B::BitArray, i::Int) = i >= length(B)

## similar, fill!, copy! etc ##

similar(B::BitArray) = BitArray(size(B))
similar(B::BitArray, dims::Int...) = BitArray(dims)
similar(B::BitArray, dims::Dims) = BitArray(dims...)

similar(B::BitArray, T::Type{Bool}, dims::Dims) = BitArray(dims)
# changing type to a non-Bool returns an Array
# (this triggers conversions like float(bitvector) etc.)
similar(B::BitArray, T::Type, dims::Dims) = Array(T, dims)

function fill!(B::BitArray, x)
    y = convert(Bool, x)
    if length(B) == 0
        return B
    end
    Bc = B.chunks
    if !y
        fill!(Bc, 0)
    else
        fill!(Bc, _msk64)
        Bc[end] &= @_msk_end length(B)
    end
    return B
end

falses(args...) = fill!(BitArray(args...), false)
trues(args...) = fill!(BitArray(args...), true)

function one(x::BitMatrix)
    m,n = size(x)
    m == n || throw(DimensionMismatch("multiplicative identity defined only for square matrices"))
    a = falses(n, n)
    for i = 1:n
        a[i,i] = true
    end
    return a
end

function copy!(dest::BitArray, src::BitArray)
    destc = dest.chunks; srcc = src.chunks
    nc_d = length(destc)
    nc_s = length(srcc)
    nc = min(nc_s, nc_d)
    if nc == 0
        return dest
    end
    @inbounds begin
        for i = 1 : nc - 1
            destc[i] = srcc[i]
        end
        if length(src) >= length(dest)
            destc[nc] = srcc[nc]
        else
            msk_s = @_msk_end length(src)
            msk_d = ~msk_s
            destc[nc] = (msk_d & destc[nc]) | (msk_s & srcc[nc])
        end
    end
    return dest
end

function copy!(dest::BitArray, pos_d::Integer, src::BitArray, pos_s::Integer, numbits::Integer)
    if pos_s+numbits-1 > length(src) || pos_d+numbits-1 > length(dest) || pos_d < 1 || pos_s < 1
        throw(BoundsError())
    end
    copy_chunks!(dest.chunks, pos_d, src.chunks, pos_s, numbits)
    return dest
end

function reshape{N}(B::BitArray, dims::NTuple{N,Int})
    if prod(dims) != length(B)
        throw(DimensionMismatch("new dimensions $(dims) must be consistent with array size $(length(B))"))
    end
    if dims == size(B)
        return B
    end
    Br = BitArray{N}(ntuple(N,i->0)...)
    Br.chunks = B.chunks
    Br.len = prod(dims)
    if N != 1
        Br.dims = dims
    end
    return Br
end

## Conversions ##

convert{T,N}(::Type{Array{T}}, B::BitArray{N}) = convert(Array{T,N},B)
function convert{T,N}(::Type{Array{T,N}}, B::BitArray{N})
    A = Array(T, size(B))
    Bc = B.chunks
    for i = 1:length(A)
        A[i] = unsafe_bitgetindex(Bc, i)
    end
    return A
end

convert{T,N}(::Type{BitArray}, A::AbstractArray{T,N}) = convert(BitArray{N},A)
function convert{T,N}(::Type{BitArray{N}}, A::AbstractArray{T,N})
    B = BitArray(size(A))
    Bc = B.chunks
    l = length(B)
    if l == 0
        return B
    end
    ind = 1
    @inbounds begin
        for i = 1:length(Bc)-1
            u = uint64(1)
            c = uint64(0)
            for j = 0:63
                if bool(A[ind])
                    c |= u
                end
                ind += 1
                u <<= 1
            end
            Bc[i] = c
        end
        u = uint64(1)
        c = uint64(0)
        for j = 0:@_mod64(l-1)
            if bool(A[ind])
                c |= u
            end
            ind += 1
            u <<= 1
        end
        Bc[end] = c
    end
    return B
end

convert{N}(::Type{BitArray{N}}, B::BitArray{N}) = B
convert{T,N}(::Type{AbstractArray{T,N}}, B::BitArray{N}) = convert(Array{T,N}, B)

reinterpret{N}(::Type{Bool}, B::BitArray, dims::NTuple{N,Int}) = reinterpret(B, dims)
reinterpret{N}(B::BitArray, dims::NTuple{N,Int}) = reshape(B, dims)

# shorthand forms BitArray <-> Array
bitunpack{N}(B::BitArray{N}) = convert(Array{Bool,N}, B)
bitpack{T,N}(A::AbstractArray{T,N}) = convert(BitArray{N}, A)

## Random ##

function bitarray_rand_fill!(B::BitArray)
    if length(B) == 0
        return B
    end
    Bc = B.chunks
    for i = 1 : length(Bc) - 1
        Bc[i] = rand(Uint64)
    end
    msk = @_msk_end length(B)
    Bc[end] = msk & rand(Uint64)
    return B
end

## Indexing: getindex ##

function unsafe_bitgetindex(Bc::Vector{Uint64}, i::Int)
    return (Bc[@_div64(i-1)+1] & (uint64(1)<<@_mod64(i-1))) != 0
end

function getindex(B::BitArray, i::Int)
    1 <= i <= length(B) || throw(BoundsError())
    return unsafe_bitgetindex(B.chunks, i)
end

getindex(B::BitArray, i::Real) = getindex(B, to_index(i))

getindex(B::BitArray) = getindex(B, 1)

# 0d bitarray
getindex(B::BitArray{0}) = unsafe_bitgetindex(B.chunks, 1)

function getindex{T<:Real}(B::BitArray, I::AbstractVector{T})
    X = BitArray(length(I))
    lB = length(B)
    Xc = X.chunks
    Bc = B.chunks
    ind = 1
    for i in I
        # faster X[ind] = B[i]
        j = to_index(i)
        1 <= j <= lB || throw(BoundsError())
        unsafe_bitsetindex!(Xc, unsafe_bitgetindex(Bc, j), ind)
        ind += 1
    end
    return X
end

# logical indexing

# (multiple signatures for disambiguation)
for IT in [AbstractVector{Bool}, AbstractArray{Bool}]
    @eval function getindex(B::BitArray, I::$IT)
        checkbounds(B, I)
        n = sum(I)
        X = BitArray(n)
        Xc = X.chunks
        Bc = B.chunks
        ind = 1
        for i = 1:length(I)
            if I[i]
                # faster X[ind] = B[i]
                unsafe_bitsetindex!(Xc, unsafe_bitgetindex(Bc, i), ind)
                ind += 1
            end
        end
        return X
    end
end

## Indexing: setindex! ##

function unsafe_bitsetindex!(Bc::Array{Uint64}, x::Bool, i::Int)
    i1, i2 = get_chunks_id(i)
    u = uint64(1) << i2
    @inbounds begin
        if x
            Bc[i1] |= u
        else
            Bc[i1] &= ~u
        end
    end
end

setindex!(B::BitArray, x) = setindex!(B, convert(Bool,x), 1)

function setindex!(B::BitArray, x::Bool, i::Int)
    1 <= i <= length(B) || throw(BoundsError())
    unsafe_bitsetindex!(B.chunks, x, i)
    return B
end

# logical indexing

function setindex!(B::BitArray, x, I::AbstractArray{Bool})
    checkbounds(B, I)
    y = convert(Bool, x)
    Bc = B.chunks
    @inbounds for i = 1:length(I)
        # faster B[i] = y
        I[i] && unsafe_bitsetindex!(Bc, y, i)
    end
    return B
end

function setindex!(B::BitArray, X::AbstractArray, I::AbstractArray{Bool})
    checkbounds(B, I)
    Bc = B.chunks
    c = 1
    @inbounds for i = 1:length(I)
        if I[i]
            # faster B[i] = X[c]
            unsafe_bitsetindex!(Bc, convert(Bool, X[c]), i)
            c += 1
        end
    end
    if length(X) != c-1
        throw(DimensionMismatch("assigned $(length(X)) elements to length $(c-1) destination"))
    end
    return B
end

## Dequeue functionality ##

function push!(B::BitVector, item)
    # convert first so we don't grow the bitarray if the assignment won't work
    item = convert(Bool, item)

    Bc = B.chunks

    l = @_mod64 length(B)
    if l == 0
        ccall(:jl_array_grow_end, Void, (Any, Uint), Bc, 1)
        Bc[end] = uint64(0)
    end
    B.len += 1
    if item
        B[end] = true
    end
    return B
end

function append!(B::BitVector, items::BitVector)
    n0 = length(B)
    n1 = length(items)
    if n1 == 0
        return B
    end
    Bc = B.chunks
    k0 = length(Bc)
    k1 = num_bit_chunks(n0 + n1)
    if k1 > k0
        ccall(:jl_array_grow_end, Void, (Any, Uint), Bc, k1 - k0)
        Bc[end] = uint64(0)
    end
    B.len += n1
    copy_chunks!(Bc, n0+1, items.chunks, 1, n1)
    return B
end

append!(B::BitVector, items::AbstractVector{Bool}) = append!(B, bitpack(items))
append!(A::Vector{Bool}, items::BitVector) = append!(A, bitunpack(items))

function prepend!(B::BitVector, items::BitVector)
    n0 = length(B)
    n1 = length(items)
    if n1 == 0
        return B
    end
    Bc = B.chunks
    k0 = length(Bc)
    k1 = num_bit_chunks(n0 + n1)
    if k1 > k0
        ccall(:jl_array_grow_end, Void, (Any, Uint), Bc, k1 - k0)
        Bc[end] = uint64(0)
    end
    B.len += n1
    copy_chunks!(Bc, 1 + n1, Bc, 1, n0)
    copy_chunks!(Bc, 1, items.chunks, 1, n1)
    return B
end

prepend!(B::BitVector, items::AbstractVector{Bool}) = prepend!(B, bitpack(items))
prepend!(A::Vector{Bool}, items::BitVector) = prepend!(A, bitunpack(items))

function sizehint(B::BitVector, sz::Integer)
    ccall(:jl_array_sizehint, Void, (Any, Uint), B.chunks, num_bit_chunks(sz))
    return B
end

function resize!(B::BitVector, n::Integer)
    if n < 0
        throw(BoundsError())
    end
    n0 = length(B)
    if n <= n0
        splice!(B, n+1:n0)
        return B
    end
    Bc = B.chunks
    k0 = length(Bc)
    k1 = num_bit_chunks(int(n))
    if k1 > k0
        ccall(:jl_array_grow_end, Void, (Any, Uint), Bc, k1 - k0)
        Bc[end] = uint64(0)
    end
    B.len = n
    return B
end

function pop!(B::BitVector)
    if isempty(B)
        error("argument must not be empty")
    end
    item = B[end]
    B[end] = false

    l = @_mod64 length(B)
    if l == 1
        ccall(:jl_array_del_end, Void, (Any, Uint), B.chunks, 1)
    end
    B.len -= 1

    return item
end

function unshift!(B::BitVector, item)
    item = convert(Bool, item)

    Bc = B.chunks

    l = @_mod64 length(B)
    if l == 0
        ccall(:jl_array_grow_end, Void, (Any, Uint), Bc, 1)
        Bc[end] = uint64(0)
    end
    B.len += 1
    if B.len == 1
        Bc[1] = item
        return B
    end
    for i = length(Bc) : -1 : 2
        Bc[i] = (Bc[i] << 1) | (Bc[i-1] >>> 63)
    end
    Bc[1] = uint64(item) | (Bc[1] << 1)
    return B
end

function shift!(B::BitVector)
    if isempty(B)
        error("argument must not be empty")
    end
    @inbounds begin
        item = B[1]

        Bc = B.chunks

        for i = 1 : length(Bc) - 1
            Bc[i] = (Bc[i] >>> 1) | (Bc[i+1] << 63)
        end

        l = @_mod64 length(B)
        if l == 1
            ccall(:jl_array_del_end, Void, (Any, Uint), Bc, 1)
        else
            Bc[end] >>>= 1
        end
        B.len -= 1
    end

    return item
end

function insert!(B::BitVector, i::Integer, item)
    if i < 1
        throw(BoundsError())
    end
    item = convert(Bool, item)

    n = length(B)
    if i > n
        x = falses(i - n)
        append!(B, x)
    else
        Bc = B.chunks

        k, j = get_chunks_id(i)

        l = @_mod64 length(B)
        if l == 0
            ccall(:jl_array_grow_end, Void, (Any, Uint), Bc, 1)
            Bc[end] = uint64(0)
        end
        B.len += 1

        for t = length(Bc) : -1 : k + 1
            Bc[t] = (Bc[t] << 1) | (Bc[t - 1] >>> 63)
        end

        msk_aft = (_msk64 << j)
        msk_bef = ~msk_aft
        Bc[k] = (msk_bef & Bc[k]) | ((msk_aft & Bc[k]) << 1)
    end
    B[i] = item
end

function _deleteat!(B::BitVector, i::Integer)

    k, j = get_chunks_id(i)

    msk_bef = _msk64 >>> (63 - j)
    msk_aft = ~msk_bef
    msk_bef >>>= 1

    Bc = B.chunks

    @inbounds begin
        Bc[k] = (msk_bef & Bc[k]) | ((msk_aft & Bc[k]) >> 1)
        if length(Bc) > k
            Bc[k] |= (Bc[k + 1] << 63)
        end

        for t = k + 1 : length(Bc) - 1
            Bc[t] = (Bc[t] >>> 1) | (Bc[t + 1] << 63)
        end

        l = @_mod64 length(B)

        if l == 1
            ccall(:jl_array_del_end, Void, (Any, Uint), Bc, 1)
        elseif length(Bc) > k
            Bc[end] >>>= 1
        end
    end

    B.len -= 1

    return B
end

function deleteat!(B::BitVector, i::Integer)
    n = length(B)
    if !(1 <= i <= n)
        throw(BoundsError())
    end

    return _deleteat!(B, i)
end

function deleteat!(B::BitVector, r::UnitRange{Int})
    n = length(B)
    i_f = first(r)
    i_l = last(r)
    if !(1 <= i_f && i_l <= n)
        throw(BoundsError())
    end

    Bc = B.chunks
    new_l = length(B) - length(r)
    delta_k = num_bit_chunks(new_l) - length(Bc)

    copy_chunks!(Bc, i_f, Bc, i_l+1, n-i_l)

    if delta_k < 0
        ccall(:jl_array_del_end, Void, (Any, Uint), Bc, -delta_k)
    end

    B.len = new_l

    if new_l > 0
        Bc[end] &= @_msk_end new_l
    end

    return B
end

function deleteat!(B::BitVector, inds)
    n = new_l = length(B)
    s = start(inds)
    done(inds, s) && return B

    Bc = B.chunks

    (p, s) = next(inds, s)
    q = p+1
    new_l -= 1
    while !done(inds, s)
        (i,s) = next(inds, s)
        if !(q <= i <= n)
            i < q && error("indices must be unique and sorted")
            throw(BoundsError())
        end
        new_l -= 1
        if i > q
            copy_chunks!(Bc, p, Bc, q, i-q)
            p += i-q
        end
        q = i+1
    end

    q <= n && copy_chunks!(Bc, p, Bc, q, n-q+1)

    delta_k = num_bit_chunks(new_l) - length(Bc)
    delta_k < 0 && ccall(:jl_array_del_end, Void, (Any, Uint), Bc, -delta_k)

    B.len = new_l

    if new_l > 0
        Bc[end] &= @_msk_end new_l
    end

    return B
end

function splice!(B::BitVector, i::Integer)
    n = length(B)
    if !(1 <= i <= n)
        throw(BoundsError())
    end

    v = B[i]   # TODO: change to a copy if/when subscripting becomes an ArrayView
    _deleteat!(B, i)
    return v
end
splice!(B::BitVector, i::Integer, ins::BitVector) = splice!(B, int(i):int(i), ins)
splice!(B::BitVector, i::Integer, ins::AbstractVector{Bool}) = splice!(B, i, bitpack(ins))

const _default_bit_splice = BitVector(0)

function splice!(B::BitVector, r::UnitRange{Int}, ins::BitVector = _default_bit_splice)
    n = length(B)
    i_f = first(r)
    i_l = last(r)
    if !(1 <= i_f <= n+1)
        throw(BoundsError())
    end
    if !(i_l <= n)
        throw(BoundsError())
    end
    if (i_f > n)
        append!(B, ins)
        return BitVector(0)
    end

    v = B[r]  # TODO: change to a copy if/when subscripting becomes an ArrayView

    Bc = B.chunks

    lins = length(ins)
    ldel = length(r)

    new_l = length(B) + lins - ldel
    delta_k = num_bit_chunks(new_l) - length(Bc)

    if delta_k > 0
        ccall(:jl_array_grow_end, Void, (Any, Uint), Bc, delta_k)
    end
    copy_chunks!(Bc, i_f+lins, Bc, i_l+1, n-i_l)
    copy_chunks!(Bc, i_f, ins.chunks, 1, lins)
    if delta_k < 0
        ccall(:jl_array_del_end, Void, (Any, Uint), Bc, -delta_k)
    end

    B.len = new_l

    if new_l > 0
        Bc[end] &= @_msk_end new_l
    end

    return v
end
splice!(B::BitVector, r::UnitRange{Int}, ins::AbstractVector{Bool}) = splice!(B, r, bitpack(ins))

function empty!(B::BitVector)
    ccall(:jl_array_del_end, Void, (Any, Uint), B.chunks, length(B.chunks))
    B.len = 0
    return B
end

## Misc functions

for f in (:iround, :itrunc, :ifloor, :iceil, :abs)
    @eval ($f)(B::BitArray) = copy(B)
end

## Unary operators ##

function (-)(B::BitArray)
    A = zeros(Int, size(B))
    l = length(B)
    if l == 0
        return A
    end
    Bc = B.chunks
    ind = 1
    for i = 1:length(Bc)-1
        u = uint64(1)
        c = Bc[i]
        for j = 1:64
            if c & u != 0
                A[ind] = -1
            end
            ind += 1
            u <<= 1
        end
    end
    u = uint64(1)
    c = Bc[end]
    for j = 0:@_mod64(l-1)
        if c & u != 0
            A[ind] = -1
        end
        ind += 1
        u <<= 1
    end
    return A
end
sign(B::BitArray) = copy(B)

function (~)(B::BitArray)
    C = similar(B)
    Bc = B.chunks
    if !isempty(Bc)
        Cc = C.chunks
        for i = 1:length(Bc)-1
            Cc[i] = ~Bc[i]
        end
        msk = @_msk_end length(B)
        Cc[end] = msk & (~Bc[end])
    end
    return C
end

function flipbits!(B::BitArray)
    Bc = B.chunks
    @inbounds if !isempty(Bc)
        for i = 1:length(B.chunks) - 1
            Bc[i] = ~Bc[i]
        end
        msk = @_msk_end length(B)
        Bc[end] = msk & (~Bc[end])
    end
    return B
end

!(B::BitArray) = ~B

## Binary arithmetic operators ##

for f in (:+, :-)
    @eval function ($f)(A::BitArray, B::BitArray)
        r = Array(Int, promote_shape(size(A), size(B)))
        ai = start(A)
        bi = start(B)
        ri = 1
        while !done(A, ai)
            a, ai = next(A, ai)
            b, bi = next(B, bi)
            @inbounds r[ri] = ($f)(a, b)
            ri += 1
        end
        return r
    end
end
for f in (:.+, :.-),
    (arg1, arg2, T, fargs) in ((:(B::BitArray), :(x::Bool)    , Int                                   , :(b, x)),
                               (:(B::BitArray), :(x::Number)  , :(promote_array_type(typeof(x), Bool)), :(b, x)),
                               (:(x::Bool)    , :(B::BitArray), Int                                   , :(x, b)),
                               (:(x::Number)  , :(B::BitArray), :(promote_array_type(typeof(x), Bool)), :(x, b)))
    @eval function ($f)($arg1, $arg2)
        r = Array($T, size(B))
        bi = start(B)
        ri = 1
        while !done(B, bi)
            b, bi = next(B, bi)
            @inbounds r[ri] = ($f)($fargs...)
            ri += 1
        end
        return r
    end
end

for f in (:/, :\)
    @eval begin
        ($f)(A::BitArray, B::BitArray) = ($f)(bitunpack(A), bitunpack(B))
    end
end
(/)(B::BitArray, x::Number) = (/)(bitunpack(B), x)
(/)(x::Number, B::BitArray) = (/)(x, bitunpack(B))

function div(A::BitArray, B::BitArray)
    shp = promote_shape(size(A), size(B))
    all(B) || throw(DivideError())
    return reshape(copy(A), shp)
end
div(A::BitArray, B::Array{Bool}) = div(A, bitpack(B))
div(A::Array{Bool}, B::BitArray) = div(bitpack(A), B)
function div(B::BitArray, x::Bool)
    return x ? copy(B) : throw(DivideError())
end
function div(x::Bool, B::BitArray)
    all(B) || throw(DivideError())
    return x ? trues(size(B)) : falses(size(B))
end
function div(x::Number, B::BitArray)
    all(B) || throw(DivideError())
    pt = promote_array_type(typeof(x), Bool)
    y = div(x, true)
    reshape(pt[ y for i = 1:length(B) ], size(B))
end

function mod(A::BitArray, B::BitArray)
    shp = promote_shape(size(A), size(B))
    all(B) || throw(DivideError())
    return falses(shp)
end
mod(A::BitArray, B::Array{Bool}) = mod(A, bitpack(B))
mod(A::Array{Bool}, B::BitArray) = mod(bitpack(A), B)
function mod(B::BitArray, x::Bool)
    return x ? falses(size(B)) : throw(DivideError())
end
function mod(x::Bool, B::BitArray)
    all(B) || throw(DivideError())
    return falses(size(B))
end
function mod(x::Number, B::BitArray)
    all(B) || throw(DivideError())
    pt = promote_array_type(typeof(x), Bool)
    y = mod(x, true)
    reshape(pt[ y for i = 1:length(B) ], size(B))
end

for f in (:div, :mod)
    @eval begin
        function ($f)(B::BitArray, x::Number)
            F = Array(promote_array_type(typeof(x), Bool), size(B))
            for i = 1:length(F)
                F[i] = ($f)(B[i], x)
            end
            return F
        end
    end
end

function (&)(B::BitArray, x::Bool)
    x ? copy(B) : falses(size(B))
end
(&)(x::Bool, B::BitArray) = B & x

function (|)(B::BitArray, x::Bool)
    x ? trues(size(B)) : copy(B)
end
(|)(x::Bool, B::BitArray) = B | x

function ($)(B::BitArray, x::Bool)
    x ? ~B : copy(B)
end
($)(x::Bool, B::BitArray) = B $ x

for f in (:&, :|, :$)
    @eval begin
        function ($f)(A::BitArray, B::BitArray)
            F = BitArray(promote_shape(size(A),size(B))...)
            Fc = F.chunks
            Ac = A.chunks
            Bc = B.chunks
            if !isempty(Ac) && !isempty(Bc)
                for i = 1:length(Fc) - 1
                    Fc[i] = ($f)(Ac[i], Bc[i])
                end
                msk = @_msk_end length(F)
                Fc[end] = msk & ($f)(Ac[end], Bc[end])
            end
            return F
        end
        ($f)(A::Array{Bool}, B::BitArray) = ($f)(bitpack(A), B)
        ($f)(B::BitArray, A::Array{Bool}) = ($f)(B, bitpack(A))
        ($f)(x::Number, B::BitArray) = ($f)(x, bitunpack(B))
        ($f)(B::BitArray, x::Number) = ($f)(bitunpack(B), x)
    end
end

function (.^)(B::BitArray, x::Bool)
    x ? copy(B) : trues(size(B))
end
function (.^)(x::Bool, B::BitArray)
    x ? trues(size(B)) : ~B
end
function (.^)(x::Number, B::BitArray)
    z = x ^ false
    u = x ^ true
    reshape([ B[i] ? u : z for i = 1:length(B) ], size(B))
end
function (.^)(B::BitArray, x::Integer)
    if x == 0
        return trues(size(B))
    elseif x < 0
        throw(DomainError())
    else
        return copy(B)
    end
end
function (.^){T<:Number}(B::BitArray, x::T)
    if x == 0
        return ones(typeof(true ^ x), size(B))
    elseif T <: Real && x > 0
        return convert(Array{T}, B)
    else
        z = nothing
        u = nothing
        zerr = nothing
        uerr = nothing
        try
            z = false^x
        catch err
            zerr = err
        end
        try
            u = true^x
        catch err
            uerr = err
        end
        if zerr == nothing && uerr == nothing
            t = promote_type(typeof(z), typeof(u))
        elseif zerr == nothing
            t = typeof(z)
        else
            t = typeof(u)
        end
        F = Array(t, size(B))
        for i = 1:length(B)
            if B[i]
                if uerr == nothing
                    F[i] = u
                else
                    throw(uerr)
                end
            else
                if zerr == nothing
                    F[i] = z
                else
                    throw(zerr)
                end
            end
        end
        return F
    end
end

(.*)(x::Bool, B::BitArray) = x & B
(.*)(B::BitArray, x::Bool) = B & x
(.*)(x::Number, B::BitArray) = x .* bitunpack(B)
(.*)(B::BitArray, x::Number) = bitunpack(B) .* x

## promotion to complex ##

# TODO?

## comparison operators ##

function (==)(A::BitArray, B::BitArray)
    if size(A) != size(B)
        return false
    end
    return A.chunks == B.chunks
end

function (!=)(A::BitArray, B::BitArray)
    if size(A) != size(B)
        return true
    end
    return A.chunks != B.chunks
end

## Data movement ##

# TODO some of this could be optimized

function slicedim(A::BitArray, d::Integer, i::Integer)
    d_in = size(A)
    leading = d_in[1:(d-1)]
    d_out = tuple(leading..., 1, d_in[(d+1):end]...)

    M = prod(leading)
    N = length(A)
    stride = M * d_in[d]

    B = BitArray(d_out)
    index_offset = 1 + (i-1)*M

    l = 1

    if M==1
        for j=0:stride:(N-stride)
            B[l] = A[j + index_offset]
            l += 1
        end
    else
        for j=0:stride:(N-stride)
            offs = j + index_offset
            for k=0:(M-1)
                B[l] = A[offs + k]
                l += 1
            end
        end
    end
    return B
end

function flipdim(A::BitArray, d::Integer)
    nd = ndims(A)
    sd = d > nd ? 1 : size(A, d)
    if sd == 1
        return copy(A)
    end

    B = similar(A)

    nnd = 0
    for i = 1:nd
        nnd += int(size(A,i)==1 || i==d)
    end
    if nnd==nd
        # flip along the only non-singleton dimension
        for i = 1:sd
            B[i] = A[sd+1-i]
        end
        return B
    end

    d_in = size(A)
    leading = d_in[1:(d-1)]
    M = prod(leading)
    N = length(A)
    stride = M * sd

    if M==1
        for j = 0:stride:(N-stride)
            for i = 1:sd
                ri = sd+1-i
                B[j + ri] = A[j + i]
            end
        end
    else
        for i = 1:sd
            ri = sd+1-i
            for j=0:stride:(N-stride)
                offs = j + 1 + (i-1)*M
                boffs = j + 1 + (ri-1)*M
                copy_chunks!(B.chunks, boffs, A.chunks, offs, M)
            end
        end
    end
    return B
end

function reverse_bits(src::Uint64)
    z    = src
    z    = ((z >>>  1) & 0x5555555555555555) | ((z <<  1) & 0xaaaaaaaaaaaaaaaa)
    z    = ((z >>>  2) & 0x3333333333333333) | ((z <<  2) & 0xcccccccccccccccc)
    z    = ((z >>>  4) & 0x0f0f0f0f0f0f0f0f) | ((z <<  4) & 0xf0f0f0f0f0f0f0f0)
    z    = ((z >>>  8) & 0x00ff00ff00ff00ff) | ((z <<  8) & 0xff00ff00ff00ff00)
    z    = ((z >>> 16) & 0x0000ffff0000ffff) | ((z << 16) & 0xffff0000ffff0000)
    return ((z >>> 32) & 0x00000000ffffffff) | ((z << 32) & 0xffffffff00000000)
end

function reverse!(B::BitVector)
    n = length(B)
    if n == 0
        return B
    end
    pnc = length(B.chunks) & 1
    hnc = (length(B.chunks) >>> 1)

    aux_chunks = Array(Uint64, 1)

    for i = 1 : hnc
        j = ((i - 1) << 6)
        aux_chunks[1] = reverse_bits(B.chunks[i])
        copy_chunks!(B.chunks, j+1, B.chunks, n-63-j, 64)
        B.chunks[i] = reverse_bits(B.chunks[i])
        copy_chunks!(B.chunks, n-63-j, aux_chunks, 1, 64)
    end

    if pnc == 0
        return B
    end

    i = hnc + 1
    j = hnc << 6
    l = (@_mod64 (n+63)) + 1
    msk = @_mskr l

    aux_chunks[1] = reverse_bits(B.chunks[i] & msk)
    aux_chunks[1] >>>= (64 - l)
    copy_chunks!(B.chunks, j+1, aux_chunks, 1, l)

    return B
end

reverse(v::BitVector) = reverse!(copy(v))

function (<<)(B::BitVector, i::Int64)
    n = length(B)
    i == 0 && return copy(B)
    A = falses(n)
    i < n && copy_chunks!(A.chunks, 1, B.chunks, i+1, n-i)
    return A
end
(<<)(B::BitVector, i::Int32) = B << int64(i)
(<<)(B::BitVector, i::Integer) = B << int64(i)

function (>>>)(B::BitVector, i::Int64)
    n = length(B)
    i == 0 && return copy(B)
    A = falses(n)
    i < n && copy_chunks!(A.chunks, i+1, B.chunks, 1, n-i)
    return A
end
(>>>)(B::BitVector, i::Int32) = B >>> int64(i)
(>>>)(B::BitVector, i::Integer) = B >>> int64(i)

(>>)(B::BitVector, i::Int32) = B >>> i
(>>)(B::BitVector, i::Integer) = B >>> i

function rol(B::BitVector, i::Integer)
    n = length(B)
    i %= n
    i == 0 && return copy(B)
    i < 0 && return ror(B, -i)
    A = BitArray(n)
    copy_chunks!(A.chunks, 1, B.chunks, i+1, n-i)
    copy_chunks!(A.chunks, n-i+1, B.chunks, 1, i)
    return A
end

function ror(B::BitVector, i::Integer)
    n = length(B)
    i %= n
    i == 0 && return copy(B)
    i < 0 && return rol(B, -i)
    A = BitArray(n)
    copy_chunks!(A.chunks, i+1, B.chunks, 1, n-i)
    copy_chunks!(A.chunks, 1, B.chunks, n-i+1, i)
    return A
end

#TODO: rol!, ror!

## countnz & find ##

function countnz(B::BitArray)
    n = 0
    Bc = B.chunks
    @inbounds for i = 1:length(Bc)
        n += count_ones(Bc[i])
    end
    return n
end

# returns the index of the next non-zero element, or 0 if all zeros
function findnext(B::BitArray, start::Integer)
    if start <= 0
        throw(BoundsError())
    elseif start > length(B)
        return 0
    end
    Bc = B.chunks

    chunk_start = @_div64(start-1)+1
    within_chunk_start = @_mod64(start-1)
    mask = _msk64 << within_chunk_start

    @inbounds begin
        if Bc[chunk_start] & mask != 0
            return (chunk_start-1) << 6 + trailing_zeros(Bc[chunk_start] & mask) + 1
        end

        for i = chunk_start+1:length(Bc)
            if Bc[i] != 0
                return (i-1) << 6 + trailing_zeros(Bc[i]) + 1
            end
        end
    end
    return 0
end
#findfirst(B::BitArray) = findnext(B, 1)  ## defined in array.jl

# aux function: same as findnext(~B, start), but performed without temporaries
function findnextnot(B::BitArray, start::Integer)
    if start <= 0
        throw(BoundsError())
    elseif start > length(B)
        return 0
    end

    Bc = B.chunks
    l = length(Bc)
    if l == 0
        return 0
    end

    chunk_start = @_div64(start-1)+1
    within_chunk_start = @_mod64(start-1)
    mask = ~(_msk64 << within_chunk_start)

    @inbounds begin
        if Bc[chunk_start] | mask != _msk64
            return (chunk_start-1) << 6 + trailing_ones(Bc[chunk_start] | mask) + 1
        end

        for i = chunk_start+1:l-1
            if Bc[i] != _msk64
                return (i-1) << 6 + trailing_ones(Bc[i]) + 1
            end
        end
        ce = Bc[end]
    end
    if ce != @_msk_end length(B)
        return (l-1) << 6 + trailing_ones(ce) + 1
    end
    return 0
end
findfirstnot(B::BitArray) = findnextnot(B,1)

# returns the index of the first matching element
function findnext(B::BitArray, v, start::Integer)
    if v == false
        return findnextnot(B, start)
    elseif v == true
        return findnext(B, start)
    else
        return 0
    end
end
#findfirst(B::BitArray, v) = findnext(B, 1, v)  ## defined in array.jl

# returns the index of the first element for which the function returns true
function findnext(testf::Function, B::BitArray, start::Integer)
    f0::Bool = testf(false)
    f1::Bool = testf(true)
    if length(B) == 0 || !(f0 || f1)
        return 0
    elseif f0 && f1
        return 1
    elseif !f0 && f1
        return findnext(B, start)
    else
        return findnextnot(B, start)
    end
end
#findfirst(testf::Function, B::BitArray) = findnext(testf, B, 1)  ## defined in array.jl

function find(B::BitArray)
    l = length(B)
    nnzB = countnz(B)
    I = Array(Int, nnzB)
    if nnzB == 0
        return I
    end
    Bc = B.chunks
    Bcount = 1
    Icount = 1
    for i = 1:length(Bc)-1
        u = uint64(1)
        c = Bc[i]
        for j = 1:64
            if c & u != 0
                I[Icount] = Bcount
                Icount += 1
            end
            Bcount += 1
            u <<= 1
        end
    end
    u = uint64(1)
    c = Bc[end]
    for j = 0:@_mod64(l-1)
        if c & u != 0
            I[Icount] = Bcount
            Icount += 1
        end
        Bcount += 1
        u <<= 1
    end
    return I
end

findn(B::BitVector) = find(B)

function findn(B::BitMatrix)
    nnzB = countnz(B)
    I = Array(Int, nnzB)
    J = Array(Int, nnzB)
    count = 1
    for j=1:size(B,2), i=1:size(B,1)
        if B[i,j]
            I[count] = i
            J[count] = j
            count += 1
        end
    end
    return (I, J)
end

function findnz(B::BitMatrix)
    I, J = findn(B)
    return (I, J, trues(length(I)))
end

nonzeros(B::BitArray) = trues(countnz(B))

## Reductions ##

sum(A::BitArray, region) = reducedim(+,A,region,0,Array(Int,reduced_dims(A,region)))

sum(B::BitArray) = countnz(B)

function all(B::BitArray)
    length(B) == 0 && return true
    Bc = B.chunks
    @inbounds begin
        for i = 1:length(Bc)-1
            if Bc[i] != _msk64
                return false
            end
        end
        if Bc[end] != @_msk_end length(B)
            return false
        end
    end
    return true
end

function any(B::BitArray)
    length(B) == 0 && return false
    Bc = B.chunks
    @inbounds begin
        for i = 1:length(Bc)
            if Bc[i] != 0
                return true
            end
        end
    end
    return false
end

minimum(B::BitArray) = isempty(B) ? error("argument must be non-empty") : all(B)
maximum(B::BitArray) = isempty(B) ? error("argument must be non-empty") : any(B)

## map over bitarrays ##

function map!(f::Callable, A::Union(StridedArray,BitArray))
    for i=1:length(A)
        A[i] = f(A[i])
    end
    return A
end

function map!(f::Callable, dest::Union(StridedArray,BitArray), A::Union(StridedArray,BitArray))
    for i=1:length(A)
        dest[i] = f(A[i])
    end
    return dest
end

function map!(f::Callable, dest::Union(StridedArray,BitArray), A::Union(StridedArray,BitArray), B::Union(StridedArray,BitArray))
    for i=1:length(A)
        dest[i] = f(A[i], B[i])
    end
    return dest
end

function map!(f::Callable, dest::Union(StridedArray,BitArray), A::Union(StridedArray,BitArray), B::Number)
    for i=1:length(A)
        dest[i] = f(A[i], B)
    end
    return dest
end

function map!(f::Callable, dest::Union(StridedArray,BitArray), A::Number, B::Union(StridedArray,BitArray))
    for i=1:length(B)
        dest[i] = f(A, B[i])
    end
    return dest
end

function map!(f::Callable, dest::Union(StridedArray,BitArray), As::Union(StridedArray,BitArray)...)
    n = length(As[1])
    i = 1
    ith = a->a[i]
    for i=1:n
        dest[i] = f(map(ith, As)...)
    end
    return dest
end

## Filter ##

function filter(f::Function, Bs::BitArray)
    boolmap::Array{Bool} = map(f, Bs)
    Bs[boolmap]
end

## Transpose ##

transpose(B::BitVector) = reshape(copy(B), 1, length(B))

# fast 8x8 bit transpose from Henry S. Warrens's "Hacker's Delight"
# http://www.hackersdelight.org/HDcode/transpose8.c.txt
function transpose8x8(x::Uint64)
    y = x
    t = (y $ (y >>> 7)) & 0x00aa00aa00aa00aa
    y = y $ t $ (t << 7)
    t = (y $ (y >>> 14)) & 0x0000cccc0000cccc
    y = y $ t $ (t << 14)
    t = (y $ (y >>> 28)) & 0x00000000f0f0f0f0
    return y $ t $ (t << 28)
end

function form_8x8_chunk(Bc::Vector{Uint64}, i1::Int, i2::Int, m::Int, cgap::Int, cinc::Int, nc::Int, msk8::Uint64)
    x = uint64(0)

    k, l = get_chunks_id(i1 + (i2 - 1) * m)
    r = 0
    for j = 1 : 8
        if k > nc
            break
        end
        x |= ((Bc[k] >>> l) & msk8) << r
        if l + 8 >= 64 && nc > k
            r0 = 8 - (@_mod64 (l + 8))
            x |= (Bc[k + 1] & (msk8 >>> r0)) << (r + r0)
        end
        k += cgap + (l + cinc >= 64 ? 1 : 0)
        l = @_mod64 (l + cinc)
        r += 8
    end
    return x
end

# note: assumes B is filled with 0's
function put_8x8_chunk(Bc::Vector{Uint64}, i1::Int, i2::Int, x::Uint64, m::Int, cgap::Int, cinc::Int, nc::Int, msk8::Uint64)
    k, l = get_chunks_id(i1 + (i2 - 1) * m)
    r = 0
    for j = 1 : 8
        if k > nc
            break
        end
        Bc[k] |= ((x >>> r) & msk8) << l
        if l + 8 >= 64 && nc > k
            r0 = 8 - (@_mod64 (l + 8))
            Bc[k + 1] |= ((x >>> (r + r0)) & (msk8 >>> r0))
        end
        k += cgap + (l + cinc >= 64 ? 1 : 0)
        l = @_mod64 (l + cinc)
        r += 8
    end
    return
end

function transpose(B::BitMatrix)
    l1 = size(B, 1)
    l2 = size(B, 2)
    Bt = falses(l2, l1)

    cgap1 = @_div64 l1
    cinc1 = @_mod64 l1

    cgap2 = @_div64 l2
    cinc2 = @_mod64 l2

    Bc = B.chunks
    Btc = Bt.chunks

    nc = length(Bc)

    for i = 1 : 8 : l1

        msk8_1 = uint64(0xff)
        if (l1 < i + 7)
            msk8_1 >>>= i + 7 - l1
        end

        for j = 1 : 8 : l2
            x = form_8x8_chunk(Bc, i, j, l1, cgap1, cinc1, nc, msk8_1)
            x = transpose8x8(x)

            msk8_2 = uint64(0xff)
            if (l2 < j + 7)
                msk8_2 >>>= j + 7 - l2
            end

            put_8x8_chunk(Btc, j, i, x, l2, cgap2, cinc2, nc, msk8_2)
        end
    end
    return Bt
end

ctranspose(B::BitArray) = transpose(B)

## Permute array dims ##

function permutedims(B::Union(BitArray,StridedArray), perm)
    dimsB = size(B)
    ndimsB = length(dimsB)
    (ndimsB == length(perm) && isperm(perm)) || error("no valid permutation of dimensions")
    dimsP = ntuple(ndimsB, i->dimsB[perm[i]])::typeof(dimsB)
    P = similar(B, dimsP)
    permutedims!(P, B, perm)
end

## Concatenation ##

function hcat(B::BitVector...)
    height = length(B[1])
    for j = 2:length(B)
        length(B[j]) == height || error("dimensions must match")
    end
    M = BitArray(height, length(B))
    for j = 1:length(B)
        copy_chunks!(M.chunks, (height*(j-1))+1, B[j].chunks, 1, height)
    end
    return M
end

function vcat(V::BitVector...)
    n = 0
    for Vk in V
        n += length(Vk)
    end
    B = BitArray(n)
    j = 1
    for Vk in V
        copy_chunks!(B.chunks, j, Vk.chunks, 1, length(Vk))
        j += length(Vk)
    end
    return B
end

function hcat(A::Union(BitMatrix,BitVector)...)
    nargs = length(A)
    nrows = size(A[1], 1)
    ncols = 0
    dense = true
    for j = 1:nargs
        Aj = A[j]
        nd = ndims(Aj)
        ncols += (nd==2 ? size(Aj,2) : 1)
        if size(Aj, 1) != nrows; error("rows must match"); end
    end

    B = BitArray(nrows, ncols)

    pos = 1
    for k=1:nargs
        Ak = A[k]
        n = length(Ak)
        copy_chunks!(B.chunks, pos, Ak.chunks, 1, n)
        pos += n
    end
    return B
end

function vcat(A::BitMatrix...)
    nargs = length(A)
    nrows = sum(a->size(a, 1), A)::Int
    ncols = size(A[1], 2)
    for j = 2:nargs
        if size(A[j], 2) != ncols; error("columns must match"); end
    end
    B = BitArray(nrows, ncols)
    Bc = B.chunks
    nrowsA = [size(a, 1) for a in A]
    Ac = [a.chunks for a in A]
    pos_d = 1
    pos_s = ones(Int, nargs)
    for j = 1:ncols
        for k = 1:nargs
            copy_chunks!(Bc, pos_d, Ac[k], pos_s[k], nrowsA[k])
            pos_s[k] += nrowsA[k]
            pos_d += nrowsA[k]
        end
    end
    return B
end

# general case, specialized for BitArrays and Integers
function cat(catdim::Integer, X::Union(BitArray, Integer)...)
    nargs = length(X)
    # using integers results in conversion to Array{Int}
    # (except in the all-Bool case)
    has_bitarray = false
    has_integer = false
    for a in X
        if isa(a, BitArray)
            has_bitarray = true
        else
            has_integer = true
        end
    end
    # just integers and no BitArrays -> general case
    has_bitarray || return invoke(cat, (Integer, Any...), catdim, X...)
    dimsX = map((a->isa(a,BitArray) ? size(a) : (1,)), X)
    ndimsX = map((a->isa(a,BitArray) ? ndims(a) : 1), X)
    d_max = maximum(ndimsX)

    if catdim > d_max + 1
        for i=1:nargs
            if dimsX[1] != dimsX[i]
                error("all inputs must have same dimensions when concatenating along a higher dimension");
            end
        end
    elseif nargs >= 2
        for d=1:d_max
            if d == catdim; continue; end
            len = d <= ndimsX[1] ? dimsX[1][d] : 1
            for i = 2:nargs
                if len != (d <= ndimsX[i] ? dimsX[i][d] : 1)
                    error("mismatch in dimension ", d)
                end
            end
        end
    end

    cat_ranges = ntuple(nargs, i->(catdim <= ndimsX[i] ? dimsX[i][catdim] : 1))

    function compute_dims(d)
        if d == catdim
            if catdim <= d_max
                return sum(cat_ranges)
            else
                return nargs
            end
        else
            if d <= ndimsX[1]
                return dimsX[1][d]
            else
                return 1
            end
        end
    end

    ndimsC = max(catdim, d_max)
    dimsC = ntuple(ndimsC, compute_dims)::(Int...)
    typeC = promote_type(map(x->isa(x,BitArray) ? eltype(x) : typeof(x), X)...)
    if !has_integer || typeC == Bool
        C = BitArray(dimsC)
    else
        C = Array(typeC, dimsC)
    end

    range = 1
    for k = 1:nargs
        nextrange = range+cat_ranges[k]
        cat_one = ntuple(ndimsC, i->(i != catdim ?
                                     (1:dimsC[i]) : (range:nextrange-1) ))
        # note: when C and X are BitArrays, this calls
        #       the special assign with ranges
        C[cat_one...] = X[k]
        range = nextrange
    end
    return C
end

# hvcat -> use fallbacks in abstractarray.jl

isequal(A::BitArray, B::BitArray) = (A == B)

# Hashing

hash(B::BitArray) = hash((size(B), B.chunks))


using JSON
using MD5
using Sockets

function formatError(p, incNr)
    write(p.socket, "You deliviered the wrong format. This incident will be reported. $incNr\n")
end

function updateNote(p::Player, codetype)
    if codetype == "bombcode"
        notename = "note"
    elseif codetype == "cheesecode"
        notename = "leaflet"
    end

    if !(isfile(string("data/.", notename)))
        #println("No notes files found")
        notes = []
    else
        #println("Notes file found")
        f = open_file_try(string("data/.", notename), "r")
        notes = JSON.parse(read(f, String))
        close(f)
    end

    current_length = length(notes)

    new_notes = [p.dimension]
    if (0 <= current_length < note_max_length)
        append!(new_notes, notes)
    elseif (current_length == note_max_length)
        append!(new_notes, notes[1:note_max_length-1])

        old_path = string("data/.", codetype, "_", notes[note_max_length])
        rm(old_path)
    else
        return
    end

    f = open_file_try(string("data/.", notename), "w")
    write(f, JSON.json(new_notes))
    close(f)

end

function updateOS(p::Player)
    print_dict(p, "cryptomat_os_update_mode")
    user_input = readline(p.socket)
    while true
        if (user_input == "💣" || user_input == "🧀")
            break
        elseif (user_input == "l")
            return
        else
            print_dict("crpytomat_os_update_invalid_mode")
            return
        end
    end
    mode = user_input

    print_dict(p, "cryptomat_os_update_1")
    new_os = readline(p.socket)

    #validate input
    try
        new_os = JSON.parse(new_os)
    catch y
        formatError(p, 0)
        return
    end

    if !(isa(new_os, Array{Any, 1})
        && length(new_os) == 2
        && isa(new_os[1], String)
        && isa(new_os[2], Array{Any, 1}))
        formatError(p, 1)
        return
    end

    for a in new_os[2]
        if !isa(a, Int)
            formatError(p, 2)
            return
        end
    end

    print_dict(p, "cryptomat_os_update_accept_format")

    new_os_json = JSON.json(new_os)

    try
        proc = run(`./cryptomat/rsa_sig.py $new_os_json`)
    catch
        write(p.socket, "Not accepting this signature\n")
        return
    end

    print_dict(p, "cryptomat_os_update_accept_signature")

    os_string = new_os[1]

    write(p.socket, "Updating...\n")
    #println(os_string)
    if mode == "💣"
        codetype = "bombcode"
    elseif mode == "🧀"
        codetype = "cheesecode"
    end
    path = string("data/.", codetype ,"_", p.dimension)
    f = open_file_try(path, "w")
    write(f, os_string)
    close(f)

    updateNote(p, codetype)
    write(p.socket, "Updated\n")

end


"""
    ransac(pc, params, setenabled; reset_rand = false)

Run the RANSAC algorithm on a pointcloud with the given parameters.

Return the extracted primitives and the time it took to run the algorithm (in seconds).

# Arguments
- `pc::RANSACCloud`: the point cloud.
- `params::NamedTuple`: parameters.
- `setenabled::Bool`: if `true`: set every point to enabled.
- `reset_rand::Bool=false`: if `true`, resets the random seed with `Random.seed!(1234)`
"""
function ransac(pc, params, setenabled; reset_rand = false)
    if setenabled
        for i in eachindex(pc.isenabled)
            pc.isenabled[i] = true
        end
    end
    ransac(pc, params, reset_rand=reset_rand)
end

"""
    ransac(pc, params; reset_rand = false)

Run the RANSAC algorithm on a pointcloud with the given parameters.

Return the extracted primitives and the time it took to run the algorithm (in seconds).

# Arguments
- `pc::RANSACCloud`: the point cloud.
- `params::NamedTuple`: parameters.
- `reset_rand::Bool=false`: if `true`, resets the random seed with `Random.seed!(1234)`
"""
function ransac(pc, params; reset_rand = false)
    reset_rand && Random.seed!(1234)

    @extract params : iter_params=iteration
    @extract iter_params : drawN minsubsetN prob_det τ
    @extract iter_params : itermax shape_types
    @extract iter_params : extract_s terminate_s

    # @unpack drawN, minsubsetN, prob_det, τ = params
    # @unpack itermax, shape_types = params
    # @unpack extract_s, terminate_s = params
    start_time = time_ns()

    # FittedShape[] is abstract array, which is bad
    candidates = FittedShape[]
    scoredshapes = IterationCandidates()
    extracted = ExtractedShape[]

    # save octree level of a FittedShape
    shape_octree_level = Int[]

    # smallest distance in the pointcloud
    #lsd = smallestdistance(pc.vertices)

    # allocate for the random selected points
    sd = Vector{Int}(undef, drawN)
    @logmsg IterInf "Iteration begins."

    # for logging
    notifit = itermax > 10 ? div(itermax,10) : 1

    # track the number of candidates for the probabilities
    # 1. :lengthC - number of candidates in a given iteration
    # 2. :allcand - number of candidates that have been ever scored
    # 3. :nofminset - number of minimal sets that have been drawn so far
    countcandidates = [0, 0, 0]

    # iterate begin
    for k in 1:itermax
        if count(pc.isenabled) < τ
            @logmsg IterInf "Break at $k iteration, because left less points, then τ"
            break
        end
        # generate minsubsetN candidate
        for i in 1:minsubsetN
            sampling_res = samplepointcloud4!(sd, pc, params)
            sampling_res[1] || continue

            # sd: indexes of the actually selected points
            #TODO: this should be something more general
            # fit plane to the selected points
            f_v = @view pc.vertices[sd]
            f_n = @view pc.normals[sd]

            forcefitshapes!(f_v, f_n, params, candidates, shape_octree_level, sampling_res[2], pc)
        end # for t

        # evaluate the compatible points, currently used as score
        # TODO: do something with octree levels and scores
        which_ = 1

        # update number of candidates
        countcandidates[2] += size(candidates,1)
        scorecandidates!(pc, scoredshapes, candidates, which_, params, shape_octree_level)
        # by now every candidate is scored into scoredshapes

        # considered so far k*minsubsetN pcs. minimal sets
        countcandidates[3] = k*minsubsetN
        # length of the candidate struct

        countcandidates[1] = length(scoredshapes)
        if ! (length(scoredshapes) < 1)
            best = findhighestscore(scoredshapes)
            bestshape = scoredshapes.shapes[best.index]
            # TODO: refine if best.overlap
            scr = E(scoredshapes.scores[best.index])
            #best_length = size(bestshape.inpoints, 1)

            s = chooseS(countcandidates, extract_s)
            ppp = prob(scr, s, pc.size, drawN)

            #info printing: there's a best
            if k%notifit == 0
                @logmsg IterInf "$k. it, best: $scr score, prob: $ppp, scored shapes: $(length(scoredshapes)) pcs."
            end

            # if the probability is large enough, extract the shape
            if ppp > prob_det

                # TODO: proper refit, not only getting the points that fit to that shape
                # what do you mean by refit?
                # refit on the whole pointcloud
                extr_shape = refit(bestshape, pc, params)

                if ! (extr_shape === nothing)
                    @logmsg IterInf "Extracting best: $(strt(bestshape)) score: $scr, size: $(size(extr_shape.inpoints,1)) ps."
                    # invalidate indexes
                    invalidate_indexes!(pc, extr_shape.inpoints)
                    # extract the shape and delete from scoredshapes
                    push!(extracted, extr_shape)
                    deleteat!(scoredshapes, best.index)
                    # remove invalid shapes
                    removeinvalidshapes!(pc, scoredshapes)
                end
            end # if extract shape
        else
            #info printing: not best
            if k%notifit == 0
                @logmsg IterInf "$k. it, no candidates."
            end
        end
        # update octree levels
        updatelevelweight(pc)

        # check exit condition
        s = chooseS(countcandidates, terminate_s)
        #if prob(τ/subsetN, s, pc.size, drawN) > prob_det
        if prob(τ, s, pc.size, drawN) > prob_det
            @logmsg IterInf "Break, at this point all shapes should be extracted: $k. iteration."
            break
        end
    end # iterate end
    fint = trunc((time_ns() - start_time)/1_000_000_000, digits=2)
    @logmsg IterInf "Iteration finished in $fint seconds with $(length(extracted)) extracted and $(length(scoredshapes)) scored shapes."
    return extracted, fint
end # ransac function


"""
    ComposedTransform(outer, inner)

A subtype of `CoordinateTransform` that lazily combines index `outer` and `inner`. Indexing a
`ComposedTransform` whith `i` is equivalent to `outer[inner[i]]`.
"""
struct ComposedTransform{NI,NO,O,I} <: CoordinateTransform{NI,NO}
    outer::O
    inner::I

    function ComposedTransform(o::CoordinateTransform{O,N}, i::CoordinateTransform{N,I}) where {O,N,I}
        @assert !(i isa ComposedTransform)
        new{O,I,typeof(o),typeof(i)}(o, i)
    end
end

@inline (t::ComposedTransform)(x) = getfield(t, :outer)(getfield(t, :inner)(x))

function Base.show(io::IO, m::MIME"text/plain", @nospecialize(x::ComposedTransform))
    show(io, m, x.outer)
    print(io, " ∘ ")
    show(io, m, x.inner)
end

@inline (t::CoordinateTransform)(::IdentityTransform) = t
@inline (t1::CoordinateTransform)(t2::CoordinateTransform) = ComposedTransform(t2, t1)

@inline function (t1::VecPermute)(t2::StrideTransform{1})
    s = getfield(getfield(t1, :strides), 1)
    StrideTransform((s, s), (static(1), getfield(getfield(t1, :offsets), 1)))
end

@inline function (t2::Permute{I1,I2,N})(t1::StrideTransform{N}) where {N,I1,I2}
    StrideTransform{N}(t2(getfield(x, :strides)), t2(getfield(x, :offsets)))
end

function (v::View{NI,NO,I})(t::StrideTransform{NO,S,O}) where {NI,NO,I,S,O}
    _combined_sub_strides(ArrayInterface.stride_preserving_index(I), t, v)
end

@inline function _combined_sub_strides(::False, t::StrideTransform, v::View{NI,NO,I}) where {NI,NO,I}
    __map_strides(View(_map_strides(
        Static.eachop(ArrayInterface._ndims_index, Static.nstatic(Val(known_length(I))), I),
        getfield(v, :indices),
        getfield(t, :strides),
        getfield(t, :offsets)
        )))
end
@inline function __map_strides(v::View{<:Any,O}) where {O}
    ComposedTransform(ComposedTransform(Offset(static(1)), ReduceTransform{O,typeof(+)}(+)), v)
end
@generated function _map_strides(::NDimsIndex, inds::I, s::S, o::O) where {NDimsIndex,I,S,O}
    t = Expr(:tuple)
    dim = 0
    ndi = known(NDimsIndex)
    for i in 1:length(ndi)
        indsexpr = :(@inbounds(getfield(inds, $i)))
        stride_expr = Expr(:tuple)
        offset_expr = Expr(:tuple)
        for j in 1:ndi[i]
            dim += 1
            push!(stride_expr.args, :(@inbounds(getfield(s, $dim))))
            push!(offset_expr.args, :(@inbounds(getfield(o, $dim))))
        end
        push!(t.args, :(StrideTransform($stride_expr, $offset_expr)($indsexpr)))
    end
    Expr(:block, Expr(:meta, :inline), t)
end
@inline function _combined_sub_strides(::True,
    x::StrideTransform{NO,S,O},
    i::View{NI,NO,I}
) where {NI,NO,S,O,I}
    pdims = ArrayInterface._to_sub_dims(I)
    o = getfield(x, :offsets)
    s = getfield(x, :strides)
    inds = getfield(i, :indices)
    out = StrideTransform(
        eachop(getmul, pdims, map(maybe_static_step, inds), s),
        Static.permute(o, pdims)
    )
    ComposedTransform(Offset(Static.reduce_tup(+, map(*, map(_diff, inds, o), s))), out)
end
getmul(x::Tuple, y::Tuple, ::StaticInt{i}) where {i} = getfield(x, i) * getfield(y, i)
@inline _diff(::Base.Slice, ::Any) = Zero()
@inline _diff(x::AbstractRange, o) = static_first(x) - o
@inline _diff(x::Integer, o) = x - o


#CircuitAnalysis: Base tools

#==Main data structures
===============================================================================#

#Object used to provide type information to something:
struct DataTag{ID, T}
	v::T
end
#=TODO: Figure out how to broadcast operations on DataTag (ex: make 5 .* ycap
work)... but not ycap .* ycap.
=#

#==Type aliases
===============================================================================#

#Tags:
const TImpedance = DataTag{:Z}
const TAdmittance = DataTag{:Y}
const TInductance = DataTag{:L}
const TCapacitance = DataTag{:C}

#Base.size(x::DataTag) = size(x.v)
#Base.length(x::DataTag) = length(x.v)
#Base.eltype(x::DataTag{ID, T}) where {ID, T} = T

#==Useful validations/assertions
===============================================================================#

#==Constructor interfaces
===============================================================================#

(::Type{DataTag{ID}})(v::T) where {ID, T} = DataTag{ID, T}(v)

#==Helper functions
===============================================================================#
#multMD: disambiguate multiplication to use
multMD(a, b) = a*b #DataMD values don't require .*
multMD(a::AbstractArray, b::AbstractArray) = a .* b
divMD(a, b) = a/b #DataMD values don't require .*
divMD(a::AbstractArray, b::AbstractArray) = a ./ b

#==Operations
===============================================================================#
#Mostly for unit testing:
Base.:(≈)(a::DataTag, b::DataTag; kwargs...) = false
Base.:(≈)(a::DataTag{ID}, b::DataTag{ID}; kwargs...) where ID =
	Base.:(≈)(a.v, b.v; kwargs...)

Base.:(==)(a::DataTag, b::DataTag) = false
Base.:(==)(a::DataTag{ID}, b::DataTag{ID}) where ID = (a.v == b.v)

#_dotop(x)=Symbol(".$x")
_operators1 = [:+, :-] #Unary operator on a tagged object
_operators2 = [:+, :-, :*, :/]

for op in _operators1; @eval begin #CODEGEN-------------------------------------

Base.$op(d1::DataTag{ID}) where ID = DataTag{ID}(Base.$op(d1.v))

end; end #CODEGEN---------------------------------------------------------------

#=

for op in _operators2; @eval begin #CODEGEN--------------------------------

Base.$op(d1::DataTag{ID}, d2::Number) where ID =
	DataTag{ID}(Base.$op(d1.v, d2))

Base.$op(d1::Number, d2::DataTag{ID}) where ID =
	DataTag{ID}(Base.$op(d1, d2.v))

Base.$op(d1::DataTag{ID}, d2::DataTag{ID}) where ID =
	DataTag{ID}(Base.$op(d1.v, d2.v))

end; end #CODEGEN---------------------------------------------------------------

=#
#==Main functions
===============================================================================#
impedance(v) = TImpedance(v)
admittance(v) = TAdmittance(v)
capacitance(v) = TCapacitance(v)
inductance(v) = TInductance(v)

impedance(y::TAdmittance) = impedance(divMD(1, y.v))
admittance(z::TImpedance) = admittance(divMD(1, z.v))

#Capacitor values:
_Y(c::TCapacitance, f) = admittance(multMD((2*pi*c.v)im, f))
_Y(c::TCapacitance, ::Nothing) = throw(ArgumentError("Missing kwarg :f"))
admittance(c::TCapacitance; f = nothing) = _Y(c, f)
impedance(c::TCapacitance, args...; kwargs...) = impedance(admittance(c, args...; kwargs...))
#TODO: create admittance(:L, value, f=x)??

#Inductor values:
_Z(l::TInductance, f) = admittance(multMD((2*pi*c.v)im, f))
_Z(l::TInductance, ::Nothing) = throw(ArgumentError("Missing kwarg :f"))
impedance(l::TInductance; f = nothing) = _Z(l, f)
admittance(l::TInductance, args...; kwargs...) = admittance(impedance(l, args...; kwargs...))

#==Helper functions
===============================================================================#

#==Traits (utility functions for instances and types)
===============================================================================#

#Last line


function get_sunits_from_subfield_data!(c::Hecke.ClassGrpCtx, UZK::Hecke.UnitGrpCtx, N::NormRelation; compact::Bool = true)
  docompact = false
  onlyp = 0

  if Hecke.isprime_power(index(N)) && compact
    docompact = true
    _, onlyp = ispower(index(N))
  end

  @vprint :NormRelation 1 "Doing something in the subfields\n"
  if !docompact
    _add_sunits_from_brauer_relation!(c, UZK, N)
  else
    @vprint :NormRelation 1 "Using the compact presentation\n"
    _add_sunits_from_brauer_relation!(c, UZK, N, compact = onlyp)
  end
end

function class_group_via_brauer(O::NfOrd, N::NormRelation; compact::Bool = true)
  K = N.K
  OK = lll(maximal_order(nf(O)))

  bound = Hecke.factor_base_bound_grh(OK)
  @vprint :NormRelation 1 "Factor base bound: $bound\n"
  S = prime_ideals_up_to(OK, bound)
  #First, we try with a smaller factor base.

  Sfirst = NfOrdIdl[]
  primes_added = fmpz[]
  threshold = min(200, div(length(S), 2))
  i = 1
  while length(Sfirst) < threshold
    p = minimum(S[i], copy = false)
    if p in primes_added
      i += 1
      continue
    end
    push!(primes_added, p)
    push!(Sfirst, S[i])
    for j = i+1:length(S)
      if minimum(S[j], copy = false) == p
        push!(Sfirst, S[j])
      end
    end
    i += 1
  end
  @vprint :NormRelation 1 "Length of the first factor base: $(length(Sfirst))\n"
  c, UZK = _setup_for_norm_relation_fun(K, Sfirst)
  get_sunits_from_subfield_data!(c, UZK, N, compact = compact)
  for (p, e) in factor(index(N))
    @vprint :NormRelation 1 "Saturating at $p \n"
    b = Hecke.saturate!(c, UZK, Int(p), 3.5, easy_root = degree(K) < 30)
    while b
      idx = Hecke._validate_class_unit_group(c, UZK)[1]
      @vprint :NormRelation 1 "Index bound from analysis $idx\n"
      b = Hecke.saturate!(c, UZK, Int(p), 3.5, easy_root = degree(K) < 30)
    end
  end
  UZK.finished = true
  idx = Hecke._validate_class_unit_group(c, UZK)[1]
  if !isone(idx)
    c = _setup_for_norm_relation_fun(K, S)[1]
    get_sunits_from_subfield_data!(c, UZK, N, compact = compact)
    for (p, e) in factor(index(N))
      @vprint :NormRelation 1 "Saturating at $p \n"
      b = Hecke.saturate!(c, UZK, Int(p), 3.5, easy_root = degree(K) < 30)
      while b
        idx = Hecke._validate_class_unit_group(c, UZK)[1]
        @vprint :NormRelation 1 "Index bound from analysis $idx\n"
        b = Hecke.saturate!(c, UZK, Int(p), 3.5, easy_root = degree(K) < 30)
      end
    end
    idx = Hecke._validate_class_unit_group(c, UZK)[1]
  end
  @vprint :NormRelation 1 "Index is $idx (should be 1)!\n"
  @assert idx == 1
  @vprint :NormRelation 1 "\n"
  c.finished = true
  Hecke._set_ClassGrpCtx_of_order(OK, c)
  Hecke._set_UnitGrpCtx_of_order(OK, UZK)
  return class_group(c, O)
end


function demo_steady_solution()
  cur_T = 18.067

  cur_reactor = demo_steady_deck(cur_T)

  cur_reactor.n_bar = 0.8746
  cur_reactor.I_P = 21.627

  cur_reactor.R_0 = 7.500
  cur_reactor.B_0 = 5.627

  cur_reactor.P_F = 3255

  cur_reactor.f_BS = 0.611
  cur_reactor.f_CD = 1.00 - cur_reactor.f_BS
  cur_reactor.f_IN = 0.00

  cur_reactor.beta_N = cur_reactor.max_beta_N
  cur_reactor.q_95 = 4.405
  cur_reactor.P_W = 1.911

  cur_reactor.volume = 2217

  cur_reactor
end


# Tests for align.jl

@testset "Align" begin
    coords_one = [
        1.0 0.0 0.0;
        0.0 1.0 0.0;
        0.0 0.0 0.0;
    ]
    coords_two = [
        0.0 -1.0 0.0;
        1.0 0.0 0.0;
        1.0 1.0 1.0;
    ]
    trans_one_test, trans_two_test, rotation_test = kabschalignment(coords_one, coords_two)
    trans_one_real = [1/3, 1/3, 0]
    trans_two_real = [-1/3, 1/3, 1]
    rotation_real = [
        0.0 -1.0 0.0;
        1.0 0.0 0.0;
        0.0 0.0 1.0;
    ]
    @test isapprox(trans_one_test, trans_one_real)
    @test isapprox(trans_two_test, trans_two_real)
    @test isapprox(rotation_test, rotation_real)

    coords_one = [
        1.0 0.0 0.0;
        0.0 1.0 0.0;
        0.0 0.0 0.0;
    ]
    coords_two = [
        0.0 -1.0 0.0;
        1.0 0.0 0.0;
        1.0 1.0 1.0;
    ]
    devs_test = displacements(coords_one, coords_two)
    devs_real = [sqrt(3), sqrt(3), 1.0]
    @test isapprox(devs_test, devs_real)

    atom_one = Atom("CA", "ALA", 'A', 20, [0.0, 0.0, 0.0], "C")
    atom_two = Atom("C", "ALA", 'A', 20, [0.0, 0.0, 0.0], "C")
    atom_thr = Atom("CA", "ALA", 'A', 21, [0.0, 0.0, 0.0], "C")
    atoms = [atom_one, atom_two, atom_thr]
    coords_one = [
        0.0 1.0 2.0;
        0.0 1.0 0.0;
        0.0 0.0 0.0;
    ]
    coords_two = [
        0.0 -1.0 0.0;
        0.0 0.0 3.0;
        3.0 3.0 3.0;
    ]
    align!(coords_one, coords_two, atoms)
    coords_one_real = [
        0.0 0.5 3.0;
        -1.0 1.5 3.0;
        0.0 2.5 3.0;
    ]
    coords_one_real = [
        0.0 -1.0 0.0;
        0.5 1.5 2.5;
        3.0 3.0 3.0;
    ]
    @test coords_one == coords_one_real

    atom_one = Atom("CA", "ALA", 'A', 20, [0.0, 0.0, 0.0], "C")
    atom_two = Atom("C", "ALA", 'A', 20, [0.0, 0.0, 0.0], "C")
    atom_thr = Atom("CA", "ALA", 'A', 21, [0.0, 0.0, 0.0], "C")
    atoms = [atom_one, atom_two, atom_thr]
    coords_one = [
        0.0 1.0 2.0;
        0.0 1.0 0.0;
        0.0 0.0 0.0;
    ]
    coords_two = [
        0.0 -1.0 0.0;
        0.0 0.0 3.0;
        3.0 3.0 3.0;
    ]
    alignsimple!(coords_one, coords_two, atoms)
    coords_one_real = [
        0.0 -1.0 0.0;
        0.5 1.5 2.5;
        3.0 3.0 3.0;
    ]
    @test coords_one == coords_one_real

    atom_one = Atom("CA", "ALA", 'A', 20, [0.0, 0.0, 0.0], "C")
    atom_two = Atom("C", "ALA", 'A', 20, [1.0, 1.0, 0.0], "C")
    atom_thr = Atom("CA", "ALA", 'A', 21, [2.0, 0.0, 0.0], "C")
    atoms = [atom_one, atom_two, atom_thr]
    atoms_new = deepcopy(atoms)
    coords_ref = [
        0.0 -1.0 0.0;
        0.0 0.0 3.0;
        3.0 3.0 3.0;
    ]
    alignatoms!(atoms, coords_ref)
    coords_real = [
        0.0 -1.0 0.0;
        0.5 1.5 2.5;
        3.0 3.0 3.0;
    ]
    @test atomcoords(atoms) == coords_real
    atoms_ref = atomcoords(atoms, coords_ref)
    alignatoms!(atoms_new, atoms_ref)
    @test atomcoords(atoms_new) == coords_real

    atom_one = Atom("CA", "ALA", 'A', 20, [0.0, 0.0, 0.0], "C")
    atom_two = Atom("CA", "ALA", 'A', 30, [0.0, 0.0, 0.0], "C")
    struc_one = ModelledStructure(0.0, [0.0 5.0; 0.0 5.0; 0.0 5.0])
    struc_two = ModelledStructure(0.0, [1.0 4.0; 1.0 4.0; 1.0 4.0])
    ensemble = ModelledEnsemble([atom_one, atom_two], [struc_one, struc_two])
    average_test = centroid(ensemble)
    average_real = [0.5 4.5; 0.5 4.5; 0.5 4.5]
    @test average_test == average_real

    coords_one = [
        0.0 3.0 13.0;
        0.0 0.0 0.0;
        0.0 0.0 0.0;
    ]
    coords_two = [
        0.0 5.0 10.0;
        1.0 0.0 0.0;
        0.0 0.0 0.0;
    ]
    atom_one = Atom("CA", "ALA", 'A', 20, [0.0, 0.0, 0.0], "C")
    atom_two = Atom("CA", "ALA", 'A', 30, [0.0, 0.0, 0.0], "C")
    atom_thr = Atom("C", "ALA", 'A', 40, [0.0, 0.0, 0.0], "C")
    atoms = [atom_one, atom_two, atom_thr]
    @test isapprox(rmsd(coords_one, coords_two, atoms), sqrt(2.5))

    coords_one = [
        0.0 3.0;
        0.0 0.0;
        0.0 0.0;
    ]
    coords_two = [
        0.0 5.0;
        1.0 0.0;
        0.0 0.0;
    ]
    @test isapprox(rmsd(coords_one, coords_two), sqrt(2.5))
end


@testset "Mergers" begin
    splitname = "splittestname"
    fbop = x -> 0.5 .* x.^2
    combineop = (x, y) -> x + σ.(y)
    m = Merger(splitname, fbop, combineop)
    x = randn(5)
    y = randn(5)
    h = Dict(splitname => y)

    @testset "constructor" begin
        @test m isa Merger{typeof(fbop), typeof(combineop)}
        @test m.splitname == splitname
        @test m.fb(x) ≈ fbop(x)
        @test m.op(x, y) ≈ combineop(x, y)
    end # @testset "constructor"

    @testset "apply" begin
        @test m(x, h) ≈ m.op(x, m.fb(y))
    end # @testset "apply"

    @testset "params" begin
        # Test that a Merger returns the parameters of its internal operations
        fb = Dense(10, 5)
        op = Chain((x, y) -> hcat(x, y), Dense(10, 5))
        m = Merger(splitname, fb, op)
        @test hcat(params(m)...) == hcat(params(fb)..., params(op)...)
    end # @testset "params"

    @testset "name" begin
        @test inputname(m) == splitname
    end # @testset "name"
end # @testset "Mergers"


module Bijectors

#=
  NOTE: Codes below are adapted from
  https://github.com/brian-j-smith/Mamba.jl/blob/master/src/distributions/transformdistribution.jl
  The Mamba.jl package is licensed under the MIT License:
  > Copyright (c) 2014: Brian J Smith and other contributors:
  >
  > https://github.com/brian-j-smith/Mamba.jl/contributors
  >
  > Permission is hereby granted, free of charge, to any person obtaining
  > a copy of this software and associated documentation files (the
  > "Software"), to deal in the Software without restriction, including
  > without limitation the rights to use, copy, modify, merge, publish,
  > distribute, sublicense, and/or sell copies of the Software, and to
  > permit persons to whom the Software is furnished to do so, subject to
  > the following conditions:
  >
  > The above copyright notice and this permission notice shall be
  > included in all copies or substantial portions of the Software.
  >
  > THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
  > EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
  > MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.
  > IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY
  > CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT,
  > TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE
  > SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
=#

using Reexport, Requires
@reexport using Distributions
using StatsFuns
using LinearAlgebra
using MappedArrays
using Base.Iterators: drop
using LinearAlgebra: AbstractTriangular
import Functors
import NonlinearSolve
import ChainRulesCore

export  TransformDistribution,
        PositiveDistribution,
        UnitDistribution,
        SimplexDistribution,
        PDMatDistribution,
        link,
        invlink,
        logpdf_with_trans,
        isclosedform,
        transform,
        forward,
        logabsdetjac,
        logabsdetjacinv,
        Bijector,
        ADBijector,
        Inverse,
        Composed,
        compose,
        Stacked,
        stack,
        Identity,
        bijector,
        transformed,
        UnivariateTransformed,
        MultivariateTransformed,
        logpdf_with_jac,
        logpdf_forward,
        PlanarLayer,
        RadialLayer,
        CouplingLayer,
        InvertibleBatchNorm

if VERSION < v"1.1"
    using Compat: eachcol
end

const DEBUG = Bool(parse(Int, get(ENV, "DEBUG_BIJECTORS", "0")))
_debug(str) = @debug str

_eps(::Type{T}) where {T} = T(eps(T))
_eps(::Type{Real}) = eps(Float64)
_eps(::Type{<:Integer}) = eps(Float64)

function _clamp(x, a, b)
    T = promote_type(typeof(x), typeof(a), typeof(b))
    ϵ = _eps(T)
    clamped_x = ifelse(x < a, convert(T, a), ifelse(x > b, convert(T, b), x))
    DEBUG && _debug("x = $x, bounds = $((a, b)), clamped_x = $clamped_x")
    return clamped_x
end

function mapvcat(f, args...)
    out = map(f, args...)
    init = vcat(out[1])
    return reshape(reduce(vcat, drop(out, 1); init = init), size(out))
end

function maphcat(f, args...)
    out = map(f, args...)
    init = reshape(out[1], :, 1)
    return reduce(hcat, drop(out, 1); init = init)
end
function eachcolmaphcat(f, x1, x2)
    out = [f(x1[:,i], x2[i]) for i in 1:size(x1, 2)]
    init = reshape(out[1], :, 1)
    return reduce(hcat, drop(out, 1); init = init)
end
function eachcolmaphcat(f, x)
    out = map(f, eachcol(x))
    init = reshape(out[1], :, 1)
    return reduce(hcat, drop(out, 1); init = init)
end
function sumeachcol(f, x1, x2)
    # Using a view below for x1 breaks Tracker
    return sum(f(x1[:,i], x2[i]) for i in 1:size(x1, 2))
end

# Distributions

link(d::Distribution, x) = bijector(d)(x)
invlink(d::Distribution, y) = inv(bijector(d))(y)
function logpdf_with_trans(d::Distribution, x, transform::Bool)
    if ispd(d)
        return pd_logpdf_with_trans(d, x, transform)
    elseif isdirichlet(d)
        l = logpdf(d, x .+ eps(eltype(x)))
    else
        l = logpdf(d, x)
    end
    if transform
        return l - logabsdetjac(bijector(d), x)
    else
        return l
    end
end

## Univariate

const TransformDistribution = Union{
    T,
    Truncated{T},
} where T <: ContinuousUnivariateDistribution
const PositiveDistribution = Union{
    BetaPrime, Chi, Chisq, Erlang, Exponential, FDist, Frechet, Gamma, InverseGamma,
    InverseGaussian, Kolmogorov, LogNormal, NoncentralChisq, NoncentralF, Rayleigh, Weibull,
}
const UnitDistribution = Union{Beta, KSOneSided, NoncentralBeta}

function logpdf_with_trans(d::UnivariateDistribution, x, transform::Bool)
    if transform
        return map(x -> logpdf(d, x), x) - logabsdetjac(bijector(d), x)
    else
        return map(x -> logpdf(d, x), x)
    end
end

## Multivariate

const SimplexDistribution = Union{Dirichlet}
isdirichlet(::SimplexDistribution) = true
isdirichlet(::Distribution) = false

###########
# ∑xᵢ = 1 #
###########

function link(
    d::Dirichlet,
    x::AbstractVecOrMat{<:Real},
    proj::Bool = true,
)
    return SimplexBijector{proj}()(x)
end

function link_jacobian(
    d::Dirichlet,
    x::AbstractVector{T},
    proj::Bool = true,
) where {T<:Real}
    return jacobian(SimplexBijector{proj}(), x)
end

function invlink(
    d::Dirichlet,
    y::AbstractVecOrMat{<:Real},
    proj::Bool = true
)
    return inv(SimplexBijector{proj}())(y)
end
function invlink_jacobian(
    d::Dirichlet,
    y::AbstractVector{T},
    proj::Bool = true
) where {T<:Real}
    return jacobian(inv(SimplexBijector{proj}()), y)
end

## Matrix

#####################
# Positive definite #
#####################

const PDMatDistribution = Union{MatrixBeta, InverseWishart, Wishart}
ispd(::Distribution) = false
ispd(::PDMatDistribution) = true

function logpdf_with_trans(
    d::MatrixDistribution,
    X::AbstractArray{<:AbstractMatrix{<:Real}},
    transform::Bool,
)
    return map(X) do x
        logpdf_with_trans(d, x, transform)
    end
end
function pd_logpdf_with_trans(d, X::AbstractMatrix{<:Real}, transform::Bool)
    T = eltype(X)
    Xcf = cholesky(X, check = false)
    if !issuccess(Xcf)
        Xcf = cholesky(X + max(eps(T), eps(T) * norm(X)) * I)
    end
    lp = getlogp(d, Xcf, X)
    if transform && isfinite(lp)
        U = Xcf.U
        d = dim(d)
        lp += sum((d .- (1:d) .+ 2) .* log.(diag(U)))
        lp += d * log(T(2))
    end
    return lp
end
function getlogp(d::MatrixBeta, Xcf, X)
    n1, n2 = params(d)
    p = dim(d)
    return ((n1 - p - 1) / 2) * logdet(Xcf) + ((n2 - p - 1) / 2) * logdet(I - X) + d.logc0
end
function getlogp(d::Wishart, Xcf, X)
    return 0.5 * ((d.df - (dim(d) + 1)) * logdet(Xcf) - tr(d.S \ X)) + d.logc0
end
function getlogp(d::InverseWishart, Xcf, X)
    Ψ = Matrix(d.Ψ)
    return -0.5 * ((d.df + dim(d) + 1) * logdet(Xcf) + tr(Xcf \ Ψ)) + d.logc0
end

include("utils.jl")
include("interface.jl")
include("chainrules.jl")

# Broadcasting here breaks Tracker for some reason
maporbroadcast(f, x::AbstractArray{<:Any, N}...) where {N} = map(f, x...)
maporbroadcast(f, x::AbstractArray...) = f.(x...)

# optional dependencies
function __init__()
    @require LazyArrays = "5078a376-72f3-5289-bfd5-ec5146d43c02" begin
        function maporbroadcast(f, x1::LazyArrays.BroadcastArray, x...)
            return copy(f.(x1, x...))
        end
        function maporbroadcast(f, x1, x2::LazyArrays.BroadcastArray, x...)
            return copy(f.(x1, x2, x...))
        end
        function maporbroadcast(f, x1, x2, x3::LazyArrays.BroadcastArray, x...)
            return copy(f.(x1, x2, x3, x...))
        end
    end
    @require ForwardDiff="f6369f11-7733-5829-9624-2563aa707210" include("compat/forwarddiff.jl")
    @require Tracker="9f7883ad-71c0-57eb-9f7f-b5c9e6d3789c" include("compat/tracker.jl")
    @require Zygote="e88e6eb3-aa80-5325-afca-941959d7151f" include("compat/zygote.jl")
    @require ReverseDiff="37e2e3b7-166d-5795-8a7a-e32c996b4267" include("compat/reversediff.jl")
    @require DistributionsAD="ced4e74d-a319-5a8a-b0ac-84af2272839c" include("compat/distributionsad.jl")
end

end # module


mutable struct DESPOTSolver{S,A,O,B,RS} <: POMDPs.Solver
    belief::DESPOTBelief{S,A,O}
    bounds::B
    random_streams::RS
    root::VNode{S,A,O,B}
    node_count::Int64
    config::DESPOTConfig
    rng::AbstractRNG
    #preallocated for simulations
    curr_reward::DESPOTReward
    next_state::S
    curr_obs::O
    print_bounds::Bool

    # default constructor
    function DESPOTSolver{S,A,O,B,RS}(  ;
                            bounds::B = B(), #TODO: fix
                            rng::AbstractRNG = Base.GLOBAL_RNG,
                            search_depth::Int64 = 90,
                            main_seed::UInt32 = convert(UInt32, 42),
                            time_per_move::Float64 = 1.0,                 # sec
                            n_particles::Int64 = 500,
                            pruning_constant::Float64 = 0.0,
                            eta::Float64 = 0.95,
                            sim_len::Int64 = -1,
                            approximate_ubound::Bool = false,
                            tiny::Float64 = 1e-6,
                            max_trials::Int64 = -1,
                            rand_max::Int64 = 2147483647,
                            debug::Int64 = 0,
                            random_streams::RS = RandomStreams(n_particles, search_depth, main_seed),
                            next_state = S(),
                            curr_obs = O(),
                            print_bounds = false
                           ) where {S,A,O,B,RS}

        this = new()

        # supplied variables
        this.bounds = bounds

        # Instantiate and initialize config
        this.config = DESPOTConfig()
        this.config.search_depth = search_depth
        this.config.main_seed = main_seed
        this.config.time_per_move = time_per_move
        this.config.n_particles = n_particles
        this.config.pruning_constant = pruning_constant
        this.config.eta = eta
        this.config.sim_len = sim_len
        this.config.approximate_ubound = approximate_ubound
        this.config.tiny = tiny
        this.config.max_trials = max_trials
        this.config.rand_max = rand_max
        this.config.debug = debug
        this.rng = rng
        this.next_state = next_state
        this.curr_obs = curr_obs
        this.curr_reward = 0.0
        this.random_streams = random_streams
        this.print_bounds = print_bounds
        return this
    end
end

function init_solver{S,A,O,B}(solver::DESPOTSolver{S,A,O,B}, pomdp::POMDPs.POMDP{S,A,O})

    # Instantiate random streams

    fill_random_streams!(solver.random_streams, solver.config.rand_max)
    init_bounds(solver.bounds, pomdp, solver.config)

    return nothing
end

function new_root{S,A,O,B}(solver::DESPOTSolver{S,A,O,B},
                  pomdp::POMDP{S,A,O},
                  belief::DESPOTBelief{S})

    solver.belief = belief
    solver.root = VNode{S,A,O,B}(
                        pomdp,
                        belief.particles,
                        solver.bounds,
                        0,
                        1.0,
                        false,
                        solver.config)
    return nothing
end

function search{S,A,O,B}(solver::DESPOTSolver{S,A,O,B}, pomdp::POMDP{S,A,O})
    n_trials::Int64 = 0
    start_time::Float64 = time()
    stop_now::Bool = false

    if solver.print_bounds
        @printf("Before: lBound = %.10f, uBound = %.10f\n", solver.root.lbound, solver.root.ubound)
    end

    while ((excess_uncertainty(solver.root.lbound,
                                solver.root.ubound,
                                solver.root.lbound,
                                solver.root.ubound,
                                0,
                                solver.config.eta,
                                discount(pomdp)) > 1e-6)
                                && !stop_now)

        trial(solver, pomdp, solver.root, n_trials)
        n_trials += 1
        if ((solver.config.max_trials > 0) && (n_trials >= solver.config.max_trials)) ||
        ((solver.config.time_per_move > 0) && ((time() - start_time) >= solver.config.time_per_move))
            stop_now = true
        end
    end

    if solver.print_bounds
        @printf("After:  lBound = %.10f, uBound = %.10f\n", solver.root.lbound, solver.root.ubound)
        @printf("Number of trials: %d\n", n_trials)
    end

    if solver.config.pruning_constant != 0.0
        total_pruned = prune(solver.root) # Number of non-child belief nodes pruned
        act = solver.root.pruned_action
        default = default_action(solver.bounds, pomdp, solver.root.particles, solver.config)
        return (act == -1 ? default : act), n_trials #TODO: fix actions
    elseif !solver.root.in_tree
        println("Root not in tree")
        default = default_action(solver.bounds, pomdp, solver.root.particles, solver.config)
        return default, n_trials
    else
        return get_lb_action(solver.root, solver.config, discount(pomdp)), n_trials
    end
    return nothing
end

function trial{S,A,O,B}(solver::DESPOTSolver{S,A,O,B}, pomdp::POMDP{S,A,O}, node::VNode{S,A,O,B}, n_trials::Int64)

    n_nodes_added::Int64 = 0
    ubound::Float64 = 0.0

    if (node.depth >= solver.config.search_depth) || isterminal(pomdp, node.particles[1].state)
      return 0 # nodes added
    end

    if isempty(node.q_nodes)
        expand_one_step(solver, pomdp, node, solver.bounds)
    end

    a_star::A = node.best_ub_action

    i_star::Int, weighted_eu_star::Float64 =
                               get_best_weuo(node.q_nodes[a_star],
                                             solver.root,
                                             solver.config,
                                             discount(pomdp)) # it's an array!

    o_star, vnode_star = node.q_nodes[a_star].obs_and_nodes[i_star]

    if weighted_eu_star > 0.0
        add(solver.belief.history, a_star, o_star)
        n_nodes_added = trial(solver,
                            pomdp,
                            vnode_star,
                            n_trials)
        remove_last(solver.belief.history)
    end
    node.n_tree_nodes += n_nodes_added

    # Backup
    potential_lbound = node.q_nodes[a_star].first_step_reward +
                    discount(pomdp) * get_lower_bound(node.q_nodes[a_star])
    node.lbound = max(node.lbound, potential_lbound)

    # As the upper bound of a_star may become smaller than the upper bound of
    # another action, we need to check all actions - unlike the lower bound.
    node.ubound = -Inf

    for a in iterator(actions(pomdp))
        ubound = node.q_nodes[a].first_step_reward +
            discount(pomdp) * get_upper_bound(node.q_nodes[a])
        if ubound > node.ubound
            node.ubound = ubound
            node.best_ub_action = a
        end
    end

    # Sanity check
    if (node.lbound > node.ubound + solver.config.tiny)
        println("depth = $(node.depth)")
        warn("Lower bound ($(node.lbound)) is higher than upper bound ($(node.ubound))")
    end

    if !node.in_tree
        node.in_tree = true
        node.n_tree_nodes += 1
        n_nodes_added += 1
    end
    return n_nodes_added
end

function expand_one_step{S,A,O,B}(solver::DESPOTSolver{S,A,O}, pomdp::POMDP{S,A,O}, node::VNode{S,A,O,B}, bounds::B)

    q_star::Float64 = -Inf
    first_step_reward::Float64 = 0.0
    remaining_reward::Float64 = 0.0
    rng = create_rng(solver.random_streams)

    for curr_action in iterator(actions(pomdp))
        obs_to_particles = Dict{O,Vector{DESPOTParticle{S}}}()

        for p in node.particles

            set_rng_state!(rng, solver.random_streams, p.id+1, node.depth+1)

            step(
                solver,
                pomdp,
                p.state,
                rng,
                curr_action)

            if !haskey(obs_to_particles, solver.curr_obs)
                obs_to_particles[solver.curr_obs] = DESPOTParticle{S}[]
            end

            push!(obs_to_particles[solver.curr_obs],
                  DESPOTParticle(solver.next_state,
                  p.id,
                  p.weight))
            first_step_reward += solver.curr_reward * p.weight
        end

        first_step_reward /= node.weight

        new_qnode = QNode{S,A,O,B}(
                          pomdp,
                          bounds,
                          obs_to_particles,
                          node.depth,
                          curr_action,
                          first_step_reward,
                          solver.belief.history,
                          solver.config)
        node.q_nodes[curr_action] = new_qnode
        remaining_reward = get_upper_bound(new_qnode)

        if (first_step_reward + discount(pomdp)*remaining_reward) > (q_star + solver.config.tiny)
            q_star = first_step_reward + discount(pomdp) * remaining_reward
            node.best_ub_action = curr_action
        end

        first_step_reward = 0.0
    end # for a
    return node
end

# fill in pre-allocated variables
function step{S,A,O,B}(solver::DESPOTSolver{S,A,O,B},
              pomdp::POMDPs.POMDP{S,A,O},
              state::S,
              rng::AbstractRNG,
              action::A)

    solver.next_state, solver.curr_obs, solver.curr_reward =
        POMDPs.generate_sor(pomdp, state, action, rng)

    return nothing
end


"""
Construct model for HydroGen with RunOfRiver Dispatch Formulation
"""
function PSI.construct_device!(
    optimization_container::PSI.OptimizationContainer,
    sys::PSY.System,
    model::PSI.DeviceModel{H, D},
    ::Type{S},
) where {H <: HydroCascade, D <: HydroDispatchRunOfRiverCascade, S <: PM.AbstractPowerModel}
    devices = PSI.get_available_components(H, sys)

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

    #Variables
    PSI.add_variables!(optimization_container, PSI.ActivePowerVariable, devices, D())
    PSI.add_variables!(optimization_container, PSI.ReactivePowerVariable, devices, D())
    PSI.add_variables!(optimization_container, PSI.SpillageVariable, 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),
    )=#
    flow_balance_cascade_constraint!(
        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, nothing)

    return
end

"""
Construct model for HydroGen with RunOfRiver Dispatch Formulation
with only Active Power.
"""
function PSI.construct_device!(
    optimization_container::PSI.OptimizationContainer,
    sys::PSY.System,
    model::PSI.DeviceModel{H, D},
    ::Type{S},
) where {
    H <: HydroCascade,
    D <: HydroDispatchRunOfRiverCascade,
    S <: PM.AbstractActivePowerModel,
}
    devices = PSI.get_available_components(H, sys)

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

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

    #Constraints
    #=
    PSI.add_constraints!(
        optimization_container,
        PSI.RangeConstraint,
        PSI.ActivePowerVariable,
        devices,
        model,
        S,
        PSI.get_feedforward(model),
    )=#
    flow_balance_cascade_constraint!(
        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, nothing)

    return
end


using Flux: @treelike
using MacroTools: @forward

using ..Basic
using ..Basic: onehot, AbstractTransformer
using ..Stacks

struct Gpt <: AbstractTransformer
    ts::Stack
    drop::Dropout
end

@treelike Gpt

@forward Gpt.ts Base.getindex, Base.length

"""
    Gpt(size::Int, head::Int, ps::Int, layer::Int;
        act = gelu, pdrop = 0.1, attn_pdrop = 0.1)
    Gpt(size::Int, head::Int, hs::Int, ps::Int, layer::Int;
        act = gelu, pdrop = 0.1, attn_pdrop = 0.1)

the Generative Pretrained Transformer(GPT) model.

    (gpt::Gpt)(x::T, mask=nothing; all::Bool=false)

eval the gpt layer on input `x`. If length `mask` is given (in shape (1, seq_len, batch_size)), mask the attention with `mask`. Moreover, set `all` to `true` to get all
outputs of each transformer layer.
"""
function Gpt(size::Int, head::Int, ps::Int, layer::Int;
             act = gelu, pdrop = 0.1, attn_pdrop = 0.1)
    rem(size, head) != 0 && error("size not divisible by head")
    Gpt(size, head, div(size, head), ps, layer; act=act, pdrop=pdrop, attn_pdrop=attn_pdrop)
end

function Gpt(size::Int, head::Int, hs::Int, ps::Int, layer::Int;
             act = gelu, pdrop = 0.1, attn_pdrop = 0.1)
    Gpt(
        Stack(
            @nntopo_str("x':x => $layer"),
            [
                Transformer(size, head, hs, ps; future=false, act=act, pdrop=attn_pdrop)
                for i = 1:layer
            ]...
        ),
        Dropout(pdrop)
    )
end

function (gpt::Gpt)(x::T, mask=nothing; all::Bool=false) where T
    e = gpt.drop(x)
    t, ts = gpt.ts(e)
    t = mask === nothing ? t : t .* mask
    if all
        t, ts
    else
        t #size(t) == (size, seq_len, batch)
    end
end

"""
    lmloss(embed, onehot, encoding, mask)

compute the language modeling loss for Gpt, onehot is the onehot array of the origin
input sentence. encoding the output of Gpt, mask is the mask between input sentences.

"""
lmloss(embed::Embed{T}, o::OneHotArray, t::AbstractArray{T}, mask) where T = lmloss(embed, tofloat(T, o), t, mask)
function lmloss(embed::Embed{T}, et, t::AbstractArray{T, N}, mask) where {T,N}
    if N == 3
        t = t[:, 1:end-1, :]
        s = size(t)
        sim = logsoftmax(transpose(embed.embedding) * reshape(t, s[1], :)) #(vocab, seq_len*batch)
        sim = reshape(sim, :, s[2], s[3])
        return logcrossentropy(et[:, 2:end, :], sim, mask[:, 2:end, :])
    elseif N == 2
        t = t[:, 1:end-1]
        sim = logsoftmax(transpose(embed.embedding) * t)
        return logcrossentropy(et[:, 2:end], sim, mask[:, 2:end])
    end
end

function Base.show(io::IO, gpt::Gpt)
    hs = div(size(gpt.ts[1].mh.iqproj.W)[1], gpt.ts[1].mh.head)
    h, ps = size(gpt.ts[1].pw.dout.W)

    print(io, "Gpt(")
    print(io, "layers=$(length(gpt.ts)), ")
    print(io, "head=$(gpt.ts[1].mh.head), ")
    print(io, "head_size=$(hs), ")
    print(io, "pwffn_size=$(ps), ")
    print(io, "size=$(h))")
end