import numpy as np
import scipy.sparse as sp
import numba
from ppr_solver import *
import sys
import os

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']

graph_names = ['wiki-talk', 'Citeseer', 'ogbn-arxiv', 'com-youtube']
path = './dataset/'
test_num = 50

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:
    all_result = {}
    graph_path = path + graph_name + '/'
    adj_matrix = sp.load_npz(graph_path + graph_name + '_csr-mat.npz')
    indices = adj_matrix.indices
    indptr = adj_matrix.indptr
    n = len(indptr) - 1
    degree = np.array(adj_matrix.sum(1)).flatten()
    np.random.seed(17)
    s_nodes = exp_result[graph_name][0]
    alphas = [0.5, 0.2, 0.05, 0.01]
    for i in range(len(alphas)):
        alpha = alphas[i]
        alpha = alpha / (2 - alpha)
        mu = (1. - alpha) / (1. + alpha)
        omega = 1. + (mu / (1. + np.sqrt(1. - mu ** 2.))) ** 2.
        alpha = 2 * alpha / (1 + alpha)
        eps = np.power(2.0, 8) / n
        mineps = np.power(2.0, -12) / n
        np.random.seed(17)
        s_nodes = exp_result[graph_name][0]
        while True:
            if eps >= 1 / degree.max():
                eps /= 2
                continue
            print('ppr', graph_name, eps, end=' ')
            graph_result = solve_a_graph_all(n, indptr, indices, degree, alpha, eps, omega, graph_name, s_nodes)
            print(graph_result[2].mean(), graph_result[4].mean(), graph_result[6].mean(),
                  (graph_result[4] / graph_result[2]).mean())
            all_result[eps] = graph_result
            np.save('./results/' + graph_name + '/ppr_exp_ratio_result_new_alpha' + str(alpha) + '.npy', all_result)
            eps /= 2
            if eps < mineps:
                break

