/**
 * @brief Classes for manipulating hopping terms.
 *
 * @author Junyi Wang (jywang)
 *
 * @copyright Copyright (c) 2024, TBPLaS develop team.
 */

#ifndef TBPLAS_BUILDER_HOPPING_H
#define TBPLAS_BUILDER_HOPPING_H

#include <algorithm>
#include <iostream>
#include <map>
#include <set>

#include "base.h"
#include "base/funcs.h"

namespace tbplas::builder {

/**
 * @brief Check whether to take the conjugate part of the hopping term.
 *
 * @tparam ct datatype of cell index
 * @tparam ot datatype of orbital index
 * @param[in] rn (ra, rb, rc) cell index
 * @param[in] pair (orb_i, orb_j) orbital index
 * @return whether to take conjugate
 */
template <typename ct, typename ot>
bool check_conj(const ct rn[], const ot pair[])
{
    for (int i = 0; i < 3; i++) {
        if (rn[i] > 0) {
            return false;
        }
        if (rn[i] < 0) {
            return true;
        }
    }
    return pair[0] > pair[1];
}

/// @brief Template for base class of hopping term containers
/// @tparam hop_t datatype of hopping terms
template <typename hop_t>
class BaseHopping {
protected:
    /// @brief Container of the hopping terms managed by this instance
    std::map<rn_t, std::map<pair_t, hop_t>> data_;

public:
    BaseHopping() = default;
    BaseHopping(const BaseHopping& other)
        : data_(other.data_) {};
    BaseHopping(BaseHopping&& other) noexcept
        : data_(std::move(other.data_)) {};
    BaseHopping& operator=(const BaseHopping& other)
    {
        if (this != &other) {
            data_ = other.data_;
        }
        return *this;
    }
    BaseHopping& operator=(BaseHopping&& other) noexcept
    {
        if (this != &other) {
            data_ = std::move(other.data_);
        }
        return *this;
    }
    ~BaseHopping() = default;

    /// @brief Print hopping terms for debugging purposes
    void print() const
    {
        auto iter_rn = data_.begin();
        for (; iter_rn != data_.end(); iter_rn++) {
            const auto& rn = (*iter_rn).first;
            auto iter_ij = ((*iter_rn).second).begin();
            for (; iter_ij != ((*iter_rn).second).end(); iter_ij++) {
                const auto& pair = (*iter_ij).first;
                const auto& t = (*iter_ij).second;
                cell_index_t ra, rb, rc;
                orbital_index_t orb_i, orb_j;
                std::tie(ra, rb, rc) = rn;
                std::tie(orb_i, orb_j) = pair;
                std::cout << ra << " " << rb << " " << rc << " " << orb_i << " " << orb_j << " " << t << "\n";
            }
        }
    }

    /**
     * @brief Remove the hopping terms corresponding to orbitals in the set, and
     *   update the remaining hopping terms.
     * 
     * @param[in] orbital_indices set of indices of orbitals to be removed
     *
     * @note The C++ compiler of Intel does not support using a reference to an iterator
     *   as the index for a 'for' loop an error in the code:
     *   auto &iter_rn = data_.begin();
     *   for(;iter_rn!=data_.end();)
     */
    void remove_orbitals(const std::set<orbital_index_t>& orbital_indices)
    {
        // Convert set to vector and sort
        std::vector<orbital_index_t> idx_vec;
        idx_vec.assign(orbital_indices.begin(), orbital_indices.end());
        std::sort(idx_vec.begin(), idx_vec.end());

        // Lambda function to remap orbital indices
        auto remap = [&idx_vec](orbital_index_t& orb_i) {
            orb_i -= (std::lower_bound(idx_vec.begin(), idx_vec.end(), orb_i) - idx_vec.begin());
        };

        auto iter_rn = data_.begin();
        for (; iter_rn != data_.end(); iter_rn++) {
            std::map<pair_t, hop_t> new_hop_rn;
            auto iter_ij = ((*iter_rn).second).begin();
            for (; iter_ij != ((*iter_rn).second).end(); iter_ij++) {
                orbital_index_t ii, jj;
                std::tie(ii, jj) = (*iter_ij).first;
                // Check if either index is in orbital_indices
                if (orbital_indices.find(ii) != orbital_indices.end() || orbital_indices.find(jj) != orbital_indices.end()) {
                    continue;
                }
                // Remap the indices
                remap(ii);
                remap(jj);
                new_hop_rn[{ ii, jj }] = (*iter_ij).second;
            }
            // Update hopping terms
            (*iter_rn).second = new_hop_rn;
        }
    }

    /// @brief Remove empty cell indices.
    void purge()
    {
        auto it = data_.begin();
        for (; it != data_.end();) {
            if (((*it).second).empty()) {
                it = data_.erase(it);
            } else {
                it++;
            }
        }
    }

    /// @brief Count the number of hopping terms.
    /// @return number of hopping terms
    [[nodiscard]] orbital_index_t get_num_hopping() const
    {
        orbital_index_t num_hop = 0;
        auto it = data_.begin();
        for (; it != data_.end(); it++) {
            num_hop += ((*it).second).size();
        }
        return num_hop;
    }

    /// @brief Interface for the 'data_' attribute.
    /// @return Reference to data_, read only
    /// @note
    /// Data exposure, packaged by the PrimitiveCell class, finally called in
    /// function merge_prim_cell and extend_prim_cell
    const std::map<rn_t, std::map<pair_t, hop_t>>& get_data() const
    {
        return data_;
    }
};

/// @brief Container class for holding hopping terms within a model
/// @tparam hop_t datatype of hopping terms
template <typename hop_t>
class IntraHopping : public BaseHopping<hop_t> {
private:
    /// @brief Normalize cell index and orbital pair.
    /// For IntraHopping, it should check whether to take the conjugation and
    /// return the status in conj.
    /// @param[inout] rn (ra, rb, rc), cell index
    /// @param[inout] pair (orb_i, orb_j), orbital index
    /// @return the flag of whether to take the conjugate of hopping energy. Incoming rn and pair
    /// are also be modified if conj is true.
    bool norm_keys(rn_t& rn, pair_t& pair) const
    {
        cell_index_t ra, rb, rc;
        orbital_index_t orb_i, orb_j;
        std::tie(ra, rb, rc) = rn;
        std::tie(orb_i, orb_j) = pair;
        cell_index_t rn_array[3] = { ra, rb, rc };
        orbital_index_t pair_array[2] = { orb_i, orb_j };
        bool conj = check_conj(rn_array, pair_array);
        if (conj) {
            rn = std::make_tuple(-ra, -rb, -rc);
            pair = std::make_tuple(orb_j, orb_i);
        }
        return conj;
    }

public:
    /// @brief Add a new hopping term or update existing term
    /// @param[in] ra cell index along a-direction
    /// @param[in] rb cell index along b-direction
    /// @param[in] rc cell index along c-direction
    /// @param[in] orb_i orbital index of bra
    /// @param[in] orb_j orbital index of ket
    /// @param[in] energy incoming hopping energy
    /// @note For IntraHopping, only store normalized hopping terms.
    void add_hopping(const cell_index_t& ra, const cell_index_t& rb,
        const cell_index_t& rc, const orbital_index_t& orb_i,
        const orbital_index_t& orb_j, const hop_t& energy)
    {
        rn_t rn = std::make_tuple(ra, rb, rc);
        pair_t pair = std::make_tuple(orb_i, orb_j);
        bool conj = norm_keys(rn, pair);
        if (conj) {
            BaseHopping<hop_t>::data_[rn][pair] = tbplas::base::conj_t(energy);
        } else {
            BaseHopping<hop_t>::data_[rn][pair] = energy;
        }
    }

    /// @brief Get an existing hopping term
    /// @param[in] ra cell index along a-direction
    /// @param[in] rb cell index along b-direction
    /// @param[in] rc cell index along c-direction
    /// @param[in] orb_i orbital index of bra
    /// @param[in] orb_j orbital index of ket
    /// @param[out] energy incoming hopping energy
    /// @return the flag of whether specified hopping item is found. The result is stored in energy.
    [[nodiscard]] bool get_hopping(const cell_index_t& ra, const cell_index_t& rb,
        const cell_index_t& rc,
        const orbital_index_t& orb_i,
        const orbital_index_t& orb_j,
        hop_t& energy) const
    {
        rn_t rn = std::make_tuple(ra, rb, rc);
        pair_t pair = std::make_tuple(orb_i, orb_j);
        bool conj = norm_keys(rn, pair);
        bool status = false;
        auto iter_rn = BaseHopping<hop_t>::data_.find(rn);
        if (iter_rn != BaseHopping<hop_t>::data_.end()) {
            const auto& hop_rn = (*iter_rn).second;
            auto iter_ij = hop_rn.find(pair);
            if (iter_ij != hop_rn.end()) {
                if (conj) {
                    energy = tbplas::base::conj_t((*iter_ij).second);
                } else {
                    energy = (*iter_ij).second;
                }
                status = true;
            }
        }

        return status;
    }

    /// @brief Remove an existing hopping term
    /// @param[in] ra cell index along a-direction
    /// @param[in] rb cell index along b-direction
    /// @param[in] rc cell index along c-direction
    /// @param[in] orb_i orbital index of bra
    /// @param[in] orb_j orbital index of ket
    /// @return the flag of whether to remove the  specified hopping item
    /// successfully
    /// @note
    /// There might be cases where the value corresponding to the cell index is
    /// empty after using the function remove_hopping. Use the function purge() to
    /// clear the hopping terms.
    [[nodiscard]] bool remove_hopping(const cell_index_t& ra,
        const cell_index_t& rb,
        const cell_index_t& rc,
        const orbital_index_t& orb_i,
        const orbital_index_t& orb_j)
    {
        rn_t rn = std::make_tuple(ra, rb, rc);
        pair_t pair = std::make_tuple(orb_i, orb_j);
        norm_keys(rn, pair);
        bool status = false;
        auto iter_rn = BaseHopping<hop_t>::data_.find(rn);
        if (iter_rn != BaseHopping<hop_t>::data_.end()) {
            auto& hop_rn = (*iter_rn).second;
            auto iter_ij = hop_rn.find(pair);
            if (iter_ij != hop_rn.end()) {
                hop_rn.erase(iter_ij);
                status = true;
            }
        }
        return status;
    }

    /// @brief Remove all the hopping terms of given cell index
    /// @param[in] ra cell index along a-direction
    /// @param[in] rb cell index along b-direction
    /// @param[in] rc cell index along c-direction
    /// @return true if hopping terms are removed, false if not found.
    [[nodiscard]] bool remove_rn(const cell_index_t& ra, const cell_index_t& rb,
        const cell_index_t& rc)
    {
        rn_t rn = std::make_tuple(ra, rb, rc);
        orbital_index_t orb_i = 0;
        pair_t pair = std::make_tuple(orb_i, orb_i);
        norm_keys(rn, pair);
        bool status = false;
        if (BaseHopping<hop_t>::data_.find(rn) != BaseHopping<hop_t>::data_.end()) {
            BaseHopping<hop_t>::data_.erase(rn);
            status = true;
        }
        return status;
    }

    /// @brief Count how many times each orbital pair emerges.
    /// @return map with keys being orbital pairs and value being the counter
    [[nodiscard]] std::map<pair_t, orbital_index_t> count_pairs() const
    {
        std::map<pair_t, orbital_index_t> pair_count;
        auto iter_rn = BaseHopping<hop_t>::data_.begin();
        for (; iter_rn != BaseHopping<hop_t>::data_.end(); iter_rn++) {
            auto iter_ij = ((*iter_rn).second).begin();
            for (; iter_ij != ((*iter_rn).second).end(); iter_ij++) {
                const auto& pair = (*iter_ij).first;
                orbital_index_t orb_i, orb_j;
                std::tie(orb_i, orb_j) = pair;
                if (orb_i <= orb_j) {
                    pair_count[pair] += 1;
                } else {
                    pair_t pair_conj = std::make_tuple(orb_j, orb_i);
                    pair_count[pair_conj] += 1;
                }
            }
        }
        return pair_count;
    }
};

/// @brief Container class for holding hopping terms between two different models
/// @tparam hop_t datatype of hopping terms
template <typename hop_t>
class InterHopping : public BaseHopping<hop_t> {
public:
    /// @brief Add a new hopping term or update existing term.
    /// @param[in] ra cell index along a-direction
    /// @param[in] rb cell index along b-direction
    /// @param[in] rc cell index along c-direction
    /// @param[in] orb_i orbital index of bra
    /// @param[in] orb_j orbital index of ket
    /// @param[in] energy hopping energy
    /// @note For InterHopping, normalization is not needed.
    void add_hopping(const cell_index_t& ra, const cell_index_t& rb,
        const cell_index_t& rc, const orbital_index_t& orb_i,
        const orbital_index_t& orb_j, const hop_t& energy)
    {
        rn_t rn = std::make_tuple(ra, rb, rc);
        pair_t pair = std::make_tuple(orb_i, orb_j);
        BaseHopping<hop_t>::data_[rn][pair] = energy;
    }

    /// @brief Get an existing hopping term
    /// @param[in] ra cell index along a-direction
    /// @param[in] rb cell index along b-direction
    /// @param[in] rc cell index along c-direction
    /// @param[in] orb_i orbital index of bra
    /// @param[in] orb_j orbital index of ket
    /// @param[out] energy outgoing hopping energy
    /// @return the flag of whether specified hopping item is found. The result is stored in energy.
    [[nodiscard]] bool get_hopping(const cell_index_t& ra, const cell_index_t& rb,
        const cell_index_t& rc,
        const orbital_index_t& orb_i,
        const orbital_index_t& orb_j,
        hop_t& energy) const
    {
        rn_t rn = std::make_tuple(ra, rb, rc);
        pair_t pair = std::make_tuple(orb_i, orb_j);
        auto iter_rn = BaseHopping<hop_t>::data_.find(rn);
        bool status = false;
        if (iter_rn != BaseHopping<hop_t>::data_.end()) {
            const auto& hop_rn = (*iter_rn).second;
            auto iter_ij = hop_rn.find(pair);
            if (iter_ij != hop_rn.end()) {
                energy = (*iter_ij).second;
                status = true;
            }
        }
        return status;
    }

    /// @brief Remove an existing hopping term
    /// @param[in] ra cell index along a-direction
    /// @param[in] rb cell index along b-direction
    /// @param[in] rc cell index along c-direction
    /// @param[in] orb_i orbital index of bra
    /// @param[in] orb_j orbital index of ket
    /// @return true if hopping terms are removed, false if not found
    [[nodiscard]] bool remove_hopping(const cell_index_t& ra,
        const cell_index_t& rb,
        const cell_index_t& rc,
        const orbital_index_t& orb_i,
        const orbital_index_t& orb_j)
    {
        rn_t rn = std::make_tuple(ra, rb, rc);
        pair_t pair = std::make_tuple(orb_i, orb_j);
        bool status = false;
        auto iter_rn = BaseHopping<hop_t>::data_.find(rn);
        if (iter_rn != BaseHopping<hop_t>::data_.end()) {
            auto& hop_rn = (*iter_rn).second;
            auto iter_ij = hop_rn.find(pair);
            if (iter_ij != hop_rn.end()) {
                hop_rn.erase(iter_ij);
                status = true;
            }
        }
        return status;
    }

    /// @brief Remove all the hopping terms of given cell index
    /// @param[in] ra cell index along a-direction
    /// @param[in] rb cell index along b-direction
    /// @param[in] rc cell index along c-direction
    /// @return true if hopping terms are removed, false if not found.
    [[nodiscard]] bool remove_rn(const cell_index_t& ra, const cell_index_t& rb,
        const cell_index_t& rc)
    {
        rn_t rn = std::make_tuple(ra, rb, rc);
        bool status = false;
        if (BaseHopping<hop_t>::data_.find(rn) != BaseHopping<hop_t>::data_.end()) {
            BaseHopping<hop_t>::data_.erase(rn);
            status = true;
        }
        return status;
    }
};

} // namespace tbplas::builder
#endif