using Printf

const N = 10001                  # total number of integration points
const Xzero = 1.0e-5             # r_0 - initial point
const Xinf = 20.0                # r_f - final point
const Z = 6.0                    # charge of the nucleus


"""
Analytical solution: hydrogen-like 1s wavefunction for charge Z.
"""
F(x, Z) = x * Z^1.5 * exp(-Z * x) / sqrt(π)


"""
Workspace for the Numerov + Thomas solver.

The arrays are allocated only once and reused in every solver call.
"""
struct NumerovWorkspace
    Al::Vector{Float64}
    Bt::Vector{Float64}
    Y::Vector{Float64}
end


"""
Create a workspace for a problem with n mesh points.
"""
NumerovWorkspace(n::Int) = NumerovWorkspace(
    zeros(Float64, n),
    zeros(Float64, n),
    zeros(Float64, n),
)


"""
Solve the Schrödinger equation using the Numerov + Thomas method.

Inputs:
    P   - radial wavefunction solution, modified in place
    V   - potential
    h   - mesh step
    En  - energy
    ws  - reusable NumerovWorkspace
"""
function Diff!(
    P::Vector{Float64},
    V::Vector{Float64},
    h::Float64,
    En::Float64,
    ws::NumerovWorkspace,
)
    n = length(P)

    @assert length(V) == n "P and V must have the same length"
    @assert length(ws.Al) == n "Workspace size does not match P"

    Al = ws.Al
    Bt = ws.Bt
    Y  = ws.Y

    h2_12 = h^2 / 12.0
    h2_56 = 5.0 * h^2 / 6.0

    fill!(Al, 0.0)
    fill!(Bt, 0.0)
    fill!(Y, 0.0)

    Al[2] = 0.0
    Bt[2] = P[1]   # boundary condition at r0

    for i in 2:n-1
        fi_1 = 2.0 * (En - V[i-1])
        fi   = 2.0 * (En - V[i])
        fi1  = 2.0 * (En - V[i+1])

        Ai = 1.0 + h2_12 * fi_1
        Bi = 1.0 + h2_12 * fi1
        Ci = -2.0 + h2_56 * fi

        denom = Ai * Al[i] + Ci

        Al[i+1] = -Bi / denom
        Bt[i+1] = -Ai * Bt[i] / denom
    end

    Y[end] = P[end]   # boundary condition at r_f

    for i in n-1:-1:2
        Y[i] = Al[i+1] * Y[i+1] + Bt[i+1]
    end

    P[2:end-1] .= Y[2:end-1]

    return nothing
end


function main()
    # mesh
    rgrid = range(Xzero, Xinf; length=N)
    r = collect(rgrid)
    h = step(rgrid)

    # analytical solution and nuclear potential
    P0 = F.(r, Z)
    Vnuc = -Z ./ r

    # analytical energy
    En0 = -Z^2 / 2.0

    # numerical solution
    P = zeros(Float64, N)
    P[1] = P0[1]
    P[end] = P0[end]

    # allocate solver workspace only once
    ws = NumerovWorkspace(N)

    # solve Schrödinger equation
    Diff!(P, Vnuc, h, En0, ws)

    # Example of workspace reuse for another energy:
    #
    # En1 = -17.9
    # fill!(P, 0.0)
    # P[1] = P0[1]
    # P[end] = P0[end]
    # Diff!(P, Vnuc, h, En1, ws)

    # write results to file
    open("data.dat", "w") do io
        @printf(
            io,
            "#  %9d     r(i)   P0(i)   P(i)   Vnuc(i)   error (P(i)-P0(i))\n",
            N,
        )

        for i in eachindex(r)
            @printf(
                io,
                "%20.15f%20.15f%20.15f%20.10f%20.12e\n",
                r[i],
                P0[i],
                P[i],
                Vnuc[i],
                P[i] - P0[i],
            )
        end
    end
end


main()