from alg_clink import alg_clink
from alg_map import alg_map
import json, datetime, random
from alg_clif_int import alg_new_int
import argparse
from alg_tomo_z3 import alg_tomo
from Topology_zoo import *
import time, os
import pickle
    
def matrics_path(True_X, infer_X, A, obs_Y):
    path_performance = np.zeros((A.shape[0], 5))
    
    for round in range(True_X.shape[1]):
        for path_idx in range(A.shape[0]):
            TP, FP, FN, TN = 0, 0, 0, 0
            if obs_Y[path_idx][round] > 0:
                link_index_set = np.where(A[path_idx, :] == 1)[0] # 找出该路径上的所有链路
                for link_idx in link_index_set:
                    if True_X[link_idx][round]==0 and infer_X[link_idx][round]==0:
                        TN+=1
                    elif True_X[link_idx][round]==0 and infer_X[link_idx][round]==1:
                        FP+=1
                    elif True_X[link_idx][round]==1 and infer_X[link_idx][round]==1:
                        TP+=1
                    elif True_X[link_idx][round]==1 and infer_X[link_idx][round]==0:
                        FN+=1
                if TP+FN == 0:
                    DR=1
                else:
                    DR=TP/(TP+FN)
                    
                if FP + TN == 0:
                    FPR=0
                else:
                    FPR=FP/(FP+TN)
                if TP+FP==0:
                    precision=1
                else:
                    precision=TP/(TP+FP)
                if TP+FN==0:
                    recall=1
                else:
                    recall=TP/(TP+FN)
                    
                if TP == 0:
                    if sum(True_X[link_index_set, round])==0:
                        F1=1;F2=1
                    else:
                        F1=0;F2=0
                else:
                    F1 = 2*precision*recall/(precision+recall)
                    F2 = 5*precision*recall/(4*precision+recall)
                path_performance[path_idx, :] += np.array([DR, FPR, F1, F2, precision])
    
    
    path_1_sum = np.sum(obs_Y, axis=1)
    path_1_sum = np.array([x for x in path_1_sum])
    # print("path_1_sum", path_1_sum)
    
    path_performance_  = path_performance /  path_1_sum.reshape(-1, 1)
    # print("path_performance_", path_performance_)
    return path_performance_
     
    
def matrics_link(True_X, infer_X, n_link):
    
    link_performance = np.zeros((n_link, 5))
    # print("infer_x", infer_X.shape)
    for round in range(True_X.shape[1]):
        for link_idx in range(n_link):
            TP, FP, FN, TN = 0, 0, 0, 0
            if True_X[link_idx][round]==0 and infer_X[link_idx][round]==0:
                TN+=1
            elif True_X[link_idx][round]==0 and infer_X[link_idx][round]==1:
                FP+=1
            elif True_X[link_idx][round]==1 and infer_X[link_idx][round]==1:
                TP+=1
            elif True_X[link_idx][round]==1 and infer_X[link_idx][round]==0:
                FN+=1
                
            if TP+FN == 0:
                DR=1
            else:
                DR=TP/(TP+FN)
                
            if FP + TN == 0:
                FPR=0
            else:
                FPR=FP/(FP+TN)
            if TP+FP==0:
                precision=1
            else:
                precision=TP/(TP+FP)
            if TP+FN==0:
                recall=1
            else:
                recall=TP/(TP+FN)
            if TP == 0:
                if True_X[link_idx][round] == 0:
                    F1=1; F2=1
                else:
                    F1=0;F2=0
            else:
                
                F1 = 2*precision*recall/(precision+recall)
                F2 = 5*precision*recall/(4*precision+recall)
                
            link_performance[link_idx, :] += np.array([DR, FPR, F1, F2, precision])
    
    link_performance /= True_X.shape[1]
    return link_performance


def matrics(True_X, infer_X, round_matrics, allow=False):
    #inferred_array由算法得出,real_array即真实信息,rounds表示测试了几轮
    ave_FPR=0
    ave_DR=0
    ave_F1=0
    ave_F2=0; ave_Prec=0
    rounds = infer_X.shape[1]; real_array = True_X; inferred_array = infer_X
    for round in range(rounds):
        TP=0
        FP=0
        FN=0
        TN=0
        for link_index in range(len(real_array.T[round])):
            if real_array.T[round][link_index]==0 and inferred_array.T[round][link_index]==0:
                TN+=1
            elif real_array.T[round][link_index]==0 and inferred_array.T[round][link_index]==1:
                FP+=1
            elif real_array.T[round][link_index]==1 and inferred_array.T[round][link_index]==1:
                TP+=1
            elif real_array.T[round][link_index]==1 and inferred_array.T[round][link_index]==0:
                FN+=1
        if TP+FN == 0:
            DR=1
        else:
            DR=TP/(TP+FN)
            
        if FP + TN == 0:
            FPR=0
        else:
            FPR=FP/(FP+TN)
        if TP+FP==0:
            precision=1
        else:
            precision=TP/(TP+FP)
        if TP+FN==0:
            recall=1
        else:
            recall=TP/(TP+FN)
        if TP == 0:
            if sum(real_array.T[round])==0:
                F1=1;F2=1
            else:
                F1=0;F2=0
        else:
            F1 = 2*precision*recall/(precision+recall)
            F2 = 5*precision*recall/(4*precision+recall)
        ave_FPR += FPR
        ave_DR += DR
        ave_F1 += F1
        ave_F2 += F2; ave_Prec += precision
        if allow:
            single_t = np.array([DR, FPR, F1, F2, precision]).reshape(-1, 1)
            round_matrics.append(single_t)
            
            # print("round_matrics", round_matrics)
    ave_FPR /=rounds
    ave_DR /=rounds
    ave_F1 /=rounds; ave_F2 /=rounds; ave_Prec /=rounds   
    
            
    
    return [ave_DR,ave_FPR, ave_F1, ave_F2, ave_Prec]
            

def get_timestamp():
    '''
    - 获取当前时间戳
    -----------------
    - return 当前时间戳
    '''
    timestamp = time.time()
    # 将时间戳转换为日期时间对象
    dt_object = datetime.datetime.fromtimestamp(timestamp)
    # 将日期时间对象转换为字符串
    date_time_str = dt_object.strftime("%Y-%m-%d_%H:%M:%S")
    return date_time_str




class Simulation(网络基类):
    def __init__(self, topology_name: str, 源节点列表: list, n_samples: int, 先验拥塞概率: int):
        super().__init__()
        self.源节点列表 = 源节点列表
        self.topology_name = topology_name
        self.n_samples = n_samples
        self.先验拥塞概率 = 先验拥塞概率
        self.pri_cong_prob = None

    def __del__(self):                  
        print("运行结束，释放其空间")

    def return_Y_A_X_p(self, seed, seed_p):
        '''
        - 根据 拓扑的相关信息与 设置的相关参数生成 矩阵
        ----------------------
        -return 观测矩阵，路由矩阵，真实链路状态矩阵，先验拥塞概率
        '''
        self.配置拓扑(f"./topology_zoo/topology_zoo数据集/{self.topology_name}.gml")
        self.部署测量路径(源节点列表=self.源节点列表)

        np.random.seed(seed_p)
        self.配置参数(异构链路先验拥塞概率 = self.先验拥塞概率)
        pri_cong_prob = np.array([i.先验拥塞概率  for  i in self.链路集和])
        self.pri_cong_prob = pri_cong_prob
        
        np.random.seed(seed)
        运行时间 = self.n_samples
        self.运行网络(运行的总时间 = 运行时间)

        
        路由矩阵_ = self.路由矩阵
        观测矩阵_ = self.运行日志['路径和状态']
        链路观测矩阵_ = self.运行日志['链路状态']

        A = np.where(路由矩阵_, 1, 0)
        Y = np.array(观测矩阵_)
        True_X = np.where(链路观测矩阵_, 0, 1)

        return Y, A, True_X, pri_cong_prob



if __name__ == '__main__':

    # 创建参数解析器
    parser = argparse.ArgumentParser(description="解析")
    # 添加参数
    parser.add_argument('prob_', type=float, help="先验拥塞概率")
    parser.add_argument('bias', type=float, help="先验拥塞概率误差")
    parser.add_argument('changjing_seed', type=int, help="场景种子")
    parser.add_argument('topology_name', type=str, help="拓扑名称")
    parser.add_argument('params', nargs='+', type=int, help="参数列表")

    # 解析命令行参数
    args = parser.parse_args()
    prob_ = args.prob_
    bias_ = args.bias
    changjing_seed = args.changjing_seed
    topology_name = args.topology_name
    源节点列表 =  [int(param) for param in args.params]

    print(f"{topology_name}- {源节点列表}")
    print(源节点列表)
    
    ### 参数设定 ###
    label = ["DR", "FPR", "F1", "F2"]
    n_samples = 5000
    先验拥塞概率 = prob_
    prop = 0.01
    bias = bias_
    run_time = 10
    seed_p = int(changjing_seed)
    ### 参数设定 ###
    
    print(f"this process {topology_name}, prob is {prob_}, source node is {源节点列表}")
    # 创建多个进程     

    result_clink, result_heu, result_tomo, result_map = [], [], [], []
    result_clink_links, result_heu_links, result_tomo_links, result_map_links = [], [], [], []
    result_clink_paths, result_heu_paths, result_tomo_paths, result_map_paths = [], [], [], []
    pri_c_prob, mean_prob = None, None
    obs_Y_ALL = []
    round_matrics = []
    # 同一场景重复运行 run_time 次
    for m in range(run_time):
        
        simulation = Simulation(topology_name=topology_name, 源节点列表=源节点列表, n_samples=n_samples, 先验拥塞概率=先验拥塞概率)
        seed_ = random.randint(0, 100)

        Y, A, True_X, pri_cong_prob = simulation.return_Y_A_X_p(seed_, seed_p)
        print("pri_cong_prob", pri_cong_prob)
        
        obs_Y_ALL.append(Y)
        
        # Add noise
        np.random.seed(seed_)
        error = bias * pri_cong_prob.max()
            # 生成随机噪声
        noise = np.random.uniform(-bias, bias, pri_cong_prob.shape)
        noise += 1
            # 添加噪声到原始数据
        pri_cong_prob = pri_cong_prob * noise
        
            # 对元素做检查，防止有小于0 的元素
        pri_cong_prob = np.array([x if x > 0 else 0.0001 for x in pri_cong_prob])
        pri_c_prob = pri_cong_prob
        
        # 添加数据
        infer_X_heu = alg_new_int(Y, A, pri_cong_prob)
        
        # 总体性能
        result_heu.append(matrics(True_X, infer_X_heu, round_matrics, True))


    # 总体性能
    result_heu = np.array(result_heu)
    result_heu = np.sum(result_heu, axis=0) / result_heu.shape[0]
 
    print(f"new: {result_heu.tolist()}")
    
    # 初始化保存数据的字典
    results = {
        "prop": prop,
        "changjing_seed_p": seed_p,
        "topology": topology_name,
        "源节点列表": 源节点列表,
        "平均拥塞概率": 先验拥塞概率,
        "实际平均拥塞概率": np.mean(pri_c_prob),
        "prob_bias": bias,
        "version": "heu",
        "run_time": run_time,
        "n_samples": n_samples,
        "statistics": {},
        "link-performance": {},
        "path-performance": {},
        "pri-cong-prob": pri_c_prob.tolist(),
    }
    
    labels = ['DR', 'FPR', 'F1', 'F2', 'Prec']
    
    # 将统计信息存储到字典中
    for cnt in range(len(labels)):
        results["statistics"][labels[cnt]] = {
            "MF-INT": result_heu[cnt]
        }
    
    output_name = get_timestamp()

    folder_path = f'./New_int/uniform/bias_{bias}/{topology_name}-{先验拥塞概率}'
    file_name = f'{output_name}.json'
    full_path = os.path.join(folder_path, file_name)

    # 检查目录是否存在，如果不存在则创建
    if not os.path.exists(folder_path):
        os.makedirs(folder_path)