###Part I: Preparing the data
rm(list=ls())
getwd()
#Load all the necessary packages (need to download and install them first)

library(rethinking)
library(rstan)
library(bbmle)
library(nlme)



#### 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)

colnames(dt)[colSums(is.na(dt)) > 0] 
dt<-dt[,!names(dt) %in% c("ht_final","rp_final","st_init")]

colnames(dt[,sapply(dt,class)=='character'])
dt <-dt[,!names(dt) %in% c("stage","fate")]

#### Part II: Analysis of three models ####

# a) no mixed effects #

sdlnfruits <- sd(dt$lfr)

m1 <- ulam(
  alist(
    lfr ~ dnorm(mu,sigma),
    mu <- a + b*lgh,
    a ~ dnorm(0,100),
    b ~ dnorm(0,10),
    sigma ~ dunif(0,10)
  ),
  data = dt,chains =3,cores=3, log_lik = TRUE,
)

precis(m1,digits=2)
plot(m1,pars=c("a","b","sigma"))
pairs(m1,horInd = c(1:3), verInd = c(1:3))

# plot sequence of generating chains
post1 <- extract.samples(m1)
par(mfrow=c(1,4))
for(i in names(post1)){
  dfpost1 <- data.frame(post1[i])
  plot(as.numeric(rownames(dfpost1)), dfpost1[,1], type = "l",
       xlab = "Index",
       ylab = i)
}

par(mfrow=c (1,1))
plot(log(dt$ht_init),log(dt$rp_init), main="",xlab="log(height (cm))",
     ylab="log(Number of fruits)",pch=16,cex=0.5, xlim=c(1,5), ylim= c(0,9),
     col="darkgrey") 
lnhgt <- seq(1, log(max(dt$ht_init)), 0.01)
mu <- link(m1,data=data.frame(lgh=lnhgt))
mu.mean <- apply(mu,2,mean)
mu.PI <- apply(mu,2,PI,prob=0.95)
lines(lnhgt,mu.mean,col="blue")
shade(mu.PI,lnhgt)

##################################################

# b) each independent model#


par(mfrow=c (1,1))
plot(log(dt$ht_init),log(dt$rp_init), main="",xlab="log(height (cm))",
     ylab="log(Number of fruits)",pch=16,cex=0.5, xlim=c(1,5), ylim= c(0,9),
     col="darkgrey") 
lnhgt <- seq(1, log(max(dt$ht_init)), 0.01)
pops <- array(0,c(4,length(site)))
colnames(pops) <- site
rownames(pops) <- c("Beta1","Beta2","Sigma","N")

# Do for i from 1 to 14
#i=11
for(i in 1:14){
  dtp <- subset(dt,bald==site[i])
  pops[4,i] <-nrow(dtp)
  m <- ulam(
    alist(
      lfr ~ dnorm(mu,sigma),
      mu <- a + b*lgh,
      a ~ dnorm(0,100),
      b ~ dnorm(0,10),
      sigma ~ dunif(0,10)
    ),
    data = dtp,chains =2,cores=2, log_lik = TRUE,
  )
  
  z <- precis(m,digits=1)
  pops[1,i] <- z[1,1]
  pops[2,i] <- z[2,1]
  pops[3,i] <- z[3,1]
  #plot(m)
  #pairs(m)
  mu <- link(m,data=data.frame(lgh=lnhgt))
  mu.mean <- apply(mu,2,mean)
  mu.PI <- apply(mu,2,PI,prob=0.95)
  lines(lnhgt,mu.mean,col=i)
  shade(mu.PI,lnhgt)
}
round(pops,2)

##################################################

# c) each population intercept #

dt$yi <- 1
dt$yi[dt$year==1995] <- 2
dt$yi[dt$year==1996] <- 3
dt$pj <- 1
dt$pj[dt$bald==29] <-2
dt$pj[dt$bald==32] <-3
dt$pj[dt$bald==42] <-4
dt$pj[dt$bald==50] <-5
dt$pj[dt$bald==57] <-6
dt$pj[dt$bald==59] <-7
dt$pj[dt$bald==62] <-8
dt$pj[dt$bald==67] <-9
dt$pj[dt$bald==87] <-10
dt$pj[dt$bald==88] <-11
dt$pj[dt$bald==91] <-12
dt$pj[dt$bald==93] <-13
dt$pj[dt$bald==103] <-14

m3 <- ulam(
  alist(
    lfr ~ dnorm(mu,sigma),
    mu <- a  + p[pj] + b*lgh,
    a ~ dnorm(0,100),
    p[pj] ~ dnorm(0,sigmap),
    b ~ dnorm(0,10),
    sigmap ~ dcauchy(0,1),
    sigma ~ dcauchy(0,1)
  ),
  data = dt,chains =3, log_lik = TRUE,
)

precis(m3,digits=2,depth=2,pars=c("a","b","sigma"))
precis(m3,digits=2,depth=2,pars=c("a","b","sigma","p","sigmap"))
plot(m3)
plot(m3,depth=2)

post <- extract.samples(m3)
par(mfrow=c(1,6))
for(i in names(post)){
  dfpost <- data.frame(post[i])
  plot(as.numeric(rownames(dfpost)), dfpost[,1], type = "l",
       xlab = "Index",
       ylab = i)
}

post <- extract.samples(m3)
total_a <- sapply(1:14,function(beta1)post$a+post$p[,beta1])
round(apply(total_a,2,mean),2)

par(mfrow=c (1,1))
plot(log(dt$ht_init),log(dt$rp_init), main="",xlab="log(height (cm))",
     ylab="log(Number of fruits)",pch=16,cex=0.5, xlim=c(1,5), ylim= c(0,9),
     col="darkgrey") 
lnhgt <- seq(1, log(max(dt$ht_init)), 0.01)
mu <- NULL
for(j in 1:14){
  mu <- link(m3,data=data.frame(lgh=lnhgt,pj=j))
  
  
  mu.mean <- apply(mu,2,mean)
  mu.PI <- apply(mu,2,PI,prob=0.95)
  
  lines(lnhgt,mu.mean,col="blue")
}
v<- array(0,c(length(post$a),length(lnhgt)))
for (i in 1: length(post$a)){
  v[i,] <- post$a[i]+post$b[i]*lnhgt}
lines(lnhgt,colSums(v)*1/length(post$a),col="red",lwd=3)

##################################################

# d) each population slope and intercept #

m4 <- ulam(
  alist(
    lfr ~ dnorm(mu,sigma),
    mu <- a  + p[pj] + (b + bp[pj])*lgh,
    a ~ dnorm(0,100),
    p[pj] ~ dnorm(0,sigmap),
    bp[pj] ~ dnorm(0,sigmabp),
    b ~ dnorm(0,100),
    sigmap ~ dcauchy(0,10),
    sigmabp ~ dcauchy(0,1),
    sigma ~ dcauchy(0,1)
  ),
  data = dt,chains =3,
  iter=4000,warmup=1000, log_lik = TRUE,
)

precis(m4,digits=2,depth=2,pars=c("a","b","sigma"))
precis(m4,digits=2,depth=2,pars=c("p","sigmap"))
precis(m4,digits=2,depth=2,pars=c("bp","sigmabp"))
plot(m4)
plot(m4,depth=2)
#pairs(m4,pars=c("a","b","sigma"))
pairs(m4,horInd=c(1,30,33), verInd=c(1,30,33))
#pairs(m4,pars=c("p","sigmap"))
pairs(m4,horInd=c(2:15,31), verInd=c(2:15,31))
#pairs(m4,pars=c("bp","sigmabp"))
pairs(m4,horInd=c(16:29,32), verInd=c(16:29,32))

post <- extract.samples(m4)
par(mfrow=c(2,4))
for(i in names(post)){
  dfpost <- data.frame(post[i])
  plot(as.numeric(rownames(dfpost)), dfpost[,1], type = "l",
       xlab = "Index",
       ylab = i)
}

post <- extract.samples(m4)
total_a <- sapply(1:14,function(beta1)post$a+post$p[,beta1])
round(apply(total_a,2,mean),2)
total_b <- sapply(1:14,function(beta2)post$b+post$bp[,beta2])
round(apply(total_b,2,mean),2)


compare(m1,m3,m4)

par(mfrow=c (1,1))
plot(log(dt$ht_init),log(dt$rp_init), main="",xlab="log(height (cm))",
     ylab="log(Number of fruits)",pch=16,cex=0.5, xlim=c(1,5), ylim= c(0,9),
     col="darkgrey") 
lnhgt <- seq(1, log(max(dt$ht_init)), 0.01)
mu <- NULL
for(j in 1:14){
  mu <- link(m4,data=data.frame(lgh=lnhgt,pj=j))
  
  
  mu.mean <- apply(mu,2,mean)
  mu.PI <- apply(mu,2,PI,prob=0.95)
  
  lines(lnhgt,mu.mean,col="blue")
}

v<- array(0,c(length(post$a),length(lnhgt)))
for (i in 1: length(post$a)){
  v[i,] <- post$a[i]+post$b[i]*lnhgt}
lines(lnhgt,colSums(v)*1/length(post$a),col="red",lwd=3)


#### GRAPHING RANDOM DEVIATIONS
par(mfrow=c(1,2))
### SIMPLE BOXPLOTS

indevs<-extract.samples(m4)$p
boxplot(x = as.list(as.data.frame(indevs)),horizontal=T,outline=F,
        names = c("P1","P2","P3","P4","P5","P6","P7",
                  "P8","P9","P10","P11","P12","P13","P14"),
        main="Random deviations on Interecept")
abline(v=0,col="red2",lty=2)

sldevs<-extract.samples(m4)$bp
boxplot(x = as.list(as.data.frame(sldevs)),horizontal=T,outline=F,
        names = c("P1","P2","P3","P4","P5","P6","P7",
                  "P8","P9","P10","P11","P12","P13","P14"),
        main="Random deviations on Slope")
abline(v=0,col="red2",lty=2)

### VIOLIN PLOTS

library(vioplot)

vioplot(x = as.list(as.data.frame(indevs)),horizontal=T,rectCol="gray50",lineCol="gray50",
        names = c("P1","P2","P3","P4","P5","P6","P7",
                  "P8","P9","P10","P11","P12","P13","P14"),
        main="Random deviations on Interecept")
abline(v=0,col="red2",lty=2)

vioplot(x = as.list(as.data.frame(sldevs)),horizontal=T,rectCol="gray50",lineCol="gray50",
        names = c("P1","P2","P3","P4","P5","P6","P7",
                  "P8","P9","P10","P11","P12","P13","P14"),
        main="Random deviations on Slope")
abline(v=0,col="red2",lty=2)

### RANDOM STRUCTURE (CENTERED IN 0) WITH UNIT DEVIATIONS
par(mfrow=c(1,2))
intDevs<-precis(m4,digits=3,depth=2,pars=c("p","sigmap"))
d<-density(rnorm(200000,0,2.25))
dd <- approxfun(d$x,d$y)

colrs<-c("gray5","gray15","gray25","gray35","gray45","gray55","gray65",
         "gray5","gray15","gray25","gray35","gray45","gray55","gray65")
lntyp<-c(2,3,4,5,2,3,4,5,2,3,4,5,2,3)

plot(d,lwd=3,main="Random deviations on Intercept",xlab="")
for(i in 1:14){
  x<-intDevs[i,1]  
  y<-dd(x)
  segments(x,0,x,y,lwd=1.5,lty=lntyp[i],col=colrs[i])
}

slpDevs<-precis(m4,digits=3,depth=2,pars=c("bp","sigmabp"))
ds<-density(rnorm(200000,0,0.615))
dds <- approxfun(ds$x,ds$y)

plot(ds,lwd=3,main="Random deviations on slope",xlab="")
for(i in 1:14){
  x<-slpDevs[i,1]  
  y<-dds(x)
  segments(x,0,x,y,lwd=1.5,lty=lntyp[i],col=colrs[i])
}


### RANDOM STRUCTURE (CENTERED IN PAR VALUE) WITH UNIT DEVIATIONS
par(mfrow=c(1,2))
intDevs<-precis(m4,digits=3,depth=2,pars=c("p","sigmap"))
d<-density(rnorm(200000,-8.21,2.25))
dd <- approxfun(d$x,d$y)

colrs<-c("gray5","gray15","gray25","gray35","gray45","gray55","gray65",
         "gray5","gray15","gray25","gray35","gray45","gray55","gray65")
lntyp<-c(2,3,4,5,2,3,4,5,2,3,4,5,2,3)

plot(d,lwd=3,main="Random deviations on Intercept",xlab="")
for(i in 1:14){
  x<-intDevs[i,1]-8.21  
  y<-dd(x)
  segments(x,0,x,y,lwd=1.5,lty=lntyp[i],col=colrs[i])
}
segments(-8.21,0,-8.21,dd(-8.21),col="red",lwd=2)

slpDevs<-precis(m4,digits=3,depth=2,pars=c("bp","sigmabp"))
ds<-density(rnorm(200000,3.46,0.615))
dds <- approxfun(ds$x,ds$y)

plot(ds,lwd=3,main="Random deviations on slope",xlab="")
for(i in 1:14){
  x<-slpDevs[i,1]+3.46 
  y<-dds(x)
  segments(x,0,x,y,lwd=1.5,lty=lntyp[i],col=colrs[i])
}
segments(3.46,0,3.46,dds(3.46),col="red",lwd=2)