rm(list=objects())
setwd("~/Desktop/Experiments/Simulations/Results")

type <- 1
path <- "Varying K"

# Load results
if(type == 1){
  path <- paste0(path, "/Assortative SBM")
}else{
  path <- paste0(path, "/Disassortative SBM")
}


K_list <- 2:15
M <- 100 # number of experiments

# record the errors 
Error_missSBM <- matrix(NA, ncol = 3, nrow = 14)
Error_Var <- matrix(NA, ncol = 3, nrow = 14)
Error_true_Z <- matrix(NA, ncol = 3, nrow = 14)
Error_softImpute <- matrix(NA, ncol = 3, nrow = 14)

for (K in 2:15){
  res <- readRDS(paste0(path, "/K_", K,".RDS"))
  Error_missSBM[K-1,] <- quantile(res$Error_missSBM, probs = c(0.25, 0.5, 0.75))
  Error_Var[K-1,] <- quantile(res$Error_Var, probs = c(0.25, 0.5, 0.75))
  Error_softImpute[K-1,] <- quantile(res$Error_softImpute, probs = c(0.25, 0.5, 0.75))
  Error_true_Z[K-1,] <- quantile(res$Error_true_Z, probs = c(0.25, 0.5, 0.75))
}

Error_av <- matrix(NA, ncol = 3, nrow = 14)

plot(x = 2:15, y = Error_softImpute[,2], type = 'l', col = 'green', 
      xlab = "", ylab = "", ylim = c(0,0.04))
lines(x = 2:15, y = Error_softImpute[,1], lty = 3, col = 'green')
lines(x = 2:15, y = Error_softImpute[,3], lty = 3, col = 'green')

lines(x = 2:15, y = Error_true_Z[,2], col = 'black')
lines(x = 2:15, y = Error_true_Z[,1], lty = 3, col = 'black')
lines(x = 2:15, y = Error_true_Z[,3], lty = 3, col = 'black')

lines(x = 2:15, y = Error_Var[,2], col = 'blue')
lines(x = 2:15, y = Error_Var[,1], lty = 3, col = 'blue')
lines(x = 2:15, y = Error_Var[,3], lty = 3, col = 'blue')

lines(x = 2:15, y = Error_missSBM[,2], col = 'red')
lines(x = 2:15, y = Error_missSBM[,1], lty = 3, col = 'red')
lines(x = 2:15, y = Error_missSBM[,3], lty = 3, col = 'red')



for (K in 2:15){
  Q <- matrix(0.1, ncol = K, nrow = K)
  diag(Q) <- 0.4
  N <- 300
  alpha <- rep(1/K, K)
  
  set.seed(0)
  start_time <- Sys.time()
  
  # record the errors 
  Error_av_M <- 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
    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 softImpute - #####################
    estimate_Theta_av <- matrix(mean(A_obs, na.rm = T), ncol = N, nrow = N)
    diag(estimate_Theta_av) <- 0
    Error_av_M[m] <- sum((estimate_Theta_av - Theta)**2, na.rm = F)/N**2
  }
  Error_av[K-1,] <- quantile(Error_av_M, probs = c(0.25, 0.5, 0.75))
  print(K)
}

lines(x = 2:15, y = Error_av[,2], col = 'orange')
lines(x = 2:15, y = Error_av[,1], lty = 3, col = 'orange')
lines(x = 2:15, y = Error_av[,3], lty = 3, col = 'orange')