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

berggren_mesh = get_mesh([0, 0.4 - 0.15im, 0.8, 6], [128, 128, 128])
csm_mesh = get_mesh([0, 8 - 3im], [512])

for (mesh, name) in zip((berggren_mesh, csm_mesh), ("beggren", "csm"))
    p, w = mesh
    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(1.1, 0.9, 5) # original: range(1.35, 0.9, 5)
    extrapolating_c = range(0.78, 0.45, 7) # original: range(0.75, 0.40, 8)

    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=true)
    extrapolate!(EC, extrapolating_c; precalculated_exact_E=exact_E)
    
    #exportCSV(EC, "temp/2b_comparison_$name.csv")
    plot(EC, "temp/2b_comparison_$name.pdf"; basis_contour=mesh_E, xlims=(-0.3,0.3), ylims=(-0.120,0.020))
end