module MakeGraph
using OrderedCollections

export ladder, loopNN, loopNNN, square2D

#graph=Vector{Vector{Tuple{Int64,Int64}}}()

# example of graph
# [(1,2),(2,3),(3,4),(4,1)]
# [(1,4),(2,3)]

order(x::Tuple{Int64,Int64})=(x[1]<x[2] ? x : (x[2],x[1]))  # orders a bond 

function Eliminate_dupplicates(x::Vector{Tuple{Int64,Int64}}) # eliminate duplicate bonds in a subgraph
    @. x.=order(x)
    return collect(OrderedSet(x))
end

function ladder(m::Int)         # ladder = 2 identical loops + rungs
    graph=Vector{Vector{Tuple{Int64,Int64}}}()
    loops=vcat([(i,mod(i,m)+1) for i=1:m],[(i,mod(i,m)+m+1) for i=m+1:2m])
    rungs=[(i,i+m) for i=1:m]
    push!(graph,loops)
    push!(graph,rungs)
end

function loopNN(m::Int)    # loop with nn exchange
    graph=Vector{Vector{Tuple{Int64,Int64}}}()
    loops=[(i,mod(i,m)+1) for i=1:m]
    push!(graph,loops)
end

function loopNNN(m::Int)    # loop with nn and nnn exchange
    graph=Vector{Vector{Tuple{Int64,Int64}}}()
    loops=[(i,mod(i,m)+1) for i=1:m]
    push!(graph,loops)
    loops=Eliminate_dupplicates([(i,mod(i+1,m)+1) for i=1:m])
    push!(graph,loops)
end

function square2D(mx::Int, my::Int)
    graph=Vector{Vector{Tuple{Int64,Int64}}}()
    loops = Vector{Tuple{Int64,Int64}}()
    for j in 0:(my-1)
        loops = vcat(loops,[(i,mod(i,mx)+j*mx+1) for i=(j*mx+1):(j+1)*mx])
    end
    for j in 1:mx
        loops = vcat(loops,[(i,mod(i+mx-1,mx*my)+1) for i=j:mx:(my*(mx-1)+j)])
    end
    push!(graph,Eliminate_dupplicates(loops))
end

end  #MakeGraph


module Operators
using LinearAlgebra

export meanv, SuscI, SuscR, ElementsSelect, WeightedElements, Z, FeeEnergy, Ene, SpecificHeat, Entropy

function Bweights!(eig::Vector{U}, T::Real; Norm::Bool = true) where U <:Real
    beta=1/T
    if Norm
        eig.-=eig[1]
    end
    eig.*=-beta 
    @. eig.=exp(eig)
    x=sum(eig)
    Norm ? eig ./=x : eig
end

function Bweights(eig::Vector{U}, T::Real; Norm::Bool = true) where U <:Real
    beta=1/T
    Norm ? y=eig.-eig[1] : y = eig
    y.*=-beta 
    @. y.=exp(y)
    x=sum(y)
    Norm ? y ./=x : y
end

function Z(eig::Vector, T::Real)
    beta=1/T
    y=-beta*eig
    @. y.=exp(y)
    return sum(y)
end   

FeeEnergy(eig::Vector, T::Real)=-T*log(Z(eig::Vector, T::Real))

Ene(eig::Vector, T::Real)=dot(eig,exp.(-1/T*eig))/Z(eig,T)

Ene2(eig::Vector, T::Real)=dot(eig.^2,exp.(-1/T*eig))/Z(eig,T)

SpecificHeat(eig::Vector, T::Real)=(Ene2(eig,T)-Ene(eig,T)^2)/T^2

Entropy(eig::Vector, T::Real)=(Ene(eig,T)-FeeEnergy(eig,T))/T

function rho(H::Matrix, T::Real; Norm::Bool = true)
    local F=eigen(H)
    Bweights!(F.values, T, Norm=Norm)
    return F.vectors*Diagonal(F.values)*F.vectors'
end

function rhotol(H::Matrix, T::Real; Norm::Bool = true, tol::Real = 1e-6)
    local F=eigen(H)
    Bweights!(F.values, T, Norm=Norm)
    i=1
    while F.values[i]>tol && i<Bsize
        i +=1
    end
    return view(F.vectors,:,1:i)*Diagonal(view(F.values,1:i))*(view(F.vectors,:,1:i))'
end

function meanv(H::Matrix, T::Real, O::Matrix...; Norm::Bool = true)
    local r=transpose(rho(H,T,Norm=Norm))
    x=zeros(length(O))
    for i in 1:length(O)
        x[i]=sum(O[i].*r)    # Tr(AB) = sum (A_ij * B^T_ij) -- elementwise product
    end
    return x
end

ILor(w::Real,x::Real,G::Real)=G/(2*pi)*1/((w-x)^2+G^2)-G/(2*pi)*1/((w+x)^2+G^2)   # imag 1/(w-E) -1/(w+E)
RLor(w::Real,x::Real,G::Real)=(w-x)/((w-x)^2+G^2)-(w+x)/((w+x)^2+G^2)      # real 1/(w-E) -1/(w+E)


function ElementsSelect(O::Matrix, U::Matrix, eig::Vector; Max::Int=100000, tol::Real = 1e-6)
   elements=Vector()
    Max=min(Max,size(U)[2])
    A=U'*O*U[:,1:Max]
    @. A=A.^2                        #     |<2|O|1>|^2     real algebra assumed
    i2s=CartesianIndices(A)
    for i in 1:length(A)
        if (A[i] > tol)
            de=eig[i2s[i][1]]-eig[i2s[i][2]]                #   (E2-E1, <1|O|2><2|O|1>)
            push!(elements,(i2s[i][1],i2s[i][2],A[i]))
        end
    end    
    return elements
end

function WeightedElements(elements::Vector, eig::Vector, T::Real; Norm::Bool = true, tol::Float64=1e-6)
     weights=Bweights(eig,T,Norm=Norm)
     Welements=[(0.,0.)]
     static = 0.
     for i in 1:length(elements)
        x=weights[elements[i][2]]*elements[i][3]
        if x > tol
            E=eig[elements[i][1]]-eig[elements[i][2]]
            if abs(E) > tol
                push!(Welements,(E,x))      #   (E2-E1, <1|O|2><2|O|1>exp(-beta*E1)/Z )       
            else
                static +=x                  #     sum_1,2 <1|O|2><2|O|1>exp(-beta*E1)/Z for E1=E2
            end                             #     non-causal (thermal) contribution to static chi, this part does not appear in KK relations
        end
    end
    return Welements,static
end
    
function SuscI(elements::Vector, w::Real, G::Real)      # sum_1,2 <1|O|2><2|O|1>exp(-beta*E1)/Z imag[1/(w+iG-E)-1/(w+iG+E)]
    sum(map(x->x[2]*ILor(w,x[1],G),elements))            #   E=E2-E1
end

function SuscR(elements::Vector, w::Real, G::Real)     # sum_1,2 <1|O|2><2|O|1>exp(-beta*E1)/Z real[1/(w+iG-E)-1/(w+iG+E)]
    sum(map(x->-x[2]*RLor(w,x[1],G),elements))            #   E=E2-E1
end

end #Operators    

module BinaryBasisFermions
using LinearAlgebra
export BaseGen, H1gen, H2gen, occbin, szbin

function bits2int(bits::BitVector,N::Int)          #  dim-N BitVector -> Int
    base_2=(2 .^(N-1:-1:0))
    dot(bits,base_2)+1   # bits -> Fock index
end

function int2bits(n::Int, N::Int)              # Int -> dim-N BitVector 
    y=string(n-1, base = 2)
    parse.(Bool,collect(lpad(y,N,"0")))
end

function c!(flavor::Int, state::BitVector)      # c(flavor) in bit basis, updates state, output is sign
    l=length(state)
    if state[flavor]
        state[flavor] = false
        isodd(sum(view(state,flavor:l))) ? sgn = -1 : sgn = 1
    else
        sgn = 0
    end
    return sgn
end

function cdag!(flavor::Int, state::BitVector)   # cdag(flavor) in bit basis, updates state, output is sign
    l=length(state)
    if !state[flavor]
        state[flavor] = true
        isodd(sum(view(state,flavor:l))) ? sgn = 1 : sgn = -1
    else
        sgn = 0
    end
    return sgn
end

occbin(state::BitVector) = sum(state)   # occupation number of a bit state

function szbin(state::BitVector)      # sz(Int) of a bit state
    m=div(length(state),2)
    return 2*sum(view(state,1:m))-sum(state)
end


function BaseGen(N::Int, n::Int)          # generates a basis in n-particle subspace of 2^N Fock space
    base = Vector{BitVector}()
    hash = Dict{BitVector,Int}()
    j=0
    for i in 1:2^N
        state = int2bits(i,N)
        if (sum(state)==n)
            j +=1
            push!(base,state)
            merge!(hash,Dict(state=>j))    
        end
    end
    return base, hash
end

function BaseGen(N::Int, func::Function, val::Int)   # generates a basis of func(state)== val subspace of 2^N Fock space       
    base = Vector{BitVector}()
    hash = Dict{BitVector,Int}()
    j=0
    for i in 1:2^N
        state = int2bits(i,N)
        if (func(state)==val)
            j +=1
            push!(base,state)
            merge!(hash,Dict(state=>j))    
        end
    end
    return base, hash
end

function BaseGen(N::Int, conds::Tuple{Function,Int}...)  # as above with multiple conditions
    base = Vector{BitVector}()
    hash = Dict{BitVector,Int}()
    j=0
    for i in 1:2^N
        state = int2bits(i,N)
        acc = true
        for cond in conds
            acc &=(cond[1](state)==cond[2])
        end
        if acc
            j +=1
            push!(base,state)
            merge!(hash,Dict(state=>j))    
        end
    end
    return base, hash
end

#generates a matrix of 1P operator in BitVector basis "base"; "index" is the distionary state => index
function H1gen(elements::Vector{Tuple{Int64, Int64, T}},base::Vector{BitVector},index::Dict{BitVector, Int64}) where T <:Number
    bsize=length(base)
    if length(index)!=bsize  
        return "incosistent basis map" 
    end
    H1=zeros(T,bsize,bsize)
    for i in 1:bsize
        s=base[i]
        for x in elements
           if s[x[2]] && (!s[x[1]] || x[1]==x[2])
                state=copy(s)
                sgn = c!(x[2],state)
                sgn2= cdag!(x[1],state)
                sgn *=sgn2 
                j=index[state]
                H1[j,i] +=sgn*x[3]
            end
        end
    end
    return H1
end

#generates a matrix of 2P operator in BitVector basis "base"; "index" is the distionary state => index
function H2gen(elements::Vector{Tuple{Int, Int, Int, Int, T}},base::Vector{BitVector},index::Dict{BitVector, Int64}) where T <:Number
    bsize=length(base)
    H2=zeros(bsize,bsize)
    for i in 1:bsize
        s=base[i]
        for x in elements
            if s[x[4]]&s[x[3]]
                state=copy(s)
                sgn = 1
                sgn *=c!(x[4],state)
                sgn *=c!(x[3],state)
                sgn *=cdag!(x[2],state)
                sgn *=cdag!(x[1],state)
                if sgn != 0
                    j = index[state]
                    H2[j,i] += sgn*x[5]
                end
            end
        end
    end
    return H2
end

end

module CouplingConstants
export Graph2Hop,CrystalField,Sz,Sx,iSy,ChainU,SlaterKanamori,Elements2Tensor,Tensor2Elements,ElementsTransform

#############################################
# 1P (bilinear) operators 
function Graph2Hop(graph::Vector{Vector{Tuple{Int,Int}}}, m::Int, ts::Real...)  # generates spin-diag. hopping on a given graph
    x=Vector{Tuple{Int,Int,Real}}()
    i=0
    for sub in graph
        i +=1
        if i <= length(ts) 
            for bond in sub
                push!(x,(bond[1],bond[2],ts[i]))
                push!(x,(bond[1]+m,bond[2]+m,ts[i]))
                push!(x,(bond[2],bond[1],ts[i]))
                push!(x,(bond[2]+m,bond[1]+m,ts[i]))
            end
        end
    end
    return x
end

function CrystalField(x::Real...) #  crystal field (spin-indepent orbital-diagonal operator)
   l=length(x)
   vcat([(i,i,Float64(x[i])) for i=1:l],[(i+l,i+l,Float64(x[i])) for i=1:l])
end

function Sz(m::Int) # total Sz: up 1,...,m  down m+1,...,2m
   vcat([(i,i,Float64(1/2)) for i=1:m],[(i+m,i+m,Float64(-1/2)) for i=1:m]) # total spin
end

function Sx(m::Int) # total Sz: up 1,...,m  down m+1,...,2m
   vcat([(i,i+m,Float64(1/2)) for i=1:m],[(i+m,i,Float64(1/2)) for i=1:m]) # total spin
end

function iSy(m::Int) # total Sz: up 1,...,m  down m+1,...,2m
   vcat([(i,i+m,Float64(1/2)) for i=1:m],[(i+m,i,Float64(-1/2)) for i=1:m])
end

Sz(m::Int,i::Int)=[(i,i,Float64(1/2)),(i+m,i+m,Float64(-1/2))]

Sx(m::Int,i::Int)=[(i,i+m,Float64(1/2)),(i+m,i,Float64(1/2))]

iSyi(m::Int,i::Int)=[(i,i+m,Float64(1/2)),(i+m,i,Float64(-1/2))]

#function Matrix2Elements(h::Matrix{T}; tol::Real = 1e-6) where T <:Number   # extracts non-zero (>1e-6) elements from a matrix
#  tmp=findall(x->abs(x)>tol,h)
#    map(x->(x[1],x[2],h[x]),tmp)
#end

#function Elements2Matrix(elements::Vector{Tuple{Int64, Int64, T}}, m::Int) where T <:Number # generate 1P matrix from set of elements
#    x=zeros(T,m,m)
#    for y in elements
#        x[y[1],y[2]]=y[3]
#    end
#    return x
#end

################################################
# 2P (biquadratic) operators 
function ChainU(m::Int,U::Real) # local interaction on m-sites (1-orbital) chain
    [(i,i+m,i+m,i,U) for  i in  1:m]
end

function ChainU(m::Int,U::Vector{Real}) # local U_m interaction on m-sites (1-orbital) chain
    if length(U) != m
        return "incosistent number of Us"
    else
        return [(i,i+m,i+m,i,U[i]) for  i in  1:m]
    end
end

function SlaterKanamori(m::Int, U::Real, J::Real; g::Real = 1) # Slate-Kanamori interaction for m orbtials
    elements=Vector{Tuple{Int, Int, Int, Int, Float64}}()
    append!(elements,[(i,i+m,i+m,i,U) for i in 1:m])   # Unn
    for i in 1:m
        append!(elements,[(i,j+m,j+m,i,U-2*J) for j in 1:m if i!=j]) # (U-2J)n_up n_dn
    end
    for i in 1:m
        append!(elements,[(i,j,j,i,U-3*J) for j in i+1:m])   # (U-3J)n_up n_up
    end
    for i in 1:m    
        append!(elements,[(i+m,j+m,j+m,i+m,U-3*J) for j in i+1:m]) # (U-3J)n_dn n_dn
    end
    for i in 1:m
     append!(elements,[(j,i+m,j+m,i,J) for j in i+1:m])   # J a*_up b*_dn a_dn b_up
    end
    for i in 1:m
     append!(elements,[(j+m,i,j,i+m,J) for j in i+1:m])   # J a*_dn b*_up a_up b_dn
    end
    for i in 1:m
     append!(elements,[(j,j+m,i+m,i,g*J) for j in 1:m if i!=j])   # J a*_up a*_dn b_dn b_up
    end
    return elements
end

##############################################
# Basis transformations
function Elements2Tensor(elements::Vector, dim::Int)  # builds a tensor out of non-zero elements
    T    = typeof(last(elements[1]))
    rank = length(elements[1])-1
    form = ntuple(x->dim,rank)
    tensor = zeros(T,form)
    for x in elements
       tensor[CartesianIndex(x[1:rank])]=last(x)
    end
   return tensor
end

function Tensor2Elements(tensor::Array{T}; tol::Real = 1e-6) where T <:Number  # builds a list on non-zero elements
    tmp=findall(x->abs(x)>tol,tensor)
    map(i->(Tuple(i)...,tensor[i]), tmp)
end
 
using TensorOperations 
function transform(A::Array{T,4}, U::Matrix) where T <:Number  # transformation of 2P (4-index) tensors
    CU = conj.(U)
    @tensor D[a,b,c,d]:= CU[e,a]*CU[f,b]*U[g,c]*U[h,d]*A[e,f,g,h]
end

function transform(A::Array{T,2}, U::Matrix) where T <:Number   # transformation of 1P (2-index) tensors
     U'*A*U
end

# object of type cdag*c or cdag*cdag*c*c is expected
function ElementsTransform(elements::Vector, U::Matrix)  # transformation of 1P and 2P objects elements -> elements
    dim = size(U)[1]
    A=Elements2Tensor(elements, dim)
    D=transform(A,U)
    Tensor2Elements(D)
end

end