using DFTK
using Unitful
using UnitfulAtomic

d = 10u"Å"
a = 2.641u"Å"  # Graphene Lattice constant
lattice = [a  -a/2    0;
           0  √3*a/2  0;
           0     0    d]

C = ElementPsp(:C, psp=load_psp("hgh/pbe/c-q4"))
atoms     = [C, C]
positions = [[0.0, 0.0, 0.0], [1//3, 2//3, 0.0]]
model  = model_PBE(lattice, atoms, positions)
basis  = PlaneWaveBasis(model; Ecut=15, kgrid=[5, 5, 1])
nbandsalg = AdaptiveBands(basis.model; n_bands_converge=15)
scfres = self_consistent_field(basis; nbandsalg, tol=1e-5);

plot_bandstructure(scfres; kline_density=10)

using wannier90_jll  # Needed to make run_wannier90 available
run_wannier90(scfres;
              fileprefix="wannier/graphene",
              n_wannier=5,
              num_print_cycles=25,
              num_iter=200,
              #
              dis_win_max=19.0,
              dis_froz_max=0.1,
              dis_num_iter=300,
              dis_mix_ratio=1.0,
              #
              wannier_plot=true,
              wannier_plot_format="cube",
              wannier_plot_supercell=5,
              write_xyz=true,
              translate_home_cell=true,
             );

rm("wannier", recursive=true)