using DFTK
using LinearAlgebra

a = 10.
lattice = a * I(3)  # cube of ``a`` bohrs
He = ElementPsp(:He, psp=load_psp("hgh/lda/He-q2"))
atoms = [He => [[1/2; 1/2; 1/2]]]  # Helium at the center of the box

kgrid = [1, 1, 1]  # no kpoint sampling for an isolated system
Ecut = 30
tol = 1e-8

# dipole moment of a given density (assuming the current geometry)
function dipole(basis, ρ)
    rr = [a * (r[1] - 1/2) for r in r_vectors(basis)]
    d = sum(rr .* ρ) * basis.dvol
end;

model = model_LDA(lattice, atoms; symmetries=false)
basis = PlaneWaveBasis(model, Ecut; kgrid=kgrid)
res = self_consistent_field(basis, tol=tol)
μref = dipole(basis, res.ρ)

ε = .01
model_ε = model_LDA(lattice, atoms; extra_terms=[ExternalFromReal(r -> -ε * (r[1] - a/2))],
                    symmetries=false)
basis_ε = PlaneWaveBasis(model_ε, Ecut; kgrid=kgrid)
res_ε = self_consistent_field(basis_ε, tol=tol)
με = dipole(basis_ε, res_ε.ρ)

polarizability = (με - μref) / ε

println("Reference dipole:  $μref")
println("Displaced dipole:  $με")
println("Polarizability :   $polarizability")

using KrylovKit

# Apply (1- χ0 K)
function dielectric_operator(dρ)
    dv = apply_kernel(basis, dρ; ρ=res.ρ)
    χ0dv = apply_χ0(res.ham, res.ψ, res.εF, res.eigenvalues, dv)
    dρ - χ0dv
end

# dVext is the potential from a uniform field interacting with the dielectric dipole
# of the density.
dVext = [-a * (r[1] - 1/2) for r in r_vectors(basis)]
dVext = cat(dVext; dims=4)

# Apply χ0 once to get non-interacting dipole
dρ_nointeract = apply_χ0(res.ham, res.ψ, res.εF, res.eigenvalues, dVext)

# Solve Dyson equation to get interacting dipole
dρ = linsolve(dielectric_operator, dρ_nointeract, verbosity=3)[1]

println("Non-interacting polarizability: $(dipole(basis, dρ_nointeract))")
println("Interacting polarizability:     $(dipole(basis, dρ))")