from stochastic import *
import random
import time

GraphSize = [20]
RepeatNum = 20
RandSeed = np.arange(1,RepeatNum+1)

ErGot_Sigma1= np.zeros((len(GraphSize),RepeatNum,2))
ErGot_Sigma2= np.zeros((len(GraphSize),RepeatNum,2))
ErGot_Sigma3= np.zeros((len(GraphSize),RepeatNum,2))
ErGot_Sigma4= np.zeros((len(GraphSize),RepeatNum,2))

Got_Sigma1= np.zeros((len(GraphSize),RepeatNum,2))
Got_Sigma2= np.zeros((len(GraphSize),RepeatNum,2))
Got_Sigma3= np.zeros((len(GraphSize),RepeatNum,2))
Got_Sigma4= np.zeros((len(GraphSize),RepeatNum,2))

GW = np.zeros((len(GraphSize),RepeatNum,2))

for ID in 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())
        # ## ErGOT
        # sbm_x_inv, sbm_y_inv, sbm_P, TotalTimeCost = find_permutation(sbm_x, sbm_y, it=10, tau=2, n_samples=10, epochs=1000, lr=0.5, loss_type = 's', epsilon=500, alpha = 0.2, ones = 1, PlotLoss=False, verbose=True)
        # L2Distance = evaluation(sbm_x, sbm_y, sbm_P)
        # ErGot_Sigma1[ID,IDR,0]=L2Distance ## L2 distance
        # ErGot_Sigma1[ID,IDR,1]=TotalTimeCost ## Time cost for running

        # sbm_x_inv, sbm_y_inv, sbm_P, TotalTimeCost = find_permutation(sbm_x, sbm_y, it=10, tau=2, n_samples=10, epochs=1000, lr=0.5, loss_type = 's', epsilon=500, alpha = 0.2, ones = 2, PlotLoss=False, verbose=True)
        # L2Distance = evaluation(sbm_x, sbm_y, sbm_P)
        # ErGot_Sigma2[ID,IDR,0]=L2Distance ## L2 distance
        # ErGot_Sigma2[ID,IDR,1]=TotalTimeCost ## Time cost for running

        # sbm_x_inv, sbm_y_inv, sbm_P, TotalTimeCost = find_permutation(sbm_x, sbm_y, it=10, tau=2, n_samples=10, epochs=1000, lr=0.5, loss_type = 's', epsilon=500, alpha = 0.2, ones = 3, PlotLoss=False, verbose=True)
        # L2Distance= evaluation(sbm_x, sbm_y, sbm_P)
        # ErGot_Sigma3[ID,IDR,0]=L2Distance ## L2 distance
        # ErGot_Sigma3[ID,IDR,1]=TotalTimeCost ## Time cost for running

        # sbm_x_inv, sbm_y_inv, sbm_P, TotalTimeCost = find_permutation(sbm_x, sbm_y, it=10, tau=2, n_samples=10, epochs=1000, lr=0.5, loss_type = 's', epsilon=500, alpha = 0.2, ones = 4, PlotLoss=False, verbose=True)
        # L2Distance = evaluation(sbm_x, sbm_y, sbm_P)
        # ErGot_Sigma4[ID,IDR,0]=L2Distance ## L2 distance
        # ErGot_Sigma4[ID,IDR,1]=TotalTimeCost ## Time cost for running

        ## GOT
        sbm_x_inv, sbm_y_inv, sbm_P, TotalTimeCost = find_permutation(sbm_x, sbm_y, it=10, tau=2, n_samples=10, epochs=1000, lr=0.5, loss_type = 'w', epsilon=500, alpha = 0.2, ones = 1, PlotLoss=False, verbose=True)
        L2Distance = evaluation(sbm_x, sbm_y, sbm_P)
        Got_Sigma1[ID,IDR,0]=L2Distance ## L2 distance
        Got_Sigma1[ID,IDR,1]=TotalTimeCost ## Time cost for running

        sbm_x_inv, sbm_y_inv, sbm_P, TotalTimeCost = find_permutation(sbm_x, sbm_y, it=10, tau=2, n_samples=10, epochs=1000, lr=0.5, loss_type = 'w', epsilon=500, alpha = 0.2, ones = 2, PlotLoss=False, verbose=True)
        L2Distance = evaluation(sbm_x, sbm_y, sbm_P)
        Got_Sigma2[ID,IDR,0]=L2Distance ## L2 distance
        Got_Sigma2[ID,IDR,1]=TotalTimeCost ## Time cost for running

        sbm_x_inv, sbm_y_inv, sbm_P, TotalTimeCost = find_permutation(sbm_x, sbm_y, it=10, tau=2, n_samples=10, epochs=1000, lr=0.5, loss_type = 'w', epsilon=500, alpha = 0.2, ones = 3, PlotLoss=False, verbose=True)
        L2Distance = evaluation(sbm_x, sbm_y, sbm_P)
        Got_Sigma3[ID,IDR,0]=L2Distance ## L2 distance
        Got_Sigma3[ID,IDR,1]=TotalTimeCost ## Time cost for running

        sbm_x_inv, sbm_y_inv, sbm_P, TotalTimeCost = find_permutation(sbm_x, sbm_y, it=10, tau=2, n_samples=10, epochs=1000, lr=0.5, loss_type = 'w', epsilon=500, alpha = 0.2, ones = 4, PlotLoss=False, verbose=True)
        L2Distance = evaluation(sbm_x, sbm_y, sbm_P)
        Got_Sigma4[ID,IDR,0]=L2Distance ## L2 distance
        Got_Sigma4[ID,IDR,1]=TotalTimeCost ## Time cost for running

        # ## GW
        # StartTime = time.time()
        # sbm_P = find_gw_distances(create_distance_matrix(sbm_x), create_distance_matrix(sbm_y), 2e-3)
        # CurrentTime=time.time()
        # L2Distance = evaluation(sbm_x, sbm_y, sbm_P)
        # GW[ID,IDR,0]=L2Distance ## L2 distance
        # GW[ID,IDR,1]=CurrentTime-StartTime ## Time cost for running

results = {'ErGOT L^{-1}':ErGot_Sigma1, 
            'ErGOT L':ErGot_Sigma2, 
            'ErGOT \exp(-2\alpha L)':ErGot_Sigma3, 
            'ErGOT L^2':ErGot_Sigma4, 
            'GOT L^{-1}':Got_Sigma1, 
            'GOT L':Got_Sigma2, 
            'GOT \exp(-2\alpha L)':Got_Sigma3, 
            'GOT L^2':Got_Sigma4, 
            'GW':GW}
np.save(f'exp_results-20.npy', results)