import numpy as np
import scipy.sparse as sp
import numba
from katz_solver import *
large_graph = ['com-friendster', 'ogbn-papers100M']
graph_names = ['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', 'ogbn-papers100M', 'com-friendster']

result_names = ['local_sor_opers', 'local_sor_algo_times', 'global_sor_opers',
            'global_sor_algo_times', 'local_gs_opers', 'local_gs_algo_times',
            'global_gs_opers', 'global_gs_algo_times', 'local_gd_opers',
            'local_gd_algo_times', 'global_gd_opers', 'global_gd_algo_times']
path = './dataset/'
all_result = {}
lambda_result = {}
ppr_result = np.load('./results/ppr_exp_result_new.npy', allow_pickle=True)
ppr_result = ppr_result.reshape(1, -1)[0][0]

lambda_result = np.load('./results/katz_exp_lambda.npy', allow_pickle=True)
lambda_result = lambda_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')
    degree = np.array(adj_matrix.sum(1)).flatten()
    indices = adj_matrix.indices
    indptr = adj_matrix.indptr
    n = len(indptr) - 1
    m = len(indices)
    lambda_1 = lambda_result[graph_name]
    alpha = 1. / (lambda_1 + 1)
    omega = 2. / (1 + np.sqrt(1 - (alpha * lambda_1) ** 2))
    eps = 1 / m
    np.random.seed(17)
    s_nodes = ppr_result[graph_name][0]
    print(graph_name, lambda_1, alpha, omega, end=':\n')
    graph_result = solve_a_graph_all(n, indptr, indices, degree, alpha, eps, omega, graph_name, s_nodes)
    for i in range(len(result_names)):
        print(result_names[i], ':', graph_result[i + 1].mean())
    print()
    all_result[graph_name] = graph_result
    np.save('./results/katz_exp_result_new.npy', all_result)
