rm(list=ls())
#getwd()
#setwd("pathway")

#Load all the necessary packages (need to download and install them first)
library(nlme)
library(bbmle)
library(rjags)
library(ggplot2)
library(jagsUI) 
library(glmmADMB)
library(lme4)
library(MuMIn)


#### Part I: Preparing the data ####
orig_data <- read.table("hypericum_data_94_07.txt", header=T)
dt <- subset(orig_data, !is.na(ht_init) & rp_init > 0 & year<1997) ## change variables to include only reproductive individuals
yr <- unique(dt$year)
dt$lgh <- log(dt$ht_init)
dt$lfr <- log(dt$rp_init)
site <- unique(dt$bald)

#### Part II: Frequentists analysis of three models ####
table_coef <- array(0,c(3,2)) ## create table for comparisons
colnames(table_coef) <- c("intercept","slope")
rownames(table_coef) <- c("no mixed","random intercept","random intercept & slope")

# a) no mixed effects #

m1 <- lm(lfr~lgh,data=dt) 
summary(m1)
table_coef[1,] <- m1$coefficients

Beta_1 <- Beta_2 <- array(0,c(1,length(site))) ## create table for coefficients of separate models per population (14 sites)
colnames(Beta_1)=site
colnames(Beta_2)=site

plot(dt$lgh,dt$lfr,pch=16,ylab="log(fruits)",
  xlab="log(height)",col="grey",cex=0.5,ylim=c(0,8),main="Individual populations") ## plot data points
for (j in 1:length(site)){
  MU <- lm(lfr~lgh,subset=(bald==site[j]),data=dt)
  Mi <- summary(MU)
  x1 <- dt$lgh[dt$bald==site[j]]
  K <- order(x1)
  lines(sort(x1),predict(MU)[K],col="red",lwd=1.1) ## plot lines for separate models per population
  Beta_1[j] <- Mi$coefficients[1,1]
  Beta_2[j] <- Mi$coefficients[2,1]}
I <- order(dt$lgh)
lines(sort(dt$lgh),predict(m1)[I],col="blue",lwd=3) ## plot line of general linear model

Beta_1
Beta_2
mB1 <- mean(rowSums(Beta_1)/14)
mB2 <- mean(rowSums(Beta_2)/14)
hist(mB1-Beta_1,5,main="Intercept")
hist(mB2-Beta_2,5,main="Slope")

# b) random intercept #

dt$fbald <- factor(dt$bald)
M1 <- lme(lfr~lgh,random=~1|fbald,data=dt,method ="ML")
#M1 <- lme(lfr~lgh,random=~1|fbald,data=dt,method ="REML")
M1a <- lmer(lfr~lgh + (1|fbald) ,data=dt)
summary(M1)
summary(M1a)
table_coef[2,] <- M1$coefficients$fixed

#calculate variance explained

r.squaredGLMM(M1)


F0 <-fitted(M1,level=0)
F1 <-fitted(M1,level=1)
lgh <- sort(dt$lgh)
plot(lgh,predict(m1)[I],lwd=1,type="l",ylab="log(fruits)",xlab="log(height)",ylim=c(0,8),main="Random Intercept",col="black")
points(dt$lgh,dt$lfr,pch=16,ylab="log(fruits)",xlab="log(height)",col="grey",cex=0.5,ylim=c(0,8))
for (j in 1:length(site)){
  x1 <- dt$lgh[ dt$bald==site[j]]
  y1 <- F1[dt$bald==site[j]]
  K <-  order(x1)
  lines(sort(x1),y1[K],col="red")}
lines(lgh,F0[I],lwd=3,type="l",col="blue")

# c) random intercept and slope #

M11 <- lme(lfr~lgh,random=~1 + lgh|fbald,data=dt,method ="ML")
#M11 <- lme(lfr~lgh,random=~1 + lgh|fbald,data=dt,method ="REML")
summary(M11)
table_coef[3,] <- M11$coefficients$fixed
r.squaredGLMM(M11)


F0 <-fitted(M11,level=0)
F1 <-fitted(M11,level=1)
lfrs <- sort(dt$lgh)

plot(lfrs,predict(m1)[I],lwd=2,type="l",ylab="log(fruits)",xlab="log(height)",ylim=c(0,8),main="Random Intercept & Slope",col="green")
points(dt$lgh,dt$lfr,pch=16,ylab="log(fruits)",xlab="log(height)",col="grey",cex=0.5,ylim=c(0,8))
for (j in 1:length(site)){
  x1 <- dt$lgh[dt$bald==site[j]]
  y1 <- F1[dt$bald==site[j]]
  K <-  order(x1)
  lines(sort(x1),y1[K],col="red")}
lines(lgh,F0[I],lwd=3,type="l",col="blue")

#plot last model with confidence intervals
cff <- M11$coefficients$fixed
MyData <- data.frame(lgh =seq(2,4.5,0.5)) 
X <- model.matrix(~lgh,MyData)
FitManual <- X %*% cff
SE <- sqrt(diag(X %*% vcov(M11) %*% t(X)))
plot(dt$lgh,dt$lfr,pch=16,ylab="log(fruits)",xlab="log(height)",col="grey",cex=0.5,ylim=c(0,8))
lines(lgh,F0[I],lwd=3,type="l",col="blue")
lines(MyData$lgh,FitManual+1.96*SE, lty=2) 
lines(MyData$lgh,FitManual-1.96*SE, lty=2)

# d) compare three models #
table_coef
AICtab(m1,M1,M11,weights=TRUE,base=TRUE)

par(mfrow=c (1,2))
plot(dt$lgh,residuals(M11))
plot(fitted(M11),residuals(M11))

par(mfrow=c (2,2))
b2 <- seq(-16,16,0.5)
b1 <- seq(-5,5,0.2)
hist(mB1-Beta_1,breaks=b2, probability=TRUE,main="Pop Intercept",ylim=c(0,1.1))
hist(mB2-Beta_2,breaks=b1, probability=TRUE,main="Pop Slope",ylim=c(0,1.1))
hist(M11$coefficients$random$fbald[,1],breaks=b2, probability=TRUE,
     main="random effects intercept",xlab="Intercept",ylim=c(0,1.1))
lines(seq(-5,5,0.5),dnorm(seq(-5,5,0.5),0,2.2037935),lwd=2)
hist(M11$coefficients$random$fbald[,2],breaks=b1, probability=TRUE,
     main="random effects slope",xlab="Slope",ylim=c(0,1.1))
lines(seq(-15,15,0.1),dnorm(seq(-15,15,0.1),0, 0.6108079),lwd=2)

#### Part III:Bayesian linear mixed models ####

## Uninformative priors for fixed and random effects ##
dt$ppn[dt$bald==1] <- 1
dt$ppn[dt$bald==29] <- 2
dt$ppn[dt$bald==32] <- 3
dt$ppn[dt$bald==42] <- 4
dt$ppn[dt$bald==50] <- 5
dt$ppn[dt$bald==57] <- 6
dt$ppn[dt$bald==59] <- 7
dt$ppn[dt$bald==62] <- 8
dt$ppn[dt$bald==67] <- 9
dt$ppn[dt$bald==87] <- 10
dt$ppn[dt$bald==88] <- 11
dt$ppn[dt$bald==91] <- 12
dt$ppn[dt$bald==93] <- 13
dt$ppn[dt$bald==103] <- 14

X <- model.matrix(~lgh,data=dt)
K <- ncol(X)
Npop <- length(unique(dt$bald))
n <- nrow(dt)

## a) random intercept
model = "LMM_rand_intercept.R"

# Bundle data
win.data=list(
  Y = dt$lfr,
  X = X,
  K = K,
  N = n,
  pp = as.numeric(dt$ppn),
  Npop = Npop
)

# Inits function
inits <- function(){
  list(Betas = rnorm(K,0,1),
  sigma.plot = runif(1,0.001,10),
  sigma.eps = runif(1,0.001,10))}

# Parameters 
params <- c("alpha.p","Betas","sigma.plot","sigma.eps")

# Call JAGS
jm=jags(win.data,inits,params,model.file=model,n.chains=3, n.iter=2000, n.burnin=1000, n.thin=1)
jm <- update(jm, n.iter=10000)
jm	

#plot(jm)

## b) random intercept and slope
model = "LMM_rand_intercept&slope.R"

# Bundle data
win.data1=list(
  Y = dt$lfr,
  X = dt$lgh,
  N = n,
  pp = as.numeric(dt$ppn),
  Npop = Npop
)

# Inits function 
inits <- function(){
  list(sigma.plot.a = runif(1,0.001,10),
  sigma.plot.b = runif(1,0.001,10),
  sigma.eps = runif(1,0.001,10))}

# Parameters 
params <- c("Betas1","Betas2","alpha.p","beta.p","sigma.plot.a","sigma.plot.b","sigma.eps")

# Call JAGS
jm=jags(win.data1,inits,params,model.file=model,n.chains=3, n.iter=2000, n.burnin=1000, n.thin=1)
jm <- update(jm, n.iter=10000)
jm			

#plot(jm)