using QuadGK
using SpecialFunctions
using LinearAlgebra

#{{{1 set of G-vectors

# set of reciprocal lattice vectors limited by Gmax - 2D lattices
#
# output: array of G-vectors
# 
function generate_G_set_2D(ltype,Gmax,a)

  if ltype=="square"
    b1=[2*pi/a,0.0]
    b2=[0.0,2*pi/a]
    M=round(Int,ceil(Gmax*a/(2*pi)))
  end

  if ltype=="hexagonal"
    b1=4*pi/(sqrt(3)*a) * [ sqrt(3)/2, -1/2 ]
    b2=4*pi/(sqrt(3)*a) * [ sqrt(3)/2, +1/2 ]
    M=round(Int,ceil(Gmax*a/(2*pi)))
  end

  if ltype=="honeycomb"
    b1=4*pi/(3*a) * [ sqrt(3)/2, -1/2 ]
    b2=4*pi/(3*a) * [ sqrt(3)/2, +1/2 ]
    M=round(Int,ceil(Gmax*a*sqrt(3)/(2*pi)))
  end

  Gset=[]

  for m1 in -M:+M
    for m2 in -M:+M
      G=m1*b1+m2*b2
      if sum(G.^2)<=Gmax^2
        push!(Gset,G)
      end
    end
  end

  return Gset
end


# set of reciprocal lattice vectors limited by Gmax - 3D lattices
#
# output: array of G-vectors
# 
function generate_G_set_3D(ltype,Gmax,a,c=nothing)
  
  if ltype=="cubic"
    b1=[2*pi/a,0.0,0.0]
    b2=[0.0,2*pi/a,0.0]
    b3=[0.0,0.0,2*pi/a]
    M=round(Int,ceil(Gmax*a/(2*pi)))
    M3=M
  end

  if ltype=="tetragonal"
    b1=[2*pi/a,0.0,0.0]
    b2=[0.0,2*pi/a,0.0]
    b3=[0.0,0.0,2*pi/c]
    M=round(Int,ceil(Gmax*a/(2*pi)))
    M3=round(Int,ceil(Gmax*c/(2*pi)))
  end

  Gset=[]

  for m1 in -M:+M
    for m2 in -M:+M
      for m3 in -M3:+M3
        G=m1*b1+m2*b2+m3*b3
        if sum(G.^2)<=Gmax^2
          push!(Gset,G)
        end
      end
    end
  end

  return Gset
end	


# convert set of G vectors from an array of vectors into 2D array
#
# output: G components organized into D x NG array
#
function extract_G_vectors(Gset)

  G=zeros(size(Gset[1],1),size(Gset,1))

  for i in eachindex(Gset)
    G[:,i]=Gset[i]
  end

  return G
end

#}}}

#{{{1 volume of primitive cells

# evaluates volume of a primitive cell for the given lattice and lattice parameter(s)
#
function eval_VPC(ltype,a,c=nothing)

  if ltype=="square"
   VPC=a^2
  end

  if ltype=="hexagonal"
   VPC=sqrt(3)/2*a^2
  end

  if ltype=="honeycomb"
   VPC=3*sqrt(3)/2*a^2
  end

  if ltype=="cubic"
   VPC=a^3
  end

  if ltype=="tetragonal"
   VPC=a^2*c
  end

  return VPC
end

#}}}

#{{{1 BZ path

# conventional path through Brillouin zone to plot band structures
#
# output: array of k-vectors, array of path-length values
# 
function generate_BZ_path(ltype,N,a,c=nothing)

  if ltype=="square"
    G=[0.0,0.0]
    X=[pi/a,0.0]
    M=[pi/a,pi/a]
    kpts=[M,G,X,M]
    kticks_lbl=["M","Γ","X","M"]
  end

  if ltype=="hexagonal"
    G=[0.0,0.0]
    K=[4*pi/(3*a),0.0]
    M=[pi/a,pi/(a*sqrt(3))]
    kpts=[K,G,M,K]
    kticks_lbl=["K","Γ","M","K"]
  end

  if ltype=="honeycomb"
    G=[0.0,0.0]
    K=[4*pi/(3*sqrt(3)*a),0.0]
    M=[pi/(sqrt(3)*a),pi/(3*a)]
    kpts=[K,G,M,K]
    kticks_lbl=["K","Γ","M","K"]
  end

  if ltype=="cubic"
    G=[0.0,0.0,0.0]
    X=[pi/a,0.0,0.0]
    M=[pi/a,pi/a,0.0]
    R=[pi/a,pi/a,pi/a]
    kpts=[M,G,X,M,R,G]
    kticks_lbl=["M","Γ","X","M","R","Γ"]
  end

  if ltype=="tetragonal"
    G=[0.0,0.0,0.0]
    X=[pi/a,0.0,0.0]
    M=[pi/a,pi/a,0.0]
    A=[pi/a,pi/a,pi/c]
    kpts=[M,G,X,M,A,G]
    kticks_lbl=["M","Γ","X","M","A","Γ"]
  end

  Npts=size(kpts,1)

  kpath=collect(LinRange(kpts[1],kpts[2],N))

  for seg in 2:Npts-1
    aux=collect(LinRange(kpts[seg],kpts[seg+1],N))
    kpath=append!(kpath,aux)
  end

  Nk=size(kpath,1)
  kpath_len=zeros(Nk)
  for j in 2:Nk
    kpath_len[j]=kpath_len[j-1]+sqrt(sum((kpath[j]-kpath[j-1]).^2))
  end

  kticks_pos=zeros(Npts)
  for seg in 1:Npts-1
    kticks_pos[seg+1]=kticks_pos[seg]+sqrt(sum((kpts[seg+1]-kpts[seg]).^2))
  end

  return kpath,kpath_len,(kticks_pos,kticks_lbl)
end

#}}}

#{{{1 iBZ coverings

# cover the irreducible Brillouin zone with homogeneously distributed 
# Bloch vectors and determine their weights used when calculating BZ average
#
# output: array of k-points, array of weights
#
function generate_iBZ_covering(ltype,N,a,N3=nothing,c=nothing)

  kset=[]
  wght=[]

  if ltype=="square"
    for i1 in 0:N
      for i2 in 0:i1
        push!(kset,[i1,i2].*pi/(a*N))
    
        w=8
    
        if i1==0 && i2==0   w=1 end # Γ
        if i1==N && i2==0   w=2 end # X
        if i1==N && i2==N   w=1 end # M
    
        if i1==N && i2>0 && i2<N   w=4 end # Z
    
        if i1>0 && i1<N && i2==0    w=4 end # Δ
        if i1>0 && i1<N && i2==i1   w=4 end # Σ
          
        push!(wght,w)
      end
    end
    wght=wght/(4*N^2)
  end

  if ltype=="hexagonal" || ltype=="honeycomb"

    if ltype=="hexagonal" dk=4*pi/(3*a*N) end
    if ltype=="honeycomb" dk=4*pi/(3*sqrt(3)*a*N) end

    if N%2==1 println("generate_iBZ_covering: N has to be even for hexagonal/honeycomb lattice!") end
    
    for i1 in 0:N
      for i2 in 0:min(i1,N-i1)
        push!(kset,[i1+i2/2,i2].*dk)
    
        w=12
    
        if i1==0 && i2==0   w=1 end # Γ
        if i1==N && i2==0   w=2 end # K
    
        if i1==N÷2 && i2==i1   w=3 end # M
    
        if i1>0 && i1<N && i2==0   w=6 end # T
        
        if i1>0 && i1<N÷2 && i2==i1   w=6 end # Σ
    
        if i1>N÷2 && i1<N && i2==N-i1    w=6 end # T'
          
        push!(wght,w)
      end
    end
    wght=wght/(3*N^2)
  end

  if ltype=="cubic"
    for i1 in 0:N
      for i2 in 0:i1
        for i3 in 0:i2
          push!(kset,[i1,i2,i3].*pi/(a*N))
    
          w=48
    
          if i1==0 && i2==0 && i3==0   w=1 end # Γ
          if i1==N && i2==0 && i3==0   w=3 end # X
          if i1==N && i2==N && i3==0   w=3 end # M
          if i1==N && i2==N && i3==N   w=1 end # R
    
          if i1==N && i2>0 && i2<N && i3==0    w=12 end # Z
          if i1==N && i2>0 && i2<N && i3==i2   w=12 end # S
          if i1==N && i2==N && i3>0 && i3<N    w=6  end # T
          
          if i1>0 && i1<N && i2==0  && i3==0    w=6  end # Δ
          if i1>0 && i1<N && i2==i1 && i3==0    w=12 end # Σ
          if i1>0 && i1<N && i2==i1 && i3==i1   w=8  end # Λ
    
          if i1==N && i2>0 && i2<N && i3>0 && i3<i2    w=24 end # MXR=ZST face
          if i1>0 && i1<N && i2>0 && i2<i1 && i3==0    w=24 end # ΓXM=ΔZΣ face
          if i1>0 && i1<N && i2>0 && i2<i1 && i3==i2   w=24 end # ΓXR=ΔSΛ face
          if i1>0 && i1<N && i2==i1 && i3>0 && i3<i2   w=24 end # ΓMR=ΣTΛ face
    
          push!(wght,w)
        end
      end
    end
    wght=wght/(8*N^3)
  end

  if ltype=="tetragonal"
    for i1 in 0:N
      for i2 in 0:i1
        for i3 in 0:N3
          push!(kset,[i1*pi/(a*N),i2*pi/(a*N),i3*pi/(c*N3)])
    
          w=16
    
          if i1==0 && i2==0 && (i3==0 || i3==N3)   w=1 end # Γ,Z
          if i1==N && i2==0 && (i3==0 || i3==N3)   w=2 end # X,R
          if i1==N && i2==N && (i3==0 || i3==N3)   w=1 end # M,A
    
          if i1==N && i2>0 && i2<N && (i3==0 || i3==N3)    w=4 end # Y,T
          if i1>0 && i1<N && i2==0  && (i3==0 || i3==N3)   w=4 end # Δ,U
          if i1>0 && i1<N && i2==i1 && (i3==0 || i3==N3)   w=4 end # Σ,S
    
          if i1==0 && i2==0 && i3>0 && i3<N3   w=2 end # Λ
          if i1==N && i2==0 && i3>0 && i3<N3   w=4 end # W
          if i1==N && i2==N && i3>0 && i3<N3   w=2 end # V
    
          if i1==N && i2>0 && i2<N && i3>0 && i3<N3   w=8 end # MXRA face
    
          if i1>0 && i1<N && i2>0 && i2<i1 && (i3==0 || i3==N3)   w=8 end # ΓXM,ZRA faces
    
          if i1>0 && i1<N && i2==0  && i3>0 && i3<N3   w=8 end # ΓXRZ face
          if i1>0 && i1<N && i2==i1 && i3>0 && i3<N3   w=8 end # ΓMAZ face
    
          push!(wght,w)
        end
      end
    end
    wght=wght/(8*N^2*N3)
  end

  return kset,wght
end

#}}}

#{{{1 Fourier components of atomic potential

# evaluates Fourier components of symmetric 2D atomic potential at given G-vectors
#
# output: dictionary of |G|->VG pairs
#
function eval_atomic_VG_2D(Gset,VPC,V0,kappa,r0)

  Vat(r)=-V0*exp(-kappa*r)*r0/(r+r0)

  VG=Dict()

  for G1 in Gset
    for G2 in Gset
      Glen=sqrt(sum((G1-G2).^2))
      Glen_rnd=round(Glen,digits=6)
      if !haskey(VG,Glen_rnd)
        val,err=quadgk(r -> r*Vat(r)*besselj0(Glen*r), 0, Inf, rtol=1e-9)
        VG[Glen_rnd]=2*pi*val/VPC
      end
    end
  end

  return VG
end


# evaluates Fourier components of symmetric 3D atomic potential at given G-vectors
#
# output: dictionary of |G|->VG pairs
#
function eval_atomic_VG_3D(Gset,VPC,V0,kappa,r0)

  Vat(r)=-V0*exp(-kappa*r)*r0/(r+r0)

  VG=Dict()

  for G1 in Gset
    for G2 in Gset
      Glen=sqrt(sum((G1-G2).^2))
      Glen_rnd=round(Glen,digits=6)
      if !haskey(VG,Glen_rnd)
        val,err=quadgk(r -> r^2*Vat(r)*sin(Glen*r)/(Glen*r), 0, Inf, rtol=1e-9)
        VG[Glen_rnd]=4*pi*val/VPC
      end
    end
  end

  return VG
end

#}}}

#{{{1 matrices entering the central equation

# evaluates kinetic part of the Hamiltonian matrix entering the central equation
#
# output: NG x NG real matrix
#
function eval_Tmtx(k,Gset)

  h22m=3.809984e-2 # hbar^2/(2*me*1nm^2*1eV)

  NG=size(Gset,1)

  H=zeros(Float64,NG,NG)

  for i in 1:NG
    H[i,i]=h22m*sum((k+Gset[i]).^2)
  end

  return H
end

# evaluates potential part of the Hamiltonian matrix entering the central equation
#
# output: NG x NG complex matrix
#
function eval_Umtx(Gset,VG,dset,anisotropy=0.0)

  NG=size(Gset,1)
  Nd=size(dset,1)

  H=zeros(Complex{Float64},NG,NG)

  for i1 in 1:NG
    for i2 in 1:NG
      G=Gset[i1]-Gset[i2]
      Glen=sqrt(sum(G.^2))
      Glen_rnd=round(Glen,digits=6)

      sfac=0
      for j in 1:Nd
        sfac+=exp(1im*sum(G.*dset[j]))
      end

      H[i1,i2]=VG[Glen_rnd]*sfac

      if i1!=i2 H[i1,i2]*=(1-anisotropy*G[1]/Glen) end
    end
  end

  return H
end


# evaluates kinetic part of the Hamiltonian matrix entering the central equation
#
# output: NG x NG real matrix
#
function eval_Hmtx(k,Gset,VG,dset,anisotropy=0.0)

 return eval_Tmtx(k,Gset) + eval_Umtx(Gset,VG,dset,anisotropy)

end

#}}}

#{{{1 calculations of the band structure

# evaluates band dispersions for a set of Bloch vectors
#
# output: array (Nk x NG) of band energies
#
function eval_bands(kset,Gset,VG,dset=nothing,anisotropy=0.0)

  Nk=size(kset,1)
  NG=size(Gset,1)

  if dset==nothing
   D=size(Gset[1],1)
   Umtx=eval_Umtx(Gset,VG,[zeros(D)])
  else
   Umtx=eval_Umtx(Gset,VG,dset)
  end

  E=zeros(Nk,NG)
  for i in 1:Nk
    Tmtx=eval_Tmtx(kset[i],Gset)
    E[i,:]=eigvals(Tmtx+Umtx)
  end
  
  return E
end


# evaluates band dispersions on a rectangular grid covering [kmin,kmax]^2
#
# output: array (Nk x Nk x NG) of band energies
#
function eval_bands_grid(kmin,kmax,Nk,Gset,VG,dset=nothing,anisotropy=0.0)

  NG=size(Gset,1)

  if dset==nothing
   D=size(Gset[1],1)
   Umtx=eval_Umtx(Gset,VG,[zeros(D)])
  else
   Umtx=eval_Umtx(Gset,VG,dset)
  end

  kx=LinRange(kmin,kmax,Nk)
  ky=LinRange(kmin,kmax,Nk)
  E=zeros(Nk,Nk,NG)
  for i in 1:Nk
    for j in 1:Nk
      Tmtx=eval_Tmtx([kx[i],ky[j]],Gset)
      E[i,j,:]=eigvals(Tmtx+Umtx)
    end
  end
  
  return kx,ky,E
end

#}}}

#{{{1 calculations of wave functions

# evaluates periodic part of the Psi(r) wave function in xy-plane (z=0 implied in 3D case)
#
# output: complex values of sum_G Psi_G exp(iG.r) at given x,y grid
#
function eval_Psi(x,y,Gset,PsiG)

  Nx=size(x,1)
  Ny=size(y,1)

  Psi=zeros(Complex{Float64},Nx,Ny)

  for ix in 1:Nx
    for iy in 1:Ny
      for j in eachindex(Gset)
        G=Gset[j]
        Psi[ix,iy]+=PsiG[j]*exp(1im*(G[1]*x[ix]+G[2]*y[iy]))
      end
    end
  end

  return Psi
end

#}}}