/**
 * @brief Functions and classes for advanced modeling
 *
 * @author: @author: Yunhai Li (yhli)
 *
 * @copyright Copyright (c) 2024, TBPLaS develop team.
 */

#ifndef TBPLAS_BUILDER_ADVANCED_H
#define TBPLAS_BUILDER_ADVANCED_H

#include <cassert>
#include <tuple>

#include "base.h"
#include "kdtree.h"
#include "primitive.h"

namespace tbplas::builder {
using namespace std::complex_literals;

// Datatype aliases
using dim_t = std::tuple<cell_index_t, cell_index_t, cell_index_t>;
using id_pc_t = std::tuple<cell_index_t, cell_index_t, cell_index_t, orbital_index_t>;

/**
 * @brief Split a vector into separate components
 *
 * @tparam vector_t datatype of vector
 * @tparam scalar_t datatype of components
 * @param[in] vec vector to split
 * @param[out] fa component on a-direction
 * @param[out] fb component on b-direction
 * @param[out] fc component on c-direction
 */
template <typename vector_t, typename scalar_t>
void split_vector3(const vector_t& vec, scalar_t& fa, scalar_t& fb, scalar_t& fc)
{
    fa = vec(0);
    fb = vec(1);
    fc = vec(2);
}

/**
 * @brief Get the floor of a vector
 *
 * @tparam vector_t datatype of vector
 * @tparam scalar_t datatype of components
 * @param[in] vec vector to split
 * @param[out] ra floor on a direction
 * @param[out] rb floor on b direction
 * @param[out] rc floor on c direction
 * @param[in] delta small displacement added to vector to avoid clipping
 */
template <typename vector_t, typename scalar_t>
void floor_vector3(const vector_t& vec, scalar_t& ra, scalar_t& rb, scalar_t& rc, const double& delta = 0.02)
{
    ra = std::floor(vec(0) + delta);
    rb = std::floor(vec(1) + delta);
    rc = std::floor(vec(2) + delta);
}

/**
 * @brief Python-compatible version of division j0 / n0
 *
 * @tparam data_t datatype of j0, n0 and results
 * @param[in] j0 dividend
 * @param[in] n0 divisor
 * @param[out] j1 = j0 % n0
 * @param[out] n1 = j0 // n0
 */
template <typename data_t>
void py_div(const data_t& j0, const data_t& n0, data_t& j1, data_t& n1)
{
    // Branches are sorted according to their matching probabilities.
    if (0 <= j0 && j0 < n0) {
        j1 = j0;
        n1 = 0;
    } else if (n0 <= j0 && j0 < 2 * n0) {
        j1 = j0 - n0;
        n1 = 1;
    } else if (-n0 <= j0 && j0 < 0) {
        j1 = j0 + n0;
        n1 = -1;
    } else {
        data_t j1_new = j0;
        while (j1_new < 0) {
            j1_new += n0;
        }
        while (j1_new >= n0) {
            j1_new -= n0;
        }
        j1 = j1_new;
        n1 = (j0 - j1_new) / n0;
    }
}

/**
 * @brief Auxiliary class for mapping orbital indices in primitive cell (pc) and supercell (sc)
 * representations in general cases.
 *
 * @todo Hilbert curve mapping
 */
class OrbitalMap {
protected:
    /// @brief orbital indices in pc representation
    std::vector<id_pc_t> sc2pc_vec_;

    /// @brief map of orbital indices from pc to sc representation
    std::map<id_pc_t, orbital_index_t> pc2sc_map_;

    /// @brief counter of the orbitals
    orbital_index_t counter_ = 0;

public:
    /**
     * @brief Construct a new Orbital Map object
     *
     * @param @param[in] num_orb_sc number of orbitals to be handled
     */
    explicit OrbitalMap(orbital_index_t num_orb_sc);

    /**
     * @brief Add an orbital from index in pc representation.
     *
     * @param[in] id_pc index of the orbital
     */
    void add(const id_pc_t& id_pc);

    /**
     * @brief Convert orbital index from pc to sc representation.
     *
     * @param[in] id_pc index of the orbital
     * @return corresponding index in sc representation
     */
    [[nodiscard]] orbital_index_t pc2sc(const id_pc_t& id_pc) const;

    /**
     * @brief Convert orbital index from sc to pc representation.
     *
     * @param[in] id_sc index of the orbital
     * @return corresponding index in pc representation
     */
    [[nodiscard]] const id_pc_t& sc2pc(const orbital_index_t& id_sc) const;

    /**
     * @brief Get the counter attribute
     *
     * @return the counter starting from 0
     */
    [[nodiscard]] const orbital_index_t& get_counter() const;
};

/**
 * @brief Similar to OrbitalMap, but the orbitals are generated using a 4-fold loops over
 * dim_a, dim_b, dim_c and num_orb_pc.
 *
 */
class RegularOrbitalMap {
private:
    /// @brief Dimension of supercell along a-direction.
    cell_index_t dim_a_ = 0;

    /// @brief Dimension of supercell along a-direction.
    cell_index_t dim_b_ = 0;

    /// @brief Dimension of supercell along c-direction.
    cell_index_t dim_c_ = 0;

    /// @brief Product of dim_b, and dim_c.
    cell_index_t dim_bc_ = 0;

    /// @brief Number of orbitals in primitive cell.
    orbital_index_t num_orb_pc_ = 0;

    /// @brief counter of the orbitals
    orbital_index_t counter_ = 0;

public:
    /**
     * @brief Construct a new Regular Orbital Map object
     *
     * @param[in] dim_a Dimension of supercell along a-direction.
     * @param[in] dim_b Dimension of supercell along b-direction.
     * @param[in] dim_c Dimension of supercell along c-direction.
     * @param[in] num_orb_pc Number of orbitals in primitive cell.
     */
    explicit RegularOrbitalMap(
        cell_index_t dim_a,
        cell_index_t dim_b,
        cell_index_t dim_c,
        orbital_index_t num_orb_pc);

    /**
     * @brief Increase counter_ by 1.
     *
     */
    void add();

    /**
     * @brief Convert orbital index from pc to sc representation.
     *
     * @param[in] i_a cell index of along a direction
     * @param[in] i_b cell index of along a direction
     * @param[in] i_c cell index of along a direction
     * @param[in] i_o orbital index
     * @return corresponding index in sc representation
     *
     * @note For efficiency we do not accept id_pc_t as input argument
     * as in OrbitalMap.
     */
    [[nodiscard]] orbital_index_t pc2sc(
        const cell_index_t& i_a,
        const cell_index_t& i_b,
        const cell_index_t& i_c,
        const orbital_index_t& i_o) const;

    /**
     * @brief Get the counter attribute
     *
     * @return the counter starting from 0
     */
    [[nodiscard]] const orbital_index_t& get_counter() const;
};

/**
 * @brief Add hopping terms to extended cell.
 *
 * @tparam hop_t datatype of hopping terms
 * @param[in] prim_cell primitive cell from which the extended cell is constructed
 * @param[in] dim dimension of the extended cell along 3 directions
 * @param[in] orb_map orbital map for converting orbital indices
 * @param[inout] extend_cell extended cell to which hopping terms are added
 */
template <typename hop_t>
void extend_add_hop(
    const PrimitiveCell<hop_t>& prim_cell,
    const dim_t& dim,
    const RegularOrbitalMap& orb_map,
    PrimitiveCell<hop_t>& extend_cell)
{
    // Unpack dim
    cell_index_t dim_a = 1;
    cell_index_t dim_b = 1;
    cell_index_t dim_c = 1;
    std::tie(dim_a, dim_b, dim_c) = dim;

    // Loop over cell indices
    const auto& hop_data = prim_cell.get_all_hopping().get_data();
    for (cell_index_t i_a = 0; i_a < dim_a; ++i_a) {
        for (cell_index_t i_b = 0; i_b < dim_b; ++i_b) {
            for (cell_index_t i_c = 0; i_c < dim_c; ++i_c) {
                // Offset for orbital i in (ia, ib, ic)
                orbital_index_t offset_sc_i = orb_map.pc2sc(i_a, i_b, i_c, 0);

                // Loop over rn of hopping terms
                auto iter_rn = hop_data.begin();
                for (; iter_rn != hop_data.end(); iter_rn++) {
                    const auto& rn = (*iter_rn).first;
                    bool rn_non_zero = (std::get<0>(rn) != 0) || (std::get<1>(rn) != 0) || (std::get<2>(rn) != 0);

                    // Default offset_j and new hop_rn assuming rn == (0, 0, 0)
                    orbital_index_t offset_sc_j = offset_sc_i;
                    cell_index_t na = 0;
                    cell_index_t nb = 0;
                    cell_index_t nc = 0;

                    // Update offset_j and new hop_rn if rn != (0, 0, 0)
                    if (rn_non_zero) {
                        cell_index_t ja0 = i_a + std::get<0>(rn);
                        cell_index_t jb0 = i_b + std::get<1>(rn);
                        cell_index_t jc0 = i_c + std::get<2>(rn);
                        cell_index_t ja = 0;
                        cell_index_t jb = 0;
                        cell_index_t jc = 0;
                        py_div(ja0, dim_a, ja, na);
                        py_div(jb0, dim_b, jb, nb);
                        py_div(jc0, dim_c, jc, nc);
                        offset_sc_j = orb_map.pc2sc(ja, jb, jc, 0);
                    }

                    // Add hopping term
                    auto iter_ij = ((*iter_rn).second).begin();
                    for (; iter_ij != ((*iter_rn).second).end(); iter_ij++) {
                        const auto& pair = (*iter_ij).first;
                        const auto& t = (*iter_ij).second;
                        orbital_index_t id_sc_i = std::get<0>(pair) + offset_sc_i;
                        orbital_index_t id_sc_j = std::get<1>(pair) + offset_sc_j;
                        extend_cell.add_hopping(na, nb, nc, id_sc_i, id_sc_j, t);
                    }
                }
            }
        }
    }
}

/**
 * @brief Extend primitive cell along a, b and c directions.
 *
 * @tparam hop_t datatype of hopping terms
 * @param[in] prim_cell primitive cell from which the extended cell is constructed
 * @param[in] dim dimension of the extended cell along 3 directions
 * @return extended cell created from primitive cell
 */
template <typename hop_t>
PrimitiveCell<hop_t> extend_prim_cell(
    const PrimitiveCell<hop_t>& prim_cell,
    const dim_t& dim)
{
    // Check the dimension of extended cell
    cell_index_t dim_a = 1;
    cell_index_t dim_b = 1;
    cell_index_t dim_c = 1;
    std::tie(dim_a, dim_b, dim_c) = dim;
    std::tie(dim_a, dim_b, dim_c) = dim;
    assert(dim_a > 0);
    assert(dim_b > 0);
    assert(dim_c > 0);
    Eigen::Vector3d dim_vec = rn2vec(dim);

    // Extend lattice vectors
    Eigen::Matrix3d lat_vec_ext = prim_cell.get_lattice();
    for (auto i = 0; i < 3; ++i) {
        lat_vec_ext.row(i) *= dim_vec(i);
    }

    // Extended cell shares the same origin
    Eigen::Vector3d origin_ext = prim_cell.get_origin();

    // Create extended cell and add orbitals
    orbital_index_t num_orb = prim_cell.get_num_orbitals();
    orbital_index_t num_orb_ext = num_orb * dim_a * dim_b * dim_c;
    PrimitiveCell<hop_t> extend_cell(num_orb_ext, lat_vec_ext, origin_ext, 1.0);
    RegularOrbitalMap orb_map(dim_a, dim_b, dim_c, num_orb);
    for (cell_index_t i_a = 0; i_a < dim_a; ++i_a) {
        for (cell_index_t i_b = 0; i_b < dim_b; ++i_b) {
            for (cell_index_t i_c = 0; i_c < dim_c; ++i_c) {
                Eigen::Vector3d rn_vec { i_a * 1.0, i_b * 1.0, i_c * 1.0 };
                for (orbital_index_t i_o = 0; i_o < num_orb; ++i_o) {
                    const Orbital& orbital = prim_cell.get_orbital(i_o);
                    Eigen::Vector3d new_pos = orbital.get_position() + rn_vec;
                    double fa = new_pos(0) / dim_vec(0);
                    double fb = new_pos(1) / dim_vec(1);
                    double fc = new_pos(2) / dim_vec(2);
                    extend_cell.set_orbital(orb_map.get_counter(), fa, fb, fc, orbital.energy, orbital.label);
                    orb_map.add(); // Must come after 'set_orbital', otherwise get_counter will be wrong!
                }
            }
        }
    }

    // Add hopping terms
    extend_add_hop(prim_cell, dim, orb_map, extend_cell);
    return extend_cell;
}

/**
 * @brief Determine the number of orbitals in reshaped cell
 *
 * @tparam hop_t datatype of hopping terms
 * @param[in] prim_cell primitive cell from which the reshaped cell is constructed
 * @param[in] conv_mat conv_mat from reshape_prim_cell
 * @param[in] delta small displacement added to orbital positions
 * @param[in] orb_pos fractional coordinates of prim_cell
 * @param[in] rn_range searching range from reshape_prim_cell
 * @return number of orbitals in reshaped cell
 */
template <typename hop_t>
orbital_index_t reshape_get_num_orb(
    const PrimitiveCell<hop_t>& prim_cell,
    const Eigen::Matrix3d& conv_mat,
    const double& delta,
    const Eigen::Matrix3Xd& orb_pos,
    const Eigen::Matrix3Xi& rn_range)
{
    cell_index_t ra_min = rn_range(0, 0);
    cell_index_t ra_max = rn_range(0, 1);
    cell_index_t rb_min = rn_range(1, 0);
    cell_index_t rb_max = rn_range(1, 1);
    cell_index_t rc_min = rn_range(2, 0);
    cell_index_t rc_max = rn_range(2, 1);
    cell_index_t ra = 0;
    cell_index_t rb = 0;
    cell_index_t rc = 0;

    orbital_index_t num_orb_prim = prim_cell.get_num_orbitals();
    orbital_index_t num_orb_res = 0;
    for (cell_index_t i_a = ra_min; i_a <= ra_max; ++i_a) {
        for (cell_index_t i_b = rb_min; i_b <= rb_max; ++i_b) {
            for (cell_index_t i_c = rc_min; i_c <= rc_max; ++i_c) {
                Eigen::Vector3d rn(i_a * 1.0, i_b * 1.0, i_c * 1.0);
                for (orbital_index_t i_o = 0; i_o < num_orb_prim; ++i_o) {
                    Eigen::Vector3d pos = orb_pos.col(i_o);
                    Eigen::Vector3d res_pos = ((rn + pos).transpose() * conv_mat).transpose();
                    floor_vector3(res_pos, ra, rb, rc, delta);
                    if ((ra == 0) && (rb == 0) && (rc == 0)) {
                        ++num_orb_res;
                    }
                }
            }
        }
    }
    return num_orb_res;
}

/**
 * @brief Update orbital attributes of reshaped cell and orbital map
 *
 * @tparam hop_t datatype of hopping terms
 * @param[in] prim_cell primitive cell from which the reshaped cell is constructed
 * @param[in] conv_mat conv_mat from reshape_prim_cell
 * @param[in] delta small displacement added to orbital positions
 * @param[in] orb_pos fractional coordinates of prim_cell
 * @param[in] rn_range searching range from reshape_prim_cell
 * @param[inout] res_cell reshaped cell to manipulate
 * @param[inout] orb_map orbital map to manipulate
 */
template <typename hop_t>
void reshape_upd_orb(
    const PrimitiveCell<hop_t>& prim_cell,
    const Eigen::Matrix3d& conv_mat,
    const double& delta,
    const Eigen::Matrix3Xd& orb_pos,
    const Eigen::Matrix3Xi& rn_range,
    PrimitiveCell<hop_t>& res_cell,
    OrbitalMap& orb_map)
{
    cell_index_t ra_min = rn_range(0, 0);
    cell_index_t ra_max = rn_range(0, 1);
    cell_index_t rb_min = rn_range(1, 0);
    cell_index_t rb_max = rn_range(1, 1);
    cell_index_t rc_min = rn_range(2, 0);
    cell_index_t rc_max = rn_range(2, 1);
    cell_index_t ra = 0;
    cell_index_t rb = 0;
    cell_index_t rc = 0;
    double fa = 0.0;
    double fb = 0.0;
    double fc = 0.0;

    orbital_index_t num_orb_prim = prim_cell.get_num_orbitals();
    for (cell_index_t i_a = ra_min; i_a <= ra_max; ++i_a) {
        for (cell_index_t i_b = rb_min; i_b <= rb_max; ++i_b) {
            for (cell_index_t i_c = rc_min; i_c <= rc_max; ++i_c) {
                Eigen::Vector3d rn(i_a * 1.0, i_b * 1.0, i_c * 1.0);
                for (orbital_index_t i_o = 0; i_o < num_orb_prim; ++i_o) {
                    Eigen::Vector3d pos = orb_pos.col(i_o);
                    Eigen::Vector3d res_pos = ((rn + pos).transpose() * conv_mat).transpose();
                    split_vector3(res_pos, fa, fb, fc);
                    floor_vector3(res_pos, ra, rb, rc, delta);
                    if ((ra == 0) && (rb == 0) && (rc == 0)) {
                        const Orbital& orbital = prim_cell.get_orbital(i_o);
                        res_cell.set_orbital(orb_map.get_counter(), fa, fb, fc, orbital.energy, orbital.label);
                        orb_map.add(std::make_tuple(i_a, i_b, i_c, i_o));
                    }
                }
            }
        }
    }
}

/**
 * @brief Add hopping terms to reshaped cell
 *
 * @tparam hop_t datatype of hopping terms
 * @param[in] prim_cell primitive cell from which the reshaped cell is constructed
 * @param[in] conv_mat conv_mat from reshape_prim_cell
 * @param[in] delta small displacement added to orbital positions
 * @param[in] orb_pos fractional coordinates of prim_cell
 * @param[inout] res_cell reshaped cell to manipulate
 * @param[in] orb_map orbital map to convert orbital indices
 */
template <typename hop_t>
void reshape_add_hop(
    const PrimitiveCell<hop_t>& prim_cell,
    const Eigen::Matrix3d& conv_mat,
    const double& delta,
    const Eigen::Matrix3Xd& orb_pos,
    PrimitiveCell<hop_t>& res_cell,
    const OrbitalMap& orb_map)
{
    cell_index_t ra = 0;
    cell_index_t rb = 0;
    cell_index_t rc = 0;
    Eigen::Matrix3d lat_frac = conv_mat.inverse();

    if (prim_cell.get_num_hopping() > 0) {
        const auto& hop_data = prim_cell.get_all_hopping().get_data();
        for (orbital_index_t id_sc_i = 0; id_sc_i < res_cell.get_num_orbitals(); ++id_sc_i) {
            const id_pc_t& id_pc_i = orb_map.sc2pc(id_sc_i);

            // Loop over hopping terms for search for matched ones
            auto iter_rn = hop_data.begin();
            for (; iter_rn != hop_data.end(); iter_rn++) {
                const auto& rn = (*iter_rn).first;
                auto iter_ij = ((*iter_rn).second).begin();
                for (; iter_ij != ((*iter_rn).second).end(); iter_ij++) {
                    const auto& pair = (*iter_ij).first;
                    const auto& t = (*iter_ij).second;
                    if (std::get<3>(id_pc_i) == std::get<0>(pair)) {
                        // Get cell index of id_sc_j in reshaped cell
                        cell_index_t ja0 = std::get<0>(id_pc_i) + std::get<0>(rn);
                        cell_index_t jb0 = std::get<1>(id_pc_i) + std::get<1>(rn);
                        cell_index_t jc0 = std::get<2>(id_pc_i) + std::get<2>(rn);
                        Eigen::Vector3d prim_rn(ja0 * 1.0, jb0 * 1.0, jc0 * 1.0);
                        Eigen::Vector3d prim_pos = orb_pos.col(std::get<1>(pair));
                        Eigen::Vector3d res_pos = ((prim_rn + prim_pos).transpose() * conv_mat).transpose();
                        floor_vector3(res_pos, ra, rb, rc, delta);

                        //  Wrap id_sc_j back into the (0, 0, 0) reshaped cell
                        Eigen::Vector3d res_rn(ra * 1.0, rb * 1.0, rc * 1.0);
                        Eigen::Vector3d prim_rn_wrap = prim_rn - (res_rn.transpose() * lat_frac).transpose();
                        id_pc_t id_pc_j = std::make_tuple(std::round(prim_rn_wrap(0)),
                            std::round(prim_rn_wrap(1)),
                            std::round(prim_rn_wrap(2)),
                            std::get<1>(pair));
                        orbital_index_t id_sc_j = orb_map.pc2sc(id_pc_j);
                        res_cell.add_hopping(ra, rb, rc, id_sc_i, id_sc_j, t);
                    }
                }
            }
        }
    }
}

/**
 * @brief Reshape primitive cell to given lattice vectors.
 *
 * @tparam hop_t datatype of hopping terms
 * @param[in] prim_cell primitive cell from which the reshaped cell is constructed
 * @param[in] lat_frac FRACTIONAL coordinates of lattice vectors of reshaped cell in
 *   basis vectors of primitive cell
 * @param[in] delta small displacement added to orbital positions such that orbitals
 *   fall on cell borders will not be clipped
 * @return reshaped cell created from primitive cell
 */
template <typename hop_t>
PrimitiveCell<hop_t> reshape_prim_cell(
    const PrimitiveCell<hop_t>& prim_cell,
    const Eigen::Matrix3d& lat_frac,
    double delta = 1.0e-5)
{
    // Conversion matrix of fractional coordinates from primitive cell to
    // reshaped cell: x_res = x_prim * conv_mat, with x_res and x_prim
    // being ROW vectors
    Eigen::Matrix3d conv_mat = lat_frac.inverse();

    // Determine searching range
    // sum_vec is actually a0+a1, a1+a2 or a2+a0 depending on j, i.e., the
    // diagonal vector. If it is not taken into consideration, orbitals and
    // hopping terms in the top right corner of reshaped cell may be missing.
    Eigen::Matrix3Xi rn_range = Eigen::Matrix3Xi::Zero(3, 2);
    for (int i_dim = 0; i_dim < 3; ++i_dim) {
        Eigen::Vector3d sum_vec = lat_frac.colwise().sum() - lat_frac.row(i_dim);
        for (int j_dim = 0; j_dim < 3; ++j_dim) {
            cell_index_t rn_min = std::floor(sum_vec[j_dim]);
            cell_index_t rn_max = std::ceil(sum_vec[j_dim]);
            if (rn_range(j_dim, 0) > rn_min) {
                rn_range(j_dim, 0) = rn_min;
            }
            if (rn_range(j_dim, 1) < rn_max) {
                rn_range(j_dim, 1) = rn_max;
            }
        }
    }
    rn_range.col(0) -= Eigen::Vector3i::Ones();
    rn_range.col(1) += Eigen::Vector3i::Ones();

    // Determine number of atoms in reshaped cell
    Eigen::Matrix3Xd orb_pos = prim_cell.get_orbital_positions();
    orbital_index_t num_orb_res = reshape_get_num_orb(prim_cell, conv_mat, delta, orb_pos, rn_range);

    // Reshape cell lattice vectors
    Eigen::Matrix3d lat_cart;
    for (int i_dim = 0; i_dim < 3; ++i_dim) {
        lat_cart.row(i_dim) = lat_frac.row(i_dim) * prim_cell.get_lattice();
    }

    // Reshaped cell shares the same origin
    Eigen::Vector3d origin = prim_cell.get_origin();

    // Create reshaped cell
    PrimitiveCell<hop_t> res_cell(num_orb_res, lat_cart, origin, 1.0);

    // Update orbital attributes
    OrbitalMap orb_map(num_orb_res);
    reshape_upd_orb(prim_cell, conv_mat, delta, orb_pos, rn_range, res_cell, orb_map);

    // Add hopping terms
    reshape_add_hop(prim_cell, conv_mat, delta, orb_pos, res_cell, orb_map);
    return res_cell;
}

/**
 * @brief Rotate and shift primitive cell with respect to z-axis.
 *
 * @tparam hop_t datatype of hopping terms
 * @param[inout] prim_cell primitive cell to rotate
 * @param[in] angle twisting angle in RADIANs, NOT degrees
 * @param[in] center Cartesian coordinates of the rotation center in NANOMETER
 * @param[in] shift  distance of shift in NANOMETER
 *
 * @note This function is reserved for compatibility with Python version.
 */
template <typename hop_t>
void spiral_prim_cell(
    PrimitiveCell<hop_t>& prim_cell,
    const double& angle = 0.0,
    const Eigen::Vector3d& center = Eigen::Vector3d::Zero(),
    const double& shift = 0.0)
{
    prim_cell.rotate(angle, center, shift);
}

/**
 * @brief Get the index of prim_cell in prim_cells.
 * 
 * @tparam hop_t datatype of hopping terms
 * @param prim_cells pointers to prim_cells
 * @param pc pointer to search
 * @return index of given prim_cell
 */
template <typename hop_t>
[[nodiscard]] int64_t index_prim_cell(
    const std::vector<const PrimitiveCell<hop_t>*>& prim_cells,
    const PrimitiveCell<hop_t>* pc)
{
    auto idx_bra = std::find(prim_cells.begin(), prim_cells.end(), pc);
    if (idx_bra == prim_cells.end()) {
        std::cerr << "prim_cell not found!" << std::endl;
        std::exit(-1);
    }
    return std::distance(prim_cells.begin(), idx_bra);
}

/**
 * @brief Merge primitive cells and inter-hopping containers to build a large primitive cell.
 *
 * @tparam hop_t datatype of hopping terms
 * @param[in] prim_cells pointers to primitive cells to merge
 * @param[in] inter_hops pointers to PCInterHopping objects to merge
 * @return merged primitive cell
 */
template <typename hop_t>
PrimitiveCell<hop_t> merge_prim_cell(
    const std::vector<const PrimitiveCell<hop_t>*>& prim_cells,
    const std::vector<const PCInterHopping<hop_t>*>& inter_hops)
{
    // Check lattice mismatch
    Eigen::Matrix3d lat_ref = prim_cells[0]->get_lattice();
    for (const auto& p : prim_cells) {
        Eigen::Matrix3d lat_diff = p->get_lattice() - lat_ref;
        assert(lat_diff.array().abs().maxCoeff() <= 1.0e-5);
    }

    // Get starting indices of orbitals of each component
    std::vector<orbital_index_t> num_orb(prim_cells.size());
    std::vector<orbital_index_t> ind_start(prim_cells.size());
    for (int i = 0; i < prim_cells.size(); ++i) {
        num_orb[i] = prim_cells[i]->get_num_orbitals();
    }
    ind_start[0] = 0;
    for (int i = 1; i < prim_cells.size(); ++i) {
        ind_start[i] = ind_start[i - 1] + num_orb[i - 1];
    }

    // Create merged primitive cell
    // The merge cell shares the same lattice vector and origin as the 0th primitive cell.
    Eigen::Vector3d origin_ref = prim_cells[0]->get_origin();
    orbital_index_t num_orb_total = 0;
    for (const auto& i : num_orb) {
        num_orb_total += i;
    }
    PrimitiveCell<hop_t> merged_cell(num_orb_total, lat_ref, origin_ref, 1.0);

    // Add orbitals
    orbital_index_t counter = 0;
    for (const auto& p : prim_cells) {
        Eigen::Matrix3Xd orb_pos = p->get_orbital_positions_nm();
        orb_pos = tbplas::base::cart2frac(lat_ref, orb_pos, origin_ref);
        auto orbital_list = p->get_all_orbitals();
        for (orbital_index_t i = 0; i < p->get_num_orbitals(); ++i) {
            const Orbital& orbital = orbital_list[i];
            merged_cell.set_orbital(counter, orb_pos(0, i), orb_pos(1, i), orb_pos(2, i),
                orbital.energy, orbital.label);
            ++counter;
        }
    }

    // Add intra-hopping terms
    for (int i = 0; i < prim_cells.size(); ++i) {
        orbital_index_t offset = ind_start[i];
        const auto& hop_data = prim_cells[i]->get_all_hopping().get_data();
        auto iter_rn = hop_data.begin();
        for (; iter_rn != hop_data.end(); iter_rn++) {
            const auto& rn = (*iter_rn).first;
            auto iter_ij = ((*iter_rn).second).begin();
            for (; iter_ij != ((*iter_rn).second).end(); iter_ij++) {
                const auto& pair = (*iter_ij).first;
                const auto& t = (*iter_ij).second;
                cell_index_t ra, rb, rc;
                std::tie(ra, rb, rc) = rn;
                orbital_index_t orb_i = std::get<0>(pair) + offset;
                orbital_index_t orb_j = std::get<1>(pair) + offset;
                merged_cell.add_hopping(ra, rb, rc, orb_i, orb_j, t);
            }
        }
    }

    // Add inter-hopping terms
    for (const auto& p : inter_hops) {
        int64_t idx_bra = index_prim_cell(prim_cells, p->get_pc_bra());
        int64_t idx_ket = index_prim_cell(prim_cells, p->get_pc_ket());
        orbital_index_t offset_bra = ind_start[idx_bra];
        orbital_index_t offset_ket = ind_start[idx_ket];
        const auto& hop_data = p->get_data();
        auto iter_rn = hop_data.begin();
        for (; iter_rn != hop_data.end(); iter_rn++) {
            const auto& rn = (*iter_rn).first;
            auto iter_ij = ((*iter_rn).second).begin();
            for (; iter_ij != ((*iter_rn).second).end(); iter_ij++) {
                const auto& pair = (*iter_ij).first;
                const auto& t = (*iter_ij).second;
                orbital_index_t orb_i = std::get<0>(pair) + offset_bra;
                orbital_index_t orb_j = std::get<1>(pair) + offset_ket;
                cell_index_t ra, rb, rc;
                std::tie(ra, rb, rc) = rn;
                merged_cell.add_hopping(ra, rb, rc, orb_i, orb_j, t);
            }
        }
    }
    return merged_cell;
}

/// @brief Make one layer in the hetero-structure by reshaping primitive cell to given lattice vectors.
/// @tparam hop_t datatype of hopping terms
/// @param prim_cell primitive cell of the layer
/// @param hetero_lattice Cartesian coordinates of hetero-structure lattice vectors in NANOMETER,
/// with each ROW being a vector
/// @param[in] delta small displacement added to orbital positions such that orbitals
/// fall on cell borders will not be clipped
template <typename hop_t>
PrimitiveCell<hop_t> make_hetero_layer(const PrimitiveCell<hop_t>& prim_cell, const Eigen::Matrix3d& hetero_lattice,
    double delta)
{
    // NOTE: the coordinates passed to cart2frac should be in column-wise order, so we need two transposes.
    Eigen::Matrix3d h_cart_t = hetero_lattice.transpose();
    Eigen::Matrix3d h_frac_t = tbplas::base::cart2frac(prim_cell.get_lattice(), h_cart_t);
    Eigen::Matrix3d h_frac = h_frac_t.transpose();
    PrimitiveCell<hop_t> hetero_layer = reshape_prim_cell(prim_cell, h_frac, delta);
    return hetero_layer;
}

/// @brief Find neighbors between the (0, 0, 0) cell of 'bra' primitive cell and nearby cells
/// of 'ket' primitive cell up to given cutoff distance.
/// @param[in] cell_bra 'bra' primitive cell
/// @param[in] cell_ket 'ket' primitive cell
/// @param[in] a_max upper bound of range on a-axis
/// @param[in] b_max upper bound of range on b-axis
/// @param[in] c_max upper bound of range on c-axis
/// @param[in] max_distance cutoff distance in NM
/// @param[in] leaf_max_size maximum leaf size of kdtree
/// @return list of hopping neighbors
/// @note Only neighbors with distance > 0 will be returned.
template <typename hop_t>
std::list<Term> find_neighbors(const PrimitiveCell<hop_t>& cell_bra,
    const PrimitiveCell<hop_t>& cell_ket,
    cell_index_t a_max = 0,
    cell_index_t b_max = 0,
    cell_index_t c_max = 0,
    double max_distance = 1.0,
    int leaf_max_size = 10)
{
    // Assemble kdtree for cell_bra
    orbital_index_t num_orb_bra = cell_bra.get_num_orbitals();
    Eigen::Matrix3Xd pos_bra = cell_bra.get_orbital_positions_nm();
    std::vector<std::vector<double>> pos_bra_tree(num_orb_bra);
    for (orbital_index_t i = 0; i < num_orb_bra; ++i) {
        pos_bra_tree[i] = { pos_bra(0, i), pos_bra(1, i), pos_bra(2, i) };
    }
    VectorOfVectorsAdapter<double> pointCloud(pos_bra_tree);
    kd_tree_t tree_bra(3, pointCloud, nanoflann::KDTreeSingleIndexAdaptorParams(leaf_max_size));
    tree_bra.buildIndex();

    // Call kdtree to find hopping terms
    Eigen::Matrix3Xd pos_ket = cell_ket.get_orbital_positions_nm();
    Eigen::Matrix3d lat_ket = cell_ket.get_lattice();
    std::list<Term> neighbors = find_neighbors_within_radius(tree_bra, pos_bra, pos_ket,
        lat_ket, a_max, b_max, c_max, max_distance);
    return neighbors;
}

/// @brief Class for evaluating hopping integrals using Slater-Koster formula.
/// The maximum supported angular momentum is l=2 (d orbitals).
/// Reference: https://journals.aps.org/pr/abstract/10.1103/PhysRev.94.1498
/// The reason why we make orbital labels as attributes is to avoid misspelling,
/// which is likely to happen as we have to repeat them many times.
class SK {
private:
    std::string s_ = "s";
    std::string px_ = "px";
    std::string py_ = "py";
    std::string pz_ = "pz";
    std::string dxy_ = "dxy";
    std::string dyz_ = "dyz";
    std::string dzx_ = "dzx";
    std::string dx2_y2_ = "dx2-y2";
    std::string dz2_ = "dz2";
    std::set<std::string> p_labels_ = { px_, py_, pz_ };
    std::set<std::string> d_labels_ = { dxy_, dyz_, dzx_, dx2_y2_, dz2_ };
    double sqrt3_ = std::sqrt(3.0);
    double half_sqrt3_ = sqrt3_ * 0.5;

    /// @brief Check the sanity of labels of p orbitals.
    /// @param[in] label labels to check
    /// @warning check label in p_labels_
    void check_p_labels_(const std::string& label);

    /// @brief Check the sanity of labels of d orbitals.
    /// @param[in] label labels to check
    /// @warning check label in d_labels_
    void check_d_labels_(const std::string& label);

    /// @brief Permute a given vector according to the new x_axis.
    /// @param[in] vector vector to permute
    /// @param[in] x_new label of the new x_axis
    /// @return permuted vector
    /// @warning check x_new in 'x', 'y', 'z'
    Eigen::Vector3d perm_vector_(const Eigen::Vector3d& vector, const char x_new);

    /// @brief Remap orbital label after permutation.
    /// @param[in] label orbital label to remap
    /// @param[in] x_new index of the new 'x' direction
    /// @return remapped orbital label
    /// @warning check x_new in 'x', 'y', 'z'
    std::string remap_label_(const std::string& label, const char x_new);

    /// @brief Evaluate direction cosines for given displacement vector.
    /// @param[in] r Cartesian coordinates of the displacement vector
    /// @return the direction cosines along x, y, and z directions
    /// @warning check the norm of r greater than 1e-15
    std::tuple<double, double, double> eval_dir_cos_(const Eigen::Vector3d& r);

public:
    /// @brief Evaluate the hopping integral <s,0|H|s,r>
    /// @param[in] v_sss V_ss_sigma
    /// @return hopping integral
    complex_t ss(const complex_t& v_sss = 1.0);

    /// @brief Evaluate the hopping integral <s,0|H|p,r>.
    /// @param[in] r Cartesian coordinates of the displacement vector
    /// @param[in] label_p label of p orbital
    /// @param[in] v_sps V_sp_sigma
    /// @return hopping integral
    /// @warning check label_p in p_labels_
    complex_t sp(const Eigen::Vector3d& r, const std::string& label_p = "px", const complex_t& v_sps = 0.0);

    /// @brief Evaluate the hopping integral <s,0|H|d,r>.
    /// @param[in] r Cartesian coordinates of the displacement vector
    /// @param[in] label_d label of d orbital
    /// @param[in] v_sds V_sd_sigma
    /// @return hopping integral
    /// @warning check label_d in d_labels_
    complex_t sd(const Eigen::Vector3d& r, const std::string& label_d = "dxy", const complex_t& v_sds = 0.0);

    /// @brief Evaluate the hopping integral <p_i,0|H|p_j,r>.
    /// @param[in] r Cartesian coordinates of the displacement vector
    /// @param[in] label_i label of p_i orbital
    /// @param[in] label_j label of p_j orbital
    /// @param[in] v_pps V_pp_sigma
    /// @param[in] v_ppp  V_pp_pi
    /// @return hopping integral
    /// @warning check label_i in p_labels_ and label_j in p_labels_
    complex_t pp(const Eigen::Vector3d& r, const std::string& label_i = "px", const std::string& label_j = "px", const complex_t& v_pps = 0.0, const complex_t& v_ppp = 0.0);

    /// @brief Evaluate the hopping integral <p,0|H|d,r>.
    /// @param[in] r Cartesian coordinates of the displacement vector
    /// @param[in] label_p label of p orbital
    /// @param[in] label_d label of d orbital
    /// @param[in] v_pds V_pd_sigma
    /// @param[in] v_pdp V_pd_pi
    /// @return hopping integral
    /// @warning check label_i in p_labels_ and label_j in p_labels_
    complex_t pd(const Eigen::Vector3d& r, const std::string& label_p = "px", const std::string& label_d = "dxy", const complex_t& v_pds = 0.0, const complex_t& v_pdp = 0.0);

    /// @brief Evaluate the hopping integral <d_i,0|H|d_j,r>.
    /// @param[in] r Cartesian coordinates of the displacement vector
    /// @param[in] label_i label of d_i orbital
    /// @param[in] label_j label of d_j orbital
    /// @param[in] v_dds  V_dd_sigma
    /// @param[in] v_ddp V_dd_pi
    /// @param[in] v_ddd V_dd_delta
    /// @return hopping integral
    /// @warning check label_i in p_labels_ and label_j in p_labels_
    complex_t dd(const Eigen::Vector3d& r, const std::string& label_i = "dxy", const std::string& label_j = "dxy", const complex_t& v_dds = 0.0, const complex_t& v_ddp = 0.0, const complex_t& v_ddd = 0.0);

    /// @brief Evaluate the hopping integral <p,0|H|s,r>.
    /// @param[in] r Cartesian coordinates of the displacement vector
    /// @param[in] label_p label of p orbital
    /// @param[in] v_sps V_sp_sigma
    /// @return hopping integral
    /// @warning check label_p in p_labels_
    complex_t ps(const Eigen::Vector3d& r, const std::string& label_p = "px", const complex_t& v_sps = 0.0);

    /// @brief Evaluate the hopping integral <d,0|H|s,r>.
    /// @param[in] r Cartesian coordinates of the displacement vector
    /// @param[in] label_d label of d orbital
    /// @param[in] v_sds V_sd_sigma
    /// @return hopping integral
    /// @warning check label_d in d_labels_
    complex_t ds(const Eigen::Vector3d& r, const std::string& label_d = "dxy", const complex_t& v_sds = 0.0);

    /// @brief Evaluate the hopping integral <d,0|H|p,r>.
    /// @param[in] r Cartesian coordinates of the displacement vector
    /// @param[in] label_p label of p orbital
    /// @param[in] label_d label of d orbital
    /// @param[in] v_pds V_pd_sigma
    /// @param[in] v_pdp V_pd_pi
    /// @return hopping integral
    /// @warning check label_p in p_labels_ and label_d in d_labels_
    complex_t dp(const Eigen::Vector3d& r, const std::string& label_p = "px", const std::string& label_d = "dxy", const complex_t& v_pds = 0.0, const complex_t& v_pdp = 0.0);

    /// @brief Common interface for evaluating hopping integral <i,0|H|j,r>.
    /// @param[in] r Cartesian coordinates of the displacement vector
    /// @param[in] label_i label for orbital i
    /// @param[in] label_j label for orbital j
    /// @param[in] v_sss V_ss_sigma
    /// @param[in] v_sps V_sp_sigma
    /// @param[in] v_sds V_sd_sigma
    /// @param[in] v_pps V_pp_sigma
    /// @param[in] v_ppp V_pp_pi
    /// @param[in] v_pds V_pd_sigma
    /// @param[in] v_pdp V_pd_pi
    /// @param[in] v_dds V_dd_sigma
    /// @param[in] v_ddp V_dd_pi
    /// @param[in] v_ddd V_dd_delta
    /// @return hopping integral
    complex_t eval(const Eigen::Vector3d& r, const std::string& label_i = "s", const std::string& label_j = "s",
        const complex_t& v_sss = 0.0, const complex_t& v_sps = 0.0, const complex_t& v_sds = 0.0,
        const complex_t& v_pps = 0.0, const complex_t& v_ppp = 0.0, const complex_t& v_pds = 0.0, const complex_t& v_pdp = 0.0,
        const complex_t& v_dds = 0.0, const complex_t& v_ddp = 0.0, const complex_t& v_ddd = 0.0);
};

/// @brief Hardcoded spin-orbit coupling glossary, including f orbitals
class SOCTable {
private:
    /// @brief  soc terms for up-up spin chanel
    std::map<std::tuple<std::string, std::string>, complex_t> up_up_ = { { { "px", "py" }, -0.5i },
        { { "py", "px" }, 0.5i },
        { { "dxy", "dx2-y2" }, 1.0i },
        { { "dx2-y2", "dxy" }, -1.0i },
        { { "dyz", "dzx" }, 0.5i },
        { { "dzx", "dyz" }, -0.5i },
        { { "fy(3x2-y2)", "fx(x2-3y2)" }, 1.5i },
        { { "fxyz", "fz(x2-y2)" }, 1.0i },
        { { "fyz2", "fxz2" }, 0.5i },
        { { "fxz2", "fyz2" }, -0.5i },
        { { "fz(x2-y2)", "fxyz" }, -1.0i },
        { { "fx(x2-3y2)", "fy(3x2-y2)" }, -1.5i } };

    /// @brief soc terms for up-down spin chanel
    std::map<std::tuple<std::string, std::string>, complex_t> up_down_ = { { { "px", "pz" }, 0.500000000000000 },
        { { "py", "pz" }, -0.5i },
        { { "pz", "px" }, -0.500000000000000 },
        { { "pz", "py" }, 0.5i },
        { { "dxy", "dyz" }, 0.500000000000000 },
        { { "dxy", "dzx" }, -0.5i },
        { { "dx2-y2", "dyz" }, 0.5i },
        { { "dx2-y2", "dzx" }, 0.500000000000000 },
        { { "dyz", "dxy" }, -0.500000000000000 },
        { { "dyz", "dx2-y2" }, -0.5i },
        { { "dyz", "dz2" }, -0.866025403784439i },
        { { "dzx", "dxy" }, 0.5i },
        { { "dzx", "dx2-y2" }, -0.500000000000000 },
        { { "dzx", "dz2" }, 0.866025403784439 },
        { { "dz2", "dyz" }, 0.866025403784439i },
        { { "dz2", "dzx" }, -0.866025403784439 },
        { { "fy(3x2-y2)", "fxyz" }, 0.612372435695794 },
        { { "fy(3x2-y2)", "fz(x2-y2)" }, -0.612372435695794i },
        { { "fxyz", "fy(3x2-y2)" }, -0.612372435695794 },
        { { "fxyz", "fyz2" }, 0.790569415042095 },
        { { "fxyz", "fxz2" }, -0.790569415042095i },
        { { "fxyz", "fx(x2-3y2)" }, -0.612372435695794i },
        { { "fyz2", "fxyz" }, -0.790569415042095 },
        { { "fyz2", "fz3" }, -1.22474487139159i },
        { { "fyz2", "fz(x2-y2)" }, -0.790569415042095i },
        { { "fz3", "fyz2" }, 1.22474487139159i },
        { { "fz3", "fxz2" }, -1.22474487139159 },
        { { "fxz2", "fxyz" }, 0.790569415042095i },
        { { "fxz2", "fz3" }, 1.22474487139159 },
        { { "fxz2", "fz(x2-y2)" }, -0.790569415042095 },
        { { "fz(x2-y2)", "fy(3x2-y2)" }, 0.612372435695794i },
        { { "fz(x2-y2)", "fyz2" }, 0.790569415042095i },
        { { "fz(x2-y2)", "fxz2" }, 0.790569415042095 },
        { { "fz(x2-y2)", "fx(x2-3y2)" }, -0.612372435695794 },
        { { "fx(x2-3y2)", "fxyz" }, 0.612372435695794i },
        { { "fx(x2-3y2)", "fz(x2-y2)" }, 0.612372435695794 } };

    /// @brief soc terms for down-up spin chanel
    std::map<std::tuple<std::string, std::string>, complex_t> down_up_ = { { { "px", "pz" }, -0.500000000000000 },
        { { "py", "pz" }, -0.5i },
        { { "pz", "px" }, 0.500000000000000 },
        { { "pz", "py" }, 0.5i },
        { { "dxy", "dyz" }, -0.500000000000000 },
        { { "dxy", "dzx" }, -0.5i },
        { { "dx2-y2", "dyz" }, 0.5i },
        { { "dx2-y2", "dzx" }, -0.500000000000000 },
        { { "dyz", "dxy" }, 0.500000000000000 },
        { { "dyz", "dx2-y2" }, -0.5i },
        { { "dyz", "dz2" }, -0.866025403784439i },
        { { "dzx", "dxy" }, 0.5i },
        { { "dzx", "dx2-y2" }, 0.500000000000000 },
        { { "dzx", "dz2" }, -0.866025403784439 },
        { { "dz2", "dyz" }, 0.866025403784439i },
        { { "dz2", "dzx" }, 0.866025403784439 },
        { { "fy(3x2-y2)", "fxyz" }, -0.612372435695794 },
        { { "fy(3x2-y2)", "fz(x2-y2)" }, -0.612372435695794i },
        { { "fxyz", "fy(3x2-y2)" }, 0.612372435695794 },
        { { "fxyz", "fyz2" }, -0.790569415042095 },
        { { "fxyz", "fxz2" }, -0.790569415042095i },
        { { "fxyz", "fx(x2-3y2)" }, -0.612372435695794i },
        { { "fyz2", "fxyz" }, 0.790569415042095 },
        { { "fyz2", "fz3" }, -1.22474487139159i },
        { { "fyz2", "fz(x2-y2)" }, -0.790569415042095i },
        { { "fz3", "fyz2" }, 1.22474487139159i },
        { { "fz3", "fxz2" }, 1.22474487139159 },
        { { "fxz2", "fxyz" }, 0.790569415042095i },
        { { "fxz2", "fz3" }, -1.22474487139159 },
        { { "fxz2", "fz(x2-y2)" }, 0.790569415042095 },
        { { "fz(x2-y2)", "fy(3x2-y2)" }, 0.612372435695794i },
        { { "fz(x2-y2)", "fyz2" }, 0.790569415042095i },
        { { "fz(x2-y2)", "fxz2" }, -0.790569415042095 },
        { { "fz(x2-y2)", "fx(x2-3y2)" }, 0.612372435695794 },
        { { "fx(x2-3y2)", "fxyz" }, 0.612372435695794i },
        { { "fx(x2-3y2)", "fz(x2-y2)" }, -0.612372435695794 } };

    /// @brief labels of atomic orbitals
    std::set<std::string> orbital_labels_ = { "s", "px", "py", "pz",
        "dxy", "dx2-y2", "dyz", "dzx", "dz2",
        "fy(3x2-y2)", "fxyz", "fyz2", "fz3", "fxz2",
        "fz(x2-y2)", "fx(x2-3y2)" };

    /// @brief directions of spins
    std::set<std::string> spin_labels_ = { "up", "down" };

public:
    SOCTable() = default;

    /// @brief Evaluate the matrix element <i,s_i|l*s|j,s_j>.
    /// @param label_i orbital label of bra
    /// @param spin_i spin direction of bra
    /// @param label_j orbital label of ket
    /// @param spin_j spin direction of ket
    /// @return turn: matrix element in h_bar**2
    complex_t eval(const std::string& label_i = "s", const std::string& spin_i = "up",
        const std::string& label_j = "s", const std::string& spin_j = "down");
};
} // namespace tbplas::builder
#endif // TBPLAS_BUILDER_ADVANCED_H