/**
 * @brief Classes for calculating Berry curvature, Chern number, etc.
 *
 * @author Yunhai Li (liyunhai1016@hotmail.com)
 *
 * @copyright Copyright (c) 2024, TBPLaS develop team.
 */

#ifndef TBPLAS_DIAG_BERRY_H
#define TBPLAS_DIAG_BERRY_H

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

namespace tbplas::diag {

/**
 * @brief Calculator of Hamiltonian derivative w.r.t Cartesian coordinates of
 * k-points, using either analytical or numerical derivation.
 *
 * @tparam model_t datatype of model
 */
template <typename model_t>
class HamiltonianDerivative {
public:
    explicit HamiltonianDerivative(
        const model_t* model,
        const DiagConfig& config)
        : model_(model)
        , config_(&config)
        , recip_vectors_(model_->get_reciprocal_lattice())
    {
        size_t num_orbitals = model_->get_num_orbitals();
        if (config_->ham_deriv_analytical) {
            dh_dka_ = Eigen::MatrixXcd::Zero(num_orbitals, num_orbitals);
            dh_dkb_ = Eigen::MatrixXcd::Zero(num_orbitals, num_orbitals);
            dh_dkc_ = Eigen::MatrixXcd::Zero(num_orbitals, num_orbitals);
        } else {
            Eigen::Matrix3Xd dk_cart(3, 2);
            double delta = config_->ham_deriv_delta;
            dk_cart.col(0) = Eigen::Vector3d(delta, 0.0, 0.0); // dkx
            dk_cart.col(1) = Eigen::Vector3d(0.0, delta, 0.0); // dky
            dk_frac_ = tbplas::base::cart2frac(recip_vectors_, dk_cart);
            h0_ = Eigen::MatrixXcd::Zero(num_orbitals, num_orbitals);
            h0_dkx_ = Eigen::MatrixXcd::Zero(num_orbitals, num_orbitals);
            h0_dky_ = Eigen::MatrixXcd::Zero(num_orbitals, num_orbitals);
        }
    }

    /**
     * @brief Calculate derivative of Hamiltonian w.r.t. kx/ky/kz.
     *
     * @param[in] k_point fractional coordinate of k-point
     * @param[out] dh_dkx dH/dkx
     * @param[out] dh_dky dH/dky
     * @param[out] dh_dkz dH/dkz
     *
     * @note 'd' means partial derivative here.
     */
    void eval(
        const Eigen::Vector3d& k_point,
        Eigen::MatrixXcd& dh_dkx,
        Eigen::MatrixXcd& dh_dky,
        Eigen::MatrixXcd& dh_dkz)
    {
        int convention = config_->convention;
        if (config_->ham_deriv_analytical) {
            model_->build_ham_der_dense(k_point, dh_dka_, dh_dkb_, dh_dkc_, convention);
            tbplas::base::df_dk_frac2cart(recip_vectors_, dh_dka_, dh_dkb_, dh_dkc_, dh_dkx, dh_dky, dh_dkz);
        } else {
            double delta = config_->ham_deriv_delta;
            model_->build_ham_dense(k_point, h0_, convention);
            model_->build_ham_dense(k_point + dk_frac_.col(0), h0_dkx_, convention);
            model_->build_ham_dense(k_point + dk_frac_.col(1), h0_dky_, convention);
            dh_dkx = (h0_dkx_ - h0_) / delta;
            dh_dky = (h0_dky_ - h0_) / delta;
        }
    }

private:
    /// @brief pointer to the model
    const model_t* model_ = nullptr;

    /// @brief pointer to config
    const DiagConfig* config_ = nullptr;

    /// @brief Cartesian coordinates of reciprocal lattice vectors in 1/NM
    /// in ROW-WISE order
    Eigen::Matrix3d recip_vectors_;

    /// @brief fractional coordinates of k-points for numerical derivation
    /// @note col(0): dkx, col(1): dky
    Eigen::Matrix3Xd dk_frac_ = Eigen::Matrix3Xd::Zero(3, 2);

    /// @brief derivatives of H w.r.t. fractional coordinates of k-point
    Eigen::MatrixXcd dh_dka_; // dH/dka
    Eigen::MatrixXcd dh_dkb_; // dH/dkb
    Eigen::MatrixXcd dh_dkc_; // dH/dkc

    /// @brief Hamiltonians for numerical derivation
    Eigen::MatrixXcd h0_; // h(k0)
    Eigen::MatrixXcd h0_dkx_; // h(k0+dkx)
    Eigen::MatrixXcd h0_dky_; // h(k0+dky)
};

/**
 * @brief Class for evaluating Berry curvature, Chern number, etc.
 *
 * @tparam model_t datatype of model assigned to this solver
 */
template <typename model_t>
class Berry : public DiagSolver<model_t> {
private:
    /// @brief F matrix for (ki, ki+1)
    Eigen::MatrixXcd fk_;

    /**
     * @brief Evaluate the F matrix from eigenvectors of hk_ at ki and ki+1.
     *
     * @param[in] vec0 eigenvectors at ki
     * @param[in] vec1 eigenvectors at ki+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 ki and ki+1.
     *
     * @param[in] vec0 eigenvectors at ki
     * @param[in] vec1 eigenvectors at ki+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 Generate k-points in extended Brillouin zone.
     *
     * @return Fractional coordinates of k-points in extended Brillouin zone defined by
     *   config.bz_size and config.bz_shift, dimension (3, num_ka*num*num_kc).
     */
    [[nodiscard]] Eigen::Matrix3Xd gen_k_points() const
    {
        const DiagConfig& config = DiagSolver<model_t>::config;
        size_t num_ka = std::get<0>(config.k_grid_size) * std::get<0>(config.bz_size);
        size_t num_kb = std::get<1>(config.k_grid_size) * std::get<1>(config.bz_size);
        size_t num_kc = std::get<2>(config.k_grid_size) * std::get<2>(config.bz_size);
        Eigen::Matrix3Xd k_points = tbplas::base::gen_kmesh({ num_ka, num_kb, num_kc });
        k_points.row(0).array() *= std::get<0>(config.bz_size);
        k_points.row(1).array() *= std::get<1>(config.bz_size);
        k_points.row(2).array() *= std::get<2>(config.bz_size);
        k_points.colwise() += config.bz_shift;
        return k_points;
    }

    /**
     * @brief Get the unit area of 1st Brillouin zone (FBZ).
     *
     * @param[in] recip_vectors Cartesian coordinates of reciprocal lattice vectors
     * in 1/NM in ROW-WISE order
     * @return unit area of FBZ in 1/NM^2
     */
    [[nodiscard]] double get_fbz_unit_area(const Eigen::Matrix3d& recip_vectors) const
    {
        const DiagConfig& config = DiagSolver<model_t>::config;
        size_t num_kpt_fbz = std::get<0>(config.k_grid_size) * std::get<1>(config.k_grid_size);
        double area_fbz = tbplas::base::get_lattice_area(recip_vectors);
        return area_fbz / num_kpt_fbz;
    }

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

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

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

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

    /**
     * @brief Evaluate Berry curvature and Chern number for each band
     * using Kubo formula.
     *
     * @return Berry curvature on k_grid in 1st Brillouin zone (FBZ)
     */
    [[nodiscard]] BerryData calc_berry_curvature_kubo()
    {
        // Extract and check params
        const DiagConfig& config = DiagSolver<model_t>::config;
        assert(std::get<0>(config.k_grid_size) > 0);
        assert(std::get<1>(config.k_grid_size) > 0);
        assert(std::get<2>(config.k_grid_size) > 0);
        assert(std::get<0>(config.bz_size) > 0);
        assert(std::get<1>(config.bz_size) > 0);
        assert(std::get<2>(config.bz_size) > 0);
        auto backend = DiagSolver<model_t>::make_backend();
        size_t num_orbitals = backend->get_num_orbitals();
        size_t num_states = backend->get_num_states();
        int convention = config.convention;

        // Generate k-mesh in FBZ
        Eigen::Matrix3Xd k_points = gen_k_points();
        size_t num_kpt = k_points.cols();

        // Derivatives of Hamiltonian, use numerical derivation by default
        Eigen::MatrixXcd dh_dkx(num_orbitals, num_orbitals);
        Eigen::MatrixXcd dh_dky(num_orbitals, num_orbitals);
        Eigen::MatrixXcd dh_dkz(num_orbitals, num_orbitals);
        Eigen::MatrixXcd prod_dkx(num_states, num_states); // <n|dH/dkx|m>/|En-Em|
        Eigen::MatrixXcd prod_dky(num_states, num_states); // <n|dH/dky|m>/|En-Em|
        HamiltonianDerivative<model_t> ham_deriv(DiagSolver<model_t>::model_, DiagSolver<model_t>::config);

        // Create result
        BerryData data(config.prefix);
        Eigen::Matrix3d recip_vectors = DiagSolver<model_t>::model_->get_reciprocal_lattice();
        data.k_cart = tbplas::base::frac2cart(recip_vectors, k_points);
        data.bands = Eigen::MatrixXd::Zero(num_states, num_kpt);
        data.omega_xy = Eigen::MatrixXd::Zero(num_states, num_kpt);

        // Fill result
        Eigen::VectorXd eigenvalues(num_states);
        Eigen::MatrixXcd eigenstates(num_orbitals, num_states);
        RangeInfo range = DiagSolver<model_t>::dist_range(num_kpt);
        for (int ik = range.k_min; ik <= range.k_max; ++ik) {
            // Evaluate eigenvalues and eigenstates.
            backend->calc_eigen_all(k_points.col(ik), convention, eigenvalues, eigenstates);
            data.bands.col(ik) = eigenvalues;

            // Evaluate matrix elements of dH/dk
            // Consider the upper-triangular part since the matrices are Hermitian
            ham_deriv.eval(k_points.col(ik), dh_dkx, dh_dky, dh_dkz);
            prod_dkx.setZero();
            prod_dky.setZero();
            for (size_t nn = 0; nn < num_states; ++nn) {
                for (size_t mm = nn + 1; mm < num_states; ++mm) {
                    double delta_eng = std::abs(eigenvalues(nn) - eigenvalues(mm));
                    if (delta_eng >= 1.0e-15) {
                        prod_dkx(nn, mm) = eigenstates.col(nn).dot(dh_dkx * eigenstates.col(mm)) / delta_eng;
                        prod_dky(nn, mm) = eigenstates.col(nn).dot(dh_dky * eigenstates.col(mm)) / delta_eng;
                        prod_dkx(mm, nn) = std::conj(prod_dkx(nn, mm));
                        prod_dky(mm, nn) = std::conj(prod_dky(nn, mm));
                    }
                }
            }

            // Sum up products
            //   \sum_m{prod_dkx(n, m) * prod_dky(m, n) - prod_dky(n, m) * prod_dkx(m, n)}
            // = \sum_m{conj(prod_dkx(m, n)) * prod_dky(m, n) - conj(prod_dky(m, n)) * prod_dkx(m, n)}
            // = prod_dkx(:,n).dot(prod_dky(:,n)) - prod_dky(:,n).dot(prod_dkx(:,n))
            // = {0, 2*prod_dkx(:,n).dot(prod_dky(:,n)).imag()}
            // Multiply it by UIMAG, we get -2 * prod_dkx.col(n).dot(prod_dky.col(n)).imag();
            for (size_t nn = 0; nn < num_states; ++nn) {
                data.omega_xy(nn, ik) = -2 * prod_dkx.col(nn).dot(prod_dky.col(nn)).imag();
            }
        }
        DiagSolver<model_t>::mpi_env_.InplaceAllreduceMatrixXd(data.bands);
        DiagSolver<model_t>::mpi_env_.InplaceAllreduceMatrixXd(data.omega_xy);

        // Evaluate Chern number
        //
        // 1. bz_size_ab is the ratio of area of extended BZ with respect to FBZ.
        //
        // 2. Since Chern number is evaluated on extended BZ, we need to divide it by bz_size_ab.
        //    That's why bz_size_ab in the denominator of scale_factor.
        size_t bz_size_ab = std::get<0>(config.bz_size) * std::get<1>(config.bz_size);
        double unit_area = get_fbz_unit_area(recip_vectors);
        double scale_factor = 1.0 / bz_size_ab * unit_area / tbplas::base::TWOPI;
        if (DiagSolver<model_t>::is_master()) {
            for (size_t nn = 0; nn < num_states; ++nn) {
                double chern_number = data.omega_xy.row(nn).sum() * scale_factor;
                std::cout << "Chern number for band " << nn << ": " << chern_number << std::endl;
            }
        }

        // Save and return data
        if (DiagSolver<model_t>::is_master()) {
            data.save();
        }
        // 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;
    }

    /**
     * @brief Evaluate total Berry curvature and Chern number using Wilson loop.
     *
     * @return Berry curvature on k_grid in 1st Brillouin zone (FBZ)
     */
    [[nodiscard]] BerryData calc_berry_curvature_wilson()
    {
        // Extract and check params
        const DiagConfig& config = DiagSolver<model_t>::config;
        assert(std::get<0>(config.k_grid_size) > 0);
        assert(std::get<1>(config.k_grid_size) > 0);
        assert(std::get<2>(config.k_grid_size) > 0);
        assert(std::get<0>(config.bz_size) > 0);
        assert(std::get<1>(config.bz_size) > 0);
        assert(std::get<2>(config.bz_size) > 0);
        auto backend = DiagSolver<model_t>::make_backend();
        size_t num_orbitals = backend->get_num_orbitals();
        size_t num_states = backend->get_num_states();
        int convention = config.convention;
        int num_occ = config.num_occ;
        assert(num_states >= num_occ);

        // Generate k-mesh in FBZ
        Eigen::Matrix3Xd k_points = gen_k_points();
        size_t num_kpt = k_points.cols();

        // 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
        fk_ = Eigen::MatrixXcd::Identity(num_occ, num_occ);
        Eigen::VectorXcd phi = Eigen::VectorXcd::Zero(num_kpt);
        Eigen::VectorXd phase = Eigen::VectorXd::Zero(num_kpt);

        // Local Wilson loop: k0, k0+dka, k0+dka+dkb, k0+dkb
        Eigen::Matrix3Xd k0_loop(3, 4);
        double dka = 1.0 / std::get<0>(config.k_grid_size);
        double dkb = 1.0 / std::get<1>(config.k_grid_size);
        k0_loop.col(0) = Eigen::Vector3d(0.0, 0.0, 0.0); // k0
        k0_loop.col(1) = Eigen::Vector3d(dka, 0.0, 0.0); // k0+dka
        k0_loop.col(2) = Eigen::Vector3d(dka, dkb, 0.0); // k0+dkb+dkb
        k0_loop.col(3) = Eigen::Vector3d(0.0, dkb, 0.0); // k0+dkb

        // Create result
        // We have only the total Berry curvature, so omega_xy has only 1 row.
        BerryData data(config.prefix);
        Eigen::Matrix3d recip_vectors = DiagSolver<model_t>::model_->get_reciprocal_lattice();
        data.k_cart = tbplas::base::frac2cart(recip_vectors, k_points);
        data.bands = Eigen::MatrixXd::Zero(num_states, num_kpt);
        data.omega_xy = Eigen::MatrixXd::Zero(1, num_kpt);

        // Evaluate phases
        RangeInfo range = DiagSolver<model_t>::dist_range(num_kpt);
        for (int ik = range.k_min; ik <= range.k_max; ++ik) {
            phi[ik] = 1.0;

            // Get and backup the eigenvectors of Hamiltonian at k0
            Eigen::Matrix3Xd k_loop = k0_loop.colwise() + k_points.col(ik);
            backend->calc_eigen_all(k_loop.col(0), convention, hk_eigenvalues, vec0);
            data.bands.col(ik) = hk_eigenvalues;
            vec0_bak = vec0;
            if (DiagSolver<model_t>::overlap_ != nullptr) {
                sk0_bak = backend->get_sk();
            }

            // Loop over k-loop
            Eigen::MatrixXcd* p0 = &vec0;
            Eigen::MatrixXcd* p1 = &vec1;
            for (ptrdiff_t jk = 1; jk < k_loop.cols(); ++jk) {
                backend->calc_eigen_all(k_loop.col(jk), convention, hk_eigenvalues, *p1);
                if (DiagSolver<model_t>::overlap_ == nullptr) {
                    eval_fk(*p0, *p1);
                } else {
                    eval_fk(*p0, *p1, backend->get_sk());
                }
                phi[ik] *= fk_.determinant();
                std::swap(p0, p1);
            }

            // Special treatment for the last k-point in k_loop
            if (DiagSolver<model_t>::overlap_ == nullptr) {
                eval_fk(*p0, vec0_bak);
            } else {
                eval_fk(*p0, vec0_bak, sk0_bak);
            }
            phi[ik] *= fk_.determinant();
        }
        phase = -1.0 * phi.cwiseArg();
        Eigen::VectorXd norm = phi.cwiseAbs();
        for (ptrdiff_t i = 0; i < phi.size(); ++i) {
            if (norm[i] < tbplas::base::INFINITESIMAL) {
                phase[i] = 0.0;
            }
        }
        DiagSolver<model_t>::mpi_env_.InplaceAllreduceVectorXd(phase);

        // Evaluate Berry curvature and Chern number
        //
        // 1. bz_size_ab is the ratio of area of extended BZ with respect to FBZ.
        //
        // 2. Since Chern number is evaluated on extended BZ, we need to divide it by bz_size_ab.
        //    That's why bz_size_ab in the denominator of scale_factor.
        size_t bz_size_ab = std::get<0>(config.bz_size) * std::get<1>(config.bz_size);
        double unit_area = get_fbz_unit_area(recip_vectors);
        double scale_factor = 1.0 / bz_size_ab / tbplas::base::TWOPI;
        data.omega_xy = phase.transpose() / unit_area;
        if (DiagSolver<model_t>::is_master()) {
            double chern_number = phase.sum() * scale_factor;
            std::cout << "Chern number for num_occ " << num_occ << ": " << chern_number << std::endl;
        }

        // Save and return data
        if (DiagSolver<model_t>::is_master()) {
            data.save();
        }
        // 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_BERRY_H
