############################################################################
## Code by Vanja Dukic, U of Colorado at Boulder
## Bayesian Statistics and Computing 2019
##
## LECTURE 7 - Genetic Linkage Example (Rao 1973, Tanner and Wong 1987)
############################################################################

# Data
# 197 animals are distributed into 4 categories: Y = (y1, y2, y3, y4) according to the genetic linkage
# model, where the probabilities of falling into the 4 categories are given as:
# ((2 + θ)/4, (1 − θ)/4, (1 − θ)/4, θ/4)

y =c(125,18,20,34)

xgrid = seq(0.01,0.99,by=0.01)

# Posterior, under the flat (Beta(1,1)) prior on theta
MultNomPost = function(theta){
((2+theta)/4)^y[1]*((1-theta)/4)^(y[2]+y[3])*(theta/4)^y[4]}

# Negative Log-Posterior, under the flat Beta(1,1) prior on theta
MultNomNegLogPost = function(theta){
-log(2+theta)*y[1]-log(1-theta)*(y[2]+y[3])-log(theta)*y[4]}

# Plot the posterior and log-posterior
par(mfrow=c(1,2))
plot(xgrid,apply(t(t(xgrid)),1,MultNomPost), 'l', lwd = 2, col='blue', xlab = expression(theta), ylab = 'Posterior')
plot(xgrid,apply(t(t(xgrid)),1,MultNomNegLogPost), 'l', lwd = 2, col='blue', xlab = expression(theta), ylab = 'Negative Log Posterior')

OptTheta = optim(0.5, MultNomNegLogPost, method='L-BFGS-B',lower=0.01,upper=0.99,hessian=TRUE)

# Normal Approximation to the posterior
ThetaHat = OptTheta$par
VarThetaHat = (1/OptTheta$hessian)

# Plot the posterior and log-posterior with the Normal approximation superimposed
x11()
plot(xgrid,apply(t(t(xgrid)),1,MultNomPost), 'l', lwd = 2, col='blue', xlab = expression(theta), ylab = 'Posterior')
par(new=TRUE)
plot(xgrid,dnorm(xgrid,ThetaHat,sqrt(VarThetaHat)), 'l', lty=2, col='red',axes='') 
legend(x=0,y=1,legend=c('Posterior','Normal Approx'), lty=c(1,2), col=c('blue','red'))

x11()
plot(xgrid,apply(t(t(xgrid)),1,MultNomNegLogPost), 'l', lwd = 2, col='blue', xlab = expression(theta), ylab = 'Log-Posterior')
par(new=TRUE)
plot(xgrid,-dnorm(xgrid,ThetaHat,sqrt(VarThetaHat),log=TRUE), 'l', lty=2, col='red',axes='') 
legend(x=0.6,y=60,legend=c('Log-Posterior','Log-Normal Approx'),lty=c(1,2), col=c('blue','red'))