import networkx as nx
import torch
import numpy as np
from fluidc import asyn_fluidc
import sys
import pickle
data = str(sys.argv[1])
if data == 'citeseer' or data == 'cora':
    num_local = 8
elif data == 'pubmed':
    num_local = 10

np.random.seed(123)
device = 'cpu'
adj = pickle.load(open('data/'+data + '_adj.pkl', 'rb'))
feature = pickle.load(open('data/'+data + '_feat.pkl', 'rb')).toarray()
net = nx.from_scipy_sparse_matrix(adj)
nets = sorted(nx.connected_components(net), key=len, reverse=True)
if len(nets) > 1:
    lgst_net = nx.subgraph(net, nets[0])
    com = asyn_fluidc(lgst_net, k=num_local)
else:
    com = asyn_fluidc(net, k=num_local)
anchor_nodes = []

for i, c in enumerate(com):
    if i == num_local:
        break
    subnet = nx.subgraph(net, c)
    anchor_node = sorted(subnet.degree(), key=lambda x : x[1])[0][0]
    anchor_nodes.append(anchor_node)


torch.save(anchor_nodes, 'data/'+data+'_anchor_async_fluid.pt')
pos_enc = []
for node in range(feature.shape[0]):
    temp = []
    if node % 1000 == 0:
        print (f"{node} node completed!!")
    for anchor_node in anchor_nodes:
        try:
            temp.append(len(nx.algorithms.shortest_path(net, node, anchor_node)))
        except:
            temp.append(0)
    pos_enc.append(np.array(temp))
pos_enc = np.array(pos_enc)
max_d = 0
for p in pos_enc:
    max_d = max(max_d, max(p))

np.where(pos_enc == 0, max_d, pos_enc)
for i, a in enumerate(anchor_nodes):
    pos_enc[a][i] = 0
torch.save(pos_enc, 'data/'+data+'_pe_async_fluid.pt')
torch.save(com, 'data/'+data+'_com_async_fluid.pt')
