# helper functions



"""
    r_lattice(prob)

Return a 2d array of size (3, N_r) where the columns are the lattice coordinates of a uniform r grid in one unit cell with size_r(prob)[i] points in direction i.
"""
function r_lattice(prob::AbstractBSEProblem)
    N_rs = size_r(prob)
    rl_1d_1 = range(0.0, stop = 1.0, length = N_rs[1] + 1)[1:N_rs[1]]
    rl_1d_2 = range(0.0, stop = 1.0, length = N_rs[2] + 1)[1:N_rs[2]]
    rl_1d_3 = range(0.0, stop = 1.0, length = N_rs[3] + 1)[1:N_rs[3]]
    rl_1 = kron(ones(N_rs[3]), ones(N_rs[2]), rl_1d_1)
    rl_2 = kron(ones(N_rs[3]), rl_1d_2, ones(N_rs[1]))
    rl_3 = kron(rl_1d_3, ones(N_rs[2]), ones(N_rs[1]))
    rl = transpose(hcat(rl_1, rl_2, rl_3))
    return rl
end

"""
    r_cartesian(prob)

Return a 2d array of size (3, N_r) where the columns are the cartesian coordinates of the r grid generated by r_lattice.
"""
function r_cartesian(prob::AbstractBSEProblem)
    a_mat = lattice_matrix(prob)
    rl = r_lattice(prob)
    rc = similar(rl)
    for ir in 1:size(rl, 2)
        rc[:, ir] = a_mat * rl[:, ir]
    end
    return rc
end

"""
    k_lattice(prob[, shift])

Return a 2d array of size `(3, N_k)` where the columns are the lattice coordinates of a uniform ``k`` grid in one reciprocal unit cell with `size_k(prob)[i]` points in direction `i`. The grid can optionally be shifted by a vector `shift`.
"""
function k_lattice(prob::AbstractBSEProblem, shift = [0.0, 0.0, 0.0])
    N_ks = size_k(prob)
    kl_1d_1 = mod.(range(0.0, stop = 1.0, length = N_ks[1] + 1)[1:N_ks[1]] .+ shift[1], 1.0)
    kl_1d_2 = mod.(range(0.0, stop = 1.0, length = N_ks[2] + 1)[1:N_ks[2]] .+ shift[2], 1.0)
    kl_1d_3 = mod.(range(0.0, stop = 1.0, length = N_ks[3] + 1)[1:N_ks[3]] .+ shift[3], 1.0)
    kl_1 = kron(ones(N_ks[3]), ones(N_ks[2]), kl_1d_1)
    kl_2 = kron(ones(N_ks[3]), kl_1d_2, ones(N_ks[1]))
    kl_3 = kron(kl_1d_3, ones(N_ks[2]), ones(N_ks[1]))
    kl = transpose(hcat(kl_1, kl_2, kl_3))
    return kl
end

"""
    k_cartesian(prob[, shift])

Return a 2d array of size `(3, N_k)` where the columns are the
lattice coordinates of a uniform ``k`` grid in the first Brillouin zone
with `size_k(prob)[i]` points in direction `i`. The grid can
optionally be shifted by a vector `shift`.
The cartesian representation is chosen to lie in the first Brillouin
zone, i.e. to have minimal distance from the origin.
"""
function k_cartesian(prob::AbstractBSEProblem, shift = [0.0, 0.0, 0.0])
    a_mat = lattice_matrix(prob)
    b_mat = 2 * π * inv(a_mat)
    kl = k_lattice(prob, b_mat \ shift)
    kc = similar(kl)
    for ik in 1:size(kl, 2)
        kc[:, ik] = k_lattice_to_reduced_bz(kl[:, ik], b_mat)
    end
    return kc
end

"""
    k_lattice_to_reduced_bz(k_l, b_mat)

Helper function to compute cartesian coordinate of reciprocal point given in lattice coordinates.
"""
function k_lattice_to_reduced_bz(k_l, b_mat)
    shifted_norms = zeros(2, 2, 2)
    for i1 in 1:2, i2 in 1:2, i3 in 1:2
        shifted_norms[i1, i2, i3] = norm(b_mat * (k_l - [i1 - 1, i2 - 1, i3 - 1]))
    end
    m, i = findmin(shifted_norms)
    return b_mat * (k_l - [i[1] - 1, i[2]  - 1, i[3] - 1])
end

"""
    size_q(prob)

Number of differences of k-points in the different directions.
"""
function size_q(prob::AbstractBSEProblem)
    size_q(size_k(prob))
end
function size_q(N_ks::Tuple{Int,Int,Int})
    return (n -> n == 1 ? 1 : 2 * n).(N_ks) # double size unless size is 1
end

"""
    q_lattice(prob)

Return a 2d array of size `(3, N_q)` where the columns are the lattice coordinates of the uniform ``q`` grid generated by a uniform ``k`` grid in one reciprocal unit cell with `size_k(prob)[i]` points in direction `i`.
"""

function q_lattice(prob::AbstractBSEProblem)
    return q_lattice(size_k(prob))
end
function q_lattice(N_ks::Tuple{Int,Int,Int})
    N_qs = size_q(N_ks)
    ql_1d_1 = vcat(range(0, stop = 1, length = N_ks[1] + 1)[1:(end - 1)], range(-1, stop = 0, length = N_ks[1] + 1)[(end - (N_qs[1] - N_ks[1])):(end - 1)])
    ql_1d_2 = vcat(range(0, stop = 1, length = N_ks[2] + 1)[1:(end - 1)], range(-1, stop = 0, length = N_ks[2] + 1)[(end - (N_qs[2] - N_ks[2])):(end - 1)])
    ql_1d_3 = vcat(range(0, stop = 1, length = N_ks[3] + 1)[1:(end - 1)], range(-1, stop = 0, length = N_ks[3] + 1)[(end - (N_qs[3] - N_ks[3])):(end - 1)])
    ql_1 = kron(ones(N_qs[3]), ones(N_qs[2]), ql_1d_1)
    ql_2 = kron(ones(N_qs[3]), ql_1d_2, ones(N_qs[1]))
    ql_3 = kron(ql_1d_3, ones(N_qs[2]), ones(N_qs[1]))
    ql = transpose(hcat(ql_1, ql_2, ql_3))
    return ql
end

"""
    Ω0_volume(prob)

Return the volume of the unit cell.
"""
function Ω0_volume(prob::AbstractBSEProblem)
    return abs(det(lattice_matrix(prob)))
end

"""
    BZ_volume(prob)

Return the volume of the first Brillouin zone. The convention is to use
the normalization b_mat^* * a_mat = 2π with the lattice vectors a_mat
and the reciprocal vectors b_mat.
"""
function BZ_volume(prob::AbstractBSEProblem)
    return (2 * π)^3 / Ω0_volume(prob)
end

"""
    compute_v_hat(prob[, cutoff])

Compute the Fourier coefficients of the Coulomb potential. The frequencies are stored in the order compatible with the 3d FFT.
"""
function compute_v_hat(prob::AbstractBSEProblem, cutoff = Inf)
    N_rs = size_r(prob)
    b_mat = 2 * π * inv(lattice_matrix(prob))
    v_hat = 4 * pi * vec(mapslices(x -> norm(x) < 1e-10 || norm(b_mat * x) > cutoff ? 0.0 : 1 / norm(b_mat * x)^2, fftfreq(N_rs...); dims = 1))
    return v_hat
end

"""
    fftfreq(n1, n2, n3)

Generate an 3d array of Fourier frequencies in the order that they appear in the 3d FFT.
"""
function fftfreq(n1, n2, n3)
    f1 = vcat(collect(0:(div(n1, 2) - (1 - n1 % 2))), collect(-div(n1, 2):-1))
    f2 = vcat(collect(0:(div(n2, 2) - (1 - n2 % 2))), collect(-div(n2, 2):-1))
    f3 = vcat(collect(0:(div(n3, 2) - (1 - n3 % 2))), collect(-div(n3, 2):-1))

    F1 = kron(ones(n3), ones(n2), f1)
    F2 = kron(ones(n3), f2, ones(n1))
    F3 = kron(f3, ones(n2), ones(n1))

    return vcat(F1', F2', F3')
end

"""
    G_vector_to_index(G, N_rs)

Return the index `i` such that `fftfreq(N_rs...)[:, i] == G`.
"""
function G_vector_to_index(G, N_rs) # TODO: make fast for all dimensions
    if length(G) == 3
        ind = mod1(G[1] + 1, N_rs[1]) + N_rs[1] * (mod1(G[2] + 1, N_rs[2]) - 1) + N_rs[1] * N_rs[2] * (mod1(G[3] + 1, N_rs[3]) - 1)
    else
        ind = mod1(G[1] + 1, N_rs[1])
        for id in 2:length(G)
            ind += prod(N_rs[1:(id - 1)]) * (mod1(G[id] + 1, N_rs[id]) - 1)
        end
    end
    return ind
end
