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

model1 <- map2stan(
  alist(
    rep_structures ~ dnorm(mu,sigma),
    mu <- a + b*height,
    a ~ dnorm(0,500),
    b ~ dnorm(0,100),
    sigma ~ dunif(0,400)
  ),
  data = dt,chains =3,
  start <- list(a = 100,b= 10, sigma=sdfruits)
)

precis(model1,digits=1)
plot(model1)
pairs(model1)

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)

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

model2 <- map2stan(
  alist(
    lnrep ~ dnorm(mu,sigma),
    mu <- a + b*lnht,
    a ~ dnorm(0,1),
    b ~ dnorm(0,1),
    sigma ~ dunif(0,10)
  ),
  data = dt,chains =3,
  start <- list(a = 10,b= 1, sigma=sdlnfruits)
)

precis(model2,digits=1)
plot(model2)
pairs(model2)

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

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

model3 <- map2stan(
  alist(
    rep_structures ~ dnorm(mu,sigma),
    mu <- exp(a)*height^exp(b),
    a ~ dnorm(0,1),
    b ~ lognormal(0,1),
    sigma ~ dunif(0,400)
  ),
  data = dt,chains =3,
  start <- list(a = 1,b= 1, sigma=sdfruits)
)

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

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))
hist(dt$rep_structures[dt$rep_structures<2000],100, main ="", xlab= "fruits")

##################################################
dt$lnrep <- log(dt$rep_structures)
dt$lnht <-  log(dt$height)
sdlnfruits <- sd(dt$lnrep)
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

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

precis(model4,digits=1,depth=2)
plot(model4)

pairs(model4)

compare(model2,model4)

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 <- NULL
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)
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(exp(lnhgt),exp(mu.mean),col="red")
shade(exp(mu.PI),exp(lnhgt))
##################################################