// Copyright 2005 The Trustees of Indiana University.

// Use, modification and distribution is subject to 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)

//  Authors: Douglas Gregor
//           Andrew Lumsdaine

// An implementation of Walter Hohberg's distributed biconnected
// components algorithm, from:
//
//   Walter Hohberg. How to Find Biconnected Components in Distributed
//   Networks. J. Parallel Distrib. Comput., 9(4):374-386, 1990.
//
#ifndef BOOST_GRAPH_DISTRIBUTED_HOHBERG_BICONNECTED_COMPONENTS_HPP
#define BOOST_GRAPH_DISTRIBUTED_HOHBERG_BICONNECTED_COMPONENTS_HPP

#ifndef BOOST_GRAPH_USE_MPI
#error "Parallel BGL files should not be included unless <boost/graph/use_mpi.hpp> has been included"
#endif

/* You can define PBGL_HOHBERG_DEBUG to an integer value (1, 2, or 3)
 * to enable debugging information. 1 includes only the phases of the
 * algorithm and messages as their are received. 2 and 3 add
 * additional levels of detail about internal data structures related
 * to the algorithm itself.
 *
 * #define PBGL_HOHBERG_DEBUG 1
*/

#include <boost/graph/graph_traits.hpp>
#include <boost/graph/parallel/container_traits.hpp>
#include <boost/graph/parallel/process_group.hpp>
#include <boost/static_assert.hpp>
#include <boost/mpi/operations.hpp>
#include <boost/type_traits/is_convertible.hpp>
#include <boost/graph/graph_concepts.hpp>
#include <boost/graph/iteration_macros.hpp>
#include <boost/optional.hpp>
#include <utility> // for std::pair
#include <cassert>
#include <algorithm> // for std::find, std::mismatch
#include <vector>
#include <boost/graph/parallel/algorithm.hpp>
#include <boost/graph/distributed/connected_components.hpp>

namespace boost { namespace graph { namespace distributed {

namespace hohberg_detail {
  enum message_kind {
    /* A header for the PATH message, stating which edge the message
       is coming on and how many vertices will be following. The data
       structure communicated will be a path_header. */
    msg_path_header,
    /* A message containing the vertices that make up a path. It will
       always follow a msg_path_header and will contain vertex
       descriptors, only. */
    msg_path_vertices,
    /* A header for the TREE message, stating the value of gamma and
       the number of vertices to come in the following
       msg_tree_vertices. */
    msg_tree_header,
    /* A message containing the vertices that make up the set of
       vertices in the same bicomponent as the sender. It will always
       follow a msg_tree_header and will contain vertex descriptors,
       only. */
    msg_tree_vertices,
    /* Provides a name for the biconnected component of the edge. */
    msg_name
  };

  // Payload for a msg_path_header message.
  template<typename EdgeDescriptor>
  struct path_header
  {
    // The edge over which the path is being communicated
    EdgeDescriptor edge;

    // The length of the path, i.e., the number of vertex descriptors
    // that will be coming in the next msg_path_vertices message.
    std::size_t    path_length;

    template<typename Archiver>
    void serialize(Archiver& ar, const unsigned int /*version*/)
    {
      ar & edge & path_length;
    }
  };

  // Payload for a msg_tree_header message.
  template<typename Vertex, typename Edge>
  struct tree_header
  {
    // The edge over which the tree is being communicated
    Edge edge;

    // Gamma, which is the eta of the sender.
    Vertex gamma;

    // The length of the list of vertices in the bicomponent, i.e.,
    // the number of vertex descriptors that will be coming in the
    // next msg_tree_vertices message.
    std::size_t    bicomp_length;

    template<typename Archiver>
    void serialize(Archiver& ar, const unsigned int /*version*/)
    {
      ar & edge & gamma & bicomp_length;
    }
  };

  // Payload for the msg_name message.
  template<typename EdgeDescriptor>
  struct name_header
  {
    // The edge being communicated and named.
    EdgeDescriptor edge;

    // The 0-based name of the component
    std::size_t name;

    template<typename Archiver>
    void serialize(Archiver& ar, const unsigned int /*version*/)
    {
      ar & edge & name;
    }
  };

  /* Computes the branch point between two paths. The branch point is
     the last position at which both paths are equivalent, beyond
     which the paths diverge. Both paths must have length > 0 and the
     initial elements of the paths must be equal. This is guaranteed
     in Hohberg's algorithm because all paths start at the
     leader. Returns the value at the branch point. */
  template<typename T>
  T branch_point(const std::vector<T>& p1, const std::vector<T>& p2)
  {
    assert(!p1.empty());
    assert(!p2.empty());
    assert(p1.front() == p2.front());

    typedef typename std::vector<T>::const_iterator iterator;

    iterator mismatch_pos;
    if (p1.size() <= p2.size())
      mismatch_pos = std::mismatch(p1.begin(), p1.end(), p2.begin()).first;
    else
      mismatch_pos = std::mismatch(p2.begin(), p2.end(), p1.begin()).first;
    --mismatch_pos;
    return *mismatch_pos;
  }

  /* Computes the infimum of vertices a and b in the given path. The
     infimum is the largest element that is on the paths from a to the
     root and from b to the root. */
  template<typename T>
  T infimum(const std::vector<T>& parent_path, T a, T b)
  {
    using std::swap;

    typedef typename std::vector<T>::const_iterator iterator;
    iterator first = parent_path.begin(), last = parent_path.end();

#if defined(PBGL_HOHBERG_DEBUG) and PBGL_HOHBERG_DEBUG > 2
    std::cerr << "infimum(";
    for (iterator i = first; i != last; ++i) {
      if (i != first) std::cerr << ' ';
      std::cerr << local(*i) << '@' << owner(*i);
    }
    std::cerr << ", " << local(a) << '@' << owner(a) << ", "
              << local(b) << '@' << owner(b) << ") = ";
#endif

    if (a == b) {
#if defined(PBGL_HOHBERG_DEBUG) && PBGL_HOHBERG_DEBUG > 2
      std::cerr << local(a) << '@' << owner(a) << std::endl;
#endif
      return a;
    }

    // Try to find a or b, whichever is closest to the end
    --last;
    while (*last != a) {
      // If we match b, swap the two so we'll be looking for b later.
      if (*last == b) { swap(a,b); break; }

      if (last == first) {
#if defined(PBGL_HOHBERG_DEBUG) && PBGL_HOHBERG_DEBUG > 2
        std::cerr << local(*first) << '@' << owner(*first) << std::endl;
#endif
        return *first;
      }
      else --last;
    }

    // Try to find b (which may originally have been a)
    while (*last != b) {
      if (last == first) {
#if defined(PBGL_HOHBERG_DEBUG) and PBGL_HOHBERG_DEBUG > 2
        std::cerr << local(*first) << '@' << owner(*first) << std::endl;
#endif
        return *first;
      }
      else --last;
    }

#if defined(PBGL_HOHBERG_DEBUG) && PBGL_HOHBERG_DEBUG > 2
    std::cerr << local(*last) << '@' << owner(*last) << std::endl;
#endif
    // We've found b; it's the infimum.
    return *last;
  }
} // end namespace hohberg_detail

/* A message that will be stored for each edge by Hohberg's algorithm. */
template<typename Graph>
struct hohberg_message
{
  typedef typename graph_traits<Graph>::vertex_descriptor Vertex;
  typedef typename graph_traits<Graph>::edge_descriptor   Edge;

  // Assign from a path message
  void assign(const std::vector<Vertex>& path)
  {
    gamma = graph_traits<Graph>::null_vertex();
    path_or_bicomp = path;
  }

  // Assign from a tree message
  void assign(Vertex gamma, const std::vector<Vertex>& in_same_bicomponent)
  {
    this->gamma = gamma;
    path_or_bicomp = in_same_bicomponent;
  }

  bool is_path() const { return gamma == graph_traits<Graph>::null_vertex(); }
  bool is_tree() const { return gamma != graph_traits<Graph>::null_vertex(); }

  /// The "gamma" of a tree message, or null_vertex() for a path message
  Vertex gamma;

  // Either the path for a path message or the in_same_bicomponent
  std::vector<Vertex> path_or_bicomp;
};

/* An abstraction of a vertex processor in Hohberg's algorithm. The
   hohberg_vertex_processor class is responsible for processing
   messages transmitted to it via its edges. */
template<typename Graph>
class hohberg_vertex_processor
{
  typedef typename graph_traits<Graph>::vertex_descriptor Vertex;
  typedef typename graph_traits<Graph>::edge_descriptor   Edge;
  typedef typename graph_traits<Graph>::degree_size_type  degree_size_type;
  typedef typename graph_traits<Graph>::edges_size_type   edges_size_type;
  typedef typename boost::graph::parallel::process_group_type<Graph>::type
    ProcessGroup;
  typedef std::vector<Vertex> path_t;
  typedef typename path_t::iterator path_iterator;

 public:
  hohberg_vertex_processor()
    : phase(1),
      parent(graph_traits<Graph>::null_vertex()),
      eta(graph_traits<Graph>::null_vertex())
  {
  }

  // Called to initialize a leader in the algorithm, which involves
  // sending out the initial path messages and being ready to receive
  // them.
  void initialize_leader(Vertex alpha, const Graph& g);

  /// Handle a path message on edge e. The path will be destroyed by
  /// this operation.
  void
  operator()(Edge e, path_t& path, const Graph& g);

  /// Handle a tree message on edge e. in_same_bicomponent will be
  /// destroyed by this operation.
  void
  operator()(Edge e, Vertex gamma, path_t& in_same_bicomponent,
             const Graph& g);

  // Handle a name message.
  void operator()(Edge e, edges_size_type name, const Graph& g);

  // Retrieve the phase
  unsigned char get_phase() const { return phase; }

  // Start the naming phase. The current phase must be 3 (echo), and
  // the offset contains the offset at which this processor should
  // begin when labelling its bicomponents. The offset is just a
  // parallel prefix sum of the number of bicomponents in each
  // processor that precedes it (globally).
  void
  start_naming_phase(Vertex alpha, const Graph& g, edges_size_type offset);

  /* Determine the number of bicomponents that we will be naming
   * ourselves.
   */
  edges_size_type num_starting_bicomponents(Vertex alpha, const Graph& g);

  // Fill in the edge component map with biconnected component
  // numbers.
  template<typename ComponentMap>
  void fill_edge_map(Vertex alpha, const Graph& g, ComponentMap& component);

 protected:
  /* Start the echo phase (phase 3) where we propagate information up
     the tree. */
  void echo_phase(Vertex alpha, const Graph& g);

  /* Retrieve the index of edge in the out-edges list of target(e, g). */
  std::size_t get_edge_index(Edge e, const Graph& g);

  /* Retrieve the index of the edge incidence on v in the out-edges
     list of vertex u. */
  std::size_t get_incident_edge_index(Vertex u, Vertex v, const Graph& g);

  /* Keeps track of which phase of the algorithm we are in. There are
   * four phases plus the "finished" phase:
   *
   *   1) Building the spanning tree
   *   2) Discovering cycles
   *   3) Echoing back up the spanning tree
   *   4) Labelling biconnected components
   *   5) Finished
   */
  unsigned char phase;

  /* The parent of this vertex in the spanning tree. This value will
     be graph_traits<Graph>::null_vertex() for the leader. */
  Vertex parent;

  /* The farthest ancestor up the tree that resides in the same
     biconnected component as we do. This information is approximate:
     we might not know about the actual farthest ancestor, but this is
     the farthest one we've seen so far. */
  Vertex eta;

  /* The number of edges that have not yet transmitted any messages to
     us. This starts at the degree of the vertex and decreases as we
     receive messages. When this counter hits zero, we're done with
     the second phase of the algorithm. In Hohberg's paper, the actual
     remaining edge set E is stored with termination when all edges
     have been removed from E, but we only need to detect termination
     so the set E need not be explicitly represented. */
  degree_size_type num_edges_not_transmitted;

  /* The path from the root of the spanning tree to this vertex. This
     vector will always store only the parts of the path leading up to
     this vertex, and not the vertex itself. Thus, it will be empty
     for the leader. */
  std::vector<Vertex> path_from_root;

  /* Structure containing all of the extra data we need to keep around
     PER EDGE. This information can not be contained within a property
     map, because it can't be shared among vertices without breaking
     the algorithm. Decreasing the size of this structure will drastically */
  struct per_edge_data
  {
    hohberg_message<Graph> msg;
    std::vector<Vertex> M;
    bool is_tree_edge;
    degree_size_type partition;
  };

  /* Data for each edge in the graph. This structure will be indexed
     by the position of the edge in the out_edges() list. */
  std::vector<per_edge_data> edge_data;

  /* The mapping from local partition numbers (0..n-1) to global
     partition numbers. */
  std::vector<edges_size_type> local_to_global_partitions;

  friend class boost::serialization::access;

  // We cannot actually serialize a vertex processor, nor would we
  // want to. However, the fact that we're putting instances into a
  // distributed_property_map means that we need to have a serialize()
  // function available.
  template<typename Archiver>
  void serialize(Archiver&, const unsigned int /*version*/)
  {
    assert(false);
  }
};

template<typename Graph>
void
hohberg_vertex_processor<Graph>::initialize_leader(Vertex alpha,
                                                   const Graph& g)
{
  using namespace hohberg_detail;

  ProcessGroup pg = process_group(g);

  typename property_map<Graph, vertex_owner_t>::const_type
    owner = get(vertex_owner, g);

  path_header<Edge> header;
  header.path_length = 1;
  BGL_FORALL_OUTEDGES_T(alpha, e, g, Graph) {
    header.edge = e;
    send(pg, get(owner, target(e, g)), msg_path_header, header);
    send(pg, get(owner, target(e, g)), msg_path_vertices, alpha);
  }

  num_edges_not_transmitted = degree(alpha, g);
  edge_data.resize(num_edges_not_transmitted);
  phase = 2;
}

template<typename Graph>
void
hohberg_vertex_processor<Graph>::operator()(Edge e, path_t& path,
                                            const Graph& g)
{
  using namespace hohberg_detail;

  typename property_map<Graph, vertex_owner_t>::const_type
    owner = get(vertex_owner, g);

#ifdef PBGL_HOHBERG_DEBUG
//  std::cerr << local(source(e, g)) << '@' << owner(source(e, g)) << " -> "
//            << local(target(e, g)) << '@' << owner(target(e, g)) << ": path(";
//  for (std::size_t i = 0; i < path.size(); ++i) {
//    if (i > 0) std::cerr << ' ';
//    std::cerr << local(path[i]) << '@' << owner(path[i]);
//  }
  std::cerr << "), phase = " << (int)phase << std::endl;
#endif

  // Get access to edge-specific data
  if (edge_data.empty())
    edge_data.resize(degree(target(e, g), g));
  per_edge_data& edata = edge_data[get_edge_index(e, g)];

  // Record the message. We'll need it in phase 3.
  edata.msg.assign(path);

  // Note: "alpha" refers to the vertex "processor" receiving the
  // message.
  Vertex alpha = target(e, g);

  switch (phase) {
  case 1:
    {
      num_edges_not_transmitted = degree(alpha, g) - 1;
      edata.is_tree_edge = true;
      parent = path.back();
      eta = parent;
      edata.M.clear(); edata.M.push_back(parent);

      // Broadcast the path from the root to our potential children in
      // the spanning tree.
      path.push_back(alpha);
      path_header<Edge> header;
      header.path_length = path.size();
      ProcessGroup pg = process_group(g);
      BGL_FORALL_OUTEDGES_T(alpha, oe, g, Graph) {
        // Skip the tree edge we just received
        if (target(oe, g) != source(e, g)) {
          header.edge = oe;
          send(pg, get(owner, target(oe, g)), msg_path_header, header);
          send(pg, get(owner, target(oe, g)), msg_path_vertices, &path[0],
               header.path_length);
        }
      }
      path.pop_back();

      // Swap the old path in, to save some extra copying. Nobody
      path_from_root.swap(path);

      // Once we have received our place in the spanning tree, move on
      // to phase 2.
      phase = 2;

      // If we only had only edge, skip to phase 3.
      if (num_edges_not_transmitted == 0)
        echo_phase(alpha, g);
      return;
    }

  case 2:
    {
      --num_edges_not_transmitted;
      edata.is_tree_edge = false;

      // Determine if alpha (our vertex) is in the path
      path_iterator pos = std::find(path.begin(), path.end(), alpha);
      if (pos != path.end()) {
        // Case A: There is a cycle alpha beta ... gamma alpha
        // M(e) <- {beta, gammar}
        edata.M.clear();
        ++pos;
        // If pos == path.end(), we have a self-loop
        if (pos != path.end()) {
          // Add beta
          edata.M.push_back(*pos);
          ++pos;
        }
        // If pos == path.end(), we have a self-loop or beta == gamma
        // (parallel edge). Otherwise, add gamma.
        if (pos != path.end()) edata.M.push_back(path.back());
      } else {
        // Case B: There is a cycle but we haven't seen alpha yet.
        // M(e) = {parent, path.back()}
        edata.M.clear();
        edata.M.push_back(path.back());
        if (parent != path.back()) edata.M.push_back(parent);

        // eta = inf(eta, bra(pi_t, pi))
        eta = infimum(path_from_root, eta, branch_point(path_from_root, path));
      }
      if (num_edges_not_transmitted == 0)
        echo_phase(alpha, g);
      break;
    }

  default:
//    std::cerr << "Phase is " << int(phase) << "\n";
    assert(false);
  }
}

template<typename Graph>
void
hohberg_vertex_processor<Graph>::operator()(Edge e, Vertex gamma,
                                            path_t& in_same_bicomponent,
                                            const Graph& g)
{
  using namespace hohberg_detail;

#ifdef PBGL_HOHBERG_DEBUG
  std::cerr << local(source(e, g)) << '@' << owner(source(e, g)) << " -> "
            << local(target(e, g)) << '@' << owner(target(e, g)) << ": tree("
            << local(gamma) << '@' << owner(gamma) << ", ";
  for (std::size_t i = 0; i < in_same_bicomponent.size(); ++i) {
    if (i > 0) std::cerr << ' ';
    std::cerr << local(in_same_bicomponent[i]) << '@'
              << owner(in_same_bicomponent[i]);
  }
  std::cerr << ", " << local(source(e, g)) << '@' << owner(source(e, g))
            << "), phase = " << (int)phase << std::endl;
#endif

  // Get access to edge-specific data
  per_edge_data& edata = edge_data[get_edge_index(e, g)];

  // Record the message. We'll need it in phase 3.
  edata.msg.assign(gamma, in_same_bicomponent);

  // Note: "alpha" refers to the vertex "processor" receiving the
  // message.
  Vertex alpha = target(e, g);
  Vertex beta = source(e, g);

  switch (phase) {
  case 2:
    --num_edges_not_transmitted;
    edata.is_tree_edge = true;

    if (gamma == alpha) {
      // Case C
      edata.M.swap(in_same_bicomponent);
    } else {
      // Case D
      edata.M.clear();
      edata.M.push_back(parent);
      if (beta != parent) edata.M.push_back(beta);
      eta = infimum(path_from_root, eta, gamma);
    }
    if (num_edges_not_transmitted == 0)
      echo_phase(alpha, g);
    break;

  default:
    assert(false);
  }
}

template<typename Graph>
void
hohberg_vertex_processor<Graph>::operator()(Edge e, edges_size_type name,
                                            const Graph& g)
{
  using namespace hohberg_detail;

#ifdef PBGL_HOHBERG_DEBUG
  std::cerr << local(source(e, g)) << '@' << owner(source(e, g)) << " -> "
            << local(target(e, g)) << '@' << owner(target(e, g)) << ": name("
            << name << "), phase = " << (int)phase << std::endl;
#endif

  assert(phase == 4);

  typename property_map<Graph, vertex_owner_t>::const_type
    owner = get(vertex_owner, g);

  // Send name messages along the spanning tree edges that are in the
  // same bicomponent as the edge to our parent.
  ProcessGroup pg = process_group(g);

  Vertex alpha = target(e, g);

  std::size_t idx = 0;
  BGL_FORALL_OUTEDGES_T(alpha, e, g, Graph) {
    per_edge_data& edata = edge_data[idx++];
    if (edata.is_tree_edge
        && find(edata.M.begin(), edata.M.end(), parent) != edata.M.end()
        && target(e, g) != parent) {
      // Notify our children in the spanning tree of this name
      name_header<Edge> header;
      header.edge = e;
      header.name = name;
      send(pg, get(owner, target(e, g)), msg_name, header);
    } else if (target(e, g) == parent) {
      // Map from local partition numbers to global bicomponent numbers
      local_to_global_partitions[edata.partition] = name;
    }
  }

  // Final stage
  phase = 5;
}

template<typename Graph>
typename hohberg_vertex_processor<Graph>::edges_size_type
hohberg_vertex_processor<Graph>::
num_starting_bicomponents(Vertex alpha, const Graph& g)
{
  edges_size_type not_mapped = (std::numeric_limits<edges_size_type>::max)();

  edges_size_type result = 0;
  std::size_t idx = 0;
  BGL_FORALL_OUTEDGES_T(alpha, e, g, Graph) {
    per_edge_data& edata = edge_data[idx++];
    if (edata.is_tree_edge
        && find(edata.M.begin(), edata.M.end(), parent) == edata.M.end()) {
      // Map from local partition numbers to global bicomponent numbers
      if (local_to_global_partitions[edata.partition] == not_mapped)
        local_to_global_partitions[edata.partition] = result++;
    }
  }

#ifdef PBGL_HOHBERG_DEBUG
  std::cerr << local(alpha) << '@' << owner(alpha) << " has " << result
            << " bicomponents originating at it." << std::endl;
#endif

  return result;
}

template<typename Graph>
template<typename ComponentMap>
void
hohberg_vertex_processor<Graph>::
fill_edge_map(Vertex alpha, const Graph& g, ComponentMap& component)
{
  std::size_t idx = 0;
  BGL_FORALL_OUTEDGES_T(alpha, e, g, Graph) {
    per_edge_data& edata = edge_data[idx++];
    local_put(component, e, local_to_global_partitions[edata.partition]);

#if defined(PBGL_HOHBERG_DEBUG) && PBGL_HOHBERG_DEBUG > 2
    std::cerr << "component("
              << local(source(e, g)) << '@' << owner(source(e, g)) << " -> "
              << local(target(e, g)) << '@' << owner(target(e, g)) << ") = "
              << local_to_global_partitions[edata.partition]
              << " (partition = " << edata.partition << " of "
              << local_to_global_partitions.size() << ")" << std::endl;
#endif
  }
}

template<typename Graph>
void
hohberg_vertex_processor<Graph>::
start_naming_phase(Vertex alpha, const Graph& g, edges_size_type offset)
{
  using namespace hohberg_detail;

  assert(phase == 4);

  typename property_map<Graph, vertex_owner_t>::const_type
    owner = get(vertex_owner, g);

  // Send name messages along the spanning tree edges of the
  // components that we get to number.
  ProcessGroup pg = process_group(g);

  bool has_more_children_to_name = false;

  // Map from local partition numbers to global bicomponent numbers
  edges_size_type not_mapped = (std::numeric_limits<edges_size_type>::max)();
  for (std::size_t i = 0; i < local_to_global_partitions.size(); ++i) {
    if (local_to_global_partitions[i] != not_mapped)
      local_to_global_partitions[i] += offset;
  }

  std::size_t idx = 0;
  BGL_FORALL_OUTEDGES_T(alpha, e, g, Graph) {
    per_edge_data& edata = edge_data[idx++];
    if (edata.is_tree_edge
        && find(edata.M.begin(), edata.M.end(), parent) == edata.M.end()) {
      // Notify our children in the spanning tree of this new name
      name_header<Edge> header;
      header.edge = e;
      header.name = local_to_global_partitions[edata.partition];
      send(pg, get(owner, target(e, g)), msg_name, header);
    } else if (edata.is_tree_edge) {
      has_more_children_to_name = true;
    }
#if defined(PBGL_HOHBERG_DEBUG) && PBGL_HOHBERG_DEBUG > 2
    std::cerr << "M[" << local(source(e, g)) << '@' << owner(source(e, g))
              << " -> " << local(target(e, g)) << '@' << owner(target(e, g))
              << "] = ";
    for (std::size_t i = 0; i < edata.M.size(); ++i) {
      std::cerr << local(edata.M[i]) << '@' << owner(edata.M[i]) << ' ';
    }
    std::cerr << std::endl;
#endif
  }

  // See if we're done.
  if (!has_more_children_to_name)
    // Final stage
    phase = 5;
}

template<typename Graph>
void
hohberg_vertex_processor<Graph>::echo_phase(Vertex alpha, const Graph& g)
{
  using namespace hohberg_detail;

  typename property_map<Graph, vertex_owner_t>::const_type
    owner = get(vertex_owner, g);

  /* We're entering the echo phase. */
  phase = 3;

  if (parent != graph_traits<Graph>::null_vertex()) {
    Edge edge_to_parent;

#if defined(PBGL_HOHBERG_DEBUG) && PBGL_HOHBERG_DEBUG > 1
     std::cerr << local(alpha) << '@' << owner(alpha) << " echo: parent = "
               << local(parent) << '@' << owner(parent) << ", eta = "
               << local(eta) << '@' << owner(eta) << ", Gamma = ";
#endif

    std::vector<Vertex> bicomp;
    std::size_t e_index = 0;
    BGL_FORALL_OUTEDGES_T(alpha, e, g, Graph) {
      if (target(e, g) == parent && parent == eta) {
        edge_to_parent = e;
        if (find(bicomp.begin(), bicomp.end(), alpha) == bicomp.end()) {
#if defined(PBGL_HOHBERG_DEBUG) && PBGL_HOHBERG_DEBUG > 1
          std::cerr << local(alpha) << '@' << owner(alpha) << ' ';
#endif
          bicomp.push_back(alpha);
        }
      } else {
        if (target(e, g) == parent) edge_to_parent = e;

        per_edge_data& edata = edge_data[e_index];

        if (edata.msg.is_path()) {
          path_iterator pos = std::find(edata.msg.path_or_bicomp.begin(),
                                        edata.msg.path_or_bicomp.end(),
                                        eta);
          if (pos != edata.msg.path_or_bicomp.end()) {
            ++pos;
            if (pos != edata.msg.path_or_bicomp.end()
                && find(bicomp.begin(), bicomp.end(), *pos) == bicomp.end()) {
#if defined(PBGL_HOHBERG_DEBUG) && PBGL_HOHBERG_DEBUG > 1
              std::cerr << local(*pos) << '@' << owner(*pos) << ' ';
#endif
              bicomp.push_back(*pos);
            }
          }
        } else if (edata.msg.is_tree() && edata.msg.gamma == eta) {
          for (path_iterator i = edata.msg.path_or_bicomp.begin();
               i != edata.msg.path_or_bicomp.end(); ++i) {
            if (find(bicomp.begin(), bicomp.end(), *i) == bicomp.end()) {
#if defined(PBGL_HOHBERG_DEBUG) && PBGL_HOHBERG_DEBUG > 1
              std::cerr << local(*i) << '@' << owner(*i) << ' ';
#endif
              bicomp.push_back(*i);
            }
          }
        }
      }
      ++e_index;
    }
#ifdef PBGL_HOHBERG_DEBUG
    std::cerr << std::endl;
#endif

    // Send tree(eta, bicomp) to parent
    tree_header<Vertex, Edge> header;
    header.edge = edge_to_parent;
    header.gamma = eta;
    header.bicomp_length = bicomp.size();
    ProcessGroup pg = process_group(g);
    send(pg, get(owner, parent), msg_tree_header, header);
    send(pg, get(owner, parent), msg_tree_vertices, &bicomp[0],
         header.bicomp_length);
  }

  // Compute the partition of edges such that iff two edges e1 and e2
  // are in different subsets then M(e1) is disjoint from M(e2).

  // Start by putting each edge in a different partition
  std::vector<degree_size_type> parent_vec(edge_data.size());
  degree_size_type idx = 0;
  for (idx = 0; idx < edge_data.size(); ++idx)
    parent_vec[idx] = idx;

  // Go through each edge e, performing a union() on the edges
  // incident on all vertices in M[e].
  idx = 0;
  BGL_FORALL_OUTEDGES_T(alpha, e, g, Graph) {
    per_edge_data& edata = edge_data[idx++];

    // Compute union of vertices in M
    if (!edata.M.empty()) {
      degree_size_type e1 = get_incident_edge_index(alpha, edata.M.front(), g);
      while (parent_vec[e1] != e1) e1 = parent_vec[e1];

      for (std::size_t i = 1; i < edata.M.size(); ++i) {
        degree_size_type e2 = get_incident_edge_index(alpha, edata.M[i], g);
        while (parent_vec[e2] != e2) e2 = parent_vec[e2];
        parent_vec[e2] = e1;
      }
    }
  }

  edges_size_type not_mapped = (std::numeric_limits<edges_size_type>::max)();

  // Determine the number of partitions
  for (idx = 0; idx < parent_vec.size(); ++idx) {
    if (parent_vec[idx] == idx) {
      edge_data[idx].partition = local_to_global_partitions.size();
      local_to_global_partitions.push_back(not_mapped);
    }
  }

  // Assign partition numbers to each edge
  for (idx = 0; idx < parent_vec.size(); ++idx) {
    degree_size_type rep = parent_vec[idx];
    while (rep != parent_vec[rep]) rep = parent_vec[rep];
    edge_data[idx].partition = edge_data[rep].partition;
  }

  // Enter the naming phase (but don't send anything yet).
  phase = 4;
}

template<typename Graph>
std::size_t
hohberg_vertex_processor<Graph>::get_edge_index(Edge e, const Graph& g)
{
  std::size_t result = 0;
  BGL_FORALL_OUTEDGES_T(target(e, g), oe, g, Graph) {
    if (source(e, g) == target(oe, g)) return result;
    ++result;
  }
  assert(false);
}

template<typename Graph>
std::size_t
hohberg_vertex_processor<Graph>::get_incident_edge_index(Vertex u, Vertex v,
                                                         const Graph& g)
{
  std::size_t result = 0;
  BGL_FORALL_OUTEDGES_T(u, e, g, Graph) {
    if (target(e, g) == v) return result;
    ++result;
  }
  assert(false);
}

template<typename Graph, typename InputIterator, typename ComponentMap,
         typename VertexProcessorMap>
typename graph_traits<Graph>::edges_size_type
hohberg_biconnected_components
  (const Graph& g,
   ComponentMap component,
   InputIterator first, InputIterator last,
   VertexProcessorMap vertex_processor)
{
  using namespace boost::graph::parallel;
  using namespace hohberg_detail;
  using boost::parallel::all_reduce;

  typename property_map<Graph, vertex_owner_t>::const_type
    owner = get(vertex_owner, g);

  // The graph must be undirected
  BOOST_STATIC_ASSERT(
    (is_convertible<typename graph_traits<Graph>::directed_category,
                    undirected_tag>::value));

  // The graph must model Incidence Graph
  function_requires< IncidenceGraphConcept<Graph> >();

  typedef typename graph_traits<Graph>::edges_size_type edges_size_type;
  typedef typename graph_traits<Graph>::degree_size_type degree_size_type;
  typedef typename graph_traits<Graph>::vertex_descriptor vertex_descriptor;
  typedef typename graph_traits<Graph>::edge_descriptor edge_descriptor;

  // Retrieve the process group we will use for communication
  typedef typename process_group_type<Graph>::type process_group_type;
  process_group_type pg = process_group(g);

  // Keeps track of the edges that we know to be tree edges.
  std::vector<edge_descriptor> tree_edges;

  // The leaders send out a path message to initiate the algorithm
  while (first != last) {
    vertex_descriptor leader = *first;
    if (process_id(pg) == get(owner, leader))
      vertex_processor[leader].initialize_leader(leader, g);
    ++first;
  }
  synchronize(pg);

  // Will hold the number of bicomponents in the graph.
  edges_size_type num_bicomponents = 0;

  // Keep track of the path length that we should expect, based on the
  // level in the breadth-first search tree. At present, this is only
  // used as a sanity check. TBD: This could be used to decrease the
  // amount of communication required per-edge (by about 4 bytes).
  std::size_t path_length = 1;

  typedef std::vector<vertex_descriptor> path_t;
  typedef typename path_t::iterator path_iterator;

  unsigned char minimum_phase = 5;
  do {
    while (optional<std::pair<int, int> > msg = probe(pg)) {
      switch (msg->second) {
      case msg_path_header:
        {
          // Receive the path header
          path_header<edge_descriptor> header;
          receive(pg, msg->first, msg->second, header);
          assert(path_length == header.path_length);

          // Receive the path itself
          path_t path(path_length);
          receive(pg, msg->first, msg_path_vertices, &path[0], path_length);

          edge_descriptor e = header.edge;
          vertex_processor[target(e, g)](e, path, g);
        }
        break;

      case msg_path_vertices:
        // Should be handled in msg_path_header case, unless we're going
        // stateless.
        assert(false);
        break;

      case msg_tree_header:
        {
          // Receive the tree header
          tree_header<vertex_descriptor, edge_descriptor> header;
          receive(pg, msg->first, msg->second, header);

          // Receive the tree itself
          path_t in_same_bicomponent(header.bicomp_length);
          receive(pg, msg->first, msg_tree_vertices, &in_same_bicomponent[0],
                  header.bicomp_length);

          edge_descriptor e = header.edge;
          vertex_processor[target(e, g)](e, header.gamma, in_same_bicomponent,
                                         g);
        }
        break;

      case msg_tree_vertices:
        // Should be handled in msg_tree_header case, unless we're
        // going stateless.
        assert(false);
        break;

      case msg_name:
        {
          name_header<edge_descriptor> header;
          receive(pg, msg->first, msg->second, header);
          edge_descriptor e = header.edge;
          vertex_processor[target(e, g)](e, header.name, g);
        }
        break;

      default:
        assert(false);
      }
    }
    ++path_length;

    // Compute minimum phase locally
    minimum_phase = 5;
    unsigned char maximum_phase = 1;
    BGL_FORALL_VERTICES_T(v, g, Graph) {
      minimum_phase = (std::min)(minimum_phase, vertex_processor[v].get_phase());
      maximum_phase = (std::max)(maximum_phase, vertex_processor[v].get_phase());
    }

#ifdef PBGL_HOHBERG_DEBUG
    if (process_id(pg) == 0)
      std::cerr << "<---------End of stage------------->" << std::endl;
#endif
    // Compute minimum phase globally
    minimum_phase = all_reduce(pg, minimum_phase, boost::mpi::minimum<char>());

#ifdef PBGL_HOHBERG_DEBUG
    if (process_id(pg) == 0)
      std::cerr << "Minimum phase = " << (int)minimum_phase << std::endl;
#endif

    if (minimum_phase == 4
        && all_reduce(pg, maximum_phase, boost::mpi::maximum<char>()) == 4) {

#ifdef PBGL_HOHBERG_DEBUG
      if (process_id(pg) == 0)
        std::cerr << "<---------Naming phase------------->" << std::endl;
#endif
      // Compute the biconnected component number offsets for each
      // vertex.
      std::vector<edges_size_type> local_offsets;
      local_offsets.reserve(num_vertices(g));
      edges_size_type num_local_bicomponents = 0;
      BGL_FORALL_VERTICES_T(v, g, Graph) {
        local_offsets.push_back(num_local_bicomponents);
        num_local_bicomponents +=
          vertex_processor[v].num_starting_bicomponents(v, g);
      }

      synchronize(pg);

      // Find our the number of bicomponent names that will originate
      // from each process. This tells us how many bicomponents are in
      // the entire graph and what our global offset is for computing
      // our own biconnected component names.
      std::vector<edges_size_type> all_bicomponents(num_processes(pg));
      all_gather(pg, &num_local_bicomponents, &num_local_bicomponents + 1,
                 all_bicomponents);
      num_bicomponents = 0;
      edges_size_type my_global_offset = 0;
      for (std::size_t i = 0; i < all_bicomponents.size(); ++i) {
        if (i == (std::size_t)process_id(pg))
          my_global_offset = num_bicomponents;
        num_bicomponents += all_bicomponents[i];
      }

      std::size_t index = 0;
      BGL_FORALL_VERTICES_T(v, g, Graph) {
        edges_size_type offset = my_global_offset + local_offsets[index++];
        vertex_processor[v].start_naming_phase(v, g, offset);
      }
    }

    synchronize(pg);
  } while (minimum_phase < 5);

  // Number the edges appropriately.
  BGL_FORALL_VERTICES_T(v, g, Graph)
    vertex_processor[v].fill_edge_map(v, g, component);

  return num_bicomponents;
}

template<typename Graph, typename ComponentMap, typename InputIterator>
typename graph_traits<Graph>::edges_size_type
hohberg_biconnected_components
  (const Graph& g, ComponentMap component,
   InputIterator first, InputIterator last)

{
  std::vector<hohberg_vertex_processor<Graph> >
    vertex_processors(num_vertices(g));
  return hohberg_biconnected_components
           (g, component, first, last,
            make_iterator_property_map(vertex_processors.begin(),
                                       get(vertex_index, g)));
}

template<typename Graph, typename ComponentMap, typename ParentMap>
typename graph_traits<Graph>::edges_size_type
hohberg_biconnected_components(const Graph& g, ComponentMap component,
                               ParentMap parent)
{
  // We need the connected components of the graph, but we don't care
  // about component numbers.
  connected_components(g, dummy_property_map(), parent);

  // Each root in the parent map is a leader
  typedef typename graph_traits<Graph>::vertex_descriptor vertex_descriptor;
  std::vector<vertex_descriptor> leaders;
  BGL_FORALL_VERTICES_T(v, g, Graph)
    if (get(parent, v) == v) leaders.push_back(v);

  return hohberg_biconnected_components(g, component,
                                        leaders.begin(), leaders.end());
}

template<typename Graph, typename ComponentMap>
typename graph_traits<Graph>::edges_size_type
hohberg_biconnected_components(const Graph& g, ComponentMap component)
{
  typedef typename graph_traits<Graph>::vertex_descriptor vertex_descriptor;
  std::vector<vertex_descriptor> parents(num_vertices(g));
  return hohberg_biconnected_components
           (g, component, make_iterator_property_map(parents.begin(),
                                                     get(vertex_index, g)));
}

} } } // end namespace boost::graph::distributed

#endif // BOOST_GRAPH_DISTRIBUTED_HOHBERG_BICONNECTED_COMPONENTS_HPP


// Copyright (c) 2020-2021 Franz Alt
// This code is licensed under MIT license (see LICENSE.txt for details).

#include <libcvpg/imageproc/algorithms/k_means.hpp>

#include <algorithm>
#include <cstdint>
#include <iterator>
#include <memory>
#include <random>
#include <vector>

#include <boost/asynchronous/algorithm/then.hpp>

#include <libcvpg/core/exception.hpp>
#include <libcvpg/core/histogram.hpp>
#include <libcvpg/imageproc/algorithms/tiling.hpp>
#include <libcvpg/imageproc/algorithms/tiling/histogram.hpp>
#include <libcvpg/imageproc/algorithms/tiling/functors/histogram.hpp>

namespace {

void histogram_if_gray_8bit(std::uint8_t * src, std::vector<std::size_t> * dst, std::uint8_t * predicate_img, std::function<bool(std::uint8_t const & cluster)> predicate, std::size_t from_x, std::size_t to_x, std::size_t from_y, std::size_t to_y, cvpg::imageproc::algorithms::tiling_parameters parameters)
{
    const std::size_t image_width = parameters.image_width;

    std::uint8_t * src_line = nullptr;
    std::uint8_t * predicate_line = nullptr;

    for (std::size_t y = from_y; y <= to_y; ++y)
    {
        src_line = src + image_width * y;
        predicate_line = predicate_img + image_width * y;

        for (std::size_t x = from_x; x <= to_x; ++x)
        {
            if (predicate(predicate_line[x]))
            {
                ++((*dst)[src_line[x]]);
            }
        }
    }
}

struct merge_histograms_task : public boost::asynchronous::continuation_task<std::shared_ptr<std::vector<cvpg::histogram<std::size_t> > >>
{
    merge_histograms_task(std::shared_ptr<std::vector<cvpg::histogram<std::size_t> > > a, std::shared_ptr<std::vector<cvpg::histogram<std::size_t> > > b)
        : boost::asynchronous::continuation_task<std::shared_ptr<std::vector<cvpg::histogram<std::size_t> > > >("k_means_task::merge_histograms_task")
        , m_a(a)
        , m_b(b)
    {}

    void operator()()
    {
        try
        {
            if (m_a->size() != m_b->size())
            {
                throw cvpg::exception("cannot merge histograms of different sizes");
            }

            std::vector<cvpg::histogram<std::size_t> > histograms;
            histograms.reserve(m_a->size());

            for (std::size_t i = 0; i < m_a->size(); ++i)
            {
                histograms.push_back(m_a->at(i) + m_b->at(i));
            }

            this_task_result().set_value(std::make_shared<std::vector<cvpg::histogram<std::size_t> > >(std::move(histograms)));
        }
        catch (...)
        {
            this_task_result().set_exception(std::current_exception());
        }
    }

private:
    std::shared_ptr<std::vector<cvpg::histogram<std::size_t> > > m_a;
    std::shared_ptr<std::vector<cvpg::histogram<std::size_t> > > m_b;
};

boost::asynchronous::detail::callback_continuation<std::shared_ptr<std::vector<cvpg::histogram<std::size_t> > > >
merge_histograms(std::shared_ptr<std::vector<cvpg::histogram<std::size_t> > > a, std::shared_ptr<std::vector<cvpg::histogram<std::size_t> > > b)
{
    return boost::asynchronous::top_level_callback_continuation<std::shared_ptr<std::vector<cvpg::histogram<std::size_t> > > >(
               merge_histograms_task(a, b)
           );
}

struct point
{
    point() = default;

    point(std::uint8_t x_, std::uint8_t y_, std::uint8_t z_)
        : x(x_)
        , y(y_)
        , z(z_)
    {}

    std::uint8_t x = 0;
    std::uint8_t y = 0;
    std::uint8_t z = 0;
};

void determine_cluster_gray_8bit(std::uint8_t * src, std::uint8_t * dst, std::size_t from_x, std::size_t to_x, std::size_t from_y, std::size_t to_y, cvpg::imageproc::algorithms::tiling_parameters parameters, std::vector<point> centers)
{
    const std::size_t image_width = parameters.image_width;
    const std::size_t clusters = centers.size();

    std::vector<std::size_t> distances(clusters, 0);

    std::uint8_t * src_line = nullptr;
    std::uint8_t * dst_line = nullptr;

    for (std::size_t y = from_y; y <= to_y; ++y)
    {
        const std::size_t offset_y = image_width * y;

        src_line = src + offset_y;
        dst_line = dst + offset_y;

        for (std::size_t x = from_x; x <= to_x; ++x)
        {
            const point p = { src_line[x], 0, 0 };

            for (std::size_t i = 0; i < clusters; ++i)
            {
                const point & c = centers[i];

                const double dxs = static_cast<double>(static_cast<std::int16_t>(p.x) - static_cast<std::int16_t>(c.x));

                // calculate the euclidian distance between current pixel and center of current cluster
                distances[i] = static_cast<std::size_t>(std::abs(dxs));
            }

            // determine index of cluster with the smallest distance to current pixel (in color space) for c=[1..k]
            const std::uint8_t c = static_cast<std::uint8_t>(std::distance(distances.begin(), std::min_element(distances.begin(), distances.end()))) + 1;

            // write cluster number to destination image
            dst_line[x] = c;
        }
    }
}

void determine_cluster_rgb_8bit(std::uint8_t * src_red, std::uint8_t * src_green, std::uint8_t * src_blue, std::uint8_t * dst, std::size_t from_x, std::size_t to_x, std::size_t from_y, std::size_t to_y, cvpg::imageproc::algorithms::tiling_parameters parameters, std::vector<point> centers)
{
    const std::size_t image_width = parameters.image_width;
    const std::size_t clusters = centers.size();

    std::vector<std::size_t> distances(clusters, 0);

    std::uint8_t * src_red_line = nullptr;
    std::uint8_t * src_green_line = nullptr;
    std::uint8_t * src_blue_line = nullptr;
    std::uint8_t * dst_line = nullptr;

    for (std::size_t y = from_y; y <= to_y; ++y)
    {
        const std::size_t offset_y = image_width * y;

        src_red_line = src_red + offset_y;
        src_green_line = src_green + offset_y;
        src_blue_line = src_blue + offset_y;
        dst_line = dst + offset_y;

        for (std::size_t x = from_x; x <= to_x; ++x)
        {
            const point p = { src_red_line[x], src_green_line[x], src_blue_line[x] };

            for (std::size_t i = 0; i < clusters; ++i)
            {
                const point & c = centers[i];

                const double dxs = static_cast<double>(static_cast<std::int16_t>(p.x) - static_cast<std::int16_t>(c.x));
                const double dys = static_cast<double>(static_cast<std::int16_t>(p.y) - static_cast<std::int16_t>(c.y));
                const double dzs = static_cast<double>(static_cast<std::int16_t>(p.z) - static_cast<std::int16_t>(c.z));

                // calculate the euclidian distance between current pixel and center of current cluster
                distances[i] = static_cast<std::size_t>(std::sqrt(dxs * dxs + dys * dys + dzs * dzs));
            }

            // determine index of cluster with the smallest distance to current pixel (in color space) for c=[1..k]
            const std::uint8_t c = static_cast<std::uint8_t>(std::distance(distances.begin(), std::min_element(distances.begin(), distances.end()))) + 1;

            // write cluster number to destination image
            dst_line[x] = c;
        }
    }
}

template<class image_type>
struct determine_cluster_task : public boost::asynchronous::continuation_task<cvpg::image_gray_8bit> {};

template<>
struct determine_cluster_task<cvpg::image_gray_8bit> : public boost::asynchronous::continuation_task<cvpg::image_gray_8bit>
{
    determine_cluster_task(std::shared_ptr<cvpg::image_gray_8bit> image, std::vector<point> centers)
        : boost::asynchronous::continuation_task<cvpg::image_gray_8bit>("k_means_task::determine_cluster_task")
        , m_image(std::move(image))
        , m_centers(std::move(centers))
    {}

    void operator()()
    {
        auto tf = cvpg::imageproc::algorithms::tiling_functors::image<cvpg::image_gray_8bit, cvpg::image_gray_8bit>({{ *m_image }});
        tf.parameters.image_width = m_image->width();
        tf.parameters.image_height = m_image->height();
        tf.parameters.cutoff_x = 512; // TODO use parameter
        tf.parameters.cutoff_y = 512; // TODO use parameter

        tf.tile_algorithm_task = [centers = m_centers](std::shared_ptr<cvpg::image_gray_8bit> src1, std::shared_ptr<cvpg::image_gray_8bit> /*src2*/, std::shared_ptr<cvpg::image_gray_8bit> dst, std::size_t from_x, std::size_t to_x, std::size_t from_y, std::size_t to_y, cvpg::imageproc::algorithms::tiling_parameters parameters)
        {
            determine_cluster_gray_8bit(src1->data(0).get(), dst->data(0).get(), from_x, to_x, from_y, to_y, std::move(parameters), std::move(centers));
        };

        boost::asynchronous::create_callback_continuation(
            [result = this->this_task_result()](auto cont_res)
            {
                try
                {
                    result.set_value(std::get<0>(cont_res).get());
                }
                catch (...)
                {
                    result.set_exception(std::current_exception());
                }
            },
            cvpg::imageproc::algorithms::tiling(std::move(tf))
        );
    }

private:
    std::shared_ptr<cvpg::image_gray_8bit> m_image;

    std::vector<point> m_centers;
};

template<>
struct determine_cluster_task<cvpg::image_rgb_8bit> : public boost::asynchronous::continuation_task<cvpg::image_gray_8bit>
{
    determine_cluster_task(std::shared_ptr<cvpg::image_rgb_8bit> image, std::vector<point> centers)
        : boost::asynchronous::continuation_task<cvpg::image_gray_8bit>("k_means_task::determine_cluster_task")
        , m_image(std::move(image))
        , m_centers(std::move(centers))
    {}

    void operator()()
    {
        auto tf = cvpg::imageproc::algorithms::tiling_functors::image<cvpg::image_rgb_8bit, cvpg::image_gray_8bit>({{ *m_image }});
        tf.parameters.image_width = m_image->width();
        tf.parameters.image_height = m_image->height();
        tf.parameters.cutoff_x = 512; // TODO use parameter
        tf.parameters.cutoff_y = 512; // TODO use parameter

        tf.tile_algorithm_task = [centers = m_centers](std::shared_ptr<cvpg::image_rgb_8bit> src1, std::shared_ptr<cvpg::image_rgb_8bit> /*src2*/, std::shared_ptr<cvpg::image_gray_8bit> dst, std::size_t from_x, std::size_t to_x, std::size_t from_y, std::size_t to_y, cvpg::imageproc::algorithms::tiling_parameters parameters)
        {
            determine_cluster_rgb_8bit(src1->data(0).get(), src1->data(1).get(), src1->data(2).get(), dst->data(0).get(), from_x, to_x, from_y, to_y, std::move(parameters), std::move(centers));
        };

        boost::asynchronous::create_callback_continuation(
            [result = this->this_task_result()](auto cont_res)
            {
                try
                {
                    result.set_value(std::get<0>(cont_res).get());
                }
                catch (...)
                {
                    result.set_exception(std::current_exception());
                }
            },
            cvpg::imageproc::algorithms::tiling(std::move(tf))
        );
    }

private:
    std::shared_ptr<cvpg::image_rgb_8bit> m_image;

    std::vector<point> m_centers;
};

template<class image_type>
boost::asynchronous::detail::callback_continuation<cvpg::image_gray_8bit> determine_cluster(std::shared_ptr<image_type> image, std::vector<point> centers)
{
    return boost::asynchronous::top_level_callback_continuation<cvpg::image_gray_8bit>(
               determine_cluster_task<image_type>(std::move(image), std::move(centers))
           );
}

template<class image_type>
struct calculate_cluster_means_task : public boost::asynchronous::continuation_task<std::pair<cvpg::image_gray_8bit, std::vector<point> > > {};

template<>
struct calculate_cluster_means_task<cvpg::image_gray_8bit> : public boost::asynchronous::continuation_task<std::pair<cvpg::image_gray_8bit, std::vector<point> > >
{
    calculate_cluster_means_task(std::shared_ptr<cvpg::image_gray_8bit> image, cvpg::image_gray_8bit clusters, std::size_t k)
        : boost::asynchronous::continuation_task<std::pair<cvpg::image_gray_8bit, std::vector<point> > >("k_means_task::calculate_cluster_means_task")
        , m_image(std::move(image))
        , m_clusters(std::make_shared<cvpg::image_gray_8bit>(std::move(clusters)))
        , m_k(k)
    {}

    void operator()()
    {
        auto tf = cvpg::imageproc::algorithms::tiling_functors::histogram<cvpg::image_gray_8bit, std::vector<cvpg::histogram<std::size_t> > >({{ *m_image }});
        tf.parameters.image_width = m_image->width();
        tf.parameters.image_height = m_image->height();
        tf.parameters.cutoff_x = 512; // TODO use parameter
        tf.parameters.cutoff_y = 512; // TODO use parameter

        tf.tile_algorithm_task = [clusters = m_clusters, k = m_k](std::shared_ptr<cvpg::image_gray_8bit> src1, std::shared_ptr<cvpg::image_gray_8bit> /*src2*/, std::shared_ptr<std::vector<cvpg::histogram<std::size_t> > > dst, std::size_t from_x, std::size_t to_x, std::size_t from_y, std::size_t to_y, cvpg::imageproc::algorithms::tiling_parameters parameters)
        {
            // for each cluster ('k' pieces) create one histograms
            std::vector<std::vector<std::size_t> > h(k, std::vector<std::size_t>(256));

            for (std::size_t i = 0; i < k; ++i)
            {
                histogram_if_gray_8bit(src1->data(0).get(), &h[i], clusters->data(0).get(), [i](std::uint8_t const & c){ return c == (i + 1); }, from_x, to_x, from_y, to_y, parameters);
            }

            // move arrays into histogram structure
            std::vector<cvpg::histogram<std::size_t> > histograms;
            histograms.reserve(h.size());

            std::for_each(h.begin(),
                          h.end(),
                          [&histograms](auto & histogram)
                          {
                              histograms.emplace_back(std::move(histogram));
                          });

            *dst = std::move(histograms);
        };

        tf.horizontal_merge_task = [](std::shared_ptr<std::vector<cvpg::histogram<std::size_t> > > dst1, std::shared_ptr<std::vector<cvpg::histogram<std::size_t> > > dst2, std::size_t /*from_x*/, std::size_t /*to_x*/, std::size_t /*from_y*/, std::size_t /*to_y*/, cvpg::imageproc::algorithms::tiling_parameters /*parameters*/)
        {
            return merge_histograms(dst1, dst2);
        };

        tf.vertical_merge_task = [](std::shared_ptr<std::vector<cvpg::histogram<std::size_t> > > dst1, std::shared_ptr<std::vector<cvpg::histogram<std::size_t> > > dst2, std::size_t /*from_x*/, std::size_t /*to_x*/, std::size_t /*from_y*/, std::size_t /*to_y*/, cvpg::imageproc::algorithms::tiling_parameters /*parameters*/)
        {
            return merge_histograms(dst1, dst2);
        };

        boost::asynchronous::create_callback_continuation(
            [result = this->this_task_result(), image = std::move(m_image), clusters = std::move(m_clusters), k = m_k](auto cont_res) mutable
            {
                try
                {
                    auto histograms = std::move(std::get<0>(cont_res).get());

                    auto sum_histogram =
                        [](const auto & histogram) -> std::pair<std::size_t, std::size_t>
                        {
                            std::size_t sum = 0;
                            std::size_t entries = 0;

                            for (std::size_t i = 0; i < histogram.bins(); ++i)
                            {
                                sum += i * histogram.at(i);
                                entries += histogram.at(i);
                            }

                            return { sum, entries };
                        };

                    // create new mean centers of all clusters
                    std::vector<point> centers;
                    centers.reserve(k);

                    for (std::size_t i = 0; i < k; ++i)
                    {
                        auto [sum, entries] = sum_histogram(histograms.at(i));

                        centers.emplace_back(
                            entries != 0 ? (sum / entries) : 0,
                            0,
                            0
                        );
                    }

                    result.set_value({ std::move(*clusters), std::move(centers) });
                }
                catch (...)
                {
                    result.set_exception(std::current_exception());
                }
            },
            cvpg::imageproc::algorithms::tiling(std::move(tf))
        );
    }

private:
    std::shared_ptr<cvpg::image_gray_8bit> m_image;

    std::shared_ptr<cvpg::image_gray_8bit> m_clusters;

    std::size_t m_k;
};

template<>
struct calculate_cluster_means_task<cvpg::image_rgb_8bit> : public boost::asynchronous::continuation_task<std::pair<cvpg::image_gray_8bit, std::vector<point> > >
{
    calculate_cluster_means_task(std::shared_ptr<cvpg::image_rgb_8bit> image, cvpg::image_gray_8bit clusters, std::size_t k)
        : boost::asynchronous::continuation_task<std::pair<cvpg::image_gray_8bit, std::vector<point> > >("k_means_task::calculate_cluster_means_task")
        , m_image(std::move(image))
        , m_clusters(std::make_shared<cvpg::image_gray_8bit>(std::move(clusters)))
        , m_k(k)
    {}

    void operator()()
    {
        auto tf = cvpg::imageproc::algorithms::tiling_functors::histogram<cvpg::image_rgb_8bit, std::vector<cvpg::histogram<std::size_t> > >({{ *m_image }});
        tf.parameters.image_width = m_image->width();
        tf.parameters.image_height = m_image->height();
        tf.parameters.cutoff_x = 512; // TODO use parameter
        tf.parameters.cutoff_y = 512; // TODO use parameter

        tf.tile_algorithm_task = [clusters = m_clusters, k = m_k](std::shared_ptr<cvpg::image_rgb_8bit> src1, std::shared_ptr<cvpg::image_rgb_8bit> /*src2*/, std::shared_ptr<std::vector<cvpg::histogram<std::size_t> > > dst, std::size_t from_x, std::size_t to_x, std::size_t from_y, std::size_t to_y, cvpg::imageproc::algorithms::tiling_parameters parameters)
        {
            // for each cluster ('k' pieces) create 3 histograms for red, green and blue parts
            std::vector<std::vector<std::size_t> > h(k * 3, std::vector<std::size_t>(256));

            for (std::size_t i = 0; i < k; ++i)
            {
                histogram_if_gray_8bit(src1->data(0).get(), &h[i * 3 + 0], clusters->data(0).get(), [i](std::uint8_t const & c){ return c == (i + 1); }, from_x, to_x, from_y, to_y, parameters);
                histogram_if_gray_8bit(src1->data(1).get(), &h[i * 3 + 1], clusters->data(0).get(), [i](std::uint8_t const & c){ return c == (i + 1); }, from_x, to_x, from_y, to_y, parameters);
                histogram_if_gray_8bit(src1->data(2).get(), &h[i * 3 + 2], clusters->data(0).get(), [i](std::uint8_t const & c){ return c == (i + 1); }, from_x, to_x, from_y, to_y, parameters);
            }

            // move arrays into histogram structure
            std::vector<cvpg::histogram<std::size_t> > histograms;
            histograms.reserve(h.size());

            std::for_each(h.begin(),
                          h.end(),
                          [&histograms](auto & histogram)
                          {
                              histograms.emplace_back(std::move(histogram));
                          });

            *dst = std::move(histograms);
        };

        tf.horizontal_merge_task = [](std::shared_ptr<std::vector<cvpg::histogram<std::size_t> > > dst1, std::shared_ptr<std::vector<cvpg::histogram<std::size_t> > > dst2, std::size_t /*from_x*/, std::size_t /*to_x*/, std::size_t /*from_y*/, std::size_t /*to_y*/, cvpg::imageproc::algorithms::tiling_parameters /*parameters*/)
        {
            return merge_histograms(dst1, dst2);
        };

        tf.vertical_merge_task = [](std::shared_ptr<std::vector<cvpg::histogram<std::size_t> > > dst1, std::shared_ptr<std::vector<cvpg::histogram<std::size_t> > > dst2, std::size_t /*from_x*/, std::size_t /*to_x*/, std::size_t /*from_y*/, std::size_t /*to_y*/, cvpg::imageproc::algorithms::tiling_parameters /*parameters*/)
        {
            return merge_histograms(dst1, dst2);
        };

        boost::asynchronous::create_callback_continuation(
            [result = this->this_task_result(), image = std::move(m_image), clusters = std::move(m_clusters), k = m_k](auto cont_res) mutable
            {
                try
                {
                    auto histograms = std::move(std::get<0>(cont_res).get());

                    if (histograms.size() % 3 != 0)
                    {
                        throw cvpg::exception("invalid amount of histograms");
                    }

                    auto sum_histogram =
                        [](const auto & histogram) -> std::pair<std::size_t, std::size_t>
                        {
                            std::size_t sum = 0;
                            std::size_t entries = 0;

                            for (std::size_t i = 0; i < histogram.bins(); ++i)
                            {
                                sum += i * histogram.at(i);
                                entries += histogram.at(i);
                            }

                            return { sum, entries };
                        };

                    // create new mean centers of all clusters
                    std::vector<point> centers;
                    centers.reserve(k);

                    for (std::size_t i = 0; i < k; ++i)
                    {
                        auto [sum_red, entries_red] = sum_histogram(histograms.at(i * 3 + 0));
                        auto [sum_green, entries_green] = sum_histogram(histograms.at(i * 3 + 1));
                        auto [sum_blue, entries_blue] = sum_histogram(histograms.at(i * 3 + 2));

                        centers.emplace_back(
                            entries_red != 0 ? (sum_red / entries_red) : 0,
                            entries_green != 0 ? (sum_green / entries_green) : 0,
                            entries_blue != 0 ? (sum_blue / entries_blue) : 0
                        );
                    }

                    result.set_value({ std::move(*clusters), std::move(centers) });
                }
                catch (...)
                {
                    result.set_exception(std::current_exception());
                }
            },
            cvpg::imageproc::algorithms::tiling(std::move(tf))
        );
    }

private:
    std::shared_ptr<cvpg::image_rgb_8bit> m_image;

    std::shared_ptr<cvpg::image_gray_8bit> m_clusters;

    std::size_t m_k;
};

template<class image_type>
boost::asynchronous::detail::callback_continuation<std::pair<cvpg::image_gray_8bit, std::vector<point> > > calculate_cluster_means(std::shared_ptr<image_type> image, cvpg::image_gray_8bit clusters, std::size_t k)
{
    return boost::asynchronous::top_level_callback_continuation<std::pair<cvpg::image_gray_8bit, std::vector<point> > >(
               calculate_cluster_means_task<image_type>(std::move(image), std::move(clusters), k)
           );
}

void create_cluster_result_gray_8bit(std::uint8_t * src, std::uint8_t * dst, std::size_t from_x, std::size_t to_x, std::size_t from_y, std::size_t to_y, cvpg::imageproc::algorithms::tiling_parameters parameters, std::vector<point> centers)
{
    const std::size_t image_width = parameters.image_width;

    std::uint8_t * src_line = nullptr;
    std::uint8_t * dst_line = nullptr;

    for (std::size_t y = from_y; y <= to_y; ++y)
    {
        const std::size_t offset_y = image_width * y;

        src_line = src + offset_y;
        dst_line = dst + offset_y;

        for (std::size_t x = from_x; x <= to_x; ++x)
        {
            const std::uint8_t cluster_number = src_line[x] - 1;
            const auto & cluster = centers.at(cluster_number);

            dst_line[x] = cluster.x;
        }
    }
}

void create_cluster_result_rgb_8bit(std::uint8_t * src, std::uint8_t * dst_red, std::uint8_t * dst_green, std::uint8_t * dst_blue, std::size_t from_x, std::size_t to_x, std::size_t from_y, std::size_t to_y, cvpg::imageproc::algorithms::tiling_parameters parameters, std::vector<point> centers)
{
    const std::size_t image_width = parameters.image_width;

    std::uint8_t * src_line = nullptr;
    std::uint8_t * dst_red_line = nullptr;
    std::uint8_t * dst_green_line = nullptr;
    std::uint8_t * dst_blue_line = nullptr;

    for (std::size_t y = from_y; y <= to_y; ++y)
    {
        const std::size_t offset_y = image_width * y;

        src_line = src + offset_y;
        dst_red_line = dst_red + offset_y;
        dst_green_line = dst_green + offset_y;
        dst_blue_line = dst_blue + offset_y;

        for (std::size_t x = from_x; x <= to_x; ++x)
        {
            const std::uint8_t cluster_number = src_line[x] - 1;
            const auto & cluster = centers.at(cluster_number);

            dst_red_line[x] = cluster.x;
            dst_green_line[x] = cluster.y;
            dst_blue_line[x] = cluster.z;
        }
    }
}

template<class image_type>
struct create_result_image_task : public boost::asynchronous::continuation_task<image_type> {};

template<>
struct create_result_image_task<cvpg::image_gray_8bit> : public boost::asynchronous::continuation_task<cvpg::image_gray_8bit>
{
    create_result_image_task(std::shared_ptr<cvpg::image_gray_8bit> clusters, std::vector<point> centers)
        : boost::asynchronous::continuation_task<cvpg::image_gray_8bit>("k_means_task::create_result_image_task")
        , m_clusters(std::move(clusters))
        , m_centers(std::move(centers))
    {}

    void operator()()
    {
        auto result_image = std::make_shared<cvpg::image_gray_8bit>(m_clusters->width(), m_clusters->height());

        auto tf = cvpg::imageproc::algorithms::tiling_functors::image<cvpg::image_gray_8bit, cvpg::image_gray_8bit>({{ *m_clusters }});
        tf.parameters.image_width = m_clusters->width();
        tf.parameters.image_height = m_clusters->height();
        tf.parameters.cutoff_x = 512; // TODO use parameter
        tf.parameters.cutoff_y = 512; // TODO use parameter

        tf.tile_algorithm_task = [centers = m_centers](std::shared_ptr<cvpg::image_gray_8bit> src1, std::shared_ptr<cvpg::image_gray_8bit> /*src2*/, std::shared_ptr<cvpg::image_gray_8bit> dst, std::size_t from_x, std::size_t to_x, std::size_t from_y, std::size_t to_y, cvpg::imageproc::algorithms::tiling_parameters parameters)
        {
            create_cluster_result_gray_8bit(src1->data(0).get(), dst->data(0).get(), from_x, to_x, from_y, to_y, std::move(parameters), std::move(centers));
        };

        boost::asynchronous::create_callback_continuation(
            [result = this->this_task_result()](auto cont_res)
            {
                try
                {
                    result.set_value(std::get<0>(cont_res).get());
                }
                catch (...)
                {
                    result.set_exception(std::current_exception());
                }
            },
            cvpg::imageproc::algorithms::tiling(std::move(tf))
        );
    }

private:
    std::shared_ptr<cvpg::image_gray_8bit> m_clusters;

    std::vector<point> m_centers;
};

template<>
struct create_result_image_task<cvpg::image_rgb_8bit> : public boost::asynchronous::continuation_task<cvpg::image_rgb_8bit>
{
    create_result_image_task(std::shared_ptr<cvpg::image_gray_8bit> clusters, std::vector<point> centers)
        : boost::asynchronous::continuation_task<cvpg::image_rgb_8bit>("k_means_task::create_result_image_task")
        , m_clusters(std::move(clusters))
        , m_centers(std::move(centers))
    {}

    void operator()()
    {
        auto result_image = std::make_shared<cvpg::image_rgb_8bit>(m_clusters->width(), m_clusters->height());

        auto tf = cvpg::imageproc::algorithms::tiling_functors::image<cvpg::image_gray_8bit, cvpg::image_rgb_8bit>({{ *m_clusters }});
        tf.parameters.image_width = m_clusters->width();
        tf.parameters.image_height = m_clusters->height();
        tf.parameters.cutoff_x = 512; // TODO use parameter
        tf.parameters.cutoff_y = 512; // TODO use parameter

        tf.tile_algorithm_task = [centers = m_centers](std::shared_ptr<cvpg::image_gray_8bit> src1, std::shared_ptr<cvpg::image_gray_8bit> /*src2*/, std::shared_ptr<cvpg::image_rgb_8bit> dst, std::size_t from_x, std::size_t to_x, std::size_t from_y, std::size_t to_y, cvpg::imageproc::algorithms::tiling_parameters parameters)
        {
            create_cluster_result_rgb_8bit(src1->data(0).get(), dst->data(0).get(), dst->data(1).get(), dst->data(2).get(), from_x, to_x, from_y, to_y, std::move(parameters), std::move(centers));
        };

        boost::asynchronous::create_callback_continuation(
            [result = this->this_task_result()](auto cont_res)
            {
                try
                {
                    result.set_value(std::get<0>(cont_res).get());
                }
                catch (...)
                {
                    result.set_exception(std::current_exception());
                }
            },
            cvpg::imageproc::algorithms::tiling(std::move(tf))
        );
    }

private:
    std::shared_ptr<cvpg::image_gray_8bit> m_clusters;

    std::vector<point> m_centers;
};

template<class image_type>
boost::asynchronous::detail::callback_continuation<image_type> create_result_image(std::shared_ptr<cvpg::image_gray_8bit> clusters, std::vector<point> centers)
{
    return boost::asynchronous::top_level_callback_continuation<image_type>(
               create_result_image_task<image_type>(std::move(clusters), std::move(centers))
           );
}

template<class image_type>
struct k_means_iteration_task : public boost::asynchronous::continuation_task<image_type>
{
    k_means_iteration_task(std::shared_ptr<image_type> image, std::vector<point> centers, std::size_t k, std::size_t iteration, std::size_t max_iterations, std::uint8_t eps)
        : boost::asynchronous::continuation_task<image_type>(std::string("k_means_task::k_means_iteration_task#iteration").append(std::to_string(iteration)))
        , m_image(std::move(image))
        , m_centers(std::move(centers))
        , m_k(k)
        , m_iteration(iteration)
        , m_max_iterations(max_iterations)
        , m_eps(eps)
    {}

    void operator()()
    {
        auto old_centers = m_centers;

        boost::asynchronous::create_callback_continuation(
            [result = this->this_task_result(), old_centers = std::move(old_centers), iteration = m_iteration, max_iterations = m_max_iterations, image = m_image, k = m_k, eps = m_eps](auto cont_res) mutable
            {
                try
                {
                    auto cluster_result = std::move(std::get<0>(cont_res).get());

                    auto clusters_image = std::make_shared<cvpg::image_gray_8bit>(std::move(cluster_result.first));
                    auto new_centers = std::move(cluster_result.second);

                    // calculate the distances of old and new center points
                    std::vector<std::uint8_t> distances;
                    distances.reserve(old_centers.size());

                    std::transform(old_centers.begin(),
                                   old_centers.end(),
                                   new_centers.begin(),
                                   std::back_inserter(distances),
                                   [](point const & a, point const & b)
                                   {
                                       const auto dx = static_cast<std::int16_t>(a.x) - static_cast<std::int16_t>(b.x);
                                       const auto dy = static_cast<std::int16_t>(a.y) - static_cast<std::int16_t>(b.y);
                                       const auto dz = static_cast<std::int16_t>(a.z) - static_cast<std::int16_t>(b.z);

                                       return static_cast<std::uint8_t>(std::sqrt(dx * dx + dy * dy + dz * dz));
                                   });

                    // check if all distances are not moving (moving below threshold)
                    const bool all_centers_steady = std::all_of(distances.begin(),
                                                                distances.end(),
                                                                [eps](auto const & distance)
                                                                {
                                                                    return distance < eps; // TODO define a parameter for this !!!
                                                                });

                    // check if no further iterations are needed
                    if ((iteration + 1) >= max_iterations || all_centers_steady)
                    {
                        boost::asynchronous::create_callback_continuation(
                            [result = std::move(result)](auto cont_res) mutable
                            {
                                try
                                {
                                    result.set_value(std::move(std::get<0>(cont_res).get()));
                                }
                                catch (...)
                                {
                                    result.set_exception(std::current_exception());
                                }
                            },
                            create_result_image<image_type>(clusters_image, std::move(new_centers))
                        );
                    }
                    else
                    {
                        boost::asynchronous::create_callback_continuation(
                            [result = std::move(result)](auto cont_res) mutable
                            {
                                try
                                {
                                    result.set_value(std::move(std::get<0>(cont_res).get()));
                                }
                                catch (...)
                                {
                                    result.set_exception(std::current_exception());
                                }
                            },
                            k_means_iteration_task(image, std::move(new_centers), k, iteration + 1, max_iterations, eps)
                        );
                    }
                }
                catch (...)
                {
                    result.set_exception(std::current_exception());
                }
            },
            boost::asynchronous::then(
                determine_cluster<image_type>(m_image, std::move(m_centers)),
                [image = m_image, k = m_k](auto cont_res)
                {
                    return calculate_cluster_means<image_type>(image, std::move(cont_res.get()), k);
                }
            )
        );
    }

private:
    std::shared_ptr<image_type> m_image;

    std::vector<point> m_centers;

    std::size_t m_k;

    std::size_t m_iteration;
    std::size_t m_max_iterations;

    std::uint8_t m_eps;
};

template<class image_type>
struct k_means_task : public boost::asynchronous::continuation_task<image_type>
{
    k_means_task(image_type image, std::size_t k, std::size_t max_iterations, std::uint8_t eps)
        : boost::asynchronous::continuation_task<image_type>("k_means_task")
        , m_image(std::make_shared<image_type>(std::move(image)))
        , m_k(k)
        , m_max_iterations(max_iterations)
        , m_eps(eps)
    {}

    void operator()()
    {
        // create K random numbers
        std::random_device rd;
        std::mt19937 gen(rd());
        std::uniform_int_distribution<> distribution(0, 255);

        std::vector<point> centers;
        centers.reserve(m_k);

        for (std::size_t i = 0; i < m_k; ++i)
        {
            centers.emplace_back(distribution(gen), distribution(gen), distribution(gen));
        }

        boost::asynchronous::create_callback_continuation(
            [result = this->this_task_result()](auto cont_res)
            {
                try
                {
                    result.set_value(std::move(std::get<0>(cont_res).get()));
                }
                catch (...)
                {
                    result.set_exception(std::current_exception());
                }
            },
            k_means_iteration_task(m_image, std::move(centers), m_k, 0, m_max_iterations, m_eps)
        );
    }

private:
    std::shared_ptr<image_type> m_image;

    std::size_t m_k;

    std::size_t m_max_iterations;

    std::uint8_t m_eps;
};

}

namespace cvpg::imageproc::algorithms {

boost::asynchronous::detail::callback_continuation<image_gray_8bit> k_means(image_gray_8bit image, std::size_t k, std::size_t max_iterations, std::uint8_t eps)
{
    return boost::asynchronous::top_level_callback_continuation<image_gray_8bit>(
               k_means_task(std::move(image), k, max_iterations, eps)
           );
}

boost::asynchronous::detail::callback_continuation<image_rgb_8bit> k_means(image_rgb_8bit image, std::size_t k, std::size_t max_iterations, std::uint8_t eps)
{
    return boost::asynchronous::top_level_callback_continuation<image_rgb_8bit>(
               k_means_task(std::move(image), k, max_iterations, eps)
           );
}

} // namespace cvpg::imageproc::algoritms


// Copyright (C) 2008-2011 Anders Logg and Ola Skavhaug
//
// This file is part of DOLFIN.
//
// DOLFIN is free software: you can redistribute it and/or modify
// it under the terms of the GNU Lesser General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// DOLFIN 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 Lesser General Public License for more details.
//
// You should have received a copy of the GNU Lesser General Public License
// along with DOLFIN. If not, see <http://www.gnu.org/licenses/>.
//
// Modified by Niclas Jansson 2009.
// Modified by Garth N. Wells 2010.
//
// First added:  2008-08-12
// Last changed: 2011-11-14

#include <ufc.h>
#include <boost/random.hpp>
#include <boost/unordered_map.hpp>

#include <dolfin/common/Timer.h>
#include <dolfin/log/log.h>
#include <dolfin/mesh/BoundaryMesh.h>
#include <dolfin/mesh/Facet.h>
#include <dolfin/mesh/Mesh.h>
#include <dolfin/mesh/Vertex.h>
#include "DofMap.h"
#include "UFCCell.h"
#include "UFCMesh.h"
#include "DofMapBuilder.h"

using namespace dolfin;

//-----------------------------------------------------------------------------
void DofMapBuilder::build(DofMap& dofmap, const Mesh& dolfin_mesh,
                          const UFCMesh& ufc_mesh, bool distributed)
{
  // Start timer for dofmap initialization
  Timer t0("Init dofmap");

  // Create space for dof map
  dofmap._dofmap.resize(dolfin_mesh.num_cells());

  dofmap._off_process_owner.clear();

  dolfin_assert(dofmap._ufc_dofmap);

  std::cout << "DofMapBuilder: initialized" << std::endl;

  // Build dofmap from ufc::dofmap
  dolfin::UFCCell ufc_cell(dolfin_mesh);
  std::cout << "Processing cells: ";
  for (dolfin::CellIterator cell(dolfin_mesh); !cell.end(); ++cell)
  {
    std::cout << cell->index() << ",";
    // Update UFC cell
    ufc_cell.update(*cell);

    // Get standard local dimension
    const unsigned int local_dim = dofmap._ufc_dofmap->local_dimension(ufc_cell);
    dofmap._dofmap[cell->index()].resize(local_dim);

    // Tabulate standard UFC dof map
    dofmap._ufc_dofmap->tabulate_dofs(&dofmap._dofmap[cell->index()][0],
                                      ufc_mesh, ufc_cell);
  }
  std::cout << std::endl;
  std::cout << "DofMapBuilder: dofmap built" << std::endl;

  // Build (renumber) dofmap when running in parallel
  if (distributed)
  {
    // Build set of global dofs
    const set global_dofs = compute_global_dofs(dofmap, dolfin_mesh);

    // Build distributed dof map
    build_distributed(dofmap, global_dofs, dolfin_mesh);
  }
  else
    dofmap._ownership_range = std::make_pair(0, dofmap.global_dimension());
}
//-----------------------------------------------------------------------------
void DofMapBuilder::build_distributed(DofMap& dofmap,
                                      const DofMapBuilder::set& global_dofs,
                                      const Mesh& mesh)
{
  // Create data structures
  DofMapBuilder::set owned_dofs, shared_owned_dofs, shared_unowned_dofs;

  // Computed owned and shared dofs (and owned and un-owned)
  compute_ownership(owned_dofs, shared_owned_dofs, shared_unowned_dofs,
                    dofmap, global_dofs, mesh);

  // Renumber dofs owned dofs and received new numbering for unowned shared dofs
  parallel_renumber(owned_dofs, shared_owned_dofs, shared_unowned_dofs,
                    dofmap, mesh);
}
//-----------------------------------------------------------------------------
void DofMapBuilder::compute_ownership(set& owned_dofs, set& shared_owned_dofs,
                                      set& shared_unowned_dofs,
                                      const DofMap& dofmap,
                                      const DofMapBuilder::set& global_dofs,
                                      const Mesh& mesh)
{
  log(TRACE, "Determining dof ownership for parallel dof map");

  // Create a radom number generator for ownership 'voting'
  boost::mt19937 engine(MPI::process_number());
  boost::uniform_int<> distribution(0, 100000000);
  boost::variate_generator<boost::mt19937&, boost::uniform_int<> > rng(engine, distribution);

  // Clear data structures
  owned_dofs.clear();
  shared_owned_dofs.clear();
  shared_unowned_dofs.clear();

  // Data structures for computing ownership
  boost::unordered_map<uint, uint> dof_vote;
  std::vector<uint> facet_dofs(dofmap.num_facet_dofs());

  // Communication buffer
  std::vector<uint> send_buffer;

  // Extract the interior boundary
  BoundaryMesh interior_boundary;
  interior_boundary.init_interior_boundary(mesh);

  // Build set of dofs on process boundary (assume all are owned by this process)
  const MeshFunction<unsigned int>& cell_map = interior_boundary.cell_map();
  if (cell_map.size() > 0)
  {
    for (CellIterator bc(interior_boundary); !bc.end(); ++bc)
    {
      // Get boundary facet
      Facet f(mesh, cell_map[*bc]);

      // Get cell to which facet belongs (pick first)
      Cell c(mesh, f.entities(mesh.topology().dim())[0]);

      // Tabulate dofs on cell
      const std::vector<uint>& cell_dofs = dofmap.cell_dofs(c.index());

      // Tabulate which dofs are on the facet
      dofmap.tabulate_facet_dofs(&facet_dofs[0], c.index(f));

      // Insert shared dofs into set and assign a 'vote'
      for (uint i = 0; i < dofmap.num_facet_dofs(); i++)
      {
        if (shared_owned_dofs.find(cell_dofs[facet_dofs[i]]) == shared_owned_dofs.end())
        {
          shared_owned_dofs.insert(cell_dofs[facet_dofs[i]]);
          dof_vote[cell_dofs[facet_dofs[i]]] = rng();

          send_buffer.push_back(cell_dofs[facet_dofs[i]]);
          send_buffer.push_back(dof_vote[cell_dofs[facet_dofs[i]]]);
        }
      }
    }
  }

  // Decide ownership of shared dofs
  const uint num_proc = MPI::num_processes();
  const uint proc_num = MPI::process_number();
  //const uint max_recv = MPI::max(send_buffer.size());
  std::vector<uint> recv_buffer;
  for (uint k = 1; k < MPI::num_processes(); ++k)
  {
    const uint src  = (proc_num - k + num_proc) % num_proc;
    const uint dest = (proc_num + k) % num_proc;
    MPI::send_recv(send_buffer, dest, recv_buffer, src);

    for (uint i = 0; i < recv_buffer.size(); i += 2)
    {
      const uint received_dof  = recv_buffer[i];
      const uint received_vote = recv_buffer[i + 1];
      if (shared_owned_dofs.find(received_dof) != shared_owned_dofs.end())
      {
        // Move dofs with higher ownership votes from shared to shared but not owned
        if (received_vote < dof_vote[received_dof])
        {
          shared_unowned_dofs.insert(received_dof);
          shared_owned_dofs.erase(received_dof);
        }
        else if (received_vote == dof_vote[received_dof])
        {
          // FIXME: Eventually replace this with a more robust condition. It's
          // good for testing that ownership of shared dofs is spread roughly
          // equally
          dolfin_error("DofMapBuilder.cpp",
                       "compute mapping of degrees of freedom",
                       "Cannot decide on dof ownership; votes are equal");
        }
      }
    }
  }

  // Add/remove global dofs to relevant sets (process 0 owns local dofs)
  if (MPI::process_number() == 0)
  {
    shared_owned_dofs.insert(global_dofs.begin(), global_dofs.end());
    for (set::const_iterator dof = global_dofs.begin(); dof != global_dofs.begin(); ++dof)
    {
      set::const_iterator _dof = shared_unowned_dofs.find(*dof);
      if (_dof != shared_unowned_dofs.end())
        shared_unowned_dofs.erase(_dof);
    }
  }
  else
  {
    shared_unowned_dofs.insert(global_dofs.begin(), global_dofs.end());
    for (set::const_iterator dof = global_dofs.begin(); dof != global_dofs.begin(); ++dof)
    {
      set::const_iterator _dof = shared_owned_dofs.find(*dof);
      if (_dof != shared_owned_dofs.end())
        shared_owned_dofs.erase(_dof);
    }
  }

  // Mark all shared and owned dofs as owned by the processes
  for (CellIterator cell(mesh); !cell.end(); ++cell)
  {
    const std::vector<uint>& cell_dofs = dofmap.cell_dofs(cell->index());
    const uint cell_dimension = dofmap.cell_dimension(cell->index());
    for (uint i = 0; i < cell_dimension; ++i)
    {
      // Mark dof as owned if in unowned set
      if (shared_unowned_dofs.find(cell_dofs[i]) == shared_unowned_dofs.end())
        owned_dofs.insert(cell_dofs[i]);
    }
  }

  // Check that sum of locally owned dofs is equal to global dimension
  const uint _owned_dim = owned_dofs.size();
  dolfin_assert(MPI::sum(_owned_dim) == dofmap.global_dimension());

  log(TRACE, "Finished determining dof ownership for parallel dof map");
}
//-----------------------------------------------------------------------------
void DofMapBuilder::parallel_renumber(const set& owned_dofs,
                             const set& shared_owned_dofs,
                             const set& shared_unowned_dofs,
                             DofMap& dofmap, const Mesh& mesh)
{
  log(TRACE, "Renumber dofs for parallel dof map");

  // FIXME: Handle double-renumbered dof map
  if (dofmap.ufc_map_to_dofmap.size() > 0)
  {
    dolfin_error("DofMapBuilder.cpp",
                 "compute parallel renumbering of degrees of freedom",
                 "The degree of freedom mapping cannot (yet) be renumbered twice");
  }

  const std::vector<std::vector<uint> >& old_dofmap = dofmap._dofmap;
  std::vector<std::vector<uint> > new_dofmap(old_dofmap.size());
  dolfin_assert(old_dofmap.size() == mesh.num_cells());

  // Compute offset for owned and non-shared dofs
  const uint process_offset = MPI::global_offset(owned_dofs.size(), true);

  // Clear some data
  dofmap._off_process_owner.clear();

  // Map from old to new index for dofs
  boost::unordered_map<uint, uint> old_to_new_dof_index;

  // Renumber owned dofs and buffer dofs that are owned but shared with another
  // process
  uint counter = 0;
  std::vector<uint> send_buffer;
  for (set_iterator owned_dof = owned_dofs.begin(); owned_dof != owned_dofs.end(); ++owned_dof, counter++)
  {
    // Set new dof number
    old_to_new_dof_index[*owned_dof] = process_offset + counter;

    // Update UFC-to-renumbered map for new number
    dofmap.ufc_map_to_dofmap[*owned_dof] = process_offset + counter;

    // If this dof is shared and owned, buffer old and new index for sending
    if (shared_owned_dofs.find(*owned_dof) != shared_owned_dofs.end())
    {
      send_buffer.push_back(*owned_dof);
      send_buffer.push_back(process_offset + counter);
    }
  }

  // Exchange new dof numbers for dofs that are shared
  const uint num_proc = MPI::num_processes();
  const uint proc_num = MPI::process_number();
  //const uint max_recv = MPI::max(send_buffer.size());
  std::vector<uint> recv_buffer;
  for (uint k = 1; k < MPI::num_processes(); ++k)
  {
    const uint src  = (proc_num - k + num_proc) % num_proc;
    const uint dest = (proc_num + k) % num_proc;
    MPI::send_recv(send_buffer, dest, recv_buffer, src);

    // Add dofs renumbered by another process to the old-to-new map
    for (uint i = 0; i < recv_buffer.size(); i += 2)
    {
      const uint received_old_dof_index = recv_buffer[i];
      const uint received_new_dof_index = recv_buffer[i + 1];

      // Check if this process has shared dof (and is not the owner)
      if (shared_unowned_dofs.find(received_old_dof_index) != shared_unowned_dofs.end())
      {
        // Add to old-to-new dof map
        old_to_new_dof_index[received_old_dof_index] = received_new_dof_index;

        // Store map from off-process dof to owner
        dofmap._off_process_owner[received_new_dof_index] = src;

        // Update UFC-to-renumbered map
        dofmap.ufc_map_to_dofmap[received_old_dof_index] = received_new_dof_index;
      }
    }
  }

  // Build new dof map
  for (CellIterator cell(mesh); !cell.end(); ++cell)
  {
    const uint cell_index = cell->index();
    const uint cell_dimension = dofmap.cell_dimension(cell_index);

    // Resize cell map and insert dofs
    new_dofmap[cell_index].resize(cell_dimension);
    for (uint i = 0; i < cell_dimension; ++i)
    {
      const uint old_index = old_dofmap[cell_index][i];
      new_dofmap[cell_index][i] = old_to_new_dof_index[old_index];
    }
  }

  // Set new dof map
  dofmap._dofmap = new_dofmap;

  // Set ownership range
  dofmap._ownership_range = std::make_pair<uint, uint>(process_offset, process_offset + owned_dofs.size());

  log(TRACE, "Finished renumbering dofs for parallel dof map");
}
//-----------------------------------------------------------------------------
DofMapBuilder::set DofMapBuilder::compute_global_dofs(const DofMap& dofmap,
                                                       const Mesh& dolfin_mesh)
{
  // Wrap UFC dof map
  boost::shared_ptr<const ufc::dofmap> _dofmap(dofmap._ufc_dofmap.get(),
                                               NoDeleter());

  // Create UFCMesh
  const UFCMesh ufc_mesh(dolfin_mesh);

  // Compute global dof indices
  uint offset = 0;
  set global_dof_indices;
  compute_global_dofs(global_dof_indices, offset, _dofmap, dolfin_mesh, ufc_mesh);

  return global_dof_indices;
}
//-----------------------------------------------------------------------------
void DofMapBuilder::compute_global_dofs(DofMapBuilder::set& global_dofs,
                            uint& offset,
                            boost::shared_ptr<const ufc::dofmap> dofmap,
                            const Mesh& dolfin_mesh, const UFCMesh& ufc_mesh)
{
  dolfin_assert(dofmap);
  const uint D = dolfin_mesh.topology().dim();

  if (dofmap->num_sub_dofmaps() == 0)
  {
    // Check if dofmap is for global dofs
    bool global_dof = true;
    for (uint d = 0; d <= D; ++d)
    {
      if (dofmap->needs_mesh_entities(d))
      {
        global_dof = false;
        break;
      }
    }

    if (global_dof)
    {
      // Check that we have just one dof
      if (dofmap->global_dimension() != 1)
      {
        dolfin_error("DofMapBuilder.cpp",
                     "compute global degrees of freedom",
                     "Global degree of freedom has dimension != 1");
      }

      boost::scoped_ptr<ufc::mesh> ufc_mesh(new ufc::mesh);
      boost::scoped_ptr<ufc::cell> ufc_cell(new ufc::cell);
      uint dof = 0;
      dofmap->tabulate_dofs(&dof, *ufc_mesh, *ufc_cell);

      // Insert global dof index
      std::pair<DofMapBuilder::set::iterator, bool> ret = global_dofs.insert(dof + offset);
      if (!ret.second)
      {
        dolfin_error("DofMapBuilder.cpp",
                     "compute global degrees of freedom",
                     "Global degree of freedom already exists");
      }
    }
  }
  else
  {
    for (uint i = 0; i < dofmap->num_sub_dofmaps(); ++i)
    {
      // Extract sub-dofmap and intialise
      boost::shared_ptr<ufc::dofmap> sub_dofmap(dofmap->create_sub_dofmap(i));
      DofMap::init_ufc_dofmap(*sub_dofmap, ufc_mesh, dolfin_mesh);

      compute_global_dofs(global_dofs, offset, sub_dofmap, dolfin_mesh,
                          ufc_mesh);

      // Get offset
      if (sub_dofmap->num_sub_dofmaps() == 0)
        offset += sub_dofmap->global_dimension();
    }
  }
}
//-----------------------------------------------------------------------------