rm(list = ls()) 
set.seed(111)
library(unmarked)
library(RPresence)
library(hunviphab)

generate_formulas <- function(vars, max_terms = 2) {
  combos <- list("1")
  vars_no1 <- setdiff(vars, "1")
  for (k in 1:max_terms) {
    if (length(vars_no1) >= k) {
      cmb <- combn(vars_no1, k, simplify = FALSE)
      for (c in cmb) {
        combos[[length(combos) + 1]] <- paste(c, collapse = " + ")
      }
    }
  }
  return(combos)
}

get_significance <- function(beta_vec, se_vec) {
  z <- beta_vec / se_vec
  p <- 2 * (1 - pnorm(abs(z)))
  data.frame(z = z, p = p)
}

beta_fit <- function(beta, se, x, intercept = 0) {
  fit_lin <- intercept + beta * x
  se_lin <- se * abs(x)
  fit <- plogis(fit_lin)
  lci <- plogis(fit_lin - 1.96 * se_lin)
  uci <- plogis(fit_lin + 1.96 * se_lin)
  
  data.frame(x = x, fit = fit, lci = lci, uci = uci)
}

umf <- readRDS("data.R")

pao <- createPao(data=umf@y,
                 unitcov=umf@siteCovs,
                 survcov=data.frame(temp=as.numeric(scale(umf@obsCovs$temp)),
                                    doy=as.numeric(scale(umf@obsCovs$doy)),
                                    skill=as.numeric(scale(umf@obsCovs$skill))),
                 nsurveyseason=rep(10,10),
                 title="test")

summary(pao)

scov <- data.frame(vv=as.numeric(scale(umf@yearlySiteCovs$vv_r)),
                   mm=as.numeric(scale(umf@yearlySiteCovs$mm_r)),
                   ca=as.numeric(scale(umf@yearlySiteCovs$ca_r)),
                   pp=as.numeric(scale(umf@yearlySiteCovs$pp_r)),
                   cc=as.numeric(scale(umf@yearlySiteCovs$cc_r)),
                   bird=as.numeric(scale(umf@yearlySiteCovs$bird_r)),
                   mammal=as.numeric(scale(umf@yearlySiteCovs$mammal_r)),
                   introduction=as.numeric(scale(umf@yearlySiteCovs$introduction)))

cor(scov)


psi_forms <- generate_formulas(c("log(area)"),max_terms=2)
p_forms <- generate_formulas(c("temp", "doy", "skill"),max_terms=3)
gamma_forms <- generate_formulas(c("vv","mm","ca","pp","cc","bird","mammal","introduction"),max_terms=1)

# Create model grid
model_grid <- expand.grid(psi=psi_forms,
                          p=p_forms,
                          gamma=gamma_forms,
                          epsilon=gamma_forms,
                          stringsAsFactors = FALSE)

##
mods=list() # create list object to contain results of each model
i <- 1
for(i in 1:nrow(model_grid)){
  try({
    m <- occMod(model=list(as.formula(paste("psi~", model_grid$psi[i])),
                           as.formula(paste("p~", model_grid$p[i])),
                           as.formula(paste("gamma~", model_grid$gamma[i])),
                           as.formula(paste("epsilon~", model_grid$epsilon[i]))),
                data=pao,
                type="do.1",
                cov.list=list(gamma.cov=scov,epsilon.cov=scov))
    mods[[i]] <- m
    
    vars <- c(sub("^[^.]*\\.[^.]*\\.", "", row.names(m$beta$psi)),
              sub("^[^.]*\\.[^.]*\\.", "", row.names(m$beta$gamma)),
              sub("^[^.]*\\.[^.]*\\.", "", row.names(m$beta$epsilon)),
              sub("^[^.]*\\.[^.]*\\.", "", row.names(m$beta$p)))
    
    coeff.i  <- c(m$beta$psi[,1],m$beta$gamma[,1],m$beta$epsilon[,1],m$beta$p[,1])
    se.i <- c(m$beta$psi[,2],m$beta$gamma[,2],m$beta$epsilon[,2],m$beta$p[,2])
    
    coeff <- as.data.frame(setNames(as.list(coeff.i), vars))
    coeff$model <- m[["modname"]]
    coeff$id <- i
    
    se <- as.data.frame(setNames(as.list(se.i), vars))
    se$model <- m[["modname"]]
    se$id <- i
    
    
    if(i==1){
      coeff.df <- coeff
      se.df <- se
    } else {
      coeff.df <- rbindallcolumns(coeff.df,coeff)
      se.df <- rbindallcolumns(se.df,se)
    }
    
  }, silent = TRUE)
  
  cat("Fitted model", i, "of", nrow(model_grid), "\n")
}

results <- createAicTable(mods)
modsel_table <- results$table

write.csv(modsel_table,"modsel_all_models.csv",row.names=F)
write.csv(modsel_table[modsel_table$DAIC<2,],"modsel_daic2.csv",row.names=F)

coeffs <- coeff.df[row.names(modsel_table[modsel_table$DAIC<2,]),]
coeffs <- coeffs[, colSums(!is.na(coeffs)) > 0]

ses <- se.df[row.names(modsel_table[modsel_table$DAIC<2,]),]
ses <- ses[, colSums(!is.na(ses)) > 0]


modav <- data.frame(var=colnames(coeffs[,-5]),
                    beta=colMeans(coeffs[,-5],na.rm=T),
                    se=colMeans(ses[,-5],na.rm=T))

modav <- cbind(modav,get_significance(modav$beta,modav$se))
modav <- modav[!modav$var=="id",]
modav[,2:5] <- round(modav[,2:5],3)

write.csv(modav,"results/modav_beta_coeff.csv",row.names=F)

# plots
gamma.ca <- beta_fit(modav$beta[modav$var=="gamma.ca1"],
                     modav$se[modav$var=="gamma.ca1"],
                     seq(-2.5,2.5,length.out=100),
                     modav$beta[modav$var=="gamma1_B1.Int"])

epsilon.cc <- beta_fit(modav$beta[modav$var=="epsilon.cc1"],
                       modav$se[modav$var=="epsilon.cc1"],
                       seq(-2.5,2.5,length.out=100),
                       modav$beta[modav$var=="epsilon1_C1.Int"])

pdf("figure2.pdf",6,4)
par(mfrow=c(1,2),mar=c(4,4,1,1))
plot(gamma.ca$x,gamma.ca$fit,type="l",lwd=2,
     ylim=c(0,1),xaxs="i", yaxs="i",axes=F,
     xlab="Golden jackal control",ylab="Meadow viper colonization probability")
box(bty="l")
axis(1)
axis(2)
lines(gamma.ca$x,gamma.ca$lci,lty=3)
lines(gamma.ca$x,gamma.ca$uci,lty=3)

plot(epsilon.cc$x,epsilon.cc$fit,type="l",lwd=2,
     ylim=c(0,1),xaxs="i", yaxs="i",axes=F,
     xlab="Hooded crow control",ylab="Meadow viper extinction probability")
box(bty="l")
axis(1)
axis(2)
lines(epsilon.cc$x,epsilon.cc$lci,lty=3)
lines(epsilon.cc$x,epsilon.cc$uci,lty=3)
dev.off()