using QuadGK
using SpecialFunctions
using Plots
pythonplot()


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

  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


function eval_atomic_VG_2D(Gset,VPC,V0,kappa,r0)

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

  VG=Dict()

  for G in Gset
    Glen=sqrt(sum(G.^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

  return VG
end


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


function eval_Fourier_series_2D(x,y,Gset,VG)

  val=0

  for G in Gset
    Glen=sqrt(sum(G.^2))
    Glen_rnd=round(Glen,digits=6)
    val+=VG[Glen_rnd]*cos(G[1]*x+G[2]*y)
  end

  return val
end


function eval_map_2D(x,y,Gset,VG)

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

  V=zeros(Nx,Ny)

  for i in 1:Nx
    for j in 1:Ny
      V[i,j]=eval_Fourier_series_2D(x[i],y[j],Gset,VG)
    end
  end

  return V
end


a=1.0
Gmax=50.0

Gset=generate_G_set_2D("hexagonal",Gmax,a)

G=extract_G_vectors(Gset)
plt1=scatter(G[1,:],G[2,:])

V0=5      # eV
kappa=0.5 # 1/nm
r0=0.5    # nm

VPC=a^2

VG=eval_atomic_VG_2D(Gset,VPC,V0,kappa,r0)

plt2=scatter(collect(keys(VG)),collect(values(VG)),ylimits=(-0.25,0.05))

N=201
x=LinRange(-2,+2,N)
y=LinRange(-2,+2,N)
Vmap=eval_map_2D(x,y,Gset,VG)
plt3=heatmap(x,y,Vmap')
plt4=plot(x,Vmap[:,N÷2+1])

plot(plt1,plt2,plt3,plt4,layout=(2,2))