R/event_model.R

In bbuchsbaum/fmrireg: R package for the anlysis of fmri data

#' @importFrom utils head
#' @importFrom plotly plot_ly layout animation_opts
#' @importFrom tidyr complete pivot_longer
#' @importFrom dplyr mutate
#' @importFrom rlang .data
NULL

## ============================================================================
## Section 1: Unified Event Model Constructor (NEW Pipeline)
## ============================================================================

#' Construct an fMRI Event Model
#' 
#' This is the main constructor for `event_model` objects. It unifies the previous
#' formula and list-based interfaces and uses a more efficient internal pipeline.
#'
#' @param formula_or_list Either a formula (e.g., `onset ~ hrf(cond) + hrf(mod)`) 
#'        or a list of pre-defined `hrfspec` objects.
#' @param data A `data.frame` containing event variables referenced in the formula 
#'        or needed by the `hrfspec` objects.
#' @param block A formula (e.g., `~ run`) or vector specifying the block/run for each event.
#' @param sampling_frame An object of class `sampling_frame` defining the scan timing (TR, block lengths).
#' @param durations Numeric vector or scalar specifying event durations (seconds). Default is 0.
#' @param drop_empty Logical indicating whether to drop empty events during term construction. Default is TRUE.
#' @param precision Numeric precision for HRF sampling/convolution. Default is 0.3.
#' @param parallel Logical indicating whether to use parallel processing for term convolution (requires `future.apply`). Default is FALSE.
#' @param progress Logical indicating whether to show a progress bar during term realisation. Default is FALSE.
#' @param ... Additional arguments (currently unused).
#' 
#' @details 
#' ### Column Naming
#' The columns in the resulting design matrix follow the naming convention: 
#' `term_tag` + `_` + `condition_tag` + `_b##` basis suffix
#' 
#' Where:
#' *   `term_tag`: The unique tag assigned to the `hrf()` term (see below).
#' *   `condition_tag`: Represents the specific factor level or continuous regressor within the term (e.g., `condition.A`, `poly_RT_01`, `condition.A_task.go`).
#' *   `_b##`: Optional suffix added for HRFs with multiple basis functions (e.g., `_b01`, `_b02`).
#' 
#' ### Term Naming and Clash Resolution
#' Each term in the model (typically defined by an `hrf()` call in a formula) gets a 
#' unique `term_tag`. This tag is used as the prefix for all columns generated by that term.
#' 
#' - **Default Naming:** If no explicit `id` (or `name`) is given in `hrf()`, the tag is derived  
#'   from the variable names (e.g., `hrf(condition)` -> `condition`, `hrf(RT, acc)` -> `RT_acc`).
#' - **Explicit Naming:** Use `id=` within `hrf()` for an explicit tag (e.g., `hrf(condition, id="CondMain")`).
#' - **Sanitization:** Dots (`.`) in tags are converted to underscores (`_`).   
#' - **Clash Resolution:** If multiple terms generate the same tag, `#` and a number are appended 
#'   to ensure uniqueness (e.g., `condition`, `condition#1`).
#' 
#' This consistent naming scheme replaces the previous `compact` and `qualified` styles.
#' 
#' @return An object of class `c("event_model", "list")` containing the terms, design matrix, 
#'         sampling frame, and other metadata.
#'         
#' @examples 
#' # Example using formula interface
#' des <- data.frame(onset = seq(0, 90, by=10),
#'                   run = rep(1:2, each=5),
#'                   cond = factor(rep(c("A","B"), 5)),
#'                   mod = rnorm(10))
#' sframe <- fmrihrf::sampling_frame(blocklens=c(50, 60), TR=2)
#' 
#' ev_model_form <- event_model(onset ~ hrf(cond) + hrf(mod, basis="spmg3"), 
#'                             data = des, block = ~run, sampling_frame = sframe)
#' print(ev_model_form)
#' head(design_matrix(ev_model_form))
#' 
#' # Example using list interface (less common)
#' # spec1 <- hrf(cond)
#' # spec2 <- hrf(mod, basis="spmg3")
#' # ev_model_list <- event_model(list(spec1, spec2), data=des, block=des$run, sampling_frame=sframe)
#' # print(ev_model_list)                         
#'                              
#' @export
#' @include event_model_helpers.R
event_model <- function(formula_or_list, data, block, sampling_frame, 
                        durations = 0, drop_empty = TRUE, precision = 0.3, 
                        parallel = FALSE, progress = FALSE, ...) {
  
  # Stage 1: Parse inputs
  parsed <- parse_event_model(formula_or_list, data, block, durations)
  
  # Stage 2: Realise event terms
  terms <- realise_event_terms(parsed, sampling_frame, drop_empty, progress)
  
  # Stage 3: Build design matrix
  dm <- build_event_model_design_matrix(terms, sampling_frame, precision, parallel)
  
  # Assemble final event_model object
  ret <- list(
    terms          = terms,
    design_matrix  = dm,
    blockids       = parsed$blockids, # Store canonicalized blockids
    sampling_frame = sampling_frame,
    contrasts      = list(), # Contrasts are now primarily within terms via hrfspec
    model_spec     = list( # Store essential original inputs for reference?
        formula_or_list = formula_or_list, # Keep original spec
        data_names = names(data),
        block_arg = block,          # Keep original block arg
        durations_arg = durations,  # Keep original duration arg
        drop_empty = drop_empty,
        precision = precision
    )
  )
  class(ret) <- c("event_model", "list")
  ret
}

## ============================================================================
## Section 1.5: Deprecated Constructors (REMOVED)
## ============================================================================

# Deprecated event_model.formula removed.
# Deprecated event_model.list removed.
# Deprecated create_event_model removed.


#' @export
blocklens.event_model <- function(x, ...) {
  fmrihrf::blocklens(x$sampling_frame)
}

#' @export
blockids.event_model <- function(x) {
  x$blockids
}

#' @export
#' @rdname design_matrix
design_matrix.event_model <- function(x, blockid = NULL, ...) {
  # New version simply returns the pre-computed matrix, subsetting if needed
  dm <- x$design_matrix
  if (is.null(blockid)) {
    dm
  } else {
    # The design matrix has one row per timepoint, not per event
    # We need to map blockid to timepoint indices using the sampling frame
    sampling_frame <- x$sampling_frame
    if (is.null(sampling_frame)) {
        warning("Sampling frame not stored in event_model; cannot subset design matrix by block.")
        return(dm)
    }
    
    # Get the block IDs for each timepoint from the sampling frame
    timepoint_blockids <- fmrihrf::blockids(sampling_frame)
    if (is.null(timepoint_blockids)) {
        warning("Cannot determine timepoint block IDs from sampling frame.")
        return(dm)
    }
    
    # Find timepoints that belong to the requested blocks
    keep_rows <- timepoint_blockids %in% blockid
    if (!any(keep_rows)) {
        warning("Specified blockid(s) not found in the sampling frame.")
        # Return empty tibble with same colnames
        return(dm[0, , drop = FALSE]) 
    }
    dm[keep_rows, , drop = FALSE]
  }
}

#' @export
terms.event_model <- function(x, ...) {
  x$terms
}

#' @export
conditions.event_model <- function(x, drop.empty = TRUE, expand_basis = FALSE, ...) {
  unlist(lapply(terms(x), function(t) {
    conditions(t, drop.empty = drop.empty, expand_basis = expand_basis, ...)
  }), use.names = FALSE)
}

#' @export
#' @rdname contrast_weights
contrast_weights.convolved_term <- function(x, ...) {
  lapply(x$contrasts, function(cspec) {
    if (!is.null(cspec))
      contrast_weights(cspec, x)
  })
}

#' @export
#' @rdname Fcontrasts
Fcontrasts.convolved_term <- function(x, ...) {
  Fcontrasts(x$evterm)
}

#' @export
#' @rdname contrast_weights
contrast_weights.fmri_model <- function(x, ...) { 
  contrast_weights(x$event_model) 
}

#' @export
term_names.event_model <- function(x) {
  xt <- terms(x)
  unlist(lapply(xt, function(term) term$varname))
}

#' @export
#' @rdname contrast_weights
contrast_weights.event_model <- function(x, ...) {
  tnames <- term_names(x)
  tind <- attr(x$design_matrix, "col_indices") 
  ncond <- ncol(x$design_matrix)
  if (is.null(tind)) {
      stop("Cannot compute contrast weights: design matrix missing 'col_indices' attribute.")
  }
  
  ret <- lapply(seq_along(terms(x)), function(i) {
    # Get contrasts defined *within* the term (from hrfspec)
    term_contrasts <- attr(terms(x)[[i]], "hrfspec")$contrasts
    term_indices_vec <- tind[[ names(terms(x))[i] ]]
    
    if (!is.null(term_contrasts) && length(term_contrasts) > 0) {
      # Ensure term_contrasts list is treated as a contrast_set for dispatch
      cset <- term_contrasts
      if (!inherits(cset, "contrast_set")) class(cset) <- c("contrast_set", class(cset))
      cwlist_local <- contrast_weights(cset, terms(x)[[i]])
      
      # Map local weights to full design matrix
      ret_mapped <- lapply(cwlist_local, function(cw) {
          if(is.null(cw)) return(NULL) # Skip if contrast calculation failed
          
          # Check consistency
          if (nrow(cw$weights) != length(term_indices_vec)){
              warning(sprintf("Contrast '%s' for term '%s' has %d rows, but term has %d columns in design matrix. Skipping.", 
                             cw$name, names(terms(x))[i], nrow(cw$weights), length(term_indices_vec)), call.=FALSE)
              return(NULL)
          }
          
          out <- matrix(0, ncond, ncol(cw$weights))
          colnames(out) <- colnames(cw$weights) # Preserve contrast names if multiple columns
          rownames(out) <- colnames(x$design_matrix)
          out[term_indices_vec, ] <- cw$weights
          
          # Add attributes/update structure
          cw$offset_weights <- out
          cw$condnames <- colnames(x$design_matrix) # Update condnames to full set
          attr(cw, "term_indices") <- as.vector(term_indices_vec)
          cw
      })
      
      # Filter out NULLs from failed mappings
      ret_mapped <- ret_mapped[!sapply(ret_mapped, is.null)]
      
      # Set names based on original contrast list
      if (length(ret_mapped) > 0) {
         cnames <- sapply(ret_mapped, function(cw) cw$name)
         names(ret_mapped) <- cnames
      }
      ret_mapped
      
    } else { # No contrasts defined within this term
        NULL
    }
  })
  
  names(ret) <- tnames
  # Filter out terms that had no contrasts defined
  ret <- ret[!sapply(ret, is.null)]
  
  # Flatten the nested structure to match test expectations
  # Convert from ret[[term]][[contrast]] to ret[["term.contrast"]]
  flattened <- unlist(ret, recursive = FALSE)
  if (length(flattened) > 0) {
    # Create flattened names like "term.contrast"  
    names(flattened) <- paste0(
      rep(names(ret), sapply(ret, length)), 
      ".", 
      unlist(sapply(ret, names))
    )
  }
  
  return(flattened)
}

#' @export
Fcontrasts.event_model <- function(x, ...) {
  tnames <- term_names(x)
  tind <- attr(x$design_matrix, "col_indices") 
  ncond <- ncol(x$design_matrix)
  if (is.null(tind)) {
      stop("Cannot compute Fcontrasts: design matrix missing 'col_indices' attribute.")
  }
  
  ret_list <- lapply(seq_along(terms(x)), function(i) {
    term_i <- terms(x)[[i]]
    term_indices_vec <- tind[[ names(terms(x))[i] ]]
    
    # Calculate Fcontrasts relative to the term itself
    fcon_local <- Fcontrasts(term_i)
    
    if (!is.null(fcon_local) && length(fcon_local) > 0) {
        ret_mapped <- lapply(names(fcon_local), function(con_name) {
            cw <- fcon_local[[con_name]]
            
            # Check consistency
            if (nrow(cw) != length(term_indices_vec)){
                warning(sprintf("F-contrast '%s' for term '%s' has %d rows, but term has %d columns in design matrix. Skipping.", 
                                con_name, names(terms(x))[i], nrow(cw), length(term_indices_vec)), call.=FALSE)
                return(NULL)
            }
            
            out <- matrix(0, ncond, ncol(cw))
            colnames(out) <- colnames(cw)
            rownames(out) <- colnames(x$design_matrix)
            out[term_indices_vec, ] <- cw
            
            # Add attributes/update structure if needed (Fcontrast might just be matrix)
            attr(out, "term_indices") <- as.vector(term_indices_vec)
            out
        })
      
        # Filter out NULLs from failed mappings
        ret_mapped <- ret_mapped[!sapply(ret_mapped, is.null)]
        
        # Set names based on original contrast names, prefixed with term name
        if (length(ret_mapped) > 0) {
            prefix <- tnames[i]
            # Check if sapply(fcon_local, non_null) is necessary or exists
            valid_local_cons <- !sapply(fcon_local, is.null) # Simple check for non-NULL
            names(ret_mapped) <- paste0(prefix, "#", names(fcon_local)[valid_local_cons]) 
        }
         ret_mapped
         
    } else { # No Fcontrasts for this term
        NULL
    }
    
  })

  # Drop NULL entries for terms without F-contrasts
  ret_list <- ret_list[!sapply(ret_list, is.null)]

  # Concatenate without modifying individual elements
  ret <- if (length(ret_list)) do.call(c, ret_list) else list()
  ret
}

#' Print an Event Model
#'
#' Provides a concise summary of an `event_model` object using `cli`.
#'
#' @param x An `event_model` object.
#' @param ... Additional arguments (unused).
#'
#' @importFrom cli cli_h1 cli_div cli_text cli_h2 cli_ul cli_li cli_end cli_verbatim cli_alert_info
#' @importFrom utils head
#' @export
#' @rdname print
print.event_model <- function(x, ...) {
  n_terms <- length(x$terms)
  term_names <- names(x$terms)
  n_events <- length(x$blockids)
  n_blocks <- length(unique(x$blockids))
  total_scans <- sum(fmrihrf::blocklens(x$sampling_frame))
  dm <- design_matrix(x)
  dm_dims <- dim(dm)

  cli::cli_h1("fMRI Event Model")
  cli::cli_div(theme = list(
    span.info = list(color = "blue"),
    span.detail = list(color = "grey")
  ))

  cli::cli_text("{.info Number of Terms:} {n_terms}")
  cli::cli_text("{.info Number of Events:} {n_events}")
  cli::cli_text("{.info Number of Blocks:} {n_blocks}")
  cli::cli_text("{.info Total Scans:} {total_scans}")
  cli::cli_text("{.info Design Matrix Dimensions:} {dm_dims[1]} x {dm_dims[2]}")

  if (n_terms > 0) {
    cli::cli_h2("Terms")
    cli::cli_ul()
    for (i in seq_len(n_terms)) {
      tclass <- class(x$terms[[i]])[1]
      cli::cli_li("{.field {term_names[i]}} ({.cls {tclass}})")
    }
    cli::cli_end()
  }

  if (prod(dm_dims) > 0) {
    cli::cli_h2("Design Matrix Preview")
    dm_preview <- head(dm[, seq_len(min(4, dm_dims[2])), drop = FALSE], 3)
    dm_char <- apply(dm_preview, 2, function(col) sprintf("%7.3f", col))
    rownames(dm_char) <- paste("Scan", seq_len(nrow(dm_char)))
    preview_text <- capture.output(print(dm_char, quote = FALSE))
    cli::cli_verbatim(paste("  ", preview_text, collapse = "\n"))
    if (dm_dims[1] > nrow(dm_preview) || dm_dims[2] > ncol(dm_preview)) {
      cli::cli_text("  {.detail ...}")
    }
  } else {
    cli::cli_alert_info("Design matrix is empty.")
  }

  contrasts_info <- tryCatch(contrast_weights(x), error = function(e) NULL)
  if (!is.null(contrasts_info) && length(contrasts_info) > 0) {
    cli::cli_h2("Contrasts")
    cli::cli_ul()
    for (con in names(contrasts_info)) {
      cli::cli_li("{.field {con}}")
    }
    cli::cli_end()
  }

  cli::cli_end()
  invisible(x)
}