#Functions to define correlations between Beals, SYNCSA, ENV
# W - Sp. composition
# X - fuzzy weighted
# Y - Beals
# B - traits
# E - Environment


library(tidyverse)
library(SYNCSA)
library(vegan)
library(abind)
library(ade4)
library(energy)

#### Function 1b - CorXY bootstrap####


get.corXY.bootstrap <- function(comm, traits, trait.sel="all", bootstrap=199){
  
  if(identical(trait.sel, "all")) {trait.sel <- 1:ncol(traits)}
  ii <- trait.sel
  lab.tmp <- paste(ii, collapse="_")
  ### delete potential missing species
  if(any(colSums(comm)==0)){
    empty <- which(colSums(comm)==0)
    traits <- traits[-empty,]
    comm <- comm[,-empty]
  }
  #Calculate correlations on full data
  n.sites <- nrow(comm)
  n.species <- ncol(comm)
  
  ## caution
  ## ALL columns with only 0-1 values are AUTOMATICALLY considered as asym.bin sensu FD:gowdis
  ##get all columns with binary variables
  ## CAUTION - function below doesn't work as expected
#  binary.traits <- which(apply(traits[,ii,drop=F], MARGIN=2, function(x)( all(na.omit(x) %in% 0:1) ))==T)
  
  syn.out.tmp <- matrix.x(comm=comm, traits=traits[,ii,drop=F], scale=T)$matrix.X #, asym.bin=binary.traits
  W.beals <- as.data.frame(beals(comm, include=T, type=2))
  # permute traits 
  ### Create vector of species to resample from, which exclude species with NAs
  traits.ii <- traits[,ii, drop=F]
  ids <- 1:nrow(traits.ii)
  id.nas <- ids[which(rowSums(is.na(traits.ii))==ncol(traits.ii))]
  toresample <- ids[!ids %in% id.nas]
  ## create list of indices of permuted traits, but without shuffling traits with missing data
  ## the length of the list is = to the number of bootstraps
  index.traits <- lapply(1:(bootstrap+1), function(x){
    out <- rep(NA, nrow(traits.ii))
    out[id.nas] <- id.nas
    out[is.na(out)] <- sample(toresample, replace=F)
    return(out)
  })
  #index.traits <- lapply(1:(bootstrap+1), function(x){sample(1:nrow(traits), replace=F)})
  syn.out.perm.tmp <- matrix.x(comm=comm, traits=traits[index.traits[[bootstrap+1]],ii,drop=F], 
                               scale=T)$matrix.X #, asym.bin=binary.traits
  
  (RD.tmp <- dcor(W.beals, as.data.frame(syn.out.tmp))^2)
  (RD.perm.tmp <- dcor(W.beals, as.data.frame(syn.out.perm.tmp))^2)
  corXY <- NULL
  corXY <- data.frame(Trait.comb=lab.tmp, bootstr.n=0, Test="RD", 
                            Coef.obs=RD.tmp, Coef.perm=RD.perm.tmp)
  
  index.bootstr <- lapply(1:bootstrap, function(x){sample(1:n.sites, replace=T)})
  for(b in 1:bootstrap){
    RD.tmp <- dcor(W.beals[index.bootstr[[b]],], as.data.frame(syn.out.tmp)[index.bootstr[[b]],])^2
    syn.out.perm.tmp <- matrix.x(comm=comm, traits=traits[index.traits[[b]],ii,drop=F], 
                                 scale=T)$matrix.X #, asym.bin=binary.traits
    RD.perm.tmp <- dcor(W.beals[index.bootstr[[b]],], as.data.frame(syn.out.perm.tmp)[index.bootstr[[b]],])^2
    corXY <- rbind(corXY,
                   data.frame(Trait.comb=lab.tmp, bootstr.n=b, Test="RD", 
                              Coef.obs=RD.tmp, Coef.perm=RD.perm.tmp))
    if(b %in% round(seq(1,bootstrap, length.out=10))){
      print(paste("trait", paste0(ii, collapse="_"), "perm", b, paste(colnames(traits[trait.sel]), collapse="+")))}
  }
  return(corXY)
}


## Cracked version of matrix.x from SYNCSA to account for asymmetric binary variables in gowdis
## and to avoid error propagation after gowdis, when species have NAs in traits

matrix.x <- function (comm, traits, scale = TRUE, ranks = TRUE, ord, notification = TRUE, 
                      ...) 
{
  comm <- as.matrix(comm)
  vartype <- var.type(traits)
  if (any(vartype == "n")) {
    stop("\n trait must contain only numeric, binary or ordinal variables \n")
  }
  if (missing(ord)) {
    for (i in 1:length(vartype)) {
      if (ranks & vartype[i] == "o") {
        traits[, i] <- rank(traits[, i], na.last = "keep")
      }
      traits[, i] <- as.numeric(traits[, i])
    }
    traits <- as.matrix(traits)
  }
  matrix.w <- sweep(comm, 1, rowSums(comm, na.rm = TRUE), "/")
  w.NA <- apply(matrix.w, 2, is.na)
  matrix.w[w.NA] <- 0
  if (notification) {
    if (any(w.NA)) {
      warning("Warning: NA in community data", call. = FALSE)
    }
  }
  x.NA <- apply(traits, 2, is.na)
  if (notification) {
    if (any(x.NA)) {
      warning("Warning: NA in traits matrix", call. = FALSE)
    }
  }
  if (scale) {
    dist.traits <- FD::gowdis(traits, ...)
    dist.traits <- as.matrix(dist.traits)
    dist.traits[which(rowSums(is.na(dist.traits))==(ncol(dist.traits)-1)),] <- NA   ### In case a species has NO trait info, replace the diag with NA to avoid propagation
    similar.traits <- 1 - as.matrix(dist.traits)
    matrix.traits <- 1/colSums(similar.traits, na.rm = TRUE)
    matrix.u <- sweep(similar.traits, 1, matrix.traits, "*")
  }
  else {
    dist.traits <- as.matrix(vegan::vegdist(traits, method = "euclidean", 
                                            diag = TRUE, upper = TRUE, na.rm = TRUE))
    similar.traits <- 1 - (dist.traits/max(dist.traits, na.rm = TRUE))
    matrix.traits <- 1/colSums(similar.traits, na.rm = TRUE)
    matrix.u <- sweep(similar.traits, 1, matrix.traits, "*")
  }
  u.NA <- apply(matrix.u, 2, is.na)
  if (notification) {
    if (any(u.NA)) {
      warning("Warning: NA in matrix U", call. = FALSE)
    }
  }
  matrix.u[u.NA] <- 0  
  matrix.X <- matrix.w %*% matrix.u
  matrix.X <- sweep(matrix.X, 1, rowSums(matrix.X), "/")  ## RE-SCALE to avoid problems propagating from NAs
  return(list(matrix.w = matrix.w, matrix.u = matrix.u, matrix.X = matrix.X))
}






##### NOT ACTIVELY MANTAINED BELOW ####


#### Function 1 - CorXY ####
get.corXY <- function(comm, traits, trait.sel="all", stat=c("mantel", "RV", "procrustes")){
  if(identical(trait.sel, "all")) {trait.sel <- 1:ncol(traits)}
  ii <- trait.sel
  lab.tmp <- paste(ii, collapse="_")
  ### delete potential missing species
  if(any(colSums(comm)==0)){
    empty <- which(colSums(comm)==0)
    traits <- traits[-empty,]
    comm <- comm[,-empty]
  }
  syn.out.tmp <- matrix.x(comm=comm, traits=traits[,ii,drop=F], scale=T)$matrix.X
  W.beals <- as.data.frame(beals(comm, include=T, type=2))
  corXY <- NULL
  if("mantel" %in% stat){
    W.beals.d <- dist(W.beals)
    mantel.tmp <- mantel(W.beals.d, dist(syn.out.tmp[]))
    corXY <- rbind(corXY, 
                   data.frame(Trait.comb=lab.tmp,  Test="Mantel", Coef=mantel.tmp$statistic, pvalue=mantel.tmp$signif))
  } 
  if("RV" %in% stat){
    RV.tmp <- RV.rtest(W.beals, as.data.frame(syn.out.tmp))
    corXY <- rbind(corXY,
                   data.frame(Trait.comb=lab.tmp,  Test="RV", Coef=RV.tmp$obs, pvalue=RV.tmp$pvalue))
  } 
  if("procrustes" %in% stat){
    prot.tmp <- protest(W.beals, syn.out.tmp)
    corXY <- rbind(corXY,
                   data.frame(Trait.comb=lab.tmp,  Test="Procrustes", Coef=prot.tmp$t0, pvalue=prot.tmp$signif))
  }
  return(corXY)
}





#### Function 2 - CorTE ####
get.corTE <- function(comm, traits, envir,  trait.sel="all", env.sel="all", stat=c("mantel", "RV")){
  if(identical(trait.sel, "all")) {trait.sel <- 1:ncol(traits)}
  if(identical(env.sel, "all" )) {env.sel <- 1:ncol(envir)}
  ii <- trait.sel
  lab.tmp <- paste(ii, collapse="_")
  ## delete potential missing species
  if(any(colSums(comm)==0)){
    empty <- which(colSums(comm)==0)
    traits <- traits[-empty,]
    comm <- comm[,-empty]
  }
  syn.out.tmp <- matrix.t(comm=comm, traits=traits, scale=T)$matrix.T
  ee <- env.sel
  lab.env <- paste(ee, collapse="_")
  corTE <- NULL
  if("mantel" %in% stat){
    mantel.tmp <- mantel(dist(envir[,ee, drop=F]), dist(syn.out.tmp[,ii,drop=F]))
    corTE <- rbind(corTE, 
                   data.frame(Trait.comb=lab.tmp, Env.comb=lab.env, Test="Mantel", Coef=mantel.tmp$statistic, pvalue=mantel.tmp$signif))
  } 
  if("RV" %in% stat){
    RV.tmp <- RV.rtest(as.data.frame(envir[,ee, drop=F]), as.data.frame(syn.out.tmp[,ii,drop=F]))
    corTE <- rbind(corTE,
                   data.frame(Trait.comb=lab.tmp, Env.comb=lab.env, Test="RV", Coef=RV.tmp$obs, pvalue=RV.tmp$pvalue))
  } 
  if("procrustes" %in% stat){
    prot.tmp <- protest(envir[,ee, drop=F], syn.out.tmp[,ii,drop=F])
    corTE <- rbind(corTE,
                   data.frame(Trait.comb=lab.tmp, Env.comb=lab.env, Test="Procrustes", Coef=prot.tmp$t0, pvalue=prot.tmp$signif))
  }
  return(corTE)
}


#### Function 3 - CorXE ####
get.corXE <- function(comm, traits, envir, trait.sel="all", env.sel="all", stat=c("mantel", "RV", "procrustes")){
  if(identical(trait.sel, "all")) {trait.sel <- 1:ncol(traits)}
  if(identical(env.sel, "all" )) {env.sel <- 1:ncol(envir)}
  ii <- trait.sel
  lab.tmp <- paste(ii, collapse="_")
  ### delete potential missing species
  if(any(colSums(comm)==0)){
    empty <- which(colSums(comm)==0)
    traits <- traits[-empty,]
    comm <- comm[,-empty]
  }
  syn.out.tmp <- matrix.x(comm=comm, traits=traits[,ii,drop=F], scale=T)$matrix.X
  ee <- env.sel
  lab.env <- paste(ee, collapse="_")
  corXE <- NULL
  if("mantel" %in% stat){
    mantel.tmp <- mantel(dist(envir[,ee, drop=F]), dist(syn.out.tmp[]))
    corXE <- rbind(corXE, 
                   data.frame(Trait.comb=lab.tmp, Env.comb=lab.env, Test="Mantel", Coef=mantel.tmp$statistic, pvalue=mantel.tmp$signif))
  } 
  if("RV" %in% stat){
    RV.tmp <- RV.rtest(as.data.frame(envir[,ee, drop=F]), as.data.frame(syn.out.tmp))
    corXE <- rbind(corXE,
                   data.frame(Trait.comb=lab.tmp, Env.comb=lab.env, Test="RV", Coef=RV.tmp$obs, pvalue=RV.tmp$pvalue))
  } 
  if("procrustes" %in% stat){
    prot.tmp <- protest(envir[,ee, drop=F], syn.out.tmp)
    corXE <- rbind(corXE,
                   data.frame(Trait.comb=lab.tmp, Env.comb=lab.env, Test="Procrustes", Coef=prot.tmp$t0, pvalue=prot.tmp$signif))
  }
  return(corXE)
}



### Get SES (both parametric and non parametric)
#obs.df = output from one of the get functions above
#perm.df = df traitCombination X permutations with the chosen statistic of correlation on permuteda data

get.SES <- function(obs.df, perm.df, stat="RV") {
  #nperm <- dim(perm.df)[2]
  SES.out <- as.data.frame.table(perm.df) %>% 
    rename(Trait.comb=Var1, perm=Var2, coef=Freq ) %>% 
    group_by(Trait.comb) %>% 
    left_join(obs.df %>% 
                filter(Test==stat) %>% 
                dplyr::select(-pvalue, -Test) %>% 
                mutate(Trait.comb=paste0("t", Trait.comb)) %>% 
                rename(obs=Coef), 
              by=c("Trait.comb")) %>% 
    summarize(q025 = quantile(coef, probs = 0.025), 
              q15 = quantile(coef, probs = 0.15865), 
              q50 = quantile(coef, .5),
              q84 = quantile(coef, 0.84135),
              q975 = quantile(coef, 0.975),
              mean.perm=mean(coef),
              sd.perm = sd(coef), 
              greater.than.obs = sum( coef >= obs )/n(), 
              smaller.than.obs = sum( coef <= obs )/n()) %>% 
    left_join(obs.df %>% 
                filter(Test==stat) %>% 
                dplyr::select(-pvalue, -Test) %>% 
                mutate(Trait.comb=paste0("t", Trait.comb)) %>% 
                rename(obs=Coef), 
              by=c("Trait.comb")) %>% 
    mutate(SES.np= ifelse(obs>=q50, (obs-q50)/(q84-q50), (obs-q50)/(q50-q15))) %>%
    mutate(SES=(obs-mean.perm)/sd.perm) %>% 
    arrange(desc(SES.np)) %>% 
    mutate(conf.m=mean.perm - 1.96*(sd.perm)) %>% #/sqrt(nperm))) %>% 
    mutate(conf.p=mean.perm + 1.96*(sd.perm)) #/sqrt(nperm)))
  return(SES.out)
}