Float = Union{Float32,Float64}

norm_square(x) = sum(x .* x)
reducedMass(m1::Float, m2::Float) = 1 / (1/m1 + 1/m2)

"A few-body system defined by its physical parameters"
struct system{T}
    d::Int
    n::Int
    N::Int
    L::T
    μs::Vector{T}
    invU::Matrix{T}

    function system{T}(d::Int, n::Int, N::Int, L::Real) where {T<:Float}
        μs = collect(1:n) ./ collect(2:(n + 1))

        U = zeros(T, n, n)
        for coord in 1:n
            U[coord, 1:coord] .= 1 / coord
            if coord != n
                U[coord, coord + 1] = -1
            end
        end
        invU = inv(U)[:, 1:(n - 1)]
        
        return new{T}(d, n, N, convert(T, L), μs, invU)
    end
end

"Eq (46): Partial derivative matrix element for 1 degree of freedom"
function ∂_1DOF(s::system{T}, k::Int, l::Int)::Complex{T} where {T<:Float}
    if k == l
        return -im * (π / s.L)
    else
        return (π / s.L) * (-1)^(k - l) * exp(-im * π * (k - l) / s.N) / sin(π * (k - l) / s.N)
    end
end

"Which index (dimension of the multidimensional array) corresponds to spatial dimension 'dim' of coordinate 'coord'?"
which_index(s::system, dim::Int, coord::Int)::Int = (dim - 1) * (s.n - 1) + coord

"Get the distance to the nearest image of the particle"
function nearest(s, Δk)
    if Δk > s.N ÷ 2
        return Δk - s.N
    elseif Δk < -s.N ÷ 2
        return Δk + s.N
    else
        return Δk
    end
end

"Calculate diagonal elements of the V matrix"
function calculate_Vs(s::system{T}, V_twobody::Function, ϕ::T, n_image::Int)::Array{Complex{T}} where {T<:Float}
    coeff² = (exp(im * ϕ) * s.L / s.N)^2
    images = collect.(Iterators.product(fill(-n_image:n_image, s.d)...)) # TODO: Learn how to use tuples instead of vectors
    Vs = zeros(Complex{T}, fill(s.N, s.d * (s.n - 1))...)
    Threads.@threads for i in CartesianIndices(Vs)
        xs = reshape(collect(Tuple(i)), s.n - 1, s.d) .- (s.N ÷ 2 - 1)
        rs = s.invU * xs
        for p1 in 1:s.n
            for p2 in 1:(p1 - 1)
                Δk = Array{T}(undef, s.d)
                for dim in 1:s.d
                    Δk[dim] = nearest(s, rs[p1, dim] - rs[p2, dim])
                end
                for image in images
                    Δk² = norm_square(Δk .- (s.N .* image))
                    Vs[i] += V_twobody(Δk² * coeff²)
                end
            end
        end
    end
    return Vs
end