#include "meshgrid.h"

namespace tbplas::builder {

void Cube::add_point(const size_t& idx)
{
    point_indices_.emplace_back(idx);
}

[[nodiscard]] size_t Cube::get_num_points() const
{
    return point_indices_.size();
}

[[nodiscard]] const std::vector<size_t>& Cube::get_point_indices() const
{
    return point_indices_;
}

void Cube::set_cube_index(const int64_t& idx)
{
    cube_index_ = idx;
}

[[nodiscard]] const int64_t& Cube::get_cube_index() const
{
    return cube_index_;
}

ReducedTerm::ReducedTerm(const pair_t& pair, const Eigen::Vector3d& rij, const double& distance)
    : pair(pair)
    , rij(rij)
    , distance(distance)
{
}

bool ReducedTerm::operator<(const ReducedTerm& other) const
{
    return distance < other.distance;
}

bool ReducedTerm::operator>(const ReducedTerm& other) const
{
    return distance > other.distance;
}

MeshGrid::MeshGrid(
    const Eigen::Matrix3Xd& pos_bra,
    const Eigen::Matrix3Xd& pos_ket,
    const double& cube_size)
    : pos_bra_(&pos_bra)
    , pos_ket_(&pos_ket)
{
    set_partition(cube_size);
    divide_points(pos_bra_, cubes_bra_);
    divide_points(pos_ket_, cubes_ket_);
}

[[nodiscard]] std::list<ReducedTerm>
MeshGrid::get_neighbors(const double& max_distance)
{
    // Filter empty cubes
    std::vector<const Cube*> filled_cubes_bra;
    for (const auto& c : cubes_bra_) {
        if (c.get_num_points() > 0) {
            filled_cubes_bra.push_back(&c);
        }
    }

    // Build neighbors
    std::list<ReducedTerm> neighbors;
    for (const Cube* cube_bra : filled_cubes_bra) {
        std::vector<const Cube*> neighbor_cubes_ket = get_neighbor_cubes_ket(cube_bra);
        for (const Cube* cube_ket : neighbor_cubes_ket) {
            add_neighbors(cube_bra, cube_ket, max_distance, neighbors);
        }
    }
    return neighbors;
}

void MeshGrid::set_partition(const double& cube_size)
{
    // Determine range of positions
    Eigen::Vector3d pos_min;
    Eigen::Vector3d pos_max;
    auto set_min_max = [&pos_min, &pos_max](
                           const Eigen::Matrix3Xd* pos_bra,
                           const Eigen::Matrix3Xd* pos_ket,
                           int i) {
        double min_bra = pos_bra->row(i).minCoeff();
        double max_bra = pos_bra->row(i).maxCoeff();
        double min_ket = pos_ket->row(i).minCoeff();
        double max_ket = pos_ket->row(i).maxCoeff();
        pos_min[i] = min_bra < min_ket ? min_bra : min_ket;
        pos_max[i] = max_bra > max_ket ? max_bra : max_ket;
    };
    for (int i = 0; i < 3; ++i) {
        set_min_max(pos_bra_, pos_ket_, i);
    }
    pos_min.array() -= 0.1;
    pos_max.array() += 0.1;
    origin_ = pos_min;

    // Determine grid size and dimension
    // Dimension should be larger than grid_size to avoid clipping points
    Eigen::Vector3d max_displace = pos_max - pos_min;
    Eigen::Vector3d grid_size = max_displace / cube_size;
    for (int i = 0; i < 3; ++i) {
        dim_[i] = static_cast<int64_t>(std::ceil(grid_size[i] + 1));
    }

    // Determine lattice vector of unit volume
    for (int i = 0; i < 3; ++i) {
        unit_vec_(i, i) = cube_size;
    }

    // Allocate cubes
    int64_t num_cubes = dim_[0] * dim_[1] * dim_[2];
    cubes_bra_.resize(num_cubes);
    cubes_ket_.resize(num_cubes);
    for (int64_t counter = 0; counter < num_cubes; ++counter) {
        cubes_bra_[counter].set_cube_index(counter);
        cubes_ket_[counter].set_cube_index(counter);
    }
}

[[nodiscard]] int64_t MeshGrid::id_pc2sc(
    const int64_t& ra,
    const int64_t& rb,
    const int64_t& rc) const
{
    return (ra * dim_[1] + rb) * dim_[2] + rc;
}

void MeshGrid::id_sc2pc(
    const int64_t& id_sc,
    int64_t& ra,
    int64_t& rb,
    int64_t& rc) const
{
    // id_sc = (ra * dim_[1] + rb) * dim_[2] + rc.
    int64_t n1 = 0;
    py_div(id_sc, dim_[2], rc, n1);

    // Then n1 = ra * dim_[1] + rb.
    py_div(n1, dim_[1], rb, ra);
}

void MeshGrid::divide_points(const Eigen::Matrix3Xd* pos, std::vector<Cube>& cubes)
{
    double delta = 0.02;
    size_t num_points = pos->cols();
    Eigen::Matrix3d conv_mat = unit_vec_.inverse();
    int64_t ra = 0;
    int64_t rb = 0;
    int64_t rc = 0;
    for (size_t i = 0; i < num_points; ++i) {
        // Get cube index from fractional coordinate.
        Eigen::Vector3d pos_frac = ((pos->col(i) - origin_).transpose() * conv_mat).transpose();
        floor_vector3(pos_frac, ra, rb, rc);
        int64_t cube_idx = id_pc2sc(ra, rb, rc);

        // Add Cartesian coordinate to cube.
        cubes[cube_idx].add_point(i);
    }
}

[[nodiscard]] std::vector<const Cube*>
MeshGrid::get_neighbor_cubes_ket(const Cube* cube_bra) const
{
    std::vector<const Cube*> result;
    int64_t cube_id_bra = cube_bra->get_cube_index();
    size_t num_cubes = cubes_ket_.size();

    // Get cube index in pc representation
    int64_t ra = 0;
    int64_t rb = 0;
    int64_t rc = 0;
    id_sc2pc(cube_id_bra, ra, rb, rc);

    // Loop over neighboring cubes_ket
    for (int da = -1; da <= 1; ++da) {
        for (int db = -1; db <= 1; ++db) {
            for (int dc = -1; dc <= 1; ++dc) {
                int64_t id_sc_tmp = id_pc2sc(ra + da, rb + db, rc + dc);
                if (0 <= id_sc_tmp && id_sc_tmp < num_cubes && cubes_ket_[id_sc_tmp].get_num_points() > 0) {
                    result.push_back(&cubes_ket_[id_sc_tmp]);
                }
            }
        }
    }
    return result;
}

void MeshGrid::add_neighbors(
    const Cube* cube_bra,
    const Cube* cube_ket,
    const double& max_distance,
    std::list<ReducedTerm>& neighbors) const
{
    const std::vector<size_t>& indices_bra = cube_bra->get_point_indices();
    const std::vector<size_t>& indices_ket = cube_ket->get_point_indices();
    double d_max2 = max_distance * max_distance;
    for (const auto& i_bra : indices_bra) {
        for (const auto& i_ket : indices_ket) {
            Eigen::Vector3d rij = pos_ket_->col(i_ket) - pos_bra_->col(i_bra);

            double dist2 = rij[0] * rij[0];
            if (dist2 <= d_max2) { // do not add 0 < dist2 here
                dist2 += rij[1] * rij[1];
                if (dist2 <= d_max2) { // do not add 0 < dist2 here
                    dist2 += rij[2] * rij[2];
                    if (0 < dist2 && dist2 <= d_max2) {
                        neighbors.emplace_back(std::make_pair(i_bra, i_ket), rij, std::sqrt(dist2));
                    }
                }
            }
        }
    }
}
} // namespace tbplas::builder