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

Λ = 0
m = 1.0
V_of_r(r) =  2 * exp(-(r-3)^2 / (1.5)^2)

vertices = [0, 6 - 0.6im]
subdivisions = [50]
jmax = 4

E_max = 40
μω_global = 0.5

H0, weights = get_3b_H_matrix(jacobi, V_of_r, vertices, subdivisions, jmax, μω_global, E_max, Λ, m)

# Vp = perturbation to make the state artificially bound
Vp_of_r(r) = -exp(-(r/3)^2)
Vp, _ = get_3b_H_matrix(jacobi, Vp_of_r, vertices, subdivisions, jmax, μω_global, E_max, Λ, m, false, true)

training_c = [2.6, 2.4, 2.2, 2.0, 1.8]
extrapolating_c = 0.0 : 0.2 : 1.2

training_ref = -2.22 # complete list not needed because identification is simple

exact_E = [4.077092809998592-0.01206085331850782im,
3.613579042377367-0.006920188044987599im,
3.145489628680764-0.003757512658877539im,
2.673033482861357-0.001939576896512454im,
2.196539134888566-0.0009597849595725841im,
1.7163902133045392-0.000456595029296216im,
1.2329696647679096-0.00019879325231064393im]

EC = affine_EC(H0, Vp, weights)
train!(EC, training_c; ref_eval=training_ref, CAEC=true)
extrapolate!(EC, extrapolating_c; precalculated_exact_E=exact_E)

exportCSV(EC, "temp/CSM_B2R.csv")
plot(EC, "temp/CSM_B2R.pdf")