#include "edmonds.h"
#include <iostream>
#include <stack>
#include <tuple>
#include "graph_attributes.h"

using namespace ED;

namespace Edmonds {

void check_integrity(GraphAttributes const & attrs)
{
   for(NodeId id = 0; id < attrs.num_nodes(); ++id)
   {
      // Check that μ encodes a valid matching
      NodeId matched = attrs.mu[id];
      if(matched != id)
      {
         assert(attrs.mu[matched] == id);
      }

      if (attrs.is_out_of_forest(id))
      {
         assert(attrs.phi[id] == id);
         assert(attrs.rho[id] == id);
      }
      else
      {
         // check for a path to the root, i.e. ρ(node)
         NodeId cur_node = id;
         while(cur_node != attrs.rho[cur_node])
         {
            // If the condition was true, then cur_node is outer, part of a blossom
            // and we want to follow its path
            // therefore, we check that both φ and μ are not the identity on this node
            // and point to vertices that have the same rho
            assert(attrs.mu[cur_node] != cur_node);
            assert(attrs.phi[cur_node] != cur_node);
            assert(attrs.rho[attrs.mu[cur_node]] == attrs.rho[cur_node]);
            assert(attrs.rho[attrs.phi[cur_node]] == attrs.rho[cur_node]);

            // now, walk along the matched edge
            cur_node = attrs.mu[cur_node];

            // now we want to walk along φ, this will again
            // - not be the identity
            // - result in a node that has the same rho
            assert(attrs.phi[cur_node] != cur_node);
            assert(attrs.rho[attrs.phi[cur_node]] == attrs.rho[cur_node]);

            cur_node = attrs.mu[attrs.phi[cur_node]];
         }
      }

      if (not attrs.is_outer(id))
      {
         assert(attrs.rho[id] == id);
      }
   }
}

/**
 * @return List of vertices of the x-r path, where r is the root of the
 * special blossom forest component x belongs to.
 *
 * @note This assumes that the values of μ, φ and ρ represent a special
 * blossom forest on the graph when this method is called.
 * **/
__attribute__((noinline))
std::vector<NodeId> path_to_forest_root(GraphAttributes const & attrs, NodeId id)
{
   std::vector<NodeId> retval;
   retval.push_back(id);
   while (attrs.mu[id] != id)
   {
      id = attrs.mu[id];
      retval.push_back(id);

      // Note that it is guaranteed that this does not produce a loop:
      // We are traversing the path to a root of the forest,
      // but we know that each root is exposed by M, so after traversing
      // the matching edge, we cannot have reached a root.
      id = attrs.phi[id];
      retval.push_back(id);
   }
   return retval;
}

void collect_exposed_vertices(GraphAttributes & attrs, std::stack<NodeId> & container)
{
   std::stack<NodeId>().swap(container);
   for(NodeId id = 0; id < attrs.num_nodes(); id++)
   {
      if (attrs.mu[id] == id)
      {
         container.push(id);
         attrs.scanned[id] = true;
      }
   }
}

__attribute__((noinline))
void augment(GraphAttributes & attrs, std::vector<NodeId> const & x_path, std::vector<NodeId> const & y_path,
             std::stack<NodeId> & outer_unvisited_nodes)
{
   //std::cout << "Augment" << std::endl;
   // Paths are disjoint -> augment
   attrs.mu[x_path.front()] = y_path.front();
   attrs.mu[y_path.front()] = x_path.front();

   // TODO: put this into own method?
   for(size_t i = 1; i < x_path.size(); i += 2)
   {
      attrs.mu[x_path[i]] = x_path[i+1];
      attrs.mu[x_path[i+1]] = x_path[i];
   }
   for(size_t i = 1; i < y_path.size(); i += 2)
   {
      attrs.mu[y_path[i]] = y_path[i+1];
      attrs.mu[y_path[i+1]] = x_path[i+1];
   }
   attrs.reset_forest();
   collect_exposed_vertices(attrs, outer_unvisited_nodes);
}

__attribute__((noinline))
std::tuple<NodeId, size_type, size_type> find_blossom_root_id(GraphAttributes const & attrs, std::vector<NodeId> const & x_path, std::vector<NodeId> const & y_path)
{
   size_t distance_from_x = x_path.size() - 1;
   size_t distance_from_y = y_path.size() - 1;
   while (distance_from_x > 0 and distance_from_y > 0 and \
                  x_path[distance_from_x - 1] == y_path[distance_from_y - 1])
   {
      --distance_from_x;
      --distance_from_y;
   }
   // found first vertex of x_path \cap y_path
   while (attrs.rho[x_path[distance_from_x]] != x_path[distance_from_x])
   {
      ++distance_from_x;
      ++distance_from_y;
   };
   // found first vertex fixed by rho
   return { x_path[distance_from_x], distance_from_x, distance_from_y };
}

__attribute__((noinline))
void update_phialong_blossom_paths(GraphAttributes & attrs, std::vector<NodeId> const & x_path, std::vector<NodeId> const & y_path,
                                    std::tuple<NodeId, size_type, size_type> const & blossom_root)
{
   auto const [blossom_root_id, distance_from_x, distance_from_y] = blossom_root;

   // Update φ along the paths to encode the ear decomposition
   for (size_t i = 1; i <= distance_from_x; i += 2)
   {
      if (attrs.rho[attrs.phi[x_path[i]]] != blossom_root_id)
      {
         attrs.phi[attrs.phi[x_path[i]]] = x_path[i];
      }
   }

   for (size_t i = 1; i <= distance_from_y; i += 2)
   {
      if (attrs.rho[attrs.phi[y_path[i]]] != blossom_root_id)
      {
         attrs.phi[attrs.phi[y_path[i]]] = y_path[i];
      }
   }

   // Link x and y
   if (attrs.rho[x_path.front()] != blossom_root_id)
   {
      attrs.phi[x_path.front()] = y_path.front();
   }
   if (attrs.rho[y_path.front()] != blossom_root_id)
   {
      attrs.phi[y_path.front()] = x_path.front();
   }
}

__attribute__((noinline))
void contract_rho(GraphAttributes & attrs, NodeId blossom_root_id)
{
   // Iterating over whole attrs.
   for (NodeId node_id = 0; node_id < attrs.num_nodes(); ++node_id)
   {
      if (attrs.rho[attrs.rho[node_id]] == blossom_root_id)
      {
         attrs.rho[node_id] = blossom_root_id;
      }
   }
}

__attribute__((noinline))
void update_rho(GraphAttributes & attrs, std::vector<NodeId> const & x_path, std::vector<NodeId> const & y_path,
                std::tuple<NodeId, size_type, size_type> const & blossom_root_description,
                std::stack<NodeId> & outer_unvisited_nodes)
{
   // Update root indices. We have to do this for all vertices v with
   // ρ(v) in the paths from x or y to r
   // We update ρ(v) first for the paths themselves, and then 'contract' ρ
   // by updating ρ(v) to r for all vertices where ρ(ρ(v)) = r
   // Also, while walking along the paths, we can add all vertices (which are now outer)
   // to the stack.
   auto const [blossom_root_id, distance_from_x, distance_from_y] = blossom_root_description;
   for (size_t i = 0; i <= distance_from_x; ++i)
   {
      attrs.rho[x_path[i]] = blossom_root_id;
      if (not attrs.scanned[x_path[i]])
      {
         outer_unvisited_nodes.push(x_path[i]);
      }
   }
   for (size_t i = 0; i <= distance_from_y; ++i)
   {
      attrs.rho[y_path[i]] = blossom_root_id;
      if (not attrs.scanned[y_path[i]])
      {
         outer_unvisited_nodes.push(y_path[i]);
      }
   }
   contract_rho(attrs, blossom_root_id);
}

__attribute__((noinline))
void contract_blossom(GraphAttributes & attrs, std::vector<NodeId> const & x_path, std::vector<NodeId> const & y_path,
                      std::stack<NodeId> & outer_unvisited_nodes)
{
   //std::cout << "Contract blossom" << std::endl;
   std::tuple<NodeId, size_type, size_type> const blossom_root_description = find_blossom_root_id(attrs, x_path, y_path);
   update_phialong_blossom_paths(attrs, x_path, y_path, blossom_root_description);

   check_integrity(attrs);
   update_rho(attrs, x_path, y_path, blossom_root_description, outer_unvisited_nodes);
}

void maximum_matching_from_initial_matching(Graph const & graph, GraphAttributes & attrs)
{
   attrs.reset_forest();
   // Over the course of the algorithm, this will maintain all outer vertices
   // that have not been scanned yet.
   // Note that at the beginning, this is exactly the exposed edges.
   // Throughout the algorithm, we push new ids whenever there are new outer vertices.
   // Also note that we separately mark each node whether it has been collected to this stack.
   // When this stack runs out, then we know that all vertices marked 'scanned' have already been processed,
   // but also all vertices not marked 'scanned' are not outer vertices, so we can in fact terminate.
   std::stack<NodeId> outer_unvisited_nodes;
   collect_exposed_vertices(attrs, outer_unvisited_nodes);
   while(not outer_unvisited_nodes.empty())
   {
      NodeId const id = outer_unvisited_nodes.top();
      outer_unvisited_nodes.pop();
      for(NodeId neighbor_id : graph.node(id).neighbors())
      {
         check_integrity(attrs);
         //std::cout << "Check passed" << std::endl;
         if (attrs.is_out_of_forest(neighbor_id))
         {
            //std::cout << "Grow forest" << std::endl;
            // Grow Forest
            attrs.phi[neighbor_id] = id;
            assert(attrs.mu[neighbor_id] != neighbor_id);
            outer_unvisited_nodes.push(attrs.mu[neighbor_id]);
         }
         else if (attrs.is_outer(neighbor_id) and attrs.rho[id] != attrs.rho[neighbor_id])
         {
            std::vector<NodeId> x_path = path_to_forest_root(attrs, id);
            std::vector<NodeId> y_path = path_to_forest_root(attrs, neighbor_id);

            if (x_path.back() != y_path.back())
            {
               // paths are disjoint -> can augment
               augment(attrs, x_path, y_path, outer_unvisited_nodes);
               break;
            }
            else
            {
               // Paths are not disjoint -> contract the new blossom
               contract_blossom(attrs, x_path, y_path, outer_unvisited_nodes);
            }
         }
      }
      attrs.scanned[id] = true;
   }
};

void find_greedy_matching(Graph const & graph, GraphAttributes & attrs)
{
   attrs.reset_matching();
   for(NodeId node_id = 0; node_id < graph.num_nodes(); ++node_id)
   {
      if (attrs.mu[node_id] == node_id) {
         for(NodeId const neighbor_id : graph.node(node_id).neighbors())
         {
            if(attrs.mu[neighbor_id] == neighbor_id)
            {
               attrs.mu[neighbor_id] = node_id;
               attrs.mu[node_id] = neighbor_id;
               break;
            }
         }
      }
   }
}

Graph maximum_matching(Graph & graph) {
   GraphAttributes attrs(graph.num_nodes());
   attrs.reset_forest();
   find_greedy_matching(graph, attrs);
   check_integrity(attrs);
   maximum_matching_from_initial_matching(graph, attrs);

   ED::Graph matching = ED::Graph(graph.num_nodes());
   for (NodeId id = 0; id < graph.num_nodes(); ++id)
   {
      if (attrs.mu[id] > id)
      {
         matching.add_edge(id, attrs.mu[id]);
      }
   }
   return matching;
}

}