/**
 * @brief Functions and classes for Python interface.
 *
 * @author Junyi Wang (jywang)
 *
 * @copyright Copyright (c) 2024, TBPLaS develop team.
 */

#ifndef TBPLAS_PYTHON_PYMODEL_H
#define TBPLAS_PYTHON_PYMODEL_H

#include <fstream>
#include <iomanip>
#include <iostream>
#include <vector>

#include <eigen3/Eigen/Dense>

#include "base/consts.h"
#include "base/datatypes.h"
#include "base/funcs.h"
#include "base/lattice.h"
#include "builder/base.h"

namespace tbplas::python {
using tbplas::base::complex_t;
using tbplas::builder::cell_index_t;
using tbplas::builder::orbital_index_t;

/**
 * @brief Deep copy X to Y.
 *
 * @tparam vx_t datatype of vector x
 * @tparam vy_t datatype of vector y
 * @param[in] x pointer to data of vector x
 * @param[out] y vector y
 * @param[in] dim dimension of x
 */
template <typename vx_t, typename vy_t>
void vec_copy(const vx_t* x, Eigen::MatrixBase<vy_t>& y, size_t dim)
{
    assert(dim == y.size());
#pragma omp parallel for simd
    for (size_t i = 0; i < dim; ++i) {
        y[i] = x[i];
    }
}

/**
 * @brief Base class for PyPrimitiveCell and PySample.
 *
 * @tparam hop_t data type of hopping terms
 */
template <typename hop_t>
class PyBaseModel {
protected:
    /// @brief number or orbitals
    orbital_index_t num_orb_ = 0;

    /// @brief number of hopping terms
    orbital_index_t num_hop_ = 0;

    /// @brief Cartesian coordinates of lattice vectors in NM
    /// @note Each ROW corresponds to one lattice vector.
    Eigen::Matrix3d lattice_ = Eigen::MatrixXd::Identity(3, 3);

public:
    /**
     * @brief Construct a new PyBaseModel object
     *
     * @param[in] num_orb number or orbitals
     * @param[in] num_hop number of hopping terms
     * @param[in] lattice pointer to lattice vectors
     */
    explicit PyBaseModel(
        uint64_t num_orb,
        uint64_t num_hop,
        const double* lattice)
        : num_orb_(num_orb)
        , num_hop_(num_hop)
    {
        // Data order in lattice
        // a0x, a0y, a0z, a1x, a1y, a1z, a2x, a2y, a2z
        int offset = 0;
        for (int i = 0; i < 3; ++i) {
            for (int j = 0; j < 3; ++j) {
                lattice_(i, j) = lattice[offset];
                ++offset;
            }
        }
    }

    /**
     * @brief Get the number of orbitals
     *
     * @return number of orbitals
     */
    [[nodiscard]] orbital_index_t get_num_orbitals() const
    {
        return num_orb_;
    }

    /**
     * @brief Get the number of hopping terms
     *
     * @return  number of hopping terms
     */
    [[nodiscard]] orbital_index_t get_num_hopping() const
    {
        return num_hop_;
    }

    /**
     * @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 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(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(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(lattice_);
    }
};

/**
 * @brief Interface for Python PrimitiveCell class.
 *
 * @tparam hop_t data type of hopping terms
 */
template <typename hop_t>
class PyPrimitiveCell : public PyBaseModel<hop_t> {
private:
    /// @brief pointer to lattice origin
    const double* origin_ = nullptr;

    /// @brief pointer to data of orbital energies in eV
    const double* orb_eng_ = nullptr;

    /// @brief pointer to data of fractional orbital positions
    /// @note x, y, z components of one atom is continuous in RAM
    const double* orb_pos_ = nullptr;

    /// @brief pointer to hopping indices
    const cell_index_t* hop_ind_ = nullptr;

    /// @brief pointer to hopping energies in eV
    const hop_t* hop_eng_ = nullptr;

    /**
     * @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 = PyBaseModel<hop_t>::num_orb_;
        orbital_index_t num_hops = PyBaseModel<hop_t>::num_hop_;
        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
        for (orbital_index_t i = 0; i < num_hops; ++i) {
            orbital_index_t offset = i * 5;
            orbital_index_t orb_i = hop_ind_[offset + 3];
            orbital_index_t orb_j = hop_ind_[offset + 4];
            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];
        }
    }

    /**
     * @brief Get map to orbital positions
     *
     * @return map to orbital positions
     */
    [[nodiscard]] Eigen::Map<Eigen::Matrix3Xd> get_orb_pos_map() const
    {
        return Eigen::Map<Eigen::Matrix3Xd>(const_cast<double*>(orb_pos_), 3, PyBaseModel<hop_t>::num_orb_);
    }

    /**
     * @brief Get map to origin
     *
     * @return map to origin
     */
    [[nodiscard]] Eigen::Map<Eigen::Vector3d> get_origin_map() const
    {
        return Eigen::Map<Eigen::Vector3d>(const_cast<double*>(origin_));
    }

public:
    /**
     * @brief Construct a new PyPrimitiveCell object
     *
     * @param[in] num_orb number or orbitals
     * @param[in] num_hop number of hopping terms
     * @param[in] orb_eng pointer to data of orbital energies
     * @param[in] orb_pos pointer to data of orbital position
     * @param[in] hop_ind pointer to hopping indices
     * @param[in] hop_eng pointer to hopping energies in eV
     * @param[in] lattice pointer to lattice vectors
     * @param[in] origin pointer to lattice origin
     */
    explicit PyPrimitiveCell(
        orbital_index_t num_orb,
        orbital_index_t num_hop,
        const double* orb_eng,
        const double* orb_pos,
        const cell_index_t* hop_ind,
        const hop_t* hop_eng,
        const double* lattice,
        const double* origin)
        : PyBaseModel<hop_t>(num_orb, num_hop, lattice)
        , origin_(origin)
        , orb_eng_(orb_eng)
        , orb_pos_(orb_pos)
        , hop_ind_(hop_ind)
        , hop_eng_(hop_eng)
    {
    }

    /**
     * @brief Save model to file
     *
     * @param[in] file_name file name of model
     */
    void save(const std::string& file_name) const
    {
        constexpr int prec = 16;
        constexpr int width = 22;
        constexpr 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";
        for (int i = 0; i < 3; ++i) {
            for (int j = 0; j < 3; ++j) {
                outfile << std::setw(width) << PyBaseModel<hop_t>::lattice_(i, j);
            }
            outfile << "\n";
        }
        outfile << "end_lattice\n";

        // PySample class uses Cartesian coordinates and origin is not essential.
        // So we make it (0, 0, 0) in saved model.
        outfile << "begin_origin\n";
        for (int i = 0; i < 3; ++i) {
            outfile << std::setw(width) << origin_[i];
        }
        outfile << "\nend_origin\n";

        // orbitals
        outfile << "begin_orbitals\n";
        for (uint64_t i = 0; i < PyBaseModel<hop_t>::num_orb_; ++i) {
            outfile << std::setw(width) << orb_pos_[i * 3]
                    << std::setw(width) << orb_pos_[i * 3 + 1]
                    << std::setw(width) << orb_pos_[i * 3 + 2]
                    << std::setw(width) << orb_eng_[i]
                    << std::setw(width) << 0 // orbital label is assumed to be 0
                    << "\n";
        }
        outfile << "end_orbitals\n";

        // hopping terms
        outfile << "begin_hopping\n";
        for (uint64_t i = 0; i < PyBaseModel<hop_t>::num_hop_; ++i) {
            size_t ind_offset = i * 5;
            outfile << std::setw(width_int) << hop_ind_[ind_offset]
                    << std::setw(width_int) << hop_ind_[ind_offset + 1]
                    << std::setw(width_int) << hop_ind_[ind_offset + 2]
                    << std::setw(width_int) << hop_ind_[ind_offset + 3]
                    << std::setw(width_int) << hop_ind_[ind_offset + 4]
                    << std::setw(width) << tbplas::base::real_part(hop_eng_[i])
                    << std::setw(width) << tbplas::base::imag_part(hop_eng_[i])
                    << "\n";
        }
        outfile << "end_hopping\n";
        outfile.close();
    }

    /**
     * @brief Extract hopping data in COO format for Lindhard AC conductivity calculation
     *
     * @param[out] hop_i index i in orbital pairs, i.e, orbital in r0
     * @param[out] hop_j index j in orbital pairs, i.e, orbital in rn
     * @param[out] hop_v hopping energy in eV
     * @param[out] dr_nm hopping distances in nm in COLUMN-WISE order
     *
     * @note Conjugate terms are not included. The caller is responsible for restoring
     * them.
     */
    void build_ham_dr_coo(
        std::vector<orbital_index_t>& hop_i,
        std::vector<orbital_index_t>& hop_j,
        std::vector<complex_t>& hop_v,
        Eigen::Matrix3Xd& dr_nm) const
    {
        // Allocate arrays
        orbital_index_t num_hop = PyBaseModel<hop_t>::num_hop_;
        hop_i = std::vector<orbital_index_t>(num_hop);
        hop_j = std::vector<orbital_index_t>(num_hop);
        hop_v = std::vector<complex_t>(num_hop);
        dr_nm = Eigen::Matrix3Xd(3, num_hop);

        // Loop over hopping terms
        auto orb_pos_map = get_orb_pos_map();
        for (size_t ptr = 0; ptr < num_hop; ++ptr) {
            // Extract cell and orbital indices
            size_t offset_hop = ptr * 5;
            cell_index_t ra = hop_ind_[offset_hop];
            cell_index_t rb = hop_ind_[offset_hop + 1];
            cell_index_t rc = hop_ind_[offset_hop + 2];
            cell_index_t idx_i = hop_ind_[offset_hop + 3];
            cell_index_t idx_j = hop_ind_[offset_hop + 4];

            // Evaluate dr
            Eigen::Vector3d dr_frac { ra * 1.0, rb * 1.0, rc * 1.0 }; // rn
            dr_frac += orb_pos_map.col(idx_j) - orb_pos_map.col(idx_i); // \tau_j - \tau_i
            Eigen::Vector3d dr_cart = (dr_frac.transpose() * PyBaseModel<hop_t>::lattice_).transpose();

            // Fill arrays
            hop_i[ptr] = idx_i;
            hop_j[ptr] = idx_j;
            hop_v[ptr] = hop_eng_[ptr];
            dr_nm.col(ptr) = dr_cart;
        }
    }

    /**
     * @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: IT MUST BE count[i]++, not ++count[i]!
        for (orbital_index_t i = 0; i < num_orb; ++i) {
            orbital_index_t ptr = indptr[i] + count[i]++;
            indices[ptr] = i;
            ham[ptr] = orb_eng_[i];
        }

        // Add hopping terms
        size_t num_hop_reduced = PyBaseModel<hop_t>::num_hop_;
        auto orb_pos_map = get_orb_pos_map();
        for (orbital_index_t ptr = 0; ptr < num_hop_reduced; ++ptr) {
            // Extract cell and orbital indices, etc.
            size_t offset_hop = ptr * 5;
            cell_index_t ra = hop_ind_[offset_hop];
            cell_index_t rb = hop_ind_[offset_hop + 1];
            cell_index_t rc = hop_ind_[offset_hop + 2];
            cell_index_t idx_i = hop_ind_[offset_hop + 3];
            cell_index_t idx_j = hop_ind_[offset_hop + 4];
            complex_t hij = hop_eng_[ptr];

            // Evaluate hopping distance
            Eigen::Vector3d dr_frac { ra * 1.0, rb * 1.0, rc * 1.0 }; // rn
            dr_frac += orb_pos_map.col(idx_j) - orb_pos_map.col(idx_i); // \tau_j - \tau_i
            Eigen::Vector3d dr = (dr_frac.transpose() * PyBaseModel<hop_t>::lattice_).transpose();

            // 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
    {
        // 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 (orbital_index_t i = 0; i < num_orb; ++i) {
            orbital_index_t ptr = indptr[i] + count[i]++;
            indices[ptr] = i;
            ham[ptr] = orb_eng_[i];
        }

        // Add hopping terms
        size_t num_hop_reduced = PyBaseModel<hop_t>::num_hop_;
        auto orb_pos_map = get_orb_pos_map();
        for (orbital_index_t ptr = 0; ptr < num_hop_reduced; ++ptr) {
            // Extract cell and orbital indices, etc.
            size_t offset_hop = ptr * 5;
            cell_index_t ra = hop_ind_[offset_hop];
            cell_index_t rb = hop_ind_[offset_hop + 1];
            cell_index_t rc = hop_ind_[offset_hop + 2];
            cell_index_t idx_i = hop_ind_[offset_hop + 3];
            cell_index_t idx_j = hop_ind_[offset_hop + 4];
            complex_t hij = hop_eng_[ptr];

            // Add phase factor for convention 1 and 2
            if ((convention == 1) || (convention == 2)) {
                Eigen::Vector3d dr_frac { ra * 1.0, rb * 1.0, rc * 1.0 }; // rn
                if (convention == 1) {
                    dr_frac += orb_pos_map.col(idx_j) - orb_pos_map.col(idx_i); // \tau_j - \tau_i
                }
                double phase = base::TWOPI * k_point.dot(dr_frac);
                hij *= 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] = tbplas::base::conj_t(hij);
        }
    }

    /**
     * @brief Core part of setting dense Hamiltonian
     *
     * @tparam kpt_t data type of k-point
     * @tparam ham_t datatype of dense Hamiltonian
     * @param[in] k_point fractional coordinate of k-point
     * @param[out] ham incoming Hamiltonian matrix
     * @param[in] convention convention of Hamiltonian
     */
    template <typename kpt_t, typename ham_t>
    void build_ham_dense_core(
        const kpt_t& k_point,
        ham_t& ham,
        int convention) const
    {
        // Reset Hamiltonian
        ham.setZero();

        // On-site energy
        orbital_index_t num_orb = PyBaseModel<hop_t>::num_orb_;
        for (orbital_index_t i = 0; i < num_orb; ++i) {
            ham(i, i) = orb_eng_[i];
        }

        // Hopping terms
        size_t num_hop = PyBaseModel<hop_t>::num_hop_;
        auto orb_pos_map = get_orb_pos_map();
        for (size_t ptr = 0; ptr < num_hop; ++ptr) {
            // Extract cell and orbital indices, etc.
            size_t offset_hop = ptr * 5;
            cell_index_t ra = hop_ind_[offset_hop];
            cell_index_t rb = hop_ind_[offset_hop + 1];
            cell_index_t rc = hop_ind_[offset_hop + 2];
            cell_index_t idx_i = hop_ind_[offset_hop + 3];
            cell_index_t idx_j = hop_ind_[offset_hop + 4];
            complex_t hij = hop_eng_[ptr];

            // Add phase factor for convention 1 and 2
            Eigen::Vector3d dr_frac { ra * 1.0, rb * 1.0, rc * 1.0 }; // rn
            if (convention == 1) {
                dr_frac += orb_pos_map.col(idx_j) - orb_pos_map.col(idx_i); // \tau_j - \tau_i
            }
            double phase = base::TWOPI * k_point.dot(dr_frac);
            hij *= base::exp_it(phase);

            // Fill Hamiltonian
            ham(idx_i, idx_j) += hij;
            ham(idx_j, idx_i) += std::conj(hij);
        }
    }

    /**
     * @brief Set up dense Hamiltonian matrix for given k-point.
     *
     * @param[in] k_point fractional coordinate of k-point
     * @param[out] ham incoming Hamiltonian matrix
     * @param[in] convention convention of Hamiltonian
     */
    void build_ham_dense(
        const Eigen::Vector3d& k_point,
        Eigen::MatrixXcd& ham,
        int convention = 1) const
    {
        assert((convention == 1) || (convention == 2));
        orbital_index_t num_orb = PyBaseModel<hop_t>::num_orb_;
        assert((ham.rows() == num_orb) && (ham.cols() == num_orb));
        build_ham_dense_core(k_point, ham, convention);
    }

    /**
     * @brief Set up dense Hamiltonian matrix for given k-point,
     * for fitting model parameters in Python.
     *
     * @param[in] k_point fractional coordinate of k-point
     * @param[out] ham incoming Hamiltonian matrix
     * @param[in] convention convention of Hamiltonian
     *
     * @note This function assumes ham to be pre-allocated in column-major
     * order. So the 'order' attribute of numpy array must be 'F' (Fortran).
     */
    void build_ham_dense(
        double* k_point,
        complex_t* ham,
        int convention = 1) const
    {
        assert((convention == 1) || (convention == 2));
        orbital_index_t num_orb = PyBaseModel<hop_t>::num_orb_;
        Eigen::Map<Eigen::Vector3d> k_point_map(k_point);
        Eigen::Map<Eigen::MatrixXcd> ham_map(ham, num_orb, num_orb);
        build_ham_dense_core(k_point_map, ham_map, convention);
    }

    /**
     * @brief Set up derivatives of dense Hamiltonian matrix for given k-point.
     * @param[in] k_point fractional coordinate of k-point
     * @param[out] dh_dka dH/dka, not dH/dkx!
     * @param[out] dh_dkb dH/dkb, not dH/dky!
     * @param[out] dh_dkc dH/dkc, not dH/dkz!
     * @param[in] convention convention of Hamiltonian
     */
    void build_ham_der_dense(
        const Eigen::Vector3d& k_point,
        Eigen::MatrixXcd& dh_dka,
        Eigen::MatrixXcd& dh_dkb,
        Eigen::MatrixXcd& dh_dkc,
        int convention = 1) const
    {
        assert((convention == 1) || (convention == 2));
        orbital_index_t num_orb = PyBaseModel<hop_t>::num_orb_;
        assert((dh_dka.rows() == num_orb) && (dh_dka.cols() == num_orb));
        assert((dh_dkb.rows() == num_orb) && (dh_dkb.cols() == num_orb));
        assert((dh_dkc.rows() == num_orb) && (dh_dkc.cols() == num_orb));
        dh_dka.setZero();
        dh_dkb.setZero();
        dh_dkc.setZero();
        complex_t i2pi = base::UIMAG * base::TWOPI;

        // No need for on-site terms since the hopping distance is zero.

        // Hopping terms
        size_t num_hop = PyBaseModel<hop_t>::num_hop_;
        auto orb_pos_map = get_orb_pos_map();
        for (size_t ptr = 0; ptr < num_hop; ++ptr) {
            // extract cell and orbital indices, etc.
            size_t offset_hop = ptr * 5;
            cell_index_t ra = hop_ind_[offset_hop];
            cell_index_t rb = hop_ind_[offset_hop + 1];
            cell_index_t rc = hop_ind_[offset_hop + 2];
            cell_index_t idx_i = hop_ind_[offset_hop + 3];
            cell_index_t idx_j = hop_ind_[offset_hop + 4];
            complex_t hij = hop_eng_[ptr];

            // add phase factor for convention 1 and 2
            Eigen::Vector3d dr_frac { ra * 1.0, rb * 1.0, rc * 1.0 }; // rn
            if (convention == 1) {
                dr_frac += orb_pos_map.col(idx_j) - orb_pos_map.col(idx_i); // \tau_j - \tau_i
            }
            double phase = base::TWOPI * k_point.dot(dr_frac);
            hij *= base::exp_it(phase);

            // dH/dk is also Hermitian due to the I2PI factor
            Eigen::Vector3cd i2pi_dr = i2pi * dr_frac;
            dh_dka(idx_i, idx_j) += i2pi_dr[0] * hij;
            dh_dkb(idx_i, idx_j) += i2pi_dr[1] * hij;
            dh_dkc(idx_i, idx_j) += i2pi_dr[2] * hij;
            dh_dka(idx_j, idx_i) -= i2pi_dr[0] * std::conj(hij);
            dh_dkb(idx_j, idx_i) -= i2pi_dr[1] * std::conj(hij);
            dh_dkc(idx_j, idx_i) -= i2pi_dr[2] * std::conj(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 The formula is eqn. 78 of CPC paper.
     */
    void build_density_coeff(
        const Eigen::Vector3d& q_point,
        Eigen::VectorXcd& coeff_q,
        Eigen::VectorXcd& coeff_mq) const
    {
        Eigen::Vector3d origin_frac = get_origin_map().transpose() * PyBaseModel<hop_t>::lattice_.inverse();
        orbital_index_t num_orb = PyBaseModel<hop_t>::num_orb_;
        coeff_q = Eigen::VectorXcd(num_orb);
        coeff_mq = Eigen::VectorXcd(num_orb);
        complex_t phase = 0.0;
        auto orb_pos_map = get_orb_pos_map();
#pragma omp parallel for private(phase)
        for (orbital_index_t i = 0; i < num_orb; ++i) {
            phase = base::exp_it(base::TWOPI * q_point.dot(orb_pos_map.col(i) + origin_frac));
            coeff_q[i] = phase;
            coeff_mq[i] = base::conj_t(phase);
        }
    }
};

/**
 * @brief Interface for Python Sample class.
 *
 * @tparam hop_t data type of hopping terms
 */
template <typename hop_t>
class PySample : public PyBaseModel<hop_t> {
private:
    /// @brief pointer to data of orbital energies in eV
    const double* orb_eng_ = nullptr;

    /// @brief pointer to data of orbital position in nm
    /// @note x, y, z components of one atom is continuous in RAM
    const double* orb_pos_ = nullptr;

    /// @brief pointer to row and column indices of hopping terms reduced by conjugate relation
    const orbital_index_t* hop_ind_ = nullptr;

    /// @brief pointer to hopping energies in eV
    const hop_t* hop_eng_ = nullptr;

    /// @brief pointer to hopping distances in nm
    const double* dr_ = nullptr;

    /**
     * @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 = PyBaseModel<hop_t>::num_orb_;
        orbital_index_t num_hops = PyBaseModel<hop_t>::num_hop_;
        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
        for (orbital_index_t i = 0; i < num_hops; ++i) {
            orbital_index_t offset = i * 2;
            orbital_index_t orb_i = hop_ind_[offset];
            orbital_index_t orb_j = hop_ind_[offset + 1];
            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];
        }
    }

    /**
     * @brief Get map to orbital positions
     *
     * @return map to orbital positions
     */
    [[nodiscard]] Eigen::Map<Eigen::Matrix3Xd> get_orb_pos_map() const
    {
        return Eigen::Map<Eigen::Matrix3Xd>(const_cast<double*>(orb_pos_), 3, PyBaseModel<hop_t>::num_orb_);
    }

    /**
     * @brief Get map to hopping distances
     *
     * @return map to hopping distances
     */
    [[nodiscard]] Eigen::Map<Eigen::Matrix3Xd> get_dr_map() const
    {
        return Eigen::Map<Eigen::Matrix3Xd>(const_cast<double*>(dr_), 3, PyBaseModel<hop_t>::num_hop_);
    }

public:
    /**
     * @brief Construct a new PySample object
     *
     * @param[in] num_orb number or orbitals
     * @param[in] num_hop number of hopping terms
     * @param[in] orb_eng pointer to data of orbital energies
     * @param[in] orb_pos pointer to data of orbital position
     * @param[in] hop_ind pointer to row and column indices of hopping terms
     * @param[in] hop_eng  pointer to hopping energies
     * @param[in] dr pointer to hopping distances
     * @param[in] lattice pointer to lattice vectors
     */
    explicit PySample(
        uint64_t num_orb,
        uint64_t num_hop,
        const double* orb_eng,
        const double* orb_pos,
        const int64_t* hop_ind,
        const hop_t* hop_eng,
        const double* dr,
        const double* lattice)
        : PyBaseModel<hop_t>(num_orb, num_hop, lattice)
        , orb_eng_(orb_eng)
        , orb_pos_(orb_pos)
        , hop_ind_((orbital_index_t*)hop_ind)
        , hop_eng_(hop_eng)
        , dr_(dr)
    {
    }

    /**
     * @brief Save model to file
     *
     * @param[in] file_name file name of model
     */
    void save(const std::string& file_name) const
    {
        constexpr int prec = 16;
        constexpr int width = 22;
        constexpr 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";
        for (int i = 0; i < 3; ++i) {
            for (int j = 0; j < 3; ++j) {
                outfile << std::setw(width) << PyBaseModel<hop_t>::lattice_(i, j);
            }
            outfile << "\n";
        }
        outfile << "end_lattice\n";

        // PySample class uses Cartesian coordinates and origin is not essential.
        // So we make it (0, 0, 0) in saved model.
        outfile << "begin_origin\n";
        Eigen::Vector3d origin(0.0, 0.0, 0.0);
        for (int i = 0; i < 3; ++i) {
            outfile << std::setw(width) << origin(i);
        }
        outfile << "\nend_origin\n";

        // orbitals
        outfile << "begin_orbitals\n";
        for (uint64_t i = 0; i < PyBaseModel<hop_t>::num_orb_; ++i) {
            outfile << std::setw(width) << orb_pos_[i * 3]
                    << std::setw(width) << orb_pos_[i * 3 + 1]
                    << std::setw(width) << orb_pos_[i * 3 + 2]
                    << std::setw(width) << orb_eng_[i]
                    << std::setw(width) << 0 // orbital label is assumed to be 0
                    << "\n";
        }
        outfile << "end_orbitals\n";

        // hopping terms
        outfile << "begin_hopping\n";
        for (uint64_t i = 0; i < PyBaseModel<hop_t>::num_hop_; ++i) {
            outfile << std::setw(width_int) << hop_ind_[i * 2]
                    << std::setw(width_int) << hop_ind_[i * 2 + 1]
                    << std::setw(width) << tbplas::base::real_part(hop_eng_[i])
                    << std::setw(width) << tbplas::base::imag_part(hop_eng_[i])
                    << "\n";
        }
        outfile << "end_hopping\n";

        outfile << "begin_dr\n";
        for (uint64_t i = 0; i < PyBaseModel<hop_t>::num_hop_; ++i) {
            outfile << std::setw(width) << dr_[i * 3]
                    << std::setw(width) << dr_[i * 3 + 1]
                    << std::setw(width) << dr_[i * 3 + 2]
                    << "\n";
        }
        outfile << "end_dr\n";
        outfile.close();
    }

    /**
     * @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 (orbital_index_t i = 0; i < num_orb; ++i) {
            orbital_index_t ptr = indptr[i] + count[i]++;
            indices[ptr] = i;
            ham[ptr] = orb_eng_[i];
            curr_coeff_x[ptr] = 0.0;
            curr_coeff_y[ptr] = 0.0;
        }

        size_t num_hop_reduced = PyBaseModel<hop_t>::num_hop_;
        auto dr_map = get_dr_map();
        for (orbital_index_t i = 0; i < num_hop_reduced; ++i) {
            orbital_index_t offset = i * 2;
            orbital_index_t idx_i = hop_ind_[offset];
            orbital_index_t idx_j = hop_ind_[offset + 1];
            complex_t hij = hop_eng_[i];
            Eigen::Vector3d dr = dr_map.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 (orbital_index_t i = 0; i < num_orb; ++i) {
            orbital_index_t ptr = indptr[i] + count[i]++;
            indices[ptr] = i;
            ham[ptr] = orb_eng_[i];
        }

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

        // Add hopping terms
        size_t num_hop_reduced = PyBaseModel<hop_t>::num_hop_;
        auto dr_map = get_dr_map();
        for (orbital_index_t i = 0; i < num_hop_reduced; ++i) {
            orbital_index_t offset = i * 2;
            orbital_index_t idx_i = hop_ind_[offset];
            orbital_index_t idx_j = hop_ind_[offset + 1];
            complex_t hij = hop_eng_[i];

            // Add phase factor for convention 1
            if (convention == 1) {
                double phase = k_cart.dot(dr_map.col(i));
                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] = tbplas::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 = PyBaseModel<hop_t>::num_orb_;
        coeff_q = Eigen::VectorXcd(num_orb);
        coeff_mq = Eigen::VectorXcd(num_orb);
        complex_t phase = 0.0;
        Eigen::Matrix3d recip_vec = PyBaseModel<hop_t>::get_reciprocal_lattice();
        Eigen::Vector3d q_cart = (q_point.transpose() * recip_vec).transpose();
        auto orb_pos_map = get_orb_pos_map();
#pragma omp parallel for private(phase)
        for (orbital_index_t i = 0; i < num_orb; ++i) {
            phase = base::exp_it(q_cart.dot(orb_pos_map.col(i)));
            coeff_q[i] = phase;
            coeff_mq[i] = base::conj_t(phase);
        }
    }
};
} // namespace tbplas::python
#endif