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


###
###  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])

###
###  Fit Each Probit Occupancy Model Individually 
###

source("occ.aux.mcmc.R")
n.mcmc=160000
s2.beta=1.5^2

tmp.time=proc.time()
tmp.out.1=occ.aux.mcmc(Y,J,X,1,1,1,n.mcmc,s2.beta,null.model=TRUE,no.print=TRUE,get.D=FALSE,no.pred=TRUE)   
tmp.out.2=occ.aux.mcmc(Y,J,X[,1],1,1,1,n.mcmc,s2.beta,null.model=FALSE,no.print=TRUE,get.D=FALSE,no.pred=TRUE)
tmp.out.3=occ.aux.mcmc(Y,J,X[,2],1,1,1,n.mcmc,s2.beta,null.model=FALSE,no.print=TRUE,get.D=FALSE,no.pred=TRUE)   
tmp.out.4=occ.aux.mcmc(Y,J,X[,1:2],1,1,1,n.mcmc,s2.beta,null.model=FALSE,no.print=TRUE,get.D=FALSE,no.pred=TRUE)   
time.1=proc.time()-tmp.time
time.1

###
###  Implement post hoc RJMCMC using Barker and Link method
###

out.list=vector("list",4)
out.list[[1]]=tmp.out.1
out.list[[2]]=tmp.out.2
out.list[[3]]=tmp.out.3
out.list[[4]]=tmp.out.4

make.theta <- function(out.list,M.idx,k){
  theta=rep(0,4)
  theta[1]=out.list[[M.idx]]$p.save[k]
  theta[2]=out.list[[M.idx]]$beta.0.save[k]
  if(M.idx==1){
    theta[-(1:2)]=rnorm(2)
  }
  if(M.idx==2){
    theta[3]=out.list[[M.idx]]$beta.save[k,1]
    theta[4]=rnorm(1)
  }
  if(M.idx==3){
    theta[3]=rnorm(1)
    theta[4]=out.list[[M.idx]]$beta.save[k,1]
  }
  if(M.idx==4){
    theta[3]=out.list[[M.idx]]$beta.save[k,1]
    theta[4]=out.list[[M.idx]]$beta.save[k,2]
  }
  theta
}

make.prob.vec <- function(out.list,k,M.idx,theta){
  prob.vec=rep(0,4)
  prior.prob=rep(1,4)
  y=out.list[[1]]$y
  J=out.list[[1]]$J
  y.0.idx=(y==0)
  n=length(y)

  p=theta[1]
  beta.0=theta[2]
  theta.left=theta[-(1:2)]

  p.beta=0
  p.left=4-(p.beta+2)
  s2.beta=out.list[[1]]$s2.beta
  psi.mix=pnorm(beta.0+rep(0,n))
  lik=prod(psi.mix[!y.0.idx]*dbinom(y[!y.0.idx],J[!y.0.idx],p))*prod(1-psi.mix[y.0.idx]+psi.mix[y.0.idx]*(1-p)^J[y.0.idx])
  prob.vec[1]=lik*1*dnorm(beta.0,0,sqrt(s2.beta))*prod(dnorm(theta.left))*prior.prob[1]

  for(l in 2:4){
    X=out.list[[l]]$X
    if(l==2){
      beta=theta[3]
      theta.left=theta[-(1:3)]
    }
    if(l==3){
      beta=theta[4]
      theta.left=theta[-c(1,2,4)]
    }
    if(l==4){
      beta=theta[3:4]
      theta.left=NULL
    }
    p.beta=length(beta)
    p.left=4-(p.beta+2)
    s2.beta=out.list[[l]]$s2.beta
    psi.mix=pnorm(beta.0+X%*%beta)
    lik=prod(psi.mix[!y.0.idx]*dbinom(y[!y.0.idx],J[!y.0.idx],p))*prod(1-psi.mix[y.0.idx]+psi.mix[y.0.idx]*(1-p)^J[y.0.idx])
    if(!is.null(theta.left)){
      prob.vec[l]=lik*1*dnorm(beta.0,0,sqrt(s2.beta))*prod(dnorm(beta,rep(0,p.beta),sqrt(s2.beta)))*prod(dnorm(rnorm(p.left)))*prior.prob[l]
    }
    if(is.null(theta.left)){
      prob.vec[l]=lik*1*dnorm(beta.0,0,sqrt(s2.beta))*prod(dnorm(beta,rep(0,p.beta),sqrt(s2.beta)))*prior.prob[l]
    }
  }
  prob.vec
}

theta.save=matrix(0,n.mcmc,4)
prob.save=matrix(0,n.mcmc,4)
M.save=rep(0,n.mcmc)
M.idx=4

for(k in 1:n.mcmc){
  if(k%%100==0) cat(k," ")
  
  M.save[k]=M.idx
  theta.save[k,]=make.theta(out.list,M.idx,k)
  prob.tmp=make.prob.vec(out.list,k,M.idx,theta.save[k,])
  prob.save[k,]=prob.tmp/sum(prob.tmp)
  M.idx=(1:4)[rmultinom(1,1,prob.save[k,])==1]

};cat("\n")

summary(theta.save)
summary(prob.save)
table(M.save)

summary(out.list[[1]]$p.save)
summary(out.list[[1]]$beta.0.save)
summary(out.list[[2]]$p.save)
summary(out.list[[2]]$beta.0.save)
summary(out.list[[2]]$beta.save)
summary(out.list[[3]]$p.save)
summary(out.list[[3]]$beta.0.save)
summary(out.list[[3]]$beta.save)
summary(out.list[[4]]$p.save)
summary(out.list[[4]]$beta.0.save)
summary(out.list[[4]]$beta.save)