import numpy as np
import torch
import scipy.sparse as sp
import numba
from ppr_solver import *
import warnings

warnings.filterwarnings("ignore")
device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')

all_graph = ['Cora', 'Citeseer', 'ogbn-arxiv', 'as-skitter', 'ogbn-proteins', 'com-orkut', 'cit-patent', 'ogbl-ppa', 'ogbn-products',
               'wiki-talk', 'com-youtube', 'ogbn-mag', 'soc-lj1', 'reddit', 'pubmed', 'wiki-en21',
               'com-friendster', 'ogbn-papers100M', 'ogb-mag240m']

graph_names = ['ogbn-arxiv', 'wiki-talk', 'Citeseer', 'wiki-en21', 'com-orkut']
path = './dataset/'
alpha = 0.1
all_result = {}
exp_result = np.load('./results/ppr_exp_result_new.npy', allow_pickle=True)
exp_result = exp_result.reshape(1, -1)[0][0]
for graph_name in graph_names:
    graph_path = path + graph_name + '/'
    adj_matrix = sp.load_npz(graph_path + graph_name + '_csr-mat.npz')
    indices = adj_matrix.indices.astype(np.int64)
    indptr = adj_matrix.indptr.astype(np.int64)
    data = adj_matrix.data
    n = len(indptr) - 1
    m = len(indices)
    degree = np.array(adj_matrix.sum(1)).flatten()
    eps = 1 / n
    s_nodes = exp_result[graph_name][0]
    print(graph_name, end=': ')
    graph_result = solve_a_graph_gpu(n, indptr, indices, degree, alpha, eps, s_nodes)
    print(graph_result[1].mean(), graph_result[2].mean())
    all_result[graph_name] = graph_result
    np.save('./results/ppr_exp_cuda_result_new.npy', all_result)
