roleR: Rules of Life Engine

#' @title Run a `roleModel` or `roleExperiment`
#' 
#' @description `runRole` takes a `role*` object and simulates all processes 
#'     specified in the given parameters
#' 
#' @param x `roleModel` or `roleExperiment` object to run
#' @param cores number of cores
#' @return `roleModel` or `roleExperiment` run for the specified number of time
#'     steps
#' 
#' @details Prior to invoking `runRole`, the `role*` input object must be 
#'     initialized with a set of parameters, e.g., as generated by `roleParams`.
#' 
#' @examples 
#' # create and run a model
#' model <- roleModel(roleParams())
#' model <- runRole(model)
#' 
#' # create and run an experiment
#' p1 <- roleParams(dispersal_prob = 0.2)
#' p2 <- roleParams(dispersal_prob = 0.3)
#' p3 <- roleParams(dispersal_prob = 0.4)
#' exp <- roleExperiment(list(p1,p2,p3))
#' exp <- runRole(exp)
#' @rdname runRole
#' @export

setGeneric('runRole', 
           def = function(x, cores=1) standardGeneric('runRole'), 
           signature = 'x')


#' runRole on roleModel
#' @name runRole
#' @aliases runRole,roleModel-method
#' @docType methods
#' @rdname runRole
setMethod('runRole', 
          signature = 'roleModel', 
          definition = function(x) {
              m <- x
              
              pvals <- getValuesFromParams(m@params, 1:m@params@niter)
              
              # augment the data in the model based on the params
              m <- .bufferModelData(m)
              
              # returns the new modelSteps (a list of roleData)
              m@modelSteps <- iterModelCpp(slot(m@modelSteps[[1]],"localComm"), 
                                           slot(m@modelSteps[[1]],"metaComm"),
                                           slot(m@modelSteps[[1]],"phylo"),
                                           pvals,
                                           print = FALSE)
              
              # trim data, removing the unused buffer
              m <- .trimModelData(m)
              
              # increment species ids and edges to shift indexing from Cpp to R
              for(d in 1:length(m@modelSteps)) {
                  m@modelSteps[[d]]@localComm@indSpecies <- 
                      m@modelSteps[[d]]@localComm@indSpecies + 1
                  m@modelSteps[[d]]@phylo@e <- m@modelSteps[[d]]@phylo@e + 1
              }
              
              # add popgen
              # m <- sim_seqs(m)
              
              return(m)
          }
)


#' runRole on roleExperiment
#' @name runRole
#' @aliases runRole,roleExperiment-method
#' @docType methods
#' @rdname runRole

setMethod('runRole', 
          signature = 'roleExperiment', 
          definition = function(x, cores = 1) {
              oldInits <- x@inits
              
              # run consituent models
              if(cores == 1){
                  modList <- lapply(x@inits, runRole)
              } else {
                  cl <- snow::makeCluster(cores, type="SOCK")
                  modList <- snow::clusterApply(cl, x@inits, runRole)
                  snow::stopCluster(cl)
              }
              
              # convert list of models to list of experiments
              expList <- lapply(modList, as, Class = 'roleExperiment')
              
              # combine list
              x <- Reduce(rbind2, expList)
              
              # add back the original inits
              x@inits <- oldInits
              
              return(x)
          }
)

# buffer model data
# augment the data of the not-yet-run model based on what is expected from the 
# params called right before the model is run in Cpp
# @param model model

.bufferModelData <- function(model) {
    p <- model@params 
    
    # calculate expected number of new species using binom
    expec_n_spec <- stats::qbinom(0.9, p@niter,
                                  prob = mean(p@speciation_local(1:p@niter)))
    el_add <- (expec_n_spec * 2 - 1) + 1
    at_add <- expec_n_spec + 1
    
    # buffer phylo ----
    
    # edges get -1s
    model@modelSteps[[1]]@phylo@e <- rbind(model@modelSteps[[1]]@phylo@e, 
                                           matrix(-1, nrow = el_add, ncol = 2)) 
    
    # lengths get 0s
    model@modelSteps[[1]]@phylo@l <- c(model@modelSteps[[1]]@phylo@l, 
                                       rep(0, el_add)) 
    
    # alives get FALSE
    model@modelSteps[[1]]@phylo@alive <- c(model@modelSteps[[1]]@phylo@alive, 
                                           rep(FALSE, at_add)) 
    
    # tipNames get 'local_speciation_j'
    model@modelSteps[[1]]@phylo@tipNames <- 
        c(model@modelSteps[[1]]@phylo@tipNames, 
          paste0('local_speciation_', 1:at_add))
    
    # calc buffer size for local species vects 
    local_add <- p@species_meta + expec_n_spec
    
    # buffer local species vectors with 0s
    model@modelSteps[[1]]@localComm@spAbund <- 
        c(model@modelSteps[[1]]@localComm@spAbund, rep(0, local_add))
    model@modelSteps[[1]]@localComm@spTrait <- 
        c(model@modelSteps[[1]]@localComm@spTrait, rep(0, local_add))
    model@modelSteps[[1]]@localComm@spAbundHarmMean <-  
        rep(0, length(model@modelSteps[[1]]@localComm@spAbund) + local_add)
    model@modelSteps[[1]]@localComm@spLastOriginStep <-  
        rep(0, length(model@modelSteps[[1]]@localComm@spAbund) + local_add)
    model@modelSteps[[1]]@localComm@spExtinctionStep <-  
        rep(0, length(model@modelSteps[[1]]@localComm@spAbund) + local_add)
    
    return(model)
}


###### should move this to rcpp ###########

# helper to trim the data of unused indices after the model is run
.trimModelData <- function(model) {
    
    for(i in 1:length(model@modelSteps)) {
        # trim phylo
        model@modelSteps[[i]]@phylo@e <- model@modelSteps[[i]]@phylo@e[model@modelSteps[[i]]@phylo@e[,1] != -1,]
        model@modelSteps[[i]]@phylo@l <- model@modelSteps[[i]]@phylo@l[model@modelSteps[[i]]@phylo@l != 0] 
        last_alive_index <- max(which(model@modelSteps[[i]]@phylo@alive == TRUE))
        model@modelSteps[[i]]@phylo@alive <- model@modelSteps[[i]]@phylo@alive[1:last_alive_index]
        model@modelSteps[[i]]@phylo@tipNames <- model@modelSteps[[i]]@phylo@tipNames[model@modelSteps[[i]]@phylo@tipNames != ''] # this MAY cause errors
        
        # trim local
        # find place where augmented 0s started, which is the number of species to ever have existed 
        # I think this is the last alive index??
        model@modelSteps[[i]]@localComm@spAbund <- model@modelSteps[[i]]@localComm@spAbund[1:last_alive_index]
        model@modelSteps[[i]]@localComm@spTrait <- model@modelSteps[[i]]@localComm@spTrait[1:last_alive_index]
        model@modelSteps[[i]]@localComm@spAbundHarmMean  <- model@modelSteps[[i]]@localComm@spAbundHarmMean [1:last_alive_index]
        model@modelSteps[[i]]@localComm@spLastOriginStep <- model@modelSteps[[i]]@localComm@spLastOriginStep[1:last_alive_index]
        model@modelSteps[[i]]@localComm@spExtinctionStep <- model@modelSteps[[i]]@localComm@spExtinctionStep[1:last_alive_index]
    }
    
    return(model)
}


# getValuesFromParams
# @description
# run iter functions over params to generate a new object of class paramValues 
# that contains ONLY vectors of values and no functions
# @param p an object of class `roleParams`
# @param i the iterations over which to calculate parameter values
# @return a list of parameter vectors, each equal in length to `length(i)`

getValuesFromParams <- function(p, i) {
    # slot names
    s <- methods::slotNames(p)
    
    # apply over slots
    pvals <- lapply(s, function(x) {
        thisP <- methods::slot(p, x)
        
        if(inherits(thisP, 'function')) {
            return(thisP(i))
        } else {
            return(rep(thisP, length(i)))
        }
    })
    
    names(pvals) <- s
    
    return(pvals)
}