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


require(tidyverse)
require(SYNCSA)
require(vegan)
require(abind)
require(ade4)
require(parallel)
require(doParallel)


source("99_HIDDEN_functions.R")


get.ci <- function(corXY){
  corXY %>% 
  filter(bootstr.n==0) %>% 
  dplyr::select(-bootstr.n) %>% 
  group_by(Trait.comb) %>% 
  slice(1) %>% 
  left_join(corXY %>% 
              filter(bootstr.n!=0) %>% 
              group_by(Trait.comb) %>% 
              summarize(q025=quantile(Coef.obs, 0.025),
                        q975=quantile(Coef.obs, 0.975),
                        greater.than.perm=mean(Coef.obs>Coef.perm), 
                        n=n())) %>% 
  #calculate significance using permuted correlations
  mutate(sign_plus=greater.than.perm>0.995) %>% 
  mutate(sign_minus=greater.than.perm<0.005) %>% 
  ungroup() %>% 
  mutate(Trait.comb=as.character(Trait.comb)) %>% 
  rowwise() %>% 
  mutate(ntraits= length(unlist(strsplit(Trait.comb, split = "_")))) %>% 
  ungroup() %>% 
  arrange(ntraits, Coef.obs)
}

get.best <- function(x, N){
  out <- x %>% 
    filter(ntraits==N) %>% 
    filter(sign_plus==TRUE) %>% 
    arrange(desc(Coef.obs)) %>% 
    slice(1) %>% 
    pull(Trait.comb)
  return(out)
}



Mesobromion <- function(species.path, traits.path, output, myfunction="get.corXY.bootstrap", 
                        combinations=c("all", "sequential"), start.round=NA, relax.round=NA, max.inter.t, 
                        chunkn, chunk.i, nperm=199, ncores){
  
  if(ncores>1){
    require(parallel)
    require(doParallel)
    cl <- makeForkCluster(ncores, outfile="" )
    registerDoParallel(cl)
    `%myinfix%` <- `%dopar%`
  } else {
    `%myinfix%` <- `%do%`
  }
  
  myfunction <- get(myfunction)
  ## calculate corr between species composition matrix and traits
  species <- read_delim(species.path, delim="\t") %>%
		as.data.frame() %>% 
    column_to_rownames("RELEVE_NR")
  traits <- read_delim(traits.path, delim="\t") %>%
		as.data.frame() #%>% 
### TEMPORARY FOR TESTING!
    #dplyr::select(1:6)
  
  
  traits <- traits %>% 
    column_to_rownames("species0") %>% 
    rename_all(.funs=~gsub(pattern=".mean$", replacement="", x=.))  %>% 
    mutate_if(~is.character(.), .funs=~as.factor(.))
#temporary    ### Use only a subset of traits
    #dplyr::select(LeafArea:Disp.unit.leng)
  
  if(combinations=="all") {
    ## create list of indices for each combination of traits up to a max number of interactions
    n.traits <- ncol(traits)
    allcomb.t <- NULL
    for(n.inter in 1:max.inter.t){
      allcomb.t <- c(allcomb.t, combn(1:n.traits, n.inter, simplify=F))
    }
    nall <- length(allcomb.t)
    
    ## Divide in chunks if requested
    if(chunkn>1 & !is.na(chunk.i)){
      print(paste("Test all Combo - divide in", chunkn, "chunks and run on chunk n.", chunk.i))
      indices <- 1:length(allcomb.t)
      chunks <- split(indices, sort(indices%%chunkn))
      allcomb.t <- allcomb.t[chunks[[chunk.i]]]
    } 
    print(paste("Running on", length(allcomb.t),"out of", nall, "combinations"))
    #names.t <- unlist(lapply(allcomb.t, paste, collapse="_"))
    
    print("Start main loop")
    corXY.output <- NULL
    corXY.output <- foreach(i=1:length(allcomb.t), .combine=rbind) %myinfix% {
      tt <- unlist(allcomb.t[i])
      corXY.output <- rbind(corXY.output, 
                            get.corXY.bootstrap(comm=species, trait=traits, 
                                                trait.sel=tt, bootstrap=nperm))
    }
    save(corXY.output, file = paste(output,"_", chunk.i, "_.RData", sep=""))
  }
  
  
  if(combinations=="sequential"){
   
    if(start.round>1){
      print(paste("Load data from previous round=", start.round-1))
      load(file = paste(output, "_round_", start.round-1, ".RData", sep=""))
      } else {corXY.ci <- NULL}
    for(nround in start.round:max.inter.t){
      corXY.output <- NULL
      ## select combination of traits based on best
      if(nround==1) {
        allcomb.t <- lapply(1:ncol(traits), function(x){return(x)})
        nall <- length(allcomb.t)
        }  else {
          if(nround >= relax.round){
            best <- corXY.ci %>% 
              filter(sign_plus==T) %>% 
              filter(ntraits==(nround-1)) %>% 
              arrange(desc(Coef.obs)) %>% 
              slice(1) %>% 
              pull(Trait.comb)
            
            new.best.row <- corXY.ci %>% 
              filter(Trait.comb==best) 
            upper <- new.best.row$q975
            lower <- new.best.row$q025
            print(paste("Assumptions relaxed - new best at round=", nround, "is trait combination", best))
          } else {
            best.row <- corXY.ci %>% 
              filter(Trait.comb==best) 
            upper <- best.row$q975
            lower <- best.row$q025
            }
          best.split <- as.numeric(unlist(strsplit(best, "_")))
          max.inter.ti <- nround-length(best.split)
          #n.traits <- ncol(traits)
          list.traits <- as.numeric(traits.sign.alone)
          list.traits <- list.traits[-which(list.traits %in% best.split)] ## list of remaining traits without best
          if(length(list.traits) < max.inter.ti) {
            save(corXY.output, file = paste(output, ".RData", sep=""))
            if(ncores>1) {stopCluster(cl)}
            stop("Tested all combo of significant traits")
          }
          allcomb.t <- NULL
          #for(n.inter in 1:max.inter.t){
          allcomb.t <- c(allcomb.t, combn(list.traits, max.inter.ti, simplify=F))
          #}
          #add best to all combo
          allcomb.t <- lapply(allcomb.t, function(x){return(c(best.split, x))})
          nall <- length(allcomb.t)
        }
      
      ## Divide in chunks if requested
#      if(chunkn>1){
#        print(paste("nround", nround, "--- divide in", chunkn, "chunks and run on chunk n=", chunk.i))
#        indices <- 1:length(allcomb.t)
#        chunks <- split(indices, sort(indices%%chunkn))
#        allcomb.t <- allcomb.t[chunks[[chunk.i]] ]
#      }  
      #names.t <- unlist(lapply(allcomb.t, paste, collapse="_"))
      
      ## start main foreach loop
      print(paste("Start main loop - round", nround))
      corXY.output <- foreach(ch=1:length(allcomb.t), .combine=rbind) %myinfix% {
        allcomb.i <- allcomb.t[ch]
        print(paste("Running on combination", ch,"out of", nall))
        
        for(i in 1:length(allcomb.i)){
          tt <- unlist(allcomb.i[i])
          corXY.output <- rbind(corXY.output, 
                              get.corXY.bootstrap(comm=species, trait=traits, 
                                                  trait.sel=tt, bootstrap=nperm))
        }
        return(corXY.output)
      }
      #Intermediate save
      #save(corXY.output, file = paste(output, "_", nround, ".RData", sep=""))
      
      print(paste("Summarize after round", nround, "and find best, significant trait combos"))
      corXY.ci <- corXY.ci %>% 
        bind_rows(get.ci(corXY.output))
      ##get best or new best
      if(nround==1) {
        traits.sign.alone <- corXY.ci %>% 
          filter(ntraits==1) %>% 
          filter(sign_plus==T) %>% 
          pull(Trait.comb)
        if(length(traits.sign.alone)==0) {
          stop("No trait is significant taken singularly")}
        print(paste("Traits significant taken singularly:", paste(traits.sign.alone, collapse=" ")))
        best <- get.best(corXY.ci, N=1)
        best.row <- corXY.ci %>% 
          filter(Trait.comb==best) 
        upper <- best.row$q975
        lower <- best.row$q025
        print(paste("new best at nround=", nround, "is trait combination", best))
      }
      
      if(nround>1){
       better <- corXY.ci %>% 
          filter(ntraits==nround) %>% 
          filter(q025>upper) %>% 
          arrange(desc(Coef.obs)) %>% 
          slice(1) %>% 
          pull(Trait.comb)
        
        if(length(better>0)){
          new.best.row <- corXY.ci %>% 
            #rowwise() %>% 
            #mutate(nmatching= sum(unlist(strsplit(Trait.comb, "_")) %in% 
            #                        unlist(strsplit(better, "_")),
            #                      na.rm=T)) %>% 
            #ungroup() %>% 
            #filter(nmatching==nround) %>% 
            filter(Trait.comb==better) 
          upper <- new.best.row$q975
          lower <- new.best.row$q025
          best <- better
          print(paste("new best at round=", nround, "is trait combination", best))
        } else {print(paste("no new best at round=", nround))}
      }
    print(paste("save intermediate results at round", nround))
    save(corXY.output, best, traits.sign.alone, corXY.ci, file = paste(output, "_round_",nround, ".RData", sep=""))
    }    
  }
  if(ncores>1){stopCluster(cl)} 
}