###################################################
### chunk number 1: libs
###################################################
options(width = 67) 
library(spBayes)
library(MBA)
library(geoR)
library(fields)
library(rgl)
library(sp)
library(maptools)
library(rgdal)
library(classInt)
library(gstat)
library(lattice)
library(xtable)


###################################################
### chunk number 2: nutSurf
###################################################
dat <- read.table("CostaRica/T4.csv", header=T, sep=",")
coords <- as.matrix(dat[,c("X","Y")])
nut.names <- c("Ca","K","Mg")
log.nut <- log(dat[,nut.names])

par(mfrow=c(2,2))
for(i in 1:length(nut.names)){
  surf <- mba.surf(cbind(coords,data=log.nut[,i]), no.X=100, no.Y=100)$xyz.est
  image.plot(surf, main=paste("Log ",nut.names[i],sep=""))
  points(coords)
}


###################################################
### chunk number 3: geoRVariog
###################################################
max <- 0.25*max(as.matrix(dist(dat[,c("X","Y")])))
bins <- c(9,8,9)

par(mfrow=c(3,1))
for(i in 1:length(nut.names)){
  vario <- variog(coords=coords,data=log.nut[,i], uvec=(seq(0,max, length=bins[i])))

  fit <-variofit(vario, ini.cov.pars=c(0.3, 20/-log(0.05)), ##sigma^2 and 1/phi 
                     cov.model="exponential", minimisation.function="nls", weights="equal")

  plot(vario, pch=19, main=paste("Log ",nut.names[i],sep=""))
  lines(fit)
  abline(h=fit$nugget, col="blue")
  abline(h=fit$cov.pars[1]+fit$nugget, col="green")
  abline(v=-log(0.05)*fit$cov.pars[2], col="red3")
}


###################################################
### chunk number 4: spMvLM
###################################################
q <- 3
##Specify starting values and collect samples
A.starting <- diag(0.5,q)[lower.tri(diag(1,q), TRUE)]
L.starting <- diag(0.05,q)[lower.tri(diag(1,q), TRUE)]
A.tuning <- matrix(0.01,q,q); diag(A.tuning) <- 0.1;
A.tuning <- A.tuning[lower.tri(diag(1,q), TRUE)]
L.tuning <- rep(0.03,length(L.starting))

n.samples <- 10

nut.spMvLM <- spMvLM(list(Ca~1,K~1,Mg~1), coords=coords, data=log.nut,
                     starting=list("beta"=rep(1,q), "phi"=rep(3/20,q),"A"=A.starting,"L"=L.starting), 
                     sp.tuning=list("phi"=rep(0.3,q),"A"=A.tuning, "L"=L.tuning),
                     priors=list("phi.Unif"=rep(c(3/200,3/10),q),
                       "K.IW"=list(q+1, diag(0.1,q)), "Psi.IW"=list(q+1, diag(0.1,q))),
                     cov.model="exponential",
                     n.samples=n.samples, sub.samples=c(1,n.samples,1),
                     verbose=FALSE, n.report=500)


##save(nut.spMvLM, file = "R-data/nut.spMvLM")
load(file = "R-data/nut.spMvLM")

p.samples <- nut.spMvLM$p.samples
p.samples[,paste("phi_",1:q,sep="")] <- 3/p.samples[,paste("phi_",1:q,sep="")]
colnames(p.samples)[colnames(p.samples) == paste("phi_",1:q,sep="")] <- paste("Eff. range ",1:q,sep="")

summary(mcmc(p.samples))$quantiles[,c(1,3,5)]


###################################################
### chunk number 5: spLMW
###################################################
Ca.resids <- resid(lm(Ca~1, data=log.nut))
K.resids <- resid(lm(K~1, data=log.nut))
Mg.resids <- resid(lm(Mg~1, data=log.nut))

w <- rowMeans(nut.spMvLM$sp.effects)
w.Ca <- w[seq(1,length(w),q)]
w.K <- w[seq(2,length(w),q)]
w.Mg <- w[seq(3,length(w),q)]

expand.range <- function(x, p=.05){
  x <- range(x, na.rm=TRUE)
  x[1] <- x[1]-p*abs(x[1]);
  x[2] <- x[2]+p*abs(x[2]);
  x
}

res <- 100
par(mfrow=c(3,2))
surf <- mba.surf(cbind(coords,Ca.resids), no.X=res, no.Y=res, extend=FALSE)$xyz.est
z.lim <- expand.range(surf[["z"]]); 
image.plot(surf, zlim=z.lim, main="Ca lm residuals")
points(coords)

surf <- mba.surf(cbind(coords,w.Ca), no.X=res, no.Y=res, extend=FALSE)$xyz.est
image.plot(surf, zlim=z.lim, main="Ca spatial effects")
points(coords)

surf <- mba.surf(cbind(coords,K.resids), no.X=res, no.Y=res, extend=FALSE)$xyz.est
z.lim <- expand.range(surf[["z"]])
image.plot(surf, zlim=z.lim, main="K lm residuals")
points(coords)

surf <- mba.surf(cbind(coords,w.K), no.X=res, no.Y=res, extend=FALSE)$xyz.est
image.plot(surf, zlim=z.lim, main="K spatial effects")
points(coords)

surf <- mba.surf(cbind(coords,K.resids), no.X=res, no.Y=res, extend=FALSE)$xyz.est
z.lim <- expand.range(surf[["z"]], 0.2)
image.plot(surf, zlim=z.lim, main="Mg lm residuals")
points(coords)

surf <- mba.surf(cbind(coords,w.Mg), no.X=res, no.Y=res, extend=FALSE)$xyz.est
image.plot(surf, zlim=z.lim, main="Mg spatial effects")
points(coords)


###################################################
### chunk number 6: spPredict
###################################################
x.range <- range(coords[,1])
y.range <- range(coords[,2])

pred.coords <- expand.grid(seq(x.range[1], x.range[2], by=4),
                           seq(y.range[1], y.range[2], by=4))

m <- nrow(pred.coords)

pred.X <- mkMvX(list(matrix(1,m,1), matrix(1,m,1), matrix(1,m,1)))

nut.pred <- spPredict(nut.spMvLM, start=1, end=2, pred.coords=pred.coords, pred.covars=pred.X)

##save(nut.pred, file = "R-data/nut.pred")
load(file = "R-data/nut.pred")


###################################################
### chunk number 7: spPredMean
###################################################
y.pred.mu <- apply(nut.pred$y.pred, 1, mean)
y.pred.sd <- apply(nut.pred$y.pred, 1, sd)

Ca.pred.mu <- y.pred.mu[seq(1,length(y.pred.mu),q)]
K.pred.mu <- y.pred.mu[seq(2,length(y.pred.mu),q)]
Mg.pred.mu <- y.pred.mu[seq(3,length(y.pred.mu),q)]

Ca.pred.sd <- y.pred.sd[seq(1,length(y.pred.sd),q)]
K.pred.sd <- y.pred.sd[seq(2,length(y.pred.sd),q)]
Mg.pred.sd <- y.pred.sd[seq(3,length(y.pred.sd),q)]

nut.pred.grid <- as.data.frame(list(x=pred.coords[,1], y=pred.coords[,2],
                                     Ca.mu=Ca.pred.mu, K.mu=K.pred.mu, Mg.mu=Mg.pred.mu,
                                     Ca.sd=Ca.pred.sd, K.sd=K.pred.sd, Mg.sd=Mg.pred.sd))
                                     
coordinates(nut.pred.grid) <- c("x", "y") # promote to SpatialPointsDataFrame
gridded(nut.pred.grid) <- TRUE            # promote to SpatialGridDataFrame

toImage <- function(x){as.image.SpatialGridDataFrame(x)}

res <- 100
par(mfrow=c(3,2))

surf <- mba.surf(cbind(coords,log.nut[,"Ca"]), no.X=res, no.Y=res, extend=FALSE)$xyz.est
z.lim <- range(surf[["z"]], na.rm=TRUE)
image.plot(surf, xaxs = "r", yaxs = "r", main="Interpolation of observed Ca")
points(coords)

image.plot(toImage(nut.pred.grid["Ca.mu"]), xaxs = "r", yaxs = "r", zlim=z.lim, main="Mean of Ca prediction")
points(coords)

surf <- mba.surf(cbind(coords,log.nut[,"K"]), no.X=res, no.Y=res, extend=FALSE)$xyz.est
z.lim <- range(surf[["z"]], na.rm=TRUE)
image.plot(surf, xaxs = "r", yaxs = "r", main="Interpolation of observed K")
points(coords)

image.plot(toImage(nut.pred.grid["K.mu"]), xaxs = "r", yaxs = "r", zlim=z.lim, main="Mean of K prediction")
points(coords)

surf <- mba.surf(cbind(coords,log.nut[,"Mg"]), no.X=res, no.Y=res, extend=FALSE)$xyz.est
z.lim <- range(surf[["z"]], na.rm=TRUE)
image.plot(surf, xaxs = "r", yaxs = "r", main="Interpolation of observed Mg")
points(coords)

image.plot(toImage(nut.pred.grid["Mg.mu"]), xaxs = "r", yaxs = "r", zlim=z.lim, main="Mean of Mg prediction")
points(coords)


###################################################
### chunk number 8: spPredSD
###################################################
par(mfrow=c(3,2))
surf <- mba.surf(cbind(coords,log.nut[,"Ca"]), no.X=res, no.Y=res, extend=FALSE)$xyz.est
image.plot(surf, main="Interpolation of observed Ca")
points(coords)

surf <- mba.surf(cbind(pred.coords,Ca.pred.sd), no.X=res, no.Y=res, extend=FALSE)$xyz.est
image.plot(surf, main="Standard deviation of Ca prediction")
points(coords)

surf <- mba.surf(cbind(coords,log.nut[,"K"]), no.X=res, no.Y=res, extend=FALSE)$xyz.est
image.plot(surf, main="Interpolation of observed K")
points(coords)

surf <- mba.surf(cbind(pred.coords,K.pred.sd), no.X=res, no.Y=res, extend=FALSE)$xyz.est
image.plot(surf, main="Standard deviation of K prediction")
points(coords)

surf <- mba.surf(cbind(coords,log.nut[,"Mg"]), no.X=res, no.Y=res, extend=FALSE)$xyz.est
image.plot(surf, main="Interpolation of observed Mg")
points(coords)

surf <- mba.surf(cbind(pred.coords,Mg.pred.sd), no.X=res, no.Y=res, extend=FALSE)$xyz.est
image.plot(surf, main="Standard deviation of Mg prediction")
points(coords)