# Edps 590BAY
# Fall 2021
# c.j.anderson
#
# Little Gibbs for normal:  mean and variance
#  The function is basically the same at Hoff (2009)
#
#  Using full conditionals or "semi-conjugate"
#
#  Example in lecture notes are after the function
#
seed <- 91921

gibbs <- function(S, y,mu.o, tau2.o, s2.o, nu.o, seed) {
 #--- set up
  set.seed(seed)
  n <- length(y)
  ybar <- mean(y)
  vary <- var(y)
  X <- matrix(NA, nrow=S, ncol=2)
  X[1,] <- c(ybar, 1/var(y))
  
 #--- the loop
   for (s in 2:S) {
   
   #--- get new value of mu
      mu.n <- (mu.o/tau2.o + n*ybar*X[(s-1),2]) / (1/tau2.o + n*X[(s-1),2])
      tau2.n <- 1/ (1/tau2.o + n*X[(s-1),2])
      X[s,1] <- rnorm(1, mu.n, sqrt(tau2.n))
	  
   #--- get new value of 1/sigma2 = precision
      nu.n <- nu.o + n
      s2.n <- (nu.o*s2.o + (n-1)*vary + n*(ybar-X[s,1])**2) /nu.n
      X[s,2] <- rgamma(1,nu.n/2, nu.n*s2.n/2)
   }
   return(X)
}


# Create some data:
mu = 4 
sigma = 2 
n = 50  
y <- rnorm(n,mu,sigma)
ybar <- mean(y)
vary <- var(y)


# Flat prior
mu.o <- 0
tau2.o <- 100
s2.o <- 1
nu.o <- 1
seed <- 2374
S <- 10000


sim1 <- gibbs(S, y, mu.o, tau2.o, s2.o, nu.o, seed)

sim1 <- as.mcmc(sim1)
par(mfrow=c(1,1))
plot(sim1)



par(mfrow=c(2,2))
# Joint trace mean and precision:
iter.n<-10
plot( sim1[1:iter.n,],type="l",xlim=range(sim1[1:S,1]), ylim=range(sim1[1:S,2]),
       lty=1,col="grey",xlab=expression(mu),ylab=expression(tilde(sigma)^2),
       main="First 10 iterations: precision x mu")
text(  sim1[1:iter.n,1], sim1[1:iter.n,2], c(1:iter.n) )

iter.n<-S
plot( sim1[1:iter.n,],type="l",xlim=range(sim1[1:S,1]), ylim=range(sim1[1:S,2]),
       lty=1,col="gray",xlab=expression(mu),ylab=expression(tilde(sigma)^2),
       main='All iterations: precision x mu')

# Joint trace mean and variance:
iter.n<-10
t <- 1/sim1[1:S,2]
plot( sim1[1:iter.n,1],t[1:iter.n], type="l",
       xlim=range(sim1[1:S,1]), ylim=range(t[1:S]),
       lty=1,col="grey",xlab=expression(mu),ylab=expression(sigma^2),
       main="First 10 iterations: variance x mu")
text(  sim1[1:iter.n,1], t[1:iter.n], c(1:iter.n) )

iter.n<-S
plot( sim1[1:iter.n,1],t[1:iter.n], type="l",
       xlim=range(sim1[1:S,1]), 
       ylim=range(t[1:S]),
       lty=1,col="gray",xlab=expression(mu),ylab=expression(sigma^2),
       main='All iterations: variance x mu')


par(mfrow=c(2,2))
hist(sim1[,1], main="Posterior Mean", xlab="mu.n")
hist(sim1[,2], main="Posterior Precision", xlab="1/sigma2.n")
hist(t, main="Posterior Variance", xlab="sigma2.n")
hist(t/50, main="Posterior of Variance of Mean", xlab="tau2.n = sigma2.n/n")



prior <- matrix(c(mu.o, tau2.o, s2.o, nu.o, S, ybar, vary, vary/n, n), nrow=9,ncol=1)
rownames(prior) <- c("mu.o", "tau2.o", "s2.o", "nu.o", "S", "ybar", "var(y)", "var(y)/n", "n")
round(prior, digits=4)

# posteior mean | y 
summary(test[,1])

# posterior precision | y
summary(test[,2])

# posterior variance | y
summary(1/ test[,2])


# Better yet, throw away the first half and use just the 2nd

save <- 1 + S/2
mu.post <- mean(sim1[save:S,1])
tau.post <- sd(sim1[save:S,1])
var.post <- 1/sim1[save:S,2]
var.post <- mean(var.post)
sd.post <- sqrt(var.post)
post.stat <- c(mu.post,tau.post,var.post, sd.post)
names(post.stat) <- c("Posterior: mean","tau","Sigma","sqrt(Sigma)")
post.stat


library(HDInterval)
# hdi
# -- mu
hdi(sim1[save:S,1])

# -- precision
hdi(sim1[save:S,2])

# -- variance
hdi(1/sim1[save:S,2])


# hdi
# -- mu
hdi(sim1[,1])

# -- precision
hdi(sim1[,2])

# -- variance
hdi(1/sim1[,2])