using LinearAlgebra
using Plots

# Parameters
a = 100     # Lattice constant
Ecut = 300  # Cutoff in Hartree
kgrid = 50  # Number of points on equally-spaced grid for k

# Derived quantities
b = 2π / a  # Reciprocal lattice

# Step 1: k-Points
kpoints = b * (collect(1:kgrid) .- ceil(Int, kgrid / 2)) ./ kgrid

# Step 2: Basis for H_k
#         Represented as one array of all valid G*b per kpoint
Gmax = ceil(Int, sqrt(2Ecut) + b)  # Rough upper bound for G
Gs = [[Gidx*b for Gidx in -Gmax:Gmax if abs2(Gidx*b + k) ≤ 2Ecut]
      for k in kpoints]

# Step 3: Build the discretised Hk. In this case it is diagonal,
#         i.e. its diagonal values (== eigenvalues) are all we need.
#         We directly determine them and sort them ascendingly
ev_Hk = [sort([abs2(G + k)/2 for G in Gs[ik]])
         for (ik, k) in enumerate(kpoints)]

# Step 4: Plot the bands
n_bands = 6
bands = [[ev_Hk[ik][iband] for ik in 1:length(kpoints)]
         for iband in 1:n_bands]

p = plot()
for iband in 1:n_bands
    plot!(p, kpoints, bands[iband], color=:blue, label="")
end
p

using DFTK

lattice = zeros(3, 3)
lattice[1, 1] = 20.;

model = Model(lattice; terms=[Kinetic()])

basis = PlaneWaveBasis(model; Ecut=300, kgrid=(1, 1, 1))

using Unitful
using UnitfulAtomic
using Plots

plot_bandstructure(basis; n_bands=6, kline_density=100)

struct ElementGaussian <: DFTK.Element
    α  # Prefactor
    L  # Extend
end

# Some default values
ElementGaussian() = ElementGaussian(0.3, 10.0)

# Real-space representation of a Gaussian
function DFTK.local_potential_real(el::ElementGaussian, r::Real)
    -el.α / (√(2π) * el.L) * exp(- (r / el.L)^2 / 2)
end

# Fourier-space representation of the Gaussian
function DFTK.local_potential_fourier(el::ElementGaussian, q::Real)
    # = ∫ -α exp(-(r/L)^2 exp(-ir⋅q) dr
    -el.α * exp(- (q * el.L)^2 / 2)
end

using Plots
using LinearAlgebra
nucleus = ElementGaussian()
plot(r -> DFTK.local_potential_real(nucleus, norm(r)), xlims=(-50, 50))

using LinearAlgebra

# Define the 1D lattice [0, 100]
lattice = diagm([100., 0, 0])

# Place them at 20 and 80 in *fractional coordinates*,
# that is 0.2 and 0.8, since the lattice is 100 wide.
nucleus   = ElementGaussian()
atoms     = [nucleus, nucleus]
positions = [[0.2, 0, 0], [0.8, 0, 0]]

# Assemble the model, discretize and build the Hamiltonian
model = Model(lattice, atoms, positions; terms=[Kinetic(), AtomicLocal()])
basis = PlaneWaveBasis(model; Ecut=300, kgrid=(1, 1, 1));
ham   = Hamiltonian(basis)

# Extract the total potential term of the Hamiltonian and plot it
potential = DFTK.total_local_potential(ham)[:, 1, 1]
rvecs = collect(r_vectors_cart(basis))[:, 1, 1]  # slice along the x axis
x = [r[1] for r in rvecs]                        # only keep the x coordinate
plot(x, potential, label="", xlabel="x", ylabel="V(x)")

using Unitful
using UnitfulAtomic

plot_bandstructure(basis; n_bands=6, kline_density=500)