import random
import math
from docplex.mp.model import Model
from network import pro

class GenerateNetwork:
    class Node:
        def __init__(self, nodeID):
            self.nodeID = nodeID
            self.adjacentNodes = []
    
    def getNode(self, nodeId):
        if nodeId - 1 >= len(self.nodeList) or nodeId - 1 < 0:
            return self.Node(-1)
        return self.nodeList[nodeId - 1]

    def __init__(self, path_for_getxt, numNodes, sdNum, alpha, beta, seed, pro):
        self.numNodes = numNodes
        self.sdNum = sdNum
        self.sdList = []
        self.nodeList = []
        try:
            outfile = open(path_for_getxt, "w")
            outfile.write(str(numNodes) + "\n")
            for i in range(1, numNodes+1):
                outfile.write(str(i) + "\n")
                self.nodeList.append(self.Node(i))
        except Exception as e:
            print("文件写入出错:", e)
        generator = random.Random(seed)
        edges = []
        for i in range(1, numNodes+1):
            for j in range(i + 1, numNodes+1):
                dist = math.sqrt((i - j) ** 2)
                prob = alpha * math.exp(-beta * dist)
                if generator.random() < prob:
                    edges.append((i, j))
        if len(edges) == 0:
            edges.append((1, 2))
        self.make_graph_connect(edges, numNodes)
        outfile.write(str(len(edges)) + "\n")
        print("开始写入")
        for tmp in edges:
            startNode = tmp[0]
            endNode = tmp[1]
            self.nodeList[startNode - 1].adjacentNodes.append(endNode)
            self.nodeList[endNode - 1].adjacentNodes.append(startNode)
            randomInt = random.randint(2*pro-100, 100)
            outfile.write(f"{tmp[0]} {tmp[1]} {randomInt}\n")
        outfile.write(str(sdNum) + "\n")
        for i in range(sdNum):
            startNodeID = random.randint(1, numNodes)
            endNodeID = random.randint(1, numNodes)
            while startNodeID == endNodeID or [startNodeID, endNodeID] in self.sdList or [endNodeID, startNodeID] in self.sdList:
                startNodeID = random.randint(1, numNodes)
                endNodeID = random.randint(1, numNodes)
            self.sdList.append([startNodeID, endNodeID])
            outfile.write(str(startNodeID) + " " + str(endNodeID) + "\n")
        outfile.close()
        print("拓扑信息写入完成")
        self.calculateEPS(path_for_getxt)
        
    def calculateEPSSub(self, selectedSu: list) -> int:
        model = Model("EPS_LP")
        nodeNum = self.numNodes + 1
        sdNum = self.sdNum
        Su = [model.continuous_var(lb=0, ub=1, name=f"EPS[{idx}]")
              for idx in range(nodeNum)]
        for i in selectedSu:
            model.add_constraint(Su[i] == 1)
        e_uvi = {}
        for u in range(nodeNum):
            e_uvi[u] = {}
            for v in range(nodeNum):
                e_uvi[u][v] = []
                for i in range(sdNum):
                    e_uvi[u][v].append(model.continuous_var(
                        lb=0, ub=1, name=f"e[{u}][{v}][{i}]"))
        for i in range(sdNum):
            sd_tmp = self.sdList[i]
            for u in range(nodeNum):
                if (u == 0):
                    continue
                flowConstraintExpr = model.linear_expr()
                if (u == sd_tmp[0]):
                    flowConstraintExpr -= 1
                if (u == sd_tmp[1]):
                    flowConstraintExpr += 1
                for v in self.getNode(u).adjacentNodes:
                    flowConstraintExpr += e_uvi[u][v][i]
                    flowConstraintExpr -= e_uvi[v][u][i]
                model.add_constraint(flowConstraintExpr == 0)
        for i in range(sdNum):
            for u in range(nodeNum):
                if (u == 0):
                    continue
                for v in self.getNode(u).adjacentNodes:
                    model.add_constraint(e_uvi[u][v][i] <= Su[u] + Su[v])
        objExpr = model.linear_expr()
        for u in range(nodeNum):
            objExpr += Su[u]
        model.minimize(objExpr)
        # # 新增约束
        # for i in selectedSu:
        #     model.add_constraint(Su[i] == 1)

        # sol = model.solve()
        # if sol:
        #     Su_max = -1
        #     Su_max_nodeID = -1
        #     for u in range(nodeNum):
        #         if (sol[Su[u]]):
        #             if (sol[Su[u]] > Su_max and u not in selectedSu):
        #                 Su_max = sol[Su[u]]
        #                 Su_max_nodeID = u
        #             # print(f"已选中：{sol[Su[u]]}", u)
        #     # print(f"Su_max:{Su_max}")
        #     if (Su_max_nodeID != -1):
        #         print(f"EPSNodeID:{Su_max_nodeID}")
        #     return Su_max_nodeID

    def calculateEPS(self, path_for_getxt):
        print("开始创建模型")
        model = Model("EPS_LP")
        print("模型创建完成")
        selectedSu = []
        # self.calculateEPSSub(model, selectedSu)
        nodeNum = self.numNodes + 1
        sdNum = self.sdNum
        Su = [model.continuous_var(lb=0, ub=1, name=f"EPS[{idx}]")
              for idx in range(nodeNum)]
        # for i in selectedSu:
        #     model.add_constraint(Su[i] == 1)
        e_uvi = {}
        for u in range(nodeNum):
            e_uvi[u] = {}
            for v in range(nodeNum):
                e_uvi[u][v] = []
                for i in range(sdNum):
                    e_uvi[u][v].append(model.continuous_var(
                        lb=0, ub=1, name=f"e[{u}][{v}][{i}]"))
        for i in range(sdNum):
            sd_tmp = self.sdList[i]
            for u in range(nodeNum):
                if (u == 0):
                    continue
                flowConstraintExpr = model.linear_expr()
                if (u == sd_tmp[0]):
                    flowConstraintExpr -= 1
                if (u == sd_tmp[1]):
                    flowConstraintExpr += 1
                for v in self.getNode(u).adjacentNodes:
                    flowConstraintExpr += e_uvi[u][v][i]
                    flowConstraintExpr -= e_uvi[v][u][i]
                model.add_constraint(flowConstraintExpr == 0)
        for i in range(sdNum):
            for u in range(nodeNum):
                if (u == 0):
                    continue
                for v in self.getNode(u).adjacentNodes:
                    model.add_constraint(e_uvi[u][v][i] <= Su[u] + Su[v])
        objExpr = model.linear_expr()
        for u in range(nodeNum):
            objExpr += Su[u]
        model.minimize(objExpr)
        print("约束添加完成")
        while True:
            # 新增约束
            for i in selectedSu:
                model.add_constraint(Su[i] == 1)
            print("已经完成新约束的添加")
            sol = model.solve()
            if sol:
                Su_max = -1
                Su_max_nodeID = -1
                for u in range(nodeNum):
                    if (sol[Su[u]]):
                        if (sol[Su[u]] > Su_max and u not in selectedSu):
                            Su_max = sol[Su[u]]
                            Su_max_nodeID = u
                        # print(f"已选中：{sol[Su[u]]}", u)
                print(f"Su_max:{Su_max}")
                if (Su_max_nodeID != -1):
                    print(f"EPSNodeID:{Su_max_nodeID}")
                    selectedSu.append(Su_max_nodeID)
                else:
                    break
            
        outfile = open(path_for_getxt, "a")
        outfile.write(str(len(selectedSu)) + "\n")
        for i in selectedSu:
            outfile.write(str(i) + "\n")
        outfile.close()

    def get_unconnect_pair(self, edges, numNodes):
        map_arr = [0] * (numNodes + 1)
        node_queue_arr = [0] * (numNodes + 1)
        start_pos = 0
        end_pos = 0
        node_num = 0
        startBfsNode = edges[0][0]
        node_queue_arr[end_pos] = startBfsNode
        end_pos += 1
        node_num = 1
        map_arr[edges[0][0]] = 1

        while start_pos < end_pos:
            tmp = node_queue_arr[start_pos]
            start_pos += 1
            for edge in edges:
                startID = edge[0]
                endID = edge[1]
                if startID == tmp and endID != tmp and map_arr[endID] == 0:
                    node_queue_arr[end_pos] = endID
                    end_pos += 1
                    node_num += 1
                    map_arr[endID] = 1
                elif endID == tmp and startID != tmp and map_arr[startID] == 0:
                    node_queue_arr[end_pos] = startID
                    end_pos += 1
                    node_num += 1
                    map_arr[startID] = 1

        if node_num == numNodes:
            return False, (0, 0)
        else:
            for nodeIndex in range(1, numNodes):
                if map_arr[nodeIndex] + map_arr[nodeIndex + 1] == 1:
                    return True, (nodeIndex, nodeIndex + 1)
            return False, (0, 0)

    def make_graph_connect(self, edges, numNodes):
        is_unconnected, unconnected_pair = self.get_unconnect_pair(
            edges, numNodes)
        while is_unconnected:
            edges.append(unconnected_pair)
            is_unconnected, unconnected_pair = self.get_unconnect_pair(
                edges, numNodes)