using DFTK
using PyCall

silicon = pyimport("ase.build").bulk("Si")
atoms = load_atoms(silicon)
atoms = [ElementPsp(el.symbol, psp=load_psp(el.symbol, functional="lda")) => position
         for (el, position) in atoms]
lattice = load_lattice(silicon);

model = model_LDA(lattice, atoms)
kgrid = [3, 3, 3]  # k-point grid
Ecut = 5           # kinetic energy cutoff in Hartree
basis = PlaneWaveBasis(model, Ecut; kgrid=kgrid);

using Plots
p = plot(yaxis=:log)
density_differences = Float64[];

using LinearAlgebra

function plot_callback(info)
    if info.stage == :finalize
        plot!(p, density_differences, label="|ρout - ρin|", markershape=:x)
    else
        push!(density_differences, norm(info.ρout - info.ρin))
    end
    info
end
callback = DFTK.ScfDefaultCallback() ∘ plot_callback;

scfres = self_consistent_field(basis, tol=1e-8, callback=callback);

p