import heapq
import time
import networkx as nx
import matplotlib.pyplot as plt
from generate import GenerateNetwork
import random

class QKDNetwork:
    def sd_demand(self):
        tmp_rate = [1]
        if(self.hete):
            tmp_rate = [2, 4]
        return random.choice(tmp_rate) # QKD请求的速率随机从 2/4/6/8 Kbps 取
    
    def ps_rate(self):
        tmp_rate = [1]
        if(self.hete):
            tmp_rate = [1, 2]
        return random.choice(tmp_rate) # 密钥发射器速率随机从 1/2 Kbps 取

    def __init__(self, showTopology = False, num_node=20, sd_num=20, alpha=0.35, no_add_ps=False, hete = False):
        self.showTopology = showTopology
        self.hete = hete
        self.no_add_ps = no_add_ps
        self.num_nodes = num_node       # 节点数量
        self.num_sd_pairs = sd_num      # 请求个数
        self.alpha = alpha              # Waxmax Alpha

        self.prepareNetwork()
        # print("网络准备完成")
        # 展示网络
        if(self.showTopology):
            self.showNetwork()
        # QKD请求的SD对
        self.sd_list = self.generate_random_sd_pairs(self.num_sd_pairs, 0, self.num_nodes - 1)
        # print("请求准备完成")
        if(not no_add_ps):
            self.getCandidatePaths()
        # print("候选路径准备完成")
        if(not no_add_ps):
            self.estimateMinTimeSlot()
        # print(f"预估时隙完成: {self.estimate_max_time_slot}")
        
    def removePS(self):
        # 去掉所有的PS
        for node in self.G.nodes:
            self.G.nodes[node]["transmitter"] = 0
        
    def clearLoad(self):
        for u, v, _ in self.G.edges(data=True):
            self.G[u][v]["load"] = 0

    def getShortestPathWithTransmitter(self, source, target, forbidden_links: list = []):
        # 获取网络中源目节点对的中继路径，要求穿过Transmitter
        processedNodes = [] #存放已处理过的节点ID
        toBeProcessedNodes = [] #待处理的节点ID

        for i in self.G.nodes: #初始化
            node = self.G.nodes[i]
            node["distance"] = 100000000000
            node["physicsDistance"] = 100000000000
            node["shortestPathLastNodeID"] = -1 #路径的上一个节点ID为-1
            node["isAddedToQueue"] = False
            node["isVisited"] = False

        source_node = self.G.nodes[source]
        source_node["distance"] = 0 if source_node["transmitter"]==1 else 1 #初始化距离，如果是transmitter就初始化为1，否则为0
        source_node["physicsDistance"] = 0
        source_node["shortestPathLastNodeID"] = source #最短路径上自己的上一个节点就是自己
        toBeProcessedNodes.append(source)
        source_node["isAddedToQueue"] = True

        while (len(toBeProcessedNodes) != 0):
            toBeProcessedNodes = sorted(toBeProcessedNodes, key=lambda x: self.G.nodes[x]["physicsDistance"]) #排序
            tmp = toBeProcessedNodes.pop(0)
            processedNodes.append(tmp)
            self.G.nodes[tmp]["isVisited"] = True
            for i in self.G.neighbors(tmp):
                if(([tmp, i] in forbidden_links) or ([i, tmp] in forbidden_links)):
                    continue # 去掉断掉的边
                node_adjacent = self.G.nodes[i]
                node_tmp = self.G.nodes[tmp]
                if(node_adjacent["isVisited"] or (node_adjacent["transmitter"]==0 and node_tmp["distance"] + 1 > 1)):
                    continue
                if(node_adjacent["distance"] == 100000000000 or node_adjacent["physicsDistance"] > node_tmp["physicsDistance"] + 1):
                    node_adjacent["distance"] = node_tmp["distance"] + 1
                    node_adjacent["physicsDistance"] = node_tmp["physicsDistance"] + 1
                    node_adjacent["shortestPathLastNodeID"] = tmp
                if(node_adjacent["transmitter"]==1):
                    node_adjacent["distance"] = 0
                if(not node_adjacent["isAddedToQueue"]):
                    toBeProcessedNodes.append(i)
                    node_adjacent["isAddedToQueue"] = True
        try:
            path = [target]
            while(target != source and target != None):
                tmpLastID = self.G.nodes[target]["shortestPathLastNodeID"]
                path.insert(0, tmpLastID)
                target = tmpLastID
        except Exception as _:
            return [] # 说明没有合适的路
        
        # print(path)
        # print([self.G.nodes[j]["transmitter"] for j in path])

        return path
    
    def getCandidatePaths(self):
        candidate_paths = []
        self.max_candidate_size = 0
        self.max_candidate_path_length = 0
        for sd in self.sd_list:
            # 计算候选路径
            paths = []
            path_1 = self.getShortestPathWithTransmitter(sd[0], sd[1])
            self.max_candidate_path_length = self.max_candidate_path_length if self.max_candidate_path_length > len(path_1) else len(path_1)
            paths.append(path_1)
            for i in range(len(path_1) - 1):
                link = [path_1[i], path_1[i+1]]
                # 逐一断开
                yen_path = self.getShortestPathWithTransmitter(sd[0], sd[1], forbidden_links=[link])
                if(len(yen_path) != 0 and yen_path not in paths):
                    paths.append(yen_path)
                    self.max_candidate_path_length = self.max_candidate_path_length if self.max_candidate_path_length > len(yen_path) else len(yen_path)
            if(len(path_1) > 2):
                for i in range(len(path_1) - 2):
                    link_1 = [path_1[i], path_1[i+1]]
                    link_2 = [path_1[i+1], path_1[i+2]]
                    # 逐二断开
                    yen_path = self.getShortestPathWithTransmitter(sd[0], sd[1], forbidden_links=[link_1, link_2])
                    if(len(yen_path) != 0 and yen_path not in paths):
                        paths.append(yen_path)
                        self.max_candidate_path_length = self.max_candidate_path_length if self.max_candidate_path_length > len(yen_path) else len(yen_path)
            # print("Yen", paths)
            self.max_candidate_size = self.max_candidate_size if self.max_candidate_size > len(paths) else len(paths)
            candidate_paths.append({
                "sd": sd,
                "paths": paths
            })
        self.candidate_paths = candidate_paths
        return

    def estimateMinTimeSlot(self):
        # 估计最小的完成时间隙数量，每个请求都按照最短路
        def addFlowOnPath(sd): # 施加流
            path = self.getShortestPathWithTransmitter(sd[0], sd[1])
            for i in range(len(path)-1):
                link_start = path[i]
                link_end = path[i + 1]
                self.G[link_start][link_end]["load"] += sd[2]
        for tmp_sd in self.sd_list:
            addFlowOnPath(tmp_sd)

        max_time_slot = 0
        while(True):
            flag = True
            for i in self.G.nodes: #初始化
                node = self.G.nodes[i]
                if(node["transmitter"] == 1):
                    # 选择一个需要密钥最多的链路进行服务
                    max_rate_link_node = 0
                    max_rate = 0
                    for neighbor in self.G.neighbors(i):
                        if(self.G[i][neighbor]["load"] > max_rate):
                            max_rate = self.G[i][neighbor]["load"]
                            max_rate_link_node = neighbor
                            flag = False
                    if(max_rate > 0):
                        self.G[max_rate_link_node][i]["load"] -= node["transmitter_rate"]
            if(flag):
                break
            else:
                max_time_slot += 1
        self.estimate_max_time_slot = max_time_slot
    
    def generate_random_sd_pairs(self, num_pairs, min_val, max_val):
        pairs = set()
        while len(pairs) < num_pairs:
            sd_start_node = random.randint(min_val, max_val)
            sd_end_node = random.randint(min_val, max_val)
            sd_request_demand = self.sd_demand()
            if sd_start_node != sd_end_node:
                pairs.add((sd_start_node, sd_end_node, sd_request_demand))
        return list(pairs)

    def showNetwork(self):
        node_colors = ['red' if self.G.nodes[node]['transmitter'] == 1 else 'skyblue' for node in self.G.nodes]
        plt.figure(figsize=(8, 6))
        nx.draw(self.G, with_labels=True, node_color=node_colors, edge_color='gray', 
                node_size=500, font_size=10)
        plt.title("Random Graph with Transmitters")
        plt.show()

    def prepareNetwork(self):
        beta = 0.1
        self.G: nx.Graph = GenerateNetwork().generate_network(self.num_nodes, self.alpha, beta)
        # 检查图是否连通
        if not nx.is_connected(self.G):
            # 获取所有连通分量
            components = list(nx.connected_components(self.G))
            # 连接每个连通分量
            for i in range(len(components) - 1):
                # 在第i个和第i+1个连通分量之间添加一条边
                # 选择每个连通分量中的任意一个节点
                a = next(iter(components[i]))
                b = next(iter(components[i + 1]))
                self.G.add_edge(a, b)
        
        # 准备网络，使得网络链路的两个端点中至少有一个为发射器节点
        for node in self.G.nodes():
            self.G.nodes[node]["transmitter"] = 0
            self.G.nodes[node]["transmitter_rate"] = 0
        
        for u, v, _ in self.G.edges(data=True):
            self.G[u][v]["load"] = 0
            
        if(self.no_add_ps):
            return

        def hasFeasiblePath(u, v) -> bool:
            max_length = 3
            def isFeasiblePath(path: list) -> bool:
                for tmpLink in [[path[i], path[i+1]] for i in range(len(path)-1)]:
                    link_start_transmitter = self.G.nodes[tmpLink[0]]["transmitter"]
                    link_end_transmitter = self.G.nodes[tmpLink[1]]["transmitter"]
                    if(link_start_transmitter==0 and link_end_transmitter==0):
                        return False
                return True
            for path in nx.all_simple_paths(self.G, source=u, target=v, cutoff=max_length):
                if len(path) <= max_length+1:  
                    if(isFeasiblePath(path)):
                        return True
            return False

        for _ in range(100):
            for u in self.G.nodes:
                if self.G.nodes[u]["transmitter"] == 1 and self.G.nodes[v]["transmitter"] == 1:
                    # 两边都配备有密钥发生器，随机去掉一个
                    random_node = random.choice([u, v])
                    self.G.nodes[random_node]["transmitter"] = 0
            for u, v in self.G.edges:
                if self.G.nodes[u]["transmitter"] == 0 and self.G.nodes[v]["transmitter"] == 0:
                    # if((not hasFeasiblePath(u, v))):
                    # 两边都没有配备有密钥发生器且无中转节点，添加一个度数较大的
                    smaller_degree_node = u if self.G.degree(u) < self.G.degree(v) else v
                    self.G.nodes[smaller_degree_node]["transmitter"] = 1

        trans_count = 0
        not_trans_count = 0
        for i in self.G.nodes:
            if(self.G.nodes[i]["transmitter"] == 1):
                trans_count += 1
                self.G.nodes[i]["transmitter_rate"] = self.ps_rate()
            else:
                not_trans_count += 1
        
        # print(f"Transmitter: {trans_count}")
        # print(f"NotTransmitter: {not_trans_count}")
        # print(f"发射器编排完成")