import networkx as nx
import matplotlib.pyplot as plt
import numpy as np

# Figure2のグラフを生成
G = nx.DiGraph()
G.add_edges_from([('s','v',{'capacity': 1}),
                  ('s','w',{'capacity': 100}),
                  ('w','v',{'capacity': 1}),
                  ('v','t',{'capacity': 100}),
                  ('w','t',{'capacity': 1})])

#前回作ったやつ
import networkx as nx
import matplotlib.pyplot as plt
from collections import deque
import numpy as np

def bfsFlowPath(G, s, e):
    '''
    Search a path from s to t all of whose points have strictly positive capacity.
    
    Inputs:
        G: a graph
        s: a start point
        e: an end point
        each edge has two attributes:
            capacity: capacity
            flow: its current flow which should be no more than its capacity
    Output:
        path: a list of edges which represents a path from s to e.
        At each edge of the path, its current flow is strictly less than its capacity.
        In case there is no path from s to t, return None.
    '''

    # 過去に訪れた点を記録
    # sは最初から入れておく
    past = [s]
    # pathを記録するためのリスト
    path = []

    # 全ての点のsからの距離の初期値を無限大に
    for p in G.nodes():
        G.node[p]['dist'] = float('inf')
    
    # node s の距離を0に
    G.node[s]['dist'] = 0
    
    # sに隣接する点をqueueに
    # queueには、今後訪れるべき点が格納される
    queue = deque()
    for p in G.successors(s):
        # current flow < capacity となるedgeだけをpathの候補に
        # flow < capacity となるedge以外は存在しないものとして扱うのと同じ
        if G[s][p]['flow'] < G[s][p]['capacity']:
            queue.append(p)
            
    # あとで条件分岐に用いる
    numberOfSuccessorsOfSource = len(queue)
    
    # sに隣接する点の距離を1に
    for p in queue:
        G.node[p]['dist'] = 1

    # BFSを用いてflow < capacityを満たすsからeへのpathがあるのか調べる
    # pastに過去に訪れた点を格納
    while len(queue)>0:
        v = queue.popleft()
        if v == e: break
        else:
            past.append(v)
            for p in G.successors(v):
                # (過去に訪れていない and flow < capacity)を満たすedge
                if ( (not p in past) and ( G[v][p]['flow'] < G[v][p]['capacity']) ):
                    if ( not p in queue):
                        queue.append(p)
                    if G.node[p]['dist'] > G.node[v]['dist'] + 1:
                        G.node[p]['dist'] = G.node[v]['dist'] + 1

    # sからeへのpathが存在しない場合はNoneを返す
    if numberOfSuccessorsOfSource == 0:
        v = s
    if v != e or v == s:
        # print ('There is no path.')
        return None
    
    # 以下、sからeへのpathが存在する場合
    # 終点から遡ってpathを形成する
    pp = e
    while (1):
        if pp == s: break
            
        pred = G.predecessors(pp)
        count = 0

        for p in pred:
            # ここに、flow < capacity の条件を追加
            if ( G.node[p]['dist'] == G.node[pp]['dist']-1 and G[p][pp]['flow'] < G[p][pp]['capacity']):
                path.insert(0, (p,pp))
                pp = p
                break
            else:
                count += 1
        
        # 条件を満たすedgeがない
        if count == len(pred):
            break

    return path

def makeResidualGraph(G):
    '''
    Input: a graph G
    Output: its residual graph Gf
    '''
    Gf = G.copy()
    edgeList = G.edges()
    
    for edge in edgeList:
        # Initialize flow
        Gf[edge[0]][edge[1]]['flow'] = 0
        
        # 逆向きのedgeがないものは追加
        if not (edge[1], edge[0]) in edgeList:
            Gf.add_edge(edge[1],edge[0])
            Gf[edge[1]][edge[0]]['capacity'] = Gf[edge[0]][edge[1]]['flow']
            Gf[edge[1]][edge[0]]['flow'] = 0
    
    return Gf

def fordFulkerson(G, s, t):
    '''
    Inputs:
        G: a graph
        s: a source vertex
        t: a sink vertex
    Outputs:
        the graph G whose flow was modified by Ford-Fulkerson algorithm
        In case there is no path from s to t, return None.
    '''

    # initialize flows
    for e in G.edges():
        G[e[0]][e[1]]['flow'] = 0

    # Forward edgesを記録
    forwardEdges = G.edges()
    
    # Residual Graphの作成
    Gf = makeResidualGraph(G)
    
    # そもそもGにおいてsからtへのパスがあるのか確認
    path = bfsFlowPath(G, s, t)
    if path == None:
        print ("There is no path from " + str(s) + " to "+ str(t) )
        return None
    
    # Gにおいてsからtへのパスがある場合
    while(1):
        # これ以上パスがみつからない場合は最適なのでそこでループを打ち切る
        path = bfsFlowPath(Gf, s, t)
        if path == None:
            break

        # まだパスがある（最適でない）場合
        else:
            # path上のedgeについて、capacity - flowの最小値を調べる
            diff = float('inf')
            for edge in path:
                diff = np.min([diff, Gf[edge[0]][edge[1]]['capacity'] - Gf[edge[0]][edge[1]]['flow'] ])
            
            # path上のedgeのflowを更新
            for edge in path:
                if edge in forwardEdges:
                    Gf[edge[0]][edge[1]]['flow'] += diff
                    # このとき、backward edgeのcapacityを更新する必要あり？
                    Gf[edge[1]][edge[0]]['capacity'] += diff
                else:
                    Gf[edge[0]][edge[1]]['flow'] -= diff
                    Gf[edge[1]][edge[0]]['capacity'] -= diff
    
    # もともと無かったedgeを消去
    for edge in Gf.edges():
        if not edge in forwardEdges:
            Gf.remove_edge(edge[0],edge[1])
    
    return Gf

T = fordFulkerson(G, 's', 't')