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)

using Plots
using LinearAlgebra
nucleus = ElementGaussian(0.3, 10.0)
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(0.3, 10.0)
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)