using DFTK
using LinearAlgebra

function aluminium_setup(repeat=1; Ecut=7.0, kgrid=[2, 2, 2])
    a = 7.65339
    lattice = a * Matrix(I, 3, 3)
    Al = ElementPsp(:Al, psp=load_psp("hgh/lda/al-q3"))
    atoms     = [Al, Al, Al, Al]
    positions = [[0.0, 0.0, 0.0], [0.0, 0.5, 0.5], [0.5, 0.0, 0.5], [0.5, 0.5, 0.0]]

    # Make supercell in ASE:
    # We convert our lattice to the conventions used in ASE
    # and then back ...
    supercell = ase_atoms(lattice, atoms, positions) * (repeat, 1, 1)
    lattice   = load_lattice(supercell)
    positions = load_positions(supercell)
    atoms = fill(Al, length(positions))

    # Construct an LDA model and discretise
    # Note: We disable symmetries explicitly here. Otherwise the problem sizes
    #       we are able to run on the CI are too simple to observe the numerical
    #       instabilities we want to trigger here.
    model = model_LDA(lattice, atoms, positions; temperature=1e-3, symmetries=false)
    PlaneWaveBasis(model; Ecut, kgrid)
end;

setup = aluminium_setup(5)
ase_atoms(setup.model).write("al_supercell.png")

is_converged = DFTK.ScfConvergenceDensity(1e-4)  # Flag convergence based on density
self_consistent_field(aluminium_setup(1); is_converged);

self_consistent_field(aluminium_setup(2); is_converged);

self_consistent_field(aluminium_setup(4); is_converged);

self_consistent_field(aluminium_setup(1); is_converged, mixing=SimpleMixing());

self_consistent_field(aluminium_setup(4); is_converged, mixing=SimpleMixing());