# Code to estimate turtle population size
# using method of Grego and Hitchcock (2013)

# Need to install 'pscl' package to run this code.

###
### Various 3-day survey data sets:
### 

#y=c(36,38,65,17,39,44) # June 1-3, 2000
y=c(77,76,71,58,87,72) # June 16-18, 2000
#y=c(48,51,48,22,65,58) # July 1-3, 2000
#y=c(45,47,34,44,54,16) # July 16-18, 2000
#y=c(51,45,41,58,47,56)  # July 9-11, 2002

## Explanation of what the elements in y represent:

# y[1] = Nest, day 1 = NOO
# y[2] = Nest, day 2 = ONO+FNO
# y[3] = Nest, day 3 = OON+OFN+FFN
# y[4] = False Crawls, day 1 = FNO+FFN+FFF
# y[5] = False Crawls, day 2 = OFF+OFN+FFN+FFF
# y[6] = False Crawls, day 3 = FFF+OFF+OOF


N.unique.observed <- sum(y[c(1,2,3,6)])  # The number of uniquely observed turtles in the 3 days (N_obs from the paper)


### The following functions should be copied into R initially.
### They will allow the calculation of omega for a given mu and sigma.

###################################
###################################
## For calculation of omega:

#joint density of x and u
joint=function(u,x,m,s){
mu=m
h=(1/mu)*(1/(s*sqrt(2*pi)))*exp(-.5*((x-u)/s)^2)*(1/(1-pnorm(-u/s)))
return(h)
}

#waiting density of x
wait=function(arg.vec,m,s){
marg.x=rep(0,length(arg.vec))
for(i in 1:length(arg.vec)){
marg.x[i]=integrate(joint,x=arg.vec[i],m=m,s=s,lower=0,upper=m)$value
}
return(marg.x)
}

pnotdd=function(parm){
m=parm[2] 
mplus=parm[2] +((1-parm[1])/parm[1]) + 1  # = mu + q/p + 1
s=parm[3]
x=0

days=max(1,(floor(m-5*s))):ceiling((m+5*s))
         #print(paste("mu=",m))
         #print(paste("s=",s))
pn1=pnorm((days-m-.5)/s)
         #print(paste("pn1=",pn1))
pn2=pnorm((days-m+.5)/s)
         #print(paste("pn2=",pn2))
dn=wait(days-.5,m,s)
         #print(paste("dn=",dn))
         #print(paste("pn2-pn1", pn2-pn1))
         #print(dn*(pn2-pn1)/(1-pn1) )
x= 1 - sum(dn*(pn2-pn1)/(1-pn1))

#plot(days,dn*(pn2-pn1)/(1-pn1))

my.fcn<-function(days,m,mplus,s){
pn1=pnorm((days-m)/s)
        
pn2=pnorm((days-m+.05)/s)
#dn=wait(days,mplus,s)  
#dn=wait(days-.5,m,s)
out<- ( dn*(pn2-pn1)/(1-pn1) )
return(out)
}

return(x)
}
############
####################################
####################################


##### Prior parameters:

## The user should specify reasonable values for the hyperparameters
## of the priors on N.OOO, mu, p, and sigma:

## Prior on N.OOO, the number of turtles offshore all 3 days:

# These were used with the June 16-18, 2000 data:
# (These are negative binomial prior parameters.  
# They imply a prior belief that N.OOO is between 645 and 1031)
pri.mean.N.OOO <- 794  #used for prior on N.OOO
pri.size.N.OOO <- 80     #used for prior on N.OOO

# These were used with the 2002 data:
# (These are negative binomial prior parameters.  
# They imply a prior belief that N.OOO is between roughly 550 and 900)
#pri.mean.N.OOO <- 708  #used for prior on N.OOO
#pri.size.N.OOO <- 100     #used for prior on N.OOO

## Prior on the mean interarrival time mu:

# (These are normal prior parameters.  
# They imply a prior belief that mu is between 11.7 and 14.3 days,
# with a prior expected value of 13.)
pri.mean.mu <- 13
pri.sd.mu <- .65    

## Prior on the nesting probability p:

# (These are beta prior parameters.  
# They imply a prior belief that p is between 0.35 and 0.85,
# with a prior expected value of 0.6.)
pri.alpha.p <- 6
pri.beta.p <- 4      # Prior expectation for p is 0.6

## Prior on the standard deviation sigma of the distribution of interarrival times:

# (These are inverse gamma prior parameters.  
# They imply a prior expected value of 0.75 for sigma.)
pri.alpha.sig <- 9
pri.beta.sig <- 6    # Prior expectation for sigma is 6/(9-1) = 0.75

library(pscl); # this package needed for the rigamma function that generates random inverse gamma variates.

# We use a burn-in of 500 iterations here,
# followed by 10000 iterations of the MCMC.
J=10500
burn.in.amt <- 500

my.p <- my.mu <- my.sig <- N.OOO <- omega <- rep(0,times=J); 

my.m <- matrix(0,nrow=J,ncol=3)

my.mu.dens.save <- NULL

# Initial values, for step 1 of the algorithm:
# (We chose these roughly equal to the prior expected values, although they don't have to be.)

my.p[1] <- 0.6

N.OOO[1] <- 800
my.mu[1] <- 13
my.sig[1] <- 0.75

# Calculating initial omega:

omega[1] <- pnotdd(c(my.p[1],my.mu[1],my.sig[1]))


######################################
######################################
#####
##### The MCMC steps follow
#####
######################################
######################################

for (j in 2:J){

#if (j/100 == round(j/100)) print(j)  #Uncomment this line if you want to track iteration number while MCMC is running

# Defining the probabilities for the 10-category multinomial
# (interior cells)

pm <- rep(0,times=10)

 pm[1] <- my.mu[j-1]*(omega[j-1]^2)/(my.mu[j-1] + 1 + (1-my.p[j-1])/my.p[j-1])  		 #OOO
 pm[2] <- my.mu[j-1]*omega[j-1]*(1-my.p[j-1])*(1-omega[j-1])/(my.mu[j-1] + 1 + (1-my.p[j-1])/my.p[j-1])  #OOF
 pm[3] <- my.mu[j-1]*omega[j-1]*my.p[j-1]*(1-omega[j-1])/(my.mu[j-1] + 1 + (1-my.p[j-1])/my.p[j-1])      #OON
 pm[4] <- my.mu[j-1]*((1-my.p[j-1])^2)*(1-omega[j-1])/(my.mu[j-1] + 1 + (1-my.p[j-1])/my.p[j-1])    #OFF
 pm[5] <- my.mu[j-1]*(1-my.p[j-1])*(1-omega[j-1])*my.p[j-1] /(my.mu[j-1] + 1 + (1-my.p[j-1])/my.p[j-1])     #OFN
 pm[6] <- my.mu[j-1]*my.p[j-1]*(1-omega[j-1])/(my.mu[j-1] + 1 + (1-my.p[j-1])/my.p[j-1])            #ONO
 pm[7] <- ( ((1-my.p[j-1])^3)/my.p[j-1] )/(my.mu[j-1] + 1 + (1-my.p[j-1])/my.p[j-1])           #FFF
 pm[8] <- ((1-my.p[j-1])^2)/(my.mu[j-1] + 1 + (1-my.p[j-1])/my.p[j-1])                 #FFN
 pm[9] <- (1-my.p[j-1])/(my.mu[j-1] + 1 + (1-my.p[j-1])/my.p[j-1])                     #FNO
 pm[10] <- 1-sum(pm[1:9])                            #NOO


pvec.new <- pm;  # Creating a "new" pvec.new

P.C0.new <- pvec.new[1]
P.C1.new <- sum(pvec.new[c(10,6,3,2)])
P.C2.new <- sum(pvec.new[c(9,5,4)])
P.C3.new <- sum(pvec.new[c(7,8)])

#m0[j] <- Ntot[j-1]-N.unique.observed # don't need


# Sample from full conditional of m = (m1, m2, m3):


my.m[j,] <- rmultinom(1, size=N.unique.observed, prob=c(P.C1.new,P.C2.new,P.C3.new) )


# Generate N.OOO from negative binomial full conditional distn:

# Uses accept-reject algorithm:

check=0
#count=0

while(check==0){

### Accept-Reject to sample from full conditional for N.OOO

# The following parameters help to define the candidate (proposal) distribution
# for the accept-reject algorithm.
# These choices work well for data similar to ours, but may need altering for different data.
size.cand <- 5
mu.cand <- pri.mean.N.OOO # For 2000 data
N.OOO.cand <- rnbinom(1,mu=mu.cand,size=size.cand)


M <- .033 # works well for 2000 data, but may need to be adjusted for other data
my.U <- runif(1)

N.OOO.dens <- dnbinom(N.OOO.cand, size=N.unique.observed, prob= 1-P.C0.new)*dnbinom(N.OOO.cand,mu=pri.mean.N.OOO,size=pri.size.N.OOO)
compare.ratio <- N.OOO.dens/(M*dnbinom(N.OOO.cand,mu=mu.cand,size=size.cand))
check<- ifelse(my.U > compare.ratio, 0, 1)
if (compare.ratio>0.85) print(paste("M for accept-reject for N.OOO:", compare.ratio))
# compare.ratio must be less than 1; this checks to see whether it is getting close at all to 1
# If compare.ratio sometimes exceeds 1, increase M to get a proper accept-reject algorithm
# If compare.ratio is never getting close to 1 (i.e., stays below 0.85 all the time), you could decrease M
# to make the algorithm faster and more efficient.
# Once M is calibrated correctly, there is no need to check and print compare.ratio
}
#print(count)


N.OOO[j] <- N.OOO.cand

# Generate mu from full conditional distn:

data.mean = (3+(1-my.p[j-1])/my.p[j-1])*(N.OOO[j]+N.unique.observed)/N.unique.observed - (1 + (1-my.p[j-1])/my.p[j-1])
data.var = (3+(1-my.p[j-1])/my.p[j-1])^2*(N.OOO[j]+N.unique.observed)*N.OOO[j]/(N.unique.observed^3)

full.cond.mean = (pri.mean.mu/(pri.sd.mu^2) + data.mean/data.var)/( 1/(pri.sd.mu^2) + 1/data.var)
full.cond.var = ((pri.sd.mu^2)*(data.var))/((pri.sd.mu^2)+(data.var))
my.mu[j] <- rnorm(1, mean = full.cond.mean, sd = sqrt(full.cond.var))

### Accept-Reject to sample from full conditional for sigma:

# Generate sigma from full conditional distn:
# Uses accept-reject algorithm:

check=0
#count=0

while(check==0){

# The following defines the candidate (proposal) distribution
# for the accept-reject algorithm.
# This choice works well for data similar to ours, but may need altering for different data.
my.sig.cand <- rigamma(1,5,3)

M <- .65  # works well for 2000 data, but may need to be adjusted for other data
my.U <- runif(1)

omega.cand <- pnotdd(c(my.p[j-1],my.mu[j],my.sig.cand))

pOOO <- my.mu[j]*(omega.cand^2)/(my.mu[j] + 1 + (1-my.p[j-1])/my.p[j-1])  

my.sig.dens <- dnorm((N.OOO[j]/(N.OOO[j]+sum(my.m[j,])) - pOOO)/sqrt(pOOO*(1-pOOO)/(N.OOO[j]+sum(my.m[j,])) ))*densigamma(my.sig.cand, alpha=pri.alpha.sig, beta=pri.beta.sig)
compare.ratio <- my.sig.dens/(M*densigamma(my.sig.cand,5,3))
check<- ifelse(my.U > compare.ratio, 0, 1)
#if (compare.ratio>0.85) print(paste("M for accept-reject for sigma:", compare.ratio))
# compare.ratio must be less than 1; this checks to see whether it is getting close at all to 1
# If compare.ratio sometimes exceeds 1, increase M to get a proper accept-reject algorithm
# If compare.ratio is never getting close to 1 (i.e., stays below 0.85 all the time), you could decrease M
# to make the algorithm faster and more efficient.
# Once M is calibrated correctly, there is no need to check and print compare.ratio
}
#print(count)
my.sig[j] <- my.sig.cand

### Accept-Reject to sample from full conditional for p:

# Generate p from beta full conditional distn:
# Uses accept-reject algorithm:

check=0
#count=0

while(check==0){

# The following defines the candidate (proposal) distribution
# for the accept-reject algorithm.
# This choice works well for data similar to ours, but may need altering for different data.
my.p.cand <- rbeta(1,3,2)

M <- 15  # works well for 2000 data, but may need to be adjusted for other data
my.U <- runif(1)

data.mean = (3+(1-my.p.cand)/my.p.cand)*(N.OOO[j]+N.unique.observed)/N.unique.observed - (1 + (1-my.p.cand)/my.p.cand)
data.var = (3+(1-my.p.cand)/my.p.cand)^2*(N.OOO[j]+N.unique.observed)*N.OOO[j]/(N.unique.observed^3)

full.cond.mean = (pri.mean.mu/(pri.sd.mu^2) + data.mean/data.var)/( 1/(pri.sd.mu^2) + 1/data.var)
full.cond.var = ((pri.sd.mu^2)*(data.var))/((pri.sd.mu^2)+(data.var))
my.mu.part <- dnorm(my.mu[j], mean = full.cond.mean, sd = sqrt(full.cond.var))

my.p.beta.part <- dbeta(my.p.cand, sum(y[1:3]) + pri.alpha.p, sum(y[4:6]) + pri.beta.p)

#my.p[j] <- rbeta(1, sum(y[1:3]) + pri.alpha.p, sum(y[4:6]) + pri.beta.p)

my.p.dens <- my.mu.part*my.p.beta.part
compare.ratio <- my.p.dens/(M*dbeta(my.p.cand,3,2))
check<- ifelse(my.U > compare.ratio, 0, 1)
if (compare.ratio>0.85) print(paste("M for accept-reject for p:", compare.ratio))
# compare.ratio must be less than 1; this checks to see whether it is getting close at all to 1
# If compare.ratio sometimes exceeds 1, increase M to get a proper accept-reject algorithm
# If compare.ratio is never getting close to 1 (i.e., stays below 0.85 all the time), you could decrease M
# to make the algorithm faster and more efficient.
# Once M is calibrated correctly, there is no need to check and print compare.ratio
}
#print(count)
my.p[j] <- my.p.cand

# Calculating omega:

omega[j] <- pnotdd(c(my.p[j],my.mu[j],my.sig[j]))

}

######################################
######################################
#####
##### End of MCMC steps.
#####
######################################
######################################


Ntot.short <- N.unique.observed + N.OOO[-(1:burn.in.amt)]  # getting sequence of N.tot values after burn-in removed.

# Thinning the sequence by taking every k-th value
# to remove serial dependence:

J.short <- length(Ntot.short)
thin.factor <- 2  # increase this if you want more thinning, change it to 1 if you want no thinning at all.
Ntot.thin <- Ntot.short[thin.factor*(1:(J.short/thin.factor))]


Ntot.est <- quantile(Ntot.thin, probs=c(.05,.5,.95))  
# getting a posterior median and quantile-based 90% credible interval for N.tot, the total population size
print(Ntot.est)

my.mu.thin <- (my.mu[-(1:burn.in.amt)])[thin.factor*(1:(J.short/thin.factor))]
my.mu.est <- quantile(my.mu.thin, probs=c(.05,.5,.95))
# getting a posterior median and quantile-based 90% credible interval for mu
print(my.mu.est)

my.p.thin <- (my.p[-(1:burn.in.amt)])[thin.factor*(1:(J.short/thin.factor))]
my.p.est <- quantile(my.p.thin, probs=c(.05,.5,.95))
# getting a posterior median and quantile-based 90% credible interval for p
print(my.p.est)

hist(Ntot.thin, main='Draws from Posterior of N.tot (thinned)', xlab='N.tot values')  
# Histogram of thinned N.tot values from the MCMC

x11()  # opens new plotting window

plot(density(my.mu[-(1:burn.in.amt)]), main="Black: Smoothed Posterior Draws for mu. Red: Prior for mu.")
# Plotting limits may need to be changed depending on your typical values for mu.  Here, 8.5 to 16.5 are reasonable:
myseq<-seq(8.5,16.5,by=.05) 
lines(myseq,dnorm(myseq,pri.mean.mu,pri.sd.mu),col='red')
# The black plot here is a smoothed density estimate of the posterior draws of mu.
# The red plot is the prior distribution for mu.

x11()   # opens new plotting window

# Some diagostic plots for the MCMC (Trace plot and autocorrelation function plot):
par(mfrow=c(2,1))
plot(Ntot.thin, type='l', main="Time Series of (Thinned) N.tot Values", ylab='N.tot', xlab='Iteration Number in Thinned Chain')
acf(Ntot.thin, main="Autocorrelation Function for (Thinned) N.tot Values")


###################
###################
###################
### END 
###################
###################
###################