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

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

# a) no mixed effects #

sdlnfruits <- sd(dt$lfr)

m1 <- map2stan(
  alist(
    lfr ~ dnorm(mu,sigma),
    mu <- a + b*lgh,
    a ~ dnorm(0,100),
    b ~ dnorm(0,10),
    sigma ~ dunif(0,10)
  ),
  data = dt,chains =2,
  start <- list(a = 10,b= 1, sigma=sdlnfruits)
)

precis(m1,digits=2)
plot(m1)
pairs(m1)

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
dtp <- subset(dt,bald==site[i])
pops[4,i] <-nrow(dtp)
m <- map2stan(
  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,
  start <- list(a = 10,b= 1, sigma=sdlnfruits)
)
z <- precis(m,digits=1)
pops[1,i] <- z@output[1,1]
pops[2,i] <- z@output[2,1]
pops[3,i] <- z@output[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 <- map2stan(
  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
)

precis(m3,digits=2,depth=2)
plot(m3)
pairs(m3)
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 <- map2stan(
  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
)

precis(m4,digits=2,depth=2)
plot(m4)
pairs(m4)

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)