// Copyright (c), ETH Zurich and UNC Chapel Hill.
// All rights reserved.
//
// Redistribution and use in source and binary forms, with or without
// modification, are permitted provided that the following conditions are met:
//
//     * Redistributions of source code must retain the above copyright
//       notice, this list of conditions and the following disclaimer.
//
//     * Redistributions in binary form must reproduce the above copyright
//       notice, this list of conditions and the following disclaimer in the
//       documentation and/or other materials provided with the distribution.
//
//     * Neither the name of ETH Zurich and UNC Chapel Hill nor the names of
//       its contributors may be used to endorse or promote products derived
//       from this software without specific prior written permission.
//
// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
// AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
// ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDERS OR CONTRIBUTORS BE
// LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
// CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
// SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
// INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
// CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
// ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
// POSSIBILITY OF SUCH DAMAGE.

#pragma once

#include "colmap/util/logging.h"

#include <unordered_map>
#include <vector>

#if defined(__GNUC__) && !defined(__clang__)
#pragma GCC diagnostic push
#pragma GCC diagnostic ignored "-Wmaybe-uninitialized"
#endif

#include <boost/graph/adjacency_list.hpp>
#include <boost/graph/boykov_kolmogorov_max_flow.hpp>
#include <boost/graph/graph_traits.hpp>
#include <boost/graph/one_bit_color_map.hpp>

#if defined(__GNUC__) && !defined(__clang__)
#pragma GCC diagnostic pop
#endif

namespace colmap {

// Compute the min-cut of a undirected graph using the Stoer Wagner algorithm.
void ComputeMinGraphCutStoerWagner(
    const std::vector<std::pair<int, int>>& edges,
    const std::vector<int>& weights,
    int* cut_weight,
    std::vector<char>* cut_labels);

// Compute the normalized min-cut of an undirected graph using Metis.
// Partitions the graph into clusters and returns the cluster labels per vertex.
std::unordered_map<int, int> ComputeNormalizedMinGraphCut(
    const std::vector<std::pair<int, int>>& edges,
    const std::vector<int>& weights,
    int num_parts);

// Compute the minimum graph cut of a directed S-T graph using the
// Boykov-Kolmogorov max-flow min-cut algorithm, as descibed in:
//   "An Experimental Comparison of Min-Cut/Max-Flow Algorithms for Energy
//    Minimization in Vision". Yuri Boykov and Vladimir Kolmogorov. PAMI, 2004.
template <typename node_t, typename value_t>
class MinSTGraphCut {
 public:
  using graph_traits_t =
      boost::adjacency_list_traits<boost::vecS, boost::vecS, boost::directedS>;
  using edge_descriptor_t = graph_traits_t::edge_descriptor;
  using vertices_size_t = graph_traits_t::vertices_size_type;

  struct Edge {
    value_t capacity;
    value_t residual;
    edge_descriptor_t reverse;
  };

  using graph_t = boost::
      adjacency_list<boost::vecS, boost::vecS, boost::directedS, size_t, Edge>;

  explicit MinSTGraphCut(size_t num_nodes);

  // Count the number of nodes and edges in the graph.
  size_t NumNodes() const;
  size_t NumEdges() const;

  // Add node to the graph.
  void AddNode(node_t node_idx, value_t source_capacity, value_t sink_capacity);

  // Add edge to the graph.
  void AddEdge(node_t node_idx1,
               node_t node_idx2,
               value_t capacity,
               value_t reverse_capacity);

  // Compute the min-cut using the max-flow algorithm. Returns the flow.
  value_t Compute();

  // Check whether node is connected to source or sink after computing the cut.
  bool IsConnectedToSource(node_t node_idx) const;
  bool IsConnectedToSink(node_t node_idx) const;

 private:
  const node_t S_node_;
  const node_t T_node_;
  graph_t graph_;
  std::vector<boost::default_color_type> colors_;
};

////////////////////////////////////////////////////////////////////////////////
// Implementation
////////////////////////////////////////////////////////////////////////////////

template <typename node_t, typename value_t>
MinSTGraphCut<node_t, value_t>::MinSTGraphCut(const size_t num_nodes)
    : S_node_(num_nodes), T_node_(num_nodes + 1), graph_(num_nodes + 2) {}

template <typename node_t, typename value_t>
size_t MinSTGraphCut<node_t, value_t>::NumNodes() const {
  return boost::num_vertices(graph_) - 2;
}

template <typename node_t, typename value_t>
size_t MinSTGraphCut<node_t, value_t>::NumEdges() const {
  return boost::num_edges(graph_);
}

template <typename node_t, typename value_t>
void MinSTGraphCut<node_t, value_t>::AddNode(const node_t node_idx,
                                             const value_t source_capacity,
                                             const value_t sink_capacity) {
  THROW_CHECK_GE(node_idx, 0);
  THROW_CHECK_LE(node_idx, boost::num_vertices(graph_));
  THROW_CHECK_GE(source_capacity, 0);
  THROW_CHECK_GE(sink_capacity, 0);

  if (source_capacity > 0) {
    const edge_descriptor_t edge =
        boost::add_edge(S_node_, node_idx, graph_).first;
    const edge_descriptor_t edge_reverse =
        boost::add_edge(node_idx, S_node_, graph_).first;
    graph_[edge].capacity = source_capacity;
    graph_[edge].reverse = edge_reverse;
    graph_[edge_reverse].reverse = edge;
  }

  if (sink_capacity > 0) {
    const edge_descriptor_t edge =
        boost::add_edge(node_idx, T_node_, graph_).first;
    const edge_descriptor_t edge_reverse =
        boost::add_edge(T_node_, node_idx, graph_).first;
    graph_[edge].capacity = sink_capacity;
    graph_[edge].reverse = edge_reverse;
    graph_[edge_reverse].reverse = edge;
  }
}

template <typename node_t, typename value_t>
void MinSTGraphCut<node_t, value_t>::AddEdge(const node_t node_idx1,
                                             const node_t node_idx2,
                                             const value_t capacity,
                                             const value_t reverse_capacity) {
  THROW_CHECK_GE(node_idx1, 0);
  THROW_CHECK_LE(node_idx1, boost::num_vertices(graph_));
  THROW_CHECK_GE(node_idx2, 0);
  THROW_CHECK_LE(node_idx2, boost::num_vertices(graph_));
  THROW_CHECK_GE(capacity, 0);
  THROW_CHECK_GE(reverse_capacity, 0);

  const edge_descriptor_t edge =
      boost::add_edge(node_idx1, node_idx2, graph_).first;
  const edge_descriptor_t edge_reverse =
      boost::add_edge(node_idx2, node_idx1, graph_).first;
  graph_[edge].capacity = capacity;
  graph_[edge_reverse].capacity = reverse_capacity;
  graph_[edge].reverse = edge_reverse;
  graph_[edge_reverse].reverse = edge;
}

template <typename node_t, typename value_t>
value_t MinSTGraphCut<node_t, value_t>::Compute() {
  const vertices_size_t num_vertices = boost::num_vertices(graph_);

  colors_.resize(num_vertices);
  std::vector<edge_descriptor_t> predecessors(num_vertices);
  std::vector<vertices_size_t> distances(num_vertices);

  return boost::boykov_kolmogorov_max_flow(
      graph_,
      boost::get(&Edge::capacity, graph_),
      boost::get(&Edge::residual, graph_),
      boost::get(&Edge::reverse, graph_),
      predecessors.data(),
      colors_.data(),
      distances.data(),
      boost::get(boost::vertex_index, graph_),
      S_node_,
      T_node_);
}

template <typename node_t, typename value_t>
bool MinSTGraphCut<node_t, value_t>::IsConnectedToSource(
    const node_t node_idx) const {
  return colors_.at(node_idx) != boost::white_color;
}

template <typename node_t, typename value_t>
bool MinSTGraphCut<node_t, value_t>::IsConnectedToSink(
    const node_t node_idx) const {
  return colors_.at(node_idx) == boost::white_color;
}

}  // namespace colmap
