using DFTK
using LinearAlgebra
using JLD2

d = 2.079  # oxygen-oxygen bondlength
a = 9.0    # size of the simulation box
lattice = diagm(a * ones(3))
O = ElementPsp(:O, psp=load_psp("hgh/pbe/O-q6.hgh"))
atoms = [O => d / 2a * [[0, 0, 1], [0, 0, -1]]]
magnetic_moments = [O => [1., 1.]]

Ecut  = 10  # Far too small to be converged
model = model_PBE(lattice, atoms, temperature=0.02, smearing=smearing=Smearing.Gaussian(),
                  magnetic_moments=magnetic_moments)
basis = PlaneWaveBasis(model, Ecut; kgrid=[1, 1, 1])

ρspin  = guess_spin_density(basis, magnetic_moments)
scfres = self_consistent_field(basis, tol=1e-2, ρspin=ρspin)
save_scfres("scfres.jld2", scfres);

scfres.energies

using DFTK
using JLD2
loaded = load_scfres("scfres.jld2")
propertynames(loaded)

loaded.energies

callback = DFTK.ScfSaveCheckpoints()
scfres = self_consistent_field(basis, tol=1e-2, ρspin=ρspin, callback=callback);

callback = DFTK.ScfDefaultCallback() ∘ DFTK.ScfSaveCheckpoints(keep=true)
scfres = self_consistent_field(basis, tol=1e-2, ρspin=ρspin, callback=callback);

oldstate = load_scfres("dftk_scf_checkpoint.jld2")
scfres   = self_consistent_field(oldstate.basis, ρ=oldstate.ρ, ρspin=oldstate.ρspin,
                                 ψ=oldstate.ψ, tol=1e-3);

rm("dftk_scf_checkpoint.jld2")
rm("scfres.jld2")