## fill_NA_with_ma.r

fill_NA_with_ma <- function(DT, NAcolTreatment=1) {
## Note: this function does NOT modify by reference, since we expect DT to be the .SD in many cases, ie, to be a locked data.table

  NAS.cols <- colSums(is.na(DT)) > 0
  if (!any(NAS.cols))
    return(DT)

  temp_DT.filled <- 
    DT[, lapply(.SD, function(col) {

      orig.class <- class(col)
      NAs <- is.na(col)

      ## If the whole row is NA, return NAcolTreatment, coerced to appropriate class
      if (all (NAs)) {
        as({rep(NAcolTreatment, length(col))}, orig.class)

      ## Otherwise, take a moving average, expanding the width of the MA on each pass, until all cols are filled
      } else {
        len <- 1
        while (any(NAs) && len < (length(col)/2 + 1)) {
          len <- len + 1
          col[NAs] <-
              sapply(which(NAs), function(i) {
                inds <- i + seq.int(from=-min(len, i), to=len)
                mean(col[inds], na.rm=TRUE)
              })
          NAs <- is.na(col)
          } ## // closes while
         ## RETURN
         as(col, orig.class)
        } ## // closes else
    }) , .SDcols=names(which(NAS.cols))]

  ## If there are no columns WITHOUT NAs in them, then the temp_DT.filled has all of the columns. Return that
  if (all(NAS.cols))
    return(temp_DT.filled)

  ## ELSE:  cbind in the original columns, and set their order
  else 
    return(setcolorder(cbind(temp_DT.filled, DT[, names(which(!NAS.cols)), with=FALSE]), names(DT)))
}
