using DFTK
using Plots
using PseudoPotentialData
using Unitful
using UnitfulAtomic

a = 3.01794  # Bohr
b = 5.22722  # Bohr
c = 9.77362  # Bohr
lattice = [[-a -a  0]; [-b  b  0]; [0   0 -c]]

pseudopotentials = PseudoFamily("cp2k.nc.sr.pbe.v0_1.largecore.gth")
Mg = ElementPsp(:Mg, pseudopotentials)
atoms     = [Mg, Mg]
positions = [[2/3, 1/3, 1/4], [1/3, 2/3, 3/4]];

kspacing = 0.945 / u"angstrom"        # Minimal spacing of k-points,
#                                      in units of wavevectors (inverse Bohrs)
Ecut = 5                              # Kinetic energy cutoff in Hartree
temperature = 0.01                    # Smearing temperature in Hartree
smearing = DFTK.Smearing.FermiDirac() # Smearing method
#                                      also supported: Gaussian,
#                                      MarzariVanderbilt,
#                                      and MethfesselPaxton(order)

model = model_DFT(lattice, atoms, positions;
                  functionals=[:gga_x_pbe, :gga_c_pbe], temperature, smearing)
kgrid = kgrid_from_maximal_spacing(lattice, kspacing)
basis = PlaneWaveBasis(model; Ecut, kgrid);

scfres = self_consistent_field(basis, damping=0.8, mixing=KerkerMixing());

scfres.occupation[1]

scfres.energies

kgrid_dos = kgrid_from_maximal_spacing(lattice, 0.7 / u"Å")
bands = compute_bands(scfres, kgrid_dos)
plot_dos(bands)