#
#  R Script to fit Occupancy model to Willow Tit Data
#  
#  Model Selection Approach: Bayesian Regularization 
#

###
###  Load Covariates and Occupancy Data
###

wt.df=read.csv("wt.csv",header=TRUE)
head(wt.df)

n=200
Y=as.matrix(wt.df[1:n,1:3])
y=apply(wt.df[1:n,1:3],1,sum,na.rm=TRUE)
J=apply(!is.na(wt.df[1:n,1:3]),1,sum)

X=matrix(1,n,2)
X[,1]=scale(wt.df$elev[1:n])
X[,2]=scale(wt.df$forest[1:n])

###
###  Use Parallel Regularization Fit with Full Data Sets 
###
###  Issue this command in shell before starting R: export OMP_NUM_THREADS=1 
###

library(parallel)
library(snowfall)
library(rlecuyer)
source("occ.aux.mcmc.R")

n.mcmc=160000
n.beta=24
s2.beta.vec=seq(.1,1.5,,n.beta)^2

cps=detectCores()
sfInit(parallel=TRUE, cpus=cps)
sfExportAll()
sfClusterSetupRNG()

tmp.fcn <- function(i){
  tmp.out=occ.aux.mcmc(Y,J,X,1,1,1,n.mcmc,s2.beta.vec[i],null.model=FALSE,no.print=TRUE,no.pred=TRUE)
  list(WAIC=tmp.out$WAIC.2,beta=apply(tmp.out$beta.save,2,mean))
}

sfExport("Y","J","X","s2.beta.vec","n.mcmc")
tmp.time=Sys.time()
score.list=sfClusterApplySR(1:n.beta,tmp.fcn,perUpdate=1)
time.1=Sys.time()-tmp.time
sfStop()

time.1
out.mat=matrix(unlist(score.list),,3,byrow=TRUE)

###
###  Fit Each Probit Occupancy Model using Parallelized C-V 
###
###  Issue this command in shell before starting R: export OMP_NUM_THREADS=1 
###

library(parallel) 
library(snowfall) 
library(rlecuyer) 
source("occ.aux.mcmc.R")

n.mcmc=160000
n.beta=24
s2.beta.vec.2=seq(.25,1.5,,n.beta)^2

K=10
cv.s2.grid=expand.grid(1:n.beta,1:K)
n.grid=dim(cv.s2.grid)[1]
fold.idx.mat=matrix(1:n,,K)

cps=detectCores()
sfInit(parallel=TRUE, cpus=cps)
sfExportAll()
sfClusterSetupRNG()

cv.fcn <- function(i){
  k=cv.s2.grid[i,2]
  fold.idx=fold.idx.mat[,k] 
  y.oos=apply(Y[fold.idx,],1,sum,na.rm=TRUE)
  tmp.out=occ.aux.mcmc(Y[-fold.idx,],J[-fold.idx],X[-fold.idx,],y.oos,J[fold.idx],X[fold.idx,],n.mcmc,s2.beta.vec.2[cv.s2.grid[i,1]],null.model=FALSE,no.print=TRUE)    
  tmp.out$score
}

sfExport("Y","J","X","n.mcmc","s2.beta.vec.2","cv.s2.grid","cv.fcn","fold.idx.mat")
tmp.time=Sys.time()
score.list=sfClusterApplySR(1:n.grid,cv.fcn,perUpdate=1)
time.2=Sys.time()-tmp.time
sfStop()

time.2
score.cv.mat=matrix(unlist(score.list),n.beta,K)

###
###  Shrinkage Trajectories and CV Score
###

score.cv.vec=apply(score.cv.mat,1,sum)
s2.beta.opt=s2.beta.vec.2[score.cv.vec==min(score.cv.vec)]
pdf(file="bayes_reg.pdf",height=10)
layout(matrix(1:2,2,1))
matplot(s2.beta.vec,out.mat[,-1],type="l",lwd=4,lty=c(1,2,rep(3,p.extra)),col=1,ylab=bquote(beta),xlab=bquote(sigma[beta]^2),xlim=c(0,2.25),main="a")
abline(h=0,col=1,lwd=1,lty=3)
abline(v=s2.beta.opt,col=8,lwd=2)
legend("right",bty="n",lwd=c(4,4,2),lty=c(1,2,1),col=c(1,1,8),legend=c("elevation","forest","optimal penalty"))
plot(s2.beta.vec.2[-c(1:2)],score.cv.vec[-c(1:2)],type="l",lwd=3,ylab="c-v score",xlab=bquote(sigma[beta]^2),xlim=c(0,2.25),main="b")
abline(v=s2.beta.opt,col=8,lwd=2)
legend("right",bty="n",lwd=c(4,2),lty=1,col=c(1,8),legend=c("c-v score","optimal penalty"))
dev.off()