module HartreePotential

export HartreeWorkspace,
       hartree_potential!


"""
    HartreeWorkspace(n)
    HartreeWorkspace(T, n)

Reusable workspace for solving the radial Poisson equation with the
Numerov discretization and Thomas algorithm.
"""
struct HartreeWorkspace{T<:AbstractFloat}
    α::Vector{T}
    β::Vector{T}
    y::Vector{T}
end


function HartreeWorkspace(::Type{T}, n::Integer) where {T<:AbstractFloat}
    n ≥ 3 || throw(ArgumentError("At least three grid points are required"))

    return HartreeWorkspace(
        zeros(T, n),
        zeros(T, n),
        zeros(T, n),
    )
end


HartreeWorkspace(n::Integer) = HartreeWorkspace(Float64, n)


"""
    hartree_potential!(Vh, P, r, nuclear_charge, h, workspace)

Compute the radial Hartree potential by solving

    d²Y/dr² = -4π P(r)² / r,

where

    Y(r) = r Vh(r),

using a Numerov discretization and the Thomas algorithm.

The boundary conditions are

    Y(r₀) = Z r₀,
    Y(r∞) = 1.

The array `Vh` is overwritten with the resulting Hartree potential.
The workspace is reused without allocating temporary arrays.
"""
function hartree_potential!(
    Vh::AbstractVector{T},
    P::AbstractVector{S},
    r::AbstractVector{R},
    nuclear_charge::Real,
    h::Real,
    workspace::HartreeWorkspace{W},
) where {
    T<:AbstractFloat,
    S<:Real,
    R<:Real,
    W<:AbstractFloat,
}

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

    n = length(r)

    n ≥ 3 ||
        throw(ArgumentError("At least three grid points are required"))

    length(workspace.α) == n ||
        throw(DimensionMismatch("Workspace size must match radial-grid size"))

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

    firstindex(r) == 1 ||
        throw(ArgumentError("Standard 1-based indexing is required"))

    r[1] > 0 ||
        throw(DomainError(r[1], "Radial grid points must be positive"))

    α = workspace.α
    β = workspace.β
    y = workspace.y

    source_factor = -(π / 3) * h * h

    α[2] = zero(W)
    β[2] = nuclear_charge * r[1]

    @inbounds for i in 2:n-1
        r[i + 1] > 0 ||
            throw(DomainError(r[i + 1], "Radial grid points must be positive"))

        source =
            abs2(P[i - 1]) / r[i - 1] +
            10 * abs2(P[i]) / r[i] +
            abs2(P[i + 1]) / r[i + 1]

        rhs = source_factor * source
        denominator = α[i] - 2

        α[i + 1] = -inv(denominator)
        β[i + 1] = (rhs - β[i]) / denominator
    end

    y[n] = one(W)

    @inbounds for i in n-1:-1:2
        y[i] = muladd(α[i + 1], y[i + 1], β[i + 1])
    end

    y[1] = β[2]

    @inbounds @simd for i in eachindex(Vh, y, r)
        Vh[i] = y[i] / r[i]
    end

    return Vh
end


end # module HartreePotential