library(xtable)
library(rstan)
library(loo)


############################
## 1) Bayes factors
############################


## Toy example

pdf(file="Toy_ex1.pdf", width=9.5, height=6.5)
par(mar=c(5,6,2,1))
curve(dnorm(x, -1, 0.7), xlim=c(-4,4), 
      ylab ="density", xlab=expression(theta), 
      cex.lab=2, ylim=c(0, 0.8), lwd=3, lty=1, col="blue")
curve(dnorm(x, 1, 0.7), xlim=c(-4,4), add= TRUE, lwd=3, lty=1)
text(-1, 0.7, "H0", col="blue", cex=2)
text(1, 0.7, "H1", cex=2)
dev.off()

pdf(file="Toy_ex2.pdf", width=9.5, height=6.5)
par(mar=c(5,6,2,1))
curve(dnorm(x, -1, 0.7), xlim=c(-4,4), 
      ylab ="density", xlab=expression(theta), 
      cex.lab=2, ylim=c(0, 0.8), lwd=3, lty=1, col="blue")
curve(dnorm(x, 1, 0.7), xlim=c(-4,4), add= TRUE, lwd=3, lty=1)
text(-1, 0.7, "H0", col="blue", cex=2)
text(1, 0.7, "H1", cex=2)
polygon(c(0.624, seq(0.624,4,0.1),4),
        c(0, dnorm(seq(0.624,4,0.1), 1, 0.7) ,0), col="red",
        border =NA)
polygon(c(-4, seq(-4, 0.624, 0.1),0.624),
        c(0, c(dnorm(seq(-4, 0.624, 0.1), -1, 0.7)),0), 
          col="darkgreen",
        border =NA)
abline(v=0.624, col="darkgray", lwd=3, lty=2)
legend(2.3, 0.7, col=c("red", "darkgreen"), lwd=3, 
       c("reject", "accept"), cex=1.5)
dev.off()


pdf(file="Toy_ex3.pdf", width=9.5, height=6.5)
par(mar=c(5,6,2,1))
curve(dnorm(x, -1, 0.7), xlim=c(-4,4), 
      ylab ="density", xlab=expression(theta), 
      cex.lab=2, ylim=c(0, 0.8), lwd=3, lty=1, col="blue")
curve(dnorm(x, 1, 0.7), xlim=c(-4,4), add= TRUE, lwd=3, lty=1)
text(-1, 0.7, "H0", col="blue", cex=2)
text(1, 0.7, "H1", cex=2)
polygon(c(0.624, seq(0.624,4,0.1),4),
        c(0, dnorm(seq(0.624,4,0.1), 1, 0.7) ,0), col="red",
        border =NA)
polygon(c(-4, seq(-4, 0.624, 0.1),0.624),
        c(0, c(dnorm(seq(-4, 0.624, 0.1), -1, 0.7)),0), 
        col="darkgreen",
        border =NA)
abline(v=0.624, col="darkgray", lwd=3, lty=2)
segments(0.3, 0, 0.3, dnorm(0.3, -1, 0.7), col="orange", lwd=3)
points(0.3,dnorm(0.3, -1, 0.7), lwd=6, pch=16 , col="violet")
segments(0.3, dnorm(0.3, -1, 0.7), 0.3, dnorm(0.3, 1, 0.7), 
         col="orange", lwd=3)
points(0.3,dnorm(0.3, 1, 0.7), lwd=6,lty=3, pch=16 , col="violet")
legend(2.1, 0.7, col=c("red", "darkgreen", "orange"), lwd=3, 
       c("reject", "accept", "observed"), cex=1.5)
dev.off()


## Major league goals

library(LearnBayes)
data(soccergoals)

y <- soccergoals$goals

# write the likelihood function via the gamma distribution


lik_pois<- function(data, theta){
  n <- length(data)
  lambda <- exp(theta)
  dgamma(lambda, shape =sum(data)+1, scale=1/n)
}

prior_gamma <- function(par, theta){
  lambda <- exp(theta)
  dgamma(lambda, par[1], rate=par[2])*lambda  
}

prior_norm <- function(npar, theta){
  lambda=exp(theta)  
  (dnorm(theta, npar[1], npar[2]))
  
}

lik_pois_v <- Vectorize(lik_pois, "theta")
prior_gamma_v <- Vectorize(prior_gamma, "theta")
prior_norm_v <- Vectorize(prior_norm, "theta")


#likelihood

pdf(file= "soccer_priors.pdf", width=10, height =6.5)
par(mar=c(5,6,3,1))
curve(lik_pois_v(theta=x, data=y), xlim=c(-1,4), 
      xlab=expression(theta), ylab = "density", lwd =2,
      cex.lab=2)
#prior 1
curve(prior_gamma_v(theta=x, par=c(4.57, 1.43)), lty =2, col="red", add = TRUE, lwd =4)
#prior 2 
curve(prior_norm_v(theta=x, npar=c(1, .5)), lty =3, col="blue", add =TRUE, lwd=4)
#prior 3 
curve(prior_norm_v(theta=x, npar=c(2, .5)), lty =4, col="green", add =TRUE, lwd =4)
#prior 4 
curve(prior_norm_v(theta=x, npar=c(1, 2)), lty =5, col="violet", add =TRUE, lwd =4)
legend(2, 1.8, c("Lik.", "(1) Ga(4.57,1.43)", "(2) N(1, 0.25)", "(3) N(2,0.25)","(4) N(1, 4)" ),
       lty=c(1,2,3,4,5), col=c("black", "red", "blue", "green", "violet"),lwd=4, cex=1.8)
dev.off()

logpoissongamma <- function(theta, datapar){
  data <- datapar$data
  par <- datapar$par
  lambda <- exp(theta)
  log_lik <- log(lik_pois(data, theta))
  log_prior <- log(prior_gamma(par, theta))
  return(log_lik+log_prior)
}

logpoissongamma.v <- Vectorize( logpoissongamma, "theta")


logpoissonnormal <- function( theta, datapar){
  data <- datapar$data
  npar <- datapar$par
  lambda <- exp(theta)
  log_lik <- log(lik_pois(data, theta))
  log_prior <- log(prior_norm(npar, theta))
  return(log_lik+log_prior)
}  
logpoissonnormal.v <- Vectorize( logpoissonnormal, "theta")

#log-likelihood
curve(log(lik_pois(y, theta=x)), xlim=c(-1,4),ylim=c(-20,2), lty =1,
      ylab="log-posteriors", xlab=expression(theta))
#log posterior 1
curve(logpoissongamma.v(theta=x, list(data=y, par=c(4.57, 1.43))), col="red", xlim=c(-1,4),ylim=c(-20,2), lty =1, add =TRUE)
#log posterior 2
curve(logpoissonnormal.v( theta=x, datapar <- list(data=y, par=c(1, .5))), lty =1, col="blue",  add =TRUE)
#log posterior 3
curve(logpoissonnormal.v( theta=x, datapar <- list(data=y, par=c(2, .5))), lty =1, col="green", add =TRUE, lwd =2)
#log posterior 4
curve(logpoissonnormal.v( theta=x, list(data=y, par=c(1, 2))), lty =1, col="violet", add =TRUE, lwd =2)
legend(2.6, 1.3, c( "loglik", "lpost 1", "lpost 2", "lpost 3", "lpost 4" ),
       lty=1, col=c("black", "red", "blue", "green", "violet"),lwd=2, cex=0.9)


datapar <- list(data=y, par=c(4.57, 1.43))
fit1 <- laplace(logpoissongamma, .5, datapar)
datapar <- list(data=y, par=c(1, .5))
fit2 <- laplace(logpoissonnormal, .5, datapar)
datapar <- list(data=y, par=c(2, .5))
fit3 <- laplace(logpoissonnormal, .5, datapar)
datapar <- list(data=y, par=c(1, 2))
fit4 <- laplace(logpoissonnormal, .5, datapar)

postmode <- c(fit1$mode, fit2$mode, fit3$mode, fit4$mode )
postsds <- sqrt(c(fit1$var, fit2$var, fit3$var, fit4$var))
logmarg <- c(fit1$int, fit2$int, fit3$int, fit4$int)
cbind(postmode, postsds, logmarg)

BF_matrix <- matrix(1, 4,4)
for (i in 1:3){
  for (j in 2:4){
    BF_matrix[i,j]<- exp(logmarg[i]-logmarg[j])
    BF_matrix[j,i]=(1/BF_matrix[i,j]) 
  }
}

round_bf <- round(BF_matrix,3)
xtable(round_bf)

######################################
## 2) Predictive information criteria
######################################

## Eight schools

y <- c(28,8,-3,7,-1,1,18,12)
sigma <- c(15,10,16,11,9,11,10,18)
J <- 8

data <- list(y = y, sigma=sigma, J = J)
fit_1 <- stan("8schools.stan", 
            data = data, iter=200,
            cores = 4, chains =4)

fit_2 <- stan("8schools_invgamma.stan", 
              data = data, iter=200,
              cores = 4, chains =4)

fit_3 <- stan("8schools_halfcauchy.stan", 
              data = data, iter=200,
              cores = 4, chains =4)

data_fake <- list(y = y, sigma=sigma, J = J, tau =0.001)
fit_4 <- stan("8schools_fake.stan", 
              data = data_fake, iter=200,
              cores = 4, chains =4)



sims_1 <- extract(fit_1)
posterior_1 <- as.matrix(fit_1)

sims_2 <- extract(fit_2)
posterior_2 <- as.matrix(fit_2)

sims_3 <- extract(fit_3)
posterior_3 <- as.matrix(fit_3)

sims_4 <- extract(fit_4)
posterior_4 <- as.matrix(fit_4)

# Compute LOO
log_lik_1 <- extract_log_lik(fit_1)
loo_1 <- loo(log_lik_1)
print(loo_1)

log_lik_2 <- extract_log_lik(fit_2)
loo_2 <- loo(log_lik_2)
print(loo_2)

log_lik_3 <- extract_log_lik(fit_3)
loo_3 <- loo(log_lik_3)
print(loo_3)

log_lik_4 <- extract_log_lik(fit_4)
loo_4 <- loo(log_lik_4)
print(loo_4)

loo_diff <- compare(loo_1, loo_2, loo_3)

#  compute waic

waic_1 <- waic(log_lik_1)
waic_2 <- waic(log_lik_2)
waic_3 <- waic(log_lik_3)

waic_diff <- compare(waic_1, waic_2, waic_3)


## Hibbs model

y <- c(44.6, 57.76, 49.91, 61.34,49.6, 61.79, 48.95, 44.7, 59.17,
       53.94, 46.55, 54.74, 50.27, 51.24, 46.32)
X <- c(2.4, 2.89, 0.85, 4.21, 3.02, 3.62, 1.08,-0.39, 3.86,
       2.27, 0.38, 1.04, 2.36, 1.72, 0.1)

hibbs_data <- list(y=y, X=X, N=length(y))
fit_hibbs <- stan("hibbs.stan", data =hibbs_data,
                  iter =500)
log_lik_hibbs <- extract_log_lik(fit_hibbs)
loo_hibbs <- loo(log_lik_hibbs)
print(loo_hibbs)

waic(log_lik_hibbs)


## MLS soccer
library(LearnBayes)
data(soccergoals)

y <- soccergoals$goals
mls_data <- list(y=y, N=length(y))
mls_fit_1 <- stan('mls_gamma.stan', data =mls_data,
                  iter =500, cores = 4 )

mls_data <- list(y=y, N=length(y), mu=1, tau=0.5)
mls_fit_2 <- stan('mls_normal.stan', data =mls_data,
                  iter =500, cores = 4 )

mls_data <- list(y=y, N=length(y), mu=2, tau=0.5)
mls_fit_3 <- stan('mls_normal.stan', data =mls_data,
                  iter =500, cores = 4 )

mls_data <- list(y=y, N=length(y), mu=1, tau=2)
mls_fit_4 <- stan('mls_normal.stan', data =mls_data,
                  iter =500, cores = 4 )

# Compute LOO
log_lik_1 <- extract_log_lik(mls_fit_1)
loo_1 <- loo(log_lik_1)
print(loo_1)

log_lik_2 <- extract_log_lik(mls_fit_2)
loo_2 <- loo(log_lik_2)
print(loo_2)

log_lik_3 <- extract_log_lik(mls_fit_3)
loo_3 <- loo(log_lik_3)
print(loo_3)

log_lik_4 <- extract_log_lik(mls_fit_4)
loo_4 <- loo(log_lik_4)
print(loo_4)

loo_diff <- compare(loo_1, loo_2, loo_3, loo_4)