/**
 * @brief Functions and classes for solvers based on exact diagonalization
 *
 * @author Yunhai Li (liyunhai1016@hotmail.com)
 *
 * @copyright Copyright (c) 2024, TBPLaS develop team.
 */

#ifndef TBPLAS_DIAG_BASE_H
#define TBPLAS_DIAG_BASE_H

#include <cassert>
#include <iostream>
#include <memory>
#include <regex>
#include <string>

#include <eigen3/Eigen/Dense>

#include "backend.h"
#include "base/datatypes.h"
#include "base/kpoints.h"
#include "base/parallel.h"
#include "config.h"
#include "results.h"

namespace tbplas::diag {
using base::BaseMPIEnv;
using base::complex_t;
using base::RangeInfo;

/**
 * @brief Gaussian type broadening function.
 *
 * @param[in] x incoming x
 * @param[in] mu center of the Gaussian function
 * @param[in] sigma half-width of the Gaussian function
 * @param[out] gau_x normalized Gaussian function value at each x
 */
void gaussian(const Eigen::VectorXd& x, const double& mu, const double& sigma, Eigen::VectorXd& gau_x);

/**
 * @brief Lorentzian type broadening function.
 *
 * @param[in] x incoming x
 * @param[in] mu center of the Lorentzian function
 * @param[in] sigma half-width of the Lorentzian function
 * @param[out] lor_x normalized Lorentzian function value at each x
 */
void lorentzian(const Eigen::VectorXd& x, const double& mu, const double& sigma, Eigen::VectorXd& lor_x);

/**
 * @brief Base class for solvers based on exact diagonalization.
 *
 * @tparam model_t datatype of model assigned to this solver
 */
template <typename model_t>
class DiagSolver {
protected:
    /// @brief pointer to the model for which calculations will be performed
    const model_t* model_ = nullptr;

    /// @brief pointer to the model which holds the overlap
    const model_t* overlap_ = nullptr;

    /// @brief MPI environment fo parallel calculation
    BaseMPIEnv mpi_env_;

    /**
     * @brief Create a backend for diagonalization.
     *
     * @return unique pointer to the backend
     */
    std::unique_ptr<AbstractBackend<model_t>> make_backend() const
    {
        const std::regex pattern_feast_dense(R"(^feast_dense$)", std::regex_constants::icase);
        const std::regex pattern_feast_sparse(R"(^feast_sparse$)", std::regex_constants::icase);
        const std::regex pattern_lapacke(R"(^lapack(e)?$)", std::regex_constants::icase);

        // FEAST dense backend
        if (std::regex_match(config.diag_algo, pattern_feast_dense)) {
#ifdef WITH_FEAST
            print_stdout("Using FEAST dense backend for diagonalization.");
            return std::make_unique<FeastDenseBackend<model_t>>(model_, overlap_, &config);
#else
            print_stderr("ERROR: FEAST support disabled during compilation!");
            std::exit(-1);
#endif
        }

        // FEAST sparse backend
        if (std::regex_match(config.diag_algo, pattern_feast_sparse)) {
#ifdef WITH_FEAST
            print_stdout("Using FEAST sparse backend for diagonalization.");
            return std::make_unique<FeastCSRBackend<model_t>>(model_, overlap_, &config);
#else
            print_stderr("ERROR: FEAST support disabled during compilation!");
            std::exit(-1);
#endif
        }

        // LAPACKE backend
        if (std::regex_match(config.diag_algo, pattern_lapacke)) {
#ifdef WITH_LAPACKE
            print_stdout("Using LAPACKE backend for diagonalization.");
            return std::make_unique<LAPACKEBackend<model_t>>(model_, overlap_, &config);
#else
            print_stderr("ERROR: LAPACKE support disabled during compilation!");
            std::exit(-1);
#endif
        }

        // Default backend
        print_stdout("Using Eigen backend for diagonalization.");
        return std::make_unique<DefaultBackend<model_t>>(model_, overlap_, &config);
    }

public:
    /// @brief configurations
    DiagConfig config;

    /**
     * @brief Construct a new DiagSolver object
     *
     * @param[in] model model for which calculations will be performed
     */
    explicit DiagSolver(const model_t& model)
        : model_(&model)
        , overlap_(nullptr)
    {
    }

    /**
     * @brief Construct a new DiagSolver object with overlap
     *
     * @param[in] model model for which calculations will be performed
     * @param[in] overlap model which holds the overlap
     */
    explicit DiagSolver(const model_t& model, const model_t& overlap)
        : model_(&model)
        , overlap_(&overlap)
    {
    }

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

    /**
     * @brief Check if this process is the master process in the communicator
     *
     * @return true if this process is the master process
     * @return false otherwise
     */
    [[nodiscard]] bool is_master() const
    {
        return mpi_env_.isMaster();
    }

    /**
     * @brief Distribute range [0, n_max-1] into num_group groups and return the
     * lower and upper bounds assigned to this process.
     *
     * @param[in] n_max upper bound of the range
     * @return i_min, i_max
     */
    [[nodiscard]] RangeInfo dist_range(int n_max) const
    {
        return mpi_env_.distRange(n_max, mpi_env_.getSize());
    }

    /**
     * @brief Print message on master process to stdout.
     *
     * @param[in] msg message text
     */
    void print_stdout(const std::string& msg) const
    {
        if (is_master()) {
            std::cout << msg << std::endl;
        }
    }

    /**
     * @brief Print message on master process to stderr.
     *
     * @param[in] msg message text
     */
    void print_stderr(const std::string& msg) const
    {
        if (is_master()) {
            std::cerr << msg << std::endl;
        }
    }

    /**
     * @brief Calculate band structure along k-path
     *
     * @return band structure
     */
    [[nodiscard]] BandData calc_bands()
    {
        // Get array dimensions
        auto backend = make_backend();
        size_t num_states = backend->get_num_states();
        Eigen::Map<const Eigen::Matrix3Xd> k_points = config.get_k_points();
        size_t num_kpt = k_points.cols();

        // Allocate working arrays
        Eigen::VectorXd k_len = tbplas::base::gen_kdist(model_->get_lattice(), k_points);
        Eigen::MatrixXd bands = Eigen::MatrixXd::Zero(num_states, num_kpt);
        Eigen::VectorXd eigenvalues(num_states);

        // Evaluate band structure
        assert((config.convention == 1) || (config.convention == 2));
        RangeInfo range = dist_range(num_kpt);
        for (int ik = range.k_min; ik <= range.k_max; ++ik) {
            backend->calc_eigen_values(k_points.col(ik), config.convention, eigenvalues);
            bands.col(ik) = eigenvalues;
        }
        mpi_env_.InplaceAllreduceMatrixXd(bands);

        // Save and return results
        BandData data(config.prefix);
        data.k_len = std::move(k_len);
        data.bands = std::move(bands);
        if (is_master()) {
            data.save("bands");
        }
        // Master process needs to save data and finishes later than other processes.
        // So we need to put a barrier here.
        mpi_env_.barrier();
        return data;
    }

    /**
     * @brief Calculate density of states for given energy range and step.
     *
     * @return density of states
     */
    [[nodiscard]] DOSData calc_dos()
    {
        // Determine smear type based on the first letter of config.smear_func
        bool smear_gauss = false;
        const std::regex pattern_gauss(R"(^G\w*$)", std::regex_constants::icase);
        if (std::regex_match(config.smear_func, pattern_gauss)) {
            smear_gauss = true;
        }

        // Create energy grid
        Eigen::VectorXd energies;
        config.gen_energies(energies);

        // Evaluate eigenvalues and dos
        Eigen::VectorXd dos = Eigen::VectorXd::Zero(energies.size());
        if (!config.ldos) {
            calc_dos_core(energies, smear_gauss, dos);
        } else {
            calc_ldos_core(energies, smear_gauss, dos);
        }
        mpi_env_.InplaceAllreduceVectorXd(dos);

        // Save and return results
        DOSData data(config.prefix);
        data.energies = std::move(energies);
        data.dos = std::move(dos);
        if (is_master()) {
            data.save("dos");
        }
        // Master process needs to save data and finishes later than other processes.
        // So we need to put a barrier here.
        mpi_env_.barrier();
        return data;
    }

    /**
     * @brief Calculate band structure as well as eigenstates.
     *
     */
    void calc_states()
    {
        // Get array dimensions
        auto backend = make_backend();
        size_t num_orbitals = backend->get_num_orbitals();
        size_t num_states = backend->get_num_states();
        Eigen::Map<const Eigen::Matrix3Xd> k_points = config.get_k_points();
        size_t num_kpt = k_points.cols();

        // Allocate working arrays
        Eigen::VectorXd k_len = tbplas::base::gen_kdist(model_->get_lattice(), k_points);
        Eigen::MatrixXd bands = Eigen::MatrixXd::Zero(num_states, num_kpt);
        Eigen::VectorXd eigenvalues(num_states);
        Eigen::MatrixXcd eigenstates(num_orbitals, num_states);

        // Evaluate band structure and eigen states
        assert((config.convention == 1) || (config.convention == 2));
        RangeInfo range = dist_range(num_kpt);
        EigenStateData data(config.prefix);
        for (int ik = range.k_min; ik <= range.k_max; ++ik) {
            backend->calc_eigen_all(k_points.col(ik), config.convention, eigenvalues, eigenstates);
            bands.col(ik) = eigenvalues;
            data.save_eigen_states(ik, eigenstates);
        }
        mpi_env_.InplaceAllreduceMatrixXd(bands);

        // Save band structure.
        // Wave functions are saved during evaluation to save RAM usage.
        data.k_len = std::move(k_len);
        data.bands = std::move(bands);
        if (is_master()) {
            data.save_eigen_values();
        }
        // Master process needs to save data and finishes later than other processes.
        // So we need to put a barrier here.
        mpi_env_.barrier();
    }

private:
    /**
     * @brief Core function for calculating DOS
     *
     * @param[in] energies energy grid
     * @param[in] smear_gauss whether to use Gaussian smearing
     * @param[out] dos DOS
     */
    void calc_dos_core(
        const Eigen::VectorXd& energies,
        bool smear_gauss,
        Eigen::VectorXd& dos)
    {
        // Get array dimensions
        auto backend = make_backend();
        size_t num_states = backend->get_num_states();
        Eigen::Map<const Eigen::Matrix3Xd> k_points = config.get_k_points();
        size_t num_kpt = k_points.cols();

        // Check params
        assert((config.convention == 1) || (config.convention == 2));

        // Allocate working arrays
        Eigen::VectorXd dos_work = Eigen::VectorXd::Zero(energies.size());
        Eigen::VectorXd eigenvalues(num_states);

        // Evaluate DOS
        RangeInfo range = dist_range(num_kpt);
        for (int ik = range.k_min; ik <= range.k_max; ++ik) {
            backend->calc_eigen_values(k_points.col(ik), config.convention, eigenvalues);
            for (size_t ib = 0; ib < num_states; ++ib) {
                if (smear_gauss) {
                    gaussian(energies, eigenvalues(ib), config.sigma, dos_work);
                } else {
                    lorentzian(energies, eigenvalues(ib), config.sigma, dos_work);
                }
                dos += dos_work;
            }
        }

        // Re-normalize dos
        dos *= (static_cast<double>(config.num_spin) / static_cast<double>(num_states * num_kpt));
    }

    /**
     * @brief Core function for calculating LDOS
     *
     * @param[in] energies energy grid
     * @param[in] smear_gauss whether to use Gaussian smearing
     * @param[out] dos DOS
     */
    void calc_ldos_core(
        const Eigen::VectorXd& energies,
        bool smear_gauss,
        Eigen::VectorXd& dos)
    {
        // Get array dimensions
        auto backend = make_backend();
        size_t num_orbitals = backend->get_num_orbitals();
        size_t num_states = backend->get_num_states();
        Eigen::Map<const Eigen::Matrix3Xd> k_points = config.get_k_points();
        size_t num_kpt = k_points.cols();

        // Check params
        assert((config.convention == 1) || (config.convention == 2));
        assert(!config.ldos_orbital_indices.empty());
        for (const auto& i : config.ldos_orbital_indices) {
            assert((num_orbitals > i) && (i >= 0));
        }

        // Allocate working arrays
        Eigen::VectorXd dos_work = Eigen::VectorXd::Zero(energies.size());
        Eigen::VectorXd eigenvalues(num_states);
        Eigen::MatrixXcd eigenstates(num_orbitals, num_states);
        Eigen::VectorXcd sk_psi(num_orbitals); // Sk|psi>

        // Evaluate DOS
        RangeInfo range = dist_range(num_kpt);
        for (int ik = range.k_min; ik <= range.k_max; ++ik) {
            backend->calc_eigen_all(k_points.col(ik), config.convention, eigenvalues, eigenstates);
            for (size_t ib = 0; ib < num_states; ++ib) {
                const complex_t* eig_vec = eigenstates.col(ib).data();
                if (overlap_ != nullptr) {
                    sk_psi = backend->get_sk() * eigenstates.col(ib);
                    eig_vec = sk_psi.data();
                }
                double proj = 0.0;
                for (const auto& io : config.ldos_orbital_indices) {
                    proj = proj + std::norm(eig_vec[io]); // std::norm is actually norm2.
                }
                if (smear_gauss) {
                    gaussian(energies, eigenvalues(ib), config.sigma, dos_work);
                } else {
                    lorentzian(energies, eigenvalues(ib), config.sigma, dos_work);
                }
                dos += proj * dos_work;
            }
        }

        // Re-normalize dos
        dos *= (static_cast<double>(config.num_spin) / static_cast<double>(num_states * num_kpt));
    }
};

} // namespace tbplas::diag

#endif
