######
# Main code for running MCMC algorithm for capture-recapture data
# - white stork data with RJ - created by RK (Nov 08)
######

# Define the function: with input parameters:
# nt = number of iterations
# nburn  = burn-in

storkcodeRJ <- function(nt,nburn){

# Define the parameter values:
# ni = number of release years
# nj = number of recapture years
# nparam = maximum number of parameters
# ncov = maximum number of covariates

    ni = 16
    nj = 16
    nparam = 12
    ncov = 10

# Read in the data:

    data <- matrix(c(
19,2,0,0,0,0,0,0,0,0,0,0,0,0,0,0,5,
0,33,3,0,0,0,0,0,0,0,0,0,0,0,0,0,14,
0,0,35,4,0,0,0,0,0,0,0,0,0,0,0,0,14,
0,0,0,42,1,0,0,0,0,0,0,0,0,0,0,0,26,
0,0,0,0,42,1,0,0,0,0,0,0,0,0,0,0,30,
0,0,0,0,0,32,2,1,0,0,0,0,0,0,0,0,36,
0,0,0,0,0,0,46,2,0,0,0,0,0,0,0,0,16,
0,0,0,0,0,0,0,33,3,0,0,0,0,0,0,0,28,
0,0,0,0,0,0,0,0,44,2,0,0,0,0,0,0,20,
0,0,0,0,0,0,0,0,0,43,1,0,0,1,0,0,10,
0,0,0,0,0,0,0,0,0,0,34,1,0,0,0,0,25,
0,0,0,0,0,0,0,0,0,0,0,36,1,0,0,0,16,
0,0,0,0,0,0,0,0,0,0,0,0,27,2,0,0,22,
0,0,0,0,0,0,0,0,0,0,0,0,0,22,0,0,16,
0,0,0,0,0,0,0,0,0,0,0,0,0,0,15,1,17,
0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,15,8),nrow=ni,byrow=T)

# Read in the covariate values:

    cov <- matrix(c(
 0.73,-0.70, 0.61, 0.79, 1.38, 1.60, 0.28, 1.75, 0.17,-0.73,
 0.41, 0.82,-0.23, 2.04, 0.27, 1.50, 0.32, 0.97,-0.81, 0.28,
 0.84, 2.44, 0.03, 1.04, 0.50, 0.41, 0.90,-0.13, 1.89,-0.27,
-0.24, 0.86, 1.77,-0.15,-0.41, 0.79,-0.40, 0.93, 0.79, 0.74,
 0.72,-0.33, 0.41,-0.01, 0.42,-0.50,-0.45,-0.52,-0.92, 0.44,
-0.49,-0.31,-0.13,-0.48, 1.30,-0.77, 1.45,-0.02, 0.13, 0.05,
-0.55,-0.12,-1.10, 1.72, 0.66,-1.56,-0.06, 0.02, 0.07, 3.22,
-1.64, 0.27, 1.38,-0.83,-0.07,-1.21, 1.35,-0.91,-0.01,-0.47,
 1.12, 0.96, 0.02,-0.02,-0.16, 0.80,-0.98, 1.20,-0.13, 0.09,
-0.22, 0.19, 0.58, 0.14,-0.23, 1.08, 0.86, 0.61, 1.11,-0.70,
-0.40,-1.54, 1.22,-0.71, 0.09,-0.44,-1.22, 0.20, 0.01,-0.07,
 0.94,-1.26,-1.57, 0.50, 0.40,-1.43, 1.23, 0.29, 1.38,-1.02,
-2.45, 0.46,-1.82,-0.62,-2.76,-0.65, 0.03,-1.52,-0.72, 0.32,
 1.00,-0.91,-0.53,-0.88,-0.06, 0.09,-0.34,-0.62,-1.48,-0.70,
-0.32, 0.09,-0.16,-1.00, 0.21, 0.40,-1.05,-0.15, 0.33,-0.84,
 0.55,-0.93,-0.48,-1.52,-1.55,-0.11,-1.94,-2.10,-1.81,-0.36),
nrow=16,byrow=T)

# Set up the matrices to calculate posterior mean and sd

    mn <- array(0,nparam)
    std <- array(0,nparam)

    probmodel <- array(0,2**nparam)
    model <- array(0,dim=c(2**nparam,ncov))

# Set up matrices to calculate mean acceptance rates for RJ moves

    meanacc <- array(0,nparam)
    countacc <- array(0,nparam)


# Read in the priors:

# Beta prior on recapture probability

    alphap <- 1
    betap <- 1

# Normal priors (mean and variance) on regression coefficients

    mu <- c(0,0,0,0,0,0,0,0,0,0,0)
    sig2 <- c(10,10,10,10,10,10,10,10,10,10,10)

    sig <- sqrt(sig2)

# Specify independent prior probabilities on each regression coefficient being present (including beta_0):
# Note we specify with probability 1 that the intercept beta_0 is present.

    prob <- c(1,0.5,0.5,0.5,0.5,0.5,0.5,0.5,0.5,0.5,0.5)

# Parameters for MH updates (Uniform random walk)

    delta <- c(0.05,0.1,0.1,0.1,0.1,0.1,0.1,0.1,0.1,0.1,0.1,0.1)

# Parameters for RJ updates (Normal mean and variance) for beta_0 to beta_10

    rjmu <- c(0,0,0,0,0,0,0,0,0,0,0)
    rjsig2 <- c(0.5,0.5,0.5,0.5,0.5,0.5,0.5,0.5,0.5,0.5,0.5)

    rjsig <- sqrt(rjsig2)

# Set initial parameter values:

    param <- c(0.9,0.7,0,0,0,0.3,0,0,0,0,0,0)
    probparam <- array(0,nparam)
    sum1 <- array(0,nparam)
    sum2 <- array(0,nparam)

# param[1] = recapture rate
# param[2:12] = regression coefficients for survival rates

# sample is an array in which we put the sample from the posterior distribution.

    sample <- array(0, dim=c(nt, nparam))

# Calculate log(likelihood) for initial state using a separate function "calclikhood":

    likhood <- calclikhood(ni, nj, data, param, nparam, cov, ncov)

# Set up a vector for parameter values and associated log(likelihood):

    output <- dim(nparam+3)

# MCMC updates - MH algorithm:

# Cycle through each iteration:

    for (t in 1:nt){

# Update the parameters in the model using function "updateparam":

        output <- updateparam(nparam, param, ncov, cov, ni, nj, data, likhood, alphap, betap, mu, sig, delta)

# Set parameter values and log(likelihood) value of current state to be the output from
# the MH step:

        param <- output[1:nparam]
        likhood <- output[nparam+1]

# Update the covariates in the model using function "updatemodel":

        output <- updatemodel(nparam, param, ncov, cov, ni, nj, data, likhood, mu, sig, prob, rjmu, rjsig)

# Set parameter values and log(likelihood) value of current state to be the output from
# the RJ step:

        param <- output[1:nparam]
        likhood <- output[nparam+1]

# To calculate mean acceptance probabilities for RJ steps for adding/removing each covariate

        meanacc[output[nparam+2]] <- meanacc[output[nparam+2]] + output[nparam+3]
        countacc[output[nparam+2]] <- countacc[output[nparam+2]] + 1

# To calculate the model in:

        bin <- calcmodel(ncov,param)

# Record the set of parameter values:

        for (i in 1:nparam) {
            sample[t,i] <- param[i]
        }

# To calculate the posterior mean and standard deviation for outputting to screen:

        if (t >= (nburn+1)) {
            for (i in 1:nparam) {
                if (param[i] != 0) {
                    sum1[i] <- sum1[i] + param[i]
                    sum2[i] <- sum2[i] + param[i]^2
                    probparam[i] <- probparam[i] + 1
                }
            }

# To update the model probability and record the model:

            probmodel[bin] <- probmodel[bin] + 1
            if (probmodel[bin] == 1) {
                for (i in 1:ncov) {
                    if (param[i+2] != 0) {
                        model[bin,i] <- 1
                    }
                }
            }
        }
    }

for (i in 1:nparam) {
    mn[i] <- sum1[i]/probparam[i]
    std[i] <- sqrt((sum2[i] - probparam[i]*mn[i]^2)/(probparam[i]-1))
}

# Output the posterior mean and standard deviation of the parameters
# following burn-in to the screen:

cat("Posterior summary estimates for each parameter, conditional on being present:  ", "\n")
cat("\n")
cat("mean  (SD)", "\n")
cat("p: ", "\n")
cat(mn[1], "   (", std[1], ")", "\n")
for (i in 2:nparam) {
    cat("\n")
    cat("beta_", (i-2), "     Prob:", probparam[i]/(nt-nburn), "\n")
    if (probparam[i] != 0) {
        cat(mn[i], "   (", std[i], ")", "\n")
    }
    cat("Mean acceptance probability:", meanacc[i]/countacc[i], "\n")
}

cat("\n")
cat("Overall model            Posterior probability")
cat("\n")

for (i in 1:(2**ncov)) {
    if ((probmodel[i]/(nt-nburn)) >= 0.01) {
        cat(model[i,], "      :  ", probmodel[i]/(nt-nburn), "\n")
    }
}

# Output the sample from the posterior distribution:

sample

}

######
# This function calculates the log likelihood of capture-recapture data
######

calclikhood <- function(ni, nj, data, param, nparam, cov, ncov){

# First we set the survival and recapture probs:

# Set up the size of the array containing the survival and recapture probs:
# Set up the set of cell probabilities (q), and initially set them to be all equal to zero:

    phi <- array(0,nj)
    p <- array(0,nj)
    q <- array(0,dim=c(ni,nj+1))

# Set the recapture and survival probs for each year of the study:

    for (i in 1:nj) {

        exprn <- sum(param[2],param[3:nparam]*cov[i,1:ncov])
        phi[i] <- 1/(1+exp(-exprn))

        p[i] <- param[1]

    }

# Set initial value of the log(likelihood) to be zero

    likhood <- 0

# Calculate the Multinomial cell probabilities and corresponding contribution to the
# log(likelihood):
# Cycle through the number of years that there are releases:

    for (i in 1:ni){

# For diagonal elements:

        q[i,i] <- phi[i]*p[i]
        likhood <- likhood + data[i,i]*log(q[i,i])

# Calculate the elements above the diagonal:

        if (i <= (nj-1)) {

            for (j in (i+1):nj) {

                q[i,j] <- prod(phi[i:j])*prod(1-p[i:(j-1)])*p[j]
                likhood <- likhood + data[i,j]*log(q[i,j])

            }
        }

# Probability of an animal never being seen again

        q[i,nj+1] <- 1 - sum(q[i,i:nj])
        likhood <- likhood + data[i,nj+1]*log(q[i,nj+1])

    }

# Output the log(likelihood) value:

likhood

}

######
# Function for updating the parameters values:
######

updateparam <- function(nparam, param, ncov, cov, ni, nj, data, likhood, alphap, betap, mu, sig, delta){

# Cycle through each parameter in turn and propose to update using
# random walk MH with Uniform proposal density:

    for (i in 1:nparam) {

# Update the parameter if present in model (i.e. not equal to zero)

        if (param[i] != 0) {

# Keep a record of the current parameter value being updated

            oldparam <- param[i]

# Propose a new value for the parameter using a random walk with
# Uniform proposal density

            param[i] <- runif(1, param[i]-delta[i], param[i]+delta[i])

# Automatically reject any moves where recapture prob is outside [0,1]

            if (param[1] >= 0 & param[1] <= 1) {

# Calculate the log(acceptance probability):

# Calculate the new likelihood value for the proposed move:
# Calculate the numerator (num) and denominator (den) in turn:

                newlikhood <- calclikhood(ni, nj, data, param, nparam, cov, ncov)

                if (i == 1) {

# For recapture probability add in prior (Beta) terms to the acceptance probability

                    num <- newlikhood + log(dbeta(param[i],alphap,betap))
                    den <- likhood + log(dbeta(oldparam,alphap,betap))

                }

# For regression coefficients add in prior (Normal) terms to the acceptance probability

                else {

                    num <- newlikhood + log(dnorm(param[i],mu[i-1],sig[i-1]))
                    den <- likhood + log(dnorm(oldparam,mu[i-1],sig[i-1]))

                }


# All other prior terms (for other parameters) cancel in the acceptance probability.
# Proposal terms cancel since proposal distribution is symmetric.

# Acceptance probability of MH step:

                A <- min(1,exp(num-den))

            }

            else {

                A <- 0

            }

# To do the accept/reject step of the algorithm:
# Simulate a random number in [0,1]:

                u <- runif(1)

# Accept the move with probability A:

                if (u <= A) {

# Accept the proposed move:
# Update the log(likelihood) value:

                    likhood <- newlikhood

                }

                else {

# Reject proposed move so parameter stays at current value:

                    param[i] <- oldparam

                }
            }
        }

# Set the values to be outputted from the function to be the
# set of parameter values and log(likelihood) value:

        output <- c(param, likhood)

# Output the parameter values:

    output
}

######
# Function for updating the model:
######

updatemodel <- function(nparam, param, ncov, cov, ni, nj, data, likhood, mu, sig, prob, rjmu, rjsig){

# Propose a parameter to be added/removed:
# Note that we assume the intercept is always non-zero here
# This can be relaxed by setting sample(2:12,1)
# Note would also need to change prior probability for this intercept term.

    r <- sample(3:12,1)

# Keep a record of the current parameter value being updated

        oldparam <- param[r]

# Add/remove covariate r from the model
# If the covariate is not in the model:

        if (param[r] == 0) {

# Propose a new value for the parameter using Normal proposal

            param[r] <- rnorm(1, rjmu[r-1], rjsig[r-1])

        }

        else {

# If the covariate is in the model:
            param[r] <- 0
        }

# Calculate the log(acceptance probability):

# Calculate the new likelihood value for the proposed move:
# Calculate the numerator (num) and denominator (den) in turn:

        newlikhood <- calclikhood(ni, nj, data, param, nparam, cov, ncov)

# Include the likelihood term in the acceptance probability

        if (param[r] != 0) {

# Covariate is added - so prior is in the numerator; proposal in denominator:

            num <- newlikhood + log(dnorm(param[r],mu[r-1],sig[r-1]))
            den <- likhood + log(dnorm(param[r],rjmu[r-1],rjsig[r-1]))

# Priors for each model:

            num <- num + log(prob[r-1])
            den <- den + log(1-prob[r-1])

        }

        else {

# Covariate is removed - so prior is in the denominator; proposal in numerator:

            num <- newlikhood + log(dnorm(oldparam,rjmu[r-1],rjsig[r-1]))
            den <- likhood + log(dnorm(oldparam,mu[r-1],sig[r-1]))

# Priors for each model:

            num <- num + log(1-prob[r-1])
            den <- den + log(prob[r-1])

        }

# Acceptance probability of MH step:

        A <- min(1,exp(num-den))

# To do the accept/reject step of the algorithm:
# Simulate a random number in [0,1]:

        u <- runif(1)

# Accept the move with probability A:

        if (u <= A) {

# Accept the proposed move:
# Update the log(likelihood) value:

            likhood <- newlikhood

        }

        else {

# Reject proposed move so parameter stays at current value:

            param[r] <- oldparam

        }

# Set the values to be outputted from the function to be the
# set of parameter values and log(likelihood) value:

        output <- c(param, likhood, r, A)

# Output the parameter values:

    output
}

######
# Function for calculating the model in (uses form of binary conversion)
######

calcmodel <- function(ncov,param){

    bin <- 1

    for (i in 1:ncov) {
        if (param[i+2] != 0) {
            bin <- bin + 2^(i-1)
        }
    }

# Out put model number (as binary conversion)

bin
}