######################################################################################################
#
# Gamerman and Lopes (2006)
# Example 5.1 (pages 143-5) – Figures 5.1 (page 144) and 5.2 (page 146) 
#
######################################################################################################
#
# Poisson data with change point and conditionally conjugate Gamma priors:
#
#  Model: yi |lambda,m  ~ Poi(lambda)   i=1,...,m
#         yi |phi,m         ~ Poi(phi)  i=m+1,...,n
#
#  Prior: m      ~ U{1,...,n}
#         lambda ~ Gamma(alpha,beta)
#         phi    ~ Gamma(gamma,delta)
#
#
#  Figure 5.1 shows typical trajectoriesof the Gibbs sampler for the pair (phi,lambda)
#
#  Figure 5.2 shows the data and its apparent change point plus posterior summaries 
#             from the Gibbs output
#
#  Dataset:   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
#
# Additional References:
#
#    * Carlin, Gelfand and Smith (1992) 
#    * Tanner (1996, page 147)
#
######################################################################################################
rm(list=ls())
quant025=function(x){quantile(x,0.025)}
quant975=function(x){quantile(x,0.975)}

# full conditional of m
full = function(m,lambda,phi,y,n,alpha,beta,gamma,delta){
  lambda^(alpha-1+ifelse(m>1,sum(y[1:m]),0))*exp(-(beta+m)*lambda)*phi^(gamma-1+ifelse(m<n,sum(y[(m+1):n]),0))*exp(-(delta+n-m)*phi)
}

##################################################################
#
# Real data application: Counts of coal mining disasters in 
# Great Britain by year from 1851 to 1962.
#
##################################################################
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)

# Time-series plot
pdf(file="coalminingdisasters-timeseries.pdf",width=8,height=6)
par(mfrow=c(1,1))
plot(1851:1962,y,type="h",xlab="years",ylab="",main="",pch=16)
title("Counts of coal mining disasters in Great Britain")
points(1891,0,pch=16,col=2)
text(1910,6,paste("Sample mean up to 1891 = ",round(mean(y[1:41]),2),sep=""))
text(1910,5.7,paste("Sample mean after 1891  = ",round(mean(y[42:n]),2),sep=""))
dev.off()

# hyperparameters
# ---------------
alpha = 0.001
beta  = 0.001
delta = 0.001
gamma = 0.001

##################################
# Exact solution
##################################

# Exact p(m|y)
# ------------
sn    = sum(y)
lprob = NULL
for (m in 1:n){
  sm = sum(y[1:m])
  lprob = c(lprob,lgamma(alpha+sm)+lgamma(gamma+sn-sm)-(alpha+sm)*log(m+beta)-(gamma+sn-sm)*log(n-m+delta))
}
lprob = lprob-max(lprob)
prob  = exp(lprob)
prob  = prob/sum(prob)
probm = prob
meanm = sum(1:n*probm)
varm  = sum((1:n-meanm)^2*probm)

# Exact p(lambda|y) and p(phi|y) and moments
# ------------------------------------------
sm      = cumsum(y)
sm1     = sn-sm
N       = 500
lambdas = seq(2,5,length=N)
phis    = seq(0.01,2,length=N)
f       = matrix(0,N,2)
E       = rep(0,2)
V       = rep(0,2)
for (m in 1:n){
  E[1]  = E[1]  + probm[m]*(sm[m]+alpha)/(m+beta)
  E[2]  = E[2]  + probm[m]*(sm1[m]+gamma)/(n-m+delta)
  V[1]  = V[1]  + probm[m]*(sm[m]+alpha)/((m+beta)^2)
  V[2]  = V[2]  + probm[m]*(sm1[m]+gamma)/((n-m+delta)^2)
}
for (i in 1:N){
  for (m in 1:n){
    f[i,1] = f[i,1] + dgamma(lambdas[i],sm[m]+alpha,m+beta)*probm[m]
    f[i,2] = f[i,2] + dgamma(phis[i],sm1[m]+gamma,n-m+delta)*probm[m]
  }
}
rbind(cbind(E,sqrt(V)),c(meanm,sqrt(varm)))


#########################################
# Gibbs Sampler
#########################################
set.seed(123456)

# MCMC set up
m = 41           # Initial value for m
M0     = 10000   # Burn-in
M      = 10000   # posterior draws
niter  = M+M0
draws  = matrix(0,niter,3)
time   = system.time(
for (iter in 1:niter){
  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)
  fulls  = NULL
  for (j in 1:n)
    fulls = c(fulls,full(j,lambda,phi,y,n,alpha,beta,gamma,delta))
  fulls = fulls/sum(fulls)
  m     = sample(1:n,size=1,prob=fulls)
  draws[iter,] = c(lambda,phi,m)
}
)
draws = draws[(M0+1):niter,]

summary = round(cbind(apply(draws,2,mean),sqrt(apply(draws,2,var)),
apply(draws,2,quant025),apply(draws,2,quant975)),3)
summary[3,c(1,3,4)]=1850+round(summary[3,c(1,3,4)])

# Posterior summaries
# -------------------
ind   = seq(1,M,by=M/1000)
pdf(file="coalminingdisasters-graph1.pdf",width=8,height=6)
par(mfrow=c(2,3))
plot(ind,draws[ind,1],xlab="iteration",ylab="",main=expression(lambda),type="l")
plot(ind,draws[ind,2],xlab="iteration",ylab="",main=expression(phi),type="l")
plot(ind,1850+draws[ind,3],xlab="iteration",ylab="",main="m",type="l")
acf(draws[,1],main="")
acf(draws[,2],main="")
acf(draws[,3],main="")
dev.off()

pdf(file="coalminingdisasters-graph2.pdf",width=10,height=6)
par(mfrow=c(1,2))
hist(draws[,1],xlab="",main=expression(lambda),prob=TRUE);box()
hist(draws[,2],xlab="",main=expression(phi),prob=TRUE);box()
dev.off()


pdf(file="coalminingdisasters-graph3.pdf",width=10,height=6)
par(mfrow=c(1,2))
ns = rep(0,n)
for (i in 1:n)
  ns[i] = mean(draws[,3]==i)
plot(1850+(1:n),ns,type="h",xlab="year",ylab="Probability",main="m")
abline(h=1/n,lty=2)
plot(table(1850+draws[,3])/M,type="h",xlab="year",main="m",ylab="Probability")
abline(h=1/n,lty=2)
dev.off()