/* +------------------------------------------------------------------------+
   |                     Mobile Robot Programming Toolkit (MRPT)            |
   |                          https://www.mrpt.org/                         |
   |                                                                        |
   | Copyright (c) 2005-2024, Individual contributors, see AUTHORS file     |
   | See: https://www.mrpt.org/Authors - All rights reserved.               |
   | Released under BSD License. See: https://www.mrpt.org/License          |
   +------------------------------------------------------------------------+ */

#include "graphs-precomp.h"  // Precompiled headers
//
#include <mrpt/graphs/CGraphPartitioner.h>
#include <mrpt/math/CMatrixD.h>
#include <mrpt/math/CMatrixF.h>
#include <mrpt/math/ops_matrices.h>  // laplacian()

#include <iostream>  // cout

using namespace mrpt;
using namespace mrpt::graphs;
using namespace std;

/*---------------------------------------------------------------
SpectralPartition
---------------------------------------------------------------*/
template <class GRAPH_MATRIX, typename num_t>
void CGraphPartitioner<GRAPH_MATRIX, num_t>::SpectralBisection(
    GRAPH_MATRIX& in_A,
    std::vector<uint32_t>& out_part1,
    std::vector<uint32_t>& out_part2,
    num_t& out_cut_value,
    bool forceSimetry)
{
  size_t nodeCount;  // Nodes count
  GRAPH_MATRIX Adj, eigenVectors;
  std::vector<typename GRAPH_MATRIX::Scalar> eigenValues;

  // Check matrix is square:
  if (in_A.cols() != int(nodeCount = in_A.rows()))
    THROW_EXCEPTION("Weights matrix is not square!!");

  // Shi & Malik's method
  //-----------------------------------------------------------------

  // forceSimetry?
  if (forceSimetry)
  {
    Adj.setSize(nodeCount, nodeCount);
    for (size_t i = 0; i < nodeCount; i++)
      for (size_t j = i; j < nodeCount; j++)
        Adj(i, j) = Adj(j, i) = 0.5f * (in_A(i, j) + in_A(j, i));
  }
  else
    Adj = in_A;

  // Compute eigen-vectors of laplacian:
  GRAPH_MATRIX LAPLACIAN;
  mrpt::math::laplacian(Adj, LAPLACIAN);

  LAPLACIAN.eig(eigenVectors, eigenValues);

  //  Execute the bisection
  // ------------------------------------
  // Second smallest eigen-vector
  double mean = 0;
  size_t colNo = 1;  // second smallest
  size_t nRows = eigenVectors.rows();

  // for (i=0;i<eigenVectors.cols();i++) mean+=eigenVectors(colNo,i);
  for (size_t i = 0; i < nRows; i++) mean += eigenVectors(i, colNo);
  mean /= nRows;

  out_part1.clear();
  out_part2.clear();

  for (size_t i = 0; i < nRows; i++)
  {
    if (eigenVectors(i, colNo) >= mean)
      out_part1.push_back(i);
    else
      out_part2.push_back(i);
  }

  // Special and strange case: Constant eigenvector: Split nodes in two
  //    equally sized parts arbitrarily:
  // ------------------------------------------------------------------
  if (!out_part1.size() || !out_part2.size())
  {
    out_part1.clear();
    out_part2.clear();
    // Assign 50%-50%:
    for (int i = 0; i < Adj.cols(); i++)
      if (i <= Adj.cols() / 2)
        out_part1.push_back(i);
      else
        out_part2.push_back(i);
  }

  // Compute the N-cut value
  out_cut_value = nCut(Adj, out_part1, out_part2);
}

/*---------------------------------------------------------------
RecursiveSpectralPartition
---------------------------------------------------------------*/
template <class GRAPH_MATRIX, typename num_t>
void CGraphPartitioner<GRAPH_MATRIX, num_t>::RecursiveSpectralPartition(
    GRAPH_MATRIX& in_A,
    std::vector<std::vector<uint32_t>>& out_parts,
    num_t threshold_Ncut,
    bool forceSimetry,
    bool useSpectralBisection,
    bool recursive,
    unsigned minSizeClusters,
    const bool verbose)
{
  MRPT_START

  size_t nodeCount;
  std::vector<uint32_t> p1, p2;
  num_t cut_value;
  size_t i, j;
  GRAPH_MATRIX Adj;

  out_parts.clear();

  // Check matrix is square:
  if (in_A.cols() != int(nodeCount = in_A.rows()))
    THROW_EXCEPTION("Weights matrix is not square!!");

  if (nodeCount == 1)
  {
    // Don't split, there is just a node!
    p1.push_back(0);
    out_parts.push_back(p1);
    return;
  }

  // forceSimetry?
  if (forceSimetry)
  {
    Adj.setSize(nodeCount, nodeCount);
    for (i = 0; i < nodeCount; i++)
      for (j = i; j < nodeCount; j++) Adj(i, j) = Adj(j, i) = 0.5f * (in_A(i, j) + in_A(j, i));
  }
  else
    Adj = in_A;

  // Make bisection
  if (useSpectralBisection)
    SpectralBisection(Adj, p1, p2, cut_value, false);
  else
    exactBisection(Adj, p1, p2, cut_value, false);

  if (verbose)
    std::cout << format(
        "Cut:%u=%u+%u,nCut=%.02f->", (unsigned int)nodeCount, (unsigned int)p1.size(),
        (unsigned int)p2.size(), cut_value);

  // Is it a useful partition?
  if (cut_value > threshold_Ncut || p1.size() < minSizeClusters || p2.size() < minSizeClusters)
  {
    if (verbose) std::cout << "->NO!" << std::endl;

    // No:
    p1.clear();
    for (i = 0; i < nodeCount; i++) p1.push_back(i);
    out_parts.push_back(p1);
  }
  else
  {
    if (verbose) std::cout << "->YES!" << std::endl;

    // Yes:
    std::vector<std::vector<uint32_t>> p1_parts, p2_parts;

    if (recursive)
    {
      // Split "p1":
      // --------------------------------------------
      // sub-matrix:
      GRAPH_MATRIX A_1(p1.size(), p1.size());
      for (i = 0; i < p1.size(); i++)
        for (j = 0; j < p1.size(); j++) A_1(i, j) = in_A(p1[i], p1[j]);

      RecursiveSpectralPartition(
          A_1, p1_parts, threshold_Ncut, forceSimetry, useSpectralBisection, recursive,
          minSizeClusters);

      // Split "p2":
      // --------------------------------------------
      // sub-matrix:
      GRAPH_MATRIX A_2(p2.size(), p2.size());
      for (i = 0; i < p2.size(); i++)
        for (j = 0; j < p2.size(); j++) A_2(i, j) = in_A(p2[i], p2[j]);

      RecursiveSpectralPartition(
          A_2, p2_parts, threshold_Ncut, forceSimetry, useSpectralBisection, recursive,
          minSizeClusters);

      // Build "out_parts" from "p1_parts" + "p2_parts"
      //  taken care of indexes mapping!
      // -------------------------------------------------------------------------------------
      // remap p1 nodes:
      for (i = 0; i < p1_parts.size(); i++)
      {
        for (j = 0; j < p1_parts[i].size(); j++) p1_parts[i][j] = p1[p1_parts[i][j]];
        out_parts.push_back(p1_parts[i]);
      }

      // remap p2 nodes:
      for (i = 0; i < p2_parts.size(); i++)
      {
        for (j = 0; j < p2_parts[i].size(); j++) p2_parts[i][j] = p2[p2_parts[i][j]];

        out_parts.push_back(p2_parts[i]);
      }
    }
    else
    {
      // Force bisection only:
      out_parts.clear();
      out_parts.push_back(p1);
      out_parts.push_back(p2);
    }
  }

  MRPT_END
}

/*---------------------------------------------------------------
nCut
---------------------------------------------------------------*/
template <class GRAPH_MATRIX, typename num_t>
num_t CGraphPartitioner<GRAPH_MATRIX, num_t>::nCut(
    const GRAPH_MATRIX& in_A,
    const std::vector<uint32_t>& in_part1,
    const std::vector<uint32_t>& in_part2)
{
  size_t i, j;
  const size_t size1 = in_part1.size();
  const size_t size2 = in_part2.size();

  // Compute the N-cut value
  // -----------------------------------------------
  num_t cut_AB = 0;
  for (i = 0; i < size1; i++)
    for (j = 0; j < size2; j++) cut_AB += in_A(in_part1[i], in_part2[j]);

  num_t assoc_AA = 0;
  for (i = 0; i < size1; i++)
    for (j = i; j < size1; j++)
      if (i != j) assoc_AA += in_A(in_part1[i], in_part1[j]);

  num_t assoc_BB = 0;
  for (i = 0; i < size2; i++)
    for (j = i; j < size2; j++)
      if (i != j) assoc_BB += in_A(in_part2[i], in_part2[j]);

  num_t assoc_AV = assoc_AA + cut_AB;
  num_t assoc_BV = assoc_BB + cut_AB;

  if (!cut_AB)
    return 0;
  else
    return cut_AB / assoc_AV + cut_AB / assoc_BV;
}

/*---------------------------------------------------------------
exactBisection
---------------------------------------------------------------*/
template <class GRAPH_MATRIX, typename num_t>
void CGraphPartitioner<GRAPH_MATRIX, num_t>::exactBisection(
    GRAPH_MATRIX& in_A,
    std::vector<uint32_t>& out_part1,
    std::vector<uint32_t>& out_part2,
    num_t& out_cut_value,
    bool forceSimetry)
{
  size_t nodeCount;  // Nodes count
  size_t i, j;
  GRAPH_MATRIX Adj;
  std::vector<bool> partition, bestPartition;
  std::vector<uint32_t> part1, part2;
  num_t partCutValue, bestCutValue = std::numeric_limits<num_t>::max();
  bool end = false;

  // Check matrix is square:
  if (in_A.cols() != int(nodeCount = in_A.rows()))
    THROW_EXCEPTION("Weights matrix is not square!!");

  ASSERT_(nodeCount >= 2);

  // forceSimetry?
  if (forceSimetry)
  {
    Adj.setSize(nodeCount, nodeCount);
    for (i = 0; i < nodeCount; i++)
      for (j = i; j < nodeCount; j++) Adj(i, j) = Adj(j, i) = 0.5f * (in_A(i, j) + in_A(j, i));
  }
  else
    Adj = in_A;

  // Brute force: compute all possible partitions:
  //-----------------------------------------------------------------

  // First combination: 1000...0000
  // Last combination:  1111...1110
  partition.clear();
  partition.resize(nodeCount, false);
  partition[0] = true;

  while (!end)
  {
    // Build partitions from binary vector:
    part1.clear();
    part2.clear();

    for (i = 0; i < nodeCount; i++)
    {
      if (partition[i])
        part2.push_back(i);
      else
        part1.push_back(i);
    }

    // Compute the n-cut:
    partCutValue = nCut(Adj, part1, part2);

    if (partCutValue < bestCutValue)
    {
      bestCutValue = partCutValue;
      bestPartition = partition;
    }

    // Next combo in the binary vector:
    i = 0;
    bool carry = false;
    do
    {
      carry = partition[i];
      partition[i] = !partition[i];
      i++;
    } while (carry && i < nodeCount);

    // End criterion:
    end = true;
    for (i = 0; end && i < nodeCount; i++) end = end && partition[i];

  };  // End of while

  // Return the best partition:
  out_cut_value = bestCutValue;

  out_part1.clear();
  out_part2.clear();

  for (i = 0; i < nodeCount; i++)
  {
    if (bestPartition[i])
      out_part2.push_back(i);
    else
      out_part1.push_back(i);
  }
}

// Explicit instantiations:
namespace mrpt::graphs
{
// Explicit instantiations:
template class CGraphPartitioner<mrpt::math::CMatrixDouble, double>;
template class CGraphPartitioner<mrpt::math::CMatrixFloat, float>;
template class CGraphPartitioner<mrpt::math::CMatrixD, double>;
template class CGraphPartitioner<mrpt::math::CMatrixF, float>;
}  // namespace mrpt::graphs
