rm(list=ls())
timestart<-Sys.time()
library(MASS)

gendata <- function(n, sigma_x, sigma_e, beta_0){
  #beta_0: coefficients contain intercept term
  #n: sample size
  #sigma_x: variance for x
  #sigma_e: variance for error term
  p <- length(beta_0) - 1
  x <- mvrnorm(n, rep(0,p), sigma_x*diag(1,p))
  if(n == 1){
    z <- c(1, x)
    pr <- exp(as.numeric(z%*%beta_0))/(1 + exp(as.numeric(z%*%beta_0)))
    y <- rbinom(n,1,pr)
  }else{
    z <- cbind(1, x)
    pr <- exp(z%*%beta_0)/(1 + exp(z%*%beta_0))
    y <- rbinom(n,1,pr)
  }
  dat <- list(x = x, z = z, y = y)
  return(dat)
}


# auxiliary functions
# first derivative for huber loss
invlogit <- function(u){return(1/(1+exp(-u)))}


############ Mallows weights ##############
####### min(1, 2/||x_i||_2^2) ############# 
###########################################
weightfn <- function(x,max.norm=sqrt(2)){
  return(min(1,(max.norm/norm(x,type="2"))^2))
}


Noisysgd <- function(n, x, y, private, mu, beta0, stepsize)
 {
## (x, y): one data point: 
## x: covariate that does not contain intercept term
## y: response variable
## n: sample size, if we use sgd, n=1
## private: whether add privacy
## mu: privacy budget
## s0: variance for error term
## bta0: coefficients.
   p <- length(x)
   grad_noise <- 0 
   priv_grad_traj <- NA
   # gets updated further down if private = T
   ##########################################
   if(n == 1){
    z <- c(1, x)
    yhat <- invlogit(as.numeric(z%*%beta0))
   }else{
    z <- cbind(1, x)
    yhat <- invlogit(as.vector(z%*%beta0))
   }

   diffs <- y - yhat
   weightfun <- weightfn(z, max.norm = sqrt(2))
   
    if(private == T){  
     GS <- 2*sqrt(2)   
     grad_noise <- GS/mu
    }
    noisy_grad <- diffs*weightfun*z + grad_noise*rnorm(p+1)
    #tracks the evolution of the noisy gradient (in L2 norm)
    priv_grad_traj <- sqrt(sum(noisy_grad^2)) 
 # this performs stochastic gradient descent to estimate beta
 ###############################################
    beta <- beta0 + stepsize*noisy_grad  

   out <- NULL
   out$beta <- beta
   out$grad_traj <- priv_grad_traj
   return(out)
}



########################
########################
##### run ##############
##### true value #######
# #beta_true <- c(1.5,1,-1,0.5)
# beta_true <- rep(1,5) #c(1,1,1,1)
# sigma_x <- 1
# sigma_e <- 2
# MAX <- 1
# n <- 100000
# mu <- 2
# alpha <- 0.5
# ##### null value #######
# p <- length(beta_true) - 1
# beta_all <- array(0, dim = c((p+1), n, MAX))

# beta <- array(0, dim = c(p+1, MAX))

# beta_bar <- array(0, dim = c(p+1, MAX))

# ##### initial value ####
# beta_0 <- rep(0, p+1)
# stepsize <- n^(-alpha)
# n_drop <- 0

# for(tt in 1:MAX){
#  set.seed(20230907 + tt)
#  for(i in 1:n){
#         data <- gendata(1, sigma_x, sigma_e, beta_true)
#         x <- data$x
#         y <- data$y
#         ####### calculate noisy sgd ########  
#         Noisysgd_out <- Noisysgd(1, x, y, private=T, mu, beta_0, stepsize)
#         beta_all[, i, tt] <- Noisysgd_out$beta
#         ####### iterate ####################
#         beta_0 <- beta_all[, i, tt] 
#         if(i>n_drop){
#           beta_bar[,tt] <- (i-n_drop-1)*beta_bar[,tt]/(i-n_drop) + beta_all[, i, tt]/(i-n_drop)
#         }

#       }
#   beta[, tt] <- beta_all[, n, tt]
#   print(tt)
# }






