module NumerovSolver

export NumerovWorkspace,
       solve_numerov!


"""
    NumerovWorkspace(n)

Reusable workspace for the Numerov–Thomas solver on a grid of `n` points.

The internal arrays are allocated once and can be reused in repeated
Kohn–Sham solves, for example during an SCF cycle or eigenvalue search.
"""
struct NumerovWorkspace{T<:AbstractFloat}
    α::Vector{T}
    β::Vector{T}
    y::Vector{T}
end


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

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


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


"""
    solve_numerov!(P, V, h, energy, workspace)

Solve the radial Kohn–Sham equation

    P''(r) + f(r)P(r) = 0,

with

    f(r) = 2[E - V(r)],

using the Numerov discretization and the Thomas algorithm.

The boundary values `P[firstindex(P)]` and `P[lastindex(P)]` are kept fixed.
The interior values of `P` are overwritten with the numerical solution.

# Arguments
- `P`: radial wavefunction; modified in place
- `V`: total Kohn–Sham potential
- `h`: uniform radial-grid spacing
- `energy`: trial orbital energy
- `workspace`: reusable `NumerovWorkspace`

# Returns
The modified wavefunction `P`.
"""
function solve_numerov!(
    P::AbstractVector{T},
    V::AbstractVector{S},
    h::Real,
    energy::Real,
    workspace::NumerovWorkspace{W},
) where {
    T<:AbstractFloat,
    S<:Real,
    W<:AbstractFloat,
}

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

    n = length(P)

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

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

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

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

    h²_over_12 = h * h / 12
    five_h²_over_6 = 10 * h²_over_12

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

    α[first + 1] = zero(W)
    β[first + 1] = P[first]

    @inbounds for i in first+1:last-1
        f₋ = 2 * (energy - V[i - 1])
        f₀ = 2 * (energy - V[i])
        f₊ = 2 * (energy - V[i + 1])

        A = 1 + h²_over_12 * f₋
        B = 1 + h²_over_12 * f₊
        C = -2 + five_h²_over_6 * f₀

        denominator = muladd(A, α[i], C)

        α[i + 1] = -B / denominator
        β[i + 1] = -(A * β[i]) / denominator
    end

    y[last] = P[last]

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

    @inbounds @simd for i in first+1:last-1
        P[i] = y[i]
    end

    return P
end


end # module NumerovSolver