
library(ggplot2)
library(ggridges)
library(grid)
library(data.table)
library(reshape2)

library(Hmsc)
library(reshape2)
library(abind)
library(tidyr)
library(dplyr)
#library(berryFunctions)
#library(lemon)
#library(cowplot)
library(patchwork)
library(ggh4x)
library(ggpubr)

#setwd("~/1.Data/3.ModelRun")
load('test10908cent_yearPA.RDATA')
modelyear = list()
modelyear[[1]] = models[[1]]

load('test10908cent_yearAB.RDATA')
modelyear[[2]] = models[[2]]



hM = modelyear[[1]]
colnames(hM$X)
#exclude = 20

hM$studyDesign

## define some grouping  element
sites = data.frame(Site=hM$studyDesign$sample)



group=NULL
groupnames=NULL
start=1
na.ignore=FALSE

wb=NULL


# function to get random parameters
GetRandparam = function(hM, Lambdas, Etas){
  
  nr = hM$nr
  np = hM$np
  ns = hM$ns
  
  # random information from design 
  mat.random = cbind(1:dim(hM$studyDesign)[1], hM$studyDesign)
  rand_eff = colnames(hM$studyDesign)
  mat.rand = list()
  
  for(x in rand_eff) {
    mat.random = cbind(mat.random, value=1) 
    mat.random[,x] = paste0(x, mat.random[,x])
    mat.rand[[x]] = dcast(mat.random, 
                          as.formula(paste0(paste(colnames(mat.random)[1]), '~ ', x)), fill=0)
    mat.rand[[x]] = mat.rand[[x]][, -1]
  }
  
  lapply(mat.rand, dim) # to verify
  
  
  H = array(NA, dim=c(nrow(hM$X), length(rand_eff), ns))
  Lrand_list = list()
  for (level in seq_len(nr)) {
    Lambda = Lambdas[[level]]
    nf = dim(Lambda)[[1]]
    Eta = Etas[[level]]
    
    # random effect from HMSC
    Lrand = matrix(NA, ncol = ns, nrow = np[level])
    Lr.fact = lapply(1:ns, matrix, data= NA, nrow=np[level], ncol=nf)
    for (j in 1:ns){
      for (factor in 1:nf) {
        Lr.fact[[j]][, factor] = Eta[,factor]*Lambda[factor,j]
      }
      Lrand[, j] = rowSums(Lr.fact[[j]])
    }
    
    Lrand_list[[level]] = Lrand
    # getting H
    H[, level, ] = (as.matrix(mat.rand[[level]]) %*% Lrand_list[[level]])
    ### 3random effect list of matrices (960 occasions x 30 species)
  }
  return(list(H=H, Lrand_list = Lrand_list))
}


computeVariancePartitioningbw = function(hM, group=NULL, groupnames=NULL, start=1, na.ignore=FALSE, exclude = NULL, cond=NULL, wb=NULL, all=TRUE)
{
  ny = hM$ny
  nc = hM$nc
  nt = hM$nt
  ns = hM$ns
  np = hM$np
  nr = hM$nr
  
  gp=group # keep track on if we had group submitted or not
  gpname = groupnames
  if(is.null(group)){
    ## default: use terms
    if(nc > 1){
      if (is.null(hM$XFormula))
        stop("no XFormula: you must give 'group' and 'groupnames'")
      
      if(is.null(exclude)){
        group = attr(hM$X, "assign")
        if (group[1] == 0) # assign (Intercept) to group
          group[1] <- 1
        
        groupnames = attr(terms(hM$XFormula, data = hM$XData), "term.labels")
        gpname = groupnames
      }else{
        name.exclude = colnames(hM$X)[exclude]
        
        Fterm = c(attr(terms(hM$XFormula), which = "term.labels"))
        Fvars.exclude = Fterm[Fterm!='Effort']
        
        XexcludeForm = paste('~', paste(Fvars.exclude, collapse = " + "))
        XexcludeForm = as.formula(XexcludeForm)
        
        Xnew = as.data.frame(hM$XData[, colnames(hM$XData)!= name.exclude])
        
        Xexclude = model.matrix(XexcludeForm, Xnew)
        
        group = attr(hM$X, "assign")   # for calculation we keep # the excluded in group totals
        if (group[1] == 0) # assign (Intercept) to group
          group[1] <- 1
        
        groupnames = attr(terms(XexcludeForm, data = Xexclude), "term.labels")
        groupnamesA = attr(terms(hM$XFormula, data = hM$XData), "term.labels")
        gpname = groupnamesA
      }
      
    } else {
      group = c(1)
      groupnames = hM$covNames[1]
      gpname = groupnames
    }
    
  }else{
    if(!is.null(exclude)){
      name.exclude = colnames(hM$X)[exclude]
      
      Fterm = c(attr(terms(hM$XFormula), which = "term.labels"))
      
      groupnames = groupnames[!groupnames==name.exclude]
      groupnamesA = groupnames
      
    }
  }
  
  
  ngroups = max(group)
  
  
  R2T.Y = 0
  R2T.Beta = rep(0,nc)
  
  #If na.ignore=T, convert XData to a list
  # if(na.ignore){
  #   xl=list()
  #   for(s in 1:ns){
  #     xl[[s]]=hM$X
  #   }
  #   hM$X=xl
  # }
  # 
  
  
  
  ## info over MCMC iterations
  postList=poolMcmcChains(hM$postList, start=start)
  
  
  # +++++++ if conditions => separate
  G = unique(cond)
  n_cond = length(G)
  
  ## within / between prep
  if(!is.null(wb)){  
    Hcond <- sapply(G, function(x, g) { as.integer(g == x) }, cond)
    colnames(Hcond) <- G
  }
  
  nsamp = length(postList)
  
  
  ## Initialize arrays
  if(is.null(cond)){
    fixed = array(0, dim=c(ns, 1, nsamp))
    fixedsplit = array(0, dim=c(ns, ngroups, nsamp))
    random = array(0, c(ns, nr, nsamp))
  }else{
    # to avoid more complexity separate the calculated and the final arrays
    fixed = array(0, c(ns, 1, n_cond));
    fixedsplit = array(0, c(ns, ngroups, n_cond));
    random = array(0, c(ns, 1, n_cond));
    randomsplit = array(0, c(ns, nr, n_cond));
    
    fixed.f = array(0, c(ns, 1, n_cond, nsamp));
    fixedsplit.f = array(0, c(ns, ngroups, n_cond, nsamp));
    random.f = array(0, c(ns, 1, n_cond, nsamp));
    randomsplit.f = array(0, c(ns, nr, n_cond, nsamp));
  }
  
  if(!is.null(wb)){
    fixed.wb = array(0, c(ns, 1, 2, nsamp))
    fixedsplit.wb = array(0, c(ns, ngroups, 2, nsamp))
    random.wb = array(0, c(ns, 1, 2, nsamp))
    randomsplit.wb = array(0, c(ns, nr, 2, nsamp))
  }
  
  
  ## prepare covariance - for fixed part only
  if(is.null(cond)){
    cMA = cov(hM$X)
    
  }else{
    if(is.null(wb)){
      # partition condition _covariance
      cMA = sapply(G, function(x, A, g) cov(A[g == x,]), hM$X, cond, simplify = "array")
      dim(cMA) <- c(rep(ncol(hM$X), 2), length(G))
      dimnames(cMA) <- list(colnames(hM$X), colnames(hM$X), condition = G)
    }else{
      # within and between _ covar
      A_cm <- Hcond %*% diag(1 / colSums(Hcond)) %*% t(Hcond) %*% hM$X
      K_wit <- cov(hM$X - A_cm) #cov(hM$X) - cov(A_cm) #
      K_btw <- cov(A_cm)
      
      #dim(cMA) <- c(rep(ncol(hM$X), length(G)), length(G))
      #dimnames(cMA) <- list(colnames(hM$X), colnames(hM$X), condition = G)
    }
    
  }
  
  
  
  # ## n iterations list of matrices (NB obs x NB sp) for traits
  # geta = function(a){
  #   switch(class(hM$X)[1L],
  #          matrix = {res = hM$X%*%(t(hM$Tr%*%t(a$Gamma)))},
  #          list = {
  #            res = matrix(NA,hM$ny,hM$ns)
  #            for(j in 1:hM$ns) res[,j] =  hM$X[[j]]%*%(t(hM$Tr[j,]%*%t(postList[[1]]$Gamma)))
  #          }
  #   )
  #   return(res)
  # }
  # la=lapply(postList, geta)
  # 
  # ## n iterations list of matrices (NB obs x NB sp) for envrionmental info
  # getf = function(a){
  #   switch(class(hM$X)[1L],
  #          matrix = {res = hM$X%*%(a$Beta)},
  #          list = {
  #            res = matrix(NA,hM$ny,hM$ns)
  #            for(j in 1:hM$ns) res[,j] = hM$X[[j]]%*%a$Beta[,j]
  #          }
  #   )
  #   return(res)
  # }
  # lf=lapply(postList, getf)
  # 
  # ## simple matrix for traits info (NB covar x NB sp)
  # gemu = function(a){
  #   res = t(hM$Tr%*%t(a$Gamma))
  #   return(res)
  # }
  # lmu=lapply(postList, gemu)
  # 
  # ## simple matrix for environmental info (NB covar x NB sp)
  # gebeta = function(a){
  #   res = a$Beta
  #   return(res)
  # }
  # lbeta=lapply(postList, gebeta)
  
  
  
  ##run over MCMC chains NB
  for (i in 1:nsamp){
    
    if(is.null(wb)){   # 1 # calculating posterior sample variance or conditional variance
      
      
      if(is.null(cond)){
        fixedsplit1 = matrix(0,nrow=ns,ncol=ngroups)
        fixed1 = matrix(0,nrow=ns,ncol=1)
        
        random1 = matrix(0,nrow=ns,ncol=nr)
      }else{
        ### under conditions
        fixedsplit1 = lapply(1:ngroups, matrix, data= 0, nrow=ns, ncol=n_cond)
        fixed1 = matrix(0, nrow=ns, ncol=n_cond)
        
        random1 = array(0, dim=c(ns, nr, n_cond))
      }
      
      
      Beta.i = postList[[i]]$Beta
      Lambdas.i = postList[[i]]$Lambda
      Etas.i = postList[[i]]$Eta
      
      #f = lf[[i]]
      #a = la[[i]]
      
      #a=a-matrix(rep(rowMeans(a),hM$ns),ncol=hM$ns)
      #f=f-matrix(rep(rowMeans(f),hM$ns),ncol=hM$ns)
      #res1 = sum((rowSums((a*f))/(hM$ns-1))^2)
      #res2 = sum((rowSums((a*a))/(hM$ns-1))*(rowSums((f*f))/(hM$ns-1)))
      #R2T.Y = R2T.Y + res1/res2
      for (j in 1:ns){
        switch(class(hM$X)[1L], matrix = {cM = cMA}, list = {cM = cMA[[j]]})
        
        if(is.null(cond)){
          
          Beta.x = Beta.i
          cM.x = cM
          
          # normal set up
          ftotal = t(Beta.x[,j])%*%cM.x%*%Beta.x[,j]
          fixed1[j] = fixed1[j] + ftotal
          
          for (k in 1:ngroups){
            sel = (group==k)
            fpart = t(Beta.x[sel,j])%*%cM.x[sel,sel]%*%Beta.x[sel,j]
            fixedsplit1[j,k] = fixedsplit1[j,k] + fpart
          }
        }else{
          
          ### under conditions
          Beta.x = Beta.i
          cM.x = cM
          
          # normal set up
          Beta = Beta.x[,j]
          ftotal = apply(cM.x, c('condition'), function(x, B){ t(B) %*% x %*% B }, Beta)
          ftotal = as.array(ftotal)
          dimnames(ftotal) = list(condition=G)
          
          for (cd in 1:n_cond) {
            fixed1[j, cd] = fixed1[j, cd] + ftotal[cd]  
            
          }
          
          for (k in 1:ngroups){
            sel = (group==k)
            fpart = apply(cM.x, c('condition'), function(x, B){ t(B[sel,j]) %*% x[sel,sel] %*% B[sel,j] }, Beta.x)
            fpart = as.array(fpart)
            dimnames(fpart) = list(condition=G)
            
            fixedsplit1[[k]][j,] = fixedsplit1[[k]][j,] + fpart
            dimnames(fixedsplit1[[k]]) <- list(species=colnames(Beta), condition = G)
          }
          dimnames(fixed1) = list(species=colnames(Beta), condition=G)
        }
        
      }
      
      # random part ********************** work in progress
      
      if(is.null(cond)){
        for (level in seq_len(nr)){
          Lambda = Lambdas.i[[level]]
          nf = dim(Lambda)[[1]]
          for (factor in 1:nf){
            random1[,level] = random1[,level] + t(Lambda[factor,])*Lambda[factor,]
          }
        }
        
      }else{
        
        ## under conditions
        H.i = GetRandparam(hM, Lambdas = Lambdas.i, Etas = Etas.i)
        
        H.i.lev = H.i$H
        dim(H.i.lev)
        
        #Lrd = H.i$Lrand_list
        #length(Lrd)
        #dim(Lrd[[1]])
        
        dimnames(H.i.lev) = list(occasions=rownames(H.i$H), random=colnames(hM$studyDesign),
                                 Species=hM$sp)
        
        # variance on condition 
        
        rdsplit1 = array(0, dim=c(nr, n_cond, ns))
        rd1 = matrix(0, nrow=ns, ncol=n_cond)
        
        # random
        for (j in 1:ns) {
          r.tmp = sapply(G, function(x, A, g) cov(A[g == x,]), H.i.lev[,,j], cond, simplify = "array")
          dimnames(r.tmp) = list(rand=colnames(hM$studyDesign), rand=colnames(hM$studyDesign), condition=G)
          
          rtotal = apply(r.tmp, 'condition', function(x) {sum(x)}, simplify = "array")
          rd1[j,] = unlist(rtotal)
          
          rdsplit1[,,j] = apply(r.tmp, 'condition', function(x){diag(x)})
          
        }
        rdsplit1 = aperm(rdsplit1, c(3,1,2))
        
      }
      
      # ****
      
      
      if(is.null(cond)){
        if (nr>0){
          tot = fixed1 + rowSums(random1)
          fixed[,,i] = fixed[,,i] + fixed1/tot
          for (level in seq_len(nr)){
            random[,level,i] = random[,level,i] + random1[,level]/tot
          }
        }else{
          fixed[,,i] = fixed[,,i] + matrix(1,nrow=ns,ncol=1)
        }
        
        for (k in 1:ngroups){
          fixedsplit[,k,i] =  fixedsplit[,k,i] + fixedsplit1[,k]/rowSums(fixedsplit1)
        }
        
      }else{
        
        ### under conditions
        if (nr>0){
          
          ## fix +++
          #RSrand1 = colSums(rd1) # sum for all condition 
          tot = fixed1 +  rd1
          fix.div.tot = fixed1/tot
          dim(fix.div.tot) = c(ns, n_cond, 1)
          fix.div.tot = aperm((fix.div.tot), c(1,3,2))
          
          fixed = fix.div.tot  
          
          ## random +++
          rd.div.tot = rd1/tot
          dim(rd.div.tot) = c(ns, n_cond, 1)
          rd.div.tot = aperm((rd.div.tot), c(1,3,2))
          
          random = rd.div.tot  
          
        }else{
          # not changed
          fixed = fixed + matrix(1,nrow=ns,ncol=1)
        }
        
        dimnames(fixed) = list(species=colnames(Beta), fx='fx1', condition=G)
        
        
        ## fixsplit +++
        fixedsplit1 = abind(fixedsplit1, along=3)
        fixedsplit1 = aperm(fixedsplit1, c(1,3,2))
        for (cd in 1:n_cond) {
          for (k in 1:ngroups){
            fixedsplit[, k, cd] = fixedsplit1[, k, cd]/rowSums(fixedsplit1[,,cd])
          } 
        }
        
        ## randomsplit +++
        for (cd in 1:n_cond) {
          for (k in 1:nr){
            randomsplit[, k, cd] = rdsplit1[, k, cd]/rowSums(rdsplit1[,,cd])
          } 
        }
        
        
        # get ready for posterior
        fixed.f[,,,i] = fixed
        fixedsplit.f[,,,i] = fixedsplit
        random.f[,,,i] = random
        randomsplit.f[,,,i] = randomsplit
      }
      
      
      
    }else{  # 2 # calculating posterior sample btw - within variance
      
      for (wbind in 1:2) {
        
        fixedsplit1 = matrix(0,nrow=ns,ncol=ngroups)
        fixed1 = matrix(0,nrow=ns,ncol=1)
        
        rdsplit1 = matrix(0, nrow=ns, ncol=nr)
        rd1 = matrix(0, nrow=ns, ncol=1)
        
        Beta = postList[[i]]$Beta
        Lambdas.i = postList[[i]]$Lambda
        Etas.i = postList[[i]]$Eta
        
        H.i = GetRandparam(hM, Lambdas = Lambdas.i, Etas = Etas.i)
        H.i.lev = H.i$H
        
        Krd_wit = Krd_btw = list()
        for(j in 1:ns){
          Ard_cm <- Hcond %*% diag(1 / colSums(Hcond)) %*% t(Hcond) %*% H.i.lev[,,j]
          
          Krd_wit[[j]] <- cov(H.i.lev[,,j] - Ard_cm) 
          Krd_btw[[j]] <- cov(Ard_cm)
        }
        
        
        if(wbind==1){
          cMA = K_btw 
          vLr = Krd_btw
        }
        
        if(wbind==2){
          cMA = K_wit
          vLr = Krd_wit
        }
        
        ## fixed
        for (j in 1:ns){
          switch(class(hM$X)[1L], matrix = {cM = cMA}, list = {cM = cMA[[j]]})
          
          ftotal = t(Beta[,j])%*%cM%*%Beta[,j]
          fixed1[j] = fixed1[j] + ftotal
          
          Beta.x = Beta
          cM.x = cM
          
          for (k in 1:ngroups){
            sel = (group==k)
            fpart = t(Beta.x[sel,j])%*%cM.x[sel,sel]%*%Beta.x[sel,j]
            fixedsplit1[j,k] = fixedsplit1[j,k] + fpart
          }
          
        }
        
        
        ## random
        for (j in 1:ns) {
          r.tmp = vLr[[j]]
          dimnames(r.tmp) = list(rand=colnames(hM$studyDesign), rand=colnames(hM$studyDesign))
          
          rtotal = sum(r.tmp)
          rd1[j,] = rtotal
          
          rdsplit1[j,] = diag(r.tmp)
          
        }
        
        if (nr>0){
          tot = fixed1 + rd1
          ## fixed
          fixed.wb[,,wbind,i] = fixed.wb[,,wbind,i] + fixed1/tot
          
          ## random 
          random.wb[,,wbind,i] = random.wb[,,wbind,i] + rd1/tot
          
          
        }else{
          fixed.wb[,,wbind,i] = fixed.wb[,,wbind,i] + matrix(1,nrow=ns,ncol=1)
        }
        
        
        ## fixsplit +++
        for (k in 1:ngroups){
          fixedsplit.wb[,k,wbind,i] =  fixedsplit.wb[,k,wbind,i] + fixedsplit1[,k]/rowSums(fixedsplit1)
        }
        
        ## randomsplit +++
        for (k in 1:nr){
          randomsplit.wb[,k,wbind,i] =  randomsplit.wb[,k,wbind,i] + rdsplit1[,k]/rowSums(rdsplit1)
        }
        
        
      } # for btw and within conditions
      
      
    } # #2# loop
  } # MCMC loop
  
  
  
  if(is.null(wb)){
    if(is.null(cond)){
      fixed = fixed#/nsamp
      random = random#/nsamp
      fixedsplit = fixedsplit#/nsamp
      
      vals = array(0, dim=c(ngroups+nr, ns, nsamp))
      v.int = array(0, dim=c(ngroups+nr, ns, nsamp))
      dimnames(v.int)[[1]] = c(gpname, colnames(hM$studyDesign))
      for (i in 1:nsamp) {
        for (j in 1:ngroups){
          v.int[j,,i] = fixed[,,i]*fixedsplit[,j,i]
          
        }
        for (j in seq_len(nr)){
          v.int[ngroups+j,,i] = random[,j,i]
        }
      }
      
    }else{
      fixed = fixed.f#/nsamp
      random = random.f#/nsamp
      fixedsplit = fixedsplit.f#/nsamp
      randomsplit = randomsplit.f
      
      vals = array(0, dim=c(ngroups+nr, ns, n_cond, nsamp))
      v.int = array(0, dim=c(ngroups+nr, ns, n_cond, nsamp))
      dimnames(v.int)[[1]] = c(gpname, colnames(hM$studyDesign))
      dimnames(v.int)[[3]] = G
      for (cd in 1:n_cond) {
        for (i in 1:ngroups){
          v.int[i,,cd,] = fixed[,,cd,]*fixedsplit[,i,cd,]
        }
        for (i in seq_len(nr)){
          v.int[ngroups+i,,cd,] = random[,,cd,]*randomsplit[,i,cd,]
        }
      }
      
      
    }
    
  }else{
    fixed = fixed.wb#/nsamp
    random = random.wb#/nsamp
    fixedsplit = fixedsplit.wb#/nsamp
    randomsplit = randomsplit.wb
    
    vals = array(0, dim=c(ngroups+nr, ns, 2, nsamp))
    v.int = array(0, dim=c(ngroups+nr, ns, 2, nsamp))
    dimnames(v.int)[[1]] = c(gpname, colnames(hM$studyDesign))
    dimnames(v.int)[[3]] = c('btw', 'wit')
    for (wb.ind in 1:2) {
      for (i in 1:nsamp) {
        for (j in 1:ngroups){
          v.int[j,,wb.ind,i] = fixed[,,wb.ind,i]*fixedsplit[,j,wb.ind,i]
          
        }
        for (j in seq_len(nr)){
          v.int[ngroups+j,,wb.ind,i] = random[,,wb.ind,i]*randomsplit[,j,wb.ind,i]
        }
      }
    }
    
  }
  
  
  #R2T.Y = R2T.Y/hM$samples
  #R2T.Beta = R2T.Beta/hM$samples
  
  if(is.null(exclude)){
    vals = v.int
  }else{
    if(isTRUE(all)){
      Lbef = v.int[1:ngroups,,]
      Totbef = apply(Lbef, c(3), colSums) #  sum accross covar
      
      pos = which(dimnames(Lbef)[[1]]==name.exclude) # new position
      Lcov = Lbef[-pos,,] # without Eff
      Totfix = apply(Lcov, c(3), colSums) #  sum accross covar without Effort
      
      Ratio = Totbef/Totfix
      
      LR =  array(0, dim=c(ngroups-length(exclude), ns, nsamp))
      
      for (i in 1:(ngroups-length(exclude))) {
        LR[i,,] = Lcov[i,,]*Ratio #
      }
      valsnew = array(0, dim=c((ngroups-length(exclude))+nr, ns, nsamp))
      
      valsnew[1:(ngroups-length(exclude)),,] = LR
      valsnew[(ngroups-length(exclude)+1):(nr+ngroups-length(exclude)),,] = v.int[(ngroups+1):(ngroups+nr),,]
      
    }else{
      Lbef = v.int[1:ngroups,,,]
      Totbef = apply(Lbef, c(3), colSums) #  sum accross covar
      
      pos = which(dimnames(Lbef)[[1]]==name.exclude) # new position
      Lcov = Lbef[-pos,,,] # without Eff
      Totfix = apply(Lcov, c(3,4), colSums) #  sum accross covar without Effort
      
      Ratio = Totbef/Totfix
      
      if(is.null(wb)){
        LR =  array(0, dim=c(ngroups-length(exclude), ns, n_cond, nsamp))
      }else{
        LR =  array(0, dim=c(ngroups-length(exclude), ns, 2, nsamp))
      }
      for (i in 1:(ngroups-length(exclude))) {
        LR[i,,,] = Lcov[i,,,]*Ratio #
      }
      
      if(is.null(wb)){
        valsnew = array(0, dim=c((ngroups-length(exclude))+nr, ns, n_cond, nsamp))
      }else{
        valsnew = array(0, dim=c((ngroups-length(exclude))+nr, ns, 2, nsamp))
      }
      
      valsnew[1:(ngroups-length(exclude)),,,] = LR
      valsnew[(ngroups-length(exclude)+1):(nr+ngroups-length(exclude)),,,] = v.int[(ngroups+1):(ngroups+nr),,,]
      
    }
    
    vals=valsnew
  }
  
  names(R2T.Beta)=hM$covNames
  colnames(vals)=hM$spNames
  
  if(is.null(exclude)){
    leg = groupnames
    for (r in seq_len(nr)) {
      leg = c(leg, paste("Random: ", hM$rLNames[r], sep = ""))
    }
    rownames(vals) = leg
    
  }else{
    leg = groupnamesA#[- which(groupnamesA == name.exclude)]
    
    for (r in seq_len(nr)) {
      leg = c(leg, paste("Random: ", hM$rLNames[r], sep = ""))
    }
    dimnames(vals)[[1]] = leg
    
  }
  
  
  VP = list()
  VP$vals = vals
  #VP$R2T = list(Beta=R2T.Beta,Y=R2T.Y)
  VP$group = group
  VP$groupnames = groupnames
  
  if(!is.null(cond)){
    VP$G = G
  }
  
  
  return(VP)
}


# yearly model
#####  species
# VP all ----------------------------------------------------------------------------------------------------------------
group1 = c(1,1,1, 2,2, 3,3, 4,4, 5,6, 7,7,8,8,9,9, 10,10, 11)
groupnames1 = c("prec", 'tempW', 'tempS', "snow", 'clumpy', 'ent',
               'BLforest', 'Cforest', 'Mforest', 
               'SNHerb', 'Effort')

MeanA.vp = list()
for (i in 1:2) {
  #####  species
  # VP all
  VP.a = computeVariancePartitioningbw(modelyear[[i]], group=group1, groupnames=groupnames1, start=1, 
                                      na.ignore=FALSE, exclude = 20, cond=NULL, wb=NULL, all=TRUE)
  
  dim(VP.a)
  # mean over species and MCMC samples
  dim(VP.a$vals)  # covar / species /  / MCMC sample
  MeanA.vp[[i]] = apply(VP.a$vals, c(1, 2), mean)
  dim(MeanA.vp[[i]])  # covar / btw-with
  
}


# Try to change order of species per trait
modelyear[[1]]$TrData[order(modelyear[[1]]$TrData$HostPlantForm),]

spT_order = rownames(modelyear[[1]]$TrData[order(modelyear[[1]]$TrData$HostPlantForm),])
spT_order

spT_order = data.frame(species = spT_order, 
                       HostPlantForm = modelyear[[1]]$TrData$HostPlantForm[order(modelyear[[1]]$TrData$HostPlantForm)])


lab1 = list('PA hurdle 1', 'AB hurdle 2')

multistack.bar <- function(x=list(x), betweenspace = 2, withinspace=0, group.names = NULL,...){
  # note that number of rows in each component matrix of the list
  # should be equal
  # if you have missing value use "NA", to make it complete 
  # number of column can differ
  mylist <- x
  space = list()
  space[[1]] <- c(rep(withinspace, ncol(mylist[[1]])), betweenspace) 
  
  for ( i in 2:length(mylist)){
    if(i == length(mylist)){
      space[[i]] <- c(rep(withinspace, ncol(mylist[[i]])-1)) 
    } else{
      if(is.null(ncol(mylist[[i]]))){
        space[[i]] <- c(betweenspace) 
      }else{
        space[[i]] <- c(rep(withinspace, ncol(mylist[[i]])-1), betweenspace) 
      }
    }
    
  }
  un.space <- c(unlist(space))
  newdata <- do.call("cbind", mylist)
  b.pos= barplot(newdata, space= un.space, ...)
  
  if(is.null(group.names)){}else{
    tg = list()
    
    for (i in 1:length(unique(group.names))) {
      t.i = group.names[group.names == unique(group.names)[i]]
      if((length(t.i) %% 2) == 0){
        tg[[i]] = c(rep('', ceiling(length(t.i)-1)/2), rep(as.character(unique(group.names)[i]), 1), rep('', ceiling(length(t.i))/2))
      }else{
        tg[[i]] = c(rep('', floor(length(t.i)-1)/2), rep(as.character(unique(group.names)[i]), 1), rep('', ceiling(length(t.i)-1)/2))
        
      }
    }
    
    lab = unlist(tg)
    
    par(mgp=c(2,1,0))
    #axis(4, at=b.pos, labels = lab, tick = F, cex.axis=0.6, las=2)
    mapply(axis, side = 4, at=b.pos, labels = lab, tick = F, cex.axis=0.8, las=2)
    par(mgp=c(3,1,0))
  }
  
}


par(mfrow=c(1,2), mar=c(2,8,4,14), xpd=T)

par(mfrow=c(1,2), mar=c(8,1,2,12), xpd=T)

coul = c('turquoise3', 'orange', 'red', 
                     'white', #'purple', 
                     'mediumpurple1', 'plum',
                     'green', 'forestgreen', 'lightgreen','brown', #"grey48", 
                     "#E5C494", 'grey78','blue3')
                     



head(spT_order)

levels(spT_order$HostPlantForm)[levels(spT_order$HostPlantForm) == '5'] = '4'


levels(spT_order$HostPlantForm)[levels(spT_order$HostPlantForm) == '1'] = 'Herbaceous'
levels(spT_order$HostPlantForm)[levels(spT_order$HostPlantForm) == '2'] = 'Graminoid'
levels(spT_order$HostPlantForm)[levels(spT_order$HostPlantForm) == '4'] = 'Other'
levels(spT_order$HostPlantForm)[levels(spT_order$HostPlantForm) == '31'] = 'Dwarf shrubs'
levels(spT_order$HostPlantForm)[levels(spT_order$HostPlantForm) == '32'] = 'Decideous'
levels(spT_order$HostPlantForm)[levels(spT_order$HostPlantForm) == '33'] = 'Coniferous'



###-------------------------------------------------------------
par(mfrow=c(1,2), mar=c(1,3,2,1), oma=c(1,0,1,0))
layout(matrix(c(1, 2), ncol=2))

par(mar = c(1,1,7.3,4))
VP2_order = MeanA.vp[[1]][, c(spT_order$species)]

ng = dim(VP2_order)[1]
leg2 = rownames(VP2_order)

mylist2 <- list(VP2_order[,which(spT_order$HostPlantForm == 'Herbaceous')], VP2_order[,which(spT_order$HostPlantForm == 'Graminoid')], 
                VP2_order[,which(spT_order$HostPlantForm == 'Other')], #VP2_order[,which(spT_order$HostPlantForm == '5')],
                VP2_order[,which(spT_order$HostPlantForm == 'Dwarf shrubs')],
                VP2_order[,which(spT_order$HostPlantForm == 'Decideous')], VP2_order[,which(spT_order$HostPlantForm == 'Coniferous')])
mylist2neg = lapply(mylist2, function(x){(-1)*x})
mylist2neg = lapply(mylist2, function(y) apply(y, 2, function(x) (-1)*x / sum(x, na.rm = TRUE)))


mn = lapply(mylist2neg, function(x) (-1)*rowMeans(x))
means2 = round(100 * Reduce('+', mn)/length(mn), 1)
for (j in 1:ng) {
  leg2[j] = paste(leg2[j], " (mean = ", toString(means2[j]),
                  ")", sep = "")
}
multistack.bar(mylist2neg, betweenspace=1.2, withinspace=0.1, las=1, group.names = as.character(spT_order$HostPlantForm),
               col = coul, xlab="", ylab="", xaxt='n', legend = leg2,  horiz=T, cex.lab =1.3, cex.axis=1.3, yaxt='n',
               main = '', cex.main = 1.3, cex.names = 2,
               args.legend=list(x ='topleft', bty='n', cex=1.05, ncol=2, inset=c(0,-0.1)))
axis(1, at = c(-1, -0.8, -0.6, -0.4, -0.2, 0), labels = c(1, 0.8, 0.6, 0.4, 0.2, 0))
title(lab1[1], adj = 0.8, line = 6.2, cex.main = 1.3)

par(mar = c(1,4,7.3,1))
VPAB_order = MeanA.vp[[2]][, c(spT_order$species)]

ng = dim(VPAB_order)[1]
leg2 = rownames(VPAB_order)

mylist2AB <- list(VPAB_order[,which(spT_order$HostPlantForm == 'Herbaceous')], VPAB_order[,which(spT_order$HostPlantForm == 'Graminoid')], 
                  VPAB_order[,which(spT_order$HostPlantForm == 'Other')], #VPAB_order[,which(spT_order$HostPlantForm == '5')],
                  VPAB_order[,which(spT_order$HostPlantForm == 'Dwarf shrubs')],
                  VPAB_order[,which(spT_order$HostPlantForm == 'Decideous')], VPAB_order[,which(spT_order$HostPlantForm == 'Coniferous')])

mylist2AB = lapply(mylist2AB, function(y) apply(y, 2, function(x) x / sum(x, na.rm = TRUE)))
mn = lapply(mylist2AB, function(x) rowMeans(x))
means2 = round(100 * Reduce('+', mn)/length(mn), 1)
for (j in 1:ng) {
  leg2[j] = paste(leg2[j], " (mean = ", toString(means2[j]),
                  ")", sep = "")
}


multistack.bar(mylist2AB, betweenspace=1.2, withinspace=0.1, las=1, group.names = rep('', 77), yaxt='n',
               col = coul, xlab="", ylab="", xaxt='n', legend = leg2,  horiz=T, cex.lab =1.3, cex.axis=1.3, 
               main = '', cex.main = 1.3, cex.names = 1.5,
               args.legend=list(x ='topright', bty='n', cex=1.05, ncol=2, inset=c(0,-0.1)))
axis(1, at = c(0, 0.2, 0.4, 0.6, 0.8, 1), labels = c(0, 0.2, 0.4, 0.6, 0.8, 1))
title(lab1[2], adj = 0.2, line = 6.2, cex.main = 1.3)



# VP all ----------------------------------------------------------------------------------------------------------------
group1 = c(1,1,1, 1,1, 1,1, 1,1, 2,2, 2,2,2,2,2,2, 2,2, 3)
groupnames1 = c('Climate','Habitat', 'Effort')

MeanA.vp = list()
conf.up = conf.low = list()

for (i in 1:2) {
  #####  species
  # VP all
  VP.a = computeVariancePartitioningbw(modelyear[[i]], group=group1, groupnames=groupnames1, start=1, 
                                      na.ignore=FALSE, exclude = 20, cond=NULL, wb=NULL, all=TRUE)
  
  dim(VP.a)
  # mean over species and MCMC samples
  dim(VP.a$vals)  # covar / species /  / MCMC sample
  MeanA.vp[[i]] = apply(VP.a$vals, c(1, 2), mean)
  dim(MeanA.vp[[i]])  # covar / btw-with
  conf.up[[i]] = apply(VP.a$vals, c(1, 2), function(x){quantile(x, 0.975)})
  conf.low[[i]] = apply(VP.a$vals, c(1, 2), function(x){quantile(x, 0.025)})
}

lapply(conf.up, rowMeans)
lapply(conf.low, rowMeans)
Mean.Sites = lapply(MeanA.vp, rowMeans)
Mean.Sites



Mean.vp = list()
confvp.up = confvp.low = medvp = list()
max = matrix(NA, nrow=1000, ncol = 2)
for (i in 1:2) {
  #####  species
  # VP all
  VP.a = computeVariancePartitioningbw(modelyear[[i]], group=group1, groupnames=groupnames1, start=1, 
                                       na.ignore=FALSE, exclude = 20, cond=NULL, wb=NULL, all=TRUE)
  
  for (j in 1:1000) {
    # mean over species and MCMC samples
    dim(VP.a$vals)  # covar / species /  / MCMC sample
    Vp.calc = VP.a$vals[,,j]
    
    max[j,] = table(apply(Vp.calc[1:2,], 2, which.max))
  }
  Mean.vp[[i]] = apply(max, c(2), mean)
  
  medvp[[i]] = apply(max, c(2), function(x){quantile(x, 0.5)})
  confvp.up[[i]] = apply(max, c(2), function(x){quantile(x, 0.975)})
  confvp.low[[i]] = apply(max, c(2), function(x){quantile(x, 0.025)})
}




# Try to change order of species per trait
modelyear[[1]]$TrData[order(modelyear[[1]]$TrData$HostPlantForm),]

spT_order = rownames(modelyear[[1]]$TrData[order(modelyear[[1]]$TrData$HostPlantForm),])
spT_order

spT_order = data.frame(species = spT_order, 
                       HostPlantForm = modelyear[[1]]$TrData$HostPlantForm[order(modelyear[[1]]$TrData$HostPlantForm)])




multistack.bar <- function(x=list(x), betweenspace = 2, withinspace=0, group.names = NULL,...){
  # note that number of rows in each component matrix of the list
  # should be equal
  # if you have missing value use "NA", to make it complete 
  # number of column can differ
  mylist <- x
  space = list()
  space[[1]] <- c(rep(withinspace, ncol(mylist[[1]])), betweenspace) 
  
  for ( i in 2:length(mylist)){
    if(i == length(mylist)){
      space[[i]] <- c(rep(withinspace, ncol(mylist[[i]])-1)) 
    } else{
      if(is.null(ncol(mylist[[i]]))){
        space[[i]] <- c(betweenspace) 
      }else{
        space[[i]] <- c(rep(withinspace, ncol(mylist[[i]])-1), betweenspace) 
      }
    }
    
  }
  un.space <- c(unlist(space))
  newdata <- do.call("cbind", mylist)
  b.pos= barplot(newdata, space= un.space, ...)
  
  if(is.null(group.names)){}else{
    tg = list()
    
    for (i in 1:length(unique(group.names))) {
      t.i = group.names[group.names == unique(group.names)[i]]
      if((length(t.i) %% 2) == 0){
        tg[[i]] = c(rep('', ceiling(length(t.i)-1)/2), rep(as.character(unique(group.names)[i]), 1), rep('', ceiling(length(t.i))/2))
      }else{
        tg[[i]] = c(rep('', floor(length(t.i)-1)/2), rep(as.character(unique(group.names)[i]), 1), rep('', ceiling(length(t.i)-1)/2))
        
      }
    }
    
    lab = unlist(tg)
    
    par(mgp=c(2,1,0))
    #axis(4, at=b.pos, labels = lab, tick = F, cex.axis=0.6, las=2)
    mapply(axis, side = 4, at=b.pos, labels = lab, tick = F, cex.axis=1.1, las=2)
    par(mgp=c(3,1,0))
  }
  
}


par(mfrow=c(1,2), mar=c(2,8,4,14), xpd=T)

par(mfrow=c(1,2), mar=c(8,1,2,12), xpd=T)

coul = c("#66C2A5", 'purple',
                     "#E5C494", 'grey78','blue3')
                     

head(spT_order)

levels(spT_order$HostPlantForm)[levels(spT_order$HostPlantForm) == '5'] = '4'


levels(spT_order$HostPlantForm)[levels(spT_order$HostPlantForm) == '1'] = 'Herbaceous'
levels(spT_order$HostPlantForm)[levels(spT_order$HostPlantForm) == '2'] = 'Graminoids'
levels(spT_order$HostPlantForm)[levels(spT_order$HostPlantForm) == '4'] = 'Others'
levels(spT_order$HostPlantForm)[levels(spT_order$HostPlantForm) == '31'] = 'Dwarf shrubs'
levels(spT_order$HostPlantForm)[levels(spT_order$HostPlantForm) == '32'] = 'Decideous'
levels(spT_order$HostPlantForm)[levels(spT_order$HostPlantForm) == '33'] = 'Coniferous'


lab1 = list('A - PA model Sites', 'B - AB model Sites')
###-------------------------------------------------------------
par(mfrow=c(1,2), mar=c(1,3,2,1), oma=c(1,0,1,0))
layout(matrix(c(1, 2), ncol=2))

par(mar = c(1,1,4,4))
VP2_order = MeanA.vp[[1]][, c(spT_order$species)]

ng = dim(VP2_order)[1]
leg2 = rownames(VP2_order)

mylist2 <- list(VP2_order[,which(spT_order$HostPlantForm == 'Herbaceous')], VP2_order[,which(spT_order$HostPlantForm == 'Graminoids')], 
                VP2_order[,which(spT_order$HostPlantForm == 'Others')], #VP2_order[,which(spT_order$HostPlantForm == '5')],
                VP2_order[,which(spT_order$HostPlantForm == 'Dwarf shrubs')],
                VP2_order[,which(spT_order$HostPlantForm == 'Decideous')], VP2_order[,which(spT_order$HostPlantForm == 'Coniferous')])
mylist2neg = lapply(mylist2, function(x){(-1)*x})
mylist2neg = lapply(mylist2, function(y) apply(y, 2, function(x) (-1)*x / sum(x, na.rm = TRUE)))


mn = Mean.Sites[[1]]
means2 = round(100*mn, 1)
for (j in 1:ng) {
  leg2[j] = paste(leg2[j], " (mean = ", toString(means2[j]),
                  ")", sep = "")
}
multistack.bar(mylist2neg, betweenspace=1.2, withinspace=0.1, las=1, group.names = as.character(spT_order$HostPlantForm),
               col = coul, xlab="", ylab="", xaxt='n', legend = leg2,  horiz=T, cex.lab =1.3, cex.axis=1.3, yaxt='n',
               main = '', cex.main = 1.3, cex.names = 2,
               args.legend=list(x ='topleft', bty='n', cex=1.05, inset=c(0,-0.11)))
axis(1, at = c(-1, -0.8, -0.6, -0.4, -0.2, 0), labels = c(1, 0.8, 0.6, 0.4, 0.2, 0))
title(lab1[1], adj = 0.8, line = 0, cex.main = 1.3)


par(mar = c(1,4,4,1))
VPAB_order = MeanA.vp[[2]][, c(spT_order$species)]

ng = dim(VPAB_order)[1]
leg2 = rownames(VPAB_order)

mylist2AB <- list(VPAB_order[,which(spT_order$HostPlantForm == 'Herbaceous')], VPAB_order[,which(spT_order$HostPlantForm == 'Graminoids')], 
                  VPAB_order[,which(spT_order$HostPlantForm == 'Others')], #VPAB_order[,which(spT_order$HostPlantForm == '5')],
                  VPAB_order[,which(spT_order$HostPlantForm == 'Dwarf shrubs')],
                  VPAB_order[,which(spT_order$HostPlantForm == 'Decideous')], VPAB_order[,which(spT_order$HostPlantForm == 'Coniferous')])

mylist2AB = lapply(mylist2AB, function(y) apply(y, 2, function(x) x / sum(x, na.rm = TRUE)))
mn = Mean.Sites[[2]]
means2 = round(100*mn, 1)
for (j in 1:ng) {
  leg2[j] = paste(leg2[j], " (mean = ", toString(means2[j]),
                  ")", sep = "")
}

multistack.bar(mylist2AB, betweenspace=1.2, withinspace=0.1, las=1, group.names = rep('', 77), yaxt='n',
               col = coul, xlab="", ylab="", xaxt='n', legend = leg2,  horiz=T, cex.lab =1.3, cex.axis=1.3, 
               main = '', cex.main = 1.3, cex.names = 1.5,
               args.legend=list(x ='topright', bty='n', cex=1.05, inset=c(0,-0.11)))
axis(1, at = c(0, 0.2, 0.4, 0.6, 0.8, 1), labels = c(0, 0.2, 0.4, 0.6, 0.8, 1))
title(lab1[2], adj = 0.2, line = 0, cex.main = 1.3)


## saving for sites
mylist2neg_sites = mylist2neg
VP2_order_sites = VP2_order
mylist2AB_sites = mylist2AB
VPAB_order_sites = VPAB_order


save(mylist2neg_sites, mylist2AB_sites, VP2_order_sites, VPAB_order_sites, Mean.Sites, file='VPall_Sites0904.RDATA')



### which is the most important
## PA
table(apply(VP2_order[1:2,], 2, which.max))#/length(apply(VP2_order[1:2,], 2, which.max))
table(apply(conf.up[[1]][1:2,], 2, which.max))#/length(apply(conf.up[[1]][1:2,], 2, which.max))
table(apply(conf.low[[1]][1:2,], 2, which.max))#/length(apply(conf.low[[1]][1:2,], 2, which.max))


## AB
table(apply(VPAB_order[1:2,], 2, which.max))/length(apply(VPAB_order[1:2,], 2, which.max))
table(apply(conf.up[[2]][1:2,], 2, which.max))/length(apply(conf.up[[2]][1:2,], 2, which.max))
table(apply(conf.low[[2]][1:2,], 2, which.max))/length(apply(conf.low[[2]][1:2,], 2, which.max))
