import random
from dinic_maxflow import max_flow as dinic_flow


def ek_flow(n, edges, s, t):
    # Edmonds-Karp (BFS augmenting paths) reference
    from collections import deque
    cap = [[0]*n for _ in range(n)]
    for u, v, c in edges:
        cap[u][v] += c
    flow = 0
    while True:
        parent = [-1]*n
        parent[s] = s
        q = deque([s])
        while q and parent[t] == -1:
            u = q.popleft()
            for v in range(n):
                if parent[v] == -1 and cap[u][v] > 0:
                    parent[v] = u
                    q.append(v)
        if parent[t] == -1:
            break
        # bottleneck
        aug = 10**18
        v = t
        while v != s:
            u = parent[v]
            aug = min(aug, cap[u][v])
            v = u
        v = t
        while v != s:
            u = parent[v]
            cap[u][v] -= aug
            cap[v][u] += aug
            v = u
        flow += aug
    return flow


def test_simple():
    assert dinic_flow(4, [(0,1,10),(1,3,10),(0,2,10),(2,3,10)], 0, 3) == 20


def test_single():
    assert dinic_flow(2, [(0,1,5)], 0, 1) == 5


def test_zero_cap():
    assert dinic_flow(2, [(0,1,0)], 0, 1) == 0


def test_disconnected():
    assert dinic_flow(4, [(0,1,5)], 0, 3) == 0


def test_bottleneck():
    assert dinic_flow(3, [(0,1,10),(1,2,3)], 0, 2) == 3


def test_multi_edge():
    assert dinic_flow(2, [(0,1,3),(0,1,4)], 0, 1) == 7


def test_random_vs_ek():
    rng = random.Random(42)
    for _ in range(300):
        n = rng.randint(2, 7)
        s, t = 0, n-1
        m = rng.randint(0, 12)
        edges = []
        for _ in range(m):
            u = rng.randrange(n)
            v = rng.randrange(n)
            c = rng.randint(0, 6)
            edges.append((u, v, c))
        a = dinic_flow(n, edges, s, t)
        b = ek_flow(n, edges, s, t)
        assert a == b, (edges, a, b)


def test_self_loop_ignored():
    # self-loops shouldn't inflate flow
    assert dinic_flow(2, [(0,0,100),(0,1,5)], 0, 1) == 5