using DFTK
using LinearAlgebra

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


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

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

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

ε = .01
model_ε = model_LDA(lattice, atoms, positions;
                    extra_terms=[ExternalFromReal(r -> -ε * (r[1] - a/2))],
                    symmetries=false)
basis_ε = PlaneWaveBasis(model_ε; Ecut, kgrid)
res_ε   = self_consistent_field(basis_ε; 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(δρ)
    δV = apply_kernel(basis, δρ; res.ρ)
    χ0δV = apply_χ0(res, δV)
    δρ - χ0δV
end

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

# Apply $χ_0$ once to get non-interacting dipole
δρ_nointeract = apply_χ0(res, δVext)

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

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