using LinearAlgebra, DelimitedFiles

"Sum over array while minimizing catastrophic cancellation as much as possible"
function better_sum(arr::Array{T}) where T<:Real
    pos_arr = arr[arr .> 0]
    neg_arr = arr[arr .< 0]

    sort!(pos_arr)
    sort!(neg_arr, rev=true)

    return sum(pos_arr) + sum(neg_arr)
end

better_sum(arr::Array{ComplexF64}) = better_sum(real.(arr)) + 1im * better_sum(imag.(arr))

"The triangle inequality for angular momenta"
triangle_ineq(l1, l2, L) = abs(l1 - l2) ≤ L && L ≤ (l1 + l2)

"Index of the nearest value in a list to a given reference point"
nearestIndex(list::Array, ref) = argmin(norm.(list .- ref))

"Nearest value in a list to a given reference point"
nearest(list::Array, ref) = list[nearestIndex(list, ref)]

"Simple implementation of the Kronecker sum"
function kron_sum(A::AbstractMatrix, B::AbstractMatrix)
    @assert size(A, 1) == size(A, 2) && size(B, 1) == size(B, 2) "Matrices should be square"
    return kron(A, I(size(B, 1))) + kron(I(size(A, 1)), B)
end

"Flattened vector version of Iterators.product(...) with index hierachy reversed -- leftmost index has the highest hierachy"
iter_prod(args...) = reverse.(collect(Iterators.product(reverse(args)...))[:])

"Export CSV data for a scatter plot taking in data as a list of complex vectors (x=Re, y=Im), and a list of corresponding labels, typically used for EC results"
function exportCSV(file::String, data, labels=nothing)
    columns = ["x" "y" "label"]

    if isnothing(labels)
        columns[end] = ""
        labels = fill("", length(data))
    end
    
    open(file, "w") do f
        writedlm(f, columns)
        for (d, label) in zip(data, labels)
            l = fill(label, length(d))
            writedlm(f, hcat(reim(d)..., l))
        end
    end
end

"In-place c-orthonormalization via (modified) Gram-Schmidt. Most significant vectors determined by the singular values (compared to the threshold) are returned."
function gram_schmidt!(vecs, ws, threshold=1e-5; verbose=false)
    c_product(i, j) = sum(vecs[i] .* ws .* vecs[j])
    norm(i) = c_product(i, i)
    proj(i, j) = (c_product(i, j) / norm(i)) .* vecs[i] # component of vec[i] in vec[j]

    # initial normalization
    for i in eachindex(vecs)
        vecs[i] ./= sqrt(norm(i))
    end

    target_dim = c_rank(vecs, ws, threshold)

    for i in eachindex(vecs)
        for j in (i + 1):length(vecs)
            vecs[j] .-= proj(i, j)
        end
    end
    
    abs_norms = abs.(norm.(eachindex(vecs)))
    verbose && println("Absolute norms = $(round.(abs_norms; sigdigits=1))")
    selected_vecs_i = partialsortperm(abs_norms, 1:target_dim; rev=true)

    for i in selected_vecs_i
        vecs[i] ./= sqrt(norm(i))
    end

    return vecs[selected_vecs_i]
end

"Rank of a basis set (pre-normalized) under c-product"
c_rank(vecs, ws, threshold=1e-8) = count(c_singular_values(vecs, ws) .> threshold)

"Singular values (magnitudes) of a basis set (pre-normalized) under c-product"
function c_singular_values(vecs, ws)
    basis = hcat(vecs...)
    N = transpose(basis) * spdiagm(ws) * basis
    singular_values = eigvals(N) .|> abs .|> sqrt
    return singular_values
end