######
# Main code for running MCMC algorithm for capture-recapture data
# - dippers when only considering model C/C2. Created August 07 (RK)
# Uses single-update Uniform random walk step
######

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

dippercodeMH1 <- function(nt,nburn){

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

    ni = 6
    nj = 6
    nparam = 3

# Read in the data:

data <- matrix(c(
11,2,0,0,0,0,9,
0,24,1,0,0,0,35,
0,0,34,2,0,0,42,
0,0,0,45,1,2,32,
0,0,0,0,51,0,37,
0,0,0,0,0,52,46),nrow=ni,byrow=T)

# Read in the priors:

# Beta priors on the parameters

    alpha <- c(1,1,1)
    beta <- c(1,1,1)

# Parameters for MH updates (Uniform random walk):

    delta <- c(0.1,0.05,0.05)

# Set initial parameter values:

    param <- c(0.5,0.5, 0.5)

# param[1] = recapture rate
# param[2] = survival rate - non-flood years
# param[3] = survival rate - flood years

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

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

    sample <- data.frame(sample)

# Label the columns of the array "sample"

    names(sample) <- c("p", "phin", "phif")

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

    likhood <- calclikhood(ni, nj, data, param)

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

    output <- dim(nparam+1)

# 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, ni, nj, data, likhood, alpha, beta, 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]

# Record the set of parameter values:

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

    }

# Calculate the mean and standard deviation of the parameters
# following burn-in:

subsample <- sample[(nburn+1):nt,]

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

for (i in 1:nparam) {
    mn[i] <- mean(subsample[,i])
    std[i] <- sqrt(var(subsample[,i]))
}

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

cat("Posterior summary estimates for each model:  ", "\n")
cat("\n")
cat("mean  (SD)", "\n")
cat("p: ", "\n")
cat(mn[1], "   (", std[1], ")", "\n")
cat("\n")
cat("phi_n: ", "\n")
cat(mn[2], "   (", std[2], ")", "\n")
cat("\n")
cat("phi_f: ", "\n")
cat(mn[3], "   (", std[3], ")", "\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){

# 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 survival and recapture probs for each year of the study:
# allowing for flood and non-flood years:

    for (i in 1:nj) {

        p[i] <- param[1]
        phi[i] <- param[2]
    }

# For flood years:

    for (i in 2:3) {
        phi[i] <- param[3]
    }

# 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, ni, nj, data, likhood, alpha, beta, 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

# 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 that are propose values outside [0,1]

        if (param[i] >= 0 & param[i] <= 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)

# Include the likelihood and prior (Beta) terms in the acceptance probability

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

# All other prior terms (for other parameters) cancel in the acceptance probability.

# 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
}