#######################
### SETUP WORKSPACE ###
#######################

library(maptools)
library(gstat)
library(spatstat)
library(raster)
library(foreach)
library(doParallel)

## Clear the workspace
rm(list=ls())

##############################################################
### NOTE: This line should be changed to point to the folder
###       where you have unzipped the contents of the
###       replication archive.
##############################################################

## Set working directory
setwd("C:/warren_jpr2015_rep")

## Import data from load_data.R 
load('load_data.RData')


###################################
### SET GLOBAL PARAMETER VALUES ###
###################################

#Define number of processor cores 
#to use for parallel calculations
num_cores <- 2


###############
### Kriging ###
###############

#Convert raster template to SpatialPoints
r.pnts <- as(r, "SpatialPoints")

#Loop over variables
#(Radios, Cell Phones, Electrification)
varlist <- c('rad', 'cell', 'elec')
for (var in varlist) {
  
  cat(paste0('Creating [', var, '.im] \n'))
  
  #Select variable to be interpolated
  dhs$value <- dhs[[var]]
  
  #Exclude follow-up Uganda survey
  surveys <- dhs[dhs$dhsyear <= 2010, ]

  #Calculate sample variogram using all surveyed countries
  v <- variogram(value~1, surveys, width=1000, cutoff=500000)
  
  #Define function to estimate variogram,
  #given values for the parameters:
  #rn: "range"
  #k:  "kappa"
  v.fit <- function(rn, k) {
    fit.variogram(v, 
      vgm(
        psill = 1, 
        model = "Ste", 
        range = rn,
        kappa = k, 
        nugget = 1),
      fit.ranges=FALSE, debug.level=0)
  }

  #Define function to esimtate the sum of errors
  #for the fitted variogram
  v.fit.err <- function(x) {
    rn <- x[1]
    k <- x[2]
    fit <- v.fit(rn, k)
    err <- attr(fit, "SSErr")
    err
  }

  #Use numerical optimization to find 
  #best fitting parameters for "range" and "kappa"
  opt.v.fit <- optim(c(1000, 0.1), v.fit.err)
  rn_best <- opt.v.fit$par[1]   
  k_best <- opt.v.fit$par[2]   
  
  #Calculate variogram at optimal parameter values
  cat("Fitting variogram using: \n range =", rn_best, 
                                "\n kappa =", k_best, '\n\n')
  
  v.fit.best <- v.fit(rn_best, k_best) 
  
  rm(v, opt.v.fit, surveys)
  
  #Initialize multi-core parallel processing
  registerDoParallel(num_cores)
    
  #Loop through surveys by country-year
  krig_rbind <- foreach(i = c(1:nrow(dhs.cowcode.year)), 
                  .combine='rbind', 
                  .packages=c('foreach', 'doParallel', 'sp', 
                              'gstat', 'spatstat', 'maptools')) %dopar% {
                          
      library(gstat)
      
      #Define country code and year
      c <- dhs.cowcode.year$cowcode[i]
      y <- dhs.cowcode.year$dhsyear[i]
      
      #Extract survey
      survey <- dhs[ (dhs$cowcode == c &
                      dhs$dhsyear == y), ]
      
      #Extract survey country shape
      shape <- cshapes[cshapes$cowcode == c,]
      
      #Define prediction grid using raster template
      pred.grid <-  r.pnts[over(shape, r.pnts, returnList=TRUE)[[1]]]
      
      #Define function to calculate kriging prediction errors,
      #for neighborhood size n
      krig.err <- function (n) {
        n <- round(n)
        krig <- krige.cv(value~1, survey, v.fit.best, nmax=n)
        rmse <- sqrt(mean((krig$observed - krig$var1.pred)^2))
        rmse
      }
      
      #Set minium and maximum values for n
      n_min <- 1
      n_max <- 200
      
      #Set optimization tolerance
      opt_tol <- 0.5
      
      #Use numerical optimization to calculate
      #optimal neighborhood size
      opt.krig.err <- optimize(krig.err, c(n_min, n_max), tol = opt_tol)
      
      #Record opimal value, and it's associated root mean squared error
      n_best <- round(opt.krig.err$minimum)
      rmse_best <- opt.krig.err$objective
           
      #Calculate interpolated surface using optimal paramters
      krig.opt <- krige(value~1, survey, pred.grid, v.fit.best, nmax=n_best)
      
      #Constrain values to [0,1]
      krig.opt$var1.pred[krig.opt$var1.pred < 0] <- 0
      krig.opt$var1.pred[krig.opt$var1.pred > 1] <- 1
      
      #Assign country-year identifiers
      krig.opt$cowcode <- c
      krig.opt$dhsyear <- y
      
      #Record optimized parameters
      krig.opt$krig_rn <- rn_best
      krig.opt$krig_k <- k_best
      krig.opt$krig_n <- n_best
      krig.opt$krig_rmse <- rmse_best
      
      #Estimate quadratic IDW interpolation for comparison
      detach(package:gstat)
      library(spatstat)
      survey.ppp <- as.ppp(survey[,'value'])
      idw.cv <- idw(survey.ppp, power=2, at="points")
      
      #Calculate errors from IDW
      idw.err <- idw.cv - survey.ppp$marks
      rmse.idw <- sqrt(mean((idw.err)^2))
      krig.opt$idw_rmse <- rmse.idw
      
      #Output predicted surfaces
      #(results from parallel iterations will be
      #combined via 'rbind' into a single data frame)
      krig.opt
        
    } #end foreach i = country-year
  
  #Pull out 2nd Uganda survey
  krig_ug2 <- krig_rbind[krig_rbind$dhsyear > 2010,]
  krig_rbind <- krig_rbind[krig_rbind$dhsyear <= 2010,]
  name.rbind <- paste0('krig.', var)
  assign(name.rbind, krig_rbind)
  
  #Combine countries to create single raster surface
  krig_rast <- rasterize(krig_rbind, r, "var1.pred") 
  krig_rast <- as(krig_rast, "SpatialGridDataFrame")
  krig_rast <- as.im.SpatialGridDataFrame(krig_rast)
  name.im <- paste0(var, ".im")
  assign(name.im, krig_rast)
  
  #Create raster surface for 2nd Uganda survey
  krig_rast <- rasterize(krig_ug2, r, "var1.pred") 
  krig_rast <- as(krig_rast, "SpatialGridDataFrame")
  krig_rast <- as.im.SpatialGridDataFrame(krig_rast)
  name.im <- paste0(var, ".ug2.im")
  assign(name.im, krig_rast)
  
  #Save results data frame
  df.name <- paste0('krig.', var)
  assign(df.name, krig_rbind)
  
} #end for var


####################
### Save Results ###
####################
varlist_out <- c(paste0('krig.', varlist), 
                 paste0(varlist, '.im'), 
                 paste0(varlist, '.ug2.im'))
save(list=varlist_out, file="krig_data.RData")


##########################
### Kriging Validation ###
##########################

##Mean Squared Prediction Error
##(Kriging vs. Quadratic IDW)

#Radios 
mspe.rad <- cbind(krig.rad$idw_rmse^2, 
                  krig.rad$krig_rmse^2)

#Cell Phones
mspe.cell <- cbind(krig.cell$idw_rmse^2, 
                   krig.cell$krig_rmse^2)

#Plot comparison
lim <- max(mspe.rad, mspe.cell)
plot(NA, type='n', xlim=c(0,lim), ylim=c(0,lim),
     xlab="IDW Errors",
     ylab="Kriging Errors")
points(mspe.rad, pch=16, col="purple")
points(mspe.cell,pch=16, col="green")
vals <- seq(0, 0.5, 0.01)
lines(vals, vals, lwd=5, col='darkgray')

