#include <cassert>
#include <cmath>
#include <numeric>
#include <tuple>
#include <vector>

#include "kpoints.h"
#include "lattice.h"

namespace tbplas::base {
std::tuple<Eigen::Matrix3Xd, Eigen::VectorXi> gen_kpath(const Eigen::Matrix3Xd& hs_kpoints, const std::vector<int>& num_interp)
{
    auto num_kpt_hs = hs_kpoints.cols();
    auto num_path = num_interp.size();
    assert(num_kpt_hs == num_path + 1);
    auto num_kpt_total = std::accumulate(num_interp.begin(), num_interp.end(), 1);
    Eigen::Matrix3Xd kpath(3, num_kpt_total);
    Eigen::VectorXi hs_index(num_kpt_hs);
    int ptr = 0;
    for (int idx = 0; idx < num_path; idx++) {
        hs_index(idx) = ptr;
        Eigen::Vector3d k0 = hs_kpoints.col(idx);
        Eigen::Vector3d k1 = hs_kpoints.col(idx + 1);
        int nk = num_interp[idx];
        for (int idx_j = 0; idx_j < nk; idx_j++) {
            kpath.col(ptr) = k0 + idx_j * 1.0 / nk * (k1 - k0);
            ++ptr;
        }
    }
    hs_index(num_kpt_hs-1) = ptr;
    kpath.col(num_kpt_total-1) = hs_kpoints.col(num_kpt_hs-1);
    return std::tie(kpath, hs_index);
}

Eigen::VectorXd gen_kdist(const Eigen::Matrix3d& lat_vec, const Eigen::Matrix3Xd& kpoints)
{
    // per row
    Eigen::Matrix3d recip_vec = gen_reciprocal_vectors(lat_vec);
    // per col
    Eigen::Matrix3Xd cart_coord = frac2cart(recip_vec, kpoints);
    auto kpo_cols = kpoints.cols();
    Eigen::VectorXd kdist(kpo_cols);
    kdist(0) = 0.0;
    for (int idx = 1; idx < kpo_cols; idx++) {
        Eigen::VectorXd dk = cart_coord.col(idx) - cart_coord.col(idx - 1);
        kdist(idx) = kdist(idx - 1) + dk.norm();
    }
    return kdist;
}

Eigen::Matrix3Xd gen_kmesh(const std::tuple<int, int, int>& grid_size)
{
    int na = std::get<0>(grid_size);
    int nb = std::get<1>(grid_size);
    int nc = std::get<2>(grid_size);
    int idx_col = 0;
    Eigen::Matrix3Xd kmesh(3, na * nb * nc);
    for (const auto& idx_na : Eigen::VectorXd::LinSpaced(na, 0, 1 - 1.0 / na)) {
        for (const auto& idx_nb : Eigen::VectorXd::LinSpaced(nb, 0, 1 - 1.0 / nb)) {
            for (const auto& idx_nc : Eigen::VectorXd::LinSpaced(nc, 0, 1 - 1.0 / nc)) {
                kmesh.col(idx_col) << idx_na, idx_nb, idx_nc;
                idx_col++;
            }
        }
    }
    return kmesh;
}

} // namespace tbplas::base
