######################################################################
# STAT 7630 Bayesian Statistics 
# Peng Zeng @ Auburn University
# 09-09-2025
######################################################################

y = c(-67, -48, 6, 8, 14, 16, 23, 24, 28, 29, 41, 49, 67, 60, 75);
n = length(y);
size = 20000; 

######################################################################
# 1. normal approximation by Laplace method 
######################################################################

logpost = function(theta, y)
{
    mu = theta[1]; 
    sigma = exp(theta[2]); 
    sum(dt((y - mu)/sigma, df = 1, log = TRUE) - log(sigma)); 
}

fit = optim(c(mean(y), log(sd(y))), logpost, hessian = TRUE,
            y = y, control = list(fnscale = -1));
V = solve(-fit$hessian)

######################################################################
# 2. brute-force on grid 
######################################################################

mu.grid = seq(-10, 60, length = 201)
logs.grid = seq(1, 4.5, length = 201)
theta.grid = as.matrix(expand.grid(mu = mu.grid, logs = logs.grid))
n.grid = nrow(theta.grid);
den.grid = numeric(n.grid)
for(i in 1:n.grid) 
    den.grid[i] = exp(logpost(theta.grid[i, ], y)); 

pos = sample(1:n.grid, size, prob = den.grid, replace = TRUE);
thmat = theta.grid[pos, ];
method2 = thmat[10001:size, ]; 

par(mfrow = c(1, 2))
image(mu.grid, logs.grid, matrix(den.grid, nrow = 201))
image(MASS::kde2d(thmat[, 1], thmat[, 2]))

######################################################################
# 3. random walk 
######################################################################

theta = c(25, 3); 
thmat = matrix(0, nrow = size, ncol = 2); 
count = 0; 
for(i in 1:size)
{
    theta.new = theta + 2.5 * MASS::mvrnorm(1, c(0, 0), V);
    logratio = logpost(theta.new, y) - logpost(theta, y);
    if(log(runif(1)) < logratio)
    {
        count = count + 1; 
        theta = theta.new; 
    }

    thmat[i, ] = theta; 
}
method3 = thmat[10001:size, ]; 

par(mfrow = c(1, 2))
plot(thmat[, 1], type = "l")
plot(thmat[, 2], type = "l")

count / size; 

######################################################################
# 4. independence chain 
######################################################################

library(mvtnorm);

theta = c(25, 3); 
logd = logpost(theta, y) - dmvnorm(theta, fit$par, V, log = TRUE);
thmat = matrix(0, nrow = size, ncol = 2);
count = 0; 
for(i in 1:size)
{
    theta.new = MASS::mvrnorm(1, fit$par, V); 
    logd.new = logpost(theta.new, y) - dmvnorm(theta.new, fit$par, V, log = TRUE);
    if(log(runif(1)) < logd.new - logd)
    {
        count = count + 1; 
        theta = theta.new; 
        logd = logd.new; 
    }

    thmat[i, ] = theta; 
}
method4 = thmat[10001:size, ]; 

count / size; 

par(mfrow = c(1, 2))
plot(thmat[, 1], type = "l")
plot(thmat[, 2], type = "l")

######################################################################
# 5. Metroplis-within-Gibbs 
######################################################################

theta = c(25, 3); 
thmat = matrix(0, nrow = size, ncol = 2); 
count = c(0, 0); 
for(i in 1:size)
{
    theta.new = c(rnorm(1, theta[1], 18), theta[2]);
    logratio = logpost(theta.new, y) - logpost(theta, y);
    if(log(runif(1)) < logratio)
    {
        count[1] = count[1] + 1; 
        theta = theta.new; 
    }

    theta.new = c(theta[1], rnorm(1, theta[2], 1.2));
    logratio = logpost(theta.new, y) - logpost(theta, y);
    if(log(runif(1)) < logratio)
    {
        count[2] = count[2] + 1; 
        theta = theta.new; 
    }

    thmat[i, ] = theta; 
}
method5 = thmat[10001:size, ]; 

count / size; 

par(mfrow = c(1, 2))
plot(thmat[, 1], type = "l")
plot(thmat[, 2], type = "l")

######################################################################
# 6. Gibbs sampling with latent variable 
######################################################################

theta = c(25, exp(3));  # (mu, sigma)
lambda = rep(1, n)
thmat = matrix(0, nrow = size, ncol = 2); 
ldamat = matrix(0, nrow = size, ncol = n)
for(i in 1:size)
{
    theta[1] = rnorm(1, sum(lambda * y) / sum(lambda), theta[2] / sqrt(sum(lambda)))
    theta[2] = 1 / sqrt(rgamma(1, n/2, sum(lambda * (y - theta[1])^2) /2))
    lambda = rgamma(n, 1, 0.5 + ((y - theta[1]) / theta[2])^2 / 2)

    thmat[i, ] = theta; 
    ldamat[i, ] = lambda; 
}
method6 = cbind(thmat[10001:size, 1], log(thmat[10001:size, 2])); 

par(mfrow = c(1, 2))
plot(thmat[, 1], type = "l")
plot(thmat[, 2], type = "l")

######################################################################
# Metropolis-Hasting algorithm - draw parameters one-by-one
######################################################################

par(mfrow = c(1, 2), mar = c(4.5, 4.5, 0.5, 0.5))
xgrid = seq(-10, 60, length = 201)
xden = dnorm(xgrid, fit$par[1], sqrt(V[1, 1]))
plot(xgrid, xden, type = "l", col = "black", xlab = "mu", ylab = "density");
lines(density(method2[, 1], ), col = "red")
lines(density(method3[, 1], ), col = "blue")
lines(density(method4[, 1], ), col = "green")
lines(density(method5[, 1], ), col = "yellow")
lines(density(method6[, 1], ), col = "cyan")

xgrid = seq(1, 4.5, length = 201)
xden = dnorm(xgrid, fit$par[2], sqrt(V[2, 2]))
plot(xgrid, xden, type = "l", col = "black", xlab = "log(sigma)", ylab = "density");
lines(density(method2[, 2], ), col = "red")
lines(density(method3[, 2], ), col = "blue")
lines(density(method4[, 2], ), col = "green")
lines(density(method5[, 2], ), col = "yellow")
lines(density(method6[, 2], ), col = "cyan")


get.summary = function(a) 
{  c(mean(a), sd(a), quantile(a, prob = c(0.025, 0.5,0.975)));  }

fit$par 
sqrt(diag(V))
qnorm(c(0.025, 0.5, 0.975), fit$par[1], sqrt(V[1, 1]))
qnorm(c(0.025, 0.5, 0.975), fit$par[2], sqrt(V[2, 2]))

get.summary(method2[, 1])
get.summary(method3[, 1])
get.summary(method4[, 1])
get.summary(method5[, 1])
get.summary(method6[, 1])

get.summary(method2[, 2])
get.summary(method3[, 2])
get.summary(method4[, 2])
get.summary(method5[, 2])
get.summary(method6[, 2])

intervals = apply(ldamat[10001:size, ], 2, quantile, prob = c(0.025, 0.975))
plot(c(1, 15), range(intervals), type = "n", xlab = "index", ylab = "lambda")
for(i in 1:15)
{
    lines(c(i, i), intervals[, i], type = "l", lwd = 2)
}

######################################################################
# THE END
######################################################################