#################################################################################
################ CLUSTER RANDOMIZED EXPERIMENT
################ PATE AND VARIANCE ESTIMATION FUNCTIONS
################################################################################

############################# KEY
#n.clust: number of clusters in population
#n.units: number of units in population
#clustersum: sum of the responses for each sampled cluster
#clustervar: variance of responses for each sampled cluster
#trtindex: a vector of 1s and 0s to indicate treatment assignment of clusters
#           1=treated, 0=control
#clustersize: vector indicating how many units are in each sampled cluster
#withinsampsize: vector indicating how many units are sampled within each sampled cluster
#                if all units are sampled within clusters, withinsampsize = clustersize

### ONLY NEED FOR PPSest function
#ppsvar.type: character indicating variance estimator for PPSWOR
#              "wr" (default) for with-replacement; "syg" for SYG
#jointprob: matrix indicating the joint cluster inclusion probabilities for sampled clusters
#           only need for PPSest function if using ppsvar.type="wor"

### ADDITIONAL NOTES
# clustersum: sum of the responses for each cluster
# clustertotal: clustersize/withinsampsize*clustersum

################################ ESTIMATION FUNCTIONS ###########################################################################

getHTvar <- function(n.clust, n.units, 
                     no.sampledclust, no.trtedclust, no.contclust=no.sampledclust-no.trtedclust,
                     clustertotal, clustervar, clustersize, withinsampsize, 
                     indices1, indices0){
  
  
  
  ########### SYG VARIANCE ESTIMATION OF HT-PATE 
  
    #### between cluster variance estimation
  #creates a matrix for (CTOTAL-CTOTAL)^2
  sqtotaldiff1 = outer(n.clust/no.trtedclust*clustertotal[indices1], n.clust/no.trtedclust*clustertotal[indices1], "-")^2
  #first term of VAR(MU_T)
  var1.HT = 1/(n.units^2)*sum((no.trtedclust*(n.clust-1)/(n.clust*(no.trtedclust-1))-1)*sqtotaldiff1[lower.tri(sqtotaldiff1)]) 
  
  # same code as previous
  sqtotaldiff0 = outer(n.clust/no.contclust*clustertotal[indices0], n.clust/no.contclust*clustertotal[indices0], "-")^2
  var0.HT = 1/(n.units^2)*sum((no.contclust*(n.clust-1)/(n.clust*(no.contclust-1))-1)*sqtotaldiff0[lower.tri(sqtotaldiff0)]) 
  
  # creates matrix of product between mu_c1 and mu_c'0
  totalprod10 = outer(clustertotal[indices1], clustertotal[indices0], "*")
  # first term of cov: sum over the whole matrix
  first10.HT = 1/(n.units^2)*(n.clust/(no.trtedclust*no.contclust))*sum(totalprod10)
  # second term of cov 
  second10.HT = 1/2*sum(n.clust/(n.units^2*no.trtedclust)*clustertotal[indices1]^2) 
  # third term of cov 
  third10.HT = 1/2*sum(n.clust/(n.units^2*no.contclust)*clustertotal[indices0]^2) 
  # estimated cov
  cov.HT = first10.HT - second10.HT - third10.HT 
  
  # estimated variance of HT estimator using SYG formula
  var.est = var0.HT + var1.HT - 2*cov.HT
  
  return(var.est)
}

HTest <- function(n.clust, n.units, 
                  clustersum, clustervar, trtindex, clustersize, withinsampsize){
  
  no.sampclust = length(clustersum)
  no.trtclust = sum(trtindex)
  clustertotal = clustersum*clustersize/withinsampsize
  
  indices1 = trtindex==1 #treatment indices
  indices0 = trtindex==0 #control indices
  
  ht1 = n.clust*sum(1/(no.trtclust*n.units)*clustertotal[indices1])
  ht0 = n.clust*sum(1/((no.sampclust-no.trtclust)*n.units)*clustertotal[indices0])
  #ht1 = n.clust/n.units*mean(clustertotal[indices1])  #same as previous
  #ht0 = n.clust/n.units*mean(clustertotal[indices0])
  pate.est = ht1-ht0
  
  var.est = getHTvar(n.clust = n.clust, n.units = n.units, indices1=indices1, indices0=indices0,
                     no.sampledclust = no.sampclust, no.trtedclust = no.trtclust,
                     clustertotal = clustertotal, clustervar = clustervar,
                     clustersize = clustersize, withinsampsize = withinsampsize)
  
  return(c(PATE.EST = pate.est, SE.EST = sqrt(var.est)))
  
}

STRTHTest <- function(n.clust, n.units, 
                      stratum, c.stratsize, strat.sampsize,
                      clustersum, clustervar, trtindex, clustersize, withinsampsize){
  
  clustertotal = clustersum*clustersize/withinsampsize
  
  strtdata = as.data.frame(cbind(clustertotal, stratum=stratum, trtindex=trtindex, withinsampsize=withinsampsize, 
                                 clustersize=clustersize, clustervar=clustervar))
  strtdata1 = strtdata %>% group_by(stratum) %>% summarise(
    strt.nounits = sum(clustersize),
    strt.notrt = sum(trtindex==1),
    strt.nocont = sum(trtindex==0),
    strt.trtsum = sum(clustertotal[trtindex==1]),
    strt.contsum = sum(clustertotal[trtindex==0]),
    diff = strt.trtsum/strt.notrt - strt.contsum/strt.nocont
  )
  
  pate.est = sum(c.stratsize/n.units*strtdata1$diff) 
  
  varest.stratum = c()
  for (i in 1:length(unique(stratum))) {
    
    varest.stratum[i] = getHTvar(n.clust=c.stratsize[i], n.units=sum(strtdata$clustersize[strtdata$stratum==i]), 
                                 no.sampledclust=strat.sampsize[i], no.trtedclust=sum(strtdata$trtindex[strtdata$stratum==i]),
                                 clustertotal=strtdata$clustertotal[strtdata$stratum==i], clustervar=strtdata$clustervar[strtdata$stratum==i], 
                                 clustersize=strtdata$clustersize[strtdata$stratum==i], withinsampsize=strtdata$withinsampsize[strtdata$stratum==i], 
                                 indices1=strtdata$trtindex[strtdata$stratum==i]==1, indices0=strtdata$trtindex[strtdata$stratum==i]==0) 
  }
  
  var.est = sum(strtdata1$strt.nounits/sum(strtdata1$strt.nounits)*varest.stratum)
  
  return(c(PATE.EST = pate.est, SE.EST = sqrt(var.est)))
}


getPPSvar <- function(n.clust, n.units, jointprob,
                      no.sampclust, no.trtclust, no.contclust=no.sampclust-no.trtclust,
                      clustertotal, clustervar, clustersize, withinsampsize, 
                      indices1, indices0){
  
  
  ##### VARIANCE ESTIMATION
  jointprob.trt = jointprob[indices1,indices1]
  jointprob.trt = jointprob.trt[lower.tri(jointprob.trt)]
  jointprob.cont = jointprob[indices0,indices0]
  jointprob.cont = jointprob.cont[lower.tri(jointprob.cont)]
  
  jointsize.trt = outer(clustersize[indices1], clustersize[indices1], "*")
  jointsize.trt = jointsize.trt[lower.tri(jointsize.trt)]
  jointsize.cont = outer(clustersize[indices0], clustersize[indices0], "*")
  jointsize.cont = jointsize.cont[lower.tri(jointsize.cont)]
  
  #within-variance calculation (var-hat of tau_ct-hat)
  withinclust.var1 = clustersize[indices1]^2/withinsampsize[indices1]*(1-withinsampsize[indices1]/clustersize[indices1])*clustervar[indices1]
  withinclust.var0 = clustersize[indices0]^2/withinsampsize[indices0]*(1-withinsampsize[indices0]/clustersize[indices0])*clustervar[indices0]
  
  #creates a matrix for (CTOTAL-CTOTAL)^2
  sqtotaldiff1 = outer(clustertotal[indices1]/clustersize[indices1], clustertotal[indices1]/clustersize[indices1], "-")^2
  sqtotaldiff1 = sqtotaldiff1[lower.tri(sqtotaldiff1)]
  #sqtotaldiff1 = as.vector(t(sqtotaldiff1))
  var1.PPS = sum((jointsize.trt*no.sampclust*(no.sampclust-1)/(jointprob.trt*n.units^2*(no.trtclust-1)*no.trtclust)-1/(no.trtclust^2))*sqtotaldiff1) 
  + sum(withinclust.var1/(n.units*clustersize[indices1]*no.trtclust))
  sqtotaldiff0 = outer(clustertotal[indices0]/clustersize[indices0], clustertotal[indices0]/clustersize[indices0], "-")^2
  sqtotaldiff0 = sqtotaldiff0[lower.tri(sqtotaldiff0)]
  var0.PPS = sum((jointsize.cont*no.sampclust*(no.sampclust-1)/(jointprob.cont*n.units^2*(no.contclust-1)*no.contclust)-1/(no.contclust^2))*sqtotaldiff0) 
  + sum(withinclust.var0/(n.units*clustersize[indices0]*no.contclust))
  
  # creates matrix of product between mu_c1 and mu_c'0
  totalprod10 = outer(clustertotal[indices1], clustertotal[indices0], "*")
  jointsize10 = outer(clustersize[indices1], clustersize[indices0], "*")
  jointprob10 = jointprob[indices1, indices0]
  
  # first term of cov: sum over the whole matrix
  first10.PPS = sum((1/(jointsize10*no.trtclust*no.contclust) - no.sampclust*(no.sampclust-1)/(jointprob10*n.units^2*no.trtclust*no.contclust))*
                      totalprod10)
  # second term of cov
  second10.PPS = sum(1/(2*n.units*no.trtclust)*clustertotal[indices1]^2) - sum(1/(2*n.units*clustersize[indices1]*no.trtclust)*withinclust.var1)
  # third term of cov
  third10.PPS = sum(1/(2*n.units*no.contclust)*clustertotal[indices0]^2) - sum(1/(2*n.units*clustersize[indices0]*no.contclust)*withinclust.var0)
  # estimated cov
  cov.PPS = first10.PPS - second10.PPS - third10.PPS 
  
  # estimated variance of HT estimator using SYG formula
  var.est = var0.PPS + var1.PPS - 2*cov.PPS
  
  return(var.est)
}


PPSvar <- function(n.clust, n.units, jointprob,
                   clustersum, clustervar, trtindex, clustersize, withinsampsize){
  
  no.sampclust = length(clustersum)
  no.trtclust = sum(trtindex)
  no.contclust = no.sampclust-no.trtclust
  
  clustertotal = clustersum*clustersize/withinsampsize
  
  ##### HT-PPS ESTIMATE OF PATE
  
  indices1 = trtindex==1 #treatment indices
  indices0 = trtindex==0 #control indices
  
  ht1 = sum(1/(clustersize[indices1]*no.trtclust)*clustertotal[indices1])
  ht0 = sum(1/(clustersize[indices0]*(no.sampclust-no.trtclust))*clustertotal[indices0])
  pate.est = ht1-ht0
  
  var.est = getPPSvar(n.clust = n.clust, n.units = n.units, jointprob = jointprob,
                      indices1=indices1, indices0=indices0, 
                      no.sampclust = no.sampclust, no.trtclust = no.trtclust,
                      clustertotal = clustertotal, clustervar = clustervar,
                      clustersize = clustersize, withinsampsize = withinsampsize)
  
  return(c(PATE.EST = pate.est, SE.EST = sqrt(var.est)))
  
}


getPPSvar.wr <- function(n.clust, n.units, indices1, indices0,
                         no.sampclust, no.trtclust, no.contclust=no.sampclust-no.trtclust,
                         ht1,ht0,
                         clustertotal, clustervar, clustersize, withinsampsize){
  
  clustermeans = clustertotal/clustersize
  
  #within-variance calculation (var-hat of tau_ct-hat)
  withinclust.var1 = clustersize[indices1]^2/withinsampsize[indices1]*(1-withinsampsize[indices1]/clustersize[indices1])*clustervar[indices1]
  withinclust.var0 = clustersize[indices0]^2/withinsampsize[indices0]*(1-withinsampsize[indices0]/clustersize[indices0])*clustervar[indices0]
  
  var1 = 1/no.trtclust/(no.trtclust-1)*sum((clustermeans[indices1]-ht1)^2) 
  var0 = 1/no.contclust/(no.contclust-1)*sum((clustermeans[indices0]-ht0)^2) 
  
  
  totalprod10 = outer(clustermeans[indices1], clustermeans[indices0], "*")
  first10 = 1/(no.sampclust*no.trtclust*no.contclust)*sum(totalprod10)
  second10 = 1/2*sum(1/(n.units*clustersize[indices1]*no.trtclust)*clustertotal[indices1]^2) -  
    1/2*sum(1/(n.units*clustersize[indices1]*no.trtclust)*withinclust.var1)
  third10 = 1/2*sum(1/(n.units*clustersize[indices0]*no.contclust)*clustertotal[indices0]^2) -
    1/2*sum(1/(n.units*clustersize[indices0]*no.contclust)*withinclust.var0)
  cov10 = first10 - second10 - third10
  cov10.2 = 0 - second10 - third10
  
  var.est0 = var1 + var0 - 2*cov10
  
  return(var.est0)
}



PPSvar.wr = function(n.clust, n.units, 
                     clustersum, clustervar, trtindex, clustersize, withinsampsize){
  
  no.sampclust = length(clustersum)
  no.trtclust = sum(trtindex)
  no.contclust = no.sampclust-no.trtclust
  
  clustertotal = clustersum*clustersize/withinsampsize
  
  ##### HT-PPS ESTIMATE OF PATE
  
  indices1 = trtindex==1 #treatment indices
  indices0 = trtindex==0 #control indices
  
  ht1 = sum(1/(clustersize[indices1]*no.trtclust)*clustertotal[indices1])
  ht0 = sum(1/(clustersize[indices0]*(no.sampclust-no.trtclust))*clustertotal[indices0])
  pate.est = ht1-ht0
  
  var.est = getPPSvar.wr(n.clust=n.clust, n.units=n.units, indices1=indices1, indices0=indices0,
                         no.sampclust=no.sampclust, no.trtclust=no.trtclust, 
                         ht1=ht1,ht0=ht0,
                         clustertotal=clustertotal, clustervar=clustervar, clustersize=clustersize, withinsampsize=withinsampsize)
  
  return(c(PATE.EST = pate.est, SE.EST = sqrt(var.est))) 
  
}

PPSest = function(n.clust, n.units, 
                  clustersum, clustervar, trtindex, clustersize, withinsampsize,
                  ppsvar.type = "wr", jointprob = NULL){
  
  if(ppsvar.type == "syg"){
    est = PPSvar(n.clust=n.clust, n.units=n.units, clusterid = clusterid, 
                 clustersum = clustersum, clustervar = clustervar, trtindex = trtindex,
                 clustersize = clustersize, withinsampsize = withinsampsize)
  }
  
  else{
    est = PPSvar.wr(n.clust=n.clust, n.units=n.units, 
                    clustersum = clustersum, clustervar = clustervar, trtindex = trtindex,
                    clustersize = clustersize, withinsampsize = withinsampsize)
  }
  
  return(est)
}

getHJvar <- function(n.clust, indices1, indices0,
                     no.sampclust, no.trtclust, no.contclust=no.sampclust-no.trtclust,
                     mu1,mu0,
                     clustertotal, clustervar, clustersize, withinsampsize){
  
  withinclust.var1 = clustersize[indices1]^2/withinsampsize[indices1]*(1-withinsampsize[indices1]/clustersize[indices1])*clustervar[indices1]
  withinclust.var0 = clustersize[indices0]^2/withinsampsize[indices0]*(1-withinsampsize[indices0]/clustersize[indices0])*clustervar[indices0]
  betclust.var1 = 1/(n.clust-1)*sum((clustertotal[indices1]-mu1*clustersize[indices1])^2)
  betclust.var0 = 1/(n.clust-1)*sum((clustertotal[indices0]-mu0*clustersize[indices0])^2)
  
  var1 = 1/sum(clustersize[indices1])^2*n.clust^2/no.trtclust*(1-no.trtclust/n.clust)*betclust.var1 
  var0= 1/sum(clustersize[indices0])^2*n.clust^2/no.contclust*(1-no.contclust/n.clust)*betclust.var0 
  
  totalprod10 = outer((clustertotal[indices1]-mu1*clustersize[indices1]), (clustertotal[indices0]-mu0*clustersize[indices0]), "*")
  first10 = 1/(sum(clustersize[indices1])*sum(clustersize[indices0]))/(n.clust-1)*sum(totalprod10)
  second10 = 1/2*1/(sum(clustersize[indices1])*sum(clustersize[indices0]))*sum((clustertotal[indices1]-mu1*clustersize[indices1])^2)
  third10 = 1/2*1/(sum(clustersize[indices1])*sum(clustersize[indices0]))*sum((clustertotal[indices0]-mu0*clustersize[indices0])^2)
  cov10 = first10 - second10 - third10
  
  var.est = var1 + var0 - 2*cov10
  
  return(var.est)
}

HJest <- function(n.clust, n.units, 
                  clustersum, clustervar, trtindex, clustersize, withinsampsize){
  
  no.sampclust = length(clustersum)
  no.trtclust = sum(trtindex)
  no.contclust = no.sampclust-no.trtclust
  
  clustertotal = clustersum*clustersize/withinsampsize
  
  indices1 = trtindex==1 #treatment indices
  indices0 = trtindex==0 #control indices
  
  ##### HJ ESTIMATE OF PATE
  
  mu1 = sum(clustertotal[indices1])/sum(clustersize[indices1])
  mu0 = sum(clustertotal[indices0])/sum(clustersize[indices0])
  pate.est = mu1-mu0
  
  var.est = getHJvar(n.clust=n.clust, indices1=indices1, indices0=indices0,
                     no.sampclust = no.sampclust, no.trtclust=no.trtclust, 
                     mu1=mu1, mu0=mu0,
                     clustertotal = clustertotal, clustervar=clustervar, clustersize=clustersize, withinsampsize = withinsampsize)
  
  return(c(PATE.EST = pate.est, SE.EST = sqrt(var.est)))
  
}

getDIMvar <- function(n.clust, indices1, indices0,
                       no.sampclust, no.trtclust, no.contclust=no.sampclust-no.trtclust,
                       mu1,mu0,
                       clustersum, clustervar, clustersize, withinsampsize){
  
  withinclust.var1 = withinsampsize[indices1]*(1-withinsampsize[indices1]/clustersize[indices1])*clustervar[indices1]
  withinclust.var0 = withinsampsize[indices0]*(1-withinsampsize[indices0]/clustersize[indices0])*clustervar[indices0]
  betclust.var1 = 1/(no.trtclust-1)*sum((withinsampsize[indices1]/clustersize[indices1]*clustersum[indices1]-mu1*withinsampsize[indices1])^2)
  betclust.var0 = 1/(no.contclust-1)*sum((withinsampsize[indices0]/clustersize[indices0]*clustersum[indices0]-mu0*withinsampsize[indices0])^2)
  
  var1 = 1/sum(withinsampsize)^2*n.clust^2/no.trtclust*(1-no.trtclust/n.clust)*betclust.var1 
  var0= 1/sum(withinsampsize)^2*n.clust^2/no.contclust*(1-no.contclust/n.clust)*betclust.var0 
  
  totalprod10 = outer((withinsampsize[indices1]/clustersize[indices1]*clustersum[indices1]-mu1*withinsampsize[indices1]), 
                      (withinsampsize[indices0]/clustersize[indices0]*clustersum[indices0]-mu0*withinsampsize[indices0]), "*")
  first10 = n.clust/sum(withinsampsize)^2/(no.trtclust*no.contclust)*sum(totalprod10)
  second10 = 1/2*n.clust/no.trtclust/(sum(withinsampsize)^2)*sum((withinsampsize[indices1]/clustersize[indices1]*clustersum[indices1]-mu1*withinsampsize[indices1])^2)
  third10 = 1/2*n.clust/no.contclust/(sum(withinsampsize)^2)*sum((withinsampsize[indices0]/clustersize[indices0]*clustersum[indices0]-mu0*withinsampsize[indices0])^2)
  cov10 = first10 - second10 - third10
  
  var.est = var1 + var0 - 2*cov10
  
  return(var.est)
  
}

DIMest <- function(n.clust, n.units, 
                   clustersum, clustervar, trtindex, clustersize, withinsampsize){
  
  no.sampclust = length(clustersum)
  no.trtclust = sum(trtindex)
  no.contclust = no.sampclust-no.trtclust
  
  indices1 = trtindex==1 #treatment indices
  indices0 = trtindex==0 #control indices
  
  ##### HJ ESTIMATE OF PATE
  
  mu1 = sum(clustersum[indices1])/sum(withinsampsize[indices1])
  mu0 = sum(clustersum[indices0])/sum(withinsampsize[indices0])
  pate.est = mu1-mu0
  
  var.est = getDIMvar(n.clust=n.clust, indices1=indices1, indices0=indices0,
                       no.sampclust = no.sampclust, no.trtclust=no.trtclust, 
                       mu1=mu1, mu0=mu0,
                       clustersum = clustersum, clustervar=clustervar, clustersize=clustersize, withinsampsize = withinsampsize)
  
  return(c(PATE.EST = pate.est, SE.EST = sqrt(var.est)))
  
}


getMPvar <- function(n.units, n.clust, no.sampclust, 
                     pair.est, strt.trtsum, strt.contsum,
                     clustersize, withinsampsize,
                     clustervar,
                     indices1, indices0){
  
  #This calculates the within cluster variability.
  withinclust.var1 = withinsampsize[indices1]*(1-withinsampsize[indices1]/clustersize[indices1])*clustervar[indices1]
  withinclust.var0 = withinsampsize[indices0]*(1-withinsampsize[indices0]/clustersize[indices0])*clustervar[indices0]
  total.withinvar1 = 4*sum(clustersize[indices1]^2/n.units^2*withinclust.var1)
  total.withinvar0 = 4*sum(clustersize[indices0]^2/n.units^2*withinclust.var0)
  
  #This calculates the within-pair variabilility.
  crosst = 2*sum(1/(n.units^2)*strt.trtsum*strt.contsum)
  bound1 = sum(1/n.units^2*(strt.trtsum^2-clustersize[indices1]^2*withinclust.var1))
  bound0 = sum(1/n.units^2*(strt.contsum^2-clustersize[indices0]^2*withinclust.var0))
  cov10 = crosst - bound1 - bound0
  withinpair.var = n.clust/no.sampclust*(total.withinvar0 + total.withinvar1 - 2*cov10)
  
  #This calculates the across pair variability.
  sqdeltadiff = outer(1/n.units*n.clust/no.sampclust*pair.est, 1/n.units*n.clust/no.sampclust*pair.est, "-")^2
  across.pairvar = -(1/2)*sum( (1-no.sampclust/n.clust*((n.clust/2-1)/(no.sampclust/2-1)))*sqdeltadiff )
  
  var.est = across.pairvar + withinpair.var
  
  return(var.est)
}

MPHTest <- function(n.clust, n.units,
                    stratum, trtindex, clustersize, withinsampsize, 
                    clustersum, clustervar){
  
  no.sampclust = length(stratum)
  
  clustertotal = clustersum*clustersize/withinsampsize
  
  strtdata = as.data.frame(cbind(clustertotal, stratum, trtindex))
  strtdata2 = strtdata %>% group_by(stratum) %>% summarise(
    strt.trtsum = sum(clustertotal[trtindex==1]),
    strt.contsum = sum(clustertotal[trtindex==0]),
    strt.diff = strt.trtsum-strt.contsum)
  pair.est = 2*n.clust/n.units/no.sampclust*strtdata2$strt.diff
  pate.est = sum(pair.est)
  
  indices1 = trtindex==1
  indices0 = trtindex==0
  
  var.est = getMPvar(n.units=n.units, n.clust=n.clust, no.sampclust=no.sampclust, 
                     pair.est=pair.est, strt.trtsum=strtdata2$strt.trtsum, strt.contsum=strtdata2$strt.contsum, 
                     clustervar = clustervar, clustersize = clustersize, withinsampsize = withinsampsize,
                     indices1=indices1, indices0=indices0)
  
  return(c(PATE.EST = pate.est, SE.EST = sqrt(var.est)))
  
}

DRest <- function(n.clust, n.units, 
                  stratum = NULL, c.stratsize = NULL, stratsize = NULL, strat.sampsize = NULL,
                  clustersum, clustervar, trtindex, clustersize, withinsampsize, 
                  opt.theta = NULL, samp.method = "SRS"){
  
  
  no.sampclust = length(clustersum)
  no.trtclust = sum(trtindex==1)
  no.contclust = no.sampclust-no.trtclust
  clustertotal = clustersum*clustersize/withinsampsize
  
  indices1 = trtindex==1 #treatment indices
  indices0 = trtindex==0 #control indices
  
  if(samp.method=="STRTSRS-BLK" | samp.method == "STRTSRS"){
    
    if(is.null(opt.theta)){
      modelt = lm(clustersum[indices1] ~ clustersize[indices1])
      t.theta = as.numeric(modelt$coefficients[2])
      
      modelc = lm(clustersum[indices0] ~ clustersize[indices0])
      c.theta = as.numeric(modelc$coefficients[2])
      
      opt.theta = sum(trtindex==1)*c.theta/no.sampclust + sum(trtindex==0)*t.theta/no.sampclust
      
    }
    
    stratsize.byclust = rep(stratsize, strat.sampsize)
    c.stratsize.byclust = rep(c.stratsize, strat.sampsize)
    cov.diff = clustersize - stratsize.byclust/c.stratsize.byclust
    
    strtdata = as.data.frame(cbind(stratum, clustertotal, trtindex, clustersize, withinsampsize, cov.diff))
    strtdata1 = strtdata %>% group_by(stratum) %>% summarise(
      strt.notrt = sum(trtindex==1),
      strt.nocont = sum(trtindex==0),
      strt.trtsum = sum(clustertotal[trtindex==1]),
      strt.contsum = sum(clustertotal[trtindex==0]),
      diff = (sum(strt.trtsum/strt.notrt) - sum(strt.contsum/strt.nocont)) - (sum(cov.diff[trtindex==1]/strt.notrt) - sum(cov.diff[trtindex==0]/strt.nocont)),
      diff2 = sum(strt.trtsum/no.trtclust) - sum(strt.contsum/(no.sampclust-no.trtclust)) - (sum(cov.diff[trtindex==1]/no.trtclust) - sum(cov.diff[trtindex==0]/(no.sampclust-no.trtclust))) 
    )
    
    if(samp.method == "STRTSRS-BLK"){
      pate.est = sum(c.stratsize/n.units*opt.theta*strtdata1$diff)
    }
    else{
      pate.est = sum(c.stratsize/n.units*no.sampclust/strat.sampsize*opt.theta*strtdata1$diff2)
    }
    
  }
  else{
    
    if(is.null(opt.theta)){
      # estimate theta
      modelt = lm(clustersum[indices1] ~ clustersize[indices1])
      t.theta = modelt$coefficients[2]
    }
    if(is.null(opt.theta)){
      modelc = lm(clustersum[indices0] ~ clustersize[indices0])
      c.theta = modelc$coefficients[2]
    }
    
    opt.theta = sum(trtindex==1)*c.theta/no.sampclust + sum(trtindex==0)*t.theta/no.sampclust
    
    # DIFFERENCE ESTIMATOR
    
    ht1 = n.clust*sum(1/(no.trtclust*n.units)*clustertotal[indices1])
    ht0 = n.clust*sum(1/(no.contclust*n.units)*clustertotal[indices0])
    dr1 = ht1 - n.clust*sum(1/(no.trtclust*n.units)*opt.theta*(clustersize[indices1]-n.units/n.clust))
    dr0 = ht0 - n.clust*sum(1/(no.contclust*n.units)*opt.theta*(clustersize[indices0]-n.units/n.clust))
    pate.est = dr1 - dr0
    
  }
  
  trans.clustertotal = clustertotal - opt.theta*(clustersize-n.units/n.clust)
  
  var.est = getHTvar(n.clust = n.clust, n.units = n.units, indices1=indices1, indices0=indices0,
                     no.sampledclust = no.sampclust, no.trtedclust = no.trtclust,
                     clustertotal = trans.clustertotal, clustervar = clustervar,
                     clustersize = clustersize, withinsampsize = withinsampsize)
  
  return(c(PATE.EST = pate.est, SE.EST = sqrt(var.est)))
  
}

######################################### DESIGN FUNCTIONS ####################################################

############################## KEY
#clusterid: vector of cluster IDs
#clustsize: vector of cluster population sizes
#stratum: vector of stratum IDs for clusters
#c.stratsize: vector of the number of clusters in each stratum
#stratsize: vector of the number of individuals in each stratum
#strat.sampsize: vector of the number of clusters wanted to sample in each stratum
#blockid: vector of block IDs for clusters; default is stratum
#blkvar: the variable to use for pairing up clusters if blockid is unknown; default: pair based on clustsize as a covariate
#no.sampclust: number of clusters wanted in sample
#no.trtclust: number of treated clusters; default is half of no.sampclust in CRD
#perc: percentage of cluster size to subsample; default=1 (no subsampling)
#withinsampsize: vector of sample sizes of units within each cluster; optional to perc
#unitid: vector of unit IDs if known
#resp.trt: vector of treatment responses if known
#resp.cont: vector of control responses if known
#samp.method: pick sampling method for clusters ("SRS", "PPS", "STRTPPS", "STRTPPS-BLK")
#est.method: pick estimation method for PATE ("HT", "DIM", "DESRAJ", "HAJEK")
#            if sampling of clusters is done with PPS or STRTPPS, or STRTPPS-BLK, est.method must be HT

### ONLY NEED FOR getTrtClust function
#sampledclust: vector of index for sampled clusters
#sampClustID: vector of cluster labels for sampled clusters
#             if clusters aren't originally labeled, sampClustID = sampledclust

### ONLY NEED FOR getWithinSample() and getSampleData()
#clust.data: dataframe with variables sampClustID and treatment indicator (optional: stratum and block)




# This code will randomly sample the clusters using either SRS, stratified SRS, PPS (via Sunter sampling), or stratified PPS.  It gives the index of the 
# clusters sampled.

getSampleClust <- function(clusterid, clustsize,
                           stratum = NULL, c.stratsize = NULL, strat.sampsize = NULL,
                           no.sampclust = sum(strat.sampsize), samp.method = "SRS"){
  
  
  library(TeachingSampling) #Sunter sampling
  
  # Identifying some key parameters
  n.clust = length(unique(clusterid)) #total number of clusters in population
  
  if(samp.method == "STRTPPS" | samp.method == "STRTPPS-BLK"){
    sampledclust = S.STpiPS(stratum, clustsize, strat.sampsize)
    sampledclust = sampledclust[,1]
  } else if(samp.method == "PPS" | samp.method == "PPS-BLK"){
    sampledclust = S.piPS(no.sampclust, clustsize) #RANDOMLY SAMPLE s CLUSTERS with PPS
    sampledclust = sampledclust[,1] #gives the location of sampled clusters
    #sampledclust = sunter(clustsize, no.sampclust)
  } else if(samp.method == "STRTSRS" | samp.method == "STRTSRS-BLK"){
    sampledclust = S.STSI(stratum, c.stratsize, strat.sampsize)
  } else{ 
    #sampling with SRSWOR
    sampledclust = sample(1:n.clust, no.sampclust, replace = FALSE) #RANDOMLY SAMPLE s CLUSTERS
  }
  
  
  return(sampledclust)
  
}

# This code is used in getTrtClust for trt randomization within a block for doing blocked design.  It gives the trtindex for clusters within a block.  

blk.trtrand <- function(x){
  
  trtrand = sample(1:length(x), size = ceiling(length(x)/2), replace = FALSE)
  trtindex = rep(0, length(x))
  trtindex[trtrand] = 1
  return(trtindex)
}

# This code is for trt randomization using either a blocked design or CRD.  It gives a dataframe of cluster and trtindex.

getTrtClust <- function(clusterid, sampledclust, sampClustID, 
                        stratum = NULL, no.trtclust, blockid = NULL, blkvar = NULL, 
                        samp.method = "SRS"){
  
  library(nbpMatching) #matching 
  
  if(samp.method == "STRTPPS-BLK" | samp.method == "STRTSRS-BLK" | samp.method == "BLK"){ 
    
    blk = blockid[sampledclust]
    trtindex = tapply(1:length(blk), blk, blk.trtrand)
    trtindex1 = unlist(trtindex)
    stratum = stratum[sampledclust]
    clust.data = as.data.frame(cbind(sampClustID, stratum, blk))
    clust.data = cbind(clust.data[order(clust.data$blk, clust.data$sampClustID),], trtindex = trtindex1)
    clust.data = clust.data[order(clust.data$sampClustID),]
    
  } else{ #CRD
    trtclust = sample(sampClustID, no.trtclust, replace = FALSE) #RANDOMLY ASSIGNED TREATMENT TO n.trt OF SAMPLED CLUSTERS
    trtindex = ifelse(clusterid %in% trtclust, 1, 0) #TREATMENT INDICATOR FOR CLUSTER
    clust.data = as.data.frame(cbind(sort(sampClustID), trtindex[clusterid %in% sampClustID]))
    colnames(clust.data) = c("sampClustID", "trtindex")
    
    if (samp.method == "STRTSRS" | samp.method == "STRTPPS"){
      stratum = stratum[sampledclust]
      clust.data = cbind(clust.data, stratum)
    }
  }
  return(clust.data)
}

# This code is for subsampling within clusters using SRS.  It outputs a data frame called sample.data that contains the unit-level information 
# on sampled units, including cluster IDs, trt assignment, and cluster size.  If resp.trt and resp.cont are known,
# it will also give the observed responses based on trt assignment.

getWithinSample <- function(clusterid, clustsize, sampClustID, trtclust,
                             perc = 1, withinsampsize = NULL, unitid = NULL,
                             resp.trt = NULL, resp.cont = NULL,
                             clust.data){
  
  library(dplyr) 
  library(purrr)
  library(tidyr)
  
  
  # Identifying some key parameters
  n.clust = length(unique(clusterid)) #total number of clusters in population
  n.units = sum(clustsize) #total number of individuals in population
  
  clusterid.unit = rep(clusterid, clustsize)
  if(is.null(unitid)){unitid = 1:n.units}
  sampclustindex = ifelse(clusterid.unit %in% sampClustID, 1, 0) #INDEX OF SAMPLED CLUSTERS FOR EACH UNIT
  trtclustindex = ifelse(clusterid.unit %in% trtclust, 1, 0) # TREATED CLUSTERS INDICATOR FOR UNITS
  clustsize.byunit = rep(clustsize, clustsize)
  inSamp = clusterid.unit %in% sampClustID
  
  
  if(!is.null(resp.trt) & !is.null(resp.cont)){
    sampdata = as.data.frame(cbind(unitid = unitid[inSamp], clusterid.unit = clusterid.unit[inSamp], resp.trt = resp.trt[inSamp], 
                                   resp.cont = resp.cont[inSamp], clustsize = clustsize.byunit[inSamp], sampclustindex = sampclustindex[inSamp], trtclustindex = trtclustindex[inSamp]))
  }
  else{
    sampdata = as.data.frame(cbind(unitid = unitid[inSamp], clusterid.unit = clusterid.unit[inSamp], clustsize = clustsize.byunit[inSamp], 
                                   sampclustindex = sampclustindex[inSamp], trtclustindex = trtclustindex[inSamp]))
  }
  
  
  # within-cluster sampling using srs
  if(is.null(withinsampsize)){
    withinclust.sampsize = round(perc*clustsize[unique(clusterid) %in% sampClustID], 0)
  }
  else{
    withinclust.sampsize = withinsampsize[clusterid %in% sampClustID]
    #sampdata = cbind(sampdata, withinclust.sampsize)
  }
  nesteddata = sampdata %>% group_by(clusterid.unit) %>% nest() %>% mutate(n = withinclust.sampsize)
  nesteddata = nesteddata %>% mutate(samp.data = map2(data, n, sample_n))
  sample.data = as.data.frame(nesteddata %>% select(clusterid.unit, samp.data) %>% unnest())
  sample.data = sample.data[order(sample.data[,1], sample.data[,2]),]
  
  if(!is.null(resp.trt) & !is.null(resp.cont)){
    obsresp = ifelse(sample.data$trtclustindex==1, sample.data$resp.trt, sample.data$resp.cont) #OBSERVED RESPONSE BASED ON TREATMENT INDICATOR
  }
  
  sample.data = cbind(sample.data[,c(1,2,5,7)], obsresp)
  
  return(sample.data)
  
}

########################################## CODE TO GET SAMPLE DATA #############################################################################

#This code will randomly sample the clusters (using SRS, stratified SRS, PPS, or stratified PPS),
#conduct trt assignment to clusters (using blocked design-only for stratified or CRD),
#and randomly sample units within clusters via SRS.

getSampleData <- function(clusterid, clustsize, 
                           stratum = NULL, c.stratsize = NULL, strat.sampsize = NULL,
                           blockid = stratum, blkvar = clustsize, 
                           no.sampclust = sum(strat.sampsize), no.trtclust = ceiling(no.sampclust/2), 
                           perc=1, withinsampsize = NULL, unitid = NULL, 
                           resp.trt = NULL, resp.cont = NULL, samp.method = "SRS"){
  
  
  # Identifying some key parameters
  n.clust = length(unique(clusterid)) #total number of clusters in population
  n.units = sum(clustsize) #total number of individuals in population
  #clustsize = table(fulldata$clusterid) #cluster size
  
  
  ####################################### SAMPLING CLUSTERS
  
  sampledclust <- getSampleClust(clusterid = clusterid, clustsize = clustsize,
                                 stratum = stratum, c.stratsize = c.stratsize, strat.sampsize = strat.sampsize,
                                 no.sampclust = no.sampclust, samp.method = samp.method)
  
  # The clusters that are sampled
  sampClustID = clusterid[sampledclust]
  
  ############## TREATMENT ASSIGNMENT
  
  clust.data = getTrtClust(clusterid = clusterid, sampledclust = sampledclust, sampClustID = sampClustID, 
                           stratum = stratum, 
                           no.trtclust = no.trtclust,
                           blockid = blockid, blkvar = blkvar, 
                           samp.method = samp.method)
  
  trtclust = clust.data$sampClustID[clust.data$trtindex==1]
  
  ########## WITHIN CLUSTER SAMPLING
  
  sample.data = getWithinSample(clusterid = clusterid, clustsize = clustsize, sampClustID = sampClustID, trtclust = trtclust, 
                                 perc = perc, withinsampsize = withinsampsize, unitid = unitid, resp.trt = resp.trt, resp.cont = resp.cont,
                                 clust.data = clust.data)  
  
  return(sample.data)
  
}