module KSEnergy

export orbital_energy


"""
    orbital_energy(P, V, h)

Compute the Kohn-Sham orbital energy expectation value on a uniform
radial grid,

    E = 4π ∫ P(r) [ -½ d²P(r)/dr² + V(r)P(r) ] dr.

The second derivative is evaluated with a three-point central finite
difference. At the left boundary, the missing point is taken to have
zero wavefunction value, preserving the boundary treatment of the
original discretization.

# Arguments
- `P`: radial wavefunction values
- `V`: total Kohn-Sham potential on the same grid
- `h`: uniform radial-grid spacing

# Returns
The orbital energy.
"""
function orbital_energy(
    P::AbstractVector{T},
    V::AbstractVector{S},
    h::Real,
) where {T<:Real, S<:Real}

    axes(P) == axes(V) ||
        throw(DimensionMismatch("P and V must have matching axes"))

    length(P) ≥ 2 ||
        throw(ArgumentError("At least two grid points are required"))

    h > 0 ||
        throw(ArgumentError("Grid spacing h must be positive"))

    first = firstindex(P)
    last  = lastindex(P)

    inv_h² = inv(h * h)

    # Left boundary: P(r₀ - h) = 0
    d²P = (P[first + 1] - 2 * P[first]) * inv_h²

    energy_sum =
        P[first] * (-0.5 * d²P + V[first] * P[first])

    @inbounds @simd for i in first+1:last-1
        d²P = (P[i + 1] - 2 * P[i] + P[i - 1]) * inv_h²

        energy_sum +=
            P[i] * (-0.5 * d²P + V[i] * P[i])
    end

    return 4π * h * energy_sum
end


end # module KSEnergy