################################################################################
# Filename: RunwayStrikeBayesSharedV02.R
# Author:	Peter Caley (CSIRO Data61)
# Does what? Bayesian implementation of runway strike model
library(MASS)
library(coda)
library(truncnorm)

# Read me -- need to set path on next line
BASE.PATH <- "Y:/OUR/PATH"

################################################################################
# Data -- http://www.atsb.gov.au/media/3913013/ar-2012-031_.pdf
# Australian Transport Safety Bureau (2012). Australian aviation wildlife strike statistics: Bird and animal 
# strikes 2002 to 2011.  Aviation Research and Analysis Report AR-2012-031. (Australian Transport Safey Bureau: 
# Canberra).

StrikeData <- read.csv(file.path(BASE.PATH,"StrikeData.csv"))
StrikeData <- within(StrikeData, {
 NoTas <- !State%in%c("TAS")
 FoxWidespread <- ifelse(State%in%c("TAS"),1,16)
})
StrikeData_NoTas <- StrikeData[StrikeData$NoTas,]

################################################################################################################
## Calculate Poisson intensity
# Holling type III
calc.mu.Hol <- function(p,dat=StrikeData,x=NULL) {
  a <- p[1]
  b <- p[2]

  if(is.null(x)) {
    x <- dat$HareRabbit
  }
  mu <- a*x^2/(b^2 + x^2)
  mu
}

# Likelihood
lnL <- function(p, dat=StrikeData, neg=!TRUE, type="calc.mu.Pwr", ...) {
 mu <- do.call(type,list(p=p,dat=dat))
   Fox <- dat$Fox
  if(!neg) {
 	return(sum(dpois(Fox,mu,log=TRUE)))
  } else {
 	return(-sum(dpois(Fox,mu,log=TRUE)))
  }
}

# Priors
# for Michaelis menton & Holling type III
a.prior <- function(x) {log(dtruncnorm(x,0,Inf,0,100))}
b.prior <- function(x) {dunif(x,0,1E6,log=T)}

# Normal proposal
norm.sd <- 2.0
propose.norm <- function(x,sd=norm.sd) {
	rnorm(1,x,sd)
}

###################################################################################################################################
# Sampler with Gibbs step
sampler.gs <- function(iter=100,burn.in=1000,start=c(11,9),dat=StrikeData_NoTas, type="calc.mu.MM") {
 iter <- iter + burn.in
 res <- matrix(NA,nrow=iter,ncol=5)

 # Lagomorphs strikes for Tasmania
 x.tas <- StrikeData$HareRabbit[StrikeData$State=="TAS"]

 p.curr <- start
 for (i in 1:iter) {	#i=1  i=i+1
   p.prop <- c(propose.norm(p.curr[1]),propose.norm(p.curr[2]))

   # Calculate mu.tas | p.prop
   mu.tas.prop <- do.call(type,list(p=p.prop,x=x.tas))
   # Impute fox runway strikes
   z.tas <- rpois(1,mu.tas.prop)	
   # Augmented data
   dat.aug <- rbind(dat[,c(1,2,4)],data.frame(State="TAS",HareRabbit=x.tas,Fox=z.tas))	

  MR <- exp(lnL(p=p.prop, dat=dat.aug, type=type) - lnL(p=p.curr, dat=dat.aug, type=type) +
   a.prior(p.prop[1]) - a.prior(p.curr[1]) + b.prior(p.prop[2]) - b.prior(p.curr[2])
  )

  if(runif(1) < MR & !is.na(MR)) {
   p.curr <- p.prop
   res[i,5] <- 1 # acceptances
  } else {
   res[i,5] <- 0 # rejections
  }
  res[i,1:length(p.curr)] <- p.curr
  res[i,4] <- z.tas
 } # end of i loop
 res[-(1:burn.in),]
} # end of function sampler.gs()


###################################################################################################################################
# Run sampler
# Gibbs version if Holling type III

chain1 <- sampler.gs(iter=1E4,burn.in=5000,start=c(5,5),type="calc.mu.Hol")
chain2 <- sampler.gs(iter=1E4,burn.in=5000,start=c(1,10),type="calc.mu.Hol")
chain3 <- sampler.gs(iter=1E4,burn.in=5000,start=c(10,1),type="calc.mu.Hol")
chain4 <- sampler.gs(iter=1E4,burn.in=5000,start=c(10,10),type="calc.mu.Hol")

res <- rbind(chain1,chain2,chain3,chain4)
res.thin10 <- res[seq(1,nrow(res),10),]

###################################################################################################################################
# Diagnostics
HT3.chains <- (mcmc.list(list(mcmc(chain1[,1:2]),mcmc(chain2[,1:2]),mcmc(chain3[,1:2]),mcmc(chain4[,1:2]))))
win.graph()
par(mfrow=c(3,3))
cumuplot(HT3.chains)

# Acceptance rate
mean(res[,5])

# Plots
win.graph(h=7,w=10)
par(mfrow=c(2,1))
plot(ts(res[,1]))
plot(ts(res[,2]))


# Quantiles for runway strike intensity for Tasmania
quantile(res[,2],c(0,0.005,0.025,0.5,0.975,0.995))
a.median <- median(res[,1])
a.mean <- mean(res[,1])

b.median <- median(res[,2])
b.mean <- mean(res[,2])

# Prediction intervals for expected Tasmanian fox strikes
quantile(res[,4],c(0,0.005,0.025,0.5,0.975,0.995),na.rm=T)

# Plots
win.graph(h=7,w=10)
par(mfrow=c(2,1))

# Calculate intensities
mu.Hol <- sapply(1:nrow(res),function(i) 
calc.mu.Hol(p=c(res[i,1],res[i,2]),dat=StrikeData,x=0:50),simplify=T)
mu.Hol.thin10 <- sapply(1:nrow(res.thin10),function(i) 
calc.mu.Hol(p=c(res.thin10[i,1],res.thin10[i,2]),dat=StrikeData,x=0:50),simplify=T)

# Tassie intensity
tas.mu <- mu.Hol[16,]
mean(tas.mu)
quantile(tas.mu,c(0,0.005,0.025,0.5,0.975,0.995))

# Predict for Tassie
tas.pred <- sapply(1:length(tas.mu),function(i) rpois(1,tas.mu[i]))
hist(tas.pred)	# check is smooth
mean(tas.pred)
quantile(tas.pred,c(0,0.005,0.025,0.5,0.975,0.995))

# Parametric simulations from posterior for intensities
par.boot.data <- sapply(1:nrow(mu.Hol.thin10), function(i) {replicate(10000, 
rpois(1,sample(mu.Hol.thin10[i,])))})
par.boot.PIs <- t(sapply(1:ncol(par.boot.data), function(i) 
quantile(par.boot.data[,i],c(0.005,0.025,0.5,0.975,0.995))))

#######################################################################################################
# Recreate plot
win.graph(h=7,w=10)
par(cex.lab=2,cex.axis=2,mar=c(6,6,2,2),mgp=c(3,1,0))
with(StrikeData,plot(HareRabbit,Fox,cex=3,pch=FoxWidespread,xlab="Lagomorph strikes",ylab="Fox 
strikes",axes=F,
ylim=c(0,15),xlim=c(0,50)))
axis(1,lwd=3,tck=0.02)
axis(2,las=2,lwd=3,tck=0.02)

# Label points
sapply(1:nrow(StrikeData),function(i) {
 with(StrikeData,text(HareRabbit[i]+3,Fox[i],State[i],cex=2))
}
)

# For Holling type III model
lines(0:50, calc.mu.Hol(c(a.median,b.median),x=0:50),lwd=3)

lines(0:50,par.boot.PIs[,2],lty=2,lwd=3)
lines(0:50,par.boot.PIs[,4],lty=2,lwd=3)

lines(0:50,par.boot.PIs[,1],lty=3,lwd=3)
lines(0:50,par.boot.PIs[,5],lty=3,lwd=3)

legend(35,15,c("95% P.I.","99% P.I."),lty=c(2,3),lwd=3,bty='n',cex=2)
################################################################################