module ExchangeCorrelationPotential

export xc_potential!


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


"""
    xc_potential!(Vxc, P, r)

Compute the local-density approximation (LDA) exchange-correlation
potential on a radial grid.

The density is constructed from the radial wavefunction as

    ρ(r) = |P(r) / r|².

The output array `Vxc` is overwritten in place, making the routine
suitable for repeated calls inside an SCF cycle.
"""
function xc_potential!(
    Vxc::AbstractVector{T},
    P::AbstractVector{S},
    r::AbstractVector{R},
) where {
    T<:AbstractFloat,
    S<:Real,
    R<:Real,
}

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

    @inbounds for i in eachindex(Vxc, P, r)
        ri = r[i]

        ri > 0 ||
            throw(DomainError(ri, "Radial grid points must be positive"))

        ρ = abs2(P[i] / ri)

        Vxc[i] = _lda_xc_potential(ρ)
    end

    return Vxc
end


@inline function _lda_xc_potential(ρ::Real)
    ρ ≥ 0 ||
        throw(DomainError(ρ, "Electron density must be non-negative"))

    iszero(ρ) && return 0.0

    rₛ = cbrt(3 / (4π * ρ))
    Vx = -cbrt(3ρ / π)

    if rₛ < 1
        log_rₛ = log(rₛ)

        Vc =
            A * log_rₛ +
            (B - A / 3) +
            (2C / 3) * rₛ * log_rₛ +
            ((2D - C) / 3) * rₛ

        return Vx + Vc
    end

    sqrt_rₛ = sqrt(rₛ)
    denominator = 1 + β1 * sqrt_rₛ + β2 * rₛ

    Vc =
        γ *
        (1 + (7β1 / 6) * sqrt_rₛ + (4β2 / 3) * rₛ) /
        denominator^2

    return Vx + Vc
end


end # module ExchangeCorrelationPotential