# This code is used to calculate the asymptotic variances of GMM estimators.
# Key parameters for this code:
# n=number of primary cluster;
# ni=number of secondary cluster within each primary cluster;
# nij=number of observations within each secondary cluster;
# beta=covariates of interest;
# sigma1=variance of random effect at level 3;
# sigma2=variance of random effect at level 2;

################################################################

library(rootSolve) # Package used for root-finding by Newton Raphson method
library(compiler) # Package used for speeding up the code
enableJIT(3) # From the library "complier"
library(foreach) # Package used for parallel processing
library(doParallel) # Package used for parallel processing
# Find out how many cores are available (if you don't already know)
detectCores()
## [1] 4
# Create cluster with desired number of cores
cl <- makeCluster(3)
# Register cluster
registerDoParallel(cl)
# Find out how many cores are being used
getDoParWorkers()

random1<-function(n){ # Function used to generate realizations of standard normal distribution for random effect at level 3.
  r<-matrix(0,n,2500)
  for (i in 1:n){
    r[i,]<-rnorm(2500)
  }
  return(r)
}

random2<-function(n,ni){  # Function used to generate realizations of standard normal distribution for random effect at level 2.
  r<-rep(0,n*ni*2500)
  dim(r)<-c(n,ni,2500)
  for (i in 1:n){
    for (j in 1:ni){
      r[i,j,]<-rnorm(2500)
    }
  }
  return (r)
}

gendata<-function(n,ni,nij,sigma1,sigma2,seed){ # Function used to generate simulated data.
  set.seed(seed)
  N<-n*ni*nij
  x1<-vector(length=N)
  x2<-vector(length=N)
  x3<-vector(length=N)
  x4<-vector(length=N)
  x5<-vector(length=N)
  xmatrix<-cbind(x1,x2,x3,x4,x5)
  for (i in 1:n){
    xmatrix[((i-1)*ni*nij+1):(i*ni*nij),1]<-i
  }
  x21<-vector(length=ni*nij)
  for (i in 1:ni){
    x21[((i-1)*nij+1):(i*nij)]<-i
  }
  xmatrix[,2]<-rep(x21,n)
  for (i in 1:N){
    xmatrix[i,3]<-i%%nij
    if (xmatrix[i,3]==0){
      xmatrix[i,3]<-nij
    }
  }
  for (i in 1:N){
    if ((xmatrix[i,2]<=(ni/2))&(xmatrix[i,3]<=(nij/2))){
      xmatrix[i,4]<-1 
      xmatrix[i,5]<-0.5
    }else if ((xmatrix[i,2]<=(ni/2))&(xmatrix[i,3]>(nij/2))){
      xmatrix[i,4]<-0 
      xmatrix[i,5]<--0.5
    }else if ((xmatrix[i,2]>(ni/2))&(xmatrix[i,3]<=(nij/2))){
      xmatrix[i,4]<-0
      xmatrix[i,5]<--1.5
    }else if ((xmatrix[i,2]>(ni/2))&(xmatrix[i,3]>(nij/2))){
      xmatrix[i,4]<-1 
      xmatrix[i,5]<--1
    }
  }
  r1<-rnorm(n)
  r2<-rnorm(n*ni)
  ri1<-vector(length=N)
  ri2<-vector(length=N)
  p<-vector(length=N)
  for (i in 1:n){
    ri1[((i-1)*ni*nij+1):(i*ni*nij)]<-r1[i]
  }
  for (i in 1:(n*ni)){
    ri2[((i-1)*nij+1):(i*nij)]<-r2[i]
  }
  for (i in 1:N){
    p[i]<-exp(xmatrix[i,4]+xmatrix[i,5]+sigma1*ri1[i]+sigma2*ri2[i])/(1+exp(xmatrix[i,4]+xmatrix[i,5]+sigma1*ri1[i]+sigma2*ri2[i])) # We generate simulated probabilities here.
  }
  y<-rbinom(N,1,p) # We generate our responses based on simulated probabilities.
  return (list(y=y,X=xmatrix))
}

solve1<-function(data,beta,n,ni,nij,sigma1,sigma2,seed){  # This function corresponds to the first GMM estimating eqaution mentioned in the model. It will be used to update beta.
  set.seed(seed)
  r1<-random1(n)
  r2<-random2(n,ni)
  miu1<-rep(0,2500)
  miu<-rep(0,n*ni*nij)
  dim(miu)<-c(n,ni,nij)
  devy<-rep(0,n*ni*(ncol(data$X)-3)*nij)
  dim(devy)<-c(n,ni,(ncol(data$X)-3),nij)
  key<-rep(0,ni*nij*2500)
  dim(key)<-c(ni,nij,2500)
  
  result<-foreach(i=1:n, .combine='+') %dopar% { # We use parallel processing to increase the efficiency of the code.
    faiy<-matrix(0,ncol(data$X)-3,2*(ncol(data$X)-3))
    for (j in 1:ni){
      for (k in 1:nij){
        key[j,k,]<-1-1/(1+exp(data$X[(i-1)*ni*nij+(j-1)*nij+k,-(1:3)]%*%beta+sigma1*r1[i,]+sigma2*r2[i,j,]))
      }
    }
    for (j in 1:ni){
      for (k in 1:nij){
        miu1<-key[j,k,]
        miu[i,j,k]<-sum(miu1)/2500
        devy[i,j,,k]<-(sum(miu1*(1-miu1))/2500)*data$X[(i-1)*ni*nij+(j-1)*nij+k,-(1:3)]
      }
      for (k in 1:nij){
        faiy[,1:(ncol(data$X)-3)]<-faiy[,1:(ncol(data$X)-3)]-t(t(devy[i,j,,k]))%*%t(devy[i,j,,k])
      }
      faiy[,-(1:(ncol(data$X)-3))]<-faiy[,-(1:(ncol(data$X)-3))]+(devy[i,j,,]%*%(data$y[((i-1)*ni*nij+(j-1)*nij+1):((i-1)*ni*nij+j*nij)]-miu[i,j,]))%*%t(devy[i,j,,]%*%(data$y[((i-1)*ni*nij+(j-1)*nij+1):((i-1)*ni*nij+j*nij)]-miu[i,j,]))
    }
    return(faiy)
  }
  
  return(solve(result[,1:(ncol(data$X)-3)]%*%solve(result[,-(1:(ncol(data$X)-3))])%*%t(result[,1:(ncol(data$X)-3)])))
}

solve2<-function(data,beta,n,ni,nij,sigma1,sigma2,seed){  # This function corresponds to the second GMM estimating eqaution mentioned in the model. It will be used to update sigma2.
  set.seed(seed)
  r1<-random1(n)
  r2<-random2(n,ni)
  miu1<-rep(0,2500)
  miu<-rep(0,n*ni*nij)
  dim(miu)<-c(n,ni,nij)
  thoij<-rep(NA,n*ni*nij*nij)
  dim(thoij)<-c(n,ni,nij,nij)
  theta<-vector(length=n*ni*nij*(nij+1)/2)
  devg<-rep(0,n*ni*nij*(nij+1)/2)
  dim(devg)<-c(n,ni,nij*(nij+1)/2)
  g<-vector(length=n*ni*nij*(nij+1)/2)
  corg<-matrix(0,n*ni*nij*(nij+1)/2,4)
  dijk<-rep(0,n*ni*nij)
  dim(dijk)<-c(n,ni,nij)
  miu2<-rep(0,2500)
  dijj<-rep(NA,n*ni*nij*nij)
  dim(dijj)<-c(n,ni,nij,nij)
  key<-rep(0,ni*nij*2500)
  dim(key)<-c(ni,nij,2500)
  
  seq<-seq(nij,1,by=-1)
  for (i in 1:(nij-1)){
    seq[i+1]<-seq[i]+seq[i+1]
  }
  
  result<-foreach(i=1:n, .combine='+') %dopar% {
    faig<-c(0,0)
    for (j in 1:ni){
      for (k in 1:nij){
        key[j,k,]<-1-1/(1+exp(data$X[(i-1)*ni*nij+(j-1)*nij+k,-(1:3)]%*%beta+sigma1*r1[i,]+sigma2*r2[i,j,]))
      }
    }
    for (j in 1:ni){
      rij<-r2[i,j,]
      for (k in 1:nij){
        miu1<-key[j,k,]
        dijk[i,j,k]<-sum(miu1*(1-miu1)*rij)/2500
        miu[i,j,k]<-sum(miu1)/2500
        for (l in 1:nij){
          if (k<l){
            miu2<-key[j,l,]
            thoij[i,j,k,l]<-sum(miu1*miu2)/2500
            dijj[i,j,k,l]<-sum(miu1*miu2*(1-miu1+1-miu2)*rij)/2500
          }
        }
      }
      for (k in 1:(nij*(nij+1)/2)){
        if (k<=nij){
          g[(i-1)*ni*nij*(nij+1)/2+(j-1)*nij*(nij+1)/2+k]<-data$y[(i-1)*ni*nij+(j-1)*nij+k]*data$y[(i-1)*ni*nij+(j-1)*nij+k]
          corg[k,]<-c(i,j,k,k)
        } else { 
          for (l in 2:nij){
            if ((k<=seq[l])&(k>seq[l-1])){
              g[(i-1)*ni*nij*(nij+1)/2+(j-1)*nij*(nij+1)/2+k]<-data$y[(i-1)*ni*nij+(j-1)*nij+(l-1)]*data$y[(i-1)*ni*nij+(j-1)*nij+(k-seq[l-1]+l-1)]
              corg[k,]<-c(i,j,(l-1),(k-seq[l-1]+l-1))
              break
            }
          }
        }
        if(corg[k,3]==corg[k,4]){
          theta[(i-1)*ni*nij*(nij+1)/2+(j-1)*nij*(nij+1)/2+k]<-miu[i,j,corg[k,3]]
          devg[i,j,k]<-dijk[i,j,k]
        } else {
          theta[(i-1)*ni*nij*(nij+1)/2+(j-1)*nij*(nij+1)/2+k]<-thoij[i,j,corg[k,3],corg[k,4]]
          devg[i,j,k]<-dijj[i,j,corg[k,3],corg[k,4]] 
        }
      }
      faig[1]<-faig[1]-t(devg[i,j,])%*%t(t(devg[i,j,]))
      faig[2]<-faig[2]+(t(devg[i,j,])%*%(g[((i-1)*ni*nij*(nij+1)/2+(j-1)*nij*(nij+1)/2+1):((i-1)*ni*nij*(nij+1)/2+j*nij*(nij+1)/2)]-theta[((i-1)*ni*nij*(nij+1)/2+(j-1)*nij*(nij+1)/2+1):((i-1)*ni*nij*(nij+1)/2+j*nij*(nij+1)/2)]))^2
    }
    return(faig)
  }
  
  return(solve(result[1]*solve(result[2])*result[1]))
}

solve3<-function(data,beta,n,ni,nij,sigma1,sigma2,seed){ # This function corresponds to the third GMM estimating eqaution mentioned in the model. It will be used to update sigma1.
  set.seed(seed)
  r1<-random1(n)
  r2<-random2(n,ni)
  miu<-rep(0,n*ni*nij)
  dim(miu)<-c(n,ni,nij)
  dijk<-rep(0,n*ni*nij)
  dim(dijk)<-c(n,ni,nij)
  miu1<-rep(0,2500)
  miu3<-rep(0,2500)
  thoijj<-rep(NA,n*ni*ni*nij*nij)
  dim(thoijj)<-c(n,ni,ni,nij,nij)
  thoijj1<-rep(0,2500)
  dijj<-rep(NA,n*ni*ni*nij*nij)
  dim(dijj)<-c(n,ni,ni,nij,nij)
  theta<-vector(length=n*ni*(ni+1)/2)
  devg<-rep(0,n*ni*(ni+1)/2)
  dim(devg)<-c(n,ni*(ni+1)/2)
  g<-vector(length=n*ni*(ni+1)/2)
  corg<-matrix(0,n*ni*(ni+1)/2,3)
  key<-rep(0,ni*nij*2500)
  dim(key)<-c(ni,nij,2500)
  
  seq<-seq(ni,1,by=-1)
  for (i in 1:(ni-1)){
    seq[i+1]<-seq[i]+seq[i+1]
  }
  
  result<-foreach(i=1:n, .combine='+') %dopar% {
    faig<-c(0,0)
    for (j in 1:ni){
      for (k in 1:nij){
        key[j,k,]<-1-1/(1+exp(data$X[(i-1)*ni*nij+(j-1)*nij+k,-(1:3)]%*%beta+sigma1*r1[i,]+sigma2*r2[i,j,]))
      }
    }
    ri<-r1[i,]
    for (j in 1:ni){
      for (k in 1:nij){
        miu1<-key[j,k,]
        dijk[i,j,k]<-sum(miu1*(1-miu1)*ri)/2500
        miu[i,j,k]<-sum(miu1)/2500
        for (l in 1:nij){
          for (m in 1:ni){
            if(j<=m){
              rim<-r2[i,m,]
              miu3<-key[m,l,]
              thoijj[i,j,m,k,l]<-sum(miu1*miu3)/2500
              dijj[i,j,m,k,l]<-sum(miu1*miu3*(1-miu1+1-miu3)*ri)/2500
            }
          }  
        }
      }
    }
    for (j in 1:(ni*(ni+1)/2)){
      if (j<=ni){
        g[(i-1)*ni*(ni+1)/2+j]<-sum(data$y[((i-1)*ni*nij+(j-1)*nij+1):((i-1)*ni*nij+j*nij )])/nij*sum(data$y[((i-1)*ni*nij+(j-1)*nij+1):((i-1)*ni*nij+j*nij )])/nij
        corg[j,]<-c(i,j,j)
      } else { 
        for (l in 2:ni){
          if ((j<=seq[l])&(j>seq[l-1])){
            g[(i-1)*ni*(ni+1)/2+j]<-sum(data$y[((i-1)*ni*nij+(l-2)*nij+1):((i-1)*ni*nij+(l-1)*nij )])/nij*sum(data$y[((i-1)*ni*nij+(j-seq[l-1]+l-2)*nij+1):((i-1)*ni*nij+(j-seq[l-1]+l-1)*nij )])/nij
            corg[j,]<-c(i,(l-1),(j-seq[l-1]+l-1))
            break
          }
        }
      }
      if(corg[j,2]==corg[j,3]){
        theta[(i-1)*ni*(ni+1)/2+j]<-(sum(miu[i,j,])+sum(thoijj[i,j,j,,]*(rep(1,nij)%*%t(rep(1,nij))-diag(nij))))/nij/nij
        devg[i,j]<-(sum(dijk[i,j,])+sum(dijj[i,j,j,,]*(rep(1,nij)%*%t(rep(1,nij))-diag(nij))))/nij/nij
      } else {
        theta[(i-1)*ni*(ni+1)/2+j]<-sum(thoijj[i,corg[j,2],corg[j,3],,])/nij/nij
        devg[i,j]<-sum(dijj[i,corg[j,2],corg[j,3],,])/nij/nij
      }
    }
    faig[1]<--t(devg[i,])%*%t(t(devg[i,]))
    faig[2]<-(t(devg[i,])%*%(g[((i-1)*ni*(ni+1)/2+1):(i*ni*(ni+1)/2)]-theta[((i-1)*ni*(ni+1)/2+1):(i*ni*(ni+1)/2)]))^2
    return(faig)
  }
  
  return(solve(result[1]*solve(result[2])*result[1]))
}


gmmsd<-function(data,beta,n,ni,nij,sigma1,sigma2,seed){ # We use this function to get standard errors
  betasd<-solve1(data,beta,n,ni,nij,sigma1,sigma2,seed)
  sigma2sd<-solve2(data,beta,n,ni,nij,sigma1,sigma2,seed)
  sigma1sd<-solve3(data,beta,n,ni,nij,sigma1,sigma2,seed)
  return(c(sqrt(betasd[1,1]),sqrt(betasd[2,2]),sqrt(sigma1sd),sqrt(sigma2sd)))
}

i<-NULL
data1<-gendata(100,6,6,0.8,1.2,1)
gmmsd(data1,c(1,1),100,6,6,0.8,1.2,32)
stopCluster(cl)