using Plots
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)

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

training_points = range(0.75, 0.45, 5)
training_E = Vector{ComplexF64}(undef, length(training_points))
EC_basis = Matrix{ComplexF64}(undef, length(p), length(training_points))

for (j, c) in enumerate(training_points)
    evals, evecs = eigen(get_H_matrix(V_system(c), p, w))
    i = identify_pole_i(p, evals)
    training_E[j] = evals[i]
    EC_basis[:, j] = evecs[:, i]
end

extrapolate_points = range(0.40, 0.25, 5)
ref_E = 0.2 - 0.1im

exact_E = Vector{ComplexF64}(undef, length(extrapolate_points))
extrapolate_E = Vector{ComplexF64}(undef, length(extrapolate_points))

EC_basis_w = EC_basis .* w
N_EC = transpose(EC_basis_w) * EC_basis

for (j, c) in enumerate(extrapolate_points)
    exact_E[j] = quick_pole_E(V_system(c))

    H = get_H_matrix(V_system(c), p, w)
    H_EC = transpose(EC_basis_w) * H * EC_basis
    evals = eigvals(H_EC, N_EC)
    i = argmin(abs.(evals .- ref_E))
    global ref_E = evals[i]
    extrapolate_E[j] = evals[i]
end

scatter(real.(training_E), imag.(training_E), label="training")
scatter!(real.(exact_E), imag.(exact_E), label="exact")
scatter!(real.(extrapolate_E), imag.(extrapolate_E), label="extrapolated")
plot!(real.(mesh_E), imag.(mesh_E), label="contour")
xlims!(0,1)