using Plots

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

function solve_E(H0, H1, c_vals, near_E)
    out_E = ComplexF64[]
    for c in c_vals
        println("Solving for c = $c")

        H = H0 + c .* H1

        evals, _ = eigs(H, sigma=near_E, maxiter=5000, ritzvec=false, check=1)
        new_E = nearest(evals, near_E)
        push!(out_E, new_E)
    end
    return out_E
end

# 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.05 * randn(n_EC)) - 1im * (0.2 .+ 0.05 * randn(n_EC))
near_E = 0.2 + 0.2im

# HO basis

HO_Es = zeros(ComplexF64, n_EC)
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)

    HO_Es .= solve_E(T, V, train_cs, near_E)    
end

# p-space

p_space_Es = zeros(ComplexF64, n_EC)
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, μ)

    p_space_Es .= solve_E(T, V, train_cs, near_E)
end

# Plotting

scatter(real.(HO_Es), imag.(HO_Es), label="HO basis", marker=:circle, color=:blue)
scatter!(real.(p_space_Es), imag.(p_space_Es), label="p-space", marker=:circle, color=:red)
xlabel!("Re(E)")
ylabel!("Im(E)")
savefig("temp/2b_p_space_vs_HO.pdf")