#ifndef TBPLAS_COMMON_H
#define TBPLAS_COMMON_H

#include <cstdint>

namespace common {
/**
 * @brief Calculate displacement vector:
 *   dr = pos_ket[j, :] - pos_bra[i, :] + [na, nb, nc] * lat_vac
 *
 * @tparam orbital_index_t datatype of orbital indices
 * @tparam cell_index_t datatype of cell indices
 * @param[in] pos_bra (num_orb_bra, 3) array in row-major order
 * @param[in] pos_ket (num_orb_ket, 3) array in row-major order
 * @param[in] orb_i orbital index i
 * @param[in] orb_j orbital index j
 * @param[in] lat_vec (3, 3) array in row-major order
 * @param[in] na cell index along a direction
 * @param[in] nb cell index along b direction
 * @param[in] nc cell index along c direction
 * @param[inout] dr (3,) array, displacement vector
 *
 * @note
 *   1. All matrices must be in ROW-MAJOR order!
 *   2. This function will be caused many times. DO NOT use MemoryView.
 *   3. This functions may be applied on both primitive and super cells.
 *      So we make it a template with orbital_index_ and cell_index_t as
 *      template parameters.
 */
template <typename orbital_index_t, typename cell_index_t>
void calc_dr(
    const double* pos_bra,
    const double* pos_ket,
    const orbital_index_t& orb_i,
    const orbital_index_t& orb_j,
    const double* lat_vec,
    const cell_index_t& na,
    const cell_index_t& nb,
    const cell_index_t& nc,
    double* dr)
{
    int64_t offset_i = orb_i * 3;
    int64_t offset_j = orb_j * 3;
    dr[0] = pos_ket[offset_j] - pos_bra[offset_i];
    dr[1] = pos_ket[offset_j + 1] - pos_bra[offset_i + 1];
    dr[2] = pos_ket[offset_j + 2] - pos_bra[offset_i + 2];
    dr[0] += na * lat_vec[0] + nb * lat_vec[3] + nc * lat_vec[6];
    dr[1] += na * lat_vec[1] + nb * lat_vec[4] + nc * lat_vec[7];
    dr[2] += na * lat_vec[2] + nb * lat_vec[5] + nc * lat_vec[8];
}
} // namespace common

#endif