library(tidyverse)
library(lme4)
library(jagsUI)

setwd("")

rm(list = ls())

# read in data
data <- read.csv('simple_data.csv')
data$genus <- unlist(lapply(strsplit(data$Species, "_"), function(x) x[1]))

#------------------------------------------------------------------------------------
# analysis for environmental variables
sc <- function(x) (x - mean(x, na.rm = T)) / sd(x, na.rm = T)

# function for BACI analysis
baci.analysis <- function(var, var_sd=NULL, log) {
  
  baci <- data[, c("genus", "Species", "BinaryThreat", var)] %>%
    rename(Before = var[1],
           After = var[2]) %>%
    pivot_longer(cols = c("Before", "After"), names_to = "BA", values_to = "response") 
  
  if(!is.null(var_sd)) {
    baci_sd <- data[, c("genus", "Species", "BinaryThreat", var_sd)] %>%
      rename(Before = var_sd[1],
             After = var_sd[2]) %>%
      pivot_longer(cols = c("Before", "After"), names_to = "BA", values_to = "response_sd") 
    
    baci <- full_join(baci, baci_sd)
  }
  
  if(is.null(var_sd)) {
    baci$response_sd <- rep(0, nrow(baci))
  }
  
  baci$response_var <- baci$response_sd^2
  if(log == TRUE) {
    baci$response_var <- baci$response_var / baci$response^2   # Taylor approximation
    baci$response <- log(baci$response)
  }
  
  # scaline  
  baci$response <- sc(baci$response)
  
  # variance accounting for the scaling
  baci$response_var <- (1/sd(baci$response, na.rm = T))^2 * baci$response_var
  
  # set factor levels so that cytrid-nonimpacted + Before will be the baseline for analysis
  baci$BA <- factor(baci$BA, levels = c("Before", "After"))
  baci$BinaryThreat <- factor(baci$BinaryThreat, levels = c("chytrid-nonimpacted", "chytrid-impacted"))
  
  p1 <<- ggplot(baci, aes(y = response, x = BA)) +
    #  geom_violin() +
    geom_line(aes(group = Species), colour = "grey") +
    geom_point(aes(colour = genus), size = 3) +
    facet_wrap(~ BinaryThreat) +
    xlab("") +
    ylab(var[1]) +
    labs(colour = "Genus", title = var[1]) +
    theme_classic() +
    theme(legend.position = "bottom")
  
  # jags analysis
  # data
  N <- baci$response
  VN <- baci$response_var
  threat <- ifelse(baci$BinaryThreat == "chytrid-impacted", 1, 0)
  ba <- ifelse(baci$BA == "After", 1, 0)
  int <- ifelse(baci$BinaryThreat == "chytrid-impacted" & baci$BA == "After", 1, 0)
  species <- as.numeric(factor(baci$Species))
  N.obs <- length(N)
  N.spp <- max(species)
  genus <- as.numeric(apply(table(species, baci$genus), 1, function(x) which(x > 0)))
  N.genus <- max(genus)
  
  model <- "model
  {
    for(i in 1:N.obs) {
        N[i] ~ dnorm(mu[i], prec[i])      
        prec[i] <- 1 / (V + VN[i])        # variances are additive
        mu[i] <- b[1] + b[2]*threat[i] + b[3]*ba[i] + b[4]*int[i] + b.spp[species[i]]
    }

    for(i in 1:4) {
      b[i] ~ dnorm(0, 0.01)
    }
    
    for(i in 1:N.spp) {
      b.spp[i] ~ dnorm(mu.genus[genus[i]], tau.spp)
    }
    
    for(i in 1:N.genus) {
      mu.genus[i] ~ dnorm(0, tau.genus)
    }
    V <- sigma * sigma        
    sigma ~ dunif(0, 100)     
    tau.spp <- 1/(sigma.spp*sigma.spp)
    sigma.spp ~ dunif(0, 100)
    tau.genus <- 1/(sigma.genus*sigma.genus)
    sigma.genus ~ dunif(0, 100)

}"
  
  # write model
  write(model, "niche model.txt")
  
  mod <- jags(model.file = "niche model.txt",
              data = list(N = N, VN = VN, threat = threat, ba = ba, int = int,
                          species = species, genus = genus, N.obs = N.obs,
                          N.spp = N.spp, N.genus = N.genus),
              parameters.to.save = c("b", "sigma", "sigma.spp", "sigma.genus"),
              n.chains = 3,
              n.iter = 50000,
              n.burnin = 10000,
              n.thin = 1,
              parallel = T)
  
  mod.sum <- mod$summary
  mod.sum[, c(1, 2, 3, 7, 8, 9)]
  
  # the results
  out <- data.frame(param = mod.sum[3:4, 1], lcl = mod.sum[3:4, 3], ucl = mod.sum[3:4, 7], rhat = mod.sum[3:4, 8])
  out$var <- c("Change in chytrid non-impacted", "Change in chytrid impacted relative to non-impacted")
  rownames(out) <- NULL
  return(out)
}

#---------------------------------------------------------------------------------
set.seed(8763)
niche <- baci.analysis(var = c("Full_Niche", "Contemporary_Niche"), 
                       var_sd = c("Full_Niche_SD", "Contempory_Niche_SD"), log = TRUE)

eoo <- baci.analysis(var = c("Full_EOO", "Contemporary_EOO"), log = TRUE)


bio1 <- baci.analysis(var = c("Bio1_Mean_Null", "Bio1_Mean_Con"),
                      var_sd = c("Bio1_Mean_Null_SD", "Bio1_Mean_Con_SD"), log = FALSE)

bio9 <- baci.analysis(var = c("Bio9_Mean_Null", "Bio9_Mean_Con"),
                      var_sd = c("Bio9_Mean_Null_SD", "Bio9_Mean_Con_SD"), log = FALSE)

bio15 <- baci.analysis(var = c("Bio15_Mean_Null", "Bio15_Mean_Con"),
                       var_sd = c("Bio15_Mean_Null_SD", "Bio15_Mean_Con_SD"), log = FALSE)

elev <- baci.analysis(var = c("Elevation_Mean_Null", "Elevation_Mean_Con"),
                      var_sd = c("Elevation_Mean_Null_SD", "Elevation_Mean_Con_SD"), log = FALSE)


niche$type <- "Niche breadth"
eoo$type <- "EOO"
bio1$type <- "Annual mean temperature"
bio9$type <- "Mean temperature driest quarter"
bio15$type <- "Precipitation seasonality"
elev$type <- "Elevation"

out <- bind_rows(niche, eoo, bio1, bio9, bio15, elev)
out

# set factor levels for plotting
out$type <- factor(out$type, levels = c("EOO", "Niche breadth", "Mean temperature driest quarter",
                                        "Annual mean temperature", "Precipitation seasonality",
                                        "Elevation"))

plot.var <- c("Change in chytrid non-impacted", "Change in chytrid impacted relative to non-impacted")

sub.out <- filter(out, var %in% plot.var)
sub.out$var <- factor(sub.out$var, levels = plot.var)

pdf("Figure 1.pdf")

ggplot(sub.out, aes(y = param, x = type)) +
  geom_pointrange(aes(ymin = lcl, ymax = ucl), size = 1) +
  facet_wrap(~ var, ncol = 2) +
  geom_hline(yintercept = 0) +
  xlab("") +
  ylab("Change in niche variable") +
  theme_bw() +
  theme(panel.grid.major = element_blank(), panel.grid.minor = element_blank(),
        panel.background = element_blank(),
        axis.text.x = element_text(angle = 35, vjust = 1, hjust=1, size = 10),
        axis.text.y = element_text(size=10),
        axis.title.y = element_text(size=10),
        strip.text.x = element_text(size = 9))

dev.off()

write.csv(out, "Niche variables.csv", row.names = F)

#-------------------------------------------------------------------------------
# niche similarity
# set factor levels so the chytrid non-impacted is the baseline reference
data$BinaryThreat <- factor(data$BinaryThreat, levels = c("chytrid-nonimpacted", "chytrid-impacted"))

# function to fit model
niche.sim <- function(var, var_sd) {
  N <- data[, var]
  VN <- data[, var_sd]
  threat <- ifelse(data$BinaryThreat == "chytrid-impacted", 1, 0)
  genus <- as.numeric(factor(data$genus))
  N.obs <- length(N)
  N.genus <- max(genus)
  
  model <- "model
  {
    for(i in 1:N.obs) {
        N[i] ~ dnorm(mu[i], prec[i])      
        prec[i] <- 1 / (V + VN[i])        # variances are additive
        mu[i] <- b[1] + b[2]*threat[i] + b.genus[genus[i]]
    }

    for(i in 1:2) {
      b[i] ~ dnorm(0, 0.01)
    }
    
    for(i in 1:N.genus) {
      b.genus[i] ~ dnorm(0, tau.genus)
    }
    V <- sigma * sigma        
    sigma ~ dunif(0, 100)     
    tau.genus <- 1/(sigma.genus*sigma.genus)
    sigma.genus ~ dunif(0, 100)

}"
  
  # write model
  write(model, "niche model.txt")
  
  mod <- jags(model.file = "niche model.txt",
              data = list(N = N, VN = VN, threat = threat, genus = genus, 
                          N.obs = N.obs, N.genus = N.genus),
              parameters.to.save = c("b", "sigma", "sigma.genus"),
              n.chains = 3,
              n.iter = 50000,
              n.burnin = 10000,
              n.thin = 1,
              parallel = T)
  
  mod.sum <- mod$summary
  mod.sum[, c(1, 2, 3, 7, 8, 9)]
  
  # the results
  out.sim <- data.frame(param = mod.sum[2, 1], lcl = mod.sum[2, 3], ucl = mod.sum[2, 7], rhat = mod.sum[2, 8])
  out.sim$var <- var
  rownames(out.sim) <- NULL
  return(out.sim)
}

out.sim <- bind_rows(niche.sim(var = "Btotal", var_sd = "BtotalSD"),
                     niche.sim(var = "Brich", var_sd = "BrichSD"),
                     niche.sim(var = "Brepl", var_sd = "BreplSD"))

out.sim

pdf("Figure 2.pdf")

ggplot(out.sim, aes(y = param, x = var)) +
  geom_pointrange(aes(ymin = lcl, ymax = ucl), size = 1) +
  geom_hline(yintercept = 0) +
  xlab("") +
  ylab("Change in niche similarity metric") +
  theme_bw() +
  theme(panel.grid.major = element_blank(), panel.grid.minor = element_blank(),
        panel.background = element_blank(),
        axis.text.x = element_text(angle = 35, vjust = 1, hjust=1, size = 10),
        axis.text.y = element_text(size=10),
        axis.title.y = element_text(size=10),
        strip.text.x = element_text(size = 9))

dev.off()

write.csv(out.sim, "Niche similarity.csv", row.names = F)
