rm(list=ls())

library(mcmc)
library(MASS)
library(GenKern)

set.seed(1234)

########################
#make some data
########################
n <- 50
x <- runif(n, 0, 100)
y <- runif(n, 0, 100)

D <- as.matrix(dist(cbind(x, y)))

phi <- 3/100
sigmasq <- 50
tausq <- 20
mu <- 150

s <- (sigmasq*exp(-phi*D))
w <-  mvrnorm(1, rep(0, n), s)
Y <- mvrnorm(1, rep(mu, n) + w, tausq*diag(n))
X <- as.matrix(rep(1, length(Y)))

lmObject <- lm(Y ~ X)
lmObject.cov <- vcov(lmObject)

write(t(Y), "Ytest.txt", ncolumns=1)
write(t(X), "Xtest.txt", ncolumns=1)
write(t(cbind(x,y)), "Sitestest.txt", ncolumns=2)

########################
#some priors for sigma^2
#tau^2 and phi
########################
#IG sigma^2 and tau^2
a.sig <- 2 
b.sig <- 100
a.tau <- 2
b.tau <- 100

#U phi
a.phi <- 3/3
b.phi <- 3/150

########################
##on to the metrop
########################

##function to evaluate spatial correlation
Sigma.marginal <- function(theta)
{
  sigmasq.cand <- exp(theta[2])
  tausq.cand <- exp(theta[3])
  phi.cand <- exp(theta[4])
  
  Sigma <- sigmasq.cand*exp(-phi.cand*D)+tausq.cand*diag(n)
  SigmaInv <- chol2inv(chol(Sigma))
  logDetSigma <- determinant(Sigma, log=TRUE)$modulus[1]
  
  out = list(SigmaInv, logDetSigma)
  return(out)
}

##metrop target
target <- function (theta)
{
  mu.cand <- theta[1]
  sigmasq.cand <- exp(theta[2])
  tausq.cand <- exp(theta[3])
  phi.cand <- exp(theta[4])
  
  if(phi.cand > b.phi && phi.cand < a.phi){

    SigmaMarginal = Sigma.marginal(theta)
    SigmaInv = SigmaMarginal[[1]]
    logDetSigma = SigmaMarginal[[2]]

##    Sigma <- sigmasq.cand*exp(-phi.cand*D)+tausq.cand*diag(n)
##    SigmaInv <- chol2inv(chol(Sigma))
##    logDetSigma <- determinant(Sigma, log=TRUE)$modulus[1]
    
    out <- -(a.sig+1)*log(sigmasq.cand)-b.sig/sigmasq.cand-(a.tau+1)*log(tausq.cand)-b.tau/tausq.cand-0.5*logDetSigma-0.5*(t(Y-X%*%mu.cand)%*%SigmaInv%*%(Y-X%*%mu.cand))
    
    return(out)
  } else{
    return(-Inf) #this tells metrop to skip and propose another
  }
}

#metrop stuff
NITER <- 50000
BatchLength <- 1

#tuning
tuning <- diag(c(diag(lmObject.cov), 0.25, 0.25, 0.25))
tuning <- chol(tuning)

#starting values
inits <- c(100, log(1000), log(1000), log(0.03))

metrop.out <- metrop(target, initial=inits, nbatch=NITER, blen=BatchLength, scale=tuning, debug=T)

########################
#summary stuff
########################
#accept rate
metrop.out$accept


myqnt <- c(0.50, 0.025, 0.975)
start <- 30001
beta.hat <- quantile(metrop.out$batch[start:nrow(metrop.out$batch),1],myqnt)
sigmasq.hat <- quantile(exp(metrop.out$batch[start:nrow(metrop.out$batch),2]),myqnt)
tausq.hat <- quantile(exp(metrop.out$batch[start:nrow(metrop.out$batch),3]),myqnt)
phi.hat <- quantile(3/metrop.out$batch[start:nrow(metrop.out$batch),4],myqnt)

beta.hat
sigmasq.hat
tausq.hat
phi.hat

##pdf("ConvergencePlots.pdf")
par(mfrow=c(2,2))
plot(metrop.out$batch[start:nrow(metrop.out$batch),1],type="l",main="mu",ylab="",xlab="samples")
abline(h=mu, col="red")
plot(exp(metrop.out$batch[start:nrow(metrop.out$batch),2]),type="l",main="sigma^2",ylab="",xlab="samples")
abline(h=sigmasq, col="red")
plot(exp(metrop.out$batch[start:nrow(metrop.out$batch),3]),type="l",main="tau^2",ylab="",xlab="samples")
abline(h=tausq, col="red")
plot(3/exp(metrop.out$batch[start:nrow(metrop.out$batch),4]),type="l",main="3/phi",ylab="",xlab="samples")
abline(h=3/phi, col="red")
##dev.off()