### Part I: Preparing the data
rm(list=ls())
getwd()
library(rethinking)
library(rstan)
library(bbmle)
#setwd("pathway")

# Read the Hypericum survival data from file
orig_data <- read.table("Hypericum_survival.txt", header=T)

# Change variables to include only reproductive individuals
dd  <- subset(orig_data, orig_data$rep_structures>0)
dd$rep <- log(dd$rep_structures)

### Part II: Logistic regression with a Bayesian approach 

model1.b.stan <- ulam(
  alist(
    survival ~ dbinom(1,p),
    logit(p) <- a + b*rep_structures,
    a ~ dnorm(0,10),
    b ~ dnorm(0,10)
  ),
  data = dd,chains=3, log_lik = TRUE,
)

precis(model1.b.stan,digits=3)

#### plots for the first model
post1 <- extract.samples(model1.b.stan)
par(mfrow=c(1,2))
plot(density(post1[[1]]),col="blue", xlim= c(-5,5),
     main="")
lines(density(rnorm(10000,0,10)),col="red")
plot(density(post1[[2]]),col="blue", xlim= c(-5,5),
     main="")
lines(density(rnorm(10000,0,10)),col="red")

par(mfrow=c(1,1))
plot(model1.b.stan)
pairs(model1.b.stan,horInd=c(1,2),verInd=c(1,2))
# pairs(model1.b.stan)

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

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

model2.b.stan <- ulam(
  alist(
    survival ~ dbinom(1,p),
    logit(p) <- a + b*rep,
    a ~ dnorm(0,10),
    b ~ dnorm(0,10)
  ),
  data = dd,chains=3, log_lik = TRUE,
)


precis(model2.b.stan)

#### plots for the second model
post2 <- extract.samples(model2.b.stan)
par(mfrow=c(1,2))
plot(density(post2[[1]]),col="blue", xlim= c(-5,5),
     main="")
lines(density(rnorm(10000,0,10)),col="red")
plot(density(post2[[2]]),col="blue", xlim= c(-5,5),
     main="")
lines(density(rnorm(10000,0,10)),col="red")

par(mfrow=c(1,1))
plot(model2.b.stan)
pairs(model2.b.stan,horInd=c(1,2),verInd=c(1,2))
# pairs(model1.b.stan)

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

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

compare(model1.b.stan,model2.b.stan)


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

##  plot comparing both models

par(mfrow=c(1,1))
plot(dd$rep_structures, dd$survival,type="n",ylab="P(survival)",xlab="# reproductive structures")
rug(jitter(dd$rep_structures[dd$survival==0]),ticksize = 0.03,)
rug(jitter(dd$rep_structures[dd$survival==1]),ticksize = 0.03, side=3)

b <- seq(0, max(dd$rep_structures),30)
z <- cut(dd$rep_structures, b)
prebyden <- tapply(dd$survival,z,sum)
tab <- table(z)
probs <-prebyden/tab
probs <- as.vector(probs)
points(b[2:length(b)],probs,pch=16,cex=1)
se2<-sqrt(probs*(1-probs)/tab)
up2 <-probs+as.vector(se2)
down2 <-probs-as.vector(se2)
for (i in 1:length(b-1)){
  lines(c(b[i+1],b[i+1]),c(up2[i],down2[i]))}

ln_fruits <- seq(0, max(dd$rep), 0.1)
mu <- link(model2.b.stan,data=data.frame(rep=ln_fruits))
mu.mean <- apply(mu,2,mean)
mu.PI <- apply(mu,2,PI,prob=0.95)
lines(exp(ln_fruits),mu.mean, ylim=c(0,1),col="red")
shade(mu.PI,exp(ln_fruits))

mu1 <- link(model1.b.stan,data=data.frame(rep_structures=exp(ln_fruits)))
mu1.mean <- apply(mu1,2,mean)
mu1.PI <- apply(mu1,2,PI,prob=0.95)
lines(exp(ln_fruits),mu1.mean, ylim=c(0,1),col="blue")
shade(mu1.PI,exp(ln_fruits))


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