### ### ### Part I: Preparing and plotting the data ### ### ###
rm(list=ls())
getwd()

library(rethinking)
library(rstan)
library(bbmle)
library(nlme)

#setwd("pathway")
reg_data <- read.table("wg_data.txt", header=T)
par(mfrow=c(1,2))
hist(reg_data$lg,xlab="height",main="height")
hist(reg_data$lf,xlab="number of tillers",main="number of tillers")
par(mfrow=c(1,1))
plot(reg_data$lf,reg_data$lg,pch=16,cex=0.55,xlab="number of tillers",ylab="height")


max <- 17
reg_data$x<- reg_data$lf
reg_data$x2 <- reg_data$x^2
reg_data$y <- reg_data$lg
reg_data$x0 <- na.omit(reg_data$x)
ord <- order(reg_data$x0)
reg_data <-subset(reg_data,lf<max)

###################Three dimensional plot ##########
library(plot3D)
gc <- seq(0,80,2)
clf <- cut(reg_data$lg,gc)
dat <- table(clf,reg_data$lf)
persp(t(dat),xlab="number of leaves",ylab="height",zlab="frequency",col="lightgrey",
      theta = 40, phi = 30, border = NULL, box = TRUE)

### ### ### Part II: Initial Analysis  with diffuse priors ### ### ###

reg_data$x_c <- scale(reg_data$x)
reg_data$x2_c <- scale(reg_data$x2)
reg_data$tp_c <- scale(reg_data$tp)

m1 <- ulam(
  alist(
    y ~ dnorm(mu,sigma),
    mu <- a + b1*x_c +b2*x2_c + c*tp_c,
    a ~ dnorm(0,100),
    b1 ~ dnorm(0,10),
    b2 ~ dnorm(0,10),
    c ~ dnorm(0,10),
    sigma ~ dunif(0,10)
  ),
  data = reg_data,chains =3, log_lik = TRUE,
)

precis(m1,digits=2)
plot(m1)
pairs(m1,horInd=c(1:5),verInd=c(1:5))
summary(model1 <- lm(y ~ x_c + x2_c +tp_c,data=reg_data))

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


m2 <- ulam(
  alist(
    y ~ dnorm(mu,sigma),
    mu <- a + b1*x_c +b2*x2_c + c*tp_c + d1*x_c*tp_c + d2*x2_c*tp_c,
    a ~ dnorm(0,100),
    b1 ~ dnorm(0,10),
    b2 ~ dnorm(0,10),
    d1 ~ dnorm(0,10),
    d2 ~ dnorm(0,10),
    c ~ dnorm(0,10),
    sigma ~ dunif(0,10)
  ),
  data = reg_data,chains =3, log_lik = TRUE,
)

precis(m2,digits=2)
par(mfrow=c(1,1))
plot(m2)
pairs(m2,horInd=c(1:7),verInd=c(1:7))
summary(model2 <-lm(y ~ x_c + x2_c + tp_c + x_c:tp_c + x2_c:tp_c ,data=reg_data))

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


compare(m1,m2)
tp.p <- unique(reg_data$tp_c)
x.p <- reg_data$x_c[order(reg_data$x_c)]
x2.p <- reg_data$x2_c[order(reg_data$x2_c)]

par(mfrow=c(1,1))
plot(reg_data$x_c, reg_data$lg,type="n", main="plants per pot",xlab="scaled(number of tillers)", ylab="height")
points(reg_data$x_c[reg_data$tp==1], reg_data$lg[reg_data$tp==1],col="blue",pch=1,cex=0.55)
points(reg_data$x_c[reg_data$tp==2], reg_data$lg[reg_data$tp==2],col="red",pch=1,cex=0.55)
mu <- link(m1,data=data.frame(x_c=x.p,x2_c=x2.p,tp_c=tp.p[1]))
mu.mean <- apply(mu,2,mean)
mu.PI <- apply(mu,2,PI,prob=0.95)
lines(x.p,mu.mean,col="blue")
shade(mu.PI,x.p)
mu <- link(m1,data=data.frame(x_c=x.p,x2_c=x2.p,tp_c=tp.p[2]))
mu.mean <- apply(mu,2,mean)
mu.PI <- apply(mu,2,PI,prob=0.95)
lines(x.p,mu.mean,col="red")
shade(mu.PI,x.p)


### ### ### Part III: Bayesian analysis with informed priors ### ### ###


m1_inf <- ulam(
  alist(
    y ~ dnorm(mu,sigma),
    mu <- a + b1*x_c +b2*x2_c + c*tp_c,
    a ~ dnorm(22.98,8.85),
    b1 ~ dnorm(12.84,5.64),
    b2 ~ dnorm(-8.69,4.92),
    c ~ dnorm(0,10),
    sigma ~ dunif(0,10)
  ),
  data = reg_data,chains =3,
)

precis(m1_inf,digits=2)
plot(m1_inf)
pairs(m1_inf,horInd=c(1:5),verInd=c(1:5))

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

tp.p <- unique(reg_data$tp_c)
x.p <- reg_data$x_c[order(reg_data$x_c)]
x2.p <- reg_data$x2_c[order(reg_data$x2_c)]

par(mfrow=c(1,1))
plot(reg_data$x_c, reg_data$lg,type="n", main="plants per pot",xlab="scaled(number of tillers)", ylab="height")
points(reg_data$x_c[reg_data$tp==1], reg_data$lg[reg_data$tp==1],col="blue",pch=1,cex=0.55)
points(reg_data$x_c[reg_data$tp==2], reg_data$lg[reg_data$tp==2],col="red",pch=1,cex=0.55)
mu <- link(m1_inf,data=data.frame(x_c=x.p,x2_c=x2.p,tp_c=tp.p[1]))
mu.mean <- apply(mu,2,mean)
mu.PI <- apply(mu,2,PI,prob=0.95)
lines(x.p,mu.mean,col="blue")
shade(mu.PI,x.p)
mu <- link(m1_inf,data=data.frame(x_c=x.p,x2_c=x2.p,tp_c=tp.p[2]))
mu.mean <- apply(mu,2,mean)
mu.PI <- apply(mu,2,PI,prob=0.95)
lines(x.p,mu.mean,col="red")
shade(mu.PI,x.p)

post <- extract.samples((m1))
post1 <- extract.samples(m1_inf)


### ### ### Part V: Compare approaches ### ### ###
## NOTE: the code is currently set up to make the graph for beta2, but can be modified to make those of b0 and b1
library(ggplot2)
s_p <- 3000

###### a
prior_m_no <-rnorm(s_p, 0,sqrt(10))
prior_m_in <-rnorm(s_p,22.98,8.85)
post_m_no <- post$a
post_m_in <- post1$a

priors <-c(prior_m_no,prior_m_in,post_m_no,post_m_in)
length(priors)
type <- c(rep(1,s_p),rep(2,s_p),rep(3,s_p),rep(4,s_p))
levels(type) <- c("diffused","informed","posterior_no","posterior_in")
length(type)
p.graph <- cbind(type,priors)
p.graph <-as.data.frame(p.graph)
p.graph$type <- factor(p.graph$type)

qplot(p.graph$priors,geom="density",colour=p.graph$type,size=p.graph$type)+
  scale_color_manual(values = c("black","red","darkgreen","blue")) + scale_size_manual(values=c(0.7,0.7,0.7,0.7))

###### b1
s_p <- 3000
prior_m_no <-rnorm(s_p, 0,sqrt(10))
prior_m_in <-rnorm(s_p,12.84,5.64)
post_m_no <- post$b1
post_m_in <- post1$b1

priors <-c(prior_m_no,prior_m_in,post_m_no,post_m_in)
length(priors)
type <- c(rep(1,s_p),rep(2,s_p),rep(3,s_p),rep(4,s_p))
levels(type) <- c("diffused","informed","posterior_no","posterior_in")
length(type)
p.graph <- cbind(type,priors)
p.graph <-as.data.frame(p.graph)
p.graph$type <- factor(p.graph$type)

qplot(p.graph$priors,geom="density",colour=p.graph$type,size=p.graph$type)+
  scale_color_manual(values = c("black","red","darkgreen","blue")) + scale_size_manual(values=c(0.7,0.7,0.7,0.7))

###### b2

s_p <- 3000
prior_m_no <-rnorm(s_p, 0,sqrt(10))
prior_m_in <-rnorm(s_p,-8.69,4.92)
post_m_no <- post$b2
post_m_in <- post1$b2

priors <-c(prior_m_no,prior_m_in,post_m_no,post_m_in)
length(priors)
type <- c(rep(1,s_p),rep(2,s_p),rep(3,s_p),rep(4,s_p))
levels(type) <- c("diffused","informed","posterior_no","posterior_in")
length(type)
p.graph <- cbind(type,priors)
p.graph <-as.data.frame(p.graph)
p.graph$type <- factor(p.graph$type)

qplot(p.graph$priors,geom="density",colour=p.graph$type,size=p.graph$type)+
  scale_color_manual(values = c("black","red","darkgreen","blue")) + scale_size_manual(values=c(0.7,0.7,0.7,0.7))