library(rjags)
library(MASS)
library(sm)

df = read.csv("colorado_tmin_ave.csv")
predictors = as.matrix(df[,c('lon','lat','elev')])
response = df$tmin
n = nrow(df)
K = ncol(predictors)

### temperature as a function of lat, lon and elevation

Y = response
## Linear regression
fit = lm(tmin~lon+lat+elev,data=df)

initf = function()list(
    alpha = coefficients(fit)[1],
    betas = coefficients(fit)[-1],
    sigmap = sd(residuals(fit)))

### JAGS model
model_string <- "model{

  # Likelihood
  for(i in 1:n){
    Y[i]   ~ dnorm(mu[i],sigma)
    mu[i] <- beta[1] + beta[2]*lon[i] + beta[3]*lat[i] + beta[4]*elev[i]
  }

  # Prior for beta
  for(j in 1:4){
    beta[j] ~ dnorm(0,0.0001)
  }

  # Prior for the inverse variance

  sigma     <- 1/sqrt(inv.var)

  inv.var   ~ dgamma(0.01, 0.01)
}"

 model1.spec<-textConnection(model_string)

model <- jags.model(model1.spec, data = list(Y = Y,n=n,lat=df$lat, lon=df$lon, elev=df$elev), n.chains = 3, n.adapt= 10000)
update(model, 10000); # Burnin for 10000 samples
mcmc_samples <- coda.samples(model, variable.names=c("beta","sigma","mu"), n.iter=20000)
zz=summary(mcmc_samples)

samples <- jags.samples(model, c('beta', 'sigma','mu'), 1000)

xx=samples$mu[,,1]
yy=apply(xx,1,stats::quantile)
yy=t(yy)
ymed = yy[,3]

### Plot the posterior
dev.new()
par(mfcol=c(2,2))
beta1 = mcmc_samples[[1]][,1]
beta2 = mcmc_samples[[1]][,2]
beta3 = mcmc_samples[[1]][,3]
beta4 = mcmc_samples[[1]][,4]
sigmasim = mcmc_samples[[1]][,5]

hist(beta1,probability=T)
sm.density(beta1,add=TRUE)

hist(beta2,probability=T)
sm.density(beta2,add=TRUE)

hist(beta3,probability=T)
sm.density(beta3,add=TRUE)

hist(beta4,probability=T)
sm.density(beta4,add=TRUE)

hist(sigmasim,probability=T)
sm.density(sigmasim,add=TRUE)