using Plots

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

# paramters of the system

angle = 0.0 * pi
μ = 0.5
l = 0
V1 = -5
R1 = sqrt(3)
V2 = 2
R2 = sqrt(10)

n_EC = 8
train_cs = (0.7 .+ 0.03 * randn(n_EC)) - 1im * (0.2 .+ 0.03 * randn(n_EC))
near_E = 0.2 + 0.2im

target_c = 0.5
exact_E = 0.20845136860234303 - 0.07100640993695649im

# HO basis

global EC_HO
begin
    println("HO basis calculation")
    μω_gen = 0.5 * exp(-1im * angle)
    n_max = 40
    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)

    global EC_HO = affine_EC(T, V)
    train!(EC_HO, train_cs; ref_eval=near_E, CAEC=false)
    extrapolate!(EC_HO, [target_c]; precalculated_exact_E=[exact_E])
end

# p-space

global EC_p_space
begin
    println("p-space calculation")
    vertices = [0, 4 * exp(-1im * angle)]
    subdivisions = [256]
    ks, ws = get_mesh(vertices, subdivisions)

    V_of_r(r) = V1 * exp(-r^2 / R1^2) + V2 * exp(-r^2 / R2^2)
    V_mat_elem(k, kp) = Vl_mat_elem(V_of_r, l, k, kp; atol=10^-5, maxevals=10^5, R_cutoff=16)
    V = get_V_matrix(V_mat_elem, ks, ws)
    T = get_T_matrix(ks, μ)

    global EC_p_space = affine_EC(T, V)
    train!(EC_p_space, train_cs; ref_eval=near_E, CAEC=false)
    extrapolate!(EC_p_space, [target_c]; precalculated_exact_E=[exact_E])
end

# Plotting

scatter([real(exact_E)], [imag(exact_E)], label="Exact", marker=:circle, markercolor=:white)
scatter!(real.(EC_HO.training_E), imag.(EC_HO.training_E), label="HO basis training", marker=:circle, color=:blue)
scatter!(real.(EC_HO.extrapolated_E), imag.(EC_HO.extrapolated_E), label="HO basis extrapolated", marker=:x, color=:blue)
scatter!(real.(EC_p_space.training_E), imag.(EC_p_space.training_E), label="p-space training", marker=:circle, color=:red)
scatter!(real.(EC_p_space.extrapolated_E), imag.(EC_p_space.extrapolated_E), label="p-space extrapolated", marker=:x, color=:red)
plot!(legend=:bottomleft)
xlabel!("Re(E)")
ylabel!("Im(E)")
savefig("temp/2b_p_space_vs_HO.pdf")