#include "Rcpp.h"
#include <deque>
#include <algorithm>

/* Adapted from bluster package
 * https://bioconductor.org/packages/release/bioc/html/bluster.html
 */

/* The 'rank' version performs the original SNN clustering described by Xu and Su (2015, Bioinformatics).
 * This defines the weight between two nodes as (k - 0.5 * r), where r is the smallest sum of ranks for any node in both NN-sets.
 * The rank is computed separately in each NN set, with each node being 0-rank in its own set.
 */

// [[Rcpp::export(rng=false)]]
Rcpp::List build_snn_rank(Rcpp::IntegerMatrix neighbors) {
  const size_t k=neighbors.ncol();
  const size_t ncells=neighbors.nrow();

  // Building a host table, identifying the reverse relation from nearest neighbours to cells.
  auto mIt=neighbors.begin();
  std::deque<std::deque<std::pair<size_t, int> > > hosts(ncells);
  for (size_t i=1; i<=k; ++i) {
    for (size_t j=0; j<ncells; ++j, ++mIt) {
      hosts[*mIt - 1].push_back(std::make_pair(i, j)); // Getting to 0-based index, keeping 1-based ranks for now.
    }
  }

  std::deque<int> output_pairs;
  std::deque<double> output_weights;
  std::deque<size_t> current_added;
  std::deque<size_t> current_score(ncells);

  for (size_t j=0; j<ncells; ++j) {
    auto rowtmp=neighbors.row(j);
    auto rtIt=rowtmp.begin();

    int cur_neighbor;
    for (size_t i=0; i<=k; ++i) {
      if (i==0) {
        cur_neighbor=j;
      } else {
        // Adding the actual nearest neighbors for cell 'j'.
        cur_neighbor=*rtIt - 1;
        ++rtIt;

        if (static_cast<size_t>(cur_neighbor) < j) { // avoid duplicates from symmetry in the SNN calculations.
          const size_t& currank=i; // +0, as neighbour 'i' is rank 0 with respect to itself.
          size_t& existing_other=current_score[cur_neighbor];
          if (existing_other==0) {
            existing_other=currank;
            current_added.push_back(cur_neighbor);
          } else if (existing_other > currank) {
            existing_other=currank;
          }
        }
      }

      // Adding the cells connected by shared nearest neighbors, again recording the lowest combined rank per neighbor.
      const auto& hosted=hosts[cur_neighbor];
      for (auto hIt=hosted.begin(); hIt!=hosted.end(); ++hIt) {
        const int& othernode=hIt->second;

        if (static_cast<size_t>(othernode) < j) { // avoid duplicates from symmetry in the SNN calculations.
          size_t currank=hIt->first + i;
          size_t& existing_other=current_score[othernode];
          if (existing_other==0) {
            existing_other=currank;
            current_added.push_back(othernode);
          } else if (existing_other > currank) {
            existing_other=currank;
          }
        }
      }
    }

    for (auto othernode : current_added) {
      // Converting to edges.
      output_pairs.push_back(j + 1);
      output_pairs.push_back(othernode + 1);

      // Ensuring that an edge with a positive weight is always reported.
      size_t& otherscore=current_score[othernode];
      double finalscore = static_cast<double>(k) - 0.5 * static_cast<double>(otherscore);
      output_weights.push_back(std::max(finalscore, 1e-6));

      // Resetting all those added to zero.
      otherscore=0;
    }
    current_added.clear();
  }

  Rcpp::IntegerVector pout(output_pairs.begin(), output_pairs.end());
  Rcpp::NumericVector wout(output_weights.begin(), output_weights.end());
  return Rcpp::List::create(pout, wout);
}

/* The 'number' version performs a much simpler SNN clustering.
 * Here, the weight between two nodes is simply the number of shared nodes in both NN-sets.
 * Each node is also included in its own set, yielding a range of [0, k+1] weights.
 */

// [[Rcpp::export(rng=false)]]
Rcpp::List build_snn_number(Rcpp::IntegerMatrix neighbors) {
  const size_t k=neighbors.ncol();
  const size_t ncells=neighbors.nrow();

  // Building a host table, identifying the reverse relation from nearest neighbours to cells.
  auto mIt=neighbors.begin();
  std::deque<std::deque<size_t> > hosts(ncells);
  for (size_t i=1; i<=k; ++i) {
    for (size_t j=0; j<ncells; ++j, ++mIt) {
      hosts[*mIt - 1].push_back(j); // Getting to 0-based index.
    }
  }

  std::deque<int> output_pairs;
  std::deque<double> output_weights;
  std::deque<size_t> current_added;
  std::deque<size_t> current_score(ncells);

  for (size_t j=0; j<ncells; ++j) {
    auto rowtmp=neighbors.row(j);
    auto rtIt=rowtmp.begin();

    int cur_neighbor;
    for (size_t i=0; i<=k; ++i) {
      if (i==0) {
        cur_neighbor=j;
      } else {
        // Adding the actual nearest neighbors for cell 'j'.
        cur_neighbor=*rtIt - 1;
        ++rtIt;

        if (static_cast<size_t>(cur_neighbor) < j) { // avoid duplicates from symmetry in the SNN calculations.
          size_t& existing_other=current_score[cur_neighbor];
          if (existing_other==0) {
            current_added.push_back(cur_neighbor);
          }
          ++existing_other;
        }
      }

      // Adding the cells connected by shared nearest neighbors, recording the number.
      const auto& hosted=hosts[cur_neighbor];
      for (auto hIt=hosted.begin(); hIt!=hosted.end(); ++hIt) {
        const int& othernode=*hIt;

        if (static_cast<size_t>(othernode) < j) { // avoid duplicates from symmetry in the SNN calculations.
          size_t& existing_other=current_score[othernode];
          if (existing_other==0) {
            current_added.push_back(othernode);
          }
          ++existing_other;
        }
      }
    }

    for (auto othernode : current_added) {
      // Converting to edges.
      output_pairs.push_back(j + 1);
      output_pairs.push_back(othernode + 1);

      // Ensuring that an edge with a positive weight is always reported.
      size_t& otherscore=current_score[othernode];
      output_weights.push_back(std::max(static_cast<double>(otherscore), 1e-6));

      // Resetting all those added to zero.
      otherscore=0;
    }
    current_added.clear();
  }

  Rcpp::IntegerVector pout(output_pairs.begin(), output_pairs.end());
  Rcpp::NumericVector wout(output_weights.begin(), output_weights.end());
  return Rcpp::List::create(pout, wout);
}