// FordFulkerson.cpp
//
//     Created on: 29. 03. 2016
//         Author: Jan Broz
//    Modified by: Jan Uhlik
//
// 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

#include "FordFulkerson.hpp"

#include <queue>
#include <climits>
#include <registration/AlgoRegistration.hpp>

namespace graph {

namespace maximum_flow {

enum State { FRESH, OPEN, CLOSED };
struct Context {
  Capacity capacity;
  Flow flow;
  ext::unordered_map<node::Node, node::Node> prev;
  ext::unordered_map<node::Node, int> sgn;
  ext::unordered_map<node::Node, int> maxFlowTo;
};

static Capacity getCapacity(const DirectedGraph &graph) {
  Capacity capacity;

  for (const auto &edge : graph.getEdges()) {
    capacity[edge.first][edge.second] = edge.capacity();
  }

  return capacity;
}

static bool findPath_dir(const DirectedGraph &graph,
                         const node::Node &source,
                         const node::Node &sink,
                         Context &ctx) {
  ext::unordered_map<node::Node, State> state;
  std::queue<node::Node> open;
  node::Node actual;

  for (const node::Node &node : graph.getNodes()) {
    state[node] = FRESH;
    ctx.maxFlowTo[node] = 0;
  }
  ctx.maxFlowTo[source] = INT_MAX;
  ctx.prev[source] = source;
  state[source] = OPEN;
  open.push(source);

  do {
    actual = open.front();
    open.pop();
    state[actual] = CLOSED;
    for (const node::Node &succ : graph.successors(actual)) {
      if (state[succ] == FRESH && ctx.flow[actual][succ] < ctx.capacity[actual][succ]) {
        state[succ] = OPEN;
        open.push(succ);
        ctx.prev[succ] = actual;
        ctx.sgn[succ] = 1;
        ctx.maxFlowTo[succ] = std::min(ctx.maxFlowTo[actual], ctx.capacity[actual][succ] - ctx.flow[actual][succ]);
      }
    }
    for (const node::Node &pred : graph.predecessors(actual)) {
      if (state[pred] == FRESH && ctx.flow[pred][actual] > 0) {
        state[pred] = OPEN;
        open.push(pred);
        ctx.prev[pred] = actual;
        ctx.sgn[pred] = -1;
        ctx.maxFlowTo[pred] = std::min(ctx.maxFlowTo[actual], ctx.flow[pred][actual]);
      }
    }
  } while (!open.empty() && actual != sink);

  return actual == sink;
}

static bool findPath_undir(const DirectedGraph &graph,
                           const node::Node &source,
                           const node::Node &sink,
                           Context &ctx) {
  ext::unordered_map<node::Node, State> state;
  std::queue<node::Node> open;
  node::Node actual;

  for (const node::Node &node : graph.getNodes()) {
    state[node] = FRESH;
    ctx.maxFlowTo[node] = 0;
  }
  ctx.maxFlowTo[source] = INT_MAX;
  ctx.prev[source] = source;
  state[source] = OPEN;
  open.push(source);

  do {
    actual = open.front();
    open.pop();
    state[actual] = CLOSED;
    for (const node::Node &succ : graph.successors(actual)) {
      if (state[succ] == FRESH && ctx.flow[actual][succ] < ctx.capacity[actual][succ]) {
        state[succ] = OPEN;
        open.push(succ);
        ctx.prev[succ] = actual;
        ctx.maxFlowTo[succ] = std::min(ctx.maxFlowTo[actual], ctx.capacity[actual][succ] - ctx.flow[actual][succ]);
      }
    }
  } while (!open.empty() && actual != sink);

  return actual == sink;
}

static void updateFlow_dir(const node::Node &source, const node::Node &sink, Context &ctx) {
/* alternative for finding the increase of flow
		int path_flow = INT_MAX;
        for (v=t; v!=s; v=parent[v])
        {
            u = parent[v];
            path_flow = min(path_flow, rGraph[u][v]);
        }
*/
  int pathFlow = ctx.maxFlowTo[sink];
  node::Node actual = sink;
  while (actual != source) {
    node::Node prev = ctx.prev[actual];
    if (ctx.sgn[actual] > 0)
      ctx.flow[prev][actual] += pathFlow;
    else
      ctx.flow[actual][prev] -= pathFlow;
    actual = prev;
  }
}

static void updateFlow_undir(const node::Node &source, const node::Node &sink, Context &ctx) {
  int pathFlow = ctx.maxFlowTo[sink];
  node::Node actual = sink;
  while (actual != source) {
    node::Node prev = ctx.prev[actual];
    ctx.flow[prev][actual] += pathFlow;
    ctx.flow[actual][prev] -= pathFlow;
    actual = prev;
  }
}

static Flow fordfulkerson_impl_dir(const DirectedGraph &graph,
                                   const node::Node &source,
                                   const node::Node &sink) {
  Context ctx;

  ctx.capacity = getCapacity(graph);
  for (const auto &edge : graph.getEdges()) {
    ctx.flow[edge.first][edge.second] = 0;
    ctx.flow[edge.second][edge.first] = 0;
  }

  while (findPath_dir(graph, source, sink, ctx))
    updateFlow_dir(source, sink, ctx);

  // assign negative flow for the reversed pairs of nodes ?
  for (auto& u : ctx.flow)
    for (auto& v : u.second)
      if (ctx.flow[u.first][v.first] != 0)
        ctx.flow[v.first][u.first] = -ctx.flow[u.first][v.first];

  return ctx.flow;
}

static Flow fordfulkerson_impl_undir(const UndirectedGraph &ugraph, const node::Node& source, const node::Node& sink) {
  DirectedGraph graph;

  for (auto &node: ugraph.getNodes()) {
    graph.addNode(node);
  }
  for (auto &edge: ugraph.getEdges()) {
    graph.addEdge(edge.first, edge.second, edge.capacity());
  }

  Context ctx;
  ctx.capacity = getCapacity(graph);
  for (const auto &edge : graph.getEdges()) {
    ctx.flow[edge.first][edge.second] = 0;
  }

  while (findPath_undir(graph, source, sink, ctx))
    updateFlow_undir(source, sink, ctx);

  return ctx.flow;
}

Flow FordFulkerson::findMaximumFlow(const DirectedGraph &graph,
                                    const node::Node &source,
                                    const node::Node &sink) {
  return fordfulkerson_impl_dir(graph, source, sink);
}

Flow FordFulkerson::findMaximumFlow(const UndirectedGraph &graph,
                                    const node::Node &source,
                                    const node::Node &sink) {
  return fordfulkerson_impl_undir(graph, source, sink);
}

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

} // namespace maximum_flow

} // namespace graph

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

namespace {
auto FordFulkersonUndirected = registration::AbstractRegister<graph::maximum_flow::FordFulkerson,
                                                              graph::maximum_flow::Flow,
                                                              const graph::maximum_flow::UndirectedGraph &,
                                                              const node::Node &,
                                                              const node::Node &>(
    graph::maximum_flow::FordFulkerson::findMaximumFlow);

auto FordFulkersonDirected = registration::AbstractRegister<graph::maximum_flow::FordFulkerson,
                                                            graph::maximum_flow::Flow,
                                                            const graph::maximum_flow::DirectedGraph &,
                                                            const node::Node &,
                                                            const node::Node &>(
    graph::maximum_flow::FordFulkerson::findMaximumFlow);
}
