source('PSGD_linear_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){
    beta_true <- rep(1,(p+1))
    theta_true = c(beta_true, sigma_e)
    for(mu in mu_values){
      for(alpha in alpha_values){

          theta_all <- array(0, dim = c((p+2), length(n_values), MAX))
          theta_bar <- array(0, dim = c((p+2), length(n_values), MAX))

          ##### initial value ####
          
    
              for(tt in 1:MAX){
                set.seed(20230910 + tt)
                beta_new <- rep(0, p+1)
                beta_ave <- rep(0, p+1)
                s_ave = 1
                s_new = 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
                    
                  Noisysgd_out <- Noisysgd(1, x, y, k=1.345, fisher_beta=0.7101645, scale=T, private=T, mu, s_new, beta_new, stepsize)
                  beta_new <- Noisysgd_out$beta
                  s_new <- Noisysgd_out$s
                  if(i>n_drop){
                      beta_ave <- (i-n_drop-1)*beta_ave/(i-n_drop) + beta_new/(i-n_drop)
                      s_ave <- (i-n_drop-1)*s_ave/(i-n_drop) + s_new/(i-n_drop)
                  }
                  if(i %% gap == 0){
                    theta_all[, i%/%gap, tt] = c(beta_new, s_new)
                    theta_bar[, i%/%gap, tt] = c(beta_ave, s_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((theta_all[, j, k] - theta_true)^2))
                l2_norms_beta_bar[, k] <- sapply(1:length(n_values), function(j) sum((theta_bar[, j, k] - theta_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/linear_results_last_mu=", mu, "_alpha=", alpha, "_p=", p, ".csv")
          write.csv(df1, file=filename1, row.names=FALSE)
          filename2 <- paste0("results/linear_results_ave_mu=", mu, "_alpha=", alpha,  "_p=", p, ".csv")
          write.csv(df2, file=filename2, row.names=FALSE)    

          filename3 <- paste0("beta/linear_results_last_mu=", mu, "_alpha=", alpha,  "_p=", p, ".RData")
          save(theta_all, file=filename3)
          filename4 <- paste0("beta/linear_results_ave_mu=", mu, "_alpha=", alpha,  "_p=", p, ".RData")
          save(theta_bar, file=filename4)    



    }
  }
}