using SparseArrays, LinearAlgebra
include("../common.jl")

#### gram_schmidt ####
 
n = 64
N = 8

vecs = [rand(n) + 1im .* rand(n) for _ in 1:N]

ws = rand(n)
ws_mat = spdiagm(ws)

basis = hcat(vecs...)

H = rand(n, n)
H += transpose(H) # complex symmetric

H_EC = transpose(basis) * ws_mat * H * ws_mat * basis
N_EC = transpose(basis) * ws_mat * basis

evals = eigvals(H_EC, N_EC)
println("Eigenvalues with GEVP:")
display(evals)

ortho_basis = hcat(gram_schmidt!(vecs, ws)...)
N_ortho = transpose(ortho_basis) * ws_mat * ortho_basis
println("Norm matrix after Gram-Schmidt:")
display(round.(N_ortho; digits=2)) # should be ≈I
@assert N_ortho ≈ I(N) "Gram-Schmidt did not yield an orthogonal basis"

H_EC_ortho = transpose(ortho_basis) * ws_mat * H * ws_mat * ortho_basis

evals_ortho = eigvals(H_EC_ortho)
println("Eigenvalues after Gram-Schmidt:")
display(evals_ortho)

@assert evals ≈ evals_ortho "Gram-Schmidt did not preserve the eigenvalues"

############
println("\nRepeat with a redundant basis\n")

println("Original dimensionality = $(length(vecs))")

for pow in 6:9
    noise = rand(n) ./ 10^pow
    new_vec = vecs[1] + noise
    push!(vecs, new_vec)
end

println("Dimensionality before Gram-Schmidt = $(length(vecs))")

ortho_vecs = gram_schmidt!(vecs, ws; verbose=true)
ortho_basis = hcat(ortho_vecs...)
println("Dimensionality after Gram-Schmidt = $(length(ortho_vecs))")

H_EC_ortho = transpose(ortho_basis) * ws_mat * H * ws_mat * ortho_basis

evals_ortho = eigvals(H_EC_ortho)
println("Eigenvalues after Gram-Schmidt:")
display(evals_ortho)

@assert isapprox(evals, evals_ortho; atol=1e-3) "Gram-Schmidt did not approximately preserve the eigenvalues"

######################