// AStar.hpp
//
//     Created on: 29. 1. 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_ASTAR_HPP
#define ALIB2_ASTAR_HPP

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

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

namespace graph {

namespace shortest_path {

class AStar {
// ---------------------------------------------------------------------------------------------------------------------

 public:

  /// Find the shortest path using AStar 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).
  ///
  /// The heuristic function must be admissible and monotone.
  ///
  /// \param graph to explore.
  /// \param start initial node.
  /// \param goal final node.
  /// \param f_heuristic heuristic function which accept node and return edge_type::weight_type.
  /// \param f_user function which is called for every opened 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 an edge with a negative weight.
  ///
  template<
      typename TGraph,
      typename TNode,
      typename F1 = std::function<typename TGraph::edge_type::weight_type(const TNode &)>,
      typename F2 = 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,
           F1 f_heuristic,
           F2 f_user = [](const TNode &,
                          const typename TGraph::edge_type::weight_type &) {});

  template<
      typename TGraph,
      typename TNode,
      typename F1 = std::function<typename TGraph::edge_type::weight_type(const TNode &, const TNode &)>
  >
  static
  ext::pair<ext::vector<TNode>, typename TGraph::edge_type::weight_type>
  findPathRegistration(const TGraph &graph,
                       const TNode &start,
                       const TNode &goal,
                       F1 f_heuristic) {
    return findPath(graph, start, goal, [&](const TNode &n) { return f_heuristic(goal, n); });
  }

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

  /// Find the shortest path using AStar algorithm from the \p start node to the \p goal node in the \p graph.
  /// This algorithm is run in both direction, from \p start and also from \p goal.
  ///
  /// Whenever node is opened, \p f_user is called with two parameters (the opened node and value of currently shortest path).
  ///
  /// The heuristic function must be admissible and monotone.
  ///
  /// \param graph to explore.
  /// \param start initial node.
  /// \param goal final node.
  /// \param f_heuristic_forward front-to-end (node->goal) heuristic function which accept node and return edge_type::weight_type.
  /// \param f_heuristic_backward front-to-end (node->start) heuristic function which accept node and return edge_type::weight_type.
  /// \param f_user function which is called for every opened 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 an edge with a negative weight.
  ///
  template<
      typename TGraph,
      typename TNode,
      typename F1 = std::function<typename TGraph::edge_type::weight_type(const TNode &)>,
      typename F2 = std::function<typename TGraph::edge_type::weight_type(const TNode &)>,
      typename F3 = std::function<void(const TNode &, const typename TGraph::edge_type::weight_type &)>>
  static
  ext::pair<ext::vector<TNode>, typename TGraph::edge_type::weight_type>
  findPathBidirectional(const TGraph &graph,
                        const TNode &start,
                        const TNode &goal,
                        F1 f_heuristic_forward,
                        F2 f_heuristic_backward,
                        F3 f_user = [](const TNode &,
                                       const typename TGraph::edge_type::weight_type &) {});

  template<
      typename TGraph,
      typename TNode,
      typename F1 = std::function<typename TGraph::edge_type::weight_type(const TNode &, const TNode &)>
  >
  static
  ext::pair<ext::vector<TNode>, typename TGraph::edge_type::weight_type>
  findPathBidirectionalRegistration(const TGraph &graph,
                                    const TNode &start,
                                    const TNode &goal,
                                    F1 f_heuristic) {
    return findPathBidirectional(graph, start, goal,
                                 [&](const TNode &n) { return f_heuristic(goal, n); },
                                 [&](const TNode &n) { return f_heuristic(start, n); });
  }


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

 private:

  template<typename TNode, typename TWeight>
  struct Data {
    ext::set<ext::pair<TWeight, TNode>> queue; // priority queue
    ext::map<TNode, TWeight> g; // G score
    ext::map<TNode, TWeight> f; // F score
    ext::map<TNode, TNode> p; // parents
  };

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

  template<typename TGraph, typename TNode, typename F1, typename F2, typename F3>
  static
  ext::pair<ext::vector<TNode>, typename TGraph::edge_type::weight_type>
  impl(const TGraph &graph,
       const TNode &start,
       const TNode &goal,
       F1 f_heuristic,
       F2 f_user,
       F3 f_stop);


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

  template<typename FSuccEdge, typename TNode, typename TWeight, typename F1, typename F2, typename F3, typename F4>
  static bool relaxation(FSuccEdge successor_edges,
                         Data<TNode, TWeight> &data,
                         F1 f_heuristic,
                         F2 f_user,
                         F3 f_stop,
                         F4 f_update);

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

  template<typename TGraph, typename TNode, typename F1, typename F2, typename F3, typename F4>
  static
  ext::pair<ext::vector<TNode>, typename TGraph::edge_type::weight_type>
  implBidirectional(const TGraph &graph,
                    const TNode &start,
                    const TNode &goal,
                    F1 f_heuristic_forward,
                    F2 f_heuristic_backward,
                    F3 f_user,
                    F4 f_stop);

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

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

// ---------------------------------------------------------------------------------------------------------------------
};

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

template<typename TGraph, typename TNode, typename F1, typename F2>
ext::pair<ext::vector<TNode>, typename TGraph::edge_type::weight_type>
AStar::findPath(const TGraph &graph,
                const TNode &start,
                const TNode &goal,
                F1 f_heuristic,
                F2 f_user) {
  // We need to change user function to return false for every node in order to have only one relaxation function
  return impl(graph, start, goal,
              f_heuristic,
              f_user,
              [&goal](const TNode &n) -> bool { return goal == n; });
}

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

template<typename TGraph, typename TNode, typename F1, typename F2, typename F3>
ext::pair<ext::vector<TNode>, typename TGraph::edge_type::weight_type>
AStar::findPathBidirectional(const TGraph &graph,
                             const TNode &start,
                             const TNode &goal,
                             F1 f_heuristic_forward,
                             F2 f_heuristic_backward,
                             F3 f_user) {
  // We need to change user function to return false for every node in order to have only one relaxation function
  return implBidirectional(graph,
                           start,
                           goal,
                           f_heuristic_forward,
                           f_heuristic_backward,
                           f_user,
                           [](const TNode &) -> bool { return false; });
}

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

template<typename TGraph, typename TNode, typename F1, typename F2, typename F3>
ext::pair<ext::vector<TNode>, typename TGraph::edge_type::weight_type>
AStar::impl(const TGraph &graph,
            const TNode &start,
            const TNode &goal,
            F1 f_heuristic,
            F2 f_user,
            F3 f_stop) {
  using weight_type = typename TGraph::edge_type::weight_type;

  Data<TNode, weight_type> data;

  // Init search
  init(data, start, f_heuristic);

  while (!data.queue.empty()) {
    bool stop = relaxation([&](const auto &node) -> auto { return graph.successorEdges(node); },
                           data,
                           f_heuristic,
                           f_user,
                           f_stop,
                           [](const TNode &) -> void {});

    if (stop) {
      break;
    }
  }

  return common::ReconstructPath::reconstructWeightedPath(data.p, data.g, start, goal);
}

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

template<typename FSuccEdge, typename TNode, typename TWeight, typename F1, typename F2, typename F3, typename F4>
bool AStar::relaxation(FSuccEdge successor_edges,
                       Data<TNode, TWeight> &data,
                       F1 f_heuristic,
                       F2 f_user,
                       F3 f_stop,
                       F4 f_update) {
  TNode n = data.queue.begin()->second;
  data.queue.erase(data.queue.begin());

  // Run user's function
  f_user(n, data.g[n]);

  // Stop if reach the goal
  if (f_stop(n)) {
    return true;
  }

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

    // Check for negative edge
    if (s_edge.weight() < 0) {
      throw std::out_of_range("AStar: Detect negative weight on edge in graph.");
    }

    // Calculate new G score
    TWeight gscore = data.g.at(n) + s_edge.weight();

    // Search if the node s was already visited
    auto search_g = data.g.find(s);

    // If not or the G score can be improve do relaxation
    if (search_g == data.g.end() || data.g.at(s) > gscore) {
      // Search if the node s is in OPEN
      auto search_q = data.queue.find(ext::make_pair(data.f[s], s));
      if (search_q != data.queue.end()) {
        // Erase node from priority queue
        data.queue.erase(search_q);
      }

      data.g[s] = gscore;
      data.f[s] = gscore + f_heuristic(s);
      data.p.insert_or_assign(s, n);
      data.queue.insert(ext::make_pair(data.f[s], s));

      f_update(s); // Update currently best path
    }
  }

  return false;
}

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

template<typename TGraph, typename TNode, typename F1, typename F2, typename F3, typename F4>
ext::pair<ext::vector<TNode>, typename TGraph::edge_type::weight_type>
AStar::implBidirectional(const TGraph &graph,
                         const TNode &start,
                         const TNode &goal,
                         F1 f_heuristic_forward,
                         F2 f_heuristic_backward,
                         F3 f_user,
                         F4 f_stop) {
  using TWeight = typename TGraph::edge_type::weight_type;

  TWeight p = std::numeric_limits<TWeight>::max(); // Currently best path weight
  ext::vector<TNode> intersection_nodes; // Last one is currently best intersection node
  Data<TNode, TWeight> forward_data;
  Data<TNode, TWeight> backward_data;

  // Init forward search
  init(forward_data, start, f_heuristic_forward);
  auto f_forward_update = [&](const auto &s) -> void {
    if (backward_data.g.find(s) != backward_data.g.end()) {
      if (forward_data.g.at(s) + backward_data.g.at(s) < p) {
        p = forward_data.g.at(s) + backward_data.g.at(s);
        intersection_nodes.push_back(s);
      }
    }
  };

  // Init backward search
  init(backward_data, goal, f_heuristic_backward);
  auto f_backward_update = [&](const auto &s) -> void {
    if (forward_data.g.find(s) != forward_data.g.end()) {
      if (backward_data.g.at(s) + forward_data.g.at(s) < p) {
        p = backward_data.g.at(s) + forward_data.g.at(s);
        intersection_nodes.push_back(s);
      }
    }
  };

  while (!forward_data.queue.empty() && !backward_data.queue.empty()) {
    if (std::max(forward_data.queue.begin()->first, backward_data.queue.begin()->first) >= p) {
      return common::ReconstructPath::reconstructWeightedPath(forward_data.p,
                                                              backward_data.p,
                                                              forward_data.g,
                                                              backward_data.g,
                                                              start,
                                                              goal,
                                                              intersection_nodes.back());
    }

    // Expand the lower value
    if (forward_data.queue.begin()->first < backward_data.queue.begin()->first) {
      // Forward search relaxationBidirectional
      relaxation([&](const auto &node) -> auto { return graph.successorEdges(node); },
                 forward_data,
                 f_heuristic_forward,
                 f_user,
                 f_stop,
                 f_forward_update);
    } else {
      // Backward search relaxationBidirectional
      relaxation([&](const auto &node) -> auto { return graph.predecessorEdges(node); },
                 backward_data,
                 f_heuristic_backward,
                 f_user,
                 f_stop,
                 f_backward_update);
    }
  }

  return ext::make_pair(ext::vector<TNode>(), p);
}

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

template<typename TNode, typename TWeight, typename F>
void AStar::init(AStar::Data<TNode, TWeight> &data, const TNode &start, F f_heuristic) {
  data.g[start] = 0;
  data.f[start] = data.g[start] + f_heuristic(start);
  data.p.insert_or_assign(start, start);
  data.queue.insert(ext::make_pair(data.f[start], start));
}

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

} // namespace shortest_path

} // namespace graph

#endif //ALIB2_ASTAR_HPP
