using Plots

include("../../EC.jl")
include("../../ho_basis.jl")
include("../../p_space.jl")

# paramters of the system

angle = 0.0
μ = 0.5
l = 0
V1 = -5
R1 = sqrt(3)
V2 = 2
R2 = sqrt(10)

n_EC = 8
train_cs = (0.7 .+ 0.03 * randn(n_EC)) - 1im * (0.2 .+ 0.03 * randn(n_EC))
near_E = 0.2 + 0.2im

target_c = 0.5
exact_E = 0.20845136860234303 - 0.07100640993695649im

vertices = [0, 4 * exp(-1im * angle)]
subdivisions = [256]
ks, ws = get_mesh(vertices, subdivisions)

V_of_r(r) = V1 * exp(-r^2 / R1^2) + V2 * exp(-r^2 / R2^2)
V_mat_elem(k, kp) = Vl_mat_elem(V_of_r, l, k, kp; atol=10^-5, maxevals=10^5, R_cutoff=16)
V = get_V_matrix(V_mat_elem, ks, ws)
T = get_T_matrix(ks, μ)

EC_p_space = affine_EC(T, V)
train!(EC_p_space, train_cs; ref_eval=near_E, CAEC=false)
extrapolate!(EC_p_space, [target_c]; precalculated_exact_E=[exact_E])

# Plotting

theme(:dark) # Set the global theme to dark

scatter([real(exact_E)], [imag(exact_E)], label="exact", marker=:circle, markercolor=:white, bg = :black) # black background
scatter!(real.(EC_p_space.training_E), imag.(EC_p_space.training_E), label="training", marker=:circle, color=:blue)
scatter!(real.(EC_p_space.extrapolated_E), imag.(EC_p_space.extrapolated_E), label="extrapolated", marker=:x, color=:green)
hline!([0], color=:red, label="continuum")
plot!(legend=:bottomleft)
xlabel!("Re(E)")
ylabel!("Im(E)")
xlims!(0, 0.3)
ylims!(-0.3, 0.3)
savefig("temp/2body_p_space.pdf")