#### Setup ####

set.seed(52054325)
projdir = "C:/Users/nate/Documents/GitHub/Rcoord/"
m = 1e+00
s1database1 = data.frame(unitid=1:100*m,R0=runif(100*m))
s2database1 = data.frame(unitid=1:100*m,R0=runif(100*m))

s1frame1 = data.frame(unitid=1:100*m,stratid=c(rep(11,40*m),rep(12,60*m)))
s2frame1 = data.frame(unitid=1:100*m,stratid=c(rep(21,30*m),rep(22,70*m)))

s1design1 = data.frame(stratid=as.numeric(as.character(as.data.frame(table(s1frame1$stratid))$Var1)),
                       smallnh=as.data.frame(table(s1frame1$stratid))$Freq*0.25)
s2design1 = data.frame(stratid=as.numeric(as.character(as.data.frame(table(s2frame1$stratid))$Var1)),
                      smallnh=as.data.frame(table(s2frame1$stratid))$Freq*0.25)

write.table(s1database1,paste0(projdir,"sasdb/","s1database1"),row.names=F)
write.table(s2database1,paste0(projdir,"sasdb/","s2database1"),row.names=F)
write.table(s1frame1,paste0(projdir,"framef/","s1frame1"),row.names=F)
write.table(s2frame1,paste0(projdir,"framef/","s2frame1"),row.names=F)
write.table(s1design1,paste0(projdir,"designf/","s1design1"),row.names=F)
write.table(s2design1,paste0(projdir,"designf/","s2design1"),row.names=F)
rm(list=ls())

#### Coordination Program ####

projdir = "C:/Users/nate/Documents/GitHub/Rcoord/"
framef = paste0(projdir,"framef/") # frame files
designf = paste0(projdir,"designf/") # design files
sasdb = paste0(projdir,"sasdb/") # sasdb

cumburd <- function(b) { # calculates the cumburd for all units
  Reduce("+",b)
}

createframe <- function(i) { # build a data using strata, r, burden, unitid
  data.frame(unitid,stratid=strata[[i]],r=r[[i]],b=b[[i]])
}

updatelists <- function(frame,i=1) { # update r and b lists
  r[[i]] <<- frame$r
  b[[i]] <<- frame$b
}

init <- function(frame="s1frame1",bf=1) { # run the first survey
  
  # read in files
  frame1 <<- read.table(paste0(framef,frame),header=T)
  design1 <<- read.table(paste0(designf,gsub("frame","design",frame)),header=T)
  database1 <<- read.table(paste0(sasdb,gsub("frame","database",frame)),header=T)

  # initialize lists
  strata <<- list(frame1$stratid)
  r <<- list(database1$R0)
  b <<- list(rep(0,length(frame1$unitid)))
  sizes <<- list(as.data.frame(table(strata[[1]]))[[2]])
  smallnh <<- list(pmin(design1$smallnh,sizes[[1]]))
  unitid <<- frame1$unitid
  
  # select units 
  frame = createframe(1) # collect cols
  frame = frame[with(frame,order(stratid,b,r)),] # order cols
  frame = split(frame,frame$stratid) # split into strata
  for(i in 1:length(frame)) { # select units
    frame[[i]][which(frame[[i]]$unitid==head(frame[[i]],
          smallnh[[1]][i])$unitid),]$b = bf
  }
  frame = do.call(rbind,frame) # rebind strata
  frame = frame[with(frame,order(unitid)),] # reorder cols by unitid
  updatelists(frame,1) # update burdens and random numbers in lists
}

runsurvey <- function(frame="s1frame1",surveynum=2,bf=1,method="ch") { # run successive surveys
  getsequence <- function(surveynum,method) {
    sequence = 0
    if(method=="ch") sequence = surveynum-1
    if(method=="gn") sequence = (surveynum-1):1
    if(method=="ms") sequence = 1
    sequence
  }
  
  # read in files
  frame2 <<- read.table(paste0(framef,frame),header=T)
  design2 <<- read.table(paste0(designf,gsub("frame","design",frame)),header=T)
  database2 <<- read.table(paste0(sasdb,gsub("frame","database",frame)),header=T)
 
  # update lists
  strata[[surveynum]] <<- frame2$stratid
  r[[surveynum]] <<- r[[surveynum-1]] # use previous random numbers
  b[[surveynum]] <<- rep(0,length(frame2$unitid))
  sizes[[surveynum]] <<- as.data.frame(table(strata[[surveynum]]))[[2]]
  smallnh[[surveynum]] <<- pmin(design2$smallnh,sizes[[surveynum]])
  unitid <<- frame2$unitid
  
  # permute nums
  frame = createframe(2) # collect cols
  frame$stratid = rep("",length(unitid)) # init microstrat
  for(i in getsequence(surveynum,method)) { # global negative
    frame = frame[with(frame,order(unitid)),] # required to update col
    frame$stratid = paste0(frame$stratid,strata[[i]]) # microstrat
    frame$b = cumburd(b[i:(surveynum-1)]) # cumburd
    frame = frame[with(frame,order(stratid,b,r)),] # order cols

    frame = split(frame,frame$stratid) # split into strata
    for(j in 1:length(frame)) { # sort within previous strata
      frame[[j]]$r = sort(frame[[j]]$r) # replace R with sorted R
    }
    frame = do.call(rbind,frame) # rebind
  }
  
  # select units
  frame$stratid = strata[[surveynum]] # update strata
  frame$b = rep(0,length(unitid)) # update burdens
  frame = frame[with(frame,order(stratid,b,r)),] # order cols
  frame = split(frame,frame$stratid) # resplit into strata
  for(i in 1:length(frame)) { # select units
    h = head(frame[[i]],smallnh[[surveynum]][i])
    h$b = bf # indicate that units were selected
    t = tail(frame[[i]], -smallnh[[surveynum]][i])
    frame[[i]] = rbind(h,t) # rebind 
  }
  frame = do.call(rbind,frame) # rebind strata
  frame = frame[with(frame,order(unitid)),] # reorder cols by unitid
  updatelists(frame,surveynum) # update burdens and random numbers in lists
}

init(frame="s1frame1",bf=5)
runsurvey(frame="s2frame1",surveynum=2,bf=4,method="gn")
runsurvey(frame="s1frame1",surveynum=3,bf=2,method="gn")
runsurvey(frame="s2frame1",surveynum=4,bf=5,method="gn")
runsurvey(frame="s1frame1",surveynum=5,bf=3,method="gn")


freq = as.data.frame(table(cumburd(b)))
plot(as.numeric(as.character((freq[[1]]))),as.numeric(as.character((freq[[2]]))),type='l',col='blue',xlab="Burden",ylab="Frequency")
points(as.numeric(as.character((freq[[1]]))),as.numeric(as.character((freq[[2]]))),pch=19,col='blue')