/**
 * @brief Data containers and filters for constructing models with the Sample class.
 *
 * @author Yunhai Li (liyunhai1016@hotmail.com)
 *
 * @copyright Copyright (c) 2024, TBPLaS develop team.
 */

#ifndef TBPLAS_BUILDER_FILTERS_H
#define TBPLAS_BUILDER_FILTERS_H

#include <tuple>
#include <vector>

#include <eigen3/Eigen/Dense>

#include "kdtree.h"
#include "primitive.h"

namespace tbplas::builder {
using pbc_t = std::tuple<bool, bool, bool>;

/**
 * @brief Orbital data for building arrays in Sample class
 *
 * @note
 *   Attributes of the creator supercell are also included, e.g, the lattice vectors,
 *   boundary conditions, etc.
 */
class OrbitalData {
public:
    /// @brief Cartesian orbital positions in nm, dimension (3, num_orb)
    Eigen::Matrix3Xd orb_pos;

    /// @brief Orbital energies in eV, dimension (num_orb,)
    Eigen::VectorXd orb_eng;

    /// @brief Orbital labels, dimension (num_orb,)
    std::vector<int64_t> orb_label;

    /// @brief Cartesian lattice vectors of creator supercell in nm, with each ROW being a vector
    Eigen::Matrix3d sc_lattice;

    /// @brief Cartesian coordinate of lattice origin of creator supercell in nm
    Eigen::Vector3d sc_origin;

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

    /// @brief Offset for concatenating arrays, to be modified by the Sample class
    /// @note This is also the offset for orbital indices.
    orbital_index_t offset = 0;

    /// @brief Shortcut to get the number of orbitals
    [[nodiscard]] inline orbital_index_t get_num_orbitals() const { return orb_label.size(); }
};

/**
 * @brief Hopping data for building arrays in Sample class
 *
 */
class HoppingData {
public:
    /// @brief hopping indices, dimension (2, num_hop)
    /// @note The indices already contain the offset. For example, if supercell 0 has 10 orbitals,
    /// then orbital indices in supercell 1 start from 10.
    std::vector<std::tuple<orbital_index_t, orbital_index_t>> hop_ind;

    /// @brief hopping energies in eV, dimension (num_hop,)
    Eigen::VectorXcd hop_eng;

    /// @brief Cartesian displacement vectors in nm, dimension (3, num_hop)
    Eigen::Matrix3Xd dr;

    /// @brief Offset for concatenating arrays, to be modified by the Sample class
    /// @note This is NOT the offset for orbital indices, since sc_bra and sc_ket may have different
    /// offsets. Also, the orbital indices in hop_ind already contain the offsets.
    orbital_index_t offset = 0;

    /// @brief Shortcut to get the number of hopping terms
    [[nodiscard]] inline orbital_index_t get_num_hopping() const { return hop_ind.size(); }
};

/**
 * @brief Abstract base class for orbital filters.
 *
 */
class AbstractOrbitalFilter {
public:
    /**
     * @brief Filter orbitals.
     *
     * @param[inout] full_orbitals orbitals to filt
     *
     * @note
     *   Orbitals not satisfying user-defined rules will be removed from the vector.
     */
    virtual void act(std::vector<Orbital>& full_orbitals) const = 0;
};

/**
 * @brief Abstract base class for orbital data modifiers.
 *
 */
class AbstractOrbitalModifier {
public:
    /**
     * @brief Modify orbital positions and energies.
     *
     * @param[inout] orb_data OrbitalData instance to modify
     */
    virtual void act(OrbitalData& orb_data) const = 0;
};

/**
 * @brief Abstract base class for hopping terms filters.
 *
 */
class AbstractHoppingFilter {
public:
    /**
     * @brief Filt hopping terms
     *
     * @param[inout] full_terms hopping terms to filt
     *
     * @note
     *   Hopping terms not satisfying user-defined rules will be removed from the list;
     */
    virtual void act(std::list<Term>& full_terms) const = 0;
};

/**
 * @brief Abstract base class for hopping data modifiers.
 *
 */
class AbstractHoppingModifier {
public:
    /**
     * @brief Modify hopping terms
     *
     * @param[in] data_bra OrbitalData from the 'bra' supercell
     * @param[in] data_ket OrbitalData from the 'ket' supercell
     * @param[inout] hop_data hopping data to modify
     */
    virtual void act(
        const OrbitalData& data_bra,
        const OrbitalData& data_ket,
        HoppingData& hop_data) const
        = 0;
};

/**
 * @brief Filter for intra-hopping terms.
 *
 */
class IntraHoppingFilter : public AbstractHoppingFilter {
public:
    inline void act(std::list<Term>& full_terms) const final
    {
        std::list<Term> new_terms;
        for (const auto& n : full_terms) {
            // Only terms without taking conjugation will be kept.
            if (!check_conj(n)) {
                new_terms.push_back(n);
            }
        }
        full_terms = new_terms;
    }

private:
    /**
     * @brief Checks if this term needs conjugation
     *
     * @param[in] term term to check
     * @return whether conjugation is needed for this term
     *
     * @note See 'check_conj' of hopping.h for more details.
     */
    static bool check_conj(const Term& term)
    {
        if (std::get<0>(term.rn) > 0) {
            return false;
        }
        if (std::get<0>(term.rn) < 0) {
            return true;
        }
        if (std::get<1>(term.rn) > 0) {
            return false;
        }
        if (std::get<1>(term.rn) < 0) {
            return true;
        }
        if (std::get<2>(term.rn) > 0) {
            return false;
        }
        if (std::get<2>(term.rn) < 0) {
            return true;
        }
        return std::get<0>(term.pair) > std::get<1>(term.pair);
    }
};

/**
 * @brief Magnetic field perpendicular to xOy-plane to -z direction via Peierls substitution.
 *
 */
class MagneticField : public AbstractHoppingModifier {
public:
    /**
     * @brief Construct a new MagneticField object
     *
     * @param[in] intensity magnetic field intensity in Tesla
     * @param[in] gauge gauge of vector potential
     */
    explicit MagneticField(
        const double& intensity,
        const int& gauge = 0)
        : intensity_(intensity)
        , gauge_(gauge)
    {
    }

    void act(
        const OrbitalData& data_bra,
        const OrbitalData& data_ket,
        HoppingData& hop_data) const final
    {
        double factor = base::PEIERLS_FACTOR;
        const auto& pos_bra = data_bra.orb_pos;
        const auto& pos_ket = data_ket.orb_pos;
        const auto& offset_bra = data_bra.offset;
        const auto& offset_ket = data_ket.offset;
        const auto& hop_ind = hop_data.hop_ind;
        auto& hop_eng = hop_data.hop_eng; // will be updated so no const qualifier
        const auto& dr = hop_data.dr;
        orbital_index_t orb_i = 0;
        orbital_index_t orb_j = 0;
        for (size_t i = 0; i < hop_data.get_num_hopping(); ++i) {
            std::tie(orb_i, orb_j) = hop_ind[i];
            double dx = dr(0, i);
            double dy = dr(1, i);
            double sx = pos_ket(0, orb_j - offset_ket) + pos_bra(0, orb_i - offset_bra);
            double sy = pos_ket(1, orb_j - offset_ket) + pos_bra(1, orb_i - offset_bra);
            double phase = 0.0;
            if (gauge_ == 0) {
                phase = factor * intensity_ * dx * sy;
            } else if (gauge_ == 1) {
                phase = -factor * intensity_ * dy * sx;
            } else {
                phase = factor * 0.5 * intensity_ * (dx * sx - dy * sy);
            }
            hop_eng[i] *= base::exp_it(phase);
        }
    }

private:
    /// @brief Intensity of magnetic field in Tesla
    double intensity_ = 0.0;

    /// @brief Gauge of vector potential
    int gauge_ = 0;
};
} // namespace tbplas::builder
#endif