module FEDW
using Printf: @printf,@sprintf
using SparseArrays: spzeros, sparse 
using QuantumLattices: AbstractFockOperator,Operators,FOperator,MAJORANA,CREATION,rank
using QuantumLattices.Prerequisites: Float
using QuantumLattices.Mathematics.VectorSpaces: VectorSpace
using QuantumLattices.Mathematics.Combinatorics: Combinations

import QuantumLattices: kind,id,dimension
import QuantumLattices.Mathematics.AlgebraOverFields: idtype
import ExactDiagonalization.FED:FockBasis

export NCBasis,bkron,matrix,boplus

"""
      NCBasis{I}(vnstate::Vector{Int},vnpart::Vector{Vector{T}}) where {I<:Unsigned,T<:Tuple{Vararg{Int}}}

Generic fermionic/hard-core-bosonic basis, which uses conserved number of particles.
 e.g.
    vnstate=[m1,m2]; vnpart=[[(n1,n2),(n3,n4)],[(l1,l2),(l3,l4)]] denotes \$[(C_{m_1}^{n_1}⊕C_{m_1}^{n_2})⊗(C_{m_2}^{n_3}⊕C_{m_2}^{n_4})]⊕[(C_{m_1}^{l_1}⊕C_{m_1}^{l_2})⊗(C_{m_2}^{l_3}⊕C_{m_2}^{l_4})]\$ Hilbert space
============================================================================================================
vnstate=[4,6] denotes four and six single-particle states at the 1st and 2nd space,respectively.
vnpart=[[(2,),(3,)],[(1,),(1,)]];vnpart[1] denotes two and three particles at the 1st and 2nd space,respectively.
Thus, there are binomial(4,2)*binomial(6,3)+ binomial(4,1)*binomial(6,1) states.
"1110001010" denotes that "1010" ("111000") belongs to 1st (2nd) space.
                        
#Examples:
```
julia> vnstate=[4,6];
julia> vnpart=[[(2,),(3,)]];
julia> NCBasis{UInt}(vnstate,vnpart).table.|>bitstring
120-element Vector{String}:
 "0000000000000000000000000000000000000000000000000000000001110011"
 "0000000000000000000000000000000000000000000000000000000010110011"
 "0000000000000000000000000000000000000000000000000000000100110011"
 ⋮
 "0000000000000000000000000000000000000000000000000000001101001100"
 "0000000000000000000000000000000000000000000000000000001110001100"
```
"""
struct NCBasis{I<:Unsigned} <: FockBasis{I}
    vnstate::Vector{Int}
    vnpart::Vector{Vector{Tuple{Vararg{Int}}}}
    table::Vector{I}
    sortflag::Bool
end
Base.kron(v1::Vector{NTuple{N,Int}},v2::Vector{NTuple{M,Int}}) where {N,M} = collect((v1[i]...,v2[j]...) for i=1:length(v1) for j=1:length(v2))
Base.kron(v1::Vector{T1},v2::Vector{T}) where {T<:Tuple{Vararg{Int}},T1<:Tuple{Vararg{Int}}} = collect((v1[i]...,v2[j]...) for i=1:length(v1) for j=1:length(v2))
Base.issorted(bs::NCBasis) = bs.sortflag

function NCBasis{I}(vnstate::Vector{Int},vnpart::Vector{Vector{T}}) where {I<:Unsigned,T<:Tuple{Vararg{Int}}}
    #@assert size(vnstate)==size(vnpart) "NCBasis error: size of vnstate ($(size(vnstate))) == size of vnpart ($(size(vnpart))) "
    table0=Vector{I}(undef,0)
    for k = 1:size(vnpart,1)
        @assert length(vnstate)==length(vnpart[k]) "NCBasis error: length of vnstate ($(length(vnstate))) == length of vnpart[$(k)] ($(length(vnpart[k]))) "                   
        @assert ~in(Tuple{}(),vnpart[k]) "NCBasis error: Tuple{}() in vnpart[$(k)]"
        vst=[0,cumsum(vnstate)...]
        temp = Vector{Vector{Tuple{Vararg{Int}}}}(undef,length(vnstate))
        for i=length(vnstate):-1:1 
            temp[i]=vcat(collect.( Combinations{vnpart[k][i][j]}( vst[i]+1:vst[i+1]) for j=1:length(vnpart[k][i] ) )... )
        end
        tabl=reduce(kron,temp)
        for j=1:length(tabl)
            push!(table0,bitgen(tabl[j],I))
        end
    end
    return NCBasis{I}(vnstate,vnpart,sort(table0),true)
end

NCBasis{I}() where I<:Unsigned = NCBasis{I}([0],[[(0,)]])

function Base.show(io::IO,bs::NCBasis)
    @printf io "%s:\n" repr(bs)
    for i=1:dimension(bs)
        @printf io "  %s\n" string(bs[i],base=2,pad=sum(bs.vnstate))
    end
end
Base.repr(bs::NCBasis)=@sprintf "%s(%s)" nameof(typeof(bs)) join(id(bs),",")

"""
    bitgen(st::NTuple{N,Int},::Type{I}) where {N,I<:Integer}
Get a integer by NTuple `st`.

#Examples:
```
julia> bitgen((0,3),UInt)|>bitstring
"0000000000000000000000000000000000000000000000000000000000000100"
julia> bitgen(Tuple{}(),UInt)|>bitstring
"0000000000000000000000000000000000000000000000000000000000000000"
```          
"""
function bitgen(st::NTuple{N,Int},::Type{I}) where {N,I<:Integer}
     eye,temp=one(I),zero(I)
    for j=1:length(st)
        temp=temp|eye<<(st[j]-1)
    end
    return temp
end
    
function dimension(bs::NCBasis)
        res0=0
        for k =1:size(bs.vnpart,1)
            res=1
            for i=1:length(bs.vnstate)
                res*=sum( binomial(bs.vnstate[i],bs.vnpart[k][i][j]) for j=1:length(bs.vnpart[k][i]) )
            end
            res0+=res
        end
        return res0
end
                                    
Base.getindex(bs::NCBasis,i::Int)=bs.table[i]

Base.searchsortedfirst(bs::NCBasis{I},i::I) where I<:Unsigned=findfirst(isequal(i),bs.table)
kind(::Type{<:NCBasis})=:nc
idtype(::Type{<:NCBasis}) = Tuple{  Vector{Int}, Vector{Vector{Tuple{Vararg{Int}}}} }
id(bs::NCBasis)=(bs.vnstate,bs.vnpart)


"""
   bkron(b1::I,n1::Int,b2::I,n2::Int) where I<:Integer -> I

Construct a state |b1>⊗|b2> by state |b1> and state |b2>. n1(n2) is the number of single-particle state |b1>(|b2>).

#Examples
```
julia> string(bkron(UInt(3),2,UInt(6),3),base=2,pad=5)
"11011"
```
"""
function bkron(b1::I,n1::Int,b2::I,n2::Int) where I<:Unsigned
            @assert n1+n2 <= length(bitstring(b1)) "bkron error: n1($(n1))+n2($(n2))<=length(bitstring(b1))($(length(bitstring(b1))))"
            b2<<n1|b1
        end
            
"""
   bkron(bs1::NCBasis{I},bs2::NCBasis{I}) where I<:Unsigned -> NCBasis{I}
   
Construct a basis {|b1>}⊗{|b2>} by basis {|b1>} and basis {|b2>}. If the `table`s of bs1 and bs2 are sorted, the `table` of result is sorted.
"""                        
function bkron(bs1::NCBasis{I},bs2::NCBasis{I}) where I<:Unsigned
            @assert size(bs2.vnstate,1)==size(bs1.vnstate,1) "bkron error: length of bs2.vnstate [$(length(bs2.vnstate))]== length of bs1.vnstate [$(length(bs1.vnstate))]"
            n2=length(bs2.vnpart)
            n1=length(bs1.vnpart)
            nvp=Vector{Vector{Tuple{Vararg{Int}}}}(undef,n1*n2)
            nvst=vcat(bs1.vnstate,bs2.vnstate)
            for  i= 1:n2
                for j=1:n1
                    nvp[j+(i-1)*n1]=vcat(bs1.vnpart[i],bs2.vnpart[j])
                end
            end
            table=collect(bkron(bs1.table[i],sum(bs1.vnstate),bs2.table[j],sum(bs2.vnstate)) for i=1:length(bs1.table) for j=1:length(bs2.table))
            #sortflag = bs1.sortflag && bs2.sortflag
            sortflag = false
            NCBasis{I}(nvst,nvp,table,sortflag)
        end

"""
   boplux(bs1::NCBasis{I},bs2::NCBasis{I}) where I<:Unsigned -> NCBasis{I}
   
Construct a basis {|b1>}⊕{|b2>} by basis {|b1>} and basis {|b2>}. The `table` is not sorted.
"""    
function boplus(bs1::NCBasis{I},bs2::NCBasis{I}) where I<:Unsigned
        @assert bs2.vnstate==bs1.vnstate "boplus error: bs2.vnstate [$(bs2.vnstate)]== bs1.vnstate [$(bs1.vnstate)]"
        nvst=bs1.vnstate
        nvp=vcat(bs1.vnpart,bs2.vnpart)
        table=vcat(bs1.table,bs2.table)
        NCBasis{I}(nvst,nvp,table,false)
end



"""
    matrix(operator::AbstractFockOperator,basisbra::FockBasis,basisket::FockBasis,table) -> SparseMatrixCSC{valtype(operator),Int}
    matrix(operator::AbstractFockOperator,basis::FockBasis,table) -> SparseMatrixCSC{valtype(operator),Int}
    matrix(opts::Operators,basisbra::FockBasis,basisket::FockBasis,table,dtype::Type{<:Number}=Complex{Float}) -> SparseMatrixCSC{dtype,Int}
    matrix(opts::Operators,basis::FockBasis,table,dtype::Type{<:Number}=Complex{Float}) -> SparseMatrixCSC{dtype,Int}

Get the CSC-formed sparse matrix representation of an operator / a set of operators.
"""
function matrix(operator::AbstractFockOperator,basisbra::FockBasis,basisket::FockBasis, table)
    #@assert rank(operator)%2==0 "matrix error: two bases are needed for odd ranked Fock operators."
    ndata,data,indices,indptr=0,zeros(valtype(operator),dimension(basisket)),zeros(Int,dimension(basisbra)),zeros(Int,dimension(basisket))
    eye,temp,flag=one(basisket|>eltype),zeros(basisket|>eltype,rank(operator)+1),true   
    seqs = NTuple{rank(operator), Int}(table[id(operator)[i]]-1 for i = 1:rank(operator))
    nambus = NTuple{rank(operator), Bool}(oid.index.nambu==CREATION for oid in id(operator))

    for i=1:dimension(basisket)
        flag=true 
        temp[1]=basisket[i]
        for j=1:rank(operator)
            (temp[j]&eye<<seqs[end+1-j]>0)==nambus[end+1-j] && (flag=false; break)
            temp[j+1]=nambus[end+1-j] ? temp[j]|eye<<seqs[end+1-j] : temp[j]&~(eye<<seqs[end+1-j])
        end
        if flag
            (ndata+=1)
            nsign=0
            isa(operator,FOperator) && for j=1:rank(operator) for k=0:seqs[end+1-j]-1 (temp[j]&eye<<k>0 && (nsign+=1)) end end
            findind = searchsortedfirst(basisbra,temp[end])
            if findind == nothing
                ndata -= 1
            else
                indices[ndata]=findind
                data[ndata]=(-1)^nsign*operator.value
                indptr[ndata]=i
            end
        end
        
    end
    if ndata==0
        return spzeros(valtype(operator),dimension(basisbra),dimension(basisket))
    else          
        return sparse(indices[1:ndata],indptr[1:ndata],data[1:ndata],dimension(basisbra),dimension(basisket))
    end
end

matrix(operator::AbstractFockOperator,basis::FockBasis,table) = matrix(operator,basis,basis,table)

function matrix(opts::Operators,basisbra::FockBasis,basisket::FockBasis,table,dtype::Type{<:Number}=Complex{Float})
    result=spzeros(dtype,dimension(basisbra),dimension(basisket))
    for opt in values(opts) result=result+matrix(opt,basisbra,basisket,table) end
    return result
end
matrix(opts::Operators,basis::FockBasis,table,dtype::Type{<:Number}=Complex{Float}) = matrix( opts,basis,basis,table,dtype)                        
                        
end #module