# Generic implementation of Metropolis Hastings algorithm.
mh.update <- function(state, log.target.density, proposal, log.proposal.density = function(x,y) { 0.0 } ) {
  log.mh.ratio <- 0
  log.mh.ratio <- log.mh.ratio + log.target.density(proposal) - log.target.density(state)
  log.mh.ratio <- log.mh.ratio + log.proposal.density(state,proposal) - log.proposal.density(proposal,state)
  if ( log(runif(1)) < log.mh.ratio ) return(proposal)
  else return(state)
}


# x: Gestational age of the infant (in weeks) at the time of birth
# y: Number of infants that were breast feeding at the time of release from hospital
# n: Total number of infants
# b0: Mean of the normal prior distribution for Beta0
# b1: Mean of the normal prior distribution for Beta1
# s0: Std. dev. of the normal prior distribution for Beta0
# s1: Std. dev. of the normal prior distribution for Beta1

x <- c(28,29,30,31,32,33)
y <- c( 2, 2, 7, 7,16,14)
n <- c( 6, 5, 9, 9,20,15)

make.density <- function(x,y,n,b0,s0,b1,s1) {
  function(state) {
    p <- 1 / ( 1 + exp(-(state[1]+state[2]*x)) )
    sum(dbinom(y,n,p,log=TRUE)) + dnorm(state[1],b0,s0,log=TRUE) + dnorm(state[2],b1,s1,log=TRUE)
  }
}

log.target.density <- make.density(x,y,n,-10,10,0.4,0.4)


diagnostic.plots <- function(parameter) {
  cat("Acceptance rate:",length(unique(parameter))/length(parameter),"\n")
  plot(parameter,type="l")
  locator(1)   # To induce a pause before going to the next figure
  acf(parameter)
  locator(1)
  hist(parameter,freq=FALSE)
  locator(1)
}


nDraws <- 10000

# First propsal scheme
draws <- matrix(NA,nrow=nDraws,ncol=2)
center <- c(-16,0.6)
draws[1,] <- center
for ( i in 2:nrow(draws) ) {
  if ( i %% 2 == 0 ) {
    proposal <- draws[i-1,] + c(runif(1,-1,1),0)                             # Propose update for Beta0
    draws[i,] <- mh.update(draws[i-1,], log.target.density, proposal)
  } else {
    proposal <- draws[i-1,] + c(0,0.05*runif(1,-1,1))                        # Propose update for Beta1
    draws[i,] <- mh.update(draws[i-1,], log.target.density, proposal)
  }
}

diagnostic.plots(draws[,1])
diagnostic.plots(draws[,2])
plot(NA,xlim=range(draws[,1]),ylim=range(draws[,2]),main=cor(draws[1],draws[2]))
for ( i in 1:nrow(draws) ) {
  points(draws[i,1],draws[i,2],pch=20)
}



# Let's try another proposal scheme
A <- 0.1*matrix(c(6,-0.18,-0.18,0.006),ncol=2)
B <- chol(A)
C <- solve(B)

log.proposal.density <- function(x,y) {
  -0.5 * t(x-y) %*% C %*% (x-y)
}

draws <- matrix(NA,nrow=nDraws,ncol=2)
draws[1,] <- center
for ( i in 2:nrow(draws) ) {
  proposal <- draws[i-1,] + B %*% rnorm(2)
  draws[i,] <- mh.update(draws[i-1,], log.target.density, proposal, log.proposal.density)
}

diagnostic.plots(draws[,1])
diagnostic.plots(draws[,2])
plot(NA,xlim=range(draws[,1]),ylim=range(draws[,2]),main=cor(draws[1],draws[2]))
for ( i in 1:(nrow(draws)/10) ) {
  points(draws[i,1],draws[i,2],pch=20)
}





# Let's try an independence sampler
A <- 2*matrix(c(6,-0.18,-0.18,0.006),ncol=2)
B <- chol(A)
C <- solve(B)

log.proposal.density <- function(x,y) {
  -0.5 * t(x-center) %*% C %*% (x-center)
}

draws <- matrix(NA,nrow=nDraws,ncol=2)
draws[1,] <- center
for ( i in 2:nrow(draws) ) {
  proposal <- center + B %*% rnorm(2)
  draws[i,] <- mh.update(draws[i-1,], log.target.density, proposal, log.proposal.density)
}

diagnostic.plots(draws[,1])
diagnostic.plots(draws[,2])
plot(NA,xlim=range(draws[,1]),ylim=range(draws[,2]),main=cor(draws[1],draws[2]))
for ( i in 1:nrow(draws)/10 ) {
  points(draws[i,1],draws[i,2],pch=20)
}