using Roots, LinearAlgebra, Plots

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

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

# determining c0 with EC
temp_c = range(1.1, 0.9, 3)
p, w = get_mesh([0, 8], [256])
H0 = get_T_matrix(p, μ)
V = get_V_matrix(V_system(1), p, w)
EC = affine_EC(H0, V, w)
train!(EC, temp_c; ref_eval=-0.2, CAEC=false, verbose=false)
quick_extrapolate(c) = minimum(abs2, get_extrapolated_evals(EC.H0_EC, EC.H1_EC, EC.N_EC, c, 0))
c0 = find_zero(quick_extrapolate, 0.85)

# Calculation of training and extrapolating E
training_c = range(1.2, 0.9, 9) # original: range(1.35, 0.9, 5)
training_E = [quick_pole_E(V_system(c)) for c in training_c]
training_k = alt_sqrt.(2μ .* training_E)

extrapolating_c = range(0.78, 0.45, 7) # original: range(0.75, 0.40, 8)
exact_E = [quick_pole_E(V_system(c)) for c in extrapolating_c]

order::Int = ceil((length(training_c) - 1) / 2) # order of the Pade approximant

# Solve coefficients as a linear system
M_left_element(c, i) = alt_sqrt(c - c0)^i
M_left = M_left_element.(training_c, (0:order)')
M_right = -training_k .* M_left[:, 2:end] # remove the first column
M = hcat(M_left, M_right) # M = [M_left | M_right]
sol = M \ training_k
a = sol[1:order+1]
b = [1; sol[order+2:end]]

# Pade approximant
polynomial(a, c) = sum(i -> a[i+1] * alt_sqrt(c - c0)^i, 0:order)
pade_approx(c) = polynomial(a, c) / polynomial(b, c)

# Extrapolate
extrapolated_k = pade_approx.([training_c; extrapolating_c])
extrapolated_E = (extrapolated_k .^ 2) / (2μ)

# Plotting
scatter(real.(training_E), imag.(training_E), label="training")
scatter!(real.(exact_E), imag.(exact_E), label="exact")
scatter!(real.(extrapolated_E), imag.(extrapolated_E), label="extrapolated", m=:star5)