import numpy as np
import copy
import networkx as nx
import matplotlib.pyplot as plt
import random
import time
import math
import multiprocessing as mp
import sys
sys.setrecursionlimit(1000000)

class Graph:
    # 初始化，读取数据
    def __init__(self,file_name):
        file = open(file_name, 'r')
        lines = file.readlines() # 读取文件
        # 构造图
        self.G = nx.Graph()
        for line in lines:
            edge = line.split(' ')
            edge[0] = int(edge[0])
            edge[1] = int(edge[1].split('\n')[0])
            if (edge[0],edge[1]) not in list(self.G.edges) or (edge[1],edge[0]) not in list(self.G.edges):
                self.G.add_edge(edge[0],edge[1],weight=1)
            else:
                self.G[edge[0]][edge[1]]['weight'] += 1

    # 画图
    def show(self,best_part,best_part2):
        plt.figure(figsize=(30, 30))
        pos = nx.spring_layout(self.G)
        nx.draw_networkx_nodes(self.G, pos,nodelist=best_part ,node_size=700,node_color='r')
        nx.draw_networkx_nodes(self.G, pos, nodelist=best_part2, node_size=700, node_color='b')
        nx.draw_networkx_edges(self.G, pos, width=3)
        nx.draw_networkx_labels(self.G, pos, font_size=20, font_family="sans-serif")
        plt.axis("off")
        plt.show()

    # 缩边
    def contract(self,G,room):
        # 随机找一个边
        edges = list(G.edges)
        index = np.random.randint(0, len(edges), 1)[0]
        u, v = edges[index]
        # 合并点，weight代表的是两点之间边的个数
        for w, e in G[v].items():
            if w != u:
                if w not in G[u]:
                    G.add_edge(u, w, weight=e['weight'])
                else:
                    G[u][w]['weight'] += e['weight']
        G.remove_node(v)
        # room用于保存合并点的记录，以此来得到最小割的两部分
        room[u - 1].extend(room[v - 1])
        room[v - 1] = [-1]

    # Stoer–Wagner算法
    def st_mincut(self,G):
        PQ = dict() # 记录权重，也就是边数
        s = None
        t = None
        mincut = float('inf')
        for node in list(G.nodes):
            PQ[str(node)] = 0
        # 循环，直到找到最后两个s，t
        while len(PQ) != 0:
            k = int(max(PQ, key=PQ.get))
            s = t
            t = k
            mincut = PQ[str(k)]
            PQ.pop(str(k))
            # weight代表的是两点之间边的个数
            for w, e in G[t].items():
                if str(w) in PQ.keys():
                    PQ[str(w)] += e['weight']
        return s,t,mincut

    def global_mincut(self,G,room):
        nodes = list(G.nodes)
        # 当只有两个点时，直接输出当前图的最小割
        if len(nodes) == 2:
            # weight代表的是两点之间边的个数
            mincut = G[nodes[0]][nodes[1]]['weight']
            best_part = room[nodes[0] - 1]
            best_part2 = room[nodes[1] - 1]
            return int(mincut),best_part,best_part2
        else:
            # 得到s，t以及割的值
            s,t,mincut1 = self.st_mincut(G)
            # 得到割的两部分
            best_part11 = room[t-1]
            best_part12 = list(set(list(self.G.nodes)).difference(set(best_part11)))
            # 复制图，weight代表的是两点之间边的个数
            new_G = nx.Graph((u, v, {'weight': e.get('weight', 1)})
                             for u, v, e in G.edges(data=True) if u != v)
            # 合并s，t
            for w, e in new_G[t].items():
                if w != s:
                    if w not in new_G[s]:
                        new_G.add_edge(s, w, weight=e['weight'])
                    else:
                        new_G[s][w]['weight'] += e['weight']
            new_G.remove_node(t)
            # 记录合并点
            room[s-1].extend(room[t-1])
            room[t - 1] = [-1]
            # 递归调用
            mincut2,best_part21,best_part22 = self.global_mincut(new_G,room)
            # 比较输出
            if mincut1<mincut2:
                return mincut1,best_part11,best_part12
            else:
                return mincut2,best_part21,best_part22


    def Stoer_Wagner(self):
        mincut = float('inf') # 最小割的值
        # 复制图，weight代表的是两点之间边的个数
        G = nx.Graph((u, v, {'weight': e.get('weight', 1)})
                     for u, v, e in self.G.edges(data=True) if u != v)
        n = len(list(G.nodes))
        best_part = None # 最小割的点集合
        room = [[i+1] for i in range(n)] # 用于记录点的合并
        for i in range(n-1):
            # 随机选一个点
            u = random.choice(list(G.nodes))
            A = set([u]) # 集合A
            PQ = dict() # 记录权重（边数）
            # 记录与点u相邻的点与u的边数
            for v,e in G[u].items():
                PQ[str(v)] = e['weight']
            for j in range(n-i-2):
                u = int(max(PQ, key=PQ.get)) # 找到与集合A中的点连边数量最多的点
                PQ.pop(str(u)) # 去除点u
                A.add(u) # 把点u放入集合A中
                # 重新计算图中除集合A中的点与集合A的连边数
                for v,e in G[u].items():
                    if v not in A:
                        if str(v) in PQ:
                            PQ[str(v)] += e['weight']
                        else:
                            PQ[str(v)] = e['weight']
            v = int(max(PQ, key=PQ.get)) # 得到终点v
            w = PQ[str(v)] # 得到v与集合A的连边数（割数）
            # 如果w的值为1，则直接得到全局最小割
            if w == 1:
                mincut = 1
                best_part = room[v-1]
                best_part2 = list(set(list(self.G.nodes)).difference(set(best_part)))
                return mincut, best_part, best_part2
            # 如果w小于全局最小割的值，则令全局最小割的值为w
            if w<mincut:
                mincut=w
                best_part = room[v-1]
            # 合并点，weight代表的是两点之间边的个数
            for w, e in G[v].items():
                if w != u:
                    if w not in G[u]:
                        G.add_edge(u, w, weight=e['weight'])
                    else:
                        G[u][w]['weight'] += e['weight']
            G.remove_node(v)
            # 记录合并点
            room[u-1].extend(room[v-1])
            room[v - 1] = [-1]
        best_part2 = list(set(list(self.G.nodes)).difference(set(best_part)))
        return mincut,best_part,best_part2


    def karger_stein(self,G,room):
        # G为图，room是存放点集的列表，用于表示最小割的点集
        nodes = list(G.nodes)
        # 如果图仅有两个点，则直接输出两点间权重（边数）
        if len(nodes) == 2:
            best_part = room[nodes[0]-1]
            best_part2 = room[nodes[1]-1]
            mincut = G[nodes[0]][nodes[1]]['weight']
            return mincut,best_part,best_part2
        else:
            # 循环一定次数
            max_iters = len(nodes)/math.sqrt(2)
            while len(list(G.nodes)) > max_iters:
                # # 判断图中点的最小度数,如果为1,则直接可以确定全局最小割
                # A = nx.adjacency_matrix(G)  # 图G的邻接矩阵
                # indices = A.indices  # 点集
                # for ind in indices:
                #     num_sum = sum(A[ind].data)  # 点的度数
                #     if num_sum == 1:
                #         li = list(self.G.nodes)
                #         mincut = num_sum
                #         best_part = room[ind - 1]
                #         best_part2 = list(set(li).difference(set(best_part)))
                #         return mincut, best_part, best_part2
                self.contract(G, room)
            # 复制图
            new_G = nx.Graph((u, v, {'weight': e.get('weight', 1)})
                         for u, v, e in G.edges(data=True) if u != v)
            new_room = copy.deepcopy(room)
            # 分两次递归调用，比较最小割的值大小，输出最小值
            mincut1,best_part11,best_part21 = self.karger_stein(G,room)
            mincut2,best_part12,best_part22 = self.karger_stein(new_G,new_room)
            if mincut1<mincut2:
                return mincut1,best_part11,best_part21
            else:
                return mincut2,best_part12,best_part22

    # karger算法
    def karger(self):
        result = float('inf') # 初始化min-cut的值
        n = len(list(self.G.nodes)) # 节点数
        iters = n*(n-1)/2 # 循环次数
        iters = int(iters)
        best_part = None # 初始化全局最小割
        best_part2 = None
        for _ in range(iters):
            room2 = [[i + 1] for i in range(n)]
            # 复制
            G = nx.Graph((u, v, {'weight': e.get('weight', 1)})
                         for u, v, e in self.G.edges(data=True) if u != v)
            # 循环压缩，直到只剩两个点
            while len(list(G.nodes)) > 2:
                self.contract(G, room2)
            nodes = list(G.nodes)
            # 比较割的值，取最小值作为全局最小割
            if G[nodes[0]][nodes[1]]['weight'] < result:
                result = G[nodes[0]][nodes[1]]['weight']
                best_part = room2[nodes[0] - 1]
                best_part2 = room2[nodes[1] - 1]
        return result, best_part, best_part2

    # 并行karger算法
    def karger_pl(self,iters):
        best_part = None
        best_part2 = None
        mincut = float('inf')
        for _ in range(iters):
            # 复制
            G = nx.Graph((u, v, {'weight': e.get('weight', 1)})
                         for u, v, e in self.G.edges(data=True) if u != v)
            n = len(list(G.nodes))
            room = [[i + 1] for i in range(n)]
            # 循环压缩，直到只剩两个点
            while len(list(G.nodes)) > 2:
                self.contract(G, room)
            nodes = list(G.nodes)
            if G[nodes[0]][nodes[1]]['weight'] < mincut:
                mincut = G[nodes[0]][nodes[1]]['weight']
                best_part = room[nodes[0] - 1]
                best_part2 = room[nodes[1] - 1]
        return {'mincut':mincut,'part':[set(best_part),set(best_part2)]}


if __name__ == '__main__':

    name = 'example'
    file_name = './data/'+name+'.txt'
    num_cores = int(mp.cpu_count())
    graph = Graph(file_name)
    num_nodes = len(list(graph.G.nodes))
    num_edges = len(list(graph.G.edges))
    print(name)
    print("点数：{0}，边数：{1}".format(num_nodes,num_edges))

    # karger-pl
    graph2 = Graph(file_name)
    n = len(graph2.G.nodes)
    iters = n*(n-1)/2
    iters = int(iters/num_cores)
    it_list = [iters for _ in range(num_cores)]
    pool = mp.Pool(num_cores)
    start = time.perf_counter()
    results = [pool.apply_async(graph2.karger_pl, args=(it,)) for it in it_list]
    results = [p.get() for p in results]
    mincut = float('inf')
    best_part = None
    for i in range(num_cores):
        if results[i]['mincut'] < mincut:
            mincut = results[i]['mincut']
            best_part = results[i]['part']
    print(mincut, best_part)
    end = time.perf_counter()
    print("karger并行 = {0}".format(end - start))

    # 改进前stoer
    start = time.perf_counter()
    graph3 = Graph(file_name)
    room = [[i + 1] for i in range(len(list(graph3.G.nodes)))]
    mincut,best_part,best_part2 = graph3.global_mincut(graph3.G,room)
    print(mincut,best_part,best_part2)
    end = time.perf_counter()
    print("改进前stoer = {0}".format(end - start))

    # stoer-wagner
    start = time.perf_counter()
    graph4 = Graph(file_name)
    mincut,best_part11,best_part12 = graph4.Stoer_Wagner()
    print(mincut,best_part11,best_part12)
    end = time.perf_counter()
    print("改进后stoer = {0}".format(end - start))

    # karger-stein
    start = time.perf_counter()
    graph5 = Graph(file_name)
    room = [[i + 1] for i in range(len(list(graph5.G.nodes)))]
    mincut,best_part,best_part2 = graph.karger_stein(graph5.G,room)
    print(mincut, best_part, best_part2)
    end = time.perf_counter()
    print("改进前karger-stein = {0}".format(end - start))

    # karger
    graph6 = Graph(file_name)
    start = time.perf_counter()
    mincut1,best_part111,best_part222 = graph6.karger()
    print(mincut, best_part, best_part2)
    end = time.perf_counter()
    print("karger = {0}".format(end - start))