############################################################################################
#
# 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)
#
# 
#  Data: Real data application: Counts of coal mining disasters in 
#        Great Britain by year from 1851 to 1962.
#
#        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
#
############################################################################################
#
# HEDIBERT FREITAS LOPES
# Associate Professor of Econometrics and Statistics
# The University of Chicago Booth School of Business
# 5807 South Woodlawn Avenue
# Chicago, Illinois, 60637
# Email : hlopes@ChicagoGSB.edu
# URL: http://faculty.chicagobooth.edu/hedibert.lopes/research/
#
############################################################################################
# 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)
}

# Data
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="coalmining.pdf",width=10,height=7)
par(mfrow=c(1,1))
plot(1851:1962,y,type="h",xlab="years",ylab="Counts",lwd=4)
abline(v=1891,col=2,lwd=2)
abline(v=1886,col=4,lwd=2)
abline(v=1896,col=4,lwd=2)
dev.off()

mean(y[1:41])
mean(y[42:n])
# hyperparameters
# ---------------
alpha = 0.001
beta  = 0.001
delta = 0.001
gamma = 0.001

# True posterior distributions
# ----------------------------

# True 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)

# True 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)))

# Initial value for m
# -------------------
m = 41

# Gibbs Sampler
# -------------
set.seed(123456)
M      = 5000
draws  = NULL
system.time(
for (i in 1: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)
  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 = rbind(draws,c(lambda,phi,m))
}
)
# M=1000  64.59 7.33  72.81  
# M=5000 255.98 0.42 435.12

quant025=function(x){quantile(x,0.025)}
quant975=function(x){quantile(x,0.975)}
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
# -------------------
names  = c("lambda","phi","m")
par(mfrow=c(3,2))
for (i in 1:2){
  ts.plot(draws[,i],xlab="iteration",ylab=names[i])
  hist(draws[,i],xlab="",main=names[i],prob=T)
}
ts.plot(draws[,i],xlab="iteration",ylab=names[3])
plot(table(draws[,3]),type="h",xlab="",main=names[3])

# This is Figure 5.2 from Gamerman and Lopes (2007)
# -------------------------------------------------
par(mfrow=c(2,2))
plot(1851:1962,y,xlab="years",ylab="",main="(a)",axes=F);axis(1);axis(2)
plot(1851:1962,cumsum(y),type="l",xlab="years",ylab="",main="(b)",axes=F);axis(1);axis(2)
L = min(draws[,1:2])-0.01
U = max(draws[,1:2])+0.01
hist(c(draws[,1],draws[,2]),xlab="",ylab="",main="(c)",xlim=c(L,U),prob=T,breaks=seq(L,U,length=55),ylim=c(0,1.7))
text(1.75,1.5,expression(pi(phi)),cex=2)
text(3,1,expression(pi(lambda)),cex=2)
lines(lambdas,f[,1]/3.42*1.7,lwd=2)
lines(phis,f[,2]/3.42*1.7,lwd=2)
ind = sort(unique(draws[,3]))
probmhat = table(draws[,3])/M
plot(1850+ind,probm[ind],type="h",xlab="years",ylab="",main="(d)",axes=F);axis(1);axis(2)
for (i in 1:length(ind))
  segments(1850+ind[i]+0.25,0.0,1850+ind[i]+0.25,probmhat[i],lty=2)


# Contour plot of the joint posterior of lambda and phi
# -----------------------------------------------------
N = 50
lambdas = seq(min(draws[,1]),max(draws[,1]),length=N)
phis    = seq(min(draws[,2]),max(draws[,2]),length=N)
joint = matrix(0,N,N)
for (i in 1:M){
  a = ifelse(m>1,sum(y[1:m]),0)+alpha
  b = m+beta
  c = ifelse(m<n,sum(y[(m+1):n]),0)+gamma
  d = n-m+delta
  for (l1 in 1:N)
    for (l2 in 1:N)
      joint[l1,l2] = dgamma(lambdas[l1],a,b)*dgamma(phis[l2],c,d)
}
joint = joint/M
par(mfrow=c(1,1))
contour(lambdas,phis,joint,nlevels=5,drawlabels=FALSE,xlab=expression(lambda),ylab=expression(phi))

# Gibbs Sampler paths
# -------------------
x1 = seq(min(draws[,1]),max(draws[,1]),length=N)
x2 = seq(min(draws[,2]),max(draws[,2]),length=N)
xs = draws[,1:2]
xss = matrix(0,2*M,2)
xss[1,] = xs[1,1:2]
xss[2,2] = xs[1,2]
for (i in 2:M){
  xss[2*(i-1),1] = xs[i,1]
  xss[2*i-1  ,1] = xs[i,1]
  xss[2*i-1,2] = xs[i,2]
  xss[2*i  ,2] = xs[i,2]
}

#xs = scan("example51-chain-gibbs.txt")

# This is Figure 5.1 from Gamerman and Lopes (2007)
#--------------------------------------------------
par(mfrow=c(1,2))
contour(x1,x2,joint,xlim=c(min(xs[1:200,1]),max(xs[1:200,1])),ylim=c(min(xs[1:200,2]),max(xs[1:200,2])),drawlabels=F,xlab=expression(lambda),ylab=expression(phi),nlevels=5)
lines(xss[1:11,],col=1,lwd=2)
title("After 5 iterations")
text(xs[1,1],xs[1,2]-0.03,"(0)",col=1)
text(xs[2,1],xs[2,2]+0.03,"(1)",col=1)
text(xs[3,1]-0.04,xs[3,2],"(2)",col=1)
text(xs[4,1]+0.04,xs[4,2],"(3)",col=1)
text(xs[5,1]-0.04,xs[5,2],"(4)",col=1)
text(xs[6,1],xs[6,2]-0.03,"(5)",col=1)

contour(x1,x2,joint,xlim=c(min(xs[1:200,1]),max(xs[1:200,1])),ylim=c(min(xs[1:200,2]),max(xs[1:200,2])),drawlabels=F,xlab=expression(lambda),ylab=expression(phi),nlevels=5)
lines(xss[1:100,],col=1,lwd=2)
title("After 50 iterations")
contour(x1,x2,joint,xlim=c(min(xs[1:200,1]),max(xs[1:200,1])),ylim=c(min(xs[1:200,2]),max(xs[1:200,2])),drawlabels=F,xlab=expression(lambda),ylab=expression(phi),nlevels=5,add=T)

neff = NULL
for (i in 1:3)
  neff = c(neff,M/(1+2*sum(acf(draws[,i],lag=20,plot=F)$acf[2:20])))

pdf(file="poisson.pdf",width=10,height=7)
inds = seq(1,M,by=5)
par(mfrow=c(2,3))
plot(inds,draws[inds,1],xlab="",ylab="",main=expression(lambda),type="l",axes=F);axis(1);axis(2)
plot(inds,draws[inds,2],xlab="",ylab="",main=expression(phi),type="l",axes=F);axis(1);axis(2)
plot(inds,1850+draws[inds,3],xlab="",ylab="",main="m",type="l",axes=F);axis(1);axis(2)
hist(draws[,1],xlab="",ylab="",main="",prob=TRUE,nclass=20)
hist(draws[,2],xlab="",ylab="",main="",prob=TRUE,nclass=20)
hist(1850+draws[,3],xlab="",ylab="",main="",prob=TRUE,nclass=15)
dev.off()