### Part I: Preparing the data
rm(list=ls())
getwd()
library(rethinking)
library(rstan)
library(bbmle)
library(nlme)
library(lme4)     # for fitting GLMMs
library(lattice)
#setwd("pathway")

##Set your working directory manually or using the following commands## 
#getwd()
#setwd("pathway")
#setwd("C:/Users/Pedro/Documents/Classes/Methods in Ecology II/Demos/")

## read the Hypericum data from file

## read the Hypericum data from file
orig_data <- read.table("hypericum_data_94_07.txt", header=T)
## subset for plants with the reproductive structures and complete data
dt <- subset(orig_data, !is.na(ht_init) & !is.na(st_init) & rp_init > 0 & year<1997)
yr <- unique(dt$year) #studied years
dt$lgh <- log(dt$ht_init) #log of height
dt$lfr <- log(dt$rp_init) # log of reproductive structures
dt$stems <- dt$st_init # stems
site <- unique(dt$bald)[order(unique(dt$bald))]
# studied sites
#Survival
dt$surv <-1
dt$surv[dt$fate =="rip"] <- 0
table(dt$surv,dt$fate)
## Time-since-fire 
table(dt$bald,dt$fire_year) 
dt$TSF <- 1
dt$TSF[dt$fire_year <1987] <-2 
dt$TSF[dt$fire_year <1973] <-3 
table(dt$bald,dt$TSF)
tsf <- unique(dt$TSF)[order(unique(dt$TSF))]

dt$lghc <- dt$lgh- mean(dt$lgh) # centered for interpretation
dt$stems[dt$stems>8] <- 8

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

par(mfrow=c(1,2)) 
plot(dt$lgh ~dt$fire_year)
summary(lm(dt$lgh~factor(dt$TSF)))
plot(dt$lfr ~dt$fire_year)
summary(lm(dt$lfr ~factor(dt$TSF)))

par(mfrow=c (1,1)) 
pairs(subset(dt,select=c(lghc, lfr, stems)))

######################  RANDOM MODELS ##############################################
dt$fyear <- factor(dt$year)
dt$fbald <- factor(dt$bald)
dt$fstems <- factor(dt$stems)
dt$fTSF <- factor(dt$TSF)
dt$fate <- as.factor(dt$fate)
require(optimx)
m1 <- glm(surv~lghc*fTSF*fstems,data=dt,family =binomial)
m2 <- glmer(surv~lghc*fTSF*fstems + (1|fbald) + (1|fyear),data=dt,family =binomial)
m2_nlminb <- update(m2,control=glmerControl(optimizer="optimx",optCtrl=list(method="nlminb")))
m3 <- glmer(surv~lghc*fTSF*fstems + (lghc|fbald)+ (lghc|fyear),data=dt,family =binomial)
m3_nlminb <- update(m3,control=glmerControl(optimizer="optimx",optCtrl=list(method="nlminb")))

AICtab(m1,m2_nlminb,m3_nlminb,weights=TRUE,base = TRUE)

############################# FIXED MODELS ###########################################
M11 <- glmer(surv~lghc*fTSF*fstems + (1|fbald) + (1|fyear),data=dt,family=binomial)
M11_nlminb <- update(M11,control=glmerControl(optimizer="optimx",optCtrl=list(method="nlminb")))
M13 <- glmer(surv~lghc+fTSF*fstems + (1|fbald) + (1|fyear),data=dt,family =binomial)
M13_nlminb <- update(M13,control=glmerControl(optimizer="optimx",optCtrl=list(method="nlminb")))
M14 <- glmer(surv~lghc+fTSF+fstems + (1|fbald) + (1|fyear),data=dt,family =binomial)
M15 <- glmer(surv~lghc+fTSF + (1|fbald) + (1|fyear),data=dt,family =binomial)
M16 <- glmer(surv~lghc+fstems + (1|fbald) + (1|fyear),data=dt,family =binomial)
M17 <- glmer(surv~lghc*fstems + (1|fbald) + (1|fyear),data=dt,family =binomial)
M18 <- glmer(surv~lghc*fstems + TSF + (1|fbald) + (1|fyear),data=dt,family =binomial)
M19 <- glmer(surv~lghc*fTSF + fstems + (1|fbald) + (1|fyear),data=dt,family =binomial)
M20 <- glmer(surv~lghc*fTSF + fstems*fTSF + (1|fbald) + (1|fyear),data=dt,family =binomial)
M20_nlminb <- update(M20,control=glmerControl(optimizer="optimx",optCtrl=list(method="nlminb")))

AICtab(M11_nlminb,M13_nlminb,M14,M15,M16,M17,M18,M19,M20_nlminb,weights=TRUE,base = TRUE)

summary(M20_nlminb)


#######################################################################
color_eren <- c("yellow2","yellowgreen","forestgreen","turquoise3","dodgerblue3",
                "darkslateblue","blueviolet","red4")
par(mfrow=c (1,1)) 
plot(rep(1,8),col=color_eren,pch=19,cex=3)
min(dt$lghc)
max(dt$lghc)
x<- seq(-1.4, 0.85, 0.01)
par(mfrow=c (1,3),bg="lightgray")

for (i in 1:3){
  plot(exp(dt$lghc+mean(dt$lgh)),dt$surv,type="n",ylim=c(0,1),xlim=exp(c(-1.4, 0.85)+mean(dt$lgh)),main=paste("TSF",tsf[i]),ylab="P(survival)",xlab="height(cm)")
  rect(par("usr")[1], par("usr")[3], par("usr")[2], par("usr")[4], col = 
         "grey")
  for (j in 1:8){
    coeff <- fixef(M20_nlminb)
    if (i == 1) {
      coeff_tsf <- 0
      coeff_TSFxstems <- 0
      coeff_tsfxlgh <- 0
    }
    if (i == 2) {
      coeff_tsf <- coeff[3]
      if (j == 1) {coeff_TSFxstems <- 0}
      coeff_tsfxlgh <- coeff[12]
      if (j == 1) {coeff_TSFxstems <- 0}
      if (j == 2) {coeff_TSFxstems <- coeff[14]}
      if (j == 3) {coeff_TSFxstems <- coeff[16]}
      if (j == 4) {coeff_TSFxstems <- coeff[18]}
      if (j == 5) {coeff_TSFxstems <- coeff[20]}
      if (j == 6) {coeff_TSFxstems <- coeff[22]}
      if (j == 7) {coeff_TSFxstems <- coeff[24]}
      if (j == 8) {coeff_TSFxstems <- coeff[26]}
    }
    if (i == 3) {
      coeff_tsf <- coeff[4]
      if (j == 1) {coeff_TSFxstems <- 0}
      coeff_tsfxlgh <- coeff[13]
      if (j == 2) {coeff_TSFxstems <- coeff[15]}
      if (j == 3) {coeff_TSFxstems <- coeff[17]}
      if (j == 4) {coeff_TSFxstems <- coeff[19]}
      if (j == 5) {coeff_TSFxstems <- coeff[21]}
      if (j == 6) {coeff_TSFxstems <- coeff[23]}
      if (j == 7) {coeff_TSFxstems <- coeff[25]}
      if (j == 8) {coeff_TSFxstems <- coeff[27]}
    }
    if (j == 1) {coeff_stems <- 0}
    if (j == 2) {coeff_stems <- coeff[5]}
    if (j == 3) {coeff_stems <- coeff[6]}
    if (j == 4) {coeff_stems <- coeff[7]}
    if (j == 5) {coeff_stems <- coeff[8]}
    if (j == 6) {coeff_stems <- coeff[9]}
    if (j == 7) {coeff_stems <- coeff[10]}
    if (j == 8) {coeff_stems <- coeff[11]}
    odds_ratio <- coeff[1]+coeff[2]*x+coeff_tsf+coeff_stems+coeff_tsfxlgh*x+coeff_TSFxstems
    prob <- 1/(1 + (1/exp(odds_ratio)))
    lines(exp(x+mean(dt$lgh)), prob,col=color_eren[j],lwd=3)
  }}

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

#  Bayesian implementation of Model M19

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
dt$tsf2 <- 0
dt$tsf2[dt$fTSF==2] <- 1
dt$tsf3 <- 0
dt$tsf3[dt$fTSF==3] <- 1
dt$s2 <- dt$s3 <- dt$s4 <- dt$s5 <- dt$s6 <- dt$s7 <- dt$s8 <-0
dt$s2[dt$stems==2] <-1
dt$s3[dt$stems==3] <-1
dt$s4[dt$stems==4] <-1
dt$s5[dt$stems==5] <-1
dt$s6[dt$stems==6] <-1
dt$s7[dt$stems==7] <-1
dt$s8[dt$stems==8] <-1

dt$ht_final[is.na(dt$ht_final)] <- 0
dt$rp_final[is.na(dt$rp_final)] <- 0
# dt_sub <- subset(dt, !is.na(ht_final) & !is.na(rp_final)) 

m19B <- ulam(
  alist(
    surv ~ dbinom(1,p), 
    logit(p) <- a  + pop[pj] +  y[yi] + b*lghc + 
      d*tsf2 + e*tsf3 + f*lghc*tsf2 + g*lghc*tsf3 +
      c2*s2 + c3*s3 + c4*s4 + c5*s5 + c6*s6 + c7*s7 +
      c8*s8, 
    a ~ dnorm(0,100),
    b ~ dnorm(0,10),
    c2 ~ dnorm(0,10),
    c3 ~ dnorm(0,10),
    c4 ~ dnorm(0,10),
    c5 ~ dnorm(0,10),
    c6 ~ dnorm(0,10),
    c7 ~ dnorm(0,10),
    c8 ~ dnorm(0,10),
    d ~ dnorm(0,10),
    e ~ dnorm(0,10),
    f ~ dnorm(0,10),
    g ~ dnorm(0,10),
    pop[pj] ~ dnorm(0,sigmap),
    y[yi] ~ dnorm(0,sigmabp),
    sigmap ~ dcauchy(0,10),
    sigmabp ~ dcauchy(0,1)
  ),
  data = dt,chains =3,
  iter=4000,warmup=1000,
)

precis(m19B,digits=2)
precis(m19B,digits=2,depth=2)
par(mfrow=c(1,1))
plot(precis(m19B,digits=2,depth=2))
plot(m19B)
pairs(m19B,horInd=c(1:16),verInd=c(1:16))
pairs(m19B,horInd=c(17:32),verInd=c(17:32))

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

par(mfrow=c(1,1))
post <- extract.samples(m19B)
dens(post$sigmap,xlab="Sigma",xlim=c(-0.5,100),ylim=c(0,3))
dens(post$sigmabp,col=rangi2,add=TRUE,lwd=2)
text(5,0.6,"year",col=rangi2)
text(10,0.1,"population")

x <- seq(-1.35,0.84,length=100)
st <- 0:7
ts <- 0:2

par(mfrow=c(2,3),oma = c(0,0,2,0))

for (k in 1:14){
  for (j in 1:3){
    plot(exp(x+mean(dt$lgh)),seq(0,1,length=100),type="n",ylim=c(0,1),
         ylab= "P(survival)",xlab="height (cm)")
    for (i in 1:8){
      d.sur <- data.frame(lghc=x,
                          pj=rep(k,100),
                          yi=rep(2,100),
                          tsf2=rep((1==ts[j])*1,100),
                          tsf3=rep((2==ts[j])*1,100),
                          s2=rep((1==st[i])*1,100),
                          s3=rep((2==st[i])*1,100),
                          s4=rep((3==st[i])*1,100),
                          s5=rep((4==st[i])*1,100),
                          s6=rep((5==st[i])*1,100),
                          s7=rep((6==st[i])*1,100),
                          s8=rep((7==st[i])*1,100)
      )
      S.ensamble <- ensemble(m19B,data=d.sur)
      pred_S <- apply(S.ensamble$link,2,mean)
      pred_S.PI <- apply(S.ensamble$link,2,PI)
      lines(exp(x+mean(dt$lgh)),pred_S,type="l")
      shade(pred_S.PI,exp(x+mean(dt$lgh)))}}
  if ((k%%2) == 1) {
    mtext(paste("Populations",site[k],"and",site[k+1]), outer = TRUE, cex = 1.5)
  } else {}
}


pj.zeros <- matrix(0,1000,14)
yj.zeros <- matrix(0,1000,3)



par(mfrow=c(1,3))
for (j in 1:3){
  plot(exp(x+mean(dt$lgh)),seq(0,1,length=100),type="n",ylim=c(0,1))
  for (i in 1:8){
    d.sur <- data.frame(lghc=x,
                        pj= matrix(1,100),
                        yi= matrix(1,100),
                        tsf2=rep((1==ts[j])*1,100),
                        tsf3=rep((2==ts[j])*1,100),
                        s2=rep((1==st[i])*1,100),
                        s3=rep((2==st[i])*1,100),
                        s4=rep((3==st[i])*1,100),
                        s5=rep((4==st[i])*1,100),
                        s6=rep((5==st[i])*1,100),
                        s7=rep((6==st[i])*1,100),
                        s8=rep((7==st[i])*1,100)
    )
    S.ensamble <- ensemble(m19B,data=d.sur)
    link.m19B <- link(m19B,n=1000,data=d.sur,
                      replace=list(pj=pj.zeros,yj=yj.zeros))
    pred_S <- apply(S.ensamble$link,2,mean)
    pred_S.PI <- apply(S.ensamble$link,2,PI)
    lines(exp(x+mean(dt$lgh)),pred_S,type="l")
    lines(exp(x+mean(dt$lgh)),pred_S,type="l")
    shade(pred_S.PI,exp(x+mean(dt$lgh)))}}

precis(m19B)

fixef(M19)
summary(M19)