#### R SCRIPT TO FIT THE GAMMA REGRESSION MODEL TO THE MAGNESIUM DATA
### Script by Pablo García-Díaz, Instituto de Ecología Regional (UNT-CONICET), Tucumán, Argentina. E-mail: pablo.garciadiaz.eco@gmail.com

rm(list=ls())

set.seed(2012)

options(scipen=999)

### Load the libraries
library(nimble)
library(coda)
library(ggplot2)
library(tidyr)
library(tidyverse)
library(hrbrthemes)
library(viridis)
library(beepr)
library(nimbleHMC)

### Load the data
soil.data<-read.table("./soilmeans.csv", header=TRUE, sep=",")

summary(soil.data)

### Remove NAs from dataset
### Identify NAs
id.na<-which(is.na(soil.data$Mg)==TRUE)

id.na

#### Keep rows with Calcium measurements
soil.data<-soil.data[-id.na, ]

### Have a look
summary(soil.data)

### Define plots as a numeric variable for modelling
soil.data$plot.num<-as.numeric(as.factor(soil.data$PLOT))

soil.data$plot.num

### Number of plots
n.plot<-max(soil.data$plot.num)

n.plot

### Define depth at which the sample was taken as a numeric variable for modelling
soil.data$depth.num<-as.numeric(as.factor(soil.data$SOIL_POSITION))

soil.data$depth.num

### Number of levels of soil depth
n.depth<-max(soil.data$depth.num)

n.depth

unique(cbind(as.numeric(as.factor(soil.data$SOIL_POSITION)), soil.data$SOIL_POSITION))

### Define forest types (invaded and native) as a numeric variable for modelling
soil.data$forest.num<-as.numeric(as.factor(soil.data$FOREST_TYPE))

soil.data$forest.num

### Number of categories (two: invaded and native)
n.forest<-max(soil.data$forest.num)

n.forest

### Specify sample size
sample.size<-nrow(soil.data)

sample.size

### Specify Gamma regression function in nimble 
mg.soil<- nimbleCode({
    
    ### Define priors
    ### Plot-level random effect
    for(j in 1:n.plot){

        	alpha.plot[j]~dnorm(0, var=var.plot)

    }

    ### Variance of the random effect
    var.plot~dnorm(0, 1000)

    ### Fixed depth x forest type effects
	for(z in 1:n.depth){

            	for(n in 1:n.forest){


		                alpha.comp[z, n]~dnorm(0, var=1000)

            }

	}

    ### Extra variance in Mg in soil
    var.mg~dunif(0, 1000)

	#### Model likelihood: Gamma distribution
   	 for(i in 1:sample.size){

            ### Define the mean as a function of plot and depth x forest type
        	log(mu.mg[i])<-alpha.plot[plot[i]]+alpha.comp[depth[i], type[i]]

            ### Gamma distribution
        	soil.mg[i]~dgamma(mean=mu.mg[i], sd=sqrt(var.mg))       

            #### Calculate squared residuals for Bayesian p-value based on the Root Mean Squared Error
            soil.mg.new[i]~dgamma(mean=mu.mg[i], sd=sqrt(var.mg))       

            res.sim[i]<-pow(soil.mg.new[i]-mu.mg[i], 2)

            res.obs[i]<-pow(soil.mg[i]-mu.mg[i], 2)


    	}

    ### Differences between the different depth x forest type combinations
    eff[1]<-alpha.comp[1, 2]-alpha.comp[1, 1]

    eff[2]<-alpha.comp[2, 2]-alpha.comp[2, 1]

    eff[3]<-alpha.comp[3, 2]-alpha.comp[3, 1]

    #### Estimate the Bayesian p-value
    ### Root mean squared error
    rmse.sim<-mean(res.sim[1:sample.size])

    rmse.obs<-mean(res.obs[1:sample.size])

    ### Bayesian p-value
    bpval<-step(rmse.sim-rmse.obs)

})

### Define initial values for the model
inits<-list(alpha.comp=matrix(0, ncol=n.forest, nrow=n.depth), alpha.plot=rep(0, n.plot), var.plot=runif(1, 0, 10),
        var.mg=1)
    
### Define constant values for the Bayesian model
constants<-list(n.plot=n.plot,  n.depth=n.depth,  n.forest=n.forest, sample.size=sample.size, plot=soil.data$plot.num,
			depth=soil.data$depth.num, type=soil.data$forest.num)

### Object with the Mg in soil data
data.mg<-list(soil.mg=soil.data$Mg)

###### Run the Bayesian model  
#### Timing the model
start.time<-Sys.time() 

start.time 
    
#### Model parameters to monitor
monitors<-list('alpha.comp', 'alpha.plot', 'var.plot', 'var.mg', 'bpval', 'eff')

#### Run the model using Hamiltonian Monte-Carlo with 3 chains, 50000 iterations, and a burn-in of 10,000 iterations
m1.mg<-nimbleHMC(mg.soil, data = data.mg, inits = inits, constants=constants, monitors = monitors, 
			thin = 5, niter=50000, nburnin=10000, nchains=3, samplesAsCodaMCMC = TRUE)	

### How long did it take?
end.time<-Sys.time()

end.time

time.taken<-end.time - start.time

time.taken

### Make some noise when it is done
beep(sound=5)

#### Posterior summary
summary(m1.mg)

### Convergence diagnostics
gelman.diag(m1.mg)

raftery.diag(m1.mg)

### Plot the chains
plot(m1.mg)

#### Convert to matrix for further manipulations
m2.mg<-do.call(rbind, m1.mg)

summary(m2.mg)

#### Backtransform to original scale and estimate for each depth x forest type
post.est<-exp(m2.mg[ , c(1:6)])

head(post.est)

#### Summary statistics: mean, standard deviation, and 95% credible intervals
mean.est<-apply(post.est, 2, mean)

sd.est<-apply(post.est, 2, sd)

low.qt<-apply(post.est, 2, function(x) quantile(x, 0.0275))

upp.qt<-apply(post.est, 2, function(x) quantile(x, 0.975))

### In a data frame & export
post.summ<-data.frame(coeff=colnames(post.est), mean=mean.est, sd=sd.est, low.qt=low.qt, upp.qt=upp.qt)

write.table(post.summ, "./BacktransformedMg.csv", sep=",")

#### Differences between forest types at different depths (log-scale)
eff.col<-grep("eff", colnames(m2.mg))

effects<-m2.mg[, eff.col]

### Summarise
summary(effects)

### Estimate the proportion of posterior iterations that are higher than zero (akin to p-value; 1 - colMeans(effects.bin) will be the  the proportion of posterior iterations that are lower than zero)
effects.bin<-ifelse(effects >0, 1, 0)

colMeans(effects.bin)

1-colMeans(effects.bin)


###### END OF SCRIPT ####