rm(list=objects())
setwd("C:/Users/clément/Downloads/optimality_of_variational_infe-Supplementary Material/Supplementary Materials/Numerical experiments/Simulations/Results")
for (type in c(1,2,3)){
  ##################### -  Set parameters - #########
  K <- 3
  M <- 100 # number of experiments############
  N_list <- round(exp(seq(log(50), log(1000), length.out = 10)))
  
  # set Q the probabilities of connections between communities
  if (type == 1){ # assortative SBM
    Q <- matrix(c(0.5,   0.2,  0.2,
                  0.2,   0.5,  0.2,
                  0.2,  0.2,  0.5), nrow = 3)
  }else if (type == 2){ # disassortative SBM
    Q <- matrix(c( 0.2,  0.5,  0.5,
                   0.5,  0.2,  0.5,
                   0.5,  0.5,  0.2), nrow = 3)
  }else{ # mixed model
    Q <- matrix(c( 0.1,  0.5,  0.3,
                   0.5,  0.2,  0.4,
                   0.3,  0.4,  0.6), nrow = 3)
  }
  
  # set alpha the populations of the communities
  if (type %in% 1:2){
    alpha <- rep(1/K, K)
  }else{
    alpha <- c(0.1, 0.3, 0.6)
  }
  
  ##################### -  Experiments - #####################
  for (N in N_list){
    set.seed(0)
    start_time <- Sys.time()
    
    # record the errors 
    Prop_Error_Z <- rep(NA, M)
    for (m in 1:M){
      
      ##################### -  Draw random variables - #####################
      # set Z the label function
      Z <- base::sample(1:K, size = N, replace = T, prob = alpha)
      
      # set A the full adjacency matrix
      Theta <- sapply(1:N, function(i) sapply(1:N, function(j) Q[Z[i], Z[j]])) # SBM matrix
      diag(Theta) <- 0
      A_undir <- rbinom(n=N*(N-1)/2, size=1, prob=Theta[upper.tri(Theta)]) # draw edges
      A <- matrix(0,N,N)
      A[upper.tri(A)] <- A_undir
      A <- (A+t(A)) # adjacency matrix
      
      # Sample edges
      Omega_undir <- rbinom(n=N*(N-1)/2, size=1, prob=0.5) # draw edges
      Omega <- matrix(0,N,N)
      Omega[upper.tri(Omega)] <- Omega_undir
      Omega <- (Omega+t(Omega)) # adjacency matrix
      
      # set A_obs the observed adjacency matrix
      A_obs <- A
      diag(A_obs) <- NA
      A_obs[Omega == 0] <- NA
     
      ##################### -  Estimate network using missSBM - #####################
      estimator_missSBM <- missSBM::estimateMissSBM(
        adjacencyMatrix = A_obs, 
        vBlocks = c(K),
        sampling = "dyad",
        control = list(trace = 0))$bestModel$fittedSBM
      
      ##################### -  Estimate labels z - #####################
      z_est <- estimator_missSBM$memberships
      
      ##################### -  Best label among permutations - #####################
      perm_list <- list(c(1,2,3), c(1, 3, 2), c(3, 2, 1), c(2, 1, 3), c(3, 1, 2), c(2, 3, 1))
      erreur <- N
      for (p in 1:6){
        z_est_p <- perm_list[[p]][z_est]
        erreur <- min(erreur, sum(Z != z_est_p))
      }
      Prop_Error_Z[m] <- erreur/N
    }
    
    results <- list(N = N, K = K, fixed_K = TRUE, type = type,
                    Prop_Error_Z = Prop_Error_Z)
    
    path <- "Missclassified"
    
    if(type == 1){
      path <- paste0(path, "/Assortative SBM")
    }else if(type == 2){
      path <- paste0(path, "/Disassortative SBM")
    }else{
      path <- paste0(path, "/Mixed SBM")
    }
    path <- paste0(path, "/N_",N ,".RDS")
    saveRDS(results, file = path)
    
    print(path)
    end_time <- Sys.time()
    print(end_time - start_time)
    print(mean(Prop_Error_Z))
  }
}
