from stochastic import *
import random
import time
from tqdm import tqdm 
GraphSize = [20, 30, 40, 50, 60, 70, 80, 90, 100, 110, 120]

epsilons = np.arange(11) * 1e-4 + 1e-3
RepeatNum = 20

RandSeed = np.arange(1,RepeatNum+1)

GW = np.zeros((len(epsilons), len(GraphSize),RepeatNum,2))
for i in range(len(epsilons)):
    eps = epsilons[i]
    for ID in tqdm(range(len(GraphSize))):
        n = GraphSize[ID]
        blocks = [int(n/2),int(n/2)]
        probs = [[0.7,0.05],[0.05,0.7]]
    
        for IDR in range(RepeatNum):
            print('Graph Size: %f - Iteration: %f ' % (n, IDR))
            [sbm_x, sbm_y, sbm_P_true] = create_test_graphs(n, blocks, probs, graph_type = 'sbm', seed_nb = RandSeed[IDR].item())
    
            ## GW
            StartTime = time.time()
            sbm_P = find_gw_distances(create_distance_matrix(sbm_x), create_distance_matrix(sbm_y), epsilon=eps)
            CurrentTime=time.time()
            L2Distance = evaluation(sbm_x, sbm_y, sbm_P)
            GW[i,ID,IDR,0]=L2Distance ## L2 distance
            GW[i,ID,IDR,1]=CurrentTime-StartTime ## Time cost for running

        np.save(f'align_result/gw_parameter_sensitivity.npy', GW)