#include <falconn/falconn_global.h>
#include <falconn/lsh_nn_table.h>

#include <pybind11/numpy.h>
#include <pybind11/pybind11.h>
#include <pybind11/stl.h>

#include <Eigen/Dense>

#include <memory>

namespace falconn {
namespace python {

class PyLSHNearestNeighborTableError : public FalconnError {
 public:
  PyLSHNearestNeighborTableError(const char *msg) : FalconnError(msg) {}
};

namespace py = pybind11;

template <typename T>
using NumPyArray = py::array_t<T, py::array::c_style>;
template <typename T>
using EigenMap = Eigen::Map<DenseVector<T>>;

template <typename T>
inline EigenMap<T> numpy_to_eigen(NumPyArray<T> x) {
  py::buffer_info buf = x.request();
  if (buf.ndim != 1) {
    throw PyLSHNearestNeighborTableError("expected a one-dimensional array");
  }
  return EigenMap<T>((T *)buf.ptr, buf.shape[0]);
}

template <typename T>
PlainArrayPointSet<T> numpy_to_plain_dataset(NumPyArray<T> dataset) {
  py::buffer_info buf = dataset.request();
  if (buf.ndim != 2) {
    throw PyLSHNearestNeighborTableError("expected a two-dimensional array");
  }
  size_t num_points = buf.shape[0];
  size_t dimension = buf.shape[1];
  PlainArrayPointSet<T> converted_points;
  converted_points.data = (T *)buf.ptr;
  converted_points.num_points = num_points;
  converted_points.dimension = dimension;
  return converted_points;
}

template <typename T>
using LSHTable = LSHNearestNeighborTable<DenseVector<T>>;
template <typename T>
using LSHQueryObject = LSHNearestNeighborQuery<DenseVector<T>>;
template <typename T>
using LSHQueryPool = LSHNearestNeighborQueryPool<DenseVector<T>>;

namespace single_precision {

typedef float ScalarType;
typedef DenseVector<ScalarType> InnerVector;
typedef EigenMap<ScalarType> InnerEigenMap;
typedef PlainArrayPointSet<ScalarType> InnerPlainArrayPointSet;
typedef LSHTable<ScalarType> InnerLSHTable;
typedef LSHQueryObject<ScalarType> InnerLSHQueryObject;
typedef LSHQueryPool<ScalarType> InnerLSHQueryPool;

typedef NumPyArray<ScalarType> OuterNumPyArray;

class PyLSHNearestNeighborQueryDenseFloat {
 public:
  PyLSHNearestNeighborQueryDenseFloat(
      std::shared_ptr<InnerLSHQueryObject> query_object)
      : inner_entity_(query_object) {}

  void set_num_probes(int_fast64_t num_probes) {
    py::gil_scoped_release release;
    inner_entity_->set_num_probes(num_probes);
  }

  int_fast64_t get_num_probes() {
    py::gil_scoped_release release;
    return inner_entity_->get_num_probes();
  }

  void set_max_num_candidates(int_fast64_t max_num_candidates) {
    py::gil_scoped_release release;
    return inner_entity_->set_max_num_candidates(max_num_candidates);
  }

  int_fast64_t get_max_num_candidates() {
    py::gil_scoped_release release;
    return inner_entity_->get_max_num_candidates();
  }

  int32_t find_nearest_neighbor(OuterNumPyArray q) {
    InnerEigenMap converted_query = numpy_to_eigen(q);
    py::gil_scoped_release release;
    return inner_entity_->find_nearest_neighbor(converted_query);
  }

  std::vector<int32_t> find_k_nearest_neighbors(OuterNumPyArray q,
                                                int_fast64_t k) {
    InnerEigenMap converted_query = numpy_to_eigen(q);
    py::gil_scoped_release release;
    std::vector<int32_t> result;
    inner_entity_->find_k_nearest_neighbors(converted_query, k, &result);
    return result;
  }

  std::vector<int32_t> find_near_neighbors(OuterNumPyArray q,
                                           ScalarType threshold) {
    InnerEigenMap converted_query = numpy_to_eigen(q);
    py::gil_scoped_release release;
    std::vector<int32_t> result;
    inner_entity_->find_near_neighbors(converted_query, threshold, &result);
    return result;
  }

  std::vector<int32_t> get_unique_candidates(OuterNumPyArray q) {
    InnerEigenMap converted_query = numpy_to_eigen(q);
    py::gil_scoped_release release;
    std::vector<int32_t> result;
    inner_entity_->get_unique_candidates(converted_query, &result);
    return result;
  }

  std::vector<int32_t> get_candidates_with_duplicates(OuterNumPyArray q) {
    InnerEigenMap converted_query = numpy_to_eigen(q);
    py::gil_scoped_release release;
    std::vector<int32_t> result;
    inner_entity_->get_candidates_with_duplicates(converted_query, &result);
    return result;
  }

  void reset_query_statistics() {
    py::gil_scoped_release release;
    inner_entity_->reset_query_statistics();
  }

  QueryStatistics get_query_statistics() {
    py::gil_scoped_release release;
    return inner_entity_->get_query_statistics();
  }

 private:
  std::shared_ptr<InnerLSHQueryObject> inner_entity_;
};

class PyLSHNearestNeighborQueryPoolDenseFloat {
 public:
  PyLSHNearestNeighborQueryPoolDenseFloat(
      std::shared_ptr<InnerLSHQueryPool> query_pool)
      : inner_entity_(query_pool) {}

  void set_num_probes(int_fast64_t num_probes) {
    py::gil_scoped_release release;
    inner_entity_->set_num_probes(num_probes);
  }

  int_fast64_t get_num_probes() {
    py::gil_scoped_release release;
    return inner_entity_->get_num_probes();
  }

  void set_max_num_candidates(int_fast64_t max_num_candidates) {
    py::gil_scoped_release release;
    return inner_entity_->set_max_num_candidates(max_num_candidates);
  }

  int_fast64_t get_max_num_candidates() {
    py::gil_scoped_release release;
    return inner_entity_->get_max_num_candidates();
  }

  int32_t find_nearest_neighbor(OuterNumPyArray q) {
    InnerEigenMap converted_query = numpy_to_eigen(q);
    py::gil_scoped_release release;
    return inner_entity_->find_nearest_neighbor(converted_query);
  }

  std::vector<int32_t> find_k_nearest_neighbors(OuterNumPyArray q,
                                                int_fast64_t k) {
    InnerEigenMap converted_query = numpy_to_eigen(q);
    py::gil_scoped_release release;
    std::vector<int32_t> result;
    inner_entity_->find_k_nearest_neighbors(converted_query, k, &result);
    return result;
  }

  std::vector<int32_t> find_near_neighbors(OuterNumPyArray q,
                                           ScalarType threshold) {
    InnerEigenMap converted_query = numpy_to_eigen(q);
    py::gil_scoped_release release;
    std::vector<int32_t> result;
    inner_entity_->find_near_neighbors(converted_query, threshold, &result);
    return result;
  }

  std::vector<int32_t> get_unique_candidates(OuterNumPyArray q) {
    InnerEigenMap converted_query = numpy_to_eigen(q);
    py::gil_scoped_release release;
    std::vector<int32_t> result;
    inner_entity_->get_unique_candidates(converted_query, &result);
    return result;
  }

  std::vector<int32_t> get_candidates_with_duplicates(OuterNumPyArray q) {
    InnerEigenMap converted_query = numpy_to_eigen(q);
    py::gil_scoped_release release;
    std::vector<int32_t> result;
    inner_entity_->get_candidates_with_duplicates(converted_query, &result);
    return result;
  }

  void reset_query_statistics() {
    py::gil_scoped_release release;
    inner_entity_->reset_query_statistics();
  }

  QueryStatistics get_query_statistics() {
    py::gil_scoped_release release;
    return inner_entity_->get_query_statistics();
  }

 private:
  std::shared_ptr<InnerLSHQueryPool> inner_entity_;
};

typedef PyLSHNearestNeighborQueryDenseFloat OuterLSHQueryObject;
typedef PyLSHNearestNeighborQueryPoolDenseFloat OuterLSHQueryPool;

class PyLSHNearestNeighborTableDenseFloat {
 public:
  PyLSHNearestNeighborTableDenseFloat(std::shared_ptr<InnerLSHTable> table)
      : table_(table) {}

  std::unique_ptr<OuterLSHQueryObject> construct_query_object(
      int_fast64_t num_probes = -1,
      int_fast64_t max_num_candidates = -1) const {
    std::unique_ptr<InnerLSHQueryObject> inner_query_object =
        table_->construct_query_object(num_probes, max_num_candidates);
    return std::unique_ptr<OuterLSHQueryObject>(
        new OuterLSHQueryObject(std::move(inner_query_object)));
  }

  std::unique_ptr<OuterLSHQueryPool> construct_query_pool(
      int_fast64_t num_probes = -1, int_fast64_t max_num_candidates = -1,
      int_fast64_t num_query_objects = 0) const {
    std::unique_ptr<InnerLSHQueryPool> inner_query_pool =
        table_->construct_query_pool(num_probes, max_num_candidates,
                                     num_query_objects);
    return std::unique_ptr<OuterLSHQueryPool>(
        new OuterLSHQueryPool(std::move(inner_query_pool)));
  }

 private:
  std::shared_ptr<InnerLSHTable> table_;
};

typedef PyLSHNearestNeighborTableDenseFloat OuterLSHTable;

std::unique_ptr<OuterLSHTable> construct_table_dense_float(
    OuterNumPyArray points, const LSHConstructionParameters &params) {
  InnerPlainArrayPointSet converted_points = numpy_to_plain_dataset(points);
  std::unique_ptr<InnerLSHTable> inner_table =
      construct_table<InnerVector, int32_t, InnerPlainArrayPointSet>(
          converted_points, params);
  return std::unique_ptr<OuterLSHTable>(
      new OuterLSHTable(std::move(inner_table)));
}

}  // namespace single_precision

namespace double_precision {

typedef double ScalarType;
typedef DenseVector<ScalarType> InnerVector;
typedef EigenMap<ScalarType> InnerEigenMap;
typedef PlainArrayPointSet<ScalarType> InnerPlainArrayPointSet;
typedef LSHTable<ScalarType> InnerLSHTable;
typedef LSHQueryObject<ScalarType> InnerLSHQueryObject;
typedef LSHQueryPool<ScalarType> InnerLSHQueryPool;

typedef NumPyArray<ScalarType> OuterNumPyArray;

class PyLSHNearestNeighborQueryDenseDouble {
 public:
  PyLSHNearestNeighborQueryDenseDouble(
      std::shared_ptr<InnerLSHQueryObject> query_object)
      : inner_entity_(query_object) {}

  void set_num_probes(int_fast64_t num_probes) {
    py::gil_scoped_release release;
    inner_entity_->set_num_probes(num_probes);
  }

  int_fast64_t get_num_probes() {
    py::gil_scoped_release release;
    return inner_entity_->get_num_probes();
  }

  void set_max_num_candidates(int_fast64_t max_num_candidates) {
    py::gil_scoped_release release;
    return inner_entity_->set_max_num_candidates(max_num_candidates);
  }

  int_fast64_t get_max_num_candidates() {
    py::gil_scoped_release release;
    return inner_entity_->get_max_num_candidates();
  }

  int32_t find_nearest_neighbor(OuterNumPyArray q) {
    InnerEigenMap converted_query = numpy_to_eigen(q);
    py::gil_scoped_release release;
    return inner_entity_->find_nearest_neighbor(converted_query);
  }

  std::vector<int32_t> find_k_nearest_neighbors(OuterNumPyArray q,
                                                int_fast64_t k) {
    InnerEigenMap converted_query = numpy_to_eigen(q);
    py::gil_scoped_release release;
    std::vector<int32_t> result;
    inner_entity_->find_k_nearest_neighbors(converted_query, k, &result);
    return result;
  }

  std::vector<int32_t> find_near_neighbors(OuterNumPyArray q,
                                           ScalarType threshold) {
    InnerEigenMap converted_query = numpy_to_eigen(q);
    py::gil_scoped_release release;
    std::vector<int32_t> result;
    inner_entity_->find_near_neighbors(converted_query, threshold, &result);
    return result;
  }

  std::vector<int32_t> get_unique_candidates(OuterNumPyArray q) {
    InnerEigenMap converted_query = numpy_to_eigen(q);
    py::gil_scoped_release release;
    std::vector<int32_t> result;
    inner_entity_->get_unique_candidates(converted_query, &result);
    return result;
  }

  std::vector<int32_t> get_candidates_with_duplicates(OuterNumPyArray q) {
    InnerEigenMap converted_query = numpy_to_eigen(q);
    py::gil_scoped_release release;
    std::vector<int32_t> result;
    inner_entity_->get_candidates_with_duplicates(converted_query, &result);
    return result;
  }

  void reset_query_statistics() {
    py::gil_scoped_release release;
    inner_entity_->reset_query_statistics();
  }

  QueryStatistics get_query_statistics() {
    py::gil_scoped_release release;
    return inner_entity_->get_query_statistics();
  }

 private:
  std::shared_ptr<InnerLSHQueryObject> inner_entity_;
};

class PyLSHNearestNeighborQueryPoolDenseDouble {
 public:
  PyLSHNearestNeighborQueryPoolDenseDouble(
      std::shared_ptr<InnerLSHQueryPool> query_pool)
      : inner_entity_(query_pool) {}

  void set_num_probes(int_fast64_t num_probes) {
    py::gil_scoped_release release;
    inner_entity_->set_num_probes(num_probes);
  }

  int_fast64_t get_num_probes() {
    py::gil_scoped_release release;
    return inner_entity_->get_num_probes();
  }

  void set_max_num_candidates(int_fast64_t max_num_candidates) {
    py::gil_scoped_release release;
    return inner_entity_->set_max_num_candidates(max_num_candidates);
  }

  int_fast64_t get_max_num_candidates() {
    py::gil_scoped_release release;
    return inner_entity_->get_max_num_candidates();
  }

  int32_t find_nearest_neighbor(OuterNumPyArray q) {
    InnerEigenMap converted_query = numpy_to_eigen(q);
    py::gil_scoped_release release;
    return inner_entity_->find_nearest_neighbor(converted_query);
  }

  std::vector<int32_t> find_k_nearest_neighbors(OuterNumPyArray q,
                                                int_fast64_t k) {
    InnerEigenMap converted_query = numpy_to_eigen(q);
    py::gil_scoped_release release;
    std::vector<int32_t> result;
    inner_entity_->find_k_nearest_neighbors(converted_query, k, &result);
    return result;
  }

  std::vector<int32_t> find_near_neighbors(OuterNumPyArray q,
                                           ScalarType threshold) {
    InnerEigenMap converted_query = numpy_to_eigen(q);
    py::gil_scoped_release release;
    std::vector<int32_t> result;
    inner_entity_->find_near_neighbors(converted_query, threshold, &result);
    return result;
  }

  std::vector<int32_t> get_unique_candidates(OuterNumPyArray q) {
    InnerEigenMap converted_query = numpy_to_eigen(q);
    py::gil_scoped_release release;
    std::vector<int32_t> result;
    inner_entity_->get_unique_candidates(converted_query, &result);
    return result;
  }

  std::vector<int32_t> get_candidates_with_duplicates(OuterNumPyArray q) {
    InnerEigenMap converted_query = numpy_to_eigen(q);
    py::gil_scoped_release release;
    std::vector<int32_t> result;
    inner_entity_->get_candidates_with_duplicates(converted_query, &result);
    return result;
  }

  void reset_query_statistics() {
    py::gil_scoped_release release;
    inner_entity_->reset_query_statistics();
  }

  QueryStatistics get_query_statistics() {
    py::gil_scoped_release release;
    return inner_entity_->get_query_statistics();
  }

 private:
  std::shared_ptr<InnerLSHQueryPool> inner_entity_;
};

typedef PyLSHNearestNeighborQueryDenseDouble OuterLSHQueryObject;
typedef PyLSHNearestNeighborQueryPoolDenseDouble OuterLSHQueryPool;

class PyLSHNearestNeighborTableDenseDouble {
 public:
  PyLSHNearestNeighborTableDenseDouble(std::shared_ptr<InnerLSHTable> table)
      : table_(table) {}

  std::unique_ptr<OuterLSHQueryObject> construct_query_object(
      int_fast64_t num_probes = -1,
      int_fast64_t max_num_candidates = -1) const {
    std::unique_ptr<InnerLSHQueryObject> inner_query_object =
        table_->construct_query_object(num_probes, max_num_candidates);
    return std::unique_ptr<OuterLSHQueryObject>(
        new OuterLSHQueryObject(std::move(inner_query_object)));
  }

  std::unique_ptr<OuterLSHQueryPool> construct_query_pool(
      int_fast64_t num_probes = -1, int_fast64_t max_num_candidates = -1,
      int_fast64_t num_query_objects = 0) const {
    std::unique_ptr<InnerLSHQueryPool> inner_query_pool =
        table_->construct_query_pool(num_probes, max_num_candidates,
                                     num_query_objects);
    return std::unique_ptr<OuterLSHQueryPool>(
        new OuterLSHQueryPool(std::move(inner_query_pool)));
  }

 private:
  std::shared_ptr<InnerLSHTable> table_;
};

typedef PyLSHNearestNeighborTableDenseDouble OuterLSHTable;

std::unique_ptr<OuterLSHTable> construct_table_dense_double(
    OuterNumPyArray points, const LSHConstructionParameters &params) {
  InnerPlainArrayPointSet converted_points = numpy_to_plain_dataset(points);
  std::unique_ptr<InnerLSHTable> inner_table =
      construct_table<InnerVector, int32_t, InnerPlainArrayPointSet>(
          converted_points, params);
  return std::unique_ptr<OuterLSHTable>(
      new OuterLSHTable(std::move(inner_table)));
}

}  // namespace double_precision

PYBIND11_MODULE(_falconn, m) {
  using single_precision::PyLSHNearestNeighborTableDenseFloat;
  using single_precision::PyLSHNearestNeighborQueryDenseFloat;
  using single_precision::PyLSHNearestNeighborQueryPoolDenseFloat;
  using single_precision::construct_table_dense_float;
  using double_precision::PyLSHNearestNeighborTableDenseDouble;
  using double_precision::PyLSHNearestNeighborQueryDenseDouble;
  using double_precision::PyLSHNearestNeighborQueryPoolDenseDouble;
  using double_precision::construct_table_dense_double;

  py::enum_<LSHFamily>(m, "LSHFamily")
      .value("Unknown", LSHFamily::Unknown)
      .value("Hyperplane", LSHFamily::Hyperplane)
      .value("CrossPolytope", LSHFamily::CrossPolytope);
  py::enum_<DistanceFunction>(m, "DistanceFunction")
      .value("Unknown", DistanceFunction::Unknown)
      .value("NegativeInnerProduct", DistanceFunction::NegativeInnerProduct)
      .value("EuclideanSquared", DistanceFunction::EuclideanSquared);
  py::enum_<StorageHashTable>(m, "StorageHashTable")
      .value("Unknown", StorageHashTable::Unknown)
      .value("FlatHashTable", StorageHashTable::FlatHashTable)
      .value("BitPackedFlatHashTable", StorageHashTable::BitPackedFlatHashTable)
      .value("STLHashTable", StorageHashTable::STLHashTable)
      .value("LinearProbingHashTable",
             StorageHashTable::LinearProbingHashTable);
  // we do not expose feature_hashing_dimension, since the wrapper does
  // not support sparse datasets yet
  py::class_<LSHConstructionParameters>(m, "LSHConstructionParameters")
      .def(py::init<>())
      .def_readwrite("dimension", &LSHConstructionParameters::dimension)
      .def_readwrite("lsh_family", &LSHConstructionParameters::lsh_family)
      .def_readwrite("distance_function",
                     &LSHConstructionParameters::distance_function)
      .def_readwrite("k", &LSHConstructionParameters::k)
      .def_readwrite("l", &LSHConstructionParameters::l)
      .def_readwrite("storage_hash_table",
                     &LSHConstructionParameters::storage_hash_table)
      .def_readwrite("num_setup_threads",
                     &LSHConstructionParameters::num_setup_threads)
      .def_readwrite("seed", &LSHConstructionParameters::seed)
      .def_readwrite("last_cp_dimension",
                     &LSHConstructionParameters::last_cp_dimension)
      .def_readwrite("num_rotations",
                     &LSHConstructionParameters::num_rotations);
  // we do not expose a constructor and make all the members read-only
  py::class_<QueryStatistics>(m, "QueryStatistics")
      .def_readonly("average_total_query_time",
                    &QueryStatistics::average_total_query_time)
      .def_readonly("average_lsh_time", &QueryStatistics::average_lsh_time)
      .def_readonly("average_hash_table_time",
                    &QueryStatistics::average_hash_table_time)
      .def_readonly("average_distance_time",
                    &QueryStatistics::average_distance_time)
      .def_readonly("average_num_candidates",
                    &QueryStatistics::average_num_candidates)
      .def_readonly("average_num_unique_candidates",
                    &QueryStatistics::average_num_unique_candidates)
      .def_readonly("num_queries", &QueryStatistics::num_queries);
  // we do not expose a constructor
  py::class_<PyLSHNearestNeighborTableDenseFloat>(
      m, "PyLSHNearestNeighborTableDenseFloat")
      .def("construct_query_object",
           &PyLSHNearestNeighborTableDenseFloat::construct_query_object,
           py::arg("num_probes") = -1, py::arg("max_num_candidates") = -1)
      .def("construct_query_pool",
           &PyLSHNearestNeighborTableDenseFloat::construct_query_pool,
           py::arg("num_probes") = -1, py::arg("max_num_candidates") = -1,
           py::arg("num_query_objects") = 0);
  // we do not expose a constructor
  py::class_<PyLSHNearestNeighborTableDenseDouble>(
      m, "PyLSHNearestNeighborTableDenseDouble")
      .def("construct_query_object",
           &PyLSHNearestNeighborTableDenseDouble::construct_query_object,
           py::arg("num_probes") = -1, py::arg("max_num_candidates") = -1)
      .def("construct_query_pool",
           &PyLSHNearestNeighborTableDenseDouble::construct_query_pool,
           py::arg("num_probes") = -1, py::arg("max_num_candidates") = -1,
           py::arg("num_query_objects") = 0);
  m.def("construct_table_dense_float", &construct_table_dense_float, "");
  m.def("construct_table_dense_double", &construct_table_dense_double, "");

  // we do not expose a constructor
  py::class_<PyLSHNearestNeighborQueryDenseFloat>(
      m, "PyLSHNearestNeighborQueryDenseFloat")
      .def("set_num_probes",
           &PyLSHNearestNeighborQueryDenseFloat::set_num_probes)
      .def("get_num_probes",
           &PyLSHNearestNeighborQueryDenseFloat::get_num_probes)
      .def("set_max_num_candidates",
           &PyLSHNearestNeighborQueryDenseFloat::set_max_num_candidates)
      .def("get_max_num_candidates",
           &PyLSHNearestNeighborQueryDenseFloat::get_max_num_candidates)
      .def("find_nearest_neighbor",
           &PyLSHNearestNeighborQueryDenseFloat::find_nearest_neighbor)
      .def("find_k_nearest_neighbors",
           &PyLSHNearestNeighborQueryDenseFloat::find_k_nearest_neighbors)
      .def("find_near_neighbors",
           &PyLSHNearestNeighborQueryDenseFloat::find_near_neighbors)
      .def("get_unique_candidates",
           &PyLSHNearestNeighborQueryDenseFloat::get_unique_candidates)
      .def("get_candidates_with_duplicates",
           &PyLSHNearestNeighborQueryDenseFloat::get_candidates_with_duplicates)
      .def("reset_query_statistics",
           &PyLSHNearestNeighborQueryDenseFloat::reset_query_statistics)
      .def("get_query_statistics",
           &PyLSHNearestNeighborQueryDenseFloat::get_query_statistics);

  // we do not expose a constructor
  py::class_<PyLSHNearestNeighborQueryPoolDenseFloat>(
      m, "PyLSHNearestNeighborQueryPoolDenseFloat")
      .def("set_num_probes",
           &PyLSHNearestNeighborQueryPoolDenseFloat::set_num_probes)
      .def("get_num_probes",
           &PyLSHNearestNeighborQueryPoolDenseFloat::get_num_probes)
      .def("set_max_num_candidates",
           &PyLSHNearestNeighborQueryPoolDenseFloat::set_max_num_candidates)
      .def("get_max_num_candidates",
           &PyLSHNearestNeighborQueryPoolDenseFloat::get_max_num_candidates)
      .def("find_nearest_neighbor",
           &PyLSHNearestNeighborQueryPoolDenseFloat::find_nearest_neighbor)
      .def("find_k_nearest_neighbors",
           &PyLSHNearestNeighborQueryPoolDenseFloat::find_k_nearest_neighbors)
      .def("find_near_neighbors",
           &PyLSHNearestNeighborQueryPoolDenseFloat::find_near_neighbors)
      .def("get_unique_candidates",
           &PyLSHNearestNeighborQueryPoolDenseFloat::get_unique_candidates)
      .def("get_candidates_with_duplicates",
           &PyLSHNearestNeighborQueryPoolDenseFloat::
               get_candidates_with_duplicates)
      .def("reset_query_statistics",
           &PyLSHNearestNeighborQueryPoolDenseFloat::reset_query_statistics)
      .def("get_query_statistics",
           &PyLSHNearestNeighborQueryPoolDenseFloat::get_query_statistics);

  // we do not expose a constructor
  py::class_<PyLSHNearestNeighborQueryDenseDouble>(
      m, "PyLSHNearestNeighborQueryDenseDouble")
      .def("set_num_probes",
           &PyLSHNearestNeighborQueryDenseDouble::set_num_probes)
      .def("get_num_probes",
           &PyLSHNearestNeighborQueryDenseDouble::get_num_probes)
      .def("set_max_num_candidates",
           &PyLSHNearestNeighborQueryDenseDouble::set_max_num_candidates)
      .def("get_max_num_candidates",
           &PyLSHNearestNeighborQueryDenseDouble::get_max_num_candidates)
      .def("find_nearest_neighbor",
           &PyLSHNearestNeighborQueryDenseDouble::find_nearest_neighbor)
      .def("find_k_nearest_neighbors",
           &PyLSHNearestNeighborQueryDenseDouble::find_k_nearest_neighbors)
      .def("find_near_neighbors",
           &PyLSHNearestNeighborQueryDenseDouble::find_near_neighbors)
      .def("get_unique_candidates",
           &PyLSHNearestNeighborQueryDenseDouble::get_unique_candidates)
      .def(
          "get_candidates_with_duplicates",
          &PyLSHNearestNeighborQueryDenseDouble::get_candidates_with_duplicates)
      .def("reset_query_statistics",
           &PyLSHNearestNeighborQueryDenseDouble::reset_query_statistics)
      .def("get_query_statistics",
           &PyLSHNearestNeighborQueryDenseDouble::get_query_statistics);

  // we do not expose a constructor
  py::class_<PyLSHNearestNeighborQueryPoolDenseDouble>(
      m, "PyLSHNearestNeighborQueryPoolDenseDouble")
      .def("set_num_probes",
           &PyLSHNearestNeighborQueryPoolDenseDouble::set_num_probes)
      .def("get_num_probes",
           &PyLSHNearestNeighborQueryPoolDenseDouble::get_num_probes)
      .def("set_max_num_candidates",
           &PyLSHNearestNeighborQueryPoolDenseDouble::set_max_num_candidates)
      .def("get_max_num_candidates",
           &PyLSHNearestNeighborQueryPoolDenseDouble::get_max_num_candidates)
      .def("find_nearest_neighbor",
           &PyLSHNearestNeighborQueryPoolDenseDouble::find_nearest_neighbor)
      .def("find_k_nearest_neighbors",
           &PyLSHNearestNeighborQueryPoolDenseDouble::find_k_nearest_neighbors)
      .def("find_near_neighbors",
           &PyLSHNearestNeighborQueryPoolDenseDouble::find_near_neighbors)
      .def("get_unique_candidates",
           &PyLSHNearestNeighborQueryPoolDenseDouble::get_unique_candidates)
      .def("get_candidates_with_duplicates",
           &PyLSHNearestNeighborQueryPoolDenseDouble::
               get_candidates_with_duplicates)
      .def("reset_query_statistics",
           &PyLSHNearestNeighborQueryPoolDenseDouble::reset_query_statistics)
      .def("get_query_statistics",
           &PyLSHNearestNeighborQueryPoolDenseDouble::get_query_statistics);

  m.def("compute_number_of_hash_functions",
        &compute_number_of_hash_functions<DenseVector<float>>, "",
        py::arg("num_hash_bits"), py::arg("params"));
  m.def("get_default_parameters", &get_default_parameters<DenseVector<float>>,
        "", py::arg("num_points"), py::arg("dimension"),
        py::arg("distance") = DistanceFunction::EuclideanSquared,
        py::arg("is_sufficiently_dense") = false);
}
}  // namespace python
}  // namespace falconn


//==================================================================================================
/*!
  @file

  @copyright 2017 NumScale SAS

  Distributed under the Boost Software License, Version 1.0.
  (See accompanying file LICENSE.md or copy at http://boost.org/LICENSE_1_0.txt)
**/
//==================================================================================================
#ifndef BOOST_SIMD_ALGORITHM_FIND_HPP_INCLUDED
#define BOOST_SIMD_ALGORITHM_FIND_HPP_INCLUDED

#include <boost/simd/range/segmented_aligned_range.hpp>
#include <boost/simd/detail/is_aligned.hpp>
#include <boost/simd/function/any.hpp>
#include <boost/simd/pack.hpp>
#include <algorithm>

namespace boost { namespace simd
{
  /*!
    @ingroup group-algo

    Returns an iterator to the first element in the range [first,last)
    for which the element equals val.
    If no such element is found, the function returns last.

    @param first  Beginning of the range of elements to find
    @param last   End of the range of elements to find
    @param val    Value to search for in the range.

    @par Requirement

      - @c first, @c last and @c out must be pointer to Vectorizable type.

      - @c val must be a scalar value convertible to the pointee type of first.

    @par Example:

       @snippet find.cpp find

    @par possible output:

       @snippet find.txt find

    @return An iterator to the first element in the range that compares equal to val.
            If no elements match, the function returns last.
  **/
  template<typename T, typename U>
  T const* find(T const* first, T const* last, U const & val)
  {
    auto pr = segmented_aligned_range(first,last);

    auto r = std::find(pr.head.begin(), pr.head.end(), val);
    if (r != pr.head.end()) return r;

    pack<T> v(val);
    auto rv = std::find_if( pr.body.begin(), pr.body.end()
                          , [&v](const pack<T>& x) { return any(v == x); }
                          );

    if (rv != pr.body.end())
    {
      pack<T> vv = *rv;
      auto z = std::find(vv.begin(), vv.end(), val);

      // If there is somethign ina vector, compute the scalar iterator to it
      return  pr.head.end()                                           //   base
            + std::distance(pr.body.begin(), rv)*pack<T>::static_size // + jump to vector
            + std::distance(vv.begin(), z);                           // + jump to element
    }

    return std::find(pr.tail.begin(), pr.tail.end(), val);
  }
} }

#endif


//==================================================================================================
/*!
  @file

  @copyright 2016 NumScale SAS
  @copyright 2016 J.T. Lapreste

  Distributed under the Boost Software License, Version 1.0.
  (See accompanying file LICENSE.md or copy at http://boost.org/LICENSE_1_0.txt)
*/
//==================================================================================================
#ifndef BOOST_SIMD_FUNCTION_IS_NOT_LESS_EQUAL_HPP_INCLUDED
#define BOOST_SIMD_FUNCTION_IS_NOT_LESS_EQUAL_HPP_INCLUDED

#if defined(DOXYGEN_ONLY)
namespace boost { namespace simd
{

 /*!
    @ingroup group-predicates
    Function object implementing is_not_less_equal capabilities

    Returns @ref True  or @ref False depending on whether x is
    not less or equal to y or is.

    @par Semantic:

    @code
    as_logical_t<T> r = is_not_less_equal(x,y);
    @endcode

    is similar to:

    @code
    as_logical_t<T> r = !(x <= y);
    @endcode

    @par Note:

    Due to existence of nan, this is not equivalent to a call to @ref is_greater
    for floating types

    @par Alias:

    @c is_nle

  **/
  const boost::dispatch::functor<tag::is_not_less_equal_> is_not_less_equal = {};
} }
#endif

#include <boost/simd/function/scalar/is_not_less_equal.hpp>
#include <boost/simd/function/simd/is_not_less_equal.hpp>

#endif


//==============================================================================
//         Copyright 2003 - 2011   LASMEA UMR 6602 CNRS/Univ. Clermont II
//         Copyright 2009 - 2011   LRI    UMR 8623 CNRS/Univ Paris Sud XI
//         Copyright 2012          Domagoj Saric, Little Endian Ltd.
//
//          Distributed under the Boost Software License, Version 1.0.
//                 See accompanying file LICENSE.txt or copy at
//                     http://www.boost.org/LICENSE_1_0.txt
//==============================================================================
#define NT2_UNIT_MODULE "nt2 signal processing - test"

#include <nt2/toolbox/signal/static_fft.hpp>

#include <nt2/sdk/unit/tests.hpp>
#include <nt2/sdk/unit/module.hpp>
#include <nt2/sdk/unit/details/helpers.hpp>
#include <nt2/include/functions/ulpdist.hpp>

#include <boost/simd/sdk/memory/allocator.hpp>
#include <boost/simd/sdk/memory/aligned_array.hpp>
#include <boost/array.hpp>
#include <boost/assert.hpp>
#include <boost/foreach.hpp>
#include <boost/random/mersenne_twister.hpp>
#include <boost/random/uniform_real_distribution.hpp>

#ifdef __APPLE__
    #include "Accelerate/Accelerate.h" //vDSP.h
#endif // __APPLE__

#include <fstream>
#include <vector>

namespace
{
    namespace constants
    {
        static std::size_t const minimum_dft_size = 128 ;
        static std::size_t const maximum_dft_size = 8192;
        static std::size_t const test_dft_size    = 4096;

        static int const test_data_range_minimum = -1;
        static int const test_data_range_maximum = +1;

        /// \todo The maximum error values should be made dependent on the:
        ///  - value type (float, double, ...)
        ///  - size of the transform
        ///  - transform type (complex, real-through-complex, pure real)
        ///  - used algorithm or more precisely, its arithmetic operation count
        ///    (currently this is the split-radix algorithm)
        ///  - used complex multiplication implementation (2 multiplies + 4
        ///    additions or 3 multiplies + 3 additions) and whether FMA is used
        ///  (see the "accuracy/precision" links in static_fft.hpp).
        ///                                   (17.07.2012.) (Domagoj Saric)

        /// \todo More extensive testing:
        ///  - test multiple sizes (requires the above theoretical-limit
        ///    function)
        ///  - log/graph our actual error "response" ( f_err( N ) )
        ///  - consider/investigate maximum allowed average ULPD thresholds
        ///  - forward complex transform of real data (result must be conjugated
        ///    symmetric)
        ///  - inverse complex transform of conjugated symmetric data (result
        ///    must be pure real
        ///  - http://en.wikipedia.org/wiki/Parseval's_theorem test
        ///  - more fixed/predefined test vectors which have known/expected
        ///    results/transforms (like sawtooth, pulse, sines...)
        ///                                   (19.07.2012.) (Domagoj Saric)

        /// \todo Investigate why testing with NDEBUG defined gives much better
        /// results (on MSVC10). Without NDEBUG MSVC10, Clang3 and GCC 4.x give
        /// exactly the same results.
        ///                                   (24.07.2012.) (Domagoj Saric)

        static unsigned int const maximum_allowed_complex_nt2_ulpd   = 1030;
#if defined(BOOST_SIMD_ARCH_X86) && !defined(BOOST_SIMD_HAS_FMA4_SUPPORT)
        static unsigned int const maximum_allowed_real_nt2_ulpd      =  580;
#else
        static unsigned int const maximum_allowed_real_nt2_ulpd      = 1290;
#endif

        static unsigned int const maximum_allowed_complex_apple_ulpd = 1290;
        static unsigned int const maximum_allowed_real_apple_ulpd    =  770;
    } // namespace constants

    static std::size_t const N = constants::test_dft_size;
    #if 0 //defined( BOOST_SIMD_HAS_LRB_SUPPORT ) || defined( BOOST_SIMD_HAS_AVX_SUPPORT )
        typedef double T;
    #else //...zzz...cardinal-must-be-4 limitation...
        typedef float T;
    #endif // BOOST_SIMD_HAS_LRB_SUPPORT || BOOST_SIMD_HAS_AVX_SUPPORT
    typedef boost::simd::memory::aligned_array<T, N, BOOST_SIMD_CONFIG_ALIGNMENT> aligned_array;
    typedef boost::simd::memory::aligned_array<T, N/2 + 1, BOOST_SIMD_CONFIG_ALIGNMENT> aligned_half_complex_array;
    typedef nt2::static_fft<constants::minimum_dft_size, constants::maximum_dft_size, T> FFT;

    void randomize( aligned_array & data )
    {
        /// \note Portably reproducible pseudo "random" values.
        ///                                   (24.07.2012.) (Domagoj Saric)
        boost::random::mt19937                            prng        ( 42                                                                     );
        boost::random::uniform_real_distribution<T> const distribution( constants::test_data_range_minimum, constants::test_data_range_maximum );
        BOOST_FOREACH( T & scalar, data )
            scalar = distribution( prng );
    }

    void scale( aligned_array & data, T const scale_factor )
    {
        BOOST_FOREACH( T & scalar, data ) { scalar *= scale_factor; }
    }

    struct failed_value_t
    {
        T           value        ;
        T           desired_value;
        T           ulp_error    ;
        std::size_t index        ;
    };

    std::vector<failed_value_t> failed_values;

    //#define NT2_PSIHA_AUX_TESTING
#ifdef NT2_PSIHA_AUX_TESTING
    std::ofstream outputfile_aux( "nt_fft_test.log", std::ios_base::out | std::ios_base::trunc );
    std::ostream & debug_output( outputfile_aux );
#else
    std::ostream & debug_output( std::cout );
#endif // NT2_PSIHA_AUX_TESTING

    void analyze_values
    (
        aligned_array const &       results           ,
        aligned_array const &       desired_results   ,
        double                const max_ulp_distance  ,
        char          const * const result_description,
        char          const * const calling_function  ,
        unsigned int          const source_file_line
    )
    {
        failed_values.clear();

        double average_ulpd( 0 );
        double max_ulpd    ( 0 );

        for ( std::size_t i( 0 ); i != N; ++i )
        {
            T      const value_a( results        [ i ] );
            T      const value_b( desired_results[ i ] );
            double const ulpd   ( nt2::ulpdist( nt2::details::smallest_a( value_a, value_b ), nt2::details::smallest_b( value_a, value_b ) ) );
            average_ulpd += ulpd;
            max_ulpd      = std::max( max_ulpd, ulpd );
            if ( ulpd > max_ulp_distance )
            {
                failed_value_t const failed_value = { value_a, value_b, static_cast<T>( ulpd ), i };
                failed_values.push_back( failed_value );
            }
        }

        debug_output << std::endl;

        nt2::unit::test_count()++;
        if ( failed_values.empty() )
        {
            nt2::unit::pass( result_description );
            debug_output
                << "   * average ULPD: " << static_cast<T>( average_ulpd / N ) << std::endl
                << "   * max     ULPD: " << static_cast<T>( max_ulpd         ) << " (allowed: " << max_ulp_distance << ")" << std::endl;
        }
        else
        {
            nt2::unit::fail( result_description, source_file_line, calling_function );

            debug_output << "   because the ULP distance for the following "
                         << failed_values.size()
                         << " " << result_description
                         << " values was too large:"
                         << std::endl;

            BOOST_FOREACH( failed_value_t const & failed_value, failed_values )
            {
                debug_output << "\tvalue: "         << failed_value.value
                             << ", desired value: " << failed_value.desired_value
                             << ", ULP distance: "  << failed_value.ulp_error
                             << ", value index: "   << failed_value.index
                             << std::endl;
            }

            debug_output
                << "   * average ULPD: " << static_cast<T>( average_ulpd / N ) << std::endl
                << "   * max     ULPD: " << static_cast<T>( max_ulpd         ) << " (allowed: " << max_ulp_distance << ")" << std::endl;
        }
    } // analyze_values

    #define NT2_FFT_ANALYZE_VALUES( results, desired_results, max_ulp_distance, result_description ) \
        analyze_values( results, desired_results, max_ulp_distance, result_description, BOOST_CURRENT_FUNCTION, __LINE__ )

} // anonymous namespace

NT2_TEST_CASE( test )
{
    // Test complex transform(s):
    {
        aligned_array real_data;
        aligned_array imag_data;

        randomize( real_data );
        randomize( imag_data );

        aligned_array const original_real_data( real_data );
        aligned_array const original_imag_data( imag_data );

        FFT::forward_transform( &real_data[ 0 ], &imag_data[ 0 ], N );
        FFT::inverse_transform( &real_data[ 0 ], &imag_data[ 0 ], N );

        T const norm( nt2::complex_fft_normalization_factor<T>( N ) );
        scale( real_data, norm );
        scale( imag_data, norm );

        NT2_FFT_ANALYZE_VALUES( real_data, original_real_data, constants::maximum_allowed_complex_nt2_ulpd, "complex real data"      );
        NT2_FFT_ANALYZE_VALUES( imag_data, original_imag_data, constants::maximum_allowed_complex_nt2_ulpd, "complex imaginary data" );

    #ifdef __APPLE__
        { // Accelerate framework FFT
            real_data = original_real_data;
            imag_data = original_imag_data;

            FFTSetup const fft_instance( vDSP_create_fftsetup( boost::static_log2<N>::value, kFFTRadix2 ) );
            BOOST_ASSERT( fft_instance );
            DSPSplitComplex complex_data = { &real_data[ 0 ], &imag_data[ 0 ] };
            vDSP_fft_zip( fft_instance, &complex_data, 1, boost::static_log2<N>::value, FFT_FORWARD );
            vDSP_fft_zip( fft_instance, &complex_data, 1, boost::static_log2<N>::value, FFT_INVERSE );
            vDSP_destroy_fftsetup( fft_instance );

            T const norm( static_cast<T>( 1 ) / N );
            scale( real_data, norm );
            scale( imag_data, norm );

            NT2_FFT_ANALYZE_VALUES( real_data, original_real_data, constants::maximum_allowed_complex_apple_ulpd, "apple complex real data"      );
            NT2_FFT_ANALYZE_VALUES( imag_data, original_imag_data, constants::maximum_allowed_complex_apple_ulpd, "apple complex imaginary data" );
        }
    #endif // __APPLE__
    }

    // Test real transform(s):
    {
        aligned_array real_time_data;
        randomize( real_time_data );

        aligned_half_complex_array real_frequency_data;
        aligned_half_complex_array imag_frequency_data;
        FFT::real_forward_transform( &real_time_data[ 0 ], &real_frequency_data[ 0 ], &imag_frequency_data[ 0 ], N );

        aligned_array real_time_data2;
        FFT::real_inverse_transform( &real_frequency_data[ 0 ], &imag_frequency_data[ 0 ], &real_time_data2[ 0 ], N );

        T const norm( nt2::real_fft_normalization_factor<T>( N ) );
        scale( real_time_data2, norm );

        NT2_FFT_ANALYZE_VALUES( real_time_data2, real_time_data, constants::maximum_allowed_real_nt2_ulpd, "real data" );

    #ifdef __APPLE__
        { // Accelerate framework FFT
            FFTSetup const fft_instance( vDSP_create_fftsetup( boost::static_log2<N>::value, kFFTRadix2 ) );
            BOOST_ASSERT( fft_instance );
            aligned_array split_real_data_storage;
            DSPSplitComplex split_real_data =
            {
                &split_real_data_storage[ 0     ],
                &split_real_data_storage[ N / 2 ]
            };
            vDSP_ctoz    ( reinterpret_cast<DSPComplex const *>( &real_time_data[ 0 ] ), 2, &split_real_data, 1, N / 2 );
            vDSP_fft_zrip( fft_instance, &split_real_data, 1, boost::static_log2<N>::value, FFT_FORWARD );
            vDSP_fft_zrip( fft_instance, &split_real_data, 1, boost::static_log2<N>::value, FFT_INVERSE );
            vDSP_ztoc    ( &split_real_data, 1, reinterpret_cast<DSPComplex *>( &real_time_data2[ 0 ] ), 2, N / 2 );
            vDSP_destroy_fftsetup( fft_instance );

            T const norm( static_cast<T>( 1 ) / ( 2 * N ) );
            scale( real_time_data2, norm );

            NT2_FFT_ANALYZE_VALUES( real_time_data, real_time_data2, constants::maximum_allowed_real_apple_ulpd, "apple real data" );
        }
    #endif // __APPLE__
    }
}


//    Copyright 2017 Rainer Gemulla
//
//    Licensed under the Apache License, Version 2.0 (the "License");
//    you may not use this file except in compliance with the License.
//    You may obtain a copy of the License at
//
//        http://www.apache.org/licenses/LICENSE-2.0
//
//    Unless required by applicable law or agreed to in writing, software
//    distributed under the License is distributed on an "AS IS" BASIS,
//    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
//    See the License for the specific language governing permissions and
//    limitations under the License.
// TODO: add all kinds of arguments

#include <iostream>

#include <boost/filesystem.hpp>
#include <boost/program_options.hpp>
#include <boost/random/uniform_real.hpp>

#include <util/evaluation.h>
#include <mf/mf.h>

using namespace std;
using namespace mf;
using namespace rg;
using namespace boost::program_options;

log4cxx::LoggerPtr logger(log4cxx::Logger::getLogger("main"));

int main(int argc, char *argv[]) {
	string traceFile;
	string traceVar;
	string inputMatrixFile;
	mf_size_type epochs;

	// read command line
	options_description desc("Options");
	desc.add_options()
		("help", "produce help message")
		("epochs", value<mf_size_type>(&epochs)->default_value(10), "number of epochs to run [10]")
		("trace", value<string>(&traceFile)->default_value("trace.R"), "filename of trace [trace.R]")
		("traceVar", value<string>(&traceVar)->default_value("trace"), "variable name for trace [traceVar]")
	    ("input-file", value<string>(&inputMatrixFile), "input matrix")
	    ;

	positional_options_description pdesc;
	pdesc.add("input-file", 1);

	variables_map vm;
	store(command_line_parser(argc, argv).options(desc).positional(pdesc).run(), vm);
	notify(vm);

	if (vm.count("help") || vm.count("input-file")==0) {
		cout << "mfsgd [options] <input-file> <sample-file>" << endl << endl;
	    cout << desc << endl;
	    return 1;
	}

	// read input matrix
	SparseMatrix v;
	readMatrix(inputMatrixFile, v);
	LOG4CXX_INFO(logger, "Data matrix: "
		<< v.size1() << " x " << v.size2() << ", " << v.nnz() << " nonzeros");
	mf_size_type size1 = v.size1();
	mf_size_type size2 = v.size2();
	mf_size_type r = 50;

	// compute sample matrix
	Random32 random; // note: this takes a default seed (not randomized!)
	ProjectedSparseMatrix Vsample;
	projectRandomSubmatrix(random, v, Vsample, v.size1()/10, v.size2()/10);
	projectFrequent(Vsample, 0);
	LOG4CXX_INFO(logger, "Sample matrix: "
		<< Vsample.data.size1() << " x " << Vsample.data.size2()
		<< ", " << Vsample.data.nnz() << " nonzeros");

	// generate initial factors by sampling from a uniform[-0.5,0.5] distribution
	DenseMatrix w(size1, r);
	DenseMatrixCM h(r, size2);
	generateRandom(w, random, boost::uniform_real<>(-0.5, 0.5));
	generateRandom(h, random, boost::uniform_real<>(-0.5, 0.5));

	// parameters for SGD
	double epsMax = 0.1;
	double lambda = 50;
	typedef UpdateTruncate<UpdateNzsl> Update;
	typedef RegularizeL2 Regularize;
	typedef SumLoss<NzslLoss, L2Loss> Loss;
	Update update = Update(UpdateNzsl(), -1000, 10100); // truncate for numerical stability
	Regularize regularize = Regularize(lambda);
	Loss loss((NzslLoss()), L2Loss(lambda));

	// initialize the SGD
	Timer t;
	SgdRunner sgdRunner(random);
	SgdJob<Update,Regularize> job(v, w, h, update, regularize, SGD_ORDER_WOR);
	DecayAuto<Update,Regularize,Loss> decay(job, loss, Vsample, epsMax, 7);
	Trace trace;

	// run SGD
	t.start();
	sgdRunner.run(job, loss, epochs, decay, trace, BALANCE_NONE, BALANCE_SIMPLE);
	t.stop();
	LOG4CXX_INFO(logger, "Total time: " << t);

	// write trace to an R file
	LOG4CXX_INFO(logger, "Writing trace to " << traceFile);
	trace.toRfile(traceFile, traceVar);

	return 0;
}


// graph-tool -- a general graph modification and manipulation thingy
//
// Copyright (C) 2006-2018 Tiago de Paula Peixoto <tiago@skewed.de>
//
// This program is free software; you can redistribute it and/or
// modify it under the terms of the GNU General Public License
// as published by the Free Software Foundation; either version 3
// of the License, or (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU General Public License for more details.
//
// You should have received a copy of the GNU General Public License
// along with this program. If not, see <http://www.gnu.org/licenses/>.

#ifndef GRAPH_REWIRING_HH
#define GRAPH_REWIRING_HH

#include <tuple>
#include <iostream>
#include <boost/functional/hash.hpp>

#include "graph.hh"
#include "graph_filtering.hh"
#include "graph_util.hh"
#include "sampler.hh"

#include "random.hh"

#include "hash_map_wrap.hh"

namespace graph_tool
{
using namespace std;
using namespace boost;

template <class Graph>
typename graph_traits<Graph>::vertex_descriptor
source(const pair<size_t, bool>& e,
       const vector<typename graph_traits<Graph>::edge_descriptor>& edges,
       const Graph& g)
{
    if (e.second)
        return target(edges[e.first], g);
    else
        return source(edges[e.first], g);
}

template <class Graph>
typename graph_traits<Graph>::vertex_descriptor
target(const pair<size_t, bool>& e,
       const vector<typename graph_traits<Graph>::edge_descriptor>& edges,
       const Graph& g)
{
    if (e.second)
        return source(edges[e.first], g);
    else
        return target(edges[e.first], g);
}

template <class Nmap, class Graph>
void add_count(size_t s, size_t t, Nmap& nvmap, Graph& g)
{
    if (!graph_tool::is_directed(g) && s > t)
        std::swap(s, t);
    auto& nmap = nvmap[s];
    nmap[t]++;
}

template <class Nmap, class Graph>
void remove_count(size_t s, size_t t, Nmap& nvmap, Graph& g)
{
    if (!graph_tool::is_directed(g) && s > t)
        std::swap(s, t);
    auto& nmap = nvmap[s];
    auto iter = nmap.find(t);
    iter->second--;
    if (iter->second == 0)
        nmap.erase(iter);
}

template <class Nmap, class Graph>
size_t get_count(size_t s, size_t t, Nmap& nvmap, Graph& g)
{
    if (!graph_tool::is_directed(g) && s > t)
        std::swap(s, t);
    auto& nmap = nvmap[s];
    auto iter = nmap.find(t);
    if (iter == nmap.end())
        return 0;
    return iter->second;
    // if (s != t)
    //     return iter->second;
    // else
    //     return iter->second / 2;
}

// this functor will swap the source of the edge e with the source of edge se
// and the target of edge e with the target of te
struct swap_edge
{
    template <class Nmap, class Graph>
    static bool
    parallel_check_target (const pair<size_t, bool>& e,
                           const pair<size_t, bool>& te,
                           vector<typename graph_traits<Graph>::edge_descriptor>& edges,
                           Nmap& nmap,
                           const Graph &g)
    {
        // We want to check that if we swap the target of 'e' with the target of
        // 'te', as such
        //
        //  (s)    -e--> (t)          (s)    -e--> (nt)
        //  (te_s) -te-> (nt)   =>    (te_s) -te-> (t)
        //
        // no parallel edges are introduced.

        typename graph_traits<Graph>::vertex_descriptor
            s = source(e, edges, g),          // current source
            t = target(e, edges, g),          // current target
            nt = target(te, edges, g),        // new target
            te_s = source(te, edges, g);      // target edge source

        if (get_count(s,  nt, nmap, g) > 0)
            return true; // e would clash with an existing edge
        if (get_count(te_s, t, nmap, g) > 0)
            return true; // te would clash with an existing edge
        return false; // no parallel edges
    }

    template <class Graph>
    static void swap_target
        (const pair<size_t, bool>& e,
         const pair<size_t, bool>& te,
         vector<typename graph_traits<Graph>::edge_descriptor>& edges,
         Graph& g)
    {
        // swap the target of the edge 'e' with the target of edge 'te', as
        // such:
        //
        //  (s)    -e--> (t)          (s)    -e--> (nt)
        //  (te_s) -te-> (nt)   =>    (te_s) -te-> (t)

        if (e.first == te.first)
            return;

        // new edges which will replace the old ones
        typename graph_traits<Graph>::edge_descriptor ne, nte;
        typename graph_traits<Graph>::vertex_descriptor
            s_e = source(e, edges, g),
            t_e = target(e, edges, g),
            s_te = source(te, edges, g),
            t_te = target(te, edges, g);
        remove_edge(edges[e.first], g);
        remove_edge(edges[te.first], g);

        if (graph_tool::is_directed(g) || !e.second)
            ne = add_edge(s_e, t_te, g).first;
        else // keep invertedness (only for undirected graphs)
            ne = add_edge(t_te, s_e, g).first;
        edges[e.first] = ne;
        if (graph_tool::is_directed(g) || !te.second)
            nte = add_edge(s_te, t_e, g).first;
        else // keep invertedness (only for undirected graphs)
            nte = add_edge(t_e, s_te,  g).first;
        edges[te.first] = nte;
    }

};

// used for verbose display
void print_progress(size_t i, size_t n_iter, size_t current, size_t total,
                    stringstream& str)
{
    size_t atom = (total > 200) ? total / 100 : 1;
    if ( ( (current+1) % atom == 0) || (current + 1) == total)
    {
        size_t size = str.str().length();
        for (size_t j = 0; j < str.str().length(); ++j)
            cout << "\b";
        str.str("");
        str << "(" << i + 1 << " / " << n_iter << ") "
            << current + 1 << " of " << total << " ("
            << (current + 1) * 100 / total << "%)";
        for (int j = 0; j < int(size - str.str().length()); ++j)
            str << " ";
        cout << str.str() << flush;
    }
}

//select blocks based on in/out degrees
class DegreeBlock
{
public:
    typedef pair<size_t, size_t> block_t;

    template <class Graph>
    block_t get_block(typename graph_traits<Graph>::vertex_descriptor v,
                      const Graph& g) const
    {
        return make_pair(in_degreeS()(v, g), out_degree(v, g));
    }
};

//select blocks based on property map
template <class PropertyMap>
class PropertyBlock
{
public:
    typedef typename property_traits<PropertyMap>::value_type block_t;

    PropertyBlock(PropertyMap p): _p(p) {}

    template <class Graph>
    block_t get_block(typename graph_traits<Graph>::vertex_descriptor v,
                      const Graph&) const
    {
        return get(_p, v);
    }

private:
    PropertyMap _p;
};

// select an appropriate "null" key for densehash
template <class Type>
struct get_null_key
{
    Type operator()() const
    {
        return numeric_limits<Type>::max();
    }
};

template <>
struct get_null_key<string>
{
    string operator()() const
    {
        return lexical_cast<string>(get_null_key<size_t>()());
    }
};

template <>
struct get_null_key<boost::python::object>
{
    boost::python::object operator()() const
    {
        return boost::python::object();
    }
};

template <class Type>
struct get_null_key<vector<Type>>
{
    vector<Type> operator()() const
    {
        vector<Type> v(1);
        v[0] = get_null_key<Type>()();
        return v;
    }
};

template <class Type1, class Type2>
struct get_null_key<pair<Type1, Type2>>
{
    pair<Type1, Type2> operator()() const
    {
        return make_pair(get_null_key<Type1>()(),
                         get_null_key<Type2>()());
    }
};

// main rewire loop
template <template <class Graph, class EdgeIndexMap, class CorrProb,
                    class BlockDeg>
          class RewireStrategy>
struct graph_rewire
{

    template <class Graph, class EdgeIndexMap, class CorrProb,
              class BlockDeg, class PinMap>
    void operator()(Graph& g, EdgeIndexMap edge_index, CorrProb corr_prob,
                    PinMap pin, bool self_loops, bool parallel_edges,
                    bool configuration, pair<size_t, bool> iter_sweep,
                    std::tuple<bool, bool, bool> cache_verbose, size_t& pcount,
                    rng_t& rng, BlockDeg bd) const
    {
        typedef typename graph_traits<Graph>::edge_descriptor edge_t;
        bool persist = std::get<0>(cache_verbose);
        bool cache = std::get<1>(cache_verbose);
        bool verbose = std::get<2>(cache_verbose);

        vector<edge_t> edges;
        vector<size_t> edge_pos;
        typename graph_traits<Graph>::edge_iterator e, e_end;
        for (tie(e, e_end) = boost::edges(g); e != e_end; ++e)
        {
            if (pin[*e])
                continue;
            edges.push_back(*e);
            edge_pos.push_back(edge_pos.size());
        }

        typedef random_permutation_iterator<typename vector<size_t>::iterator,
                                            rng_t>
            random_edge_iter;

        RewireStrategy<Graph, EdgeIndexMap, CorrProb, BlockDeg>
            rewire(g, edge_index, edges, corr_prob, bd, cache, rng,
                   parallel_edges, configuration);

        size_t niter;
        bool no_sweep;
        tie(niter, no_sweep) = iter_sweep;
        pcount = 0;
        if (verbose)
            cout << "rewiring edges: ";
        stringstream str;
        for (size_t i = 0; i < niter; ++i)
        {
            random_edge_iter
                ei_begin(edge_pos.begin(), edge_pos.end(), rng),
                ei_end(edge_pos.end(), edge_pos.end(), rng);

            for (random_edge_iter ei = ei_begin; ei != ei_end; ++ei)
            {
                size_t e_pos = ei - ei_begin;
                if (verbose)
                    print_progress(i, niter, e_pos, no_sweep ? 1 : edges.size(),
                                   str);

                size_t e = *ei;

                bool success = false;
                do
                {
                    success = rewire(e, self_loops, parallel_edges);
                }
                while(persist && !success);

                if (!success)
                    ++pcount;

                if (no_sweep)
                    break;
            }
        }
        if (verbose)
            cout << endl;
    }

    template <class Graph, class EdgeIndexMap, class CorrProb, class PinMap>
    void operator()(Graph& g, EdgeIndexMap edge_index, CorrProb corr_prob,
                    PinMap pin, bool self_loops, bool parallel_edges,
                    bool configuration, pair<size_t, bool> iter_sweep,
                    std::tuple<bool, bool, bool> cache_verbose, size_t& pcount,
                    rng_t& rng) const
    {
        operator()(g, edge_index, corr_prob, pin, self_loops, parallel_edges,
                   configuration, iter_sweep, cache_verbose, pcount, rng,
                   DegreeBlock());
    }
};

// this will rewire the edges so that the resulting graph will be entirely
// random (i.e. Erdos-Renyi)
template <class Graph, class EdgeIndexMap, class CorrProb, class BlockDeg>
class ErdosRewireStrategy
{
public:
    typedef typename graph_traits<Graph>::edge_descriptor edge_t;
    typedef typename EdgeIndexMap::value_type index_t;

    ErdosRewireStrategy(Graph& g, EdgeIndexMap edge_index,
                        vector<edge_t>& edges, CorrProb, BlockDeg,
                        bool, rng_t& rng, bool, bool configuration)
        : _g(g), _edge_index(edge_index), _edges(edges),
          _vertices(HardNumVertices()(g)), _rng(rng),
          _configuration(configuration),
          _nmap(get(vertex_index, g), num_vertices(g))
    {
        decltype(_vertices.begin()) viter = _vertices.begin();
        typename graph_traits<Graph>::vertex_iterator v, v_end;
        for (tie(v, v_end) = vertices(_g); v != v_end; ++v)
            *(viter++) = *v;

        if (!configuration)
        {
            for (size_t i = 0; i < edges.size(); ++i)
                add_count(source(edges[i], g), target(edges[i], g), _nmap, g);
        }
    }

    bool operator()(size_t ei, bool self_loops, bool parallel_edges)
    {
        size_t e_s = source(_edges[ei], _g);
        size_t e_t = target(_edges[ei], _g);

        if (!graph_tool::is_directed(_g) && e_s > e_t)
            std::swap(e_s, e_t);

        //try randomly drawn pairs of vertices
        std::uniform_int_distribution<size_t> sample(0, _vertices.size() - 1);
        typename graph_traits<Graph>::vertex_descriptor s, t;
        while (true)
        {
            s = sample(_rng);
            t = sample(_rng);
            if (s == t)
            {
                if (!self_loops) // reject self-loops if not allowed
                    continue;

            }
            else if (!graph_tool::is_directed(_g) && self_loops)
            {
                // sample self-loops w/ correct probability for undirected
                // graphs
                std::bernoulli_distribution reject(.5);
                if (reject(_rng))
                    continue;
            }
            break;
        }

        if (!graph_tool::is_directed(_g) && s > t)
            std::swap(s, t);

        if (s == e_s && t == e_t)
            return false;

        // reject parallel edges if not allowed
        if (!parallel_edges && is_adjacent(s, t, _g))
            return false;

        if (!_configuration)
        {
            size_t m = get_count(s, t, _nmap, _g);
            size_t m_e = get_count(e_s, e_t, _nmap, _g);

            double a = (m + 1) / double(m_e);

            std::bernoulli_distribution accept(std::min(a, 1.));
            if (!accept(_rng))
                return false;
        }

        remove_edge(_edges[ei], _g);
        edge_t ne = add_edge(s, t, _g).first;
        _edges[ei] = ne;

        if (!_configuration)
        {
            remove_count(e_s, e_t, _nmap, _g);
            add_count(s, t, _nmap, _g);
        }

        return true;
    }

private:
    Graph& _g;
    EdgeIndexMap _edge_index;
    vector<edge_t>& _edges;
    vector<typename graph_traits<Graph>::vertex_descriptor> _vertices;
    rng_t& _rng;
    bool _configuration;
    typedef gt_hash_map<size_t, size_t> nmapv_t;
    typedef typename property_map_type::apply<nmapv_t,
                                              typename property_map<Graph, vertex_index_t>::type>
        ::type::unchecked_t nmap_t;
    nmap_t _nmap;
};

// this is the mother class for edge-based rewire strategies
// it contains the common loop for finding edges to swap, so different
// strategies need only to specify where to sample the edges from.
template <class Graph, class EdgeIndexMap, class RewireStrategy>
class RewireStrategyBase
{
public:
    typedef typename graph_traits<Graph>::edge_descriptor edge_t;
    typedef typename graph_traits<Graph>::vertex_descriptor vertex_t;

    typedef typename EdgeIndexMap::value_type index_t;

    RewireStrategyBase(Graph& g, EdgeIndexMap edge_index, vector<edge_t>& edges,
                       rng_t& rng, bool parallel_edges, bool configuration)
        : _g(g), _edge_index(edge_index), _edges(edges), _rng(rng),
          _nmap(get(vertex_index, g), num_vertices(g)),
          _configuration(configuration)
    {
        if (!parallel_edges || !configuration)
        {
            for (size_t i = 0; i < edges.size(); ++i)
                add_count(source(edges[i], g), target(edges[i], g), _nmap, g);
        }
    }

    bool operator()(size_t ei, bool self_loops, bool parallel_edges)
    {
        RewireStrategy& self = *static_cast<RewireStrategy*>(this);

        // try randomly drawn pairs of edges and check if they satisfy all the
        // consistency checks

        pair<size_t, bool> e = make_pair(ei, false);

        // rewire target
        pair<size_t, bool> et = self.get_target_edge(e, parallel_edges);

        if (et.first == ei)
            return false;

        auto s = source(e,  _edges, _g);
        auto t = target(e,  _edges, _g);
        auto ts = source(et,  _edges, _g);
        auto tt = target(et,  _edges, _g);

        if (!self_loops) // reject self-loops if not allowed
        {
            if(s == tt || ts == t)
                return false;
        }

        // reject parallel edges if not allowed
        if (!parallel_edges && (et.first != e.first))
        {
            if (swap_edge::parallel_check_target(e, et, _edges, _nmap, _g))
                return false;
        }

        double a = 0;

        if (!graph_tool::is_directed(_g))
        {
            a -= log(2 + (s == t) + (ts == tt));
            a += log(2 + (s == tt) + (ts == t));
        }

        if (!_configuration)
        {
            map<std::pair<size_t, size_t>, int> delta;

            delta[std::make_pair(s, t)] -= 1;
            delta[std::make_pair(ts, tt)] -= 1;
            delta[std::make_pair(s, tt)] += 1;
            delta[std::make_pair(ts, t)] += 1;

            for (auto& e_d : delta)
            {
                auto u = e_d.first.first;
                auto v = e_d.first.second;
                int d = e_d.second;
                size_t m = get_count(u,  v,  _nmap, _g);
                a -= lgamma(m + 1) - lgamma((m + 1) + d);
                if (!graph_tool::is_directed(_g) && u == v)
                    a += d * log(2);
            }

        }

        std::bernoulli_distribution accept(std::min(exp(a), 1.));
        if (!accept(_rng))
            return false;

        self.update_edge(e.first, false);
        self.update_edge(et.first, false);

        if (!parallel_edges || !_configuration)
        {
            remove_count(source(e, _edges, _g), target(e, _edges, _g), _nmap, _g);
            remove_count(source(et, _edges, _g), target(et, _edges, _g), _nmap, _g);
        }

        swap_edge::swap_target(e, et, _edges, _g);

        self.update_edge(e.first, true);
        self.update_edge(et.first, true);

        if (!parallel_edges || !_configuration)
        {
            add_count(source(e, _edges, _g), target(e, _edges, _g), _nmap, _g);
            add_count(source(et, _edges, _g), target(et, _edges, _g), _nmap, _g);
        }

        return true;
    }

protected:
    Graph& _g;
    EdgeIndexMap _edge_index;
    vector<edge_t>& _edges;
    rng_t& _rng;

    typedef gt_hash_map<size_t, size_t> nmapv_t;
    typedef typename property_map_type::apply<nmapv_t,
                                              typename property_map<Graph, vertex_index_t>::type>
        ::type::unchecked_t nmap_t;
    nmap_t _nmap;
    bool _configuration;
};

// this will rewire the edges so that the combined (in, out) degree distribution
// will be the same, but all the rest is random
template <class Graph, class EdgeIndexMap, class CorrProb, class BlockDeg>
class RandomRewireStrategy:
    public RewireStrategyBase<Graph, EdgeIndexMap,
                              RandomRewireStrategy<Graph, EdgeIndexMap,
                                                   CorrProb, BlockDeg> >
{
public:
    typedef RewireStrategyBase<Graph, EdgeIndexMap,
                               RandomRewireStrategy<Graph, EdgeIndexMap,
                                                    CorrProb, BlockDeg> >
        base_t;

    typedef Graph graph_t;
    typedef EdgeIndexMap edge_index_t;

    typedef typename graph_traits<Graph>::vertex_descriptor vertex_t;
    typedef typename graph_traits<Graph>::edge_descriptor edge_t;
    typedef typename EdgeIndexMap::value_type index_t;

    struct hash_index {};
    struct random_index {};

    RandomRewireStrategy(Graph& g, EdgeIndexMap edge_index,
                         vector<edge_t>& edges, CorrProb, BlockDeg,
                         bool, rng_t& rng, bool parallel_edges,
                         bool configuration)
        : base_t(g, edge_index, edges, rng, parallel_edges, configuration),
          _g(g) {}

    pair<size_t,bool> get_target_edge(pair<size_t,bool>& e, bool)
    {
        std::uniform_int_distribution<> sample(0, base_t::_edges.size() - 1);
        pair<size_t, bool> et = make_pair(sample(base_t::_rng), false);
        if (!graph_tool::is_directed(_g))
        {
            std::bernoulli_distribution coin(0.5);
            et.second = coin(base_t::_rng);
            e.second = coin(base_t::_rng);
        }
        return et;
    }

    void update_edge(size_t, bool) {}

private:
    Graph& _g;
    EdgeIndexMap _edge_index;
};

// this will rewire the edges so that the (in,out) degree distributions and the
// (in,out)->(in,out) correlations will be the same, but all the rest is random
template <class Graph, class EdgeIndexMap, class CorrProb, class BlockDeg>
class CorrelatedRewireStrategy:
    public RewireStrategyBase<Graph, EdgeIndexMap,
                              CorrelatedRewireStrategy<Graph, EdgeIndexMap,
                                                       CorrProb, BlockDeg> >
{
public:
    typedef RewireStrategyBase<Graph, EdgeIndexMap,
                               CorrelatedRewireStrategy<Graph, EdgeIndexMap,
                                                        CorrProb, BlockDeg> >
        base_t;

    typedef Graph graph_t;

    typedef typename graph_traits<Graph>::vertex_descriptor vertex_t;
    typedef typename graph_traits<Graph>::edge_descriptor edge_t;

    typedef typename BlockDeg::block_t deg_t;

    CorrelatedRewireStrategy(Graph& g, EdgeIndexMap edge_index,
                             vector<edge_t>& edges, CorrProb, BlockDeg blockdeg,
                             bool, rng_t& rng, bool parallel_edges,
                             bool configuration)
        : base_t(g, edge_index, edges, rng, parallel_edges, configuration),
          _blockdeg(blockdeg), _g(g)
    {
        for (size_t ei = 0; ei < base_t::_edges.size(); ++ei)
        {
            // For undirected graphs, there is no difference between source and
            // target, and each edge will appear _twice_ in the list below,
            // once for each different ordering of source and target.
            edge_t& e = base_t::_edges[ei];

            vertex_t t = target(e, _g);
            deg_t tdeg = get_deg(t, _g);;
            _edges_by_target[tdeg].push_back(make_pair(ei, false));

            if (!graph_tool::is_directed(_g))
            {
                t = source(e, _g);
                deg_t tdeg = get_deg(t, _g);
                _edges_by_target[tdeg].push_back(make_pair(ei, true));
            }
        }
    }

    pair<size_t,bool> get_target_edge(pair<size_t, bool>& e, bool)
    {
        if (!graph_tool::is_directed(_g))
        {
            std::bernoulli_distribution coin(0.5);
            e.second = coin(base_t::_rng);
        }

        vertex_t t = target(e, base_t::_edges, _g);
        deg_t tdeg = get_deg(t, _g);
        auto& elist = _edges_by_target[tdeg];
        std::uniform_int_distribution<> sample(0, elist.size() - 1);
        auto ep = elist[sample(base_t::_rng)];
        if (get_deg(target(ep, base_t::_edges, _g), _g) != tdeg)
            ep.second = not ep.second;
        return ep;
    }

    void update_edge(size_t, bool) {}

    deg_t get_deg(vertex_t v, const Graph& g)
    {
        return _blockdeg.get_block(v, g);
    }

private:
    BlockDeg _blockdeg;

    typedef std::unordered_map<deg_t,
                          vector<pair<size_t, bool>>>
        edges_by_end_deg_t;

    edges_by_end_deg_t _edges_by_target;

protected:
    const Graph& _g;
};

// general stochastic blockmodel
// this version is based on rejection sampling
template <class Graph, class EdgeIndexMap, class CorrProb, class BlockDeg>
class ProbabilisticRewireStrategy:
    public RewireStrategyBase<Graph, EdgeIndexMap,
                              ProbabilisticRewireStrategy<Graph, EdgeIndexMap,
                                                          CorrProb, BlockDeg> >
{
public:
    typedef RewireStrategyBase<Graph, EdgeIndexMap,
                               ProbabilisticRewireStrategy<Graph, EdgeIndexMap,
                                                           CorrProb, BlockDeg> >
        base_t;

    typedef Graph graph_t;
    typedef EdgeIndexMap edge_index_t;

    typedef typename BlockDeg::block_t deg_t;

    typedef typename graph_traits<Graph>::vertex_descriptor vertex_t;
    typedef typename graph_traits<Graph>::edge_descriptor edge_t;
    typedef typename EdgeIndexMap::value_type index_t;

    ProbabilisticRewireStrategy(Graph& g, EdgeIndexMap edge_index,
                                vector<edge_t>& edges, CorrProb corr_prob,
                                BlockDeg blockdeg, bool cache, rng_t& rng,
                                bool parallel_edges, bool configuration)
        : base_t(g, edge_index, edges, rng, parallel_edges, configuration),
          _g(g), _corr_prob(corr_prob), _blockdeg(blockdeg)
    {
        if (cache)
        {
            // cache probabilities
            _corr_prob.get_probs(_probs);

            if (_probs.empty())
            {
                std::unordered_set<deg_t> deg_set;
                for (size_t ei = 0; ei < base_t::_edges.size(); ++ei)
                {
                    edge_t& e = base_t::_edges[ei];
                    deg_set.insert(get_deg(source(e, g), g));
                    deg_set.insert(get_deg(target(e, g), g));
                }

                for (auto s_iter = deg_set.begin(); s_iter != deg_set.end(); ++s_iter)
                    for (auto t_iter = deg_set.begin(); t_iter != deg_set.end(); ++t_iter)
                    {
                        double p = _corr_prob(*s_iter, *t_iter);
                        _probs[make_pair(*s_iter, *t_iter)] = p;
                    }
            }

            for (auto iter = _probs.begin(); iter != _probs.end(); ++iter)
            {
                double& p = iter->second;
                // avoid zero probability to not get stuck in rejection step
                if (std::isnan(p) || std::isinf(p) || p <= 0)
                    p = numeric_limits<double>::min();
                p = log(p);
            }
        }
    }

    double get_prob(const deg_t& s_deg, const deg_t& t_deg)
    {
        if (_probs.empty())
        {
            double p = _corr_prob(s_deg, t_deg);
            // avoid zero probability to not get stuck in rejection step
            if (std::isnan(p) || std::isinf(p) || p <= 0)
                p = numeric_limits<double>::min();
            return log(p);
        }
        auto k = make_pair(s_deg, t_deg);
        auto iter = _probs.find(k);
        if (iter == _probs.end())
            return log(numeric_limits<double>::min());
        return iter->second;
    }

    deg_t get_deg(vertex_t v, Graph& g)
    {
        return _blockdeg.get_block(v, g);
    }

    pair<size_t, bool> get_target_edge(pair<size_t, bool>& e, bool)
    {
        if (!graph_tool::is_directed(_g))
        {
            std::bernoulli_distribution coin(0.5);
            e.second = coin(base_t::_rng);
        }

        deg_t s_deg = get_deg(source(e, base_t::_edges, _g), _g);
        deg_t t_deg = get_deg(target(e, base_t::_edges, _g), _g);

        std::uniform_int_distribution<> sample(0, base_t::_edges.size() - 1);
        size_t epi = sample(base_t::_rng);
        pair<size_t, bool> ep = make_pair(epi, false);
        if (!graph_tool::is_directed(_g))
        {
            // for undirected graphs we must select a random direction
            std::bernoulli_distribution coin(0.5);
            ep.second = coin(base_t::_rng);
        }

        if (source(e, base_t::_edges, _g) == source(ep, base_t::_edges, _g) ||
            target(e, base_t::_edges, _g) == target(ep, base_t::_edges, _g))
            return ep; // rewiring is a no-op

        deg_t ep_s_deg = get_deg(source(ep, base_t::_edges, _g), _g);
        deg_t ep_t_deg = get_deg(target(ep, base_t::_edges, _g), _g);

        double pi = get_prob(s_deg, t_deg) + get_prob(ep_s_deg, ep_t_deg);
        double pf = get_prob(s_deg, ep_t_deg) + get_prob(ep_s_deg, t_deg);

        if (pf >= pi)
            return ep;

        double a = exp(pf - pi);

        std::uniform_real_distribution<> rsample(0.0, 1.0);
        double r = rsample(base_t::_rng);
        if (r > a)
            return e; // reject
        else
            return ep;
    }

    void update_edge(size_t, bool) {}

private:
    Graph& _g;
    EdgeIndexMap _edge_index;
    CorrProb _corr_prob;
    BlockDeg _blockdeg;

    typedef std::unordered_map<pair<deg_t, deg_t>, double> prob_map_t;
    prob_map_t _probs;
};

// general "traditional" stochastic blockmodel
// this version is based on the alias method, and does not keep the degrees fixed
template <class Graph, class EdgeIndexMap, class CorrProb, class BlockDeg,
          bool micro>
class TradBlockRewireStrategy
{
public:
    typedef typename graph_traits<Graph>::vertex_descriptor vertex_t;
    typedef typename graph_traits<Graph>::edge_descriptor edge_t;
    typedef typename EdgeIndexMap::value_type index_t;
    typedef typename BlockDeg::block_t deg_t;

    TradBlockRewireStrategy(Graph& g, EdgeIndexMap edge_index,
                            vector<edge_t>& edges, CorrProb corr_prob,
                            BlockDeg blockdeg, bool, rng_t& rng,
                            bool parallel_edges, bool configuration)

        : _g(g), _edge_index(edge_index), _edges(edges), _corr_prob(corr_prob),
          _blockdeg(blockdeg), _rng(rng), _sampler(nullptr),
          _configuration(configuration),
          _nmap(get(vertex_index, g), num_vertices(g))
    {
        for (auto v : vertices_range(_g))
        {
            deg_t d = _blockdeg.get_block(v, g);
            _vertices[d].push_back(v);
        }

        if (!micro)
        {
            std::unordered_map<pair<deg_t, deg_t>, double> probs;
            _corr_prob.get_probs(probs);

            vector<double> dprobs;
            if (probs.empty())
            {
                for (auto& s : _vertices)
                {
                    for (auto& t : _vertices)
                    {
                        double p = _corr_prob(s.first, t.first);
                        if (std::isnan(p) || std::isinf(p) || p <= 0)
                            continue;

                        _items.push_back(make_pair(s.first, t.first));
                        dprobs.push_back(p * s.second.size() * t.second.size());
                    }
                }
            }
            else
            {
                for (auto& stp : probs)
                {
                    deg_t s = stp.first.first;
                    deg_t t = stp.first.second;
                    double p = stp.second;
                    // avoid zero probability to not get stuck in rejection step
                    if (std::isnan(p) || std::isinf(p) || p <= 0)
                        continue;
                    _items.push_back(make_pair(s, t));
                    dprobs.push_back(p * _vertices[s].size() * _vertices[t].size());
                }
            }

            if (_items.empty())
                throw GraphException("No connection probabilities larger than zero!");

            _sampler = new Sampler<pair<deg_t, deg_t> >(_items, dprobs);
        }

        if (!configuration || !parallel_edges)
        {
            for (size_t i = 0; i < edges.size(); ++i)
                add_count(source(edges[i], g), target(edges[i], g), _nmap, g);
        }
    }

    ~TradBlockRewireStrategy()
    {
        if (_sampler != nullptr)
            delete _sampler;
    }

    bool operator()(size_t ei, bool self_loops, bool parallel_edges)
    {
        size_t e_s = source(_edges[ei], _g);
        size_t e_t = target(_edges[ei], _g);

        typename graph_traits<Graph>::vertex_descriptor s, t;

        pair<deg_t, deg_t> deg;
        if (micro)
            deg = {_blockdeg.get_block(e_s, _g),
                   _blockdeg.get_block(e_t, _g)};

        while (true)
        {
            if (!micro)
                deg = _sampler->sample(_rng);

            vector<vertex_t>& svs = _vertices[deg.first];
            vector<vertex_t>& tvs = _vertices[deg.second];

            if (svs.empty() || tvs.empty())
                continue;

            s = uniform_sample(svs, _rng);
            t = uniform_sample(tvs, _rng);

            if (!graph_tool::is_directed(_g) &&
                deg.first == deg.second && s != t && self_loops)
            {
                // sample self-loops w/ correct probability for undirected
                // graphs
                std::bernoulli_distribution reject(.5);
                if (reject(_rng))
                    continue;
            }

            break;
        }

        // reject self-loops if not allowed
        if (!self_loops && s == t)
            return false;

        // reject parallel edges if not allowed
        if (!parallel_edges && get_count(s, t, _nmap, _g))
            return false;

        if (!_configuration)
        {
            size_t m = get_count(s, t, _nmap, _g);
            size_t m_e = get_count(e_s, e_t, _nmap, _g);

            double a = (m + 1) / double(m_e);

            std::bernoulli_distribution accept(std::min(a, 1.));
            if (!accept(_rng))
                return false;
        }

        remove_edge(_edges[ei], _g);
        edge_t ne = add_edge(s, t, _g).first;
        _edges[ei] = ne;

        if (!_configuration || !parallel_edges)
        {
            remove_count(e_s, e_t, _nmap, _g);
            add_count(s, t, _nmap, _g);
        }

        return true;
    }

private:
    Graph& _g;
    EdgeIndexMap _edge_index;
    vector<edge_t>& _edges;
    CorrProb _corr_prob;
    BlockDeg _blockdeg;
    rng_t& _rng;

    std::unordered_map<deg_t, vector<vertex_t>> _vertices;

    vector<pair<deg_t, deg_t> > _items;
    Sampler<pair<deg_t, deg_t> >* _sampler;

    bool _configuration;

    typedef gt_hash_map<size_t, size_t> nmapv_t;
    typedef typename property_map_type::apply<nmapv_t,
                                              typename property_map<Graph, vertex_index_t>::type>
        ::type::unchecked_t nmap_t;
    nmap_t _nmap;
};

template <class Graph, class EdgeIndexMap, class CorrProb, class BlockDeg>
using CanTradBlockRewireStrategy =
    TradBlockRewireStrategy<Graph, EdgeIndexMap, CorrProb, BlockDeg, false>;

template <class Graph, class EdgeIndexMap, class CorrProb, class BlockDeg>
using MicroTradBlockRewireStrategy =
    TradBlockRewireStrategy<Graph, EdgeIndexMap, CorrProb, BlockDeg, true>;

} // graph_tool namespace

#endif // GRAPH_REWIRING_HH


#include <boost/algorithm/minmax_element.hpp>
#include <array>
#include <utility>
#include <iostream>

int main()
{
  typedef std::array<int, 4> array;
  array a{{2, 3, 0, 1}};

  std::pair<array::iterator, array::iterator> p =
    boost::minmax_element(a.begin(), a.end());

  std::cout << *p.first << '\n';
  std::cout << *p.second << '\n';
}


#include <boost/timer/timer.hpp>
#include <iostream>
#include <cmath>

using namespace boost::timer;

int main()
{
  cpu_timer timer;

  for (int i = 0; i < 1000000; ++i)
    std::pow(1.234, i);

  cpu_times times = timer.elapsed();
  std::cout << times.wall << '\n';
  std::cout << times.user << '\n';
  std::cout << times.system << '\n';
}


#ifndef LIBS_H
#define LIBS_H

#include <iostream>
#include <memory>
#include <string>
#include <vector>
// Eigen
#include <Eigen/Dense>

// OpenCV
#include "opencv2/imgproc.hpp"
#include "opencv2/imgcodecs.hpp"
#include "opencv2/surface_matching/icp.hpp"

// litorch
#include <torch/torch.h>
#include <torch/script.h>
#include <ATen/Aten.h>
#include <c10/cuda/CUDACachingAllocator.h>

typedef Eigen::Matrix<float, Eigen::Dynamic, Eigen::Dynamic, Eigen::RowMajor> MatrixXf_RM;

using namespace std;

using std::unique_ptr;
using std::make_unique;
using std::shared_ptr;
using std::make_shared;

#endif LIBS_H


//=======================================================================
// Copyright 2001 Jeremy G. Siek, Andrew Lumsdaine, Lie-Quan Lee,
//
// Distributed under the Boost Software License, Version 1.0. (See
// accompanying file LICENSE_1_0.txt or copy at
// http://www.boost.org/LICENSE_1_0.txt)
//=======================================================================
#include <boost/config.hpp>
#include <iostream>
#include <fstream>
#include <boost/graph/adjacency_list.hpp>
#include <boost/graph/depth_first_search.hpp>
#include <boost/graph/graphviz.hpp>
#include <boost/graph/copy.hpp>
#include <boost/graph/reverse_graph.hpp>

int
main(int argc, char *argv[])
{
  if (argc < 3) {
    std::cerr << "usage: reachable-loop-tail.exe <in-file> <out-file>"
      << std::endl;
    return -1;
  }
  using namespace boost;
  GraphvizDigraph g_in;
  read_graphviz(argv[1], g_in);

  typedef adjacency_list < vecS, vecS, bidirectionalS,
    GraphvizVertexProperty,
    GraphvizEdgeProperty, GraphvizGraphProperty > Graph;
  Graph g;
  copy_graph(g_in, g);

  graph_traits < GraphvizDigraph >::vertex_descriptor loop_tail = 6;
  typedef color_traits < default_color_type > Color;
  default_color_type c;

  std::vector < default_color_type > reachable_to_tail(num_vertices(g));
  reverse_graph < Graph > reverse_g(g);
  depth_first_visit(reverse_g, loop_tail, default_dfs_visitor(),
                    make_iterator_property_map(reachable_to_tail.begin(),
                                               get(vertex_index, g), c));

  std::ofstream loops_out(argv[2]);
  loops_out << "digraph G {\n"
    << "  graph [ratio=\"fill\",size=\"3,3\"];\n"
    << "  node [shape=\"box\"];\n" << "  edge [style=\"bold\"];\n";

  property_map<Graph, vertex_attribute_t>::type
    vattr_map = get(vertex_attribute, g);
  graph_traits < GraphvizDigraph >::vertex_iterator i, i_end;
  for (tie(i, i_end) = vertices(g_in); i != i_end; ++i) {
    loops_out << *i << "[label=\"" << vattr_map[*i]["label"]
      << "\"";
    if (reachable_to_tail[*i] != Color::white()) {
      loops_out << ", color=\"gray\", style=\"filled\"";
    }
    loops_out << "]\n";
  }
  graph_traits < GraphvizDigraph >::edge_iterator e, e_end;
  for (tie(e, e_end) = edges(g_in); e != e_end; ++e)
    loops_out << source(*e, g) << " -> " << target(*e, g) << ";\n";
  loops_out << "}\n";
  return EXIT_SUCCESS;
}


/*=========================================================================

  Program:   Visualization Toolkit
  Module:    vtkPBGLConnectedComponents.cxx

  Copyright (c) Ken Martin, Will Schroeder, Bill Lorensen
  All rights reserved.
  See Copyright.txt or http://www.kitware.com/Copyright.htm for details.

     This software is distributed WITHOUT ANY WARRANTY; without even
     the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR
     PURPOSE.  See the above copyright notice for more information.

=========================================================================*/
/*-------------------------------------------------------------------------
  Copyright 2008 Sandia Corporation.
  Under the terms of Contract DE-AC04-94AL85000 with Sandia Corporation,
  the U.S. Government retains certain rights in this software.
-------------------------------------------------------------------------*/
/*
 * Copyright (C) 2008 The Trustees of Indiana University.
 * Use, modification and distribution is subject to the Boost Software
 * License, Version 1.0. (See http://www.boost.org/LICENSE_1_0.txt)
 */
#include <boost/graph/use_mpi.hpp> // must precede PBGL headers

#include "vtkPBGLConnectedComponents.h"

#if !defined(VTK_LEGACY_REMOVE)

#include "vtkCellArray.h"
#include "vtkCellData.h"
#include "vtkDataArray.h"
#include "vtkDirectedGraph.h"
#include "vtkDistributedGraphHelper.h"
#include "vtkFloatArray.h"
#include "vtkInformation.h"
#include "vtkInformationVector.h"
#include "vtkMath.h"
#include "vtkObjectFactory.h"
#include "vtkPBGLGraphAdapter.h"
#include "vtkPBGLDistributedGraphHelper.h"
#include "vtkPointData.h"
#include "vtkSelection.h"
#include "vtkSmartPointer.h"
#include "vtkStreamingDemandDrivenPipeline.h"
#include "vtkStringArray.h"
#include "vtkType.h"
#include "vtkUndirectedGraph.h"
#include "vtkVertexListIterator.h"

#include <boost/graph/connected_components.hpp>
#include <boost/graph/distributed/connected_components.hpp>
#include <boost/graph/strong_components.hpp>
#include <boost/graph/distributed/strong_components.hpp>

#include <utility> // for pair

using namespace boost;

vtkStandardNewMacro(vtkPBGLConnectedComponents);

// Constructor/Destructor
vtkPBGLConnectedComponents::vtkPBGLConnectedComponents()
{
  this->ComponentArrayName = 0;
  this->SetNumberOfInputPorts(1);
  this->SetNumberOfOutputPorts(1);
  VTK_LEGACY_BODY(vtkPBGLConnectedComponents::vtkPBGLConnectedComponents, "VTK 6.2");
}

vtkPBGLConnectedComponents::~vtkPBGLConnectedComponents()
{
  this->SetComponentArrayName(0);
}

int vtkPBGLConnectedComponents::RequestData(
  vtkInformation *vtkNotUsed(request),
  vtkInformationVector **inputVector,
  vtkInformationVector *outputVector)
{
  // get the info objects
  vtkInformation *inInfo = inputVector[0]->GetInformationObject(0);
  vtkInformation *outInfo = outputVector->GetInformationObject(0);

  // get the input and output
  vtkGraph *input = vtkGraph::SafeDownCast(
    inInfo->Get(vtkDataObject::DATA_OBJECT()));
  vtkGraph *output = vtkGraph::SafeDownCast(
    outInfo->Get(vtkDataObject::DATA_OBJECT()));

  // Send the data to output.
  output->ShallowCopy(input);

  // Create the component array
  vtkIdTypeArray* componentArray = vtkIdTypeArray::New();
  if (this->ComponentArrayName)
  {
    componentArray->SetName(this->ComponentArrayName);
  }
  else
  {
    componentArray->SetName("Component");
  }
  componentArray->SetNumberOfTuples(output->GetNumberOfVertices());

  vtkDistributedGraphHelper *helper = output->GetDistributedGraphHelper();
  if (!helper)
  {
    vtkErrorMacro("Distributed vtkGraph is required.");
    return 1;
  }

  // We can only deal with Parallel BGL-distributed graphs.
  vtkPBGLDistributedGraphHelper *pbglHelper
    = vtkPBGLDistributedGraphHelper::SafeDownCast(helper);
  if (!pbglHelper)
  {
    vtkErrorMacro("Can only perform parallel shortest-paths on a Parallel BGL distributed graph");
    return 1;
  }

  int myRank = outInfo->Get(vtkStreamingDemandDrivenPipeline::UPDATE_PIECE_NUMBER());

  // Distributed component map
  typedef vtkDistributedVertexPropertyMapType<vtkIdTypeArray>::type
    DistributedComponentMap;
  DistributedComponentMap componentMap
    = MakeDistributedVertexPropertyMap(output, componentArray);

  // Execute the algorithm itself
  if (vtkUndirectedGraph::SafeDownCast(output))
  {
    // Distributed parent map. Used for scratch space in the algorithm
    // itself.
    vtkIdTypeArray* parentArray = vtkIdTypeArray::New();
    parentArray->SetNumberOfTuples(output->GetNumberOfVertices());
    typedef vtkDistributedVertexPropertyMapType<vtkIdTypeArray>::type
      DistributedParentMap;
    DistributedParentMap parentMap
      = MakeDistributedVertexPropertyMap(output, parentArray);

    vtkUndirectedGraph *g = vtkUndirectedGraph::SafeDownCast(output);
    boost::graph::distributed::cc_detail::parallel_connected_components
      (g, parentMap);
    boost::graph::distributed::number_components_from_parents(g, parentMap,
                                                              componentMap);

    // Remove the temporary parent array.
    parentArray->Delete();
  }
  else
  {
    vtkDirectedGraph *g = vtkDirectedGraph::SafeDownCast(output);
    boost::graph::distributed::fleischer_hendrickson_pinar_strong_components
      (g, componentMap, MakeDistributedVertexIndexMap(g));
  }

  // Add component array to the output
  output->GetVertexData()->AddArray(componentArray);
  componentArray->Delete();

  return 1;
}

void vtkPBGLConnectedComponents::PrintSelf(ostream& os, vtkIndent indent)
{
  this->Superclass::PrintSelf(os, indent);

  os << indent << "ComponentArrayName: "
     << (this->ComponentArrayName ? this->ComponentArrayName : "(none)")
     << endl;
}

//----------------------------------------------------------------------------
int vtkPBGLConnectedComponents::FillInputPortInformation(
  int port, vtkInformation* info)
{
  // now add our info
  if (port == 0)
  {
    info->Set(vtkAlgorithm::INPUT_REQUIRED_DATA_TYPE(), "vtkGraph");
  }
  return 1;
}

//----------------------------------------------------------------------------
int vtkPBGLConnectedComponents::FillOutputPortInformation(
  int port, vtkInformation* info)
{
  // now add our info
  if (port == 0)
  {
    info->Set(vtkDataObject::DATA_TYPE_NAME(), "vtkGraph");
  }
  return 1;
}

#endif //VTK_LEGACY_REMOVE