#ifndef TBPLAS_HOP_DICT_H
#define TBPLAS_HOP_DICT_H

#include <complex>
#include <cstdint>
#include <map>
#include <tuple>

#include "common.h"
#include "memory_view.h"

namespace hop_dict {
using complex_t = std::complex<double>;
using pair_t = std::tuple<int64_t, int64_t>;
using term_t = std::tuple<complex_t, double, double, double>;
using common::calc_dr;
using memory_view::ConstMemoryView;
using memory_view::MemoryView;

/**
 * @brief Container of hopping terms in the supercell.
 *
 * @note 
 *   Since this class is for supercell only, we use int64_t for
 *   orbital indices.
 */
class HopDict {
private:
    std::map<pair_t, term_t> data_;

public:
    explicit HopDict() = default;

    /**
     * @brief Add a new hopping term or update an existing one.
     *
     * @param[in] orb_i orbital index of bra
     * @param[in] orb_j orbital index of ket
     * @param[in] eng hopping energy in eV
     * @param[in] dx displacement along x in NM
     * @param[in] dy displacement along y in NM
     * @param[in] dz displacement along z in NM
     */
    void add_term(
        const int64_t& orb_i,
        const int64_t& orb_j,
        const complex_t& eng,
        const double& dx,
        const double& dy,
        const double& dz)
    {
        // Lower triangular terms will be transformed to upper triangular
        if (orb_i <= orb_j) {
            data_[{ orb_i, orb_j }] = { eng, dx, dy, dz };
        } else {
            data_[{ orb_j, orb_i }] = { std::conj(eng), -dx, -dy, -dz };
        }
    }

    void add_hop_origin(
        const int64_t& num_hop,
        const int64_t* hop_ind,
        const complex_t* hop_eng,
        const double* dr)
    {
        ConstMemoryView<int64_t> hop_ind_map(hop_ind, 2);
        ConstMemoryView<double> dr_map(dr, 3);
        for (int64_t ih = 0; ih < num_hop; ++ih) {
            const int64_t* ptr_ind = hop_ind_map.addr(ih, 0);
            const double* ptr_dr = dr_map.addr(ih, 0);
            add_term(ptr_ind[0], ptr_ind[1], hop_eng[ih], ptr_dr[0], ptr_dr[1], ptr_dr[2]);
        }
    }

    void add_hop_modifier(
        const int64_t& num_hop,
        const int64_t* hop_ind,
        const complex_t* hop_eng,
        const double* sc_lattice,
        const double* sc_orb_pos)
    {
        ConstMemoryView<int64_t> hop_ind_map(hop_ind, 5);
        double dr[3];
        for (int64_t ih = 0; ih < num_hop; ++ih) {
            const int64_t* ptr_ind = hop_ind_map.addr(ih, 0);
            int64_t na = ptr_ind[0];
            int64_t nb = ptr_ind[1];
            int64_t nc = ptr_ind[2];
            int64_t id_sc_i = ptr_ind[3];
            int64_t id_sc_j = ptr_ind[4];
            calc_dr(
                sc_orb_pos, sc_orb_pos, id_sc_i, id_sc_j,
                sc_lattice, na, nb, nc,
                &dr[0]);
            add_term(id_sc_i, id_sc_j, hop_eng[ih], dr[0], dr[1], dr[2]);
        }
    }

    /**
     * @brief Get the number of hopping terms
     *
     * @return number of hopping terms
     */
    [[nodiscard]] uint64_t get_num_hop() const
    {
        return data_.size();
    }

    /**
     * @brief Fill pre-allocated numpy arrays.
     *
     * @param[inout] hop_ind pointer to hopping indices array, dimension (num_hop, 2)
     * @param[inout] hop_eng pointer to hopping energy array, dimension (num_hop,)
     * @param[inout] dr pointer to displacement array, dimension (num_hop, 3)
     */
    void to_array(int64_t* hop_ind, complex_t* hop_eng, double* dr) const
    {
        MemoryView<int64_t> hop_ind_map(hop_ind, 2);
        MemoryView<double> dr_map(dr, 3);
        int64_t counter = 0;
        auto iter = data_.begin();
        for (; iter != data_.end(); iter++) {
            const pair_t& pair = iter->first;
            const term_t& hop = iter->second;
            int64_t* ptr_ind = hop_ind_map.addr(counter, 0);
            double* ptr_dr = dr_map.addr(counter, 0);
            ptr_ind[0] = std::get<0>(pair);
            ptr_ind[1] = std::get<1>(pair);
            hop_eng[counter] = std::get<0>(hop);
            ptr_dr[0] = std::get<1>(hop);
            ptr_dr[1] = std::get<2>(hop);
            ptr_dr[2] = std::get<3>(hop);
            ++counter;
        }
    }
};

} // namespace hop_dict

#endif