import random
from mincost_maxflow import MinCostMaxFlow


def brute_min_cost(n, edges, s, t, maxf):
    INF = float("inf")
    g = [[] for _ in range(n)]
    for (u, v, cap, cost) in edges:
        fi = len(g[u]); ri = len(g[v])
        g[u].append([v, cap, cost, ri])
        g[v].append([u, 0, -cost, fi])
    flow = 0; cost = 0
    while flow < maxf:
        dist = [INF]*n; prevv=[-1]*n; preve=[-1]*n
        dist[s]=0
        for _ in range(n):
            upd=False
            for u in range(n):
                if dist[u]==INF: continue
                for i,e in enumerate(g[u]):
                    v,cap,cc,_=e
                    if cap>0 and dist[u]+cc<dist[v]:
                        dist[v]=dist[u]+cc; prevv[v]=u; preve[v]=i; upd=True
            if not upd: break
        if dist[t]==INF: break
        addf=maxf-flow
        v=t
        while v!=s:
            u=prevv[v]; e=g[u][preve[v]]; addf=min(addf,e[1]); v=u
        v=t
        while v!=s:
            u=prevv[v]; ei=preve[v]; e=g[u][ei]
            e[1]-=addf; g[v][e[3]][1]+=addf; cost+=addf*e[2]; v=u
        flow+=addf
    return flow, cost


def test_single_path():
    m = MinCostMaxFlow(4)
    m.add_edge(0,1,2,3); m.add_edge(1,2,2,4); m.add_edge(2,3,2,1)
    f,c = m.min_cost_flow(0,3,2)
    assert (f,c)==(2, 2*(3+4+1)), (f,c)


def test_two_parallel_paths():
    m = MinCostMaxFlow(4)
    m.add_edge(0,1,1,1); m.add_edge(1,3,1,1)
    m.add_edge(0,2,1,5); m.add_edge(2,3,1,5)
    f,c = m.min_cost_flow(0,3,2)
    assert (f,c)==(2, (1+1)+(5+5)), (f,c)


def test_negative_cost():
    m = MinCostMaxFlow(3)
    m.add_edge(0,1,1,5); m.add_edge(1,2,1,-3); m.add_edge(0,2,1,10)
    f,c = m.min_cost_flow(0,2,1)
    assert (f,c)==(1,2), (f,c)


def test_random_vs_brute():
    random.seed(42)
    for trial in range(400):
        n = random.randint(2,5)
        s=0; t=n-1
        edges=[]
        for u in range(n):
            for v in range(n):
                if u==v: continue
                if random.random()<0.4:
                    edges.append((u,v,random.randint(0,3),random.randint(-3,5)))
        maxf=random.randint(1,6)
        m=MinCostMaxFlow(n)
        for (u,v,cap,cost) in edges:
            m.add_edge(u,v,cap,cost)
        f1,c1=m.min_cost_flow(s,t,maxf)
        f2,c2=brute_min_cost(n,edges,s,t,maxf)
        assert (f1,c1)==(f2,c2), (trial,(f1,c1),(f2,c2))


def test_maxflow_cap():
    m=MinCostMaxFlow(2)
    m.add_edge(0,1,5,2)
    f,c=m.min_cost_flow(0,1,10)
    assert (f,c)==(5,10)


if __name__=="__main__":
    test_single_path(); test_two_parallel_paths(); test_negative_cost()
    test_random_vs_brute(); test_maxflow_cap()
    print("ALL OK")