import json
import random
from QKDNetwork import QKDNetwork
from add_ps import AddPS
import networkx as nx
import numpy as np
import threading
import sys

thread_num = 48

def get_sps_add_ps_num(net: QKDNetwork, init_ps: list):
    sps_count = 0
    for node in net.G.nodes:
        if(node in init_ps):
            net.G.nodes[node]["transmitter"] = 1
    for sd in net.sd_list:
        path = nx.shortest_path(net.G, sd[0], sd[1])
        for path_index, path_node in enumerate(path):
            if(path_index % 2 == 0 and net.G.nodes[path_node]["transmitter"] == 0):
                net.G.nodes[path_node]["transmitter"] = 1
                sps_count += 1
    return sps_count

def get_ras_add_ps_num(net: QKDNetwork, init_ps: list):
    ras_count = 0
    for node in net.G.nodes:
        if(node in init_ps):
            net.G.nodes[node]["transmitter"] = 1
    for edge in net.G.edges:
        if(net.G.nodes[edge[0]]["transmitter"] == 0 and net.G.nodes[edge[1]]["transmitter"] == 0):
            ras_count += 1
            if(net.G.degree(edge[0]) > net.G.degree(edge[1])):
                net.G.nodes[edge[0]]["transmitter"] = 1
            else:
                net.G.nodes[edge[1]]["transmitter"] = 1
    return ras_count

def clear_all_placed_ps(net: QKDNetwork):
    for node in net.G.nodes:
        if(net.G.nodes[node]["transmitter"] == 1):
            net.G.nodes[node]["transmitter"] = 0

# 节点个数： 20-210          默认：60      迭代：20 + cnt * 10
# 源目对数： 10-200          默认：50      迭代：10 + cnt * 10
# 阿尔法值： 0.35-1.3        默认：0.2     迭代：0.2 + cnt * 0.05
# 初始光源： 5-43            默认：5       迭代：5 + cnt * 2

default_node_num = 60
default_sd_num = 50
default_alpha = 0.35
default_init_ps_num = 5

all_fig_data = []

def simulation_run(fig_idx, fig_index, fig_data_single):
    q = None
    init_ps = []
    if(fig_idx == 0):
        q = QKDNetwork(num_node=20 + fig_index * 10, sd_num=default_sd_num, alpha=default_alpha, no_add_ps=True)
    elif(fig_idx == 1):
        q = QKDNetwork(num_node=default_node_num, sd_num=10 + fig_index * 10, alpha=default_alpha, no_add_ps=True)
    elif(fig_idx == 2):
        q = QKDNetwork(num_node=default_node_num, sd_num=default_sd_num, alpha=0.2 + fig_index * 0.05, no_add_ps=True)
    elif(fig_idx == 3):
        q = QKDNetwork(num_node=default_node_num, sd_num=default_sd_num, alpha=default_alpha, no_add_ps=True)
    init_ps = random.sample(list(range(q.num_nodes)), default_init_ps_num + (fig_idx == 3) * fig_index * 1)
    clear_all_placed_ps(q)
    RAS_data = get_ras_add_ps_num(q, init_ps)
    if(RAS_data != 0):
        fig_data_single[0].append(RAS_data)
    clear_all_placed_ps(q)
    SPS_data = get_sps_add_ps_num(q, init_ps)
    if(SPS_data != 0):
        fig_data_single[1].append(SPS_data)
    clear_all_placed_ps(q)
    ret = AddPS(q).run_expr(init_ps, remove_ilp=False)
    if(ret[0] != 0):
        fig_data_single[2].append(ret[0])
    if(ret[1] != 0):
        fig_data_single[3].append(ret[1])

def thread_function(fig_idx, fig_index, fig_data_single):
    simulation_run(fig_idx, fig_index, fig_data_single)
    
fig_idx = int(sys.argv[1])
print(f"Add PS, Fig: {fig_idx}")

fig_data = [[], [], [], []]
for fig_index in range(20):  # 每张图多少x轴点
    threads = []
    fig_data_single = [[], [], [], []]
    for _ in range(thread_num): 
        t = threading.Thread(target=thread_function, args=(fig_idx, fig_index, fig_data_single))
        threads.append(t)
        t.start()
    for t in threads:
        t.join()
    fig_data[0].append(np.mean(fig_data_single[0]))
    fig_data[1].append(np.mean(fig_data_single[1]))
    fig_data[2].append(np.mean(fig_data_single[2]))
    fig_data[3].append(np.mean(fig_data_single[3]))
    
    print([np.mean(fig_data_single[0]), np.mean(fig_data_single[1]), np.mean(fig_data_single[2]), np.mean(fig_data_single[3])])
    
all_fig_data.append(fig_data)
data_string = json.dumps(all_fig_data)
with open(f'add_ps_cloud_10_05_{fig_idx}.txt', 'a') as file:
    file.write(data_string + '\n')