using LinearAlgebra
using SpecialFunctions, FastGaussQuadrature, QuadGK

function gausslegendre_shifted(a, b, n)
    scale = (b - a) / 2
    shift = (a + b) / 2
    p, w = gausslegendre(n)
    p = p .* scale .+ shift
    w = w .* scale
    return (p, w)
end

function get_mesh(vertices::Vector, subdivs::Vector)
    p = Vector{ComplexF64}()
    w = Vector{ComplexF64}()
    for (a, b, subdiv) in zip(vertices, vertices[2:end], subdivs)
        p_new, w_new = gausslegendre_shifted(a, b, subdiv)
        append!(p, p_new)
        append!(w, w_new)
    end
    return (p, w)
end

get_V_matrix(V_pq, p, w) = V_pq.(p, transpose(p)) .* transpose(w)

get_T_matrix(p, μ) = collect(Diagonal(p.*p ./ (2*μ)))

get_H_matrix(V_pq, p, w, μ=0.5) = get_T_matrix(p, μ) + get_V_matrix(V_pq, p, w)

function identify_pole_i(p, evals, μ=0.5)
    mesh_Es = (p.*p) ./ (2*μ)
    
    current_i = 0
    current_min_ΔE = -1.0
    for i in eachindex(evals)
        min_ΔE = minimum(abs.(mesh_Es .- evals[i]))
        if min_ΔE > current_min_ΔE
            current_i = i
            current_min_ΔE = min_ΔE
        end
    end
    return current_i
end

function quick_pole_E(V_pq, μ=0.5; cs_angle=0.4, cutoff=8.0, meshpoints=256)
    p, w = get_mesh([0, cutoff * exp(-1im * cs_angle)], [meshpoints])
    evals = eigvals(get_H_matrix(V_pq, p, w, μ))
    return evals[identify_pole_i(p, evals, μ)]
end

# Gaussian potentials in momentum space
g0(R, p, q) = (exp(-(1/4)*(p + q)^2*R^2)*(-1 + exp(p*q*R^2))*R)/(2*sqrt(π))
g1(R, p, q) = (exp(-(1/4)*(p + q)^2*R^2)*(2 + p*q*R^2 + exp(p*q*R^2)*(-2 + p*q*R^2)))/(2*p*sqrt(π)*q*R)

# general potential (numerical integration)
jHat(l, z) = z * sphericalbesselj(l, z)
function Vl_mat_elem(V_of_r, l, p, q; atol=0, maxevals=10^7, R_cutoff=Inf)
    integrand(r) = jHat(l, p * r) * V_of_r(r) * jHat(l, q * r)
    (integral, _) = quadgk(integrand, 0, R_cutoff; atol=atol, maxevals=maxevals)
    return (2 / pi) * integral
end