using AtomsBuilder
using DFTK
using PseudoPotentialData
using Statistics
using Unitful
using UnitfulAtomic

a0 = 10.26u"bohr"  # Experimental lattice constant of silicon
a_list = a0 * range(0.98, 1.02, length=20)

function compute_ground_state_energy(a; Ecut, kgrid, kinetic_blowup, kwargs...)
    pseudopotentials = PseudoFamily("cp2k.nc.sr.pbe.v0_1.semicore.gth")
    model = model_DFT(bulk(:Si; a); functionals=PBE(), kinetic_blowup, pseudopotentials)
    basis = PlaneWaveBasis(model; Ecut, kgrid)
    self_consistent_field(basis; callback=identity, kwargs...).energies.total
end

Ecut  = 5          # Very low Ecut to display big irregularities
kgrid = (4, 4, 4)  # Very sparse k-grid to speed up convergence
E0_naive = compute_ground_state_energy.(a_list; kinetic_blowup=BlowupIdentity(), Ecut, kgrid);

E0_ref = [-7.867399262300442, -7.867875504884598, -7.86831005961699,
          -7.868703712435519, -7.869057235894591, -7.869371393835255,
          -7.869646937992838, -7.869884611383302, -7.870085144568824,
          -7.8702492593753135, -7.870377668337734, -7.870471072484817,
          -7.870530163168231, -7.870555622590523, -7.870548125081617,
          -7.870508333380078, -7.8704369010295006, -7.870334474350743,
          -7.870201688247285, -7.870039170777946]

using Plots
shift = mean(E0_naive - E0_ref)
p = plot(a_list, E0_naive .- shift, label="Ecut=5", xlabel="lattice parameter a",
         ylabel="Ground state energy (Ha)", color=1, legend=:topright)
plot!(p, a_list, E0_ref, label="Ecut=100", color=2)

E0_modified = compute_ground_state_energy.(a_list; kinetic_blowup=BlowupCHV(), Ecut, kgrid);

estimate_a0(E0_values) = a_list[findmin(E0_values)[2]]
a0_naive, a0_ref, a0_modified = estimate_a0.([E0_naive, E0_ref, E0_modified])

shift = mean(E0_modified - E0_ref)  # Shift for legibility of the plot
plot!(p, a_list, E0_modified .- shift, label="Ecut=5 + BlowupCHV", color=3)
vline!(p, [a0], label="experimental a0", linestyle=:dash, linecolor=:black)
vline!(p, [a0_naive], label="a0 Ecut=5", linestyle=:dash, color=1)
vline!(p, [a0_ref], label="a0 Ecut=100", linestyle=:dash, color=2)
vline!(p, [a0_modified], label="a0 Ecut=5 + BlowupCHV", linestyle=:dash, color=3)

println("Error of approximation of the reference a0 with modified kinetic term:"*
        " $(round((a0_modified - a0_ref)*100/a0_ref, digits=5))%")