##########################################################################
#   Specht, et al. MEE Conditional Occupancy Modelling
#
#   This script applies conditional occupancy design to a simulated dataset 
#   and analyzed with the package unmarked.
#
#  T Arnold, Dept of Fisheries & Wildlife, Univ of Minnesota, arnol065@umn.edu
##########################################################################

#Load libraries

library(unmarked)

#Simulate occupancy data following code in Kery & Royle, Applied Heirarchical Modelling in Ecology

E=420 #total effort

#approximate true values of psi and p for estimating optimal k & S
psi.t<-0.2
p.t<-0.5

#We're going to simulate a true psi of ~0.2 and a true p of ~0.5, so we'll use optimal k (number of visits) for that combination
k.S<-3  #optimal k for standard method 
k.R<-4  #optimal k for removal method 
k.C<-4  #optimal k for conditional method 


#Approximate number of sites to visit for each method  
S.C = trunc(E/(psi.t*p.t*k.C-psi.t*p.t+1)) # Expected value of conditional sites at given psi, p, k, E 
S.S = trunc(E/k.S) # Expected value of standard sites at given E, k 
S.R = trunc(E/((k.R*(1-psi.t))+((psi.t*(1-(1-p.t)^k.R))/p.t))) # Expected value of removal sites at given psi, p, k, E


# To create one dataset that both survey structures sample, we'll use S.C sites and k.C surveys (conditional is greater for both of these)


  # Define covariates
  ag.cover<-round(runif(max(S.C, S.S, S.R), -1, 1),3) #create a standardized continuous covariate "ag cover" affecting occupancy
  biomass<-round(array(rep(runif(max(S.C, S.S, S.R), -1, 1),max(k.C, k.S, k.R)), dim=c(max(S.C, S.S, S.R),max(k.C, k.S, k.R))),3) # create a standardized continuous covariate"biomass" affecting occupancy AND detection
  veg.den<-round(array(rep(runif(max(S.C, S.S, S.R), -1, 1),max(k.C, k.S, k.R)), dim=c(max(S.C, S.S, S.R),max(k.C, k.S, k.R))),3) # continuous covariate for detection across sites, but SAME across surveys
  
  #Define relationships between covariates, site occupancy and detection probability
  # Relationships defined such that psi.t=~ 0.2 and p.true=~0.5
  beta0.psi<- -1.5 
  beta1.psi<- -1
  beta2.psi<- 1
  psi<-plogis(beta0.psi+beta1.psi*ag.cover+beta2.psi*biomass[,1]) 
  mean(psi) #mean occupancy
  
  beta0.p<- 0
  beta1.p<- 2
  beta2.p<- -1
  
  
  p<-plogis(beta0.p+beta1.p*veg.den+beta2.p*biomass)
  mean(p) #mean detection
  
  
  #generate true presence/absences
  z<-rbinom(max(S.C, S.S, S.R), 1, psi)
  mean(z) #look at mean psi
  
  
  #create matrix with detection histories across max sites(rows) and surveys(columns), given occupancy
  y<-matrix(NA, nrow=(max(S.C, S.R, S.S)), ncol=max(k.S, k.C, k.R))
  
  for(j in 1:(max(k.C, k.S, k.R))){
    y[,j]<-rbinom((max(S.C, S.R, S.S)), z, p[,j])
  }
  
  #Check number of sites with observed occupancy 
  sum(apply(y,1,max))     
  sum(apply(y,1,max))/max(S.C, S.R, S.S)  
  
  #set up data
  dat<-cbind(y, ag.cover, biomass[,1], veg.den, biomass)
  
  #-------------------------------
  #Structure data for standardized, removal, conditional analysis at same total E
  
  #Standard: 
  
  y.s<-data.frame(dat[(1:S.S),c(1:k.S,((max(k.S, k.C,k.R)+1):(max(k.S, k.C,k.R)+2)), ((max(k.S,k.C,k.R)+3):(max(k.S,k.C,k.R)+2+k.S)),
                      ((max(k.S,k.C,k.R)*2 +3):((max(k.S,k.C,k.R)*2 +2+k.S))))])
            #change to data-frame- remove 5th and 6th observations and restrict to number of standard sites
  
  #Removal
  y.r<-data.frame(dat[(1:S.R),c(1:k.R,((max(k.S, k.C,k.R)+1):(max(k.S, k.C,k.R)+2)), ((max(k.S,k.C,k.R)+3):(max(k.S,k.C,k.R)+2+k.R)),
                                ((max(k.S,k.C,k.R)*2 +3):((max(k.S,k.C,k.R)*2 +2+k.R))))])  
  for(j in 2:k.R){
    
    y.r[,j][y.r[,(j-1)]==1]<-NA
    y.r[,j][is.na(y.r[,j-1])]<-NA
  }
  
  #Conditional: 
  
  y.c<-data.frame(dat[(1:S.C),c(1:k.C,((max(k.S, k.C,k.R)+1):(max(k.S, k.C,k.R)+2)), ((max(k.S,k.C,k.R)+3):(max(k.S,k.C,k.R)+2+k.C)),
                                ((max(k.S,k.C,k.R)*2 +3):((max(k.S,k.C,k.R)*2 +2+k.C))))])  
  for (j in 2:k.C){
    y.c[,j][y.c$V1==0]<-NA  # if nothing observed in first visit, N/a for other visits
  }
  
  #--------------------------------------------------------
  #    Prep unmarked files & run analyses for each method
  #--------------------------------------------------------
  
      #-------------------------------
      #    Standard design
      #-------------------------------
      
      
      #define unmarked frame
      
      Veg=y.s[,(k.S+3):(2+(2*k.S))]
      biomass=y.s[,(3+(2*k.S)):(2+(3*k.S))]
      
      
      umf.s<-unmarkedFrameOccu(y=y.s[,1:k.S], 
                               siteCovs=data.frame(Ag=y.s$ag.cover,biomass.psi=y.s[,k.S+2]),
                               obsCovs=list(Veg=Veg, biomass.p=biomass))
    
      #Fit models, extract estimates
      
      #Constant model
      fm0.s<-occu(~1~1, data=umf.s)
      
      #Predict at mean values of covariates
      newdat.psi.s<-data.frame(Ag=mean(y.s$ag.cover, na.rm=T), biomass.psi=mean(y.s[,(k.S+2)], na.rm=T)) # create new data frame of mean values of covariates
      newdat.p.s<-data.frame(Veg=mean(umf.s@obsCovs$Veg,na.rm=T), biomass.p=mean(umf.s@obsCovs$biomass.p,na.rm=T))
      
      predict(fm0.s, type="state", newdata=newdat.psi.s)
      predict(fm0.s, type="det", newdata=newdat.p.s)
      
      #Covariate model
      fm1.s<-occu(~Veg+biomass.p~Ag+biomass.psi, data=umf.s)
      
      predict(fm1.s, type="state", newdata=newdat.psi.s)
      predict(fm1.s, type="det", newdata=newdat.p.s)
      
     
   #------------------------------------------------------------------
  
      #-------------------------------
      #    Removal design
      #-------------------------------
      
      #define unmarked frame

      Veg=y.r[,(k.R+3):(2+(2*k.R))]
      biomass=y.r[,(3+(2*k.R)):(2+(3*k.R))]
      umf.r<-unmarkedFrameOccu(y=y.r[,1:k.R], 
                               siteCovs=data.frame(Ag=y.r$ag.cover,biomass.psi=y.r[,k.R+2]),
                               obsCovs=list(Veg=Veg, biomass.p=biomass))

      #Fit models, extract estimates
      
      #Constant model
      fm0.r<-occu(~1~1, data=umf.r)
      
      #Predict at mean values of covariates
      newdat.psi.r<-data.frame(Ag=mean(y.r$ag.cover, na.rm=T), biomass.psi=mean(y.r[,(k.R+2)], na.rm=T)) # create new data frame of mean values of covariates
      newdat.p.r<-data.frame(Veg=mean(umf.r@obsCovs$Veg,na.rm=T), biomass.p=mean(umf.r@obsCovs$biomass.p,na.rm=T))
      
      predict(fm0.r, type="state", newdata=newdat.psi.r)
      predict(fm0.r, type="det", newdata=newdat.p.r)
      
      #Covariate model
      fm1.r<-occu(~Veg+biomass.p~Ag+biomass.psi, data=umf.r)
      
      predict(fm1.r, type="state", newdata=newdat.psi.r)
      predict(fm1.r, type="det", newdata=newdat.p.r)
      
  #----------------------------------
  
      #------------------------------------
      # Conditional
      #------------------------------------
      
      #define unmarked frame

      Veg=y.c[,(k.C+3):(2+(2*k.C))]
      biomass=y.c[,(3+(2*k.C)):(2+(3*k.C))]
      umf.c<-unmarkedFrameOccu(y=y.c[,1:k.C], 
                               siteCovs=data.frame(Ag=y.c$ag.cover,biomass.psi=y.c[,k.C+2]),
                               obsCovs=list(Veg=Veg, biomass.p=biomass))
      

      #Fit models, extract estimates
      
      #Constant model
      fm0.c<-occu(~1~1, data=umf.c)
      
      #Predict at mean values of covariates
      newdat.psi.c<-data.frame(Ag=mean(y.c$ag.cover, na.rm=T), biomass.psi=mean(y.s[,(k.C+2)], na.rm=T)) # create new data frame of mean values of covariates
      newdat.p.c<-data.frame(Veg=mean(umf.c@obsCovs$Veg,na.rm=T), biomass.p=mean(umf.c@obsCovs$biomass.p,na.rm=T))
      
      predict(fm0.c, type="state", newdata=newdat.psi.c)
      predict(fm0.c, type="det", newdata=newdat.p.c)
      
      #Covariate model
      fm1.c<-occu(~Veg+biomass.p~Ag+biomass.psi, data=umf.c)
      
      predict(fm1.c, type="state", newdata=newdat.psi.c)
      predict(fm1.c, type="det", newdata=newdat.p.c)
  #------------------------------------------------------------------