#include "Graph.h"
#include <queue>
#include <map>
#include <algorithm>
#include <functional>
#include <limits>

namespace {

struct FlowEdge {
    int to;
    long long cap;
    int rev;
};

using Residual = std::vector<std::vector<FlowEdge>>;

void addArc(Residual& g, int u, int v, long long cap) {
    FlowEdge a{v, cap, (int)g[v].size()};
    FlowEdge b{u, 0, (int)g[u].size()};
    g[u].push_back(a);
    g[v].push_back(b);
}

Residual buildResidual(int n, const std::vector<std::tuple<int, int, int>>& arcs) {
    Residual g(n + 1);
    for (auto& e : arcs) addArc(g, std::get<0>(e), std::get<1>(e), std::get<2>(e));
    return g;
}

// Collapse parallel arcs by summing residual capacities so the result stays a simple graph.
Graph makeResidualGraph(int n, const Residual& g) {
    std::map<std::pair<int, int>, long long> cap;
    for (int u = 1; u <= n; ++u)
        for (auto& e : g[u])
            if (e.cap > 0) cap[{u, e.to}] += e.cap;

    Graph res(n, true, true, 'E');
    for (auto& pr : cap)
        res.addEdge(pr.first.first, pr.first.second, (int)pr.second);
    return res;
}

} // namespace

Graph Graph::flowFordFulkerson(int source, int sink) const {
    Residual g = buildResidual(n, toEdgeList());
    std::vector<int> parent(n + 1), parentEdge(n + 1);

    while (true) {
        std::fill(parent.begin(), parent.end(), -1);
        parent[source] = source;
        std::queue<int> q;
        q.push(source);
        while (!q.empty()) {
            int u = q.front();
            q.pop();
            for (int i = 0; i < (int)g[u].size(); ++i) {
                const FlowEdge& e = g[u][i];
                if (e.cap > 0 && parent[e.to] == -1) {
                    parent[e.to] = u;
                    parentEdge[e.to] = i;
                    q.push(e.to);
                }
            }
        }
        if (parent[sink] == -1) break;

        long long bottleneck = std::numeric_limits<long long>::max();
        for (int v = sink; v != source; v = parent[v])
            bottleneck = std::min(bottleneck, g[parent[v]][parentEdge[v]].cap);
        for (int v = sink; v != source; v = parent[v]) {
            FlowEdge& e = g[parent[v]][parentEdge[v]];
            e.cap -= bottleneck;
            g[v][e.rev].cap += bottleneck;
        }
    }
    return makeResidualGraph(n, g);
}

Graph Graph::flowDinitz(int source, int sink) const {
    Residual g = buildResidual(n, toEdgeList());
    std::vector<int> level(n + 1), ptr(n + 1);

    auto bfs = [&]() -> bool {
        std::fill(level.begin(), level.end(), -1);
        level[source] = 0;
        std::queue<int> q;
        q.push(source);
        while (!q.empty()) {
            int u = q.front();
            q.pop();
            for (auto& e : g[u])
                if (e.cap > 0 && level[e.to] == -1) {
                    level[e.to] = level[u] + 1;
                    q.push(e.to);
                }
        }
        return level[sink] != -1;
    };

    std::function<long long(int, long long)> dfs = [&](int u, long long pushed) -> long long {
        if (pushed == 0) return 0;
        if (u == sink) return pushed;
        for (int& i = ptr[u]; i < (int)g[u].size(); ++i) {
            FlowEdge& e = g[u][i];
            if (level[e.to] != level[u] + 1 || e.cap <= 0) continue;
            long long pushedFlow = dfs(e.to, std::min(pushed, e.cap));
            if (pushedFlow == 0) continue;
            e.cap -= pushedFlow;
            g[e.to][e.rev].cap += pushedFlow;
            return pushedFlow;
        }
        return 0;
    };

    while (bfs()) {
        std::fill(ptr.begin(), ptr.end(), 0);
        while (dfs(source, std::numeric_limits<long long>::max()) > 0) {
        }
    }
    return makeResidualGraph(n, g);
}