using Plots, Arpack

include("../helper.jl")
include("../Hamiltonian.jl")

mode = cpu_tensor
T = Float32 # single-precision mode

V_r2(c) = r2 -> c * (-5 * exp(-r2/3) + 2 * exp(-r2/10))

d = 3
n = 2
N = 48
L = 30
ϕ = pi/6
n_imag = 1
s = system{T}(d, n, N, L)

train_cs = range(0.78, 0.45, length=5)
train_ref = reverse([0.05387926313545913-0.008900278182520881im,
    0.11254295298924327-0.020515067379548786im,
    0.16060154707503538-0.03716539208626717im,
    0.19741353362674618-0.05994519982799412im,
    0.2219100763497223-0.08959449893439568im])

extrapolate_cs = range(0.38, 0.22, length=5)
extrapolate_ref = reverse([0.23165109150003316-0.12052751440975719im,
    0.23190549514995962-0.1406687118589838im,
    0.22763660218046278-0.1626190970863793im,
    0.21807104244164865-0.18635600686249373im,
    0.2020979906072586-0.21180157628258728im])

training_E = ComplexF64[]
training_vec = Array[]
exact_E = ComplexF64[]
extrapolated_E = ComplexF64[]

for c in train_cs
    println("Training c=", c)
    H = Hamiltonian{T}(s, V_r2(c), ϕ, n_imag, mode)
    @time evals, evecs, info = eig(H, 20, resonances = true)
    i = nearestIndex(evals, pop!(train_ref))
    push!(training_E, evals[i])
    push!(training_vec, evecs[i])
end

N_EC = [sum(x .* y) for (x, y) in Iterators.product(training_vec, training_vec)]

for c in extrapolate_cs
    println("Extrapolating c=", c)
    H = Hamiltonian{T}(s, V_r2(c), ϕ, n_imag, mode)
    @time evals, _, info = eig(H, 40, resonances = true)
    nearestE = nearest(evals, pop!(extrapolate_ref))
    push!(exact_E, nearestE)
    
    # EC extrapolation
    H_training_vec = H.(training_vec)
    H_EC = [sum(x .* y) for (x, y) in Iterators.product(training_vec, H_training_vec)]

    evals = eigvals(H_EC, N_EC)
    push!(extrapolated_E, nearestE)
end

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")
savefig("temp/EC-R2R-S.pdf")