# Example 1: Beta distributions graph
p = seq(0,1, length=100)
par(mfrow=c(2,2))

# plot several Beta distributions
y <- dbeta(p, 2, 10)
plot(p, dbeta(p, 2, 10), ylab='density', xlab = "theta", type ='l', col='purple', main = "Beta(2, 10)")
polygon(p,y, col = "purple")

y <- dbeta(p, 2, 2)
plot(p, dbeta(p, 2, 2), ylab='density', xlab = "theta", type ='l', col='red', main = "Beta(2, 2)")
polygon(p,y, col = "red")

y <- dbeta(p, 5, 2)
plot(p, dbeta(p, 5, 2), ylab='density', xlab = "theta", type ='l', col='blue', main = "Beta(5, 2)")
polygon(p,y, col = "blue")
lines(p, dbeta(p, 5, 2), col='blue')

y <- dbeta(p, 0.5, 0.5)
y[1] = 4
y[100] =4
plot(p, dbeta(p, 0.5, 0.5), ylab='density', xlab = "theta", type ='l', col='green', main = "Beta(0.5, 0.5)")
polygon(c(0,1,1,0), c(0,0,4,4), col = "green")
polygon(p, y, col = "white")
lines(p, dbeta(p, 0.5, 0.5), col='black')
par(mfrow=c(1, 1))


# Example 3: Heights 
y = c(66.77, 65.67, 61.85, 71.30, 62.36, 65.20, 65.63, 61.61, 66.18, 65.08)
sigma = 2.50

(y.bar = mean(y))
n = length(y)

prior.mean = 69
prior.st.dev = 2.5

(post.var = 1/(1/(prior.st.dev^2) + n/(sigma^2)))
(post.mean = post.var*(prior.mean/(prior.st.dev^2) + sum(y)/(sigma^2)))

grid = seq(62, 75, len = 1000)
plot(grid, dnorm(grid, prior.mean, prior.st.dev), type = "l", col = "blue", ylim = c(0, 0.7),
     xlab = "Height (in)", ylab = "Densities")

lines(grid, dnorm(grid, post.mean, sqrt(post.var)), col = "red")
legend(71, 0.5, legend=c("Prior", "Posterior", "Sample Mean"),
       col=c("blue", "red", "black"), lty=1, cex=0.8)
abline(v = y.bar)
# What to do when sigma is not known and we want a prior on it?

# Show example on https://stan-playground.flatironinstitute.org/
# Try student_t

# Final version with the rethinking package
library(rethinking) # make sure you have the latest version 


# alist contains the formulas that define the likelihood and priors.
heights.model <- ulam(
  alist(
    y ~ dnorm(alpha, 2.5) , # normal likelihood
    alpha ~ dnorm(69, 2.5)   # normal prior
  ) ,
  data=list(y = y) )

# display summary of quadratic approximation
precis(heights.model, prob=0.95)

stancode(heights.model)



# From Hoff's book:
# Example 4: Math scores in US public schools
# Data can be accessed here: https://www2.stat.duke.edu/~pdh10/FCBS/Inline/ 
# R code for the book here: https://www2.stat.duke.edu/~pdh10/FCBS/Replication/chapter8.R 

Y.school.mathscore<-dget("Y.school.mathscore")

#### Put data into list form. 
Y<-list()
YM<-NULL
J<-max(Y.school.mathscore[,1])
n<-ybar<-ymed<-s2<-rep(0,J)
for(j in 1:J) {
  Y[[j]]<-Y.school.mathscore[ Y.school.mathscore[,1]==j,2] 
  ybar[j]<-mean(Y[[j]])
  ymed[j]<-median(Y[[j]])
  n[j]<-length(Y[[j]])
  s2[j]<-var(Y[[j]])
  YM<-rbind( YM, cbind( rep(j,n[j]), Y[[j]] ))
}

## YM is like Y.school.mathscore in the book. 
colnames(YM)<-c("school","mathscore")


# Add Stan example:
write("
data {
  int n; // number of observations
  int n_cluster; // number of clusters
  vector[n] y;
  int<lower = 0, upper = n_cluster> school[n];  
}

parameters {
  vector[n_cluster] alpha; // vector of cluster intercepts
  real<lower = 0> tau; // between cluster st. dev.
  real<lower = 0> sigma; // model residual st. dev. (within)
  real mu_a; // mean of clusters
}

model {
  vector[n] mu; // conditional mean of the data
  mu = alpha[school];
  
  // priors
  mu_a ~ normal(50, 5);
  tau ~ inv_gamma(0.5, 50);
  sigma ~ inv_gamma(0.5, 50);
  
  // level-1 likelihood
  y ~ normal(mu, sigma);
  
  // level-2 likelihood
  alpha ~ normal(mu_a, tau);
}", "Example4.stan")

library(rstan)
mod <- stan_model("Example4.stan")
YM1 = as.data.frame(YM)
mod1_fit <- sampling(mod, cores = 4, data=list(n = nrow(YM1), n_cluster = length(unique(YM1$school)), y = YM1$mathscore, school= YM1$school))
summary(mod1_fit, pars = c("tau", "sigma"),probs = c(0.025, 0.975))$summary