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

# contour
vertices = [0, 0.4 - 0.15im, 0.8, 6]
subdivisions = [128, 128, 128]
p, w = get_mesh(vertices, subdivisions)
mesh_E = p.*p ./ (2*0.5)

μ = 0.5
V_system(c) = (p, q) -> c*(-5*g0(sqrt(3), p, q) + 2*g0(sqrt(10), p, q)) # ResonanceEC: Eq. (20)

H0 = get_T_matrix(p, μ)
V = get_V_matrix(V_system(1), p, w)

training_c = range(0.75, 0.45, 5)
extrapolating_c = range(0.40, 0.25, 5)

training_ref = [quick_pole_E(V_system(c)) for c in training_c]
exact_E = [quick_pole_E(V_system(c)) for c in extrapolating_c]

EC = affine_EC(H0, V, w)
train!(EC, training_c; ref_eval=training_ref, CAEC=false)
extrapolate!(EC, extrapolating_c; precalculated_exact_E=exact_E)

#exportCSV(EC, "temp/2b_R2R.csv")
plot(EC, "temp/2b_R2R.pdf"; basis_contour=mesh_E, xlims=(0, 1))