## * Function to fill in gaps (with NA, NaN, or Inf value) in a raster
## ** Takes the mean of the 4 or 8 finite neighbouring values, iteratively, using 8 different strategies,
## **  and returning the overall mean
library(inline)

src <- "
    Rcpp::NumericMatrix Am(A);
    int na = Rcpp::as<int>(missval);
    int neighb = Rcpp::as<int>(nneighb);
    int dir = Rcpp::as<int>(direction);
    int xfirst = Rcpp::as<int>(ythenx);
    int istart;
    int istep;
    int jstart;
    int jstep;
    int i;
    int j;
    NumericMatrix Bm = Am;
    int nrows = Am.nrow();
    int ncolumns = Am.ncol();
    int n = 0;
    float sum = 0;

    if (dir == 1) {
      istart = 1;
      istep  = 1;
      jstart = 1;
      jstep  = 1;
    }
    if (dir == 2) {
      istart = ncolumns-1;
      istep  = -1;
      jstart = 1;
      jstep  = 1;
    }
    if (dir == 3) {
      istart = 1;
      istep  = 1;
      jstart = nrows-1;
      jstep  = -1;
    }
    if (dir == 4) {
      istart = ncolumns-1;
      istep  = -1;
      jstart = nrows-1;
      jstep  = -1;
    }

    if (xfirst) {
      i = istart;
      while (i >= 1 && i <= (ncolumns-1)) {
        j = jstart;
        while (j >= 1 && j <= (nrows-1)) {
              if (Am(j,i)==na && ((Am(j+1,i) != na || Am(j-1,i) != na || Am(j,i+1) != na || Am(j,i-1) != na || (neighb == 8 && (Am(j-1,i-1) != na || Am(j-1,i+1) != na || Am(j+1, i-1) != na || Am(j+1, i+1) != na))))) {
                 n = 0;
                 sum = 0;
                 if (Am(j+1,i) != na) {
                   n = n+1;
                   sum = sum + Am(j+1,i);
                 }
                 if (Am(j-1,i) != na) {
                   n = n+1;
                   sum = sum + Am(j-1,i);
                 }
                 if (Am(j,i+1) != na) {
                   n = n+1;
                   sum = sum + Am(j,i+1);
                 }
                 if (Am(j,i-1) != na) {
                   n = n+1;
                   sum = sum + Am(j,i-1);
                 }

                 if (neighb == 8) {
                   if (Am(j-1,i-1) != na) {
                     n = n+1;
                     sum = sum + Am(j-1,i-1);
                   }
                   if (Am(j-1,i+1) != na) {
                     n = n+1;
                     sum = sum + Am(j-1,i+1);
                   }
                   if (Am(j+1,i-1) != na) {
                     n = n+1;
                     sum = sum + Am(j+1,i-1);
                   }
                   if (Am(j+1,i+1) != na) {
                     n = n+1;
                     sum = sum + Am(j+1,i+1);
                   }
                 }

                 Bm(j,i) = sum/n;
              }
            j = j + jstep;
          }
        i = i + istep;
      }
  } else {

      j = jstart;
      while (j >= 1 && j <= (nrows-1)) {
        i = istart;
        while (i >= 1 && i <= (ncolumns-1)) {
              if (Am(j,i)==na && (Am(j+1,i) != na || Am(j-1,i) != na || Am(j,i+1) != na || Am(j,i-1) != na || (neighb == 8 && (Am(j-1,i-1) != na || Am(j-1,i+1) != na || Am(j+1, i-1) != na || Am(j+1, i+1) != na)))) {
                 n = 0;
                 sum = 0;
                 if (Am(j+1,i) != na) {
                   n = n+1;
                   sum = sum + Am(j+1,i);
                 }
                 if (Am(j-1,i) != na) {
                   n = n+1;
                   sum = sum + Am(j-1,i);
                 }
                 if (Am(j,i+1) != na) {
                   n = n+1;
                   sum = sum + Am(j,i+1);
                 }
                 if (Am(j,i-1) != na) {
                   n = n+1;
                   sum = sum + Am(j,i-1);
                 }

                 if (neighb == 8) {
                   if (Am(j-1,i-1) != na) {
                     n = n+1;
                     sum = sum + Am(j-1,i-1);
                   }
                   if (Am(j-1,i+1) != na) {
                     n = n+1;
                     sum = sum + Am(j-1,i+1);
                   }
                   if (Am(j+1,i-1) != na) {
                     n = n+1;
                     sum = sum + Am(j+1,i-1);
                   }
                   if (Am(j+1,i+1) != na) {
                     n = n+1;
                     sum = sum + Am(j+1,i+1);
                   }
                 }

                 Bm(j,i) = sum/n;
              }
            i = i + istep;
          }
        j = j + jstep;
      }

}
    return Bm;
"

fillin <- cxxfunction(signature(A = "numeric", missval = "integer", nneighb = "integer",
                                direction = "integer", ythenx = "integer"),
                      body = src, plugin="Rcpp")


## * Main function to be called in scripts
fill_raster_gaps <- function(r, na.recode=-9999, n.neighb=8) {
    ## r         : input raster with gaps (cells with non-finite value: NA, NaN, Inf...)
    ## na.recode : value to replace NA cells with, before being processed in C++
    ## n.neighb  : number of neighbours to take mean from. Either 4 or 8 (including diagonals)
    ## Returns the input raster with gaps filled in
    lst <- list()
    dir=1L
    for (dir in 1L:4L) {
        y.then.x=0L
        for (y.then.x in 0L:1L) {
            r1 <- r
            r1[!is.finite(r1[])] <- na.recode
            prevvals <- rep(0, length(r1[]))
            while (!identical(prevvals, r1[])) {
                prevvals <- r1[]
                m <- fillin(as.matrix(r1), na.recode, n.neighb, dir, y.then.x)
                r1 <- setValues(r1, m)
            }
            r1 <- reclassify(r1, cbind(missval, NA))
            lst <- c(lst, r1)
        }
    }
    rout <- calc(stack(lst), mean)
    return(rout)
}