#############################################################################################################
#
# Example 6.5 (pages 212-3) – Figure 6.8 (page 214)
# 
#   Poisson model with change point  revisited.  
#   Metropolis step for change point parameter.
#
# Model : y_i ~ Poisson(lambda)   i=1,...,m 
#         y_i ~ Poisson(phi)      i=m+1,...,n
#
# Prior : lambda ~ Gamma(alpha,beta)
#         phi    ~ Gamma(gamma,delta)
#         m      ~ Uniform{1,...,n}
#
# Additional references: 
#
#      * Carlin, Gelfand and Smith (1992) 
#      * Tanner (1996, page 147)
#
##############################################################################################################
#
# Author : Hedibert Freitas Lopes
#          Graduate School of Business
#          University of Chicago
#          5807 South Woodlawn Avenue
#          Chicago, Illinois, 60637
#          Email : hlopes@ChicagoGSB.edu
#
###############################################################################################################

##################################################################
#
# Real data application: Counts of coal mining disasters in 
# Great Britain by year from 1851 to 1962.
#
##################################################################
logpost = function(lambda,phi,m){
  sm = sum(y[1:m])
  return((alpha+sm-1)*log(lambda)+(gamma+sn-sm-1)*log(phi)-(beta+m)*lambda-(delta+n-m)*phi)
}

y = c(4,5,4,1,0,4,3,4,0,6,3,3,4,0,2,6,3,3,5,4,5,3,1,4,4,1,5,5,3,4,
      2,5,2,2,3,4,2,1,3,2,2,1,1,1,1,3,0,0,1,0,1,1,0,0,3,1,0,3,2,2,
      0,1,1,1,0,1,0,1,0,0,0,2,1,0,0,0,1,1,0,2,3,3,1,1,2,1,1,1,1,2,
      4,2,0,0,0,1,4,0,0,0,1,0,0,0,0,0,1,0,0,1,0,1)
n = length(y)

# hyperparameters
# ---------------
alpha = 0.001
beta  = 0.001
delta = 0.001
gamma = 0.001
sn    = sum(y)

# METROPOLIS-HASTINGS ALGORITHM
# -----------------------------
set.seed(2079734)
M      = 5000
burnin = 0
taus   = seq(0.05,0.5,by=0.05)
neffs  = NULL
rates  = NULL
par(mfrow=c(5,2))
for (tau in taus){
  jump = matrix(0,n,n)
  for (i in 1:n){
    jump[i,] = exp(-tau*abs(i-1:n))
    jump[i,] = jump[i,]/sum(jump[i,])
  }
  ljump = log(jump)
  draws  = NULL
  m      = 41
  acceptance = 0
  system.time(
   for (i in 1:(burnin+M)){
     lambda = rgamma(1,ifelse(m>1,sum(y[1:m]),0)+alpha,m+beta)
     phi    = rgamma(1,ifelse(m<n,sum(y[(m+1):n]),0)+gamma,n-m+delta)
     m1     = sample(1:n,size=1,prob=jump[m,])
     A      = logpost(lambda,phi,m1)-logpost(lambda,phi,m)+ljump[m1,m]-ljump[m,m1]
     if (log(runif(1))<min(1,A)){
       m = m1
       acceptance = acceptance + 1
     }
     draws = rbind(draws,c(lambda,phi,m))
   }
  )
  draws = draws[(burnin+1):(burnin+M),]
  rates = c(rates,100*acceptance/(burnin+M))
  neff = NULL
  for (i in 1:3)
    neff = c(neff,M/(1+2*sum(acf(draws[,i],lag=20,plot=F)$acf[2:20])))
  neffs = rbind(neffs,neff)
  acf(draws[,3],xlab="lag",ylab="",axes=F);axis(1);axis(2)
}

par(mfrow=c(3,2))
plot(draws[,1],xlab="",ylab="",main="(a)",type="l",axes=F);axis(1);axis(2)
acf(draws[,1],xlab="lag",ylab="",main="(b)",axes=F);axis(1);axis(2)
plot(draws[,2],xlab="",ylab="",main="(c)",type="l",axes=F);axis(1);axis(2)
acf(draws[,2],xlab="lag",ylab="",main="(b)",axes=F);axis(1);axis(2)
plot(draws[,3],xlab="",ylab="",main="(e)",type="l",axes=F);axis(1);axis(2)
acf(draws[,3],xlab="lag",ylab="",main="(b)",axes=F);axis(1);axis(2)

# Figure 6.8
# ----------
gibbs = c(4800,3950,4900)
par(mfrow=c(2,2))
plot(taus,rates,xlab=expression(tau),ylab="acceptance rate",main="(a)",type="b",axes=F)
axis(1);axis(2)
plot(taus,neffs[,1],xlab=expression(tau),ylab="neff",main="(b)",type="b",ylim=c(0,M),axes=F)
axis(1);axis(2)
abline(h=gibbs[1])
plot(taus,neffs[,2],xlab=expression(tau),ylab="neff",main="(c)",type="b",ylim=c(0,M),axes=F)
axis(1);axis(2)
abline(h=gibbs[2])
plot(taus,neffs[,3],xlab=expression(tau),ylab="neff",main="(d)",type="b",ylim=c(0,M),axes=F)
axis(1);axis(2)
abline(h=gibbs[3])

# METROPOLIS-HASTINGS ALGORITHM (tau=0.15 and M=40000)
# ----------------------------------------------------
set.seed(2079734)
M      = 40000
burnin = 0
tau    = 0.15
neffs  = NULL
rates  = NULL
jump = matrix(0,n,n)
for (i in 1:n){
  jump[i,] = exp(-tau*abs(i-1:n))
  jump[i,] = jump[i,]/sum(jump[i,])
}
ljump = log(jump)
draws  = NULL
m      = 41
acceptance = 0
system.time(
   for (i in 1:(burnin+M)){
     lambda = rgamma(1,ifelse(m>1,sum(y[1:m]),0)+alpha,m+beta)
     phi    = rgamma(1,ifelse(m<n,sum(y[(m+1):n]),0)+gamma,n-m+delta)
     m1     = sample(1:n,size=1,prob=jump[m,])
     A      = logpost(lambda,phi,m1)-logpost(lambda,phi,m)+ljump[m1,m]-ljump[m,m1]
     if (log(runif(1))<min(1,A)){
       m = m1
       acceptance = acceptance + 1
     }
     draws = rbind(draws,c(lambda,phi,m))
   }
)
draws = draws[(burnin+1):(burnin+M),]
rate  = 100*acceptance/(burnin+M)
neff = NULL
for (i in 1:3)
  neff = c(neff,M/(1+2*sum(acf(draws[,i],lag=20,plot=F)$acf[2:20])))