using DFTK
using Statistics

a0 = 10.26  # Experimental lattice constant of silicon in bohr
a_list = range(a0 - 1/2, a0 + 1/2; length=20)

function compute_ground_state_energy(a; Ecut, kgrid, kinetic_blowup, kwargs...)
    lattice = a / 2 * [[0 1 1.];
                       [1 0 1.];
                       [1 1 0.]]
    Si = ElementPsp(:Si, psp=load_psp("hgh/lda/Si-q4"))
    atoms = [Si, Si]
    positions = [ones(3)/8, -ones(3)/8]
    model = model_PBE(lattice, atoms, positions; kinetic_blowup)
    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 = (2, 2, 2)  # Very sparse k-grid to speed up convergence
E0_naive = compute_ground_state_energy.(a_list; kinetic_blowup=BlowupIdentity(), Ecut, kgrid);

E0_ref = [-7.839775223322127, -7.843031658146996, -7.845961005280923,
          -7.848576991754026, -7.850892888614151, -7.852921532056932,
          -7.854675317792186, -7.85616622262217,  -7.85740584131599,
          -7.858405359984107, -7.859175611288143, -7.859727053496513,
          -7.860069804791132, -7.860213631865354, -7.8601679947736915,
          -7.859942011410533, -7.859544518721661, -7.858984032385052,
          -7.858268793303855, -7.857406769423708]

using Plots
shift = mean(abs.(E0_naive .- E0_ref))
p = plot(a_list, E0_naive .- shift, label="Ecut=5", xlabel="lattice parameter a (bohr)",
         ylabel="Ground state energy (Ha)", color=1)
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(abs.(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))%")