##############################################
###   MULTIPLE CHANGE POINT TIME SERIES    ###
##############################################
### tested on R version 2.5.1 (2007-06-27) ###
##############################################

# Package by M. Plummer, N. Best, K. Cowles and K. Vines: output analysis and diagnostics for MCMC
library(coda)
# Package by J. Albert: functions for learning Bayesian inference
library(LearnBayes)
# note that package 'coda' requires package 'lattice'

##### Bernoulli data simulated as in Chib (1998) #####
dataname <- "Simulated Bernoulli data"
simuChib98 <- scan("simuChib98.txt")
n <- length(simuChib98)
############ prior hyper-parameters ############
m <- 2
a <- 8
b <- 0.1
c <- 2
d <- 2
################ MCMCparameters ################
mclen <- 6000
burnin <- 1000
thin <- 1
################################################

# Binary multiple change points
# MCMC algorithm based on Chib (1998)
binMulCP <- function(
                        y, # binary time series to be analyzed
                        mclen, # number of simulations desired
                        m, # number of change points in the model
                        a, b, # hyper-parameters for regime (non)transition probability
                        c, d # hyper-parameters for (regime specific) success probabilities
                    )
{
    # auxiliary normalizing function
    normalize <- function(x) x / sum(x)
    # number of observations
    n <- length(y)
    # initialization
    theta <- matrix(NA, m + 1, mclen)
    rownames(theta) <- paste("theta", 1:(m+1), sep = "")
    psi <- matrix(NA, m, mclen)
    s <- matrix(NA, n, mclen)
    pred <- matrix(NA, m + 1, n)
    filt <- matrix(NA, m + 1, n)
    # starting count of ones and zeros by hidden state
    oneCount <- rep(0.5 * n / m + 1, m + 1)
    zeroCount <- rep(0.5 * n / m + 1, m + 1)
    # Gibbs sampling
    for(i in 1:mclen)
    {
        # (regime specific) success probabilites
        for(j in 1:(m+1)) theta[j,i] <- rbeta(1, c + oneCount[j], d + zeroCount[j])
        # regime (non)transition probabilities
        for(j in 1:m) psi[j,i] <- rbeta(1, a + oneCount[j] + zeroCount[j] - 2, b + 1)
        # matrix of transition probabilities
        tramat <- diag(c(psi[,i], 1))
        for(j in 1:m) tramat[j, j+1] <- 1 - psi[j,i]
        # state predictive and filtering distributions
        pred[,1] <- c(1, rep(0, m))
        for(t in 1:(n-1))
        {
            filt[,t] <- normalize(pred[,t] * dbinom(y[t], 1, theta[,i]))
            pred[,t+1] <- diag(tramat) * filt[,t] + c(0, diag(tramat[1:m,-1]) * filt[1:m,t])
        }
        filt[,n] <- normalize(pred[,n] * dbinom(y[n], 1, theta[,i]))
        # vector of hidden states
        s[n,i] <- m + 1
        for(t in (n-1):1) s[t,i] <- sample(1:(m+1), 1, prob = filt[,t] * tramat[,s[t+1,i]])
        # count of zeros and ones by hidden state
        for(j in 1:(m+1))
        {
            oneCount[j] <- sum(y[s[,i]==j] == 1)
            zeroCount[j] <- sum(y[s[,i]==j] == 0)
        }
    }
    # change points
    tau <- matrix(NA, mclen, m)
    colnames(tau) <- paste("tau", 1:m, sep = "")
    for(j in 1:m) tau[,j] <- apply(s, 2, function(x) min(which(x == j + 1)))
    # output
    return(list(theta = t(theta), tau = tau))
}

# model fit
cat(system.time(fit <- binMulCP(simuChib98, mclen = mclen, m = m, a = a, b = b, c = c, d = d)), fill = TRUE)

# posterior state probabilities
states <- matrix(m + 1, mclen, n)
for(t in 1:n) for (j in m:1) states[t < fit$tau[,j], t] <- j
sprobs <- matrix(NA, m + 1, n)
for(j in 1:(m+1)) sprobs[j, ] <- apply(states[seq(burnin+1, mclen, thin),], 2, function(x) mean(x == j))
rm(t, j, states)

# graphical output to file
pdf("changePoints.pdf", height = 12, width = 12)
plot(c(1, n), c(0, 1), type = "n", xlab = "Time", ylab = expression(paste("Pr(",s[t]==k,"|",y,")")),
     main = "Posterior state probabilities")
for(j in 1:(m+1)) lines(1:n, sprobs[j,], lty = j)
rm(j)
legend("left", legend = paste("k =", 1:(m+1)), lty = 1:(m+1), bty = "n")
for(j in 1:m) plot(table(fit$tau[,j])/sum(table(fit$tau[,j])),
                   xlab = "Time", ylab = "Probability function", main = paste("Change point", j))
rm(j)
plot(mcmc(fit$theta[seq(burnin+1, mclen, thin),]))
autocorr.plot(mcmc(fit$theta[seq(burnin+1, mclen, thin),]))
dev.off()