source('PSGD_logistic_regression.R')

p_values = c(5, 10, 20)
sigma_x <- 1
sigma_e <- 2
gap = 5000
All = 300000
n_values <- seq(gap, All, gap)
mu_values <- c(0.5,1,2,3)
alpha_values <- c(1/3, 1/2, 2/3, 1)
MAX <- 200
n_drop <- 0

for(p in p_values){
    for(mu in mu_values){
      for(alpha in alpha_values){
          
          beta_true <- rep(1,(p+1))

          beta_all <- array(0, dim = c((p+1), length(n_values), MAX))
          beta_bar <- array(0, dim = c((p+1), length(n_values), MAX))

          ##### initial value ####

          for(tt in 1:MAX){
            set.seed(20230910 + tt)
            beta_0 <- rep(0, p+1)
            beta_ave <- rep(0, p+1)

            for(i in 1:All){
              stepsize <- 0.2*i^(-alpha)
              data <- gendata(1, sigma_x, sigma_e, beta_true)
              x <- data$x
              y <- data$y

              ####### calculate noisy sgd ########  
              Noisysgd_out <- Noisysgd(1, x, y, TRUE, mu, beta_0, stepsize)

              beta_new <- Noisysgd_out$beta
              beta_0 <- beta_new
              #beta_all[, i, tt] <- Noisysgd_out$beta

              ####### iterate ####################
              #beta_0 <- beta_all[, i, tt] 
              if(i>n_drop){
                beta_ave <- (i-n_drop-1)*beta_ave/(i-n_drop) + beta_new/(i-n_drop)
              }
              if(i%%gap==0){
                  beta_all[, i%/%gap, tt] = beta_new
                  beta_bar[, i%/%gap, tt] = beta_ave
              }
            }

          }
            

          # L2 norms at specific time points for beta_all and beta_bar
            l2_norms_beta_all <- matrix(0, length(n_values), MAX)
            l2_norms_beta_bar <- matrix(0, length(n_values), MAX)

            for (k in 1:MAX) {
                l2_norms_beta_all[, k] <- sapply(1:length(n_values), function(j) sum((beta_all[, j, k] - beta_true)^2))
                l2_norms_beta_bar[, k] <- sapply(1:length(n_values), function(j) sum((beta_bar[, j, k] - beta_true)^2))
            }

          # Create a dataframe for writing
          df1 <- data.frame(l2_norms_beta_all)
          df2 <- data.frame(l2_norms_beta_bar)

          # Write results to CSV for this combination of parameters
          filename1 <- paste0("results/logistic_results_last_mu=", mu, "_alpha=", alpha, "_p=", p, ".csv")
          write.csv(df1, file=filename1, row.names=FALSE)
          filename2 <- paste0("results/logistic_results_ave_mu=", mu, "_alpha=", alpha,  "_p=", p, ".csv")
          write.csv(df2, file=filename2, row.names=FALSE)    

          filename3 <- paste0("beta/logistic_results_last_mu=", mu, "_alpha=", alpha,  "_p=", p, ".RData")
          save(beta_all, file=filename3)
          filename4 <- paste0("beta/logistic_results_ave_mu=", mu, "_alpha=", alpha,  "_p=", p, ".RData")
          save(beta_bar, file=filename4)    



    }
  }
}