using LinearAlgebra, Random, Plots
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)

training_c = range(1.2, 0.9, 9) # original: range(1.35, 0.9, 5)
extrapolating_c = range(0.78, 0.45, 7) # original: range(0.75, 0.40, 8)

# calculate training data
data_c = training_c
data_E = [quick_pole_E(V_system(c)) for c in data_c]

# hyperparameters
N = 9

# initialize random Hamiltonians
H0 = randn(N, N)
H1 = randn(N, N)

# training
Es = ComplexF64[]
ψrs = Vector{ComplexF64}[]
ψls = Vector{ComplexF64}[]

lr = 0.05
epochs = 100000

for epoch in 1:epochs
    empty!(Es)
    empty!(ψrs)
    empty!(ψls)
    for (c, E) in zip(data_c, data_E)
        H = H0 + c * H1
        r_evals, r_evecs = eigen(H)
        l_evals, l_evecs = eigen(transpose(H))
        @assert all(r_evals .≈ l_evals) "Right/left eigenvalues do not match"
        i = nearestIndex(r_evals, E) # TODO: more robust way to identify the eigenvector
        push!(Es, r_evals[i])
        push!(ψrs, r_evecs[:, i])
        push!(ψls, l_evecs[:, i])
    end

    if epoch % 1000 == 0
        loss = sum(abs2, Es .- data_E)
        println("Epoch:$epoch/$epochs \t Loss: $loss")
    end

    # gradient of the loss function
    function grad(c_order=0)
        out = zeros(ComplexF64, N, N)
        for (c, E_target, ψr, ψl, E) in zip(data_c, data_E, ψrs, ψls, Es)
            out .+= (c^c_order * conj(E - E_target)) .* (ψl * transpose(ψr))
        end
        return 2 .* real.(out)
    end
    H0 .-= lr .* grad(0) # update H0
    H1 .-= lr .* grad(1) # update H1    
end

# evaluate for all points
all_c = vcat(training_c, extrapolating_c)
exact_E = [quick_pole_E(V_system(c)) for c in all_c]
extrapolated_E = ComplexF64[]
for (c, ref) in zip(all_c, exact_E)
    H = H0 + c * H1
    evals, evecs = eigen(H)
    evals = vcat(evals, conj.(evals)) # include complex conjugates
    push!(extrapolated_E, nearest(evals, ref))
end

# plot results
scatter(real.(data_E), imag.(data_E), label="training", title="PMM", xlabel="Re E", ylabel="Im E")
scatter!(real.(exact_E), imag.(exact_E), label="exact", m=:+)
scatter!(real.(extrapolated_E), imag.(extrapolated_E), label="predicted", m=:x)