using SparseArrays
include("ho_basis.jl")

"Basis transformation from HO to momentum space"
function get_W_matrix(basis_p, E_max, Λ, μ1ω1, μ2ω2=μ1ω1; weights=true)
    Es, n1s, l1s, n2s, l2s = get_2p_basis(E_max, Λ)
    W = zeros(ComplexF64, length(basis_p), length(Es))
    Threads.@threads for idx in CartesianIndices(W)
        (i1, i2) = Tuple(idx)
        ((j1, j2), (k1, w1), (k2, w2)) = basis_p[i1]
        if j1 == l1s[i2] && j2 == l2s[i2]
            elem1 = 1/sqrt(sqrt(μ1ω1)) * (-1)^n1s[i2] * ho_basis(j1, n1s[i2], 1/sqrt(μ1ω1) * k1)
            elem2 = 1/sqrt(sqrt(μ2ω2)) * (-1)^n2s[i2] * ho_basis(j2, n2s[i2], 1/sqrt(μ2ω2) * k2)
            W[idx] = elem1 * elem2 * (weights ? w1 * w2 : 1)
        end
    end
    return sparse(W)
end