/**
 * @brief Functions and classes for constructing samples.
 *
 * @author Yunhai Li (liyunhai1016@hotmail.com)
 *
 * @copyright Copyright (c) 2024, TBPLaS develop team.
 */

#ifndef TBPLAS_BUILDER_SAMPLE_H
#define TBPLAS_BUILDER_SAMPLE_H

#include <iterator>
#include <memory>

#include "base/lattice.h"
#include "builder/filters.h"
#include "super.h"

namespace tbplas::builder {
using base::cart2frac;
using base::imag_part;
using base::real_part;

/**
 * @brief Container of supercells and sc_hops.
 *
 * @tparam hop_t data type of hopping terms
 */
template <typename hop_t>
class Sample {
public:
    Sample() = default;

    Sample(const Sample& other)
        : super_cells_(other.super_cells_)
        , sc_hops_(other.sc_hops_)
    {
        // orb_data_ and hop_data will be generated rather than copied.
        init_array();
    }

    Sample(Sample&& other) noexcept
        : super_cells_(std::move(other.super_cells_))
        , sc_hops_(std::move(other.sc_hops_))
    {
        // Clear other to reduce memory usage
        other.clear_array();
        // orb_data_ and hop_data will be generated rather than copied.
        init_array();
    }

    Sample& operator=(const Sample& other)
    {
        if (this != &other) {
            super_cells_ = other.super_cells_;
            sc_hops_ = other.sc_hops_;
            init_array();
        }
        return *this;
    }

    Sample& operator=(Sample&& other) noexcept
    {
        if (this != &other) {
            super_cells_ = std::move(other.super_cells_);
            sc_hops_ = std::move(other.sc_hops_);
            other.clear_array();
            init_array();
        }
        return *this;
    }

    ~Sample() = default;

    /**
     * @brief Add a new supercell to the sample.
     *
     * @param[in] prim_cell primitive cell from which supercell is built
     * @param[in] sc_lattice Cartesian coordinates of supercell lattice in NM,
     *   with each row being a vector
     * @param[in] pbc boundary condition along 3 directions, true for pbc and
     *   false for non-pbc
     * @param[in] delta small displacement added to vector to avoid clipping
     * @return shared_ptr to the created supercell
     *
     * @note See the constructor for SuperCell for more details
     */
    std::shared_ptr<SuperCell<hop_t>>
    add_super_cell(
        const PrimitiveCell<hop_t>& prim_cell,
        const Eigen::Matrix3d& sc_lattice,
        const pbc_t& pbc,
        const double& delta = 1.0e-5)
    {
        super_cells_.push_back(std::make_shared<SuperCell<hop_t>>(prim_cell, sc_lattice, pbc, delta));
        return super_cells_.back();
    }

    /**
     * @brief Add a new supercell to the sample.
     *
     * @param[in] prim_cell primitive cell from which supercell is built
     * @param[in] dim dimension of supercell along 3 directions
     * @param[in] pbc boundary condition along 3 directions, true for pbc and
     *   false for non-pbc
     * @param[in] delta small displacement added to vector to avoid clipping
     * @return shared_ptr to the created supercell
     */
    std::shared_ptr<SuperCell<hop_t>>
    add_super_cell(
        const PrimitiveCell<hop_t>& prim_cell,
        const dim_t& dim,
        const pbc_t& pbc,
        const double& delta = 1.0e-5)
    {
        super_cells_.push_back(std::make_shared<SuperCell<hop_t>>(prim_cell, dim, pbc, delta));
        return super_cells_.back();
    }

    /**
     * @brief Add a new SCHopping object to the sample.
     *
     * @param[in] sc_bra 'bra' supercell
     * @param[in] sc_ket 'ket' supercell
     * @param[in] max_distance cutoff distance in NM
     * @param[in] default_hop default hopping parameter in eV
     * @return shared_ptr to the created SCHopping object
     */
    std::shared_ptr<SCHopping<hop_t>>
    add_sc_hopping(
        const std::shared_ptr<SuperCell<hop_t>>& sc_bra,
        const std::shared_ptr<SuperCell<hop_t>>& sc_ket,
        const double& max_distance = 1.0,
        const hop_t& default_hop = 1.0)
    {
        sc_hops_.push_back(std::make_shared<SCHopping<hop_t>>(sc_bra, sc_ket, max_distance, default_hop));
        return sc_hops_.back();
    }

    /**
     * @brief Get supercell according to index
     *
     * @param[in] idx index of supercell
     * @return shared_ptr to the supercell
     */
    [[nodiscard]] std::shared_ptr<SuperCell<hop_t>>
    get_super_cell(const size_t& idx = 0)
    {
        assert(0 <= idx && idx < super_cells_.size());
        return super_cells_[idx];
    }

    /**
     * @brief Get sc_hop according to index
     *
     * @param[in] idx index of sc_hop
     * @return shared_ptr to the sc_hop
     */
    [[nodiscard]] std::shared_ptr<SCHopping<hop_t>>
    get_sc_hopping(const size_t& idx = 0)
    {
        assert(0 <= idx && idx < sc_hops_.size());
        return sc_hops_[idx];
    }

    /**
     * @brief Get the number of orbitals
     *
     * @return number of orbitals
     *
     * @note This function must be called after calling 'build_orb_data'.
     */
    [[nodiscard]] orbital_index_t get_num_orbitals() const
    {
        return orb_data_.back().offset + orb_data_.back().get_num_orbitals();
    }

    /**
     * @brief Get the number of hopping terms
     *
     * @return  number of hopping terms
     *
     * @note This function must be called after calling 'build_hop_data'.
     */
    [[nodiscard]] orbital_index_t get_num_hopping() const
    {
        return hop_data_.back().offset + hop_data_.back().get_num_hopping();
    }

    /**
     * @brief Get lattice vectors
     *
     * @return Cartesian coordinates of lattice vectors in NM
     *
     * @note Each ROW of lattice is a lattice vector.
     */
    [[nodiscard]] const Eigen::Matrix3d& get_lattice() const
    {
        return super_cells_[0]->get_sc_lattice();
    }

    /**
     * @brief Get lattice area
     *
     * @param[in] direction direction along which the area is evaluated, should be in ("a", "b", "c")
     * @return lattice area in nm^2
     */
    [[nodiscard]] double get_lattice_area(const std::string& direction = "c") const
    {
        return tbplas::base::get_lattice_area(get_lattice(), direction);
    }

    /**
     * @brief Get the lattice volume attribute
     *
     * @return lattice volume in nm^3
     */
    [[nodiscard]] double get_lattice_volume() const
    {
        return tbplas::base::get_lattice_volume(get_lattice());
    }

    /**
     * @brief Get reciprocal lattice vectors
     *
     * @return Cartesian coordinates of reciprocal lattice vectors in 1/NM.
     */
    [[nodiscard]] Eigen::Matrix3d get_reciprocal_lattice() const
    {
        return tbplas::base::gen_reciprocal_vectors(get_lattice());
    }

    /**
     * @brief Build the attributes of sparse Hamiltonian and current operator
     * in CSR format.
     *
     * @param[out] indptr the 'indptr' attribute (outerIndex of eigen)
     * @param[out] indices the 'indices' attribute (innerIndex of eigen)
     * @param[out] ham non-zero elements of Hamiltonian in eV
     * @param[out] curr_coeff_x coefficients of current operator along along x direction
     * @param[out] curr_coeff_y coefficients of current operator along along x direction
     *
     * @note See eqn. 54 of CPC paper for the formula of current operator.
     */
    void build_ham_curr_csr(
        std::vector<orbital_index_t>& indptr,
        std::vector<orbital_index_t>& indices,
        Eigen::VectorXcd& ham,
        Eigen::VectorXcd& curr_coeff_x,
        Eigen::VectorXcd& curr_coeff_y) const
    {
        // Allocate indices and data
        build_indptr(indptr);
        orbital_index_t num_orb = indptr.size() - 1;
        orbital_index_t num_hop_full = indptr.back();
        indices = std::vector<orbital_index_t>(num_hop_full);
        ham = Eigen::VectorXcd(num_hop_full);
        curr_coeff_x = Eigen::VectorXcd(num_hop_full);
        curr_coeff_y = Eigen::VectorXcd(num_hop_full);
        std::vector<orbital_index_t> count(num_orb, 0);

        // Add on-site terms
        // WARNING:
        // 1. IT MUST BE count[i]++, not ++count[i]!
        // 2. DO NOT forget to set curr_coeff_x and curr_coeff_y, otherwise the diagonal terms
        //    will be undefined. This bug has caused us hours to debug.
        for (const auto& d : orb_data_) {
            orbital_index_t num_orb_this = d.get_num_orbitals();
            for (orbital_index_t i = 0; i < num_orb_this; ++i) {
                orbital_index_t overall_idx = i + d.offset;
                orbital_index_t ptr = indptr[overall_idx] + count[overall_idx]++;
                indices[ptr] = overall_idx;
                ham[ptr] = d.orb_eng[i];
                curr_coeff_x[ptr] = 0.0;
                curr_coeff_y[ptr] = 0.0;
            }
        }

        // Orbital indices in d.hop_ind already contains offset. So there's no need
        // for adding offset to idx_i and idx_j.
        orbital_index_t num_hop_reduced = get_num_hopping();
        for (const auto& d : hop_data_) {
            orbital_index_t num_hop_this = d.get_num_hopping();
            for (orbital_index_t i = 0; i < num_hop_this; ++i) {
                orbital_index_t idx_i = std::get<0>(d.hop_ind[i]);
                orbital_index_t idx_j = std::get<1>(d.hop_ind[i]);
                complex_t hij = d.hop_eng[i];
                Eigen::Vector3d dr = d.dr.col(i);

                // Fill vectors
                // Original part
                orbital_index_t ptr_i = indptr[idx_i] + count[idx_i]++;
                orbital_index_t ptr_j = indptr[idx_j] + count[idx_j]++;
                indices[ptr_i] = idx_j;
                ham[ptr_i] = hij;
                complex_t factor = -base::UIMAG * hij;
                curr_coeff_x[ptr_i] = factor * dr[0];
                curr_coeff_y[ptr_i] = factor * dr[1];

                // Conjugate part
                indices[ptr_j] = idx_i;
                ham[ptr_j] = base::conj_t(hij);
                factor = tbplas::base::UIMAG * tbplas::base::conj_t(hij);
                curr_coeff_x[ptr_j] = factor * dr[0];
                curr_coeff_y[ptr_j] = factor * dr[1];
            }
        }
    }

    /**
     * @brief Build the attributes of sparse Hamiltonian for given k-point
     * in CSR format.
     *
     * @tparam csr_idx_t datatype of CSR matrix indices
     * @param[out] indptr the 'indptr' attribute (outerIndex of eigen)
     * @param[out] indices the 'indices' attribute (innerIndex of eigen)
     * @param[out] ham non-zero elements of Hamiltonian in eV
     * @param[in] k_point fractional coordinate of k-point
     * @param[in] convention convention of Hamiltonian
     *
     * @note
     * 1. In most cases we need only the Hamiltonian at Gamma point. So we make
     *    it that if convention != 1 or 2, k_point will be discarded and Gamma point
     *    is assumed.
     * 2. The FEAST library requires the indices to be int32_t, while orbital_index_t
     *    defaults to uint64_t. So we make it a template argument.
     */
    template <typename csr_idx_t>
    void build_ham_csr(
        std::vector<csr_idx_t>& indptr,
        std::vector<csr_idx_t>& indices,
        Eigen::VectorXcd& ham,
        const Eigen::Vector3d& k_point = Eigen::Vector3d::Zero(3),
        int convention = -1) const
    {
        // Convention 2 is not supported since we do not have rn.
        assert(convention != 2);

        // Allocate indices and data
        build_indptr(indptr);
        orbital_index_t num_orb = indptr.size() - 1;
        orbital_index_t num_hop_full = indptr.back();
        indices = std::vector<csr_idx_t>(num_hop_full);
        ham = Eigen::VectorXcd(num_hop_full);
        std::vector<orbital_index_t> count(num_orb, 0);

        // Add on-site terms
        // WARNING: IT MUST BE count[i]++, not ++count[i]!
        for (const auto& d : orb_data_) {
            orbital_index_t num_orb_this = d.get_num_orbitals();
            for (orbital_index_t i = 0; i < num_orb_this; ++i) {
                orbital_index_t overall_idx = i + d.offset;
                orbital_index_t ptr = indptr[overall_idx] + count[overall_idx]++;
                indices[ptr] = overall_idx;
                ham[ptr] = d.orb_eng[i];
            }
        }

        // Convert kpt from fractional to Cartesian
        Eigen::Matrix3d recip_vec = get_reciprocal_lattice();
        Eigen::Vector3d k_cart = (k_point.transpose() * recip_vec).transpose();

        // Add hopping terms
        // Orbital indices in d.hop_ind already contains offset.
        orbital_index_t num_hop_reduced = get_num_hopping();
        for (const auto& d : hop_data_) {
            orbital_index_t num_hop_this = d.get_num_hopping();
            for (orbital_index_t i = 0; i < num_hop_this; ++i) {
                orbital_index_t idx_i = std::get<0>(d.hop_ind[i]);
                orbital_index_t idx_j = std::get<1>(d.hop_ind[i]);
                complex_t hij = d.hop_eng[i];
                Eigen::Vector3d dr = d.dr.col(i);

                // Add phase factor for convention 1
                if (convention == 1) {
                    double phase = k_cart.dot(dr);
                    hij *= tbplas::base::exp_it(phase);
                }

                // Fill vectors
                // Original part
                orbital_index_t ptr_i = indptr[idx_i] + count[idx_i]++;
                orbital_index_t ptr_j = indptr[idx_j] + count[idx_j]++;
                indices[ptr_i] = idx_j;
                ham[ptr_i] = hij;

                // Conjugate part
                indices[ptr_j] = idx_i;
                ham[ptr_j] = base::conj_t(hij);
            }
        }
    }

    /**
     * @brief Build density operator coefficients for given q-point
     *
     * @param[in] q_point FRACTIONAL coordinate of q-point
     * @param[out] coeff_q coefficients for q point
     * @param[out] coeff_mq coefficients for -q point
     *
     * @note
     * 1. The formula is eqn. 78 of CPC paper.
     * 2. There's no need to consider lattice origin since it is included in orb_pos.
     */
    void build_density_coeff(
        const Eigen::Vector3d& q_point,
        Eigen::VectorXcd& coeff_q,
        Eigen::VectorXcd& coeff_mq) const
    {
        orbital_index_t num_orb = get_num_orbitals();
        coeff_q = Eigen::VectorXcd(num_orb);
        coeff_mq = Eigen::VectorXcd(num_orb);

        complex_t phase = 0.0;
        Eigen::Matrix3d recip_vec = get_reciprocal_lattice();
        Eigen::Vector3d q_cart = (q_point.transpose() * recip_vec).transpose();

        for (const auto& d : orb_data_) {
            orbital_index_t num_orb_this = d.get_num_orbitals();
#pragma omp parallel for private(phase)
            for (orbital_index_t i = 0; i < num_orb_this; ++i) {
                phase = base::exp_it(q_cart.dot(d.orb_pos.col(i)));
                coeff_q[d.offset + i] = phase;
                coeff_mq[d.offset + i] = base::conj_t(phase);
            }
        }
    }

    /**
     * @brief Initialize the array attributes.
     *
     * @param[in] force_init whether to force initializing the arrays from scratch
     *   even if they have already been initialized
     *
     * @note This method should be explicitly called by the user, since the get_*
     *   and build_* methods cannot have the 'const' qualifier if they call this
     *   method automatically, which will eventually cause all solvers based on
     *   this class fail to work.
     */
    void init_array(const bool& force_init = false)
    {
        if (force_init || orb_data_.empty() || hop_data_.empty()) {
            check_sanity();
            build_orb_data();
            build_hop_data();
        }

        // // Concatenate orbital arrays
        // orbital_index_t num_orb_total = orb_data_.back().offset + orb_data_.back().get_num_orbitals();
        // orb_pos_ = Eigen::Matrix3Xd(3, num_orb_total);
        // orb_eng_ = Eigen::VectorXd(num_orb_total);
        // for (const auto& d : orb_data_) {
        //     orbital_index_t num_orb_this = d.get_num_orbitals();
        //     std::memcpy(orb_pos_.data() + 3 * d.offset, d.orb_pos.data(), 3 * num_orb_this * sizeof(double));
        //     std::memcpy(orb_eng_.data() + d.offset, d.orb_eng.data(), num_orb_this * sizeof(double));
        // }

        // // Concatenate hopping arrays
        // orbital_index_t num_hop_total = hop_data_.back().offset + hop_data_.back().get_num_hopping();
        // hop_ind_ = std::vector<std::tuple<orbital_index_t, orbital_index_t>>();
        // hop_eng_ = Eigen::VectorXcd(num_hop_total);
        // dr_ = Eigen::Matrix3Xd(3, num_hop_total);
        // for (const auto& d : hop_data_) {
        //     orbital_index_t num_hop_this = d.get_num_hopping();
        //     hop_ind_.insert(hop_ind_.end(), d.hop_ind.begin(), d.hop_ind.end());
        //     std::memcpy(hop_eng_.data() + d.offset, d.hop_eng.data(), num_hop_this * sizeof(hop_t));
        //     std::memcpy(dr_.data() + 3 * d.offset, d.dr.data(), 3 * num_hop_this * sizeof(double));
        // }
    }

    /**
     * @brief Reset all array attributes.
     *
     */
    void reset_array()
    {
        init_array(true);
    }

    /**
     * @brief Clear all array attributes.
     *
     */
    void clear_array()
    {
        orb_data_.clear();
        hop_data_.clear();
    }

    /**
     * @brief Save sample to file.
     *
     * @param[in] file_name name of file
     */
    void save(const std::string& file_name = "sample") const
    {
        const int prec = 16;
        const int width = 22;
        const int width_int = 6;
        std::ofstream outfile;
        outfile.open(file_name, std::ios::out);
        outfile.setf(std::ios::fixed);
        outfile.precision(prec);

        // lattice vectors
        outfile << "begin_lattice\n";
        Eigen::Matrix3d lattice = super_cells_[0]->get_sc_lattice();
        for (int i = 0; i < 3; ++i) {
            for (int j = 0; j < 3; ++j) {
                outfile << std::setw(width) << lattice(i, j);
            }
            outfile << "\n";
        }
        outfile << "end_lattice\n";

        outfile << "begin_origin\n";
        Eigen::Vector3d origin = super_cells_[0]->get_sc_origin();
        for (int i = 0; i < 3; ++i) {
            outfile << std::setw(width) << origin(i);
        }
        outfile << "\nend_origin\n";

        // orbitals
        outfile << "begin_orbitals\n";
        for (const auto& d : orb_data_) {
            for (orbital_index_t i = 0; i < d.orb_pos.cols(); ++i) {
                outfile << std::setw(width) << d.orb_pos(0, i)
                        << std::setw(width) << d.orb_pos(1, i)
                        << std::setw(width) << d.orb_pos(2, i)
                        << std::setw(width) << d.orb_eng[i]
                        << std::setw(width) << d.orb_label[i]
                        << "\n";
            }
        }
        outfile << "end_orbitals\n";

        // hopping terms
        outfile << "begin_hopping\n";
        for (const auto& d : hop_data_) {
            for (orbital_index_t i = 0; i < d.get_num_hopping(); ++i) {
                outfile << std::setw(width_int) << std::get<0>(d.hop_ind[i])
                        << std::setw(width_int) << std::get<1>(d.hop_ind[i])
                        << std::setw(width) << real_part(d.hop_eng[i])
                        << std::setw(width) << imag_part(d.hop_eng[i])
                        << "\n";
            }
        }
        outfile << "end_hopping\n";

        // hopping distances
        outfile << "begin_dr\n";
        for (const auto& d : hop_data_) {
            for (orbital_index_t i = 0; i < d.get_num_hopping(); ++i) {
                outfile << std::setw(width) << d.dr(0, i)
                        << std::setw(width) << d.dr(1, i)
                        << std::setw(width) << d.dr(2, i)
                        << "\n";
            }
        }
        outfile << "end_dr\n";
        outfile.close();
    }

    /**
     * @brief Apply magnetic field perpendicular to xOy-plane to -z direction via Peierls substitution.
     *
     * @param[in] intensity magnetic B field in Tesla
     * @param[in] gauge gauge of vector potential of the magnetic field to -z\n
     *   0 for (By, 0, 0), 1 for (0, -Bx, 0), 2 for (0.5By, -0.5Bx, 0)
     */
    void apply_magnetic_field(double intensity, int gauge = 0)
    {
        init_array();
        size_t num_sc_hops = sc_hops_.size();
        for (size_t i = 0; i < num_sc_hops; ++i) {
            std::shared_ptr<SCHopping<hop_t>> hop = sc_hops_[i];
            int64_t idx_bra = index_super_cell(hop->get_sc_bra().lock());
            int64_t idx_ket = index_super_cell(hop->get_sc_ket().lock());
            const OrbitalData& data_bra = orb_data_[idx_bra];
            const OrbitalData& data_ket = orb_data_[idx_ket];
            HoppingData& hop_data = hop_data_[i]; // will be updated so no const qualifier

            double factor = base::PEIERLS_FACTOR;
            const auto& pos_bra = data_bra.orb_pos;
            const auto& pos_ket = data_ket.orb_pos;
            const auto& offset_bra = data_bra.offset;
            const auto& offset_ket = data_ket.offset;
            const auto& hop_ind = hop_data.hop_ind;
            auto& hop_eng = hop_data.hop_eng; // will be updated so no const qualifier
            const auto& dr = hop_data.dr;
            orbital_index_t orb_i = 0;
            orbital_index_t orb_j = 0;
            for (size_t i = 0; i < hop_data.get_num_hopping(); ++i) {
                std::tie(orb_i, orb_j) = hop_ind[i];
                double dx = dr(0, i);
                double dy = dr(1, i);
                double sx = pos_ket(0, orb_j - offset_ket) + pos_bra(0, orb_i - offset_bra);
                double sy = pos_ket(1, orb_j - offset_ket) + pos_bra(1, orb_i - offset_bra);
                double phase = 0.0;
                if (gauge == 0) {
                    phase = factor * intensity * dx * sy;
                } else if (gauge == 1) {
                    phase = -factor * intensity * dy * sx;
                } else {
                    phase = factor * 0.5 * intensity * (dx * sx - dy * sy);
                }
                hop_eng[i] *= base::exp_it(phase);
            }
        }
    }

private:
    /// @brief list of supercells in the sample
    std::vector<std::shared_ptr<SuperCell<hop_t>>> super_cells_;

    /// @brief listof SCHopping objects in the sample
    std::vector<std::shared_ptr<SCHopping<hop_t>>> sc_hops_;

    /// @brief orbital data containers produced by supercells
    std::vector<OrbitalData> orb_data_;

    /// @brief hopping data containers produced by SCHopping objects
    std::vector<HoppingData> hop_data_;

    // /// @brief arrays created by assembling orb_data_ and hop_data_
    // Eigen::Matrix3Xd orb_pos_;
    // Eigen::VectorXd orb_eng_;
    // std::vector<std::tuple<orbital_index_t, orbital_index_t>> hop_ind_;
    // Eigen::VectorXcd hop_eng_;
    // Eigen::Matrix3Xd dr_;

    /**
     * @brief Check the 'sc_bra' and 'sc_ket' attributes of SCHopping objects
     * are double counted.
     *
     */
    void check_sanity() const
    {
        using sc_t = std::shared_ptr<SuperCell<hop_t>>;
        std::map<std::tuple<sc_t, sc_t>, int> counter;
        for (const auto& h : sc_hops_) {
            sc_t sc_bra = h->get_sc_bra().lock();
            sc_t sc_ket = h->get_sc_ket().lock();
            if (sc_bra <= sc_ket) {
                counter[std::make_tuple(sc_bra, sc_ket)] += 1;
            } else {
                // Convert to upper-triangular form
                counter[std::make_tuple(sc_ket, sc_bra)] += 1;
            }
        }
        auto iter_rn = counter.begin();
        for (; iter_rn != counter.end(); iter_rn++) {
            std::tuple<sc_t, sc_t> pair = iter_rn->first;
            int num_times = iter_rn->second;
            if (num_times > 1) {
                sc_t sc_bra = std::get<0>(pair);
                sc_t sc_ket = std::get<1>(pair);
                std::cerr << "sc pair has been counted twice: " << sc_bra << "," << sc_ket << std::endl;
                std::exit(-1);
            }
        }
    }

    /**
     * @brief Get the index of supercell in super_cells_.
     *
     * @param[in] sc pointer of supercell to index
     * @return index of the supercell
     */
    [[nodiscard]] int64_t index_super_cell(const std::shared_ptr<SuperCell<hop_t>>& sc) const
    {
        auto idx_bra = std::find(super_cells_.begin(), super_cells_.end(), sc);
        if (idx_bra == super_cells_.end()) {
            std::cerr << "supercell not found!" << std::endl;
            std::exit(-1);
        }
        return std::distance(super_cells_.begin(), idx_bra);
    }

    /**
     * @brief Build the orb_data_ attribute.
     *
     */
    void build_orb_data()
    {
        orb_data_ = std::vector<OrbitalData>(super_cells_.size());
        for (orbital_index_t i = 0; i < orb_data_.size(); ++i) {
            super_cells_[i]->build_orbital_data(orb_data_[i]);
        }
        orb_data_[0].offset = 0;
        for (orbital_index_t i = 1; i < orb_data_.size(); ++i) {
            orb_data_[i].offset = orb_data_[i - 1].offset + orb_data_[i - 1].get_num_orbitals();
        }
    }

    /**
     * @brief Build the hop_data_ attribute.
     *
     */
    void build_hop_data()
    {
        hop_data_ = std::vector<HoppingData>(sc_hops_.size());
        for (orbital_index_t i = 0; i < sc_hops_.size(); ++i) {
            std::shared_ptr<SCHopping<hop_t>> hop = sc_hops_[i];
            int64_t idx_bra = index_super_cell(hop->get_sc_bra().lock());
            int64_t idx_ket = index_super_cell(hop->get_sc_ket().lock());
            hop->build_hopping_data(orb_data_[idx_bra], orb_data_[idx_ket], hop_data_[i]);
        }
        hop_data_[0].offset = 0;
        for (orbital_index_t i = 1; i < hop_data_.size(); ++i) {
            hop_data_[i].offset = hop_data_[i - 1].offset + hop_data_[i - 1].get_num_hopping();
        }
    }

    /**
     * @brief Build the 'indptr' attribute of sparse Hamiltonian and hopping distances
     * in CSR format.
     *
     * @tparam csr_idx_t datatype of CSR matrix indices
     * @param[out] indptr the 'indptr' attribute
     *
     * @note The FEAST library requires the indices to be int32_t, while orbital_index_t
     *   defaults to uint64_t. So we make it a template argument.
     */
    template <typename csr_idx_t>
    void build_indptr(std::vector<csr_idx_t>& indptr) const
    {
        orbital_index_t num_orbs = get_num_orbitals();
        indptr = std::vector<csr_idx_t>(num_orbs + 1, 0);

        // On-site terms
        for (orbital_index_t i = 0; i < num_orbs; ++i) {
            indptr[i + 1] += 1;
        }

        // Hopping terms
        // Orbital indices in h.hop_ind already contain offset.
        orbital_index_t orb_i = 0;
        orbital_index_t orb_j = 0;
        for (const auto& h : hop_data_) {
            for (const auto& p : h.hop_ind) {
                std::tie(orb_i, orb_j) = p;
                indptr[orb_i + 1] += 1;
                indptr[orb_j + 1] += 1;
            }
        }

        // Sum up
        for (orbital_index_t i = 1; i <= num_orbs; ++i) {
            indptr[i] += indptr[i - 1];
        }
    }
};
} // namespace tbplas::builder
#endif