from stochastic import *
import random
import time

n = 40
RepeatNum = 50
ps = [1 - 0., 1 - 0.1, 1 - 0.2, 1 - 0.3, 1 - 0.4, 1 - 0.5, 1 - 0.6]
RandSeed = np.arange(1,RepeatNum+1)

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

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

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

blocks = [int(n/4), int(n/4), int(n/4), int(n/4)]
for i in range(len(ps[:1])):
    ID = ps[i]
    probs = [[0.5, ID, ID, ID],[ID, 0.5, ID, ID], [ID, ID, 0.5, ID], [ID, ID, ID, 0.5]]
    
    for IDR in range(RepeatNum):
        print('Prob: %f - Iteration: %f ' % (ID, 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[i,IDR,0]=L2Distance ## L2 distance
        # Got_Sigma1[i,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[i,IDR,0]=L2Distance ## L2 distance
        # Got_Sigma2[i,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[i,IDR,0]=L2Distance ## L2 distance
        # Got_Sigma3[i,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[i,IDR,0]=L2Distance ## L2 distance
        # Got_Sigma4[i,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-2)
        CurrentTime=time.time()
        L2Distance = evaluation(sbm_x, sbm_y, sbm_P * len(sbm_P))
        GW[i,IDR,0]=L2Distance ## L2 distance
        GW[i,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'gw_exp_results_sbm.npy', results)
    print(GW.mean(axis=1))