import argparse
import os
import sys
from os.path import abspath, join, dirname

sys.path.insert(0, join(abspath(dirname(__file__)+'/startGNN_sample')))
from startGNN_sample.DistGraphLoader import DataSet,  partition_load
from startGNN_sample.Sample.neighbor_sampler import NeighborSampler
from startGNN_sample.Sample.base import NegativeSampling
#path1=os.path.abspath('.')  
import torch
from startGNN_sample.DistGraphLoader import DistGraphData
from startGNN_sample.DistGraphLoader import DistributedDataLoader
from startGNN_sample.DistGraphLoader import DistCustomPool 
import startGNN_sample.distparser as distparser
import time
import torch
import torch.nn.functional as F
import torch.distributed as dist
import torch.multiprocessing as mp
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.distributed import init_process_group, destroy_process_group
from startGNN_sample.shared_mailbox import SharedMailBox, SharedRPCMemoryManager
import os
"""
test command 
python test.py --world_size 2 --rank 0 
--world_size', default=4, type=int, metavar='W',
                    help='number of workers')
parser.add_argument('--rank', default=0, type=int, metavar='W',
                    help='rank of the worker')
parser.add_argument('--log_interval', type=int, default=10, metavar='N',
                    help='interval between training status logs')
parser.add_argument('--gamma', type=float, default=0.99, metavar='G',
                    help='how much to value future rewards')
parser.add_argument('--seed', type=int, default=1, metavar='S',
                    help='random seed  for reproducibility')
parser.add_argument('--num_sampler', type=int, default=10, metavar='S',
                    help='number of samplers')
parser.add_argument('--queue_size', type=int, default=10, metavar='S',
                    help='sampler queue size')
"""
from tgl.modules import *
from tgl.sampler import *
from tgl.utils import *
from sklearn.metrics import average_precision_score, roc_auc_score
import torch
import time
import random
import dgl
import numpy as np
from sklearn.metrics import average_precision_score, roc_auc_score
from torch.nn.parallel import DistributedDataParallel as DDP
import matplotlib.pyplot as plt
import math
def main():   
    sample_param, memory_param, gnn_param, train_param = parse_config('tgl/config/TGN.yml')
    use_cuda = True
    torch.set_num_threads(12)
    args = distparser.args
    rank = distparser._get_worker_rank()
    if use_cuda:
        torch.cuda.set_device(rank)


    DistCustomPool.init_distribution('127.0.0.1',9675,'127.0.0.1',10023,backend = "nccl")
    pdata = partition_load("./startGNN_sample/GDELT", algo="metis")    
    graph = DistGraphData(pdata = pdata,edge_index= pdata.edge_index, full_edge = False)
    print(graph.edge_index[0,:].shape[0],graph.partptr)
    value,counts = torch.unique(graph.edge_index[1,:].view(-1),return_counts=True)
    counts,ind = counts.sort(descending=True)
    #print(counts[:500])
    
    plt.plot(counts)
    plt.savefig('degree-{}.png'.format(rank))
    #print(counts)
    #node = counts
    #plt.hist(node)
    #plt.title('Histogram of Random Data')
    #plt.xlabel('Value')
    #plt.ylabel('Freuency')
    print(counts[:2048])
    ##print(node.shape[1],node.unique(dim=1).shape[1])


# 添加标题和轴标签
    #plt.title('Histogram of Random Data')
    #plt.xlabel('Value')
    #plt.ylabel('Frequency')

    DistCustomPool.close_distribution()
if __name__ == "__main__":
    main()

