/**
 * @brief Classes for calculating Z2 topological invariant.
 *
 * @author Yunhai Li (liyunhai1016@hotmail.com)
 *
 * @copyright Copyright (c) 2024, TBPLaS develop team.
 */

#ifndef TBPLAS_DIAG_Z2_H
#define TBPLAS_DIAG_Z2_H

#include "base.h"
#include "base/consts.h"
#include "results.h"

namespace tbplas::diag {

/**
 * @brief Class for evaluating the Z2 topological invariant.
 * @note Reference: 10.1103/PhysRevB.84.075119
 *
 * @tparam model_t datatype of model assigned to this solver
 */
template <typename model_t>
class Z2 : public DiagSolver<model_t> {
private:
    /// @brief D matrix for single k-point in the reference
    Eigen::MatrixXcd dk_;

    /// @brief F matrix for (ka_i, ka_i+1)
    Eigen::MatrixXcd fk_;

    /// @brief solver for getting the eigenvalues of d_mat
    Eigen::ComplexEigenSolver<Eigen::MatrixXcd> zgeev_;

    /**
     * @brief Evaluate the F matrix from eigenvectors of hk_ at ka_i and ka_i+1.
     *
     * @param[in] vec0 eigenvectors at ka_i
     * @param[in] vec1 eigenvectors at ka_i+1
     * @note No results are returned. fk_ is updated.
     */
    void eval_fk(
        const Eigen::MatrixXcd& vec0,
        const Eigen::MatrixXcd& vec1)
    {
        size_t num_occ = fk_.cols();
        for (ptrdiff_t i = 0; i < num_occ; ++i) {
            for (ptrdiff_t j = 0; j < num_occ; ++j) {
                fk_(i, j) = vec0.col(i).dot(vec1.col(j));
            }
        }
    }

    /**
     * @brief Evaluate the F matrix from eigenvectors of hk_ at ka_i and ka_i+1.
     *
     * @param[in] vec0 eigenvectors at ka_i
     * @param[in] vec1 eigenvectors at ka_i+1
     * @param[in] sk1 overlap matrix at ka_i+1
     * @note
     * 1. No results are returned. fk_ is updated.
     * 2. vec0 and vec1 are at different k, so are sk0 and sk1. We assume sk0 ~= sk1,
     *    which may be inaccurate.
     */
    void eval_fk(
        const Eigen::MatrixXcd& vec0,
        const Eigen::MatrixXcd& vec1,
        const Eigen::MatrixXcd& sk1)
    {
        size_t num_occ = fk_.cols();
        for (ptrdiff_t i = 0; i < num_occ; ++i) {
            for (ptrdiff_t j = 0; j < num_occ; ++j) {
                fk_(i, j) = vec0.col(i).dot(sk1 * vec1.col(j));
            }
        }
    }

    /**
     * @brief Diagonalize dk_ to get the eigenvalues
     *
     * @param[out] dk_eigenvalues eigenvalues of dk_, dimension (num_occ)
     */
    void diag_dk(Eigen::VectorXcd& dk_eigenvalues)
    {
        const std::regex pattern_lapacke(R"(^lapack(e)?$)", std::regex_constants::icase);
        if (std::regex_match(DiagSolver<model_t>::config.diag_algo, pattern_lapacke)) {
#ifdef WITH_LAPACKE
            lapacke_zgeev_wrapper(dk_.cols(),
                reinterpret_cast<lapack_complex_double*>(dk_.data()),
                reinterpret_cast<lapack_complex_double*>(dk_eigenvalues.data()));
#else
            DiagSolver<model_t>::print_stderr("ERROR: LAPACKE support disabled during compilation!");
            std::exit(-1);
#endif
        } else {
            zgeev_.compute(dk_);
            dk_eigenvalues = zgeev_.eigenvalues();
        }
    }

    /**
     * @brief Evaluate the phases of all the elements in an array.
     *
     * @param[in] z_array complex numbers for which the phases will be evaluated
     * @param[out] phase phase matrix, dimension (num_occ, num_kb)
     * @param[in] ikb index of kb
     */
    Eigen::VectorXd eval_phase(
        const Eigen::MatrixXcd& z_array,
        Eigen::MatrixXd& phase,
        const int& ikb)
    {
        phase.col(ikb) = z_array.cwiseArg();
        Eigen::VectorXd norm = z_array.cwiseAbs();
        for (ptrdiff_t i = 0; i < z_array.size(); ++i) {
            if (norm(i) < tbplas::base::INFINITESIMAL) {
                phase(i, ikb) = 0.0;
            }
        }
        return phase;
    }

public:
    /**
     * @brief Construct a new Z2 object
     *
     * @param[in] model model for which calculations will be performed
     */
    explicit Z2(const model_t& model)
        : DiagSolver<model_t>(model)
    {
    }

    /**
     * @brief Construct a new Z2 object with overlap
     *
     * @param[in] model model for which calculations will be performed
     * @param[in] overlap model which holds the overlap

     */
    explicit Z2(const model_t& model, const model_t& overlap)
        : DiagSolver<model_t>(model, overlap)
    {
    }

    /**
     * @brief Construct a new Z2 object
     *
     * @param[in] model pointer to model
     * @param[in] overlap pointer to overlap
     *
     * @note This constructor is for Python interface.
     */
    explicit Z2(const model_t* model, const model_t* overlap = nullptr)
        : DiagSolver<model_t>(model, overlap)
    {
    }

    /**
     * @brief Calculate the evolution of phases of WF centers as function of kb.
     *
     * @return evolution of phases
     */
    [[nodiscard]] Z2Data calc_phases()
    {
        // Extract parameters from config
        int num_occ = DiagSolver<model_t>::config.num_occ;
        int convention = DiagSolver<model_t>::config.convention;
        Eigen::VectorXd& ka_array = DiagSolver<model_t>::config.ka_array;
        Eigen::VectorXd& kb_array = DiagSolver<model_t>::config.kb_array;
        double kc = DiagSolver<model_t>::config.kc;
        assert(num_occ > 0);
        size_t num_ka = ka_array.size();
        size_t num_kb = kb_array.size();

        // Allocate working arrays
        auto backend = DiagSolver<model_t>::make_backend();
        size_t num_orbitals = backend->get_num_orbitals();
        size_t num_states = backend->get_num_states();
        // Must have enough occupied states. Otherwise eval_fk may fail.
        assert(num_states >= num_occ);

        // Hamiltonian related variables
        Eigen::VectorXd hk_eigenvalues(num_states);
        Eigen::MatrixXcd vec0(num_orbitals, num_states);
        Eigen::MatrixXcd vec0_bak(num_orbitals, num_states);
        Eigen::MatrixXcd sk0_bak(num_orbitals, num_orbitals);
        Eigen::MatrixXcd vec1(num_orbitals, num_states);

        // Phase related variables
        dk_ = Eigen::MatrixXcd::Identity(num_occ, num_occ);
        fk_ = Eigen::MatrixXcd::Identity(num_occ, num_occ);
        Eigen::VectorXcd dk_eigenvalues(num_occ);
        Eigen::MatrixXd phases = Eigen::MatrixXd::Zero(num_occ, num_kb);

        // Evaluate phases
        RangeInfo range = DiagSolver<model_t>::dist_range(num_kb);
        for (int ikb = range.k_min; ikb <= range.k_max; ++ikb) {
            // As the D matrix is evaluated through iterative matrix
            // multiplication, it should be initialized as an eye matrix.
            dk_ = Eigen::MatrixXcd::Identity(num_occ, num_occ);

            // Get and backup the eigenvectors of Hamiltonian at ka_0
            double kb = kb_array(ikb);
            Eigen::Vector3d kpt0(ka_array(0), kb, kc);
            backend->calc_eigen_all(kpt0, convention, hk_eigenvalues, vec0);
            vec0_bak = vec0;
            if (DiagSolver<model_t>::overlap_ != nullptr) {
                sk0_bak = backend->get_sk();
            }

            // Loop over a-axis to build the D matrix
            Eigen::MatrixXcd* p0 = &vec0;
            Eigen::MatrixXcd* p1 = &vec1;
            for (ptrdiff_t ika = 1; ika < num_ka; ++ika) {
                Eigen::Vector3d kpt1(ka_array(ika), kb, kc);
                backend->calc_eigen_all(kpt1, convention, hk_eigenvalues, *p1);
                if (DiagSolver<model_t>::overlap_ == nullptr) {
                    eval_fk(*p0, *p1);
                } else {
                    eval_fk(*p0, *p1, backend->get_sk());
                }
                dk_ *= fk_;
                std::swap(p0, p1);
            }

            // Special treatment for the last ka
            if (DiagSolver<model_t>::overlap_ == nullptr) {
                eval_fk(*p0, vec0_bak);
            } else {
                eval_fk(*p0, vec0_bak, sk0_bak);
            }
            dk_ *= fk_;

            // Diagonalize D matrix and evaluate phases
            diag_dk(dk_eigenvalues);
            eval_phase(dk_eigenvalues, phases, ikb);
        }
        DiagSolver<model_t>::mpi_env_.InplaceAllreduceMatrixXd(phases);

        // Assemble results and return
        Z2Data data(DiagSolver<model_t>::config.prefix);
        data.k_a = ka_array; // ka_array and kb_array is stored in config, so do not move it.
        data.k_b = kb_array;
        data.phases = std::move(phases);
        if (DiagSolver<model_t>::is_master()) {
            data.save("z2");
        }
        // Master process needs to save data and finishes later than other processes.
        // So we need to put a barrier here.
        DiagSolver<model_t>::mpi_env_.barrier();
        return data;
    }
};

} // namespace tbplas::diag

#endif // TBPLAS_DIAG_Z2_H
