// ========================================================================== //
// Class inclusion
#include "PDS/PDS.hpp"

// Local libraries
#include "Utilities/BBoxMap.h"
#include "Matrix/Matrix.hpp"

// External libraries
#include <GEL/CGLA/CGLA.h>

// Standard libraries
#include <iostream>
#include <map>
#include <queue>
#include <random>

// ========================================================================== //
// Construction and Deconstruction

// Default empty constructor.
PDS::PDS() {}
// Construct a PDS sampler from sampling radius and bounding_box.
PDS::PDS(double radius, BBox &bb) {
    Matrix<double> bounding_box(3, 2);

    bounding_box.set(bb.min_corner[0], 0, 0);
    bounding_box.set(bb.min_corner[1], 1, 0);
    bounding_box.set(bb.min_corner[2], 2, 0);

    bounding_box.set(bb.max_corner[0], 0, 1);
    bounding_box.set(bb.max_corner[1], 1, 1);
    bounding_box.set(bb.max_corner[2], 2, 1);

    // Enter input values into class
    this->pds_r = radius;
    this->bbox  = bounding_box;

    // Compute element at search sizes
    element_size = pds_r / sqrt(bbox.get_dim()[0]);
    search_size  = int(pds_r / element_size + 1);

    // Compute bbox diagonal and setup PDS dimensions
    bbox_diag = 0.0;
    double d;
    pds_dim = std::vector<int>(bbox.get_dim()[0]);
    for (int idx = 0; idx < bbox.get_dim()[0]; idx++) {
        d = bbox.get(idx, 1) - bbox.get(idx, 0);
        bbox_diag += pow(d, 2.0);
        pds_dim[idx] = int(d / element_size) + 2;
    }
    bbox_diag = sqrt(bbox_diag);
}
// Construct a PDS sampler from sampling radius and bounding_box.
PDS::PDS(double radius, Matrix<double> bounding_box) {

    // Enter input values into class
    this->pds_r = radius;
    this->bbox  = bounding_box;

    // Compute element at search sizes
    element_size = pds_r / sqrt(bbox.get_dim()[0]);
    search_size  = int(pds_r / element_size + 1);

    // Compute bbox diagonal and setup PDS dimensions
    bbox_diag = 0.0;
    double d;
    pds_dim = std::vector<int>(bbox.get_dim()[0]);
    for (int idx = 0; idx < bbox.get_dim()[0]; idx++) {
        d = bbox.get(idx, 1) - bbox.get(idx, 0);
        bbox_diag += pow(d, 2.0);
        pds_dim[idx] = int(d / element_size) + 2;
    }
    bbox_diag = sqrt(bbox_diag);
}

PDS::PDS(const PDS &original) {
    this->pds_r        = original.pds_r;
    this->pds_dim      = original.pds_dim;
    this->bbox_diag    = original.bbox_diag;
    this->element_size = original.element_size;
    this->search_size  = original.search_size;

    this->bbox = Matrix<double>(original.bbox);
    this->grid = std::map<std::vector<int>, int>(original.grid);
}
PDS::PDS(PDS &&original) {
    this->pds_r        = original.pds_r;
    this->pds_dim      = original.pds_dim;
    this->bbox_diag    = original.bbox_diag;
    this->element_size = original.element_size;
    this->search_size  = original.search_size;

    this->bbox = Matrix<double>(original.bbox);
    this->grid = std::map<std::vector<int>, int>(original.grid);
}
// Deconstruction
PDS::~PDS() {}

// ========================================================================== //
// Operator overloads

PDS &PDS::operator=(PDS &original) {
    this->pds_r        = original.pds_r;
    this->pds_dim      = original.pds_dim;
    this->bbox_diag    = original.bbox_diag;
    this->element_size = original.element_size;
    this->search_size  = original.search_size;

    this->bbox = Matrix<double>(original.bbox);
    this->grid = std::map<std::vector<int>, int>(original.grid);

    return *this;
}
PDS &PDS::operator=(PDS &&original) {
    this->pds_r        = original.pds_r;
    this->pds_dim      = original.pds_dim;
    this->bbox_diag    = original.bbox_diag;
    this->element_size = original.element_size;
    this->search_size  = original.search_size;

    this->bbox = Matrix<double>(original.bbox);
    this->grid = std::map<std::vector<int>, int>(original.grid);

    return *this;
}

// ========================================================================== //
// Global operations

// Function used to reset the pds grid for a new generation of points.
void PDS::reset() {
    grid.clear();
    points = std::vector<CGLA::Vec3d>();
}

// Function to generate a full PDS sampling in the grid
void PDS::generate(CGLA::Vec3d p0) {
    std::mt19937 rand_dev;
    this->generate(p0, rand_dev);
}
// Function to generate a full PDS sampling in the grid with specified random
// number device.
void PDS::generate(CGLA::Vec3d p0, std::mt19937 rand_dev) {

    // Prepare the RNG.
    std::uniform_real_distribution<double> random_double_0_1(0.0, 1.0);

    if (!inside_bbox(p0))
        fprintf(stderr, "%s:%d\n\t%s", __FILE__, __LINE__,
                "Error when starting PDS generation, the initial point is not "
                "inside the bounding box.");

    // Initialize the queue with the initial point.
    std::queue<CGLA::Vec3d> Q({p0});

    // Run the generation
    while (!Q.empty()) {

        CGLA::Vec3d p = Q.front();
        Q.pop();

        for (int k = 0; k < 30; k++) {
            double      l   = pds_r + pds_r * random_double_0_1(rand_dev);
            CGLA::Vec3d dir = {random_double_0_1(rand_dev) - 0.5,
                               random_double_0_1(rand_dev) - 0.5,
                               random_double_0_1(rand_dev) - 0.5};
            CGLA::Vec3d p1  = p + l * normalize(dir);
            if (!inside_bbox(p1)) continue;
            if (get(p1) != -1) continue;

            add(p1);
            Q.push(p1);
        }
    }
}

// Get the radius used for the grid
double PDS::get_radius() { return pds_r; }
// Return all points
std::vector<CGLA::Vec3d> PDS::get_points() { return points; }

// ========================================================================== //
// Local operations

// Add the point to the list and return the index.
void PDS::add(CGLA::Vec3d p) {
    int              idx = points.size();
    std::vector<int> I   = index(p);
    grid.insert(std::pair<std::vector<int>, int>(I, idx));
    points.push_back(p);
}
// Add the index idx in the pds grid associated to point p
void PDS::add(int idx, CGLA::Vec3d p) {
    std::vector<int> I = index(p);
    grid.insert(std::pair<std::vector<int>, int>(I, idx));
}
// Add the index idx in the pds grid associated to point p
void PDS::add(int idx, std::vector<double> p) {
    std::vector<int> I = index(p);
    grid.insert(std::pair<std::vector<int>, int>(I, idx));
}

// Remove the point to the list and return the index.
void PDS::remove(CGLA::Vec3d p) {
    std::vector<int> I = index(p);
    grid.erase(I);
}

// Return the index stored at position p in the grid.
int PDS::get(CGLA::Vec3d p) {
    std::vector<int> I      = index(p);
    int              result = grid.count(I) > 0 ? grid[I] : -1;
    return result;
}
// Return the index stored at position p in the grid.
int PDS::get(std::vector<double> p) {
    std::vector<int> I      = index(p);
    int              result = grid.count(I) > 0 ? grid[I] : -1;
    return result;
}

// Generate a list of neighbour points which have an index stored.
std::vector<int> PDS::neighbours(CGLA::Vec3d p) {

    // PDS settings and values
    std::vector<int> pds_idx = index(p);

    return neighbours(pds_idx, search_size);
}
// Generate a list of neighbour points which have an index stored.
std::vector<int> PDS::neighbours(std::vector<double> p) {

    // PDS settings and values
    std::vector<int> pds_idx = index(p);

    return neighbours(pds_idx, search_size);
}
// Generate a list of neighbour points which have an index stored.
std::vector<int> PDS::neighbours(CGLA::Vec3d p, double r) {

    // Compute the amount of voxels within the radius
    std::vector<int> pds_idx     = index(p);
    int              search_size = r / element_size + 1;

    return neighbours(pds_idx, search_size);
}

// ========================================================================== //
// Internal functions

// Compute grid index for a point p.
std::vector<int> PDS::index(CGLA::Vec3d p) {
    std::vector<int> result = {int((p[0] - bbox.get(0, 0)) / element_size),
                               int((p[1] - bbox.get(1, 0)) / element_size),
                               int((p[2] - bbox.get(2, 0)) / element_size)};
    return result;
}
// Compute grid index for a point p.
std::vector<int> PDS::index(std::vector<double> p) {
    int N = bbox.get_dim()[0];

    std::vector<int> result(N);
    for (int i = 0; i < N; i++)
        result[i] = int((p[i] - bbox.get(i, 0)) / element_size);

    return result;
}

// Generic function designed to extract values in nearby grid positions
std::vector<int> PDS::neighbours(std::vector<int> p_idx, int search_size) {

    // Compute ranges
    std::vector<int> range_min(p_idx.size()), range_max(p_idx.size());
    for (size_t i = 0; i < p_idx.size(); i++) {
        range_min[i] = std::max(p_idx[i] - search_size, 0);
        range_max[i] = std::min(p_idx[i] + search_size, pds_dim[i] - 1);
    }

    // Setup temporary vector
    std::vector<int> tmp;

    // Extract indices from grid
    std::vector<int> idx = range_min;
    for (size_t i = 1; i < idx.size(); i++) {
        for (idx[0] = range_min[0]; idx[0] <= range_max[0]; idx[0]++) {
            std::map<std::vector<int>, int>::iterator val = grid.find(idx);
            if (val != grid.end()) tmp.push_back(val->second);
        }

        if (++idx[i] < range_max[i]) {
            for (size_t j = 1; j < i; j++) idx[j] = range_min[j];
            i = 0;
        }
    }
    return tmp;
}

// Check if point is inside the bounding box
bool PDS::inside_bbox(CGLA::Vec3d p) {

    bool x = (p[0] >= bbox.get(0, 0) && p[0] <= bbox.get(0, 1));
    bool y = (p[1] >= bbox.get(1, 0) && p[1] <= bbox.get(1, 1));
    bool z = (p[2] >= bbox.get(2, 0) && p[2] <= bbox.get(2, 1));

    return x && y && z;
}

// End of File
// ========================================================================== //