/**
 * @brief Classes for calculating spin texture.
 *
 * @author Yunhai Li (liyunhai1016@hotmail.com)
 *
 * @copyright Copyright (c) 2024, TBPLaS develop team.
 */
#ifndef TBPLAS_DIAG_MAGNETIC_H
#define TBPLAS_DIAG_MAGNETIC_H

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

namespace tbplas::diag {
using tbplas::base::complex_t;
using tbplas::base::UIMAG;

/**
 * @brief Class for calculating spin texture.
 *
 * @tparam model_t datatype of model assigned to this solver
 */
template <typename model_t>
class SpinTexture : public DiagSolver<model_t> {
public:
    /**
     * @brief Construct a new Spin Texture object
     *
     * @param[in] model model for which calculations will be performed
     */
    explicit SpinTexture(const model_t& model)
        : DiagSolver<model_t>(model)
    {
    }

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

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

    /**
     * @brief Evaluate spin texture
     *
     * @return Spin texture on k-points
     */
    [[nodiscard]] SpinData calc_spin_texture()
    {
        // Extract parameters
        auto backend = DiagSolver<model_t>::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 = DiagSolver<model_t>::config.get_k_points();
        size_t num_kpt = k_points.cols();
        int convention = DiagSolver<model_t>::config.convention;
        bool spin_major = DiagSolver<model_t>::config.spin_major;

        // Define sequence to extract spin components
        assert(num_orbitals % 2 == 0);
        size_t num_orb_half = num_orbitals / 2;
        auto seq_up = Eigen::seq(0, num_orb_half - 1, 1);
        auto seq_down = Eigen::seq(num_orb_half, num_orbitals - 1, 1);
        if (!spin_major) {
            seq_up = Eigen::seq(0, num_orbitals - 2, 2);
            seq_down = Eigen::seq(1, num_orbitals - 1, 2);
        }

        // Create result
        SpinData data(DiagSolver<model_t>::config.prefix);
        data.k_cart = tbplas::base::frac2cart(DiagSolver<model_t>::model_->get_reciprocal_lattice(), k_points);
        data.bands = Eigen::MatrixXd::Zero(num_states, num_kpt);
        data.sigma_x = Eigen::MatrixXd::Zero(num_states, num_kpt);
        data.sigma_y = Eigen::MatrixXd::Zero(num_states, num_kpt);
        data.sigma_z = Eigen::MatrixXd::Zero(num_states, num_kpt);

        // Fill result
        Eigen::VectorXd eigenvalues(num_states);
        Eigen::MatrixXcd eigenstates(num_orbitals, num_states);
        Eigen::VectorXcd spin_up(num_orb_half); // half size
        Eigen::VectorXcd spin_down(num_orb_half); // half size
        Eigen::VectorXcd sk_psi(num_orbitals); // sk|psi>, full size
        Eigen::VectorXcd sk_spin_up(num_orb_half); // sk|spin_up>, half size
        Eigen::VectorXcd sk_spin_down(num_orb_half); // sk|spin_down>, half size
        RangeInfo range = DiagSolver<model_t>::dist_range(num_kpt);
        for (int ik = range.k_min; ik <= range.k_max; ++ik) {
            backend->calc_eigen_all(k_points.col(ik), convention, eigenvalues, eigenstates);
            data.bands.col(ik) = eigenvalues;
            for (size_t ib = 0; ib < num_states; ++ib) {
                spin_up = eigenstates(seq_up, ib);
                spin_down = eigenstates(seq_down, ib);
                complex_t ud = 0.0;
                double uu = 0.0;
                double dd = 0.0;
                if (DiagSolver<model_t>::overlap_ == nullptr) {
                    ud = spin_up.dot(spin_down);
                    uu = spin_up.squaredNorm();
                    dd = spin_down.squaredNorm();
                } else {
                    sk_psi = backend->get_sk() * eigenstates.col(ib);
                    sk_spin_up = sk_psi(seq_up);
                    sk_spin_down = sk_psi(seq_down);
                    ud = spin_up.dot(sk_spin_down);
                    uu = spin_up.dot(sk_spin_up).real();
                    dd = spin_down.dot(sk_spin_down).real();
                }
                // sigma_x = (ud + du).real() = 2 * ud.real()
                // sigma_y = (-UIMAG * ud + UIMAG * du).real() = 2 * ud.imag()
                data.sigma_x(ib, ik) = 2 * ud.real();
                data.sigma_y(ib, ik) = 2 * ud.imag();
                data.sigma_z(ib, ik) = uu - dd;
            }
        }
        DiagSolver<model_t>::mpi_env_.InplaceAllreduceMatrixXd(data.bands);
        DiagSolver<model_t>::mpi_env_.InplaceAllreduceMatrixXd(data.sigma_x);
        DiagSolver<model_t>::mpi_env_.InplaceAllreduceMatrixXd(data.sigma_y);
        DiagSolver<model_t>::mpi_env_.InplaceAllreduceMatrixXd(data.sigma_z);
        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
