using LinearAlgebra, Plots
include("ho_basis.jl")
include("p_space.jl")

ω = 0.5
μ = 0.5
l = 0
V1 = -5
R1 = sqrt(3)
V2 = 2
R2 = sqrt(10)
n_max = 15

ns = collect(0:n_max)
ls = fill(l, n_max + 1)

T = 1/μ * sp_T_matrix(ns, ls; ω=ω)
V = V1 .* V_Gaussian.(R1, l, ns, transpose(ns); ω=ω) + V2 .* V_Gaussian.(R2, l, ns, transpose(ns); ω=ω)

cs = range(1.35, 0.9, 5)

E = similar(cs)
bench_E = similar(cs)

for (j, c) in enumerate(cs)
    H = T + c .* V
    evals = eigvals(H)
    bench_E[j] = quick_pole_E((p, q) -> c*(V1*g0(R1, p, q) + V2*g0(R2, p, q)), μ; cs_angle=0)
    i = argmin(abs.(evals .- bench_E[j]))
    E[j] = evals[i]
end

scatter(real.(bench_E), imag.(bench_E), label="p-space")
scatter!(real.(E), imag.(E), label="HO basis")
xlims!(-0.7,0)
ylims!(-0.1,0.1)