module TotalEnergy

export lda_xc_energy,
       lda_xc_energy!,
       energy_corrections


# Perdew-Zunger-like LDA parameterization constants
const A  =  0.0311
const B  = -0.0480
const C  =  0.0020
const D  = -0.0116
const β1 =  1.0529
const β2 =  0.3334
const γ  = -0.1423


"""
    lda_xc_energy(ρ)

Return the local-density approximation (LDA) exchange-correlation
energy per particle for electron density `ρ`.
"""
@inline function lda_xc_energy(ρ::Real)
    ρ < 0 && throw(DomainError(ρ, "Electron density must be non-negative"))

    iszero(ρ) && return 0.0

    rₛ = cbrt(3.0 / (4.0 * π * ρ))

    εx = -0.75 * cbrt(3.0 * ρ / π)

    εc = if rₛ < 1.0
        A * log(rₛ) +
        B +
        C * rₛ * log(rₛ) +
        D * rₛ
    else
        γ / (1.0 + β1 * sqrt(rₛ) + β2 * rₛ)
    end

    return εx + εc
end


"""
    lda_xc_energy!(εxc, P, r)

Fill `εxc` with the LDA exchange-correlation energy per particle
corresponding to the radial wavefunction `P` on radial grid `r`.
"""
function lda_xc_energy!(
    εxc::AbstractVector{T},
    P::AbstractVector,
    r::AbstractVector,
) where {T<:AbstractFloat}

    axes(εxc) == axes(P) == axes(r) ||
        throw(DimensionMismatch("εxc, P, and r must have matching axes"))

    @inbounds for i in eachindex(εxc, P, r)
        ρ = abs2(P[i] / r[i])
        εxc[i] = lda_xc_energy(ρ)
    end

    return εxc
end


"""
    lda_xc_energy(P, r)

Return a newly allocated vector containing the LDA exchange-correlation
energy per particle on the radial grid.
"""
function lda_xc_energy(
    P::AbstractVector,
    r::AbstractVector,
)
    axes(P) == axes(r) ||
        throw(DimensionMismatch("P and r must have matching axes"))

    T = promote_type(Float64, eltype(P), eltype(r))
    εxc = Vector{T}(undef, length(P))

    return lda_xc_energy!(εxc, P, r)
end


"""
    energy_corrections(P, r, h, Vee, Vxc)

Compute the electron-electron and exchange-correlation corrections
to the total electronic energy.

Returns a named tuple with fields:

- `hartree_correction`
- `xc_energy`
- `xc_potential_correction`
"""
function energy_corrections(
    P::AbstractVector,
    r::AbstractVector,
    h::Real,
    Vee::AbstractVector,
    Vxc::AbstractVector,
)
    axes(P) == axes(r) == axes(Vee) == axes(Vxc) ||
        throw(DimensionMismatch("P, r, Vee, and Vxc must have matching axes"))

    hartree_sum = 0.0
    xc_sum = 0.0
    xc_potential_sum = 0.0

    first_i = firstindex(P)
    last_i = lastindex(P) - 1

    @inbounds for i in first_i:last_i
        probability = abs2(P[i])
        ρ = abs2(P[i] / r[i])
        εxc = lda_xc_energy(ρ)

        hartree_sum += probability * Vee[i]
        xc_sum += probability * εxc
        xc_potential_sum += probability * Vxc[i]
    end

    integration_factor = π * h

    return (
        hartree_correction = -0.5 * integration_factor * hartree_sum,
        xc_energy = integration_factor * xc_sum,
        xc_potential_correction = -integration_factor * xc_potential_sum,
    )
end


end # module TotalEnergy