using DFTK

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

model = Model(lattice; n_electrons=0, 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=30)

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 => [[0.2, 0, 0], [0.8, 0, 0]]]

# Assemble the model, discretize and build the Hamiltonian
model = Model(lattice; atoms=atoms, 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=15)