# Most of these methods are derived from:
# Martins, F. C., Godinho, S., Guiomar, N., Medinas, D., Rebelo, H.,
# Segurado, P., & Marques, J. T. (2024). Vegetation canopy height shapes
# bats’ occupancy: a remote sensing approach. GIScience & Remote Sensing, 61(1), 2374150.

library(tidyverse)
library(optparse)
library(ebirdst)
library(auk)
library(janitor)
library(broom)
library(lubridate)
library(spOccupancy)
library(doFuture)
library(future.callr)
library(raster)
library(terra)
library(sf)
library(tidymodels)
library(fields)
library(units)
library(progressr)

options(scipen=1e9)
handlers(global = TRUE)
registerDoFuture()
options(future.rng.onMisuse = "ignore")
options(future.globals.maxSize=(4000*1024^2))

select<-dplyr::select

grid_size<-set_units(600, "m")
n.samples <- 80000
n.burn <- 3000
n.thin <- 25
n.chains <- 3
workers<-8

its<-10

setColNames<-function(x, names_x){
  colnames(x)<-names_x
}

order_by_importance<-function(preds){
  return(tibble(predictor=preds) %>%
           left_join(read_csv(list.files(file.path(getwd(), "results", species$species_code, "eBird", "encounter_rate"), recursive = T, full.names = T,
                                         pattern = "top_pred.csv"), id = "filename", show_col_types = FALSE) %>%
                       mutate(run=as.numeric(str_extract(filename, "(?<=_)[0-9]+"))) %>%
                       arrange(run) %>% mutate(run=factor(run)) %>%
                       select(-filename) %>% relocate(run) %>%
                       group_by(predictor) %>%
                       summarise(importance=mean(importance)), by = "predictor") %>%
           arrange(desc(importance)) %>% pull(predictor))
}

add_quadratic_terms <- function(formula) {
  # Extract terms from formula
  terms <- all.vars(formula)
  linear_terms<-paste0("scale(", terms, ")")
  
  # Construct new quadratic terms
  quadratic_terms <- paste0("I(scale(", terms, ")^2)")
  
  # Create new formula
  new_formula <- as.formula(paste("~", paste(c(linear_terms, quadratic_terms), collapse = " + ")))
  
  return(new_formula)
}

remove_intercept<-function(terms){
  return(terms[terms != "(Intercept)"])
}

terms_format<-function(x){
  unique(str_remove(str_remove(str_remove(str_remove(x[!grepl("Intercept", x)], "scale\\("), "I\\("), "(?<!Intercept)\\)"), "\\^2\\)"))
}

# Hastie, Tibshirani, & Friedman (2009). The Elements of Statistical Learning. Chapter 14.
# Zhao & Xie (2019). A correlation-based feature selection method for machine learning classification problems. Journal of Big Data.
# Guyon & Elisseeff (2003). An introduction to variable and feature selection. Journal of Machine Learning Research.

remove_correlated_vars<-function(data, cutoff=0.9, order_by_importance=T){
  cor_matrix<-cor(data, use="pairwise.complete.obs")
  dist_matrix<-as.dist(1-abs(cor_matrix))
  hc<-hclust(dist_matrix, method="complete")
  clusters<-cutree(hc, h=1-cutoff)
  selected_vars<-sapply(unique(clusters), function(cluster){
    if(order_by_importance){
      order_by_importance(names(which(clusters==cluster)))[1]
    } else{
      names(which(clusters==cluster))[1]
    }
  })
  return(data[,selected_vars,drop=F])
}

tidy.PGOcc<-function(x, prob=F, conf.int = FALSE, conf.level = 0.95, HDI=FALSE, det = FALSE, ...){
  if(det){
    msd<-t(apply(x$alpha.samples, 2, function(x) c(mean(x), sd(x))))
    par_names<-rownames(msd)
    ret <- tibble(term=par_names, estimate=msd[,1], std.error=msd[,2])
    
    if (conf.int) {
      if(HDI){
        ci <- coda::HPDinterval(x$alpha.samples, prob = conf.level)
        cols <- c("term", "lower", "upper")
      } else {
        ci <- t(apply(x$alpha.samples, 2, function(x) quantile(x, prob = c((1-conf.level)/2, 1-((1-conf.level)/2)))))
        cols<-c("term", "conf.low", "conf.high")
      }
      ci<-tibble(term=rownames(ci), ci[,1], ci[,2])
      colnames(ci)<-cols
      ret <- dplyr::left_join(ret, ci, by = "term")
    }
    if(prob){
      ret<-ret %>% mutate(across(where(is.numeric), plogis))
    }
    ret
  } else{
    msd<-t(apply(x$beta.samples, 2, function(x) c(mean(x), sd(x))))
    par_names<-rownames(msd)
    ret <- tibble(term=par_names, estimate=msd[,1], std.error=msd[,2])
    if (conf.int) {
      if(HDI){
        ci <- coda::HPDinterval(x$beta.samples, prob = conf.level)
        cols <- c("term", "lower", "upper")
      } else {
        ci <- t(apply(x$beta.samples, 2, function(x) quantile(x, prob = c((1-conf.level)/2, 1-((1-conf.level)/2)))))
        cols<-c("term", "conf.low", "conf.high")
      }
      ci<-tibble(term=rownames(ci), ci[,1], ci[,2])
      colnames(ci)<-cols
      ret <- dplyr::left_join(ret, ci, by = "term")
    }
    if(prob){
      ret<-ret %>% mutate(across(where(is.numeric), plogis))
    }
    ret
  }
}

all_multivariate_formulas <- function(formula) {
  # Extract response and predictors
  terms_obj <- terms(formula)
  predictors <- attr(terms_obj, "term.labels") # Extract predictor names
  
  # Generate all combinations of predictors
  all_formulas <- list(as.formula("~1"))
  for (i in 1:length(predictors)) {
    predictor_combinations <- combn(predictors, i, simplify = FALSE)
    for (preds in predictor_combinations) {
      new_formula <- as.formula(paste("~", paste(preds, collapse = " + ")))
      all_formulas <- append(all_formulas, list(new_formula))
    }
  }
  return(all_formulas)
}

all_univariate_formulas <- function(formula){
  # Extract response and predictors
  terms_obj <- terms(formula)
  predictors <- attr(terms_obj, "term.labels") # Extract predictor names
  
  # Generate all combinations of predictors
  all_formulas <- list(as.formula("~ 1"))
  for (i in 1:length(predictors)) {
    new_formula <- as.formula(paste("~", predictors[i]))
    all_formulas <- append(all_formulas, list(new_formula))
  }
  return(all_formulas)
}

make.occ.formula<-function(occ_covs=NULL, add_quadratics=F){
  occ_covs<-occ_covs[!grepl("(Intercept)", occ_covs)]
  if(is.null(occ_covs) | length(occ_covs)==0){
    occ.formula <- ~ 1
  } else if(length(occ_covs)==1){
    if(occ_covs==1){
      occ.formula <- ~ 1
    } else{
      if(sum(grepl("scale", occ_covs)>0)){
        occ.formula<-formula(paste0("~", occ_covs))
      } else{
        occ.formula<-formula(paste0("~", paste0("scale(", occ_covs, ")")))
      }
    }
  } else{
    if(sum(grepl("scale", occ_covs)>0)){
      occ.formula<-formula(paste0("~", paste(occ_covs, collapse = " + ")))
    } else{
      occ.formula<-formula(paste0("~", paste0("scale(", occ_covs, ")", collapse = " + ")))
    }
    if(add_quadratics){
      occ.formula<-add_quadratic_terms(occ.formula)
    }
  }
  return(occ.formula)
}

make.det.formula<-function(det_covs=NULL, add_quadratics=F){
  det_covs<-det_covs[!grepl("(Intercept)", det_covs)]
  if(is.null(det_covs) | length(det_covs)==0){
    det.formula <- ~ 1
  } else if(length(det_covs)==1){
    if(det_covs==1){
      det.formula <- ~ 1
    } else{
      if(sum(grepl("scale", det_covs)>0)){
        det.formula<-formula(paste0("~", det_covs))
      } else{
        det.formula<-formula(paste0("~", paste0("scale(", det_covs, ")")))
      }
    }
  } else{
    if(sum(grepl("scale", det_covs)>0)){
      det.formula<-formula(paste0("~", paste(det_covs[det_covs!="protocol_type"], collapse = " + ")))
    } else{
      det.formula <- formula(paste0("~", paste0("scale(", det_covs[det_covs!="protocol_type"], ")", collapse = " + ")))
    }
    if(add_quadratics){
      det.formula<-add_quadratic_terms(det.formula)
    }
    if("protocol_type" %in% det_covs){
      det.formula<-update(det.formula, ~ . + protocol_type)
    }
  }
  return(det.formula)
}

make.data.ss<-function(occ_covs=NULL, det_covs=NULL, wide_data){
  # Occupancy covariates
  occ.covs <- wide_data %>% select(contains(occ_covs)) %>% as.matrix
  occ.formula<-make.occ.formula(occ_covs)
  
  y<-wide_data %>% select(matches("^y.")) %>% remove_empty(which="cols") %>% as.matrix
  
  covs<-wide_data %>% remove_empty(which="cols") %>% select(-day_of_year)
  
  det.covs<-lapply(setNames(det_covs, det_covs), function(x){
    covs %>% select(matches(x)) %>% as.matrix
  })
  
  det.formula <- make.det.formula(det_covs)
  
  return(list(y=y, occ.covs = occ.covs, det.covs = det.covs))
}

univariate_model_tests<-function(occ_covs=NULL, det_covs=NULL, wide_data, cutoff=0.95){
  occ.formula<-make.occ.formula(occ_covs, add_quadratics = T)
  det.formula<-make.det.formula(det_covs, add_quadratics = T)
  
  data.ss<-make.data.ss(occ_covs, det_covs, wide_data)
  univariate_occ_formulas<-all_univariate_formulas(occ.formula)
  univariate_det_formulas<-all_univariate_formulas(det.formula)
  
  plan(callr, workers=8)
  
  occ_preds<-do.call(rbind, plyr::llply(univariate_occ_formulas, .fun = function(x){
    inits <- list(alpha = 0, 
                  beta = 0, 
                  z = apply(data.ss$y, 1, max, na.rm = TRUE))
    priors <- list(alpha.normal = list(mean = 0, var = 2.72), 
                   beta.normal = list(mean = 0, var = 2.72))
    
    out <- PGOcc(occ.formula = x,
                 det.formula = ~1, 
                 data = data.ss, 
                 inits = inits, 
                 n.samples = n.samples, 
                 priors = priors, 
                 n.omp.threads = 1, 
                 verbose = F, 
                 n.report = 1000, 
                 n.burn = n.burn, 
                 n.thin = n.thin, 
                 n.chains = n.chains)
    
    return(out %>% tidy.PGOcc(conf.int = T, conf.level = cutoff, HDI=T) %>% mutate(Rhat=out$rhat$beta))
  }, .progress = "progressr", .parallel = T))
  
  det_preds<-do.call(rbind, plyr::llply(univariate_det_formulas, .fun = function(x){
    inits <- list(alpha = 0, 
                  beta = 0, 
                  z = apply(data.ss$y, 1, max, na.rm = TRUE))
    priors <- list(alpha.normal = list(mean = 0, var = 2.72), 
                   beta.normal = list(mean = 0, var = 2.72))
    out <- PGOcc(occ.formula = ~1,
                 det.formula = x,
                 data = data.ss, 
                 inits = inits, 
                 n.samples = n.samples, 
                 priors = priors, 
                 n.omp.threads = 1, 
                 verbose = F, 
                 n.report = 1000, 
                 n.burn = n.burn,
                 n.thin = n.thin, 
                 n.chains = n.chains)
    return(out %>% tidy.PGOcc(conf.int = T, conf.level = cutoff, det = T, HDI=T) %>% mutate(Rhat=out$rhat$beta))
  }, .progress = "progressr", .parallel = T))
  
  plan(sequential)
  
  return(list(occ_preds=occ_preds, det_preds=det_preds))
}

multivariate_model_tests2<-function(univariate_model_output, occ_covs, det_covs, wide_data){
  
  best_model<-NULL
  best_waic<-Inf
  
  informative_occ_terms<-univariate_model_output$occ_preds %>% filter(term != "(Intercept)") %>%
    mutate(informative = !(0 >= lower & 0 <= upper)) %>%
    filter(informative==T) %>%
    pull(term)
  if(length(informative_occ_terms)==0){
    informative_occ_terms<-"1"
  }
  
  informative_det_terms<-univariate_model_output$det_preds %>% filter(term != "(Intercept)") %>%
    mutate(informative = !(0 >= lower & 0 <= upper)) %>%
    filter(informative==T) %>%
    pull(term)
  if(length(informative_det_terms)==0){
    informative_det_terms<-"1"
  }
  
  all_terms<-tibble(pred=c(informative_occ_terms, informative_det_terms),
                    type=c(rep("occ", length(informative_occ_terms)), rep("det", length(informative_det_terms))))
  all_terms<-all_terms %>% mutate(pred=if_else(grepl("protocol_type", pred), str_replace(pred, "Traveling|Stationary", ""), pred))
  if(sum(all_terms$type=="occ")==0){
    all_terms<-c(pred=1, type="occ") %>% bind_rows(all_terms)
  }
  if(sum(all_terms$type=="det")==0){
    all_terms<-all_terms %>% bind_rows(c(pred=1, type="det"))
  }
  
  models<-list()
  waic_values<-c()
  
  while(nrow(all_terms)>1){
    plan(callr, workers=workers)
    
    res_collector<-plyr::alply(.data = all_terms$pred, .margins = 1, .fun = function(i){
      current_pred<-i
      new.occ.formula<-make.occ.formula(informative_occ_terms[informative_occ_terms %in% setdiff(all_terms$pred, i)])
      new.det.formula<-make.det.formula(informative_det_terms[informative_det_terms %in% setdiff(all_terms$pred, i)]) # not a typo
      data.ss<-make.data.ss(occ_covs, det_covs, wide_data)
      inits <- list(alpha = 0, 
                    beta = 0, 
                    z = apply(data.ss$y, 1, max, na.rm = TRUE))
      priors <- list(alpha.normal = list(mean = 0, var = 2.72), 
                     beta.normal = list(mean = 0, var = 2.72))
      out <- PGOcc(occ.formula = new.occ.formula,
                   det.formula = new.det.formula,
                   data = data.ss, 
                   inits = inits, 
                   n.samples = n.samples, 
                   priors = priors, 
                   n.omp.threads = 1, 
                   verbose = F, 
                   n.report = 1000, 
                   n.burn = n.burn, 
                   n.thin = n.thin, 
                   n.chains = n.chains)
      
      waic_score<-waicOcc(out)[3]
      return(list(ith_model=out, waic_value=waic_score))
    }, .parallel = T, .progress = "progressr")
    
    plan(sequential)
    
    names(res_collector)<-all_terms$pred
    plyr::a_ply(all_terms$pred, .margins = 1, function(i){
      models[[i]]<<-res_collector[i][[1]]$ith_model
      waic_values[i]<<-res_collector[i][[1]]$waic_value
    })
    
    best_pred<-names(which.min(waic_values))
    message("Selected for removal: ", best_pred)
    best_waic_new<-min(waic_values)
    if(best_waic_new >= best_waic){
      break
    } else{
      message("**NEW BEST WAIC SCORE: ", round(best_waic_new,3), "**")
      best_waic <- best_waic_new
      best_model<-models[[best_pred]]
      all_terms<-all_terms[all_terms$pred != best_pred,]
    }
  }
  
  occ_estimates<-best_model %>% tidy.PGOcc(conf.int = T,conf.level = 0.95, HDI=T) %>% mutate(Rhat=best_model$rhat$beta)
  det_estimates<-best_model %>% tidy.PGOcc(conf.int = T, conf.level = 0.95, HDI=T, det=T) %>% mutate(Rhat=best_model$rhat$alpha)
  
  ret<-list(model=best_model,
            WAIC=waicOcc(best_model),
            occ_estimates=occ_estimates,
            det_estimates=det_estimates)
  return(list(ret))
}

get_best_model<-function(mods){
  best_model<-mods[[as.numeric(as_tibble(t(sapply(mods, `[[`, 2))) %>%
                                 filter(!is.infinite(WAIC) & !is.na(WAIC)) %>%
                                 rownames_to_column("id") %>%
                                 slice(which.min(WAIC)) %>% pull(id))]]
  best_model$WAIC<-t(best_model$WAIC) %>% as_tibble
  return(best_model)
}

#####

option_list <- list(
  make_option(c("-s", "--species"), action = 'store', type="character", default=NA, help="Which species do you want to run?",
              metavar="character")
)

opt_parser = OptionParser(option_list=option_list)
args = parse_args(opt_parser)

if(is.null(args$species)){
  print_help(opt_parser)
  stop("Please specify the species you want to run the model for...", call.=FALSE)
}

#####

setwd("/Volumes/T7 Shield/1. SDM.obsessive.edition/")

species<-args$s
#species<-"Formicarius analis"

species<-auk::ebird_species(species, "all") %>%
  select(scientific_name, common_name, species_code) %>% as.data.frame()

already_finished<-str_extract(list.files(file.path(getwd(), "results", species$species_code, "eBird", "occupancy")), "[0-9]+$")
start<-ifelse(length(already_finished)>0, max(as.numeric(already_finished)), 1)

for(it in start:its){
  results_folder<-file.path(getwd(), "results", species$species_code, "eBird", "occupancy", paste0("occ_", it))
  
  if(!dir.exists(results_folder)){
    dir.create(results_folder, recursive = T)
  }
  
  # Let's get the checklist and covariate data loaded
  
  env_variables<-read_csv("data/environmental-variables_checklists.csv") %>%
    group_by(checklist_id) %>%
    slice_head() %>% ungroup
  
  checklists<-read_csv(paste0("data/checklists-zf_CICRA_relOct-2024-", species$species_code, ".csv")) %>%
    filter(source=="eBird") %>%
    inner_join(env_variables, by = "checklist_id") %>%
    mutate(species_observed = as.integer(species_observed)) %>%
    filter(complete.cases(.)) %>%
    mutate(protocol_type=factor(protocol_type))
  
  # This weeds out predictors that have zero entropy (e.g. all the values are the same)
  
  all_preds<-env_variables %>% ungroup %>%
    mutate(source=ifelse(grepl("^G", checklist_id), "eBird", "Audio")) %>%
    filter(source=="eBird") %>%
    select(-checklist_id, -source) %>%
    select(contains("_1k")) %>%
    summarise(across(everything(), \(x) sd(x, na.rm=T))) %>%
    pivot_longer(cols = everything(), names_to = "pred", values_to = "sd") %>%
    filter(sd!=0 & !grepl("elevation", pred)) %>%
    pull(pred)
  
  # These are the covariates (occurrence and detection that we're going to start with)
  
  site_covs<-remove_correlated_vars(env_variables %>% select(!!all_preds)) %>% colnames
  det_covs<-c("day_of_year", "hours_of_day", "effort_hours", "protocol_type", "effort_distance_km", "number_observers")
  
  while(!file.exists(file.path(results_folder, "occupancy-plot.tif"))){
    try({
      
      # How informative do the predictors in the univariate models need to be to be included in the multivariate model?
      n_informative_occ<-0
      n_informative_det<-0
      cutoff<-0.95
      
      while(n_informative_occ==0 | n_informative_det==0){
        message("Cutoff is currently: ", cutoff)
        
        # Filter checklists to ensure evenness in survey effort
        
        checklists_filtered <- filter(checklists,
                                      effort_hours <= 5,
                                      effort_distance_km <= 1,
                                      number_observers <= 5) %>%
          filter(!is.na(observation_count))
        
        # Lump crops and built
        checklists_filtered$pland_c11_built_area_300m<-checklists_filtered$pland_c09_crops_300m+checklists_filtered$pland_c11_built_area_300m
        checklists_filtered$pland_c11_built_area_1k<-checklists_filtered$pland_c09_crops_1k+checklists_filtered$pland_c11_built_area_1k
        
        occ <- filter_repeat_visits(checklists_filtered,
                                    min_obs = 2, max_obs = 10,
                                    annual_closure = F,
                                    n_days=90,
                                    date_var = "observation_date",
                                    site_vars = c("locality_id", "observer_id"), ll_digits = 2)
        
        occ_wide <- format_unmarked_occu(occ, 
                                         site_id = "site",
                                         response = "species_observed",
                                         site_covs = c("n_observations", "latitude", "longitude", site_covs),
                                         obs_covs = det_covs)
        
        occ_wide<-occ_wide %>% mutate(obs=rowSums(across(starts_with("y.")), na.rm=T)>0)
        occ_wide<-occ_wide %>% mutate(source=if_else(grepl("^L[0-9]", site), "eBird", "Audio"))
        
        # for compatibility with eBird data later on
        occ_wide$day_of_year <- 1
        
        p<-0.8
        
        # Sample built and non-built areas separately
        occ_ss <- grid_sample_stratified(occ_wide %>% filter(pland_c11_built_area_1k<=median(occ_wide$pland_c11_built_area_1k)),
                               res = c(as.numeric(grid_size), as.numeric(grid_size), 1),
                               by_year = F, case_control = T, obs_column = "obs")
        
        occ_ss<-occ_ss %>% bind_rows(
          tryCatch({
            grid_sample_stratified(occ_wide %>% filter(pland_c11_built_area_1k>median(occ_wide$pland_c11_built_area_1k)),
                                   res = c(as.numeric(grid_size), as.numeric(grid_size), 1),
                                   by_year = F, case_control = T, obs_column = "obs")
          }, error = function(msg){
            grid_sample_stratified(occ_wide %>% filter(pland_c11_built_area_1k>median(occ_wide$pland_c11_built_area_1k)),
                                   res = c(as.numeric(grid_size), as.numeric(grid_size), 1),
                                   by_year = F, case_control = F)
          }))

        occ_ss<-bind_rows(occ_ss %>% filter(obs & pland_c11_built_area_1k<=median(occ_ss$pland_c11_built_area_1k)) %>% slice_sample(prop=p),
                          occ_ss %>% filter(!obs & pland_c11_built_area_1k<=median(occ_ss$pland_c11_built_area_1k)) %>% slice_sample(prop=p),
                          occ_ss %>% filter(obs & pland_c11_built_area_1k>median(occ_ss$pland_c11_built_area_1k)) %>% slice_sample(prop=p),
                          occ_ss %>% filter(!obs & pland_c11_built_area_1k>median(occ_ss$pland_c11_built_area_1k)) %>% slice_sample(prop=p)) %>%
          group_by(latitude, longitude) %>% slice_head(n = 1) %>% ungroup()
        
        occ_wide %>% filter(!(site %in% occ_ss$site)) %>% write_csv(file.path(results_folder, "unused_sites.csv"))
        occ_ss %>% write_csv(file.path(results_folder, "model_data.csv"))
        make.data.ss(site_covs, det_covs, occ_ss) %>% saveRDS(file.path(results_folder, "model_data_formatted.RDS"))
        
        univariate_model_output<-univariate_model_tests(site_covs, det_covs, wide_data=occ_ss, cutoff = cutoff)
        
        n_informative_occ<-length(univariate_model_output$occ_preds %>% filter(term != "(Intercept)") %>%
                                    mutate(informative = !(0 >= lower & 0 <= upper)) %>%
                                    filter(informative==T) %>%
                                    pull(term))
        n_informative_det<-length(univariate_model_output$det_preds %>% filter(term != "(Intercept)") %>%
                                    mutate(informative = !(0 >= lower & 0 <= upper)) %>%
                                    filter(informative==T) %>%
                                    pull(term))
        if(n_informative_occ==0 | n_informative_det==0){
          cutoff<<-cutoff-0.05
        }
      }
      
      message("Performing model selection using the following occupancy predictor set:")
      message(make.occ.formula(univariate_model_output$occ_preds %>% filter(term != "(Intercept)") %>%
                         mutate(informative = !(0 >= lower & 0 <= upper)) %>%
                         filter(informative==T) %>%
                         pull(term)))
      message("Performing model selection using the following detection predictor set:")
      message(make.det.formula(univariate_model_output$det_preds %>% filter(term != "(Intercept)") %>%
                         mutate(informative = !(0 >= lower & 0 <= upper)) %>%
                         filter(informative==T) %>%
                         pull(term)))
      
      models2<-multivariate_model_tests2(univariate_model_output, site_covs, det_covs, occ_ss)
      
      best_model<-get_best_model(models2)
      
      summary(best_model$model)
      
      # Save remaining model data
      
      best_model$model %>%
        saveRDS(file.path(results_folder, "occupancy_model.RDS"))
      
      best_model$WAIC %>%
        write_csv(file.path(results_folder, "model_performance.csv"))
      
      best_model$occ_estimates %>%
        write_csv(file.path(results_folder, "model_estimates_occ.csv"))
      
      best_model$det_estimates %>%
        write_csv(file.path(results_folder, "model_estimates_det.csv"))
      
      ppc.out <- ppcOcc(best_model$model, fit.stat = 'freeman-tukey', group = 1)
      
      plot_ppc<-function(x, filename){
        ppc.df <- data.frame(fit = x$fit.y, 
                             fit.rep = x$fit.y.rep, 
                             color = 'lightskyblue1')
        ppc.df$color[ppc.df$fit.rep > ppc.df$fit] <- 'lightsalmon'
        tiff(filename, units="px", width=2400, height=1600*2, res=300, pointsize = 12)
        plot(ppc.df$fit, ppc.df$fit.rep, bg = ppc.df$color, pch = 21, 
             ylab = 'Fit', xlab = 'True')
        lines(ppc.df$fit, ppc.df$fit, col = 'black')
        dev.off()
      }
      
      plot_diff_fit<-function(x, filename){
        tiff(filename, units="px", width=2400, height=1600*2, res=300, pointsize = 12)
        diff.fit <- x$fit.y.rep.group.quants[3, ] - x$fit.y.group.quants[3, ]
        plot(diff.fit, pch = 19, xlab = 'Site ID', ylab = 'Replicate - True Discrepancy')
        dev.off()
      }
      
      ppc.out %>% saveRDS(file.path(results_folder, "occupancy-model_gof_tukey.rds"))
      
      plot_ppc(ppc.out, file.path(results_folder, "ppc_plot_tukey.tiff"))
      plot_diff_fit(ppc.out, file.path(results_folder, "ppc_diff_fit_tukey.tiff"))
      
      ppc.out_2 <- ppcOcc(best_model$model, fit.stat = 'chi-square', group = 1)
      
      plot_ppc(ppc.out_2, file.path(results_folder, "ppc_plot_chisq.tiff"))
      plot_diff_fit(ppc.out_2, file.path(results_folder, "ppc_diff_fit_chisq.tiff"))
      
      ppc.out_2 %>% saveRDS(file.path(results_folder, "occupancy-model_gof_chisq.rds"))
      
      data.ss<-readRDS(file.path(results_folder, "model_data_formatted.RDS"))
      
      inits <- list(alpha = 0, 
                    beta = 0, 
                    z = apply(data.ss$y, 1, max, na.rm = TRUE))
      
      priors <- list(alpha.normal = list(mean = 0, var = 2.72), 
                     beta.normal = list(mean = 0, var = 2.72))
      
      out.null <- PGOcc(occ.formula = ~ 1,
                        det.formula = ~ 1, 
                        data = data.ss, 
                        inits = inits, 
                        n.samples = n.samples, 
                        priors = priors, 
                        n.omp.threads = 1, 
                        verbose = FALSE, 
                        n.burn = n.burn, 
                        n.thin = n.thin, 
                        n.chains = n.chains,
                        k.fold = 3,
                        k.fold.threads = 3)
      
      bind_rows(waicOcc(out.null), waicOcc(best_model$model)) %>% mutate(model=c("null", "eBird")) %>%
        arrange(WAIC) %>%
        write_csv(file.path(results_folder, "model_performance.csv"))
      
      out.k.fold.eBird <- PGOcc(occ.formula = make.occ.formula(best_model$occ_estimates$term),
                                det.formula = make.det.formula(best_model$det_estimates$term), 
                                data = data.ss, 
                                inits = inits, 
                                n.samples = n.samples, 
                                priors = priors, 
                                n.omp.threads = 1, 
                                verbose = TRUE, 
                                n.report = 1000, 
                                n.burn = n.burn, 
                                n.thin = n.thin, 
                                n.chains = n.chains,
                                k.fold = 3,
                                k.fold.threads = 3, k.fold.only = T)
      
      tibble(model="eBird", k.fold.deviance=out.k.fold.eBird$k.fold.deviance) %>%
        bind_rows(tibble(model="null", k.fold.deviance=out.null$k.fold.deviance)) %>%
        write_csv(file.path(results_folder, "k.fold.deviance.csv"))
      
      # load gis data for making maps
      
      map_proj <- 32719
      pred_grid <- read_csv("data/environmental-variables_prediction-grid.csv")
      
      pred_grid$pland_c11_built_area_300m<-pred_grid$pland_c09_crops_300m+pred_grid$pland_c11_built_area_300m
      pred_grid$pland_c11_built_area_1k<-pred_grid$pland_c09_crops_1k+pred_grid$pland_c11_built_area_1k
      
      study_area <- read_sf("gis/gis-data.gpkg", "study_area") %>%
        st_simplify(preserveTopology = T, dTolerance = 500) %>%
        st_transform(crs = 32719) %>% 
        st_geometry()
      r <- rast("data/prediction-grid.tif") %>% 
        # this second rast() call removes all the values from the raster template
        rast() %>%
        crop(study_area %>%
               st_simplify(preserveTopology = T, dTolerance = 500) %>%
               st_transform(crs = map_proj))
      crs <- st_crs(r)
      study_area <- study_area %>%
        st_transform(crs = crs)
      ne_land <- read_sf("gis/gis-data.gpkg", "ne_land") %>% 
        st_transform(crs = crs) %>% 
        st_geometry()
      # state boundaries for Peru
      ne_states <- st_read("gis/per_admbnda_adm1_ign_20200714.shp") %>%
        filter(ADM1_ES=="Madre de Dios") %>%
        select(ADM1_ES, geometry) %>%
        st_transform(crs = map_proj)
      
      if(ncol(best_model$model$X)==1 && colnames(best_model$model$X)=="(Intercept)"){
        X.0<-tibble("(Intercept)"=rep(1, nrow(pred_grid))) %>% as.data.frame()
      } else{
        pred_grid_scaled<-pred_grid %>%
          select(cell_id, x, y) %>%
          bind_cols(recipe(formula(paste0("~", paste(terms_format(best_model$occ_estimates$term), collapse = " + "))), data = occ_ss) %>%
                      step_poly(matches(terms_format(best_model$occ_estimates$term[grepl("\\^2\\)", best_model$occ_estimates$term)]))) %>%
                      step_center(matches(terms_format(best_model$occ_estimates$term))) %>%
                      step_scale(matches(terms_format(best_model$occ_estimates$term))) %>%
                      prep(occ_ss) %>%
                      bake(pred_grid) %>%
                      rename_with(~paste0("scale(", .x, ")", recycle0 = TRUE), -contains("_poly_2")) %>%
                      rename_with(~paste0("I(scale(", str_remove(.x, "_poly_2"), ")^2)", recycle0 = TRUE), contains("_poly_2")) %>%
                      rename_with(~str_remove(.x, "_poly_1")) %>%
                      mutate("(Intercept)"=1) %>%
                      select(!!best_model$occ_estimates$term))
        
        X.0<-pred_grid_scaled %>% select(!!best_model$occ_estimates$term) %>% as.data.frame()
      }
      
      shards<-split(X.0, (seq(nrow(X.0))-1) %/% 10000)
      names(shards)<-seq_along(names(shards))
      
      pred_occ_shards<-plyr::ldply(shards, function(x){
        cat(paste0("Working on rows ", as.numeric(rownames(x))[which(as.numeric(rownames(x))==min(as.numeric(rownames(x))))],
                   " through ", as.numeric(rownames(x))[which(as.numeric(rownames(x))==max(as.numeric(rownames(x))))], "...\n"))
        out.pred_shard<- predict(best_model$model, as.matrix(x))
        pred_occ_shard<-data.frame(occ_prob = apply(out.pred_shard$psi.0.samples, 2, mean),
                                   occ_se = apply(out.pred_shard$psi.0.samples, 2, sd))
        return(pred_occ_shard)
      })
      
      pred_occ<-pred_grid %>% bind_cols(pred_occ_shards %>% select(-.id))
      rm(pred_occ_shards)
      
      saveRDS(pred_occ, file.path(results_folder, "occupancy-model_predictions.rds"))
      
      # Make summary plots
      
      r_pred <- pred_occ %>%
        # convert to spatial features
        st_as_sf(coords = c("x", "y"), crs = crs) %>%
        rasterize(r, field = c("occ_prob", "occ_se")) %>%
        setNames(c("occ_prob", "occ_se")) %>%
        mask(vect(study_area %>%
                    st_simplify(preserveTopology = T, dTolerance = 500) %>%
                    st_transform(crs = crs)))
      
      writeRaster(r_pred[["occ_prob"]], 
                  filename = file.path(results_folder, "occupancy-model_prob.tif"),
                  overwrite = TRUE)
      writeRaster(r_pred[["occ_se"]], 
                  filename = file.path(results_folder, "occupancy-model_se.tif"), 
                  overwrite = TRUE)
      
      tiff(file.path(results_folder, "occupancy-plot.tif"), units="px", width=2400, height=1600*2, res=300, pointsize = 12)
      par(mfrow = c(2, 1))
      for (nm in names(r_pred)) {
        r_plot <- r_pred[[nm]]
        
        par(mar = c(3.5, 0.25, 0.25, 0.4))
        # set up plot area
        plot(study_area, col = NA, border = NA)
        plot(ne_land, col = "#dddddd", border = "#888888", lwd = 0.5, add = TRUE)
        plot(ne_states, col = "white", border = "black", lwd = 0.5, add = TRUE)
        
        # occupancy probability or standard error
        if (nm == "occ_prob") {
          title <- paste(species$common_name, "Occupancy Probability", sep=" ")
          brks <- seq(0, 1, length.out = 21)
          lbl_brks <- seq(0, 1, length.out = 11) %>% 
            round(2)
        } else {
          title <- paste(species$common_name, "Occupancy Uncertainty (SE)", sep=" ")
          mx <- ceiling(1000 * as.numeric(global(r_plot, "max", na.rm=T))) / 1000
          brks <- seq(0, mx, length.out = 21)
          lbl_brks <- seq(0, mx, length.out = 11) %>% 
            round(2)
        }
        pal <- ebirdst_palettes(length(brks) - 1, type = "weekly")
        
        # borders
        plot(ne_land, col = "#dddddd", lwd = 1.5, add = TRUE)
        plot(ne_states %>% st_transform(crs(ne_land)), col = "white", lwd = 0.75, add = TRUE)
        box()
        plot(r_plot, 
             col = pal, breaks = brks, 
             maxpixels = ncell(r_plot),
             legend = FALSE, add = TRUE)
        plot(study_area %>%
               st_simplify(preserveTopology = T, dTolerance = 500) %>%
               st_transform(crs = crs), col = NA, border="black", lwd = 0.75, add = TRUE)
        
        # legend
        par(new = TRUE, mar = c(0, 0, 0, 0))
        image.plot(zlim = range(brks), legend.only = TRUE, 
                   breaks = brks, col = pal,
                   smallplot = c(0.25, 0.75, 0.06, 0.09),
                   horizontal = TRUE,
                   axis.args = list(at = lbl_brks, labels = lbl_brks,
                                    fg = "black", col.axis = "black",
                                    cex.axis = 0.75, lwd.ticks = 0.5,
                                    padj = -1.5),
                   legend.args = list(text = title,
                                      side = 3, col = "black",
                                      cex = 1, line = 0))
      }
      dev.off()
    })
  }
}