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

#ifndef TBPLAS_BUILDER_SUPER_H
#define TBPLAS_BUILDER_SUPER_H

#include <memory>

#include <eigen3/Eigen/Dense>

#include "advanced.h"
#include "base.h"
#include "base/lattice.h"
#include "filters.h"
#include "kdtree.h"
#include "meshgrid.h"
#include "primitive.h"

namespace tbplas::builder {
using base::cart2frac;
using base::get_lattice_volume;

/**
 * @brief Model built by reshaping a primitive cell.
 *
 * @tparam hop_t datatype of hopping terms of primitive cell
 */
template <typename hop_t>
class SuperCell {
public:
    /**
     * @brief Construct a new SuperCell from lattice vectors
     *
     * @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
     *
     * @note
     * 1. The primitive cell is copied since we use shared_ptr to manage it.
     * 2. The supercell is assumed to have the same origin as the primitive cell.
     */
    explicit SuperCell(
        const PrimitiveCell<hop_t>& prim_cell,
        const Eigen::Matrix3d& sc_lattice,
        const pbc_t& pbc,
        const double& delta = 1.0e-5)
        : prim_cell_(std::make_shared<PrimitiveCell<hop_t>>(prim_cell))
        , sc_lattice_(sc_lattice)
        , sc_origin_(prim_cell.get_origin())
        , pbc_(pbc)
        , delta_(delta)
    {
    }

    /**
     * @brief Construct a new SuperCell from dimension
     *
     * @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
     *
     * @note
     * 1. The primitive cell is copied since we use shared_ptr to manage it.
     * 2. The supercell is assumed to have the same origin as the primitive cell.
     */
    explicit SuperCell(
        const PrimitiveCell<hop_t>& prim_cell,
        const dim_t& dim,
        const pbc_t& pbc,
        const double& delta = 1.0e-5)
        : prim_cell_(std::make_shared<PrimitiveCell<hop_t>>(prim_cell))
        , sc_lattice_(prim_cell.get_lattice())
        , sc_origin_(prim_cell.get_origin())
        , pbc_(pbc)
        , delta_(delta)
        , dim_(dim)
        , extend_mode_(true)
    {
        sc_lattice_.row(0) *= std::get<0>(dim);
        sc_lattice_.row(1) *= std::get<1>(dim);
        sc_lattice_.row(2) *= std::get<2>(dim);
    }

    SuperCell(const SuperCell& other)
        : prim_cell_(other.prim_cell_)
        , sc_lattice_(other.sc_lattice_)
        , sc_origin_(other.sc_origin_)
        , pbc_(other.pbc_)
        , delta_(other.delta_)
        , filters_(other.filters_)
        , modifiers_(other.modifiers_)
        , dim_(other.dim_)
        , extend_mode_(other.extend_mode_)
    {
    }

    SuperCell(SuperCell&& other) noexcept
        : prim_cell_(std::move(other.prim_cell_))
        , sc_lattice_(std::move(other.sc_lattice_))
        , sc_origin_(std::move(other.sc_origin_))
        , pbc_(std::move(other.pbc_))
        , delta_(std::move(other.delta_))
        , filters_(std::move(other.filters_))
        , modifiers_(std::move(other.modifiers_))
        , dim_(std::move(other.dim_))
        , extend_mode_(std::move(other.extend_mode_))
    {
    }

    SuperCell& operator=(const SuperCell& other)
    {
        if (this != &other) {
            prim_cell_ = other.prim_cell_;
            sc_lattice_ = other.sc_lattice_;
            sc_origin_ = other.sc_origin_;
            pbc_ = other.pbc_;
            delta_ = other.delta_;
            filters_ = other.filters_;
            modifiers_ = other.modifiers_;
            dim_ = other.dim_;
            extend_mode_ = other.extend_mode_;
        }
        return *this;
    }

    SuperCell& operator=(SuperCell&& other) noexcept
    {
        if (this != &other) {
            prim_cell_ = std::move(other.prim_cell_);
            sc_lattice_ = std::move(other.sc_lattice_);
            sc_origin_ = std::move(other.sc_origin_);
            pbc_ = std::move(other.pbc_);
            delta_ = std::move(other.delta_);
            filters_ = std::move(other.filters_);
            modifiers_ = std::move(other.modifiers_);
            dim_ = std::move(other.dim_);
            extend_mode_ = std::move(other.extend_mode_);
        }
        return *this;
    }

    ~SuperCell() = default;

    /**
     * @brief Add a new orbital filter.
     *
     * @tparam filter_t datatype of the filter
     * @param[in] filter filter to add
     *
     * @note
     * 1. filter_t should be derived from AbstractOrbitalFilter.
     * 2. The filter is copied since we use shared_ptr to manage it.
     */
    template <typename filter_t>
    void add_filter(const filter_t& filter)
    {
        filters_.push_back(std::make_shared<filter_t>(filter));
    }

    /**
     * @brief Add a new orbital modifier.
     *
     * @tparam modifier_t datatype of the modifier
     * @param[in] modifier modifier to add
     *
     * @note
     * 1. modifier_t should be derived from AbstractOrbitalModifier.
     * 2. The modifier is copied since we use shared_ptr to manage it.
     */
    template <typename modifier_t>
    void add_modifier(const modifier_t& modifier)
    {
        modifiers_.push_back(std::make_shared<modifier_t>(modifier));
    }

    /**
     * @brief Build orbital data of this supercell
     *
     * @param[inout] orb_data OrbitalData instance to hold the results
     *
     * @note This functions reuses code from reshape_prim_cell.
     */
    void build_orbital_data(OrbitalData& orb_data) const
    {
        // Determine conversion matrix
        // Both lattice and prim_cell.lattice are in ROW-WISE order. So we need to transpose
        // the lattice before calling cart2frac, and transpose the result back afterwards.
        Eigen::Matrix3d lat_frac = cart2frac(prim_cell_->get_lattice(), sc_lattice_.transpose()).transpose();
        Eigen::Matrix3d conv_mat = lat_frac.inverse();

        // Determine searching range
        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();

        // Get the full list of orbitals and apply filters
        std::vector<Orbital> orbitals = extend_mode_ ? get_full_orbitals_extend() : get_full_orbitals_reshape(conv_mat, rn_range);
        for (const auto& f : filters_) {
            f->act(orbitals);
        }

        // Fill the arrays in orb_data
        orbital_index_t num_orb = orbitals.size();
        orb_data.orb_pos = Eigen::Matrix3Xd(3, num_orb);
        orb_data.orb_eng = Eigen::VectorXd(num_orb);
        orb_data.orb_label = std::vector<int64_t>(num_orb);
        orb_data.sc_lattice = sc_lattice_;
        orb_data.sc_origin = sc_origin_;
        orb_data.pbc = pbc_;
        orb_data.offset = 0;
        for (orbital_index_t i = 0; i < num_orb; ++i) {
            const Orbital& orb = orbitals[i];
            orb_data.orb_pos.col(i) = (orb.get_position().transpose() * sc_lattice_).transpose() + sc_origin_;
            orb_data.orb_eng[i] = orb.energy;
            orb_data.orb_label[i] = orb.label;
        }

        // Apply modifiers
        for (const auto& m : modifiers_) {
            m->act(orb_data);
        }
    }

    /**
     * @brief Get the lattice attribute
     *
     * @return lattice vectors of supercell
     */
    [[nodiscard]] const Eigen::Matrix3d& get_sc_lattice() const
    {
        return sc_lattice_;
    }

    /**
     * @brief Get the origin attribute
     *
     * @return origin of supercell
     */
    [[nodiscard]] const Eigen::Vector3d& get_sc_origin() const
    {
        return sc_origin_;
    }

private:
    /// @brief Pointer to primitive cell
    std::shared_ptr<PrimitiveCell<hop_t>> prim_cell_ = nullptr;

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

    /// @brief Cartesian coordinate of supercell origin in NM
    Eigen::Vector3d sc_origin_ = Eigen::Vector3d(0.0, 0.0, 0.0);

    /// @brief Boundary condition of supercell, true for pbc, false for non-pbc
    pbc_t pbc_ = { false, false, false };

    /// @brief Small displacement added to vector to avoid clipping
    double delta_ = 1.0e-5;

    /// @brief Orbital filters assigned to the supercell
    std::vector<std::shared_ptr<AbstractOrbitalFilter>> filters_;

    /// @brief Orbital modifiers assigned to the supercell
    std::vector<std::shared_ptr<AbstractOrbitalModifier>> modifiers_;

    /// @brief Dimension of the supercell w.r.t primitive cell, active in extend_mode only
    dim_t dim_ = { 1, 1, 1 };

    /// @brief Whether the supercell can be constructed by extending the primitive cell
    /// @note If false, then the supercell is constructed by reshaping the primitive cell.
    bool extend_mode_ = false;

    /**
     * @brief Get the full list of orbitals in the supercell by extending the primitive cell.
     *
     * @return full list of orbitals in the supercell
     */
    [[nodiscard]] std::vector<Orbital> get_full_orbitals_extend() const
    {
        std::vector<Orbital> full_orbitals;
        orbital_index_t num_orb_prim = prim_cell_->get_num_orbitals();

        // 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_;
        assert(dim_a > 0);
        assert(dim_b > 0);
        assert(dim_c > 0);

        // Estimate number of orbitals
        orbital_index_t num_orb_super = num_orb_prim * dim_a * dim_b * dim_c;
        full_orbitals.reserve(num_orb_super);

        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_prim; ++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_a;
                        double fb = new_pos(1) / dim_b;
                        double fc = new_pos(2) / dim_c;
                        full_orbitals.emplace_back(fa, fb, fc, orbital.energy, orbital.label);
                    }
                }
            }
        }
        full_orbitals.shrink_to_fit();
        return full_orbitals;
    }

    /**
     * @brief Get the full list of orbitals in the supercell by reshaping the primitive cell
     *
     * @param[in] conv_mat conversion matrix from build_orbital_data
     * @param[in] rn_range searching range from build_orbital_data
     * @return full list of orbitals in the supercell
     */
    [[nodiscard]] std::vector<Orbital> get_full_orbitals_reshape(
        const Eigen::Matrix3d& conv_mat,
        const Eigen::Matrix3Xi& rn_range) const
    {
        std::vector<Orbital> full_orbitals;
        orbital_index_t num_orb_prim = prim_cell_->get_num_orbitals();

        // Estimate number of orbitals based on lattice volume
        double vol_ratio = get_lattice_volume(sc_lattice_) / prim_cell_->get_lattice_volume();
        auto num_orb_super = static_cast<orbital_index_t>(num_orb_prim * vol_ratio);
        full_orbitals.reserve(num_orb_super);

        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;
        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) {
                        const Orbital& orbital = prim_cell_->get_orbital(i_o);
                        Eigen::Vector3d res_pos = ((rn + orbital.get_position()).transpose() * conv_mat).transpose();
                        floor_vector3(res_pos, ra, rb, rc, delta_);
                        if ((ra == 0) && (rb == 0) && (rc == 0)) {
                            full_orbitals.emplace_back(res_pos[0], res_pos[1], res_pos[2], orbital.energy, orbital.label);
                        }
                    }
                }
            }
        }
        full_orbitals.shrink_to_fit();
        return full_orbitals;
    }
};

/**
 * @brief Container of intra and inter-supercell hopping terms.
 *
 * @tparam hop_t template parameter for SuperCel
 */
template <typename hop_t>
class SCHopping {
public:
    /**
     * @brief Construct a new SCHopping object
     *
     * @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
     */
    explicit SCHopping(
        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_bra_(sc_bra)
        , sc_ket_(sc_ket)
        , max_distance_(max_distance)
        , default_hop_(default_hop)
    {
        // In intra-mode, add an intra_filter to filters_.
        if (sc_bra_.lock() == sc_ket_.lock()) {
            filters_.push_back(std::make_shared<IntraHoppingFilter>());
        }
    }

    SCHopping(const SCHopping& other)
        : sc_bra_(other.sc_bra_)
        , sc_ket_(other.sc_ket_)
        , filters_(other.filters_)
        , modifiers_(other.modifiers_)
        , max_distance_(other.max_distance_)
        , default_hop_(other.default_hop_)
    {
    }

    SCHopping(SCHopping&& other) noexcept
        : sc_bra_(std::move(other.sc_bra_))
        , sc_ket_(std::move(other.sc_ket_))
        , filters_(std::move(other.filters_))
        , modifiers_(std::move(other.modifiers_))
        , max_distance_(std::move(other.max_distance_))
        , default_hop_(std::move(other.default_hop_))
    {
    }

    SCHopping& operator=(const SCHopping& other)
    {
        if (this != &other) {
            sc_bra_ = other.sc_bra_;
            sc_ket_ = other.sc_ket_;
            filters_ = other.filters_;
            modifiers_ = other.modifiers_;
            max_distance_ = other.max_distance_;
            default_hop_ = other.default_hop_;
        }
        return *this;
    }

    SCHopping& operator=(SCHopping&& other) noexcept
    {
        if (this != &other) {
            sc_bra_ = std::move(other.sc_bra_);
            sc_ket_ = std::move(other.sc_ket_);
            filters_ = std::move(other.filters_);
            modifiers_ = std::move(other.modifiers_);
            max_distance_ = std::move(other.max_distance_);
            default_hop_ = std::move(other.default_hop_);
        }
        return *this;
    }

    ~SCHopping() = default;

    /**
     * @brief Add a new hopping filter.
     *
     * @tparam filter_t datatype of filter
     * @param[in] filter filter to add
     *
     * @note
     * 1. filter_t should be derived from AbstractHoppingFilter.
     * 2. The filter is copied since we use shared_ptr to manage it.
     */
    template <typename filter_t>
    void add_filter(const filter_t& filter)
    {
        filters_.push_back(std::make_shared<filter_t>(filter));
    }

    /**
     * @brief Add a new hopping modifier.
     *
     * @tparam modifier_t datatype of modifier
     * @param[in] modifier modifier to add
     *
     * @note
     * 1. modifier_t should be derived from AbstractHoppingModifier.
     * 2. The modifier is copied since we use shared_ptr to manage it.
     */
    template <typename modifier_t>
    void add_modifier(modifier_t& modifier)
    {
        modifiers_.push_back(std::make_shared<modifier_t>(modifier));
    }

    /**
     * @brief Get the hopping data attribute
     *
     * @param[in] data_bra orbital data of sc_bra
     * @param[in] data_ket orbital data of sc_ket
     * @param[inout] hop_data HoppingData instance holding the results
     */
    void build_hopping_data(
        const OrbitalData& data_bra,
        const OrbitalData& data_ket,
        HoppingData& hop_data) const
    {
        // Get full hopping terms and apply filters
        std::list<Term> neighbors = get_neighbors_kdtree(data_bra, data_ket);
        for (const auto& f : filters_) {
            f->act(neighbors);
        }

        // Fill arrays
        orbital_index_t num_hop = neighbors.size();
        hop_data.hop_ind = std::vector<std::tuple<orbital_index_t, orbital_index_t>>(num_hop);
        hop_data.hop_eng = Eigen::VectorXcd(num_hop);
        hop_data.dr = Eigen::Matrix3Xd(3, num_hop);
        hop_data.offset = 0;
        orbital_index_t orb_i = 0;
        orbital_index_t orb_j = 0;
        orbital_index_t counter = 0;
        for (const Term& n : neighbors) {
            std::tie(orb_i, orb_j) = n.pair;
            hop_data.hop_ind[counter] = { orb_i + data_bra.offset, orb_j + data_ket.offset };
            // set to default value, updated by hopping modifiers then.
            hop_data.hop_eng[counter] = default_hop_;
            hop_data.dr.col(counter) = n.rij;
            ++counter;
        }

        // Apply modifiers
        for (const auto& m : modifiers_) {
            m->act(data_bra, data_ket, hop_data);
        }
    }

    /**
     * @brief Get the pointer to 'bra' supercell
     *
     * @return the pointer to 'bra' supercell
     */
    [[nodiscard]] std::weak_ptr<SuperCell<hop_t>> get_sc_bra() const
    {
        return sc_bra_;
    }

    /**
     * @brief Get the pointer to 'ket' supercell
     *
     * @return pointer to 'ket' supercell
     */
    [[nodiscard]] std::weak_ptr<SuperCell<hop_t>> get_sc_ket() const
    {
        return sc_ket_;
    }

private:
    /// @brief Pointer to 'bra' supercell
    std::weak_ptr<SuperCell<hop_t>> sc_bra_ = nullptr;

    /// @brief Pointer to 'ket' supercell
    std::weak_ptr<SuperCell<hop_t>> sc_ket_ = nullptr;

    /// @brief Hopping filters
    std::vector<std::shared_ptr<AbstractHoppingFilter>> filters_;

    /// @brief Hopping modifiers
    std::vector<std::shared_ptr<AbstractHoppingModifier>> modifiers_;

    /// @brief Cutoff distance for finding hopping terms
    double max_distance_ = 1.0;

    /// @brief Default hopping parameter
    hop_t default_hop_ = 1.0;

    /**
     * @brief Find neighbors using KDTree.
     *
     * @param[in] data_bra orbital data of sc_bra
     * @param[in] data_ket orbital data of sc_ket
     * @return list of neighbors
     */
    [[nodiscard]] std::list<Term> get_neighbors_kdtree(
        const OrbitalData& data_bra,
        const OrbitalData& data_ket) const
    {
        // Assemble kdtree for cell_bra
        const Eigen::Matrix3Xd& pos_bra = data_bra.orb_pos;
        orbital_index_t num_orb_bra = pos_bra.cols();
        int leaf_max_size = 10;
        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> point_cloud(pos_bra_tree);
        kd_tree_t tree_bra(3, point_cloud, nanoflann::KDTreeSingleIndexAdaptorParams(leaf_max_size));
        tree_bra.buildIndex();

        // Call kdtree to find hopping terms
        int a_max = std::get<0>(data_ket.pbc) ? 1 : 0;
        int b_max = std::get<1>(data_ket.pbc) ? 1 : 0;
        int c_max = std::get<2>(data_ket.pbc) ? 1 : 0;
        const Eigen::Matrix3Xd& pos_ket = data_ket.orb_pos;
        return find_neighbors_within_radius(
            tree_bra,
            pos_bra,
            pos_ket,
            data_ket.sc_lattice,
            a_max,
            b_max,
            c_max,
            max_distance_);
    }

    /**
     * @brief Find neighbors using MeshGrid.
     *
     * @param[in] data_bra orbital data of sc_bra
     * @param[in] data_ket orbital data of sc_ket
     * @return list of neighbors
     *
     * @note This function is slower than get_neighbors_kdtree and is for testing only.
     */
    [[nodiscard]] std::list<Term> get_neighbors_meshgrid(
        const OrbitalData& data_bra,
        const OrbitalData& data_ket) const
    {
        std::list<Term> full_terms;
        int a_max = std::get<0>(data_ket.pbc) ? 1 : 0;
        int b_max = std::get<1>(data_ket.pbc) ? 1 : 0;
        int c_max = std::get<2>(data_ket.pbc) ? 1 : 0;
        const Eigen::Matrix3Xd& pos_bra = data_bra.orb_pos;
        for (int ia = -a_max; ia <= a_max; ++ia) {
            for (int ib = -b_max; ib <= b_max; ++ib) {
                for (int ic = -c_max; ic <= c_max; ++ic) {
                    Eigen::Vector3d rn { ia * 1.0, ib * 1.0, ic * 1.0 };
                    Eigen::Matrix3Xd pos_ket = data_ket.orb_pos.colwise() + (rn.transpose() * data_ket.sc_lattice).transpose();
                    MeshGrid mesh_grid(pos_bra, pos_ket, max_distance_);
                    std::list<ReducedTerm> terms = mesh_grid.get_neighbors(max_distance_);
                    for (const auto& t : terms) {
                        full_terms.emplace_back(std::make_tuple(ia, ib, ic), t.pair, t.rij, t.distance);
                    }
                }
            }
        }
        full_terms.sort();
        return full_terms;
    }
};
} // namespace tbplas::builder
#endif