// BellmanFord.hpp
//
//     Created on: 05. 02. 2018
//         Author: Jan Uhlik
//    Modified by:
//
// Copyright (c) 2017 Czech Technical University in Prague | Faculty of Information Technology. All rights reserved.
// Git repository: https://gitlab.fit.cvut.cz/algorithms-library-toolkit/automata-library

#ifndef ALIB2_BELLMANFORD_HPP
#define ALIB2_BELLMANFORD_HPP

#include <alib/list>
#include <alib/set>
#include <alib/map>
#include <alib/vector>
#include <queue>
#include <stdexcept>
#include <functional>

#include <common/ReconstructPath.hpp>
#include <common/SupportFunction.hpp>

namespace graph {

namespace shortest_path {

class BellmanFord {
// ---------------------------------------------------------------------------------------------------------------------
 public:

  /// Run BellmanFord algorithm from the \p start node in the \p graph.
  ///
  /// Whenever node is opened, \p f_user is called with two parameters (the opened node and value of currently shortest path).
  ///
  /// \param graph to explore.
  /// \param start initial node.
  /// \param f_user function which is called for every opened node with value of currently shortest path.
  ///
  /// \note TEdge of \p graph must follow graph::edge::WeightedEdge interface.
  /// \sa graph::edge_type::WeightedEdge.
  ///
  /// \throws std::out_of_range if \p graph contains negative cycle.
  ///
  template<
      typename TGraph,
      typename TNode,
      typename F = std::function<void(const TNode &,
                                      const typename TGraph::edge_type::weight_type &)> >
  static
  void
  run(const TGraph &graph,
      const TNode &start,
      F f_user = [](const TNode &,
                    const typename TGraph::edge_type::weight_type &) -> void {});

// ---------------------------------------------------------------------------------------------------------------------

  /// Find the shortest path using BellmanFord algorithm from the \p start node to the \p goal node in the \p graph.
  ///
  /// Whenever node is opened, \p f_user is called with two parameters (the opened node and value of currently shortest path).
  ///
  /// \param graph to explore.
  /// \param start initial node.
  /// \param goal final node.
  /// \param f_user function which is called for every open node with value of currently shortest path.
  ///
  /// \returns pair where first := shortest path := distance of path, if there is no such path vector is empty and distance std::numeric_limits<edge_type:weight_type>::max().
  ///
  /// \note TEdge of \p graph must follow graph::edge::WeightedEdge interface.
  /// \sa graph::edge_type::WeightedEdge.
  ///
  /// \throws std::out_of_range if \p graph contains negative cycle.
  ///
  template<typename TGraph, typename TNode, typename F = std::function<void(const TNode &,
                                                                            const typename TGraph::edge_type::weight_type &)>>
  static
  ext::pair<ext::vector<TNode>, typename TGraph::edge_type::weight_type>
  findPath(const TGraph &graph,
           const TNode &start,
           const TNode &goal,
           F f_user = [](const TNode &,
                         const typename TGraph::edge_type::weight_type &) {});

  template<typename TGraph, typename TNode>
  static
  ext::pair<ext::vector<TNode>, typename TGraph::edge_type::weight_type>
  findPathRegistration(const TGraph &graph,
                       const TNode &start,
                       const TNode &goal) {
    return findPath(graph, start, goal);
  }

// =====================================================================================================================
 private:

  template<typename TNode, typename TWeight>
  struct Data {
    ext::map<TNode, TWeight> g; // distance (aka G score)
    ext::map<TNode, TNode> p; // parents
    ext::set<TNode> state1; // optimization Yen
    ext::set<TNode> state2; // optimization Yen
  };

// ---------------------------------------------------------------------------------------------------------------------

  template<typename TGraph, typename TNode, typename F>
  static
  Data<TNode, typename TGraph::edge_type::weight_type>
  impl(const TGraph &graph,
       const TNode &start,
       F f_user);

// ---------------------------------------------------------------------------------------------------------------------

  template<typename TGraph, typename TNode, typename F>
  static
  void
  relaxation(const TGraph &graph,
             ext::set<TNode> &nodes,
             BellmanFord::Data<TNode, typename TGraph::edge_type::weight_type> &data,
             ext::set<TNode> &state1,
             ext::set<TNode> &state2,
             F f_user);

// ---------------------------------------------------------------------------------------------------------------------

  template<typename TNode, typename TWeight>
  inline static void init(BellmanFord::Data<TNode, TWeight> &data, const TNode &start);

// ---------------------------------------------------------------------------------------------------------------------

};

// =====================================================================================================================

template<typename TGraph, typename TNode, typename F>
void BellmanFord::run(const TGraph &graph, const TNode &start, F f_user) {
  impl(graph, start, f_user);
}

// ---------------------------------------------------------------------------------------------------------------------

template<typename TGraph, typename TNode, typename F>
ext::pair<ext::vector<TNode>, typename TGraph::edge_type::weight_type>
BellmanFord::findPath(const TGraph &graph,
                      const TNode &start,
                      const TNode &goal,
                      F f_user) {
  using weight_type = typename TGraph::edge_type::weight_type;

  Data<TNode, weight_type> data = impl(graph, start, f_user);

  if (data.g.find(goal) == data.g.end()) {
    return ext::make_pair(ext::vector<TNode>(), std::numeric_limits<weight_type>::max());
  }

  return ext::make_pair(common::ReconstructPath::reconstructPath(data.p, start, goal), data.g[goal]);
}

// ---------------------------------------------------------------------------------------------------------------------

template<typename TGraph, typename TNode, typename F>
BellmanFord::Data<TNode, typename TGraph::edge_type::weight_type>
BellmanFord::impl(const TGraph &graph, const TNode &start, F f_user) {

  using weight_type = typename TGraph::edge_type::weight_type;

  Data<TNode, weight_type> data;

  // Init data
  init(data, start);

  // Run user's function
  f_user(start, 0);

  auto nodes = graph.getNodes();
  size_t nodes_cnt = nodes.size();

  for (size_t i = 1; i < nodes_cnt; ++i) {
    if (i % 2 == 1) {
      data.state2.clear();
      relaxation(graph, nodes, data, data.state1, data.state2, f_user);
    } else {
      data.state1.clear();
      relaxation(graph, nodes, data, data.state2, data.state1, f_user);
    }

    if (((i % 2 == 1) && data.state2.empty()) ||
        ((i % 2 == 0) && data.state1.empty()))
      break; // Optimization early stop
  }

  for (const auto &n: nodes) {
    for (const auto &s_edge: graph.successorEdges(n)) {
      const TNode &s = common::SupportFunction::other(s_edge, n); // successor

      auto search = data.g.find(s);

      if (search != data.g.end() && data.g.at(n) + s_edge.weight() < data.g.at(s)) {
        throw std::out_of_range("BellmanFord: Detected negative weight cycle.");
      }
    }
  }

  return data;
}

// ---------------------------------------------------------------------------------------------------------------------

template<typename TGraph, typename TNode, typename F>
void BellmanFord::relaxation(const TGraph &graph,
                             ext::set<TNode> &nodes,
                             BellmanFord::Data<TNode, typename TGraph::edge_type::weight_type> &data,
                             ext::set<TNode> &state1,
                             ext::set<TNode> &state2,
                             F f_user) {
  using weight_type = typename TGraph::edge_type::weight_type;

  for (const auto &n: nodes) {

    if (state1.find(n) == state1.end()) {
      continue;
    }

    for (const auto &s_edge: graph.successorEdges(n)) {
      const TNode &s = common::SupportFunction::other(s_edge, n); // successor

      auto search = data.g.find(s);

      // Relaxation
      if (search == data.g.end() || data.g.at(n) + s_edge.weight() < data.g.at(s)) {
        weight_type gscore = data.g.at(n) + s_edge.weight();

        // Run user's function
        f_user(s, gscore);

        data.g[s] = gscore;
        data.p.insert_or_assign(s, n);

        state1.insert(s); // Yen optimization
        state2.insert(s); // Yen optimization
      }
    }
  }

}

// ---------------------------------------------------------------------------------------------------------------------

template<typename TNode, typename TWeight>
void BellmanFord::init(BellmanFord::Data<TNode, TWeight> &data, const TNode &start) {
  data.g[start] = 0;
  data.p.insert(std::make_pair(start, start));
  data.state1.insert(start);
}

// ---------------------------------------------------------------------------------------------------------------------

// =====================================================================================================================

} // namespace shortest_path

} // namespace graph

#endif //ALIB2_BELLMANFORD_HPP
