include("orbital_data.jl")
orb_2p_radial = LinearInterpolation(rs_DFT, orb_2p_radial_DFT)

nPts = size(coulomb_precompute)[1]
Rmax = 30;
Rs = range(0, Rmax, length = nPts)
τs = range(0, π / 2, length = nPts)
interp_call = LinearInterpolation((Rs, τs), coulomb_precompute)

"""
    Ψ_pz(r::Float64, θ::Float64)

2pz orbital in spherical coordinates.

The radial part is extracted from DFT, while the angular portion is the
spherical harmonic for 2pz (n=2, l=1, m=0) is cos(θ). The wave function does not
depend on ϕ, as expected for a pz orbital.

# Arguments
* `r`: radial distance from the center of the orbital in Bohr radii
* `θ`: polar angle

Note that the function returns 0 for r > 40 since the radial component at
r = 40 is ≈4.5×10⁻¹¹.
"""
@inline function Ψ_pz(r::Float64, θ::Float64)
    if r == 0.0 || r > 40
        return 0.0
    else
        # Note that the DFT values for the radial portion are r R(r)
        # so we divide by an extra factor of r when computing the WF
        return (orb_2p_radial(r) * cos(θ) / r * √(3 / (4 * π)))
    end
end

"""
    coulomb_potential_pz(R::Float64, τ::Float64)

Coulomb potential generated by a 2pz orbital at distance R and polar angle τ
from the center of the orbital.
"""
function coulomb_potential_pz(R::Float64, τ::Float64)
    # The cubature uses spherical coordinates with x[1] = r, x[2] = θ, x[3] = ϕ
    res = hcubature(
        r ->
            Ψ_pz(r[1], r[2]) .^ 2 * r[1]^2 * sin(r[2]) ./ (√(
                r[1]^2 + R^2 -
                2 *
                r[1] *
                R *
                (cos(τ) * cos(r[2]) + cos(r[3]) * sin(τ) * sin(r[2])),
            )),
        [0, 0, 0],
        [40, π, 2 * π],
        reltol = 1e-4,
    )
end

"""
    coulomb_potential_pz_interp(R::Float64, τ::Float64)

Interpolated version of [`coulomb_potential_pz`](@ref).

Since calculating [`coulomb_potential_pz`](@ref) involves computing a 3D
integral, multiple calculations can be costly. Therefore, this function uses
a precomputed grid to approximate [`coulomb_potential_pz`](@ref) for scenarios
when the direct calculation becomes not feasible.
"""
function coulomb_potential_pz_interp(R::Float64, τ::Float64)
    # Because of the pz shape, the potential is symmetric across the xy plane.
    # Hence, if the polar angle τ > π / 2,  we use τ → π - τ
    # (since the precalculation was performed for 0 < τ < π / 2)
    if τ > π / 2
        τ = π - τ
    end
    if R > Rmax
        return 1 / R
    else
        return interp_call(R, τ)
    end
end

@inline function coulomb_energy_pz_pz_Integrand(
    r::Vector{Float64},
    ρ::Tuple{Float64,Float64,Float64},
)
    # ρ is a CARTESIAN vector between the two pz orbitals
    # r is a SPHERICAL vector pointing from the center of one of the orbitals to
    # a point in space. it is the integration variable
    # dist is a CARTESIAN vector pointing from the center of the second pz to
    # the same point
    dist =
        [ρ[1], ρ[2], ρ[3]] - [
            r[1] * sin(r[2]) * cos(r[3]),
            r[1] * sin(r[2]) * sin(r[3]),
            r[1] * cos(r[2]),
        ]
    res =
        Ψ_pz(r[1], r[2]) .^ 2 *
        r[1]^2 *
        sin(r[2]) *
        coulomb_potential_pz_interp(norm(dist), acos((dist[3]) / norm(dist)))[1]
end


"""
    coulomb_energy_pz_pz(ρ::Tuple{Float64, Float64, Float64})

Coulomb interaction between two graphene orbitals.

# Arguments
* `ρ`: a 3-tuple separation vector (x,y,z) between the orbital centers
"""
function coulomb_energy_pz_pz(ρ::Tuple{Float64,Float64,Float64})
    hcubature(
        r -> coulomb_energy_pz_pz_Integrand(r, ρ),
        [0, 0, 0],
        [40, π, 2 * π],
        reltol = 1e-5,
    )
end
