###Part I: Preparing the data
rm(list=ls())
getwd()
library(rethinking)
library(rstan)
library(bbmle)
##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
orig_data <- read.table("hypericum_data_94_07.txt", header=T)

## change variables to include only reproductive individuals

dt <- subset(orig_data, !is.na(ht_init) & rp_init > 0 & year<1997 )
yr <- unique(dt$year)
dt$height <- dt$ht_init
dt$rep_structures <- dt$rp_init
dt$id <- dt$tag
sdfruits <- sd(dt$rep_structures)

m1 <- lm(rep_structures~height, data=dt)
summary(m1)


dts<-dt[,colSums(is.na(dt))==0]
dts<-dts[, !sapply(dts, is.character)]


model1 <- ulam(
  alist(
    rep_structures ~ dnorm(mu,sigma),
    mu <- a + b*height,
    a ~ dnorm(0,500),
    b ~ dnorm(0,100),
    sigma ~ dunif(0,400)
  ),
  data = dts,chains =3,cores=3, log_lik=TRUE
)

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

post <- extract.samples(model1)

par(mfrow=c(1,3))
plot(post$a, type = "l")
plot(post$b, type = "l")
plot(post$sigma, type = "l")

ht <- seq(0, max(dt$ht_init), 0.1)
mu <- link(model1,data=data.frame(height=ht))
mu.mean <- apply(mu,2,mean)
mu.PI <- apply(mu,2,PI,prob=0.95)

par(mfrow=c (1,1))
plot(dt$height,dt$rep_structures, main="",xlab="height (cm)",
     ylab="Number of fruits",pch=16,cex=0.55, xlim=c(0,80), ylim= c(0,1000) )
lines(ht,mu.mean,col="red")
shade(mu.PI,ht)

length(ht)

##################################################
dts$lnrep <- log(dts$rep_structures)
dts$lnht <-  log(dts$height)
sdlnfruits <- sd(dts$lnrep)

model2 <- ulam(
  alist(
    lnrep ~ dnorm(mu,sigma),
    mu <- a + b*lnht,
    a ~ dnorm(0,10),
    b ~ dnorm(0,10),
    sigma ~ dunif(0,10)
  ),
  data = dts,chains =3,cores=3, log_lik=TRUE
)

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

post2 <- extract.samples(model2)

par(mfrow=c(1,3))
plot(post2$a, type = "l")
plot(post2$b, type = "l")
plot(post2$sigma, type = "l")

par(mfrow=c (1,1))
plot(dt$height,dt$rep_structures, main="",xlab="height (cm)",
     ylab="Number of fruits",pch=16,cex=0.55, xlim=c(0,80), ylim= c(0,1000) ) 
lnhgt <- seq(0, log(max(dt$ht_init)), 0.01)
mu <- link(model2,data=data.frame(lnht=lnhgt))
mu.mean <- apply(mu,2,mean)
mu.PI <- apply(mu,2,PI,prob=0.95)
lines(exp(lnhgt),exp(mu.mean),col="red")
shade(exp(mu.PI),exp(lnhgt))  #SEE ABOVE

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

model3 <- ulam(
  alist(
    rep_structures ~ dnorm(mu,sigma),
    mu <- exp(a)*height^exp(b),
    a ~ dnorm(0,10),
    b ~ lognormal(0,1),
    sigma ~ dunif(0,400)
  ),
  data = dts,chains =3,cores=3, log_lik=TRUE
)

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


post3 <- extract.samples(model3)
par(mfrow=c(1,3))
plot(post3$a, type = "l")
plot(post3$b, type = "l")
plot(post3$sigma, type = "l")

ht <- seq(0, max(dt$ht_init), 0.1)
munl <- link(model3,data=data.frame(height=ht))
munl.mean <- apply(munl,2,mean)
munl.PI <- apply(munl,2,PI,prob=0.95)

par(mfrow=c (1,1))
plot(dt$height,dt$rep_structures, main="",xlab="height (cm)",
     ylab="Number of fruits",pch=16,cex=0.55, xlim=c(0,80), ylim= c(0,1000) ) 
lines(exp(lnhgt),exp(mu.mean),col="red")
shade(exp(mu.PI),exp(lnhgt))
lines(ht,munl.mean,col="blue")
shade(munl.PI,ht)

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

compare(model1,model3)

par(mfrow=c (1,1),mar=c(4,4,2,2))
hist(dt$rep_structures[dt$rep_structures<2000],
     100, main ="", xlab= "fruits",mgp=c(1.6,0.6,0))

# compare with residuals

# predicted values
dtsdata <- dts
mu1.dts <- link(model1, data=data.frame(height = dtsdata$height))
mu3.dts <- link(model3, data=data.frame(height = dtsdata$height))

# create new column for predicted values
dtsdata$predicted_rep1 <- apply(mu1.dts,2,mean)
dtsdata$predicted_rep3 <- apply(mu3.dts,2,mean)

# calculate residuals
dtsdata$residuals1 <- dtsdata$rep_structures - dtsdata$predicted_rep1
dtsdata$residuals3 <- dtsdata$rep_structures - dtsdata$predicted_rep3

# residuals plot
par(mfrow=c(1,2),mar=c(4,4,3,2),mgp=c(2.2,0.6,0))
plot(dtsdata$height, dtsdata$residuals1, 
     main = "Model 1 Residuals plot", xlab = "height",ylab="residuals")
abline(h=0)

plot(dtsdata$height, dtsdata$residuals3, 
     main = "Model 3 Residuals plot", xlab = "height",ylab="residuals")
abline(h=0)

# standardized residuals
dtsdata$stdresiduals1 <- dtsdata$residuals1/sd(dtsdata$residuals1)
dtsdata$stdresiduals3 <- dtsdata$residuals3/sd(dtsdata$residuals3)

# standardized residuals plot
par(mfrow=c(1,2),mar=c(4,4,3,2),mgp=c(2.2,0.6,0))
plot(dtsdata$predicted_rep1, dtsdata$stdresiduals1, 
     main = "Model 1 Residuals plot", xlab = "predicted reproductive structures",ylab="standardized residuals",
   #  ylim=c(-4,4)
   )
abline(h=0)

plot(dtsdata$predicted_rep3,
     dtsdata$stdresiduals3, 
     main = "Model 3 Residuals plot", 
     xlab = "predicted reproductive structures",
     ylab="standardized residuals",
     #  ylim = c(-4,4)
)
abline(h=0)

##################################################
dts$lnrep <- log(dts$rep_structures)
dts$lnht <-  log(dts$height)
sdlnfruits <- sd(dts$lnrep)
dts$yi <- 1
dts$yi[dts$year==1995] <- 2
dts$yi[dts$year==1996] <- 3
dts$pj <- 1
dts$pj[dts$bald==29] <-2
dts$pj[dts$bald==32] <-3
dts$pj[dts$bald==42] <-4
dts$pj[dts$bald==50] <-5
dts$pj[dts$bald==57] <-6
dts$pj[dts$bald==59] <-7
dts$pj[dts$bald==62] <-8
dts$pj[dts$bald==67] <-9
dts$pj[dts$bald==87] <-10
dts$pj[dts$bald==88] <-11
dts$pj[dts$bald==91] <-12
dts$pj[dts$bald==93] <-13
dts$pj[dts$bald==103] <-14

model4 <- ulam(
  alist(
    lnrep ~ dnorm(mu,sigma),
    mu <- a  + y[yi] + p[pj] + b*lnht,
    a ~ dnorm(0,10),
    y[yi] ~ dnorm(0,sigmay),
    p[pj] ~ dnorm(0,sigmap),
    b ~ dnorm(0,10),
    sigmay ~ dcauchy(0,1),
    sigmap ~ dcauchy(0,1),
    sigma ~ dcauchy(0,1)
  ),
  data = dts,
  iter=4000,warmup=1000,chains =3,cores=3, log_lik = TRUE
)

precis(model4,digits=1)
par(mfrow=c(1,1))
plot(model4,pars=c("a","b","sigma"))
pairs(model4,horInd=c(1,19,22),verInd=c(1,19,22))
# pairs(model4,pars=c("a","b","sigma"))

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

compare(model2,model4)

# compare using residuals

# predicted values
dtsdata <- dts
mu2.dts <- link(model2, data=data.frame(lnht = dtsdata$lnht))
mu4.dts <- link(model4, data=data.frame(yi = dtsdata$yi, pj = dtsdata$pj, lnht = dtsdata$lnht))

# create new column for predicted values
dtsdata$predicted_rep2 <- apply(mu2.dts,2,mean)
dtsdata$predicted_rep4 <- apply(mu4.dts,2,mean)

# calculate residuals
dtsdata$residuals2 <- log(dtsdata$rep_structures) - dtsdata$predicted_rep2
dtsdata$residuals4 <- log(dtsdata$rep_structures) - dtsdata$predicted_rep4

# residuals plot
par(mfrow=c(1,2),mar=c(4,4,3,2),mgp=c(2.2,0.6,0))
plot(dtsdata$lnht, dtsdata$residuals2, 
     main = "Model 2 Residuals plot", xlab = "height (ln)",ylab="residuals")
abline(h=0)

plot(dtsdata$lnht, dtsdata$residuals4, 
     main = "Model 4 Residuals plot", xlab = "height (ln)",ylab="residuals")
abline(h=0)

# standardized residuals
dtsdata$stdresiduals2 <- dtsdata$residuals2/sd(dtsdata$residuals2)
dtsdata$stdresiduals4 <- dtsdata$residuals4/sd(dtsdata$residuals4)

# standardized residuals plot
par(mfrow=c(1,2),mar=c(4,4,3,2),mgp=c(2.2,0.6,0))
plot(dtsdata$predicted_rep2, dtsdata$stdresiduals2, 
     main = "Model 2 Residuals plot", xlab = "predicted reproductive structures (ln)",ylab="standardized residuals")
abline(h=0)

plot(dtsdata$predicted_rep4, dtsdata$stdresiduals4, 
     main = "Model 4 Residuals plot", xlab = "predicted reproductive structures (ln)",ylab="standardized residuals")
abline(h=0)


#### PLOT GENERAL EFFECT:
par(mfrow=c (1,1))
plot(dt$height,dt$rep_structures, main="",xlab="height (cm)",
     ylab="Number of fruits",pch=16,cex=0.55, xlim=c(0,80), ylim= c(0,1000) ) 
lnhgt <- seq(0, log(max(dt$ht_init)), 0.01)

parIt<-extract.samples(model4)
mu<- rep(parIt$a,length(lnhgt))+outer(parIt$b,lnhgt, FUN = '*')

mu.mean <- apply(mu,3,mean)
mu.PI <- apply(mu,3,PI,prob=0.95)

par(mfrow=c (1,2),mar=c(3,3,2,2),mgp=c(1.7,0.7,0))
plot(dt$height,dt$rep_structures, main="Overall effect",xlab="height (cm)",
     ylab="Number of fruits",pch=16,cex=0.55, xlim=c(0,80), ylim= c(0,1000) ) 

lines(exp(lnhgt),exp(mu.mean),col="red")
shade(exp(mu.PI),exp(lnhgt))

plot(dt$height,dt$rep_structures, main="By site and year",xlab="height (cm)",
     ylab="Number of fruits",pch=16,cex=0.55, xlim=c(0,80), ylim= c(0,1000) ) 

for (i in 1:3){
  for(j in 1:14){
    mu <- link(model4,data=data.frame(lnht=lnhgt,yi=i,pj=j))
    mu.mean <- apply(mu,2,mean)
    lines(exp(lnhgt),exp(mu.mean),col=j+1,lty=i)
  }
}




############# GRAPHING CI SOLUTIONS ################

# # OPTION 1
# shades<-function (object, lim, label = NULL, col = rgb(0,0,0,0.15), border = NA, ...) 
# {
#   if (is.matrix(object) & length(dim(object)) == 2) {
#     y <- c(object[1, ], object[2, ][ncol(object):1])
#     x <- c(lim, lim[length(lim):1])
#     polygon(x, y, col = col, border = border, ...)
#   }
# }
# 
# # OPTION 2
# 
# polygon(c(ht,rev(ht)),c(mu.PI[1,],rev(mu.PI[2,])),
#         col=rgb(0,0,0,0.15), border = NA)
# 
# # OPTION 3
# 
# lines(ht,mu.PI[1,],col="red",lty=2)
# lines(ht,mu.PI[2,],col="red",lty=2)