using Plots

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

angle = 0.47 * pi
μω_gen = 0.5 * exp(-1im * angle)
μ = 0.5
l = 0
V1 = -5
R1 = sqrt(3)
V2 = 2
R2 = sqrt(10)
n_max = 15

ns = collect(0:n_max)
ls = fill(l, n_max + 1)

T = get_sp_T_matrix(ns, ls; μω_gen=μω_gen, μ=μ)
V = V1 .* V_Gaussian.(R1, l, ns, transpose(ns); μω_gen=μω_gen) + V2 .* V_Gaussian.(R2, l, ns, transpose(ns); μω_gen=μω_gen)

n_EC = 8
train_cs = (0.7 .+ 0.05 * randn(n_EC)) - 1im * (0.2 .+ 0.05 * randn(n_EC))
target_cs = [0.5]

near_E = 0.2 + 0.2im
exact_E = [0.20845136860234303 - 0.07100640993695649im]

EC = affine_EC(T, V; ensemble_size=32)
train!(EC, train_cs; ref_eval=near_E, CAEC=false)
extrapolate!(EC, target_cs; precalculated_exact_E=exact_E)

plot(EC; xlims=(0,0.3), ylims=(-0.3,0.3))
hline!([0], color=:red, label="continuum")
xlabel!("Re(E)")
ylabel!("Im(E)")
plot!(legend=:bottomleft)
savefig("temp/XZ.pdf")