export ReorderedBasis, reorder, invorder

using LuxurySparse, LinearAlgebra

"""
    ReorderedBasis{N, T}

Lazy reorderd basis.
"""
struct ReorderedBasis{N, T <: Integer}
    orders::NTuple{N, T}
    takers::NTuple{N, T}
    differ::NTuple{N, T}
end

"""
    ReorderedBasis(orders::NTuple{N, <:Integer})

Returns a lazy set of reordered basis.
"""
ReorderedBasis(orders::NTuple{N, T}) where {N, T <: Integer} =
    ReorderedBasis{N, T}(orders, bmask.(orders), unsafe_sub(1:N, orders))

Base.length(::ReorderedBasis{N}) where N = 1 << N
Base.eltype(::ReorderedBasis{N, T}) where {N, T} = T
Base.getindex(it::ReorderedBasis, k::Int) = next_reordered_basis(k - 1, it.takers, it.differ)

function Base.iterate(it::ReorderedBasis{N}, state = 1) where N
    if state - 1 == 1 << N
        return nothing
    else
        return next_reordered_basis(state - 1, it.takers, it.differ), state + 1
    end
end

"""
    next_reordered_basis(basis, takers, differ)

Returns the next reordered basis accroding to current basis.
"""
function next_reordered_basis(basis::T, takers::NTuple{N, T}, differ::NTuple{N, T}) where {N, T}
    out = zero(T)
    # NOTE: do not add simd here, this will cause other loop become less efficient
    #       have some confidence with the compiler :-)
    for i in 1:N
        @inbounds out += (basis & takers[i]) << differ[i]
    end
    return out
end

"""
    unsafe_reorder(X::AbstractArray, orders)

Reorder `X` according to `orders`.

!!! warning

    `unsafe_reorder` won't check whether the length of `orders` and
    the size of first dimension of `X` match, use at your own risk.
"""
function unsafe_reorder end

function unsafe_reorder(v::AbstractVector, orders::NTuple{N, <:Integer}) where N
    nv = similar(v)
    for (k, b) in enumerate(ReorderedBasis(orders))
        @inbounds nv[k] = v[b + 1]
    end
    return nv
end

function unsafe_reorder(A::AbstractMatrix, orders::NTuple{N, <:Integer}) where N
    od = Vector{Int}(undef, 1 << length(orders))
    for (i, b) in enumerate(ReorderedBasis(orders))
        @inbounds od[i] = 1 + b
    end

    od = invperm(od)
    return A[od, od]
end

unsafe_reorder(A::IMatrix, orders::NTuple{N, <:Integer}) where N = A

function unsafe_reorder(A::PermMatrix, orders::NTuple{N, <:Integer}) where N
    od = Vector{Int}(undef, 1 << length(orders))
    for (i, b) in enumerate(ReorderedBasis(orders))
        @inbounds od[i] = 1 + b
    end

    perm = similar(A.perm)
    vals = similar(A.vals)

    @simd for i = 1:length(perm)
        @inbounds perm[od[i]] = od[A.perm[i]]
        @inbounds vals[od[i]] = A.vals[i]
    end

    return PermMatrix(perm, vals)
end

function unsafe_reorder(A::Diagonal, orders::NTuple{N, <:Integer}) where N
    diag = similar(A.diag)
    for (i, b) in enumerate(ReorderedBasis(orders))
        @inbounds diag[b+1] = A.diag[i]
    end
    return Diagonal(diag)
end

function _reorder(v::AbstractArray, orders::NTuple{N, <:Integer}) where N
    length(orders) == log2dim1(v) || throw(DimensionMismatch("size of array not match length of order"))
    return unsafe_reorder(v, orders)
end

"""
    reorder(X::AbstractArray, orders)

Reorder `X` according to `orders`.

!!! tip
    Although `orders` can be any iterable, `Tuple` is preferred inorder to
    gain as much performance as possible. But the conversion won't take much
    anyway.
"""
function reorder end

reorder(v::AbstractArray, itr) = _reorder(v, Tuple(itr))
reorder(v::AbstractArray, orders::Tuple) = _reorder(v, promote(orders...))
reorder(v::AbstractArray, orders::NTuple{N, <:Integer}) where N = _reorder(v, orders)

"""
    invorder(X::AbstractVecOrMat)

Inverse the order of given vector/matrix `X`.
"""
invorder(X::AbstractVecOrMat) = reorder(X, Tuple(log2dim1(X):-1:1))
