#' @include deprecated.R

# Generics ----

#' @title Constructors for `bru_mapper` objects
#' @param \dots Arguments passed on to sub-methods, or used for special
#'   purposes, see details for each function below.
#' @export
#' @seealso [bru_mapper_generics] for generic methods,
#' the individual mapper pages for special method implementations, and
#' [bru_get_mapper] for hooks to extract mappers from latent model object
#' class objects.
#' @describeIn bru_mapper
#' Generic mapper S3 constructor, used for constructing mappers for special
#' objects. See below for details of the default constructor
#' [bru_mapper_define()] that can be used to define new mapper clesses in user
#' code. To extracting mappers for latent component models, see
#' [bru_get_mapper()].
#'
#' @returns
#' A `bru_mapper` object.
#' @family mappers
#' @examples
#' mapper <- bm_index(5)
#' ibm_jacobian(mapper, input = c(1, 3, 4, 5, 2))
bru_mapper <- function(...) {
  UseMethod("bru_mapper")
}

#' Generic methods for bru_mapper objects
#'
#' @description
#' A `bru_mapper` sub-class implementation must provide an
#' [ibm_jacobian()] method. If the model size 'n' and definition
#' values 'values' are stored in the object itself, default methods [ibm_n()]
#' and [ibm_values()] are
#' available. Otherwise the
#' [ibm_n()] and [ibm_values()] methods also need to be provided.
#'
#' @seealso [bru_mapper] for constructor methods, and
#' [bru_get_mapper] for hooks to extract mappers from latent model object
#' class objects.
#'
#' @name bru_mapper_generics
#' @family mapper methods
NULL


#' @param mapper A mapper S3 object, inheriting from `bru_mapper`.
#' @param inla_f logical; when `TRUE` for `ibm_n()` and `ibm_values()`, the
#' result must be compatible with the `INLA::f(...)` and corresponding
#' `INLA::inla.stack(...)` constructions.  For `ibm_{eval,jacobian,linear}`,
#' the `input` interpretation may be different.
#' Implementations do not normally need to do anything different, except
#' for mappers of the type needed for hidden multicomponent models such
#' as "bym2", which can be handled by `bm_collect`.
#' @param \dots Arguments passed on to other methods
#' @title Size of the latent vector of a mapping
#' @description
#' Implementations must return the size of the latent vector
#' being mapped to.
#' @export
#' @inheritParams ibm_n_output
#' @family mapper methods
ibm_n <- function(mapper, inla_f = FALSE, ...) {
  UseMethod("ibm_n")
}

#' @title Output size of a mapping
#' @description
#' Implementations must return an integer denoting the
#' mapper output length.
#' The default implementation returns `NROW(input)`.
#' Mappers such as `bm_multi` and `bm_collect`,
#' that can accept `list()` inputs require their own method implementations.
#' @param state A vector of latent state values for the mapping,
#' of length `ibm_n(mapper, inla_f = FALSE)`
#' @param input Data input for the mapper.
#' @export
#' @family mapper methods
#' @inheritParams ibm_n
#' @inheritParams ibm_n_output
ibm_n_output <- function(mapper, input, state = NULL, inla_f = FALSE, ...) {
  UseMethod("ibm_n_output")
}


#' @title Value vector for a mapping
#' @description
#' When `inla_f=TRUE`, implementations must return a vector that
#' would be interpretable by an `INLA::f(..., values = ...)` specification.
#' The exception is the method for `bm_multi`, that returns a
#' multi-column data frame if `multi=TRUE`.
#' @export
#' @family mapper methods
#' @inheritParams ibm_n
ibm_values <- function(mapper, inla_f = FALSE, ...) {
  UseMethod("ibm_values")
}


#' @title Check if a mapper is linear/affine
#' @description
#' Implementations must return `TRUE` or `FALSE`.
#' If `TRUE` (returned by the default method unless the mapper
#' contains an `is_linear` variable), users of the mapper
#' may assume the mapper is linear/affine.
#' @export
#' @family mapper methods
#' @inheritParams ibm_n
ibm_is_linear <- function(mapper, ...) {
  UseMethod("ibm_is_linear")
}
#' @title Check if a mapper is rowwise
#' @description
#' Implementations must return `TRUE` or `FALSE`.
#' If `TRUE` (returned by the default method unless the mapper
#' contains an `is_rowwise` variable), users of the mapper
#' may assume the mapper uses its inputs in "rowwise" manner, so that
#' blockwise evaluation is always possible.
#' @export
#' @family mapper methods
#' @inheritParams ibm_n
ibm_is_rowwise <- function(mapper, ...) {
  UseMethod("ibm_is_rowwise")
}

#' @title Jacobian of a mapper
#' @description
#' Implementations must return a (sparse) matrix of size
#' `ibm_n_output(mapper, input, inla_f)`
#' by `ibm_n(mapper, inla_f = FALSE)`. The `inla_f=TRUE` argument should
#' only affect the allowed type of input format.
#' @export
#' @family mapper methods
#' @param sub_lin Internal, optional pre-computed sub-mapper information
#' @inheritParams ibm_n
#' @inheritParams ibm_n_output
ibm_jacobian <- function(mapper, input, state = NULL, inla_f = FALSE, ...) {
  UseMethod("ibm_jacobian")
}


#' @title Compute a mapper linearisation
#' @description
#' Implementations must return a [bm_taylor] object
#' The linearisation information includes `offset`, `jacobian`, and `state0`.
#' The state information indicates for which state the `offset` was evaluated,
#' with `NULL` meaning all-zero.
#' The linearised mapper output is defined as
#' ```
#' effect(input, state) =
#'   offset(input, state0) + jacobian(input, state0) %*% (state - state0)
#' ```
#' The default method calls [ibm_eval()] and [ibm_jacobian()] to generate
#' the needed information.
#' @export
#' @family mapper methods
#' @inheritParams ibm_n
#' @inheritParams ibm_n_output
ibm_linear <- function(mapper, input, state = NULL, ...) {
  UseMethod("ibm_linear")
}

#' @title Simplify a mapper
#' @description
#' Implementations must return a [bru_mapper] object.
#' The default method returns the result of [ibm_linear()] for linear/affine
#' mappers, and the original `mapper` for non-linear mappers.
#' @export
#' @family mapper methods
#' @inheritParams ibm_n
#' @inheritParams ibm_n_output
ibm_simplify <- function(mapper, input = NULL, state = NULL, ...) {
  UseMethod("ibm_simplify")
}

#' @title Evaluate a mapping
#' @description
#' Implementations must return a vector of length [ibm_n_output()].
#' The `input` contents must
#' be in a format accepted by [ibm_jacobian()]
#' for the mapper.
#' @export
#' @family mapper methods
#' @inheritParams ibm_n
#' @inheritParams ibm_n_output
ibm_eval <- function(mapper, input, state = NULL, ...) {
  UseMethod("ibm_eval")
}

#' @title Evaluate a mapper and its Jacobian
#' @description
#' Implementations must return a list with elements `offset` and `jacobian`.
#' The `input` contents must
#' be in a format accepted by [ibm_jacobian()]
#' for the mapper.
#' @export
#' @family mapper methods
#' @inheritParams ibm_n
#' @inheritParams ibm_n_output
ibm_eval2 <- function(mapper, input, state = NULL, ...) {
  UseMethod("ibm_eval2")
}

#' @title Names of submapper
#' @description
#' Implementations must return a character vector of sub-mapper names, or
#' `NULL`. Intended for providing information about multi-mappers and mapper
#' collections.
#' @export
#' @examples
#' # ibm_names
#' mapper <- bm_multi(list(A = bm_index(2), B = bm_index(2)))
#' ibm_names(mapper)
#' ibm_names(mapper) <- c("new", "names")
#' ibm_names(mapper)
#' @family mapper methods
#' @inheritParams ibm_n
ibm_names <- function(mapper) {
  UseMethod("ibm_names")
}

#' @param value a character vector of the same length as the number
#' of sub-mappers in the mapper
#' @describeIn ibm_names Set mapper names.
#' @export
#' @family mapper methods
#' @inheritParams ibm_n
`ibm_names<-` <- function(mapper, value) {
  UseMethod("ibm_names<-")
}


#' @title Find index subset of INLA visible states
#' @description
#' Implementations must return a logical vector of `TRUE/FALSE` for
#' the subset such that, given the full A matrix and values output,
#' `A[, subset, drop = FALSE]` and `values[subset]`
#' (or `values[subset, , drop = FALSE]` for data.frame values) are equal
#' to the `inla_f = TRUE` version of A and values. The default method uses
#' the `ibm_values` output to construct the subset indexing.
#' @export
#' @family mapper methods
#' @inheritParams ibm_n
ibm_inla_subset <- function(mapper, ...) {
  UseMethod("ibm_inla_subset")
}

#' @title Detect invalid input to a mapper
#' @description
#' Implementations should return a logical vector of length
#' `ibm_n_output(mapper, input, state, ...)` indicating which, if any,
#' output elements of `ibm_eval(mapper, input, state, ...)` are known to be
#' invalid.
#' For for multi/collect mappers, a list, when given a `multi=TRUE` argument.
#' @export
#' @family mapper methods
#' @inheritParams ibm_n
#' @inheritParams ibm_n_output
ibm_invalid_output <- function(mapper, input, state, ...) {
  UseMethod("ibm_invalid_output")
}


#' Methods for mapper lists
#'
#' `bru_mapper` lists can be combined into `bm_list` lists.
#' @name bm_list
#' @param \dots Objects to be combined.
#' @param x Object to convert
#' @export
as_bm_list <- function(x) {
  UseMethod("as_bm_list")
}

#' @rdname bm_list
#' @export
as_bm_list.list <- function(x) {
  stopifnot(all(vapply(
    x,
    function(x) inherits(x, "bru_mapper"),
    TRUE
  )))
  class(x) <- c("bm_list", "list")
  x
}

#' @rdname bm_list
#' @export
as_bm_list.bm_list <- function(x) {
  x
}

#' @rdname bm_list
#' @export
as_bm_list.bru_comp_list <- function(x) {
  if (is.null(names(x))) {
    names(x) <- vapply(x, function(y) {
      y[["label"]]
    }, "")
  }
  mappers <- as_bm_list(lapply(x, function(y) {
    y[["mapper"]]
  }))
  mappers
}

#' Methods for mapper extraction
#'
#' Extract a mapper from another object
#' @name as_bru_mapper
#' @param x Object to convert/extract
#' @returns A `bru_mapper` object
#' @export
as_bru_mapper <- function(x) {
  UseMethod("as_bru_mapper")
}

#' @rdname as_bru_mapper
#' @export
as_bru_mapper.bru_mapper <- function(x) {
  x
}

#' @rdname as_bru_mapper
#' @export
as_bru_mapper.bru_comp <- function(x) {
  x[["mapper"]]
}

#' @rdname as_bru_mapper
#' @export
#' @examples
#' # Extract a mapper from a `bru_subcomp` object
#' as_bru_mapper(bru_comp("x", x, mapper = bm_index(4))$main)
#'
as_bru_mapper.bru_subcomp <- function(x) {
  x[["mapper"]]
}

#' @export
#' @describeIn bm_list The `...` arguments should be `bru_mapper`
#' objects.
#' @returns A `bm_list` object
#' @examples
#' m <- c(A = bm_const(), B = bm_scale())
#' str(m)
#' str(m[2])
`c.bru_mapper` <- function(...) {
  as_bm_list(list(...))
}

#' @export
#' @describeIn bm_list The `...` arguments should be `bm_list`
#' objects.
`c.bm_list` <- function(...) {
  stopifnot(all(vapply(
    list(...),
    function(x) inherits(x, "bm_list"),
    TRUE
  )))
  object <- NextMethod()
  class(object) <- c("bm_list", "list")
  object
}

#' @export
#' @param x `bm_list` object from which to extract element(s)
#' @param i indices specifying elements to extract
#' @describeIn bm_list Extract sub-list
`[.bm_list` <- function(x, i) {
  object <- NextMethod()
  class(object) <- c("bm_list", "list")
  object
}


# Summaries ----

ibm_shortname <- function(mapper, ...) {
  cls1 <- class(mapper)[1]
  if (any(grepl("^bm_", cls1))) {
    shortname <- sub("^bm_", "", cls1)
  } else if (any(grepl("^bru_mapper_", cls1))) {
    shortname <- sub("^bru_mapper_", "", cls1)
  } else {
    shortname <- character(0)
  }
  if (nchar(shortname) == 0) {
    shortname <- "default"
  }
  shortname
}


#' @title mapper object summaries
#'
#' @param x Object to format/print
#' @param object Object to summarise
#' @param \dots Unused arguments
#' @param prefix character prefix for each line. Default `""`.
#' @param initial character prefix for the first line. Default `initial=prefix`.
#' @param depth The recursion depth for multi/collection/pipe mappers. Default
#'   1, to only show the collection, and not the contents of the sub-mappers.
#' @export
#' @method format bru_mapper
#' @rdname bm_summary
format.bru_mapper <- function(x, ...,
                              prefix = "",
                              initial = prefix,
                              depth = 1) {
  mapper_new <- make_bm_class_from_old(x)
  if (!is.null(mapper_new)) {
    return(format(
      mapper_new,
      ...,
      prefix = prefix,
      initial = initial,
      depth = depth
    ))
  }
  text <- paste0(initial, ibm_shortname(x))
  if (ibm_input_available(x)) {
    inp <- ibm_input_get(x)
    text <- glue("{text}({format(inp, label.override = '')})")
  }
  as.character(text)
}

#' @export
#' @method format bm_list
#' @param collapse character or NULL, as in [base::paste()].
#' @param labels logical; if TRUE, include mapper names or numerical indices.
#'   Default `TRUE`
#'
#' @rdname bm_summary
format.bm_list <- function(x, ...,
                           prefix = "",
                           initial = prefix,
                           depth = 1,
                           collapse = ", ",
                           labels = TRUE) {
  if (labels) {
    paste0(
      vapply(
        seq_along(x),
        function(k) {
          nm <- names(x)[k]
          if (is.null(nm) || (identical(nm, ""))) {
            nm <- as.character(k)
          }
          paste0(
            nm,
            " = ",
            format(
              x[[k]],
              prefix = prefix,
              initial = initial,
              depth = depth
            )
          )
        },
        ""
      ),
      collapse = collapse
    )
  } else {
    paste0(
      vapply(
        seq_along(x),
        function(k) {
          paste0(
            format(
              x[[k]],
              prefix = prefix,
              initial = initial,
              depth = depth
            )
          )
        },
        ""
      ),
      collapse = collapse
    )
  }
}

#' @rdname bm_summary
#' @export
#' @method summary bru_mapper
summary.bru_mapper <- function(object, ...,
                               prefix = "",
                               initial = prefix,
                               depth = 1) {
  mapper_new <- make_bm_class_from_old(object)
  if (!is.null(mapper_new)) {
    return(summary(
      mapper_new,
      ...,
      prefix = prefix,
      initial = initial,
      depth = depth
    ))
  }
  summary_object <- format(
    object,
    prefix = prefix,
    initial = initial,
    depth = depth
  )
  class(summary_object) <- c(
    "summary_bru_mapper",
    class(summary_object)
  )
  summary_object
}

#' @export
#' @method format bm_multi
#' @rdname bm_summary
#' @examples
#' mapper <-
#'   bm_pipe(
#'     list(
#'       bm_multi(list(
#'         A = bm_index(2),
#'         B = bm_index(3)
#'       )),
#'       bm_index(2)
#'     )
#'   )
#' summary(mapper, depth = 2)
format.bm_multi <- function(x, ...,
                            prefix = "",
                            initial = prefix,
                            depth = 1) {
  txt <- NextMethod()
  if (depth <= 0) {
    return(txt)
  }
  sub_prefix <- paste0(prefix, "      ")
  txt <-
    paste0(
      txt,
      "(",
      format(x[["mappers"]],
        prefix = sub_prefix,
        initial = "",
        depth = depth - 1,
        collapse = ", "
      ),
      ")"
    )
  txt
}


#' @export
#' @method format bm_pipe
#' @rdname bm_summary
format.bm_pipe <- function(x, ...,
                           prefix = "",
                           initial = prefix,
                           depth = 1) {
  txt <- NextMethod()
  if (depth <= 0) {
    return(txt)
  }
  sub_prefix <- paste0(prefix, "      ")
  txt <-
    paste0(
      txt,
      " = ",
      format(x[["mappers"]],
        prefix = sub_prefix,
        initial = "",
        depth = depth - 1,
        collapse = " -> ",
        labels = FALSE
      )
    )
  txt
}

#' @export
#' @method format bm_collect
#' @rdname bm_summary
format.bm_collect <- function(x, ...,
                              prefix = "",
                              initial = prefix,
                              depth = 1) {
  txt <- NextMethod()
  if (depth <= 0) {
    return(txt)
  }
  sub_prefix <- paste0(prefix, "      ")
  txt <-
    paste0(
      txt,
      "(",
      paste0(
        vapply(
          seq_along(x[["mappers"]]),
          function(k) {
            nm <- names(x[["mappers"]])[k]
            if (is.null(nm)) {
              nm <- as.character(k)
            }
            paste0(
              nm,
              " = ",
              summary(
                x[["mappers"]][[k]],
                prefix = sub_prefix,
                initial = "",
                depth = depth - 1
              ),
              if (x[["hidden"]] && (k > 1)) {
                "(hidden)"
              } else {
                NULL
              }
            )
          },
          ""
        ),
        collapse = ", "
      ),
      ")"
    )
  txt
}

#' @export
#' @method format bm_sum
#' @rdname bm_summary
format.bm_sum <- function(x, ...,
                          prefix = "",
                          initial = prefix,
                          depth = 1) {
  txt <- NextMethod()
  if (depth <= 0) {
    return(txt)
  }
  sub_prefix <- paste0(prefix, "      ")
  txt <-
    paste0(
      txt,
      "(",
      format(
        x[["mappers"]],
        prefix = sub_prefix,
        initial = "",
        depth = depth - 1,
        collapse = ", "
      ),
      ")"
    )
  txt
}

#' @export
#' @method format bm_repeat
#' @rdname bm_summary
#' @examples
#' mapper <-
#'   bm_repeat(
#'     bm_multi(
#'       list(
#'         A = bm_index(2),
#'         B = bm_index(3)
#'       )
#'     ),
#'     3
#'   )
#' summary(mapper)
#' summary(mapper, depth = 0)
format.bm_repeat <- function(x, ...,
                             prefix = "",
                             initial = prefix,
                             depth = 1) {
  txt <- NextMethod()
  sub_prefix <- paste0(prefix, "      ")
  txt <-
    paste0(
      txt,
      "(",
      x[["n_rep"]],
      " x ",
      paste0(
        format(
          x[["mapper"]],
          prefix = sub_prefix,
          initial = "",
          depth = depth
        )
      ),
      if (isTRUE(x[["interleaved"]])) {
        ", interleaved"
      } else {
        NULL
      },
      ")"
    )
  txt
}

#' @export
#' @method format bm_reparam
#' @rdname bm_summary
#' @examples
#' mapper <-
#'   bm_reparam(
#'     bm_multi(
#'       list(
#'         A = bm_index(2),
#'         B = bm_index(3)
#'       )
#'     ),
#'     matrix(1:36, nrow = 6)
#'   )
#' summary(mapper)
#' summary(mapper, depth = 0)
format.bm_reparam <- function(x, ...,
                              prefix = "",
                              initial = prefix,
                              depth = 1) {
  txt <- NextMethod()
  sub_prefix <- paste0(prefix, "      ")
  txt <-
    paste0(
      txt,
      "(",
      paste0(
        format(
          x[["mapper"]],
          prefix = sub_prefix,
          initial = "",
          depth = depth
        )
      ),
      ")"
    )
  txt
}


#' @param sep character; separator for printing the summary.
#' @export
#' @method print summary_bru_mapper
#' @rdname bm_summary
print.summary_bru_mapper <- function(x, ..., sep = "\n") {
  cat(x, sep = sep)
  invisible(x)
}

#' @export
#' @method print bru_mapper
#' @rdname bm_summary
print.bru_mapper <- function(x, ...,
                             sep = "\n",
                             prefix = "",
                             initial = prefix,
                             depth = 1) {
  cat(
    format(x,
      prefix = "",
      initial = prefix,
      depth = depth
    ),
    sep = sep
  )
  invisible(x)
}

#' @export
#' @method print bm_list
#' @rdname bm_summary
print.bm_list <- function(x, ...,
                          sep = "\n",
                          prefix = "",
                          initial = prefix,
                          depth = 1,
                          labels = TRUE,
                          collapse = ", ") {
  cat(
    format(x,
      prefix = "",
      initial = prefix,
      depth = depth,
      collapse = collapse
    ),
    sep = sep
  )
  invisible(x)
}


# MAPPERS ----
## Constructor ----

#' @param mapper For `bru_mapper_define`, a prototype mapper object, see
#'   Details.
#' @param new_class If non-`NULL`, this is added at the front of the class
#'   definition
#' @param remove_class If non-`NULL`, this class or classes is removed from the
#'   class definition before adding the `new_class` names. Default is `"list"`.
#'
#' @describeIn bru_mapper Adds the `new_class` and `"bru_mapper"` class names to
#'   the inheritance list for the input `mapper` object, unless the object
#'   already inherits from these.
#'
#' To register mapper classes and methods in scripts, use `.S3method()`
#' to register the methods, e.g.
#' `.S3method("ibm_jacobian", "my_mapper_class", ibm_jacobian.my_mapper_class)`.
#'
#' In packages with `Suggests: inlabru`, add method information for delayed
#' registration, e.g.:
#' ```
#' #' @rawNamespace S3method(inlabru::bru_get_mapper, inla_rspde)
#' #' @rawNamespace S3method(inlabru::ibm_n, bru_mapper_inla_rspde)
#' #' @rawNamespace S3method(inlabru::ibm_values, bru_mapper_inla_rspde)
#' #' @rawNamespace S3method(inlabru::ibm_jacobian, bru_mapper_inla_rspde)
#' ```
#' or before each method, use `@exportS3Method`:
#' ```
#' #' @exportS3Method inlabru::bru_get_mapper
#' ```
#' etc., which semi-automates it.
#'
#' @export
bru_mapper_define <- function(mapper,
                              new_class = NULL,
                              remove_class = "list",
                              ...) {
  old_cls <- class(mapper)

  if ((length(remove_class) > 0) && (length(old_cls) > 0)) {
    old_cls <- old_cls[!(old_cls %in% remove_class)]
  }

  new_class <- c(new_class, "bru_mapper")
  for (cls in rev(new_class)) {
    if (!(cls %in% old_cls)) {
      old_cls <- c(cls, old_cls)
    }
  }
  class(mapper) <- old_cls

  mapper
}

## Default methods ----

# Conversion method from bru_mapper_<type> to bm_<type> class names.
# Only converts known inlabru mapper classes.
# Returns NULL if no conversion was needed.
make_bm_class_from_old <- function(mapper) {
  cls <- class(mapper)
  idx <- grepl(paste0("^bru_mapper_",
    c(
      "aggregate",
      "collect",
      "const",
      "factor",
      "fmesher",
      "harmonics",
      "index",
      "linear",
      "logsumexp",
      "marginal",
      "matrix",
      "mesh_B",
      "multi",
      "pipe",
      "repeat",
      "scale",
      "shift",
      "sum",
      "taylor",
      "fm_mesh_1d",
      "fm_mesh_2d",
      "inla_mesh_1d",
      "inla_mesh_2d"
    ),
    "$",
    collapse = "|"
  ), cls)
  if (any(idx)) {
    idx <- min(which(idx))
    # Remove the bru_mapper_ prefix, and add bm_ prefix
    cls[idx] <- sub("^bru_mapper_", "bm_", cls[idx])
    class(mapper) <- cls

    if (inherits(
      mapper,
      c("bm_aggregate", "bm_multi", "bm_collect", "bm_sum", "bm_pipe")
    )) {
      mapper[["is_rowwise"]] <- FALSE
    }

    return(mapper)
  }
  NULL
}

#' @describeIn ibm_n
#' Returns a non-null element 'n' from the
#' mapper object, and gives an error if it doesn't exist. If `inla_f=TRUE`,
#' first checks for a 'n_inla' element.
#' @export
ibm_n.default <- function(mapper, inla_f = FALSE, ...) {
  mapper_new <- make_bm_class_from_old(mapper)
  if (!is.null(mapper_new)) {
    return(ibm_n(mapper_new, inla_f = inla_f, ...))
  }
  if (inla_f && !is.null(mapper[["n_inla"]])) {
    mapper[["n_inla"]]
  } else if (!is.null(mapper[["n"]])) {
    mapper[["n"]]
  } else {
    stop(paste0(
      "Default 'ibm_n()' method called but mapper doesn't have an 'n' element;",
      "\nClass: ",
      paste0(class(mapper), collapse = ", ")
    ))
  }
}

#' @export
#' @describeIn ibm_n_output Returns `NROW(input)`
ibm_n_output.default <- function(mapper,
                                 input,
                                 state = NULL,
                                 inla_f = FALSE,
                                 ...) {
  mapper_new <- make_bm_class_from_old(mapper)
  if (!is.null(mapper_new)) {
    return(ibm_n_output(
      mapper_new,
      input = input,
      state = state,
      inla_f = inla_f,
      ...
    ))
  }
  NROW(input)
}


#' @describeIn ibm_values
#' Returns a non-null element
#' 'values' from the mapper object, and `seq_len(ibm_n(mapper))` if
#' it doesn't exist.
#' @export
ibm_values.default <- function(mapper, inla_f = FALSE, ...) {
  mapper_new <- make_bm_class_from_old(mapper)
  if (!is.null(mapper_new)) {
    return(ibm_values(mapper_new, inla_f = inla_f, ...))
  }
  if (inla_f && !is.null(mapper[["values_inla"]])) {
    mapper[["values_inla"]]
  } else if (!inla_f && !is.null(mapper[["values"]])) {
    mapper[["values"]]
  } else {
    seq_len(ibm_n(mapper, inla_f = inla_f))
  }
}

#' @describeIn ibm_is_linear
#' Returns logical
#' `is_linear` from the mapper object if it exists, and otherwise `TRUE`.
#' @export
ibm_is_linear.default <- function(mapper, ...) {
  mapper_new <- make_bm_class_from_old(mapper)
  if (!is.null(mapper_new)) {
    return(ibm_is_linear(mapper_new, ...))
  }
  if (!is.null(mapper[["is_linear"]])) {
    mapper[["is_linear"]]
  } else {
    TRUE
  }
}

#' @describeIn ibm_is_rowwise
#' Returns logical
#' `is_rowwise` from the mapper object if it exists, and otherwise `TRUE`.
#' @export
ibm_is_rowwise.default <- function(mapper, ...) {
  mapper_new <- make_bm_class_from_old(mapper)
  if (!is.null(mapper_new)) {
    return(ibm_is_rowwise(mapper_new, ...))
  }
  if (!is.null(mapper[["is_rowwise"]])) {
    mapper[["is_rowwise"]]
  } else {
    TRUE
  }
}

#' @describeIn ibm_jacobian
#' Mapper classes must implement their own `ibm_jacobian` method.
#' @export
ibm_jacobian.default <- function(mapper, input, state = NULL, ...) {
  if (is.null(mapper)) {
    stop(paste0(
      "NULL mapper detected."
    ))
  }
  mapper_new <- make_bm_class_from_old(mapper)
  if (!is.null(mapper_new)) {
    return(ibm_jacobian(mapper_new, input = input, state = state, ...))
  }
  stop(paste0(
    "Missing ibm_jacobian() method for class '",
    paste0(class(mapper), collapse = ", "),
    "'."
  ))
}

#' @describeIn ibm_linear
#' Calls [ibm_eval2()] and returns a [bm_taylor] object.
#' @returns A [bm_taylor] object.
#' The `state0` information in the affine mapper indicates for which state
#' the `offset` was evaluated; The affine mapper output is defined as
#' ```
#' effect(input, state) =
#'   offset(input, state0) + jacobian(input, state0) %*% (state - state0)
#' ```
#' @export
ibm_linear.default <- function(mapper, input, state, ...) {
  mapper_new <- make_bm_class_from_old(mapper)
  if (!is.null(mapper_new)) {
    return(ibm_linear(mapper_new, input = input, state = state, ...))
  }
  eval2 <- ibm_eval2(mapper, input = input, state = state, ...)
  bm_taylor(
    offset = eval2$offset,
    jacobian = eval2$jacobian,
    state0 = state,
    values_mapper = mapper
  )
}

#' @describeIn ibm_simplify
#' Calls [ibm_linear()] for linear mappers, and returns the original mapper
#' for non-linear mappers.
#' @returns The original mapper is returned for non-linear mappers, and the
#' output of [ibm_linear()] is returned for linear mappers.
#' @export
ibm_simplify.default <- function(mapper, input = NULL, state = NULL, ...) {
  mapper_new <- make_bm_class_from_old(mapper)
  if (!is.null(mapper_new)) {
    return(ibm_simplify(mapper_new, input = input, state = state, ...))
  }
  if (ibm_is_linear(mapper)) {
    if (is.null(state)) {
      state <- rep(0, ibm_n(mapper, input = input, state = NULL, ...))
    }
    return(ibm_linear(mapper, input = input, state = state, ...))
  }
  mapper
}


#' @describeIn ibm_eval
#' Verifies that the mapper is linear
#' with [ibm_is_linear()], and then computes a linear mapping
#' as `ibm_jacobian(...) %*% state`.  When `state` is `NULL`,
#' a zero vector of length [ibm_n_output()] is returned.
#' @param jacobian
#' For [ibm_eval()] methods, an optional pre-computed Jacobian, typically
#' supplied by internal methods that already have the Jacobian.
#' @export
ibm_eval.default <- function(mapper,
                             input,
                             state = NULL,
                             ...,
                             jacobian = NULL) {
  mapper_new <- make_bm_class_from_old(mapper)
  if (!is.null(mapper_new)) {
    return(ibm_eval(
      mapper_new,
      input = input,
      state = state,
      ...,
      jacobian = jacobian
    ))
  }
  if (!ibm_is_linear(mapper)) {
    stop("Non-linear mappers must implement their own ibm_eval() method.")
  }

  val <- numeric(ibm_n_output(mapper, input, state = state, ...))

  if ((ibm_n(mapper) > 0) && !is.null(state)) {
    if (is.null(jacobian)) {
      jacobian <- ibm_jacobian(mapper, input = input, state = state, ...)
    }
    val <- val + as.vector(jacobian %*% state)
  }

  val
}


#' @describeIn ibm_eval2
#' Calls `jacobian <- ibm_jacobian(...)` and
#' `offset <- ibm_eval(..., jacobian = jacobian)`
#' and returns a list with elements `offset` and `jacobian`, as needed by
#' [ibm_linear.default()] and similar methods. Mapper classes can implement
#' their own `ibm_eval2` method if joint construction of evaluation and Jacobian
#' is more efficient than separate or sequential construction.
#' @export
ibm_eval2.default <- function(mapper, input, state, ...) {
  mapper_new <- make_bm_class_from_old(mapper)
  if (!is.null(mapper_new)) {
    return(ibm_eval2(mapper_new, input = input, state = state, ...))
  }
  jacobian <- ibm_jacobian(mapper, input, state, ...)
  offset <- ibm_eval(
    mapper,
    input,
    state,
    ...,
    jacobian = jacobian
  )
  list(offset = offset, jacobian = jacobian)
}


#' @export
#' @describeIn ibm_names Returns `NULL`
ibm_names.default <- function(mapper, ...) {
  mapper_new <- make_bm_class_from_old(mapper)
  if (!is.null(mapper_new)) {
    return(ibm_names(mapper_new, ...))
  }
  NULL
}


#' @describeIn ibm_inla_subset
#' Uses the [ibm_values`()] output to construct the inla subset indexing as the
#' difference between `inla_f=FALSE` and `inla_f=TRUE`. Extra arguments such as
#' `multi` are passed on to [ibm_values()]. This means it supports both regular
#' vector values and `multi=1` data.frame values.
#' @export
ibm_inla_subset.default <- function(mapper, ...) {
  mapper_new <- make_bm_class_from_old(mapper)
  if (!is.null(mapper_new)) {
    return(ibm_inla_subset(mapper_new, ...))
  }
  values_full <- ibm_values(mapper, inla_f = FALSE, ...)
  values_inla <- ibm_values(mapper, inla_f = TRUE, ...)
  if (is.data.frame(values_full)) {
    subset <- logical(NROW(values_full))
    if (length(subset) > 0) {
      subset[
        plyr::match_df(
          cbind(
            .inla_subset = seq_len(NROW(values_full)),
            values_full
          ),
          values_inla,
          on = names(values_full)
        )[[".inla_subset"]]
      ] <- TRUE
    }
  } else {
    subset <- base::match(values_full, values_inla, nomatch = 0) > 0
  }
  subset
}


#' @describeIn ibm_invalid_output
#' Returns an all-`FALSE` logical vector.
#' @export
ibm_invalid_output.default <- function(mapper, input, state, ...) {
  mapper_new <- make_bm_class_from_old(mapper)
  if (!is.null(mapper_new)) {
    return(ibm_invalid_output(mapper_new, input = input, state = state, ...))
  }
  rep(FALSE, ibm_n_output(mapper, input = input, state = state, ...))
}


## _fmesher ####

#' @title Mapper for general `fmesher` function space objects
#' @param mesh An `fmesher` object to map, supported by
#' [fmesher::fm_basis]`(mesh, input)` and
#' [fmesher::fm_dof]`(mesh)`.
#' @export
#' @description Creates a mapper for general `fmesher` function space objects.
#' @details Handles indexed mapping for all `fmesher` classes that support
#'   `fm_dof()` and `fm_basis()` methods. For non-indexed mapping of
#'   `fm_mesh_1d` objects, use `bru_mapper(mesh, indexed = FALSE)` which
#'   invokes the [bru_mapper.fm_mesh_1d()] method.
#' @returns A `bm_fmesher` object.
#' @rdname bm_fmesher
#' @seealso [bru_mapper], [bru_mapper_generics]
#' @family mappers
#'
#' @examples
#' m <- bm_fmesher(fmesher::fmexample$mesh)
#' ibm_n(m)
#' ibm_eval(m, as.matrix(expand.grid(-2:2, -2:2)), seq_len(ibm_n(m)))
#'
bm_fmesher <- function(mesh) {
  if (is.null(mesh)) {
    return(NULL)
  }
  mapper <- list(mesh = mesh)
  bru_mapper_define(mapper, new_class = "bm_fmesher")
}

#' @export
#' @rdname bm_fmesher
#' @param \dots Arguments passed on to [bm_fmesher()]
bru_mapper_fmesher <- function(...) {
  bm_fmesher(...)
}

#' @export
#' @describeIn ibm_n Returns the [fmesher::fm_dof()] value of the mesh being
#' mapped.
#'
ibm_n.bm_fmesher <- function(mapper, ...) {
  fmesher::fm_dof(mapper[["mesh"]])
}
#' @export
#' @describeIn ibm_values Returns an index vector for the mesh basis functions.
#'
ibm_values.bm_fmesher <- function(mapper, ...) {
  seq_len(fmesher::fm_dof(mapper[["mesh"]]))
}
#' @export
#'
#' @describeIn ibm_jacobian Returns the [fmesher::fm_basis()] matrix of the
#' mesh being mapped.
ibm_jacobian.bm_fmesher <- function(mapper, input, ...) {
  if (is.null(input)) {
    return(Matrix::Matrix(0, 0, ibm_n(mapper)))
  }
  fmesher::fm_basis(mapper[["mesh"]], input)
}


## fm_mesh_2d ####

#' @describeIn bm_fmesher
#' Equivalent to calling [bm_fmesher()].
#' Note: Prior to version `2.12.0.9021`,
#'   this returned a `bru_mapper_fm_mesh_2d` object. Also see the note for
#'   [bru_mapper.fm_mesh_1d()].
#' @export
bru_mapper.fm_mesh_2d <- function(mesh, ...) {
  bm_fmesher(mesh, ...)
}

## The following methods are only used for old stored mapper objects

#' @title Deprecated methods
#' @description `r lifecycle::badge("deprecated")`
#'   Deprecated methods only used for old stored
#'   `bru_mapper_fm_mesh_1d` and `bru_mapper_fm_mesh_2d` objects.
#' @keywords internal
#' @name bm_fm_mesh_old
#' @rdname bm_fm_mesh_old
NULL

#' @export
#' @rdname bm_fm_mesh_old
ibm_n.bm_fm_mesh_2d <- function(mapper, ...) {
  fmesher::fm_dof(mapper[["mesh"]])
}
#' @export
#' @rdname bm_fm_mesh_old
ibm_values.bm_fm_mesh_2d <- function(mapper, ...) {
  seq_len(fmesher::fm_dof(mapper[["mesh"]]))
}
#' @export
#' @rdname bm_fm_mesh_old
ibm_jacobian.bm_fm_mesh_2d <- function(mapper, input, ...) {
  if (is.null(input)) {
    return(Matrix::Matrix(0, 0, ibm_n(mapper)))
  }
  fm_basis(mapper[["mesh"]], input)
}

## The following methods are only used for old stored mapper objects

#' @export
#' @rdname bm_fm_mesh_old
ibm_n.bm_inla_mesh_2d <- function(mapper, ...) {
  mapper[["mesh"]]$n
}
#' @export
#' @rdname bm_fm_mesh_old
ibm_values.bm_inla_mesh_2d <- function(mapper, ...) {
  seq_len(mapper[["mesh"]]$n)
}
#' @export
#' @rdname bm_fm_mesh_old
ibm_jacobian.bm_inla_mesh_2d <- function(mapper, input, ...) {
  if (is.null(input)) {
    return(Matrix::Matrix(0, 0, ibm_n(mapper)))
  }
  fm_basis(mapper[["mesh"]], input)
}


## fm_mesh_1d ####

#' @title Mapper for `fm_mesh_1d`
#' @param indexed logical; If `TRUE` (default), the `ibm_values()` output will
#'   be the integer indexing sequence for the latent variables (needed for
#'   `spde` models). If `FALSE`, points representative of the basis centres
#'   are returned (useful for an interpolator for `rw2` models and similar,
#'   for `fmesher` versions `>= 0.3.0.9002`).
#' @returns A `bm_fm_mesh_1d` or `bm_fmesher` object. The the
#'   general [bm_fmesher()] mapper handles all indexed `fmesher`
#'   objects.
#' @export
#' @description Create mapper for an `fm_mesh_1d` object
#' @param mesh An `fm_mesh_1d` object to use as a mapper
#' @rdname bm_fm_mesh_1d
#' @name bm_fm_mesh_1d
#' @inheritParams bru_mapper_generics
#' @inheritParams bm_fmesher
#' @seealso [bru_mapper], [bru_mapper_generics]
#' @family mappers
#' @examples
#' m <- bru_mapper(fmesher::fm_mesh_1d(c(1:3, 5, 7)))
#' ibm_values(m)
#' ibm_eval(m, 1:7, 1:5)
#'
#' m <- bru_mapper(fmesher::fm_mesh_1d(c(1:3, 5, 7)), indexed = FALSE)
#' ibm_values(m)
#' ibm_eval(m, 1:7, 1:5)
#'
#' m <- bru_mapper(
#'   fmesher::fm_mesh_1d(c(1:3, 5, 7), degree = 2, boundary = "free"),
#'   indexed = FALSE
#' )
#' ibm_values(m)
#' ibm_eval(m, 1:7, 1:6)
#'
bru_mapper.fm_mesh_1d <- function(mesh, indexed = TRUE, ...) {
  if (indexed) {
    mapper <- bm_fmesher(mesh)
    return(mapper)
  }
  mapper <- list(
    mesh = mesh,
    indexed = indexed
  )
  bru_mapper_define(mapper, new_class = "bm_fm_mesh_1d")
}

#' @export
#' @describeIn ibm_n Returns the [fmesher::fm_dof()] value of the mesh being
#' mapped.
#'
ibm_n.bm_fm_mesh_1d <- function(mapper, ...) {
  fmesher::fm_dof(mapper[["mesh"]])
}
#' @export
#' @describeIn ibm_values Returns an index vector into the basis functions for
#'   an indexed mapper. Otherwise, the `mid` values if present in the mesh being
#'   mapped, and otherwise returns the `loc` values of the mesh.
#'
ibm_values.bm_fm_mesh_1d <- function(mapper, ...) {
  if (mapper[["indexed"]]) {
    seq_len(fmesher::fm_dof(mapper[["mesh"]]))
  } else if (!is.null(mapper[["mesh"]][["mid"]])) {
    mapper[["mesh"]]$mid
  } else {
    mapper[["mesh"]]$loc
  }
}
#' @export
#' @describeIn ibm_jacobian Returns the [fmesher::fm_basis()] matrix of the
#' mesh being mapped.
#'
ibm_jacobian.bm_fm_mesh_1d <- function(mapper, input, ...) {
  if (is.null(input)) {
    return(Matrix::Matrix(0, 0, ibm_n(mapper)))
  }
  # Note: Handles NA input from fmesher 0.2.0.9002, with fixes in 0.4.0.9003
  fm_basis(mapper[["mesh"]], input)
}

## The following methods are only used for old stored mapper objects

#' @export
#' @rdname bm_fm_mesh_old
ibm_n.bm_inla_mesh_1d <- function(mapper, ...) {
  mapper[["mesh"]]$m
}
#' @export
#' @rdname bm_fm_mesh_old
ibm_values.bm_inla_mesh_1d <- function(mapper, ...) {
  if (mapper[["indexed"]]) {
    seq_len(mapper[["mesh"]]$m)
  } else {
    mapper[["mesh"]]$loc
  }
}
#' @export
#' @rdname bm_fm_mesh_old
ibm_jacobian.bm_inla_mesh_1d <- function(mapper, input, ...) {
  if (is.null(input)) {
    return(Matrix::Matrix(0, 0, ibm_n(mapper)))
  }
  fm_basis(mapper[["mesh"]], input)
}

## _index ####

#' @title Mapper for indexed variables
#' @param n Size of a model for `bm_index`
#' @export
#' @description Create a an indexing mapper
#' @rdname bm_index
#' @seealso [bru_mapper], [bru_mapper_generics]
#' @family mappers
#'
#' @examples
#' m <- bm_index(4)
#' ibm_eval(m, -2:6, 1:4)
bm_index <- function(n = 1L) {
  bru_mapper_define(list(n = n), new_class = "bm_index")
}

#' @export
#' @rdname bm_index
#' @param \dots Arguments passed on to [bm_index()]
bru_mapper_index <- function(...) {
  bm_index(...)
}

#' @export
#' @describeIn ibm_invalid_output Returns `TRUE` out-of-range input values.
#'   Returns `FALSE` for in-range and `NA` input values (to support NA giving
#'   zero effect).
#'
ibm_invalid_output.bm_index <- function(mapper, input, state, ...) {
  # Allow NA to give effect zero
  nok <- !is.na(input)
  nok[nok] <-
    !((input[nok] >= 1) &
      (input[nok] <= ibm_n(mapper)))
  nok
}


#' @export
#' @rdname ibm_jacobian
#'
ibm_jacobian.bm_index <- function(mapper, input, state, ...) {
  if (is.null(input)) {
    return(Matrix::Matrix(0, 0, ibm_n(mapper)))
  }
  nok <- ibm_invalid_output(mapper, input = input, state = state, ...)
  ok <- which(!is.na(input) & !nok)
  Matrix::sparseMatrix(
    i = ok,
    j = input[ok],
    x = rep(1, length(ok)),
    dims = c(NROW(input), ibm_n(mapper))
  )
}

## _taylor ####

#' @param offset For `bm_taylor`, an offset vector giving the value of
#' the linearisation at `state0`.
#' May be `NULL`, interpreted as an all-zero vector of length determined by
#' a non-null Jacobian.
#' @param jacobian For `bm_taylor()`, the Jacobian matrix,
#' evaluated at `state0`, or, a named list of such matrices.
#' May be `NULL` or an empty list, for a constant mapping.
#' @param state0 For `bm_taylor`, the reference `state` for the
#' linearisation, or a list of such states matching the `jacobian` list.
#' `NULL` is interpreted as 0.
#' @param values_mapper mapper object to be used for `ibm_n` and
#' `ibm_values` for `inla_f=TRUE` (experimental, currently unused)
#' @title Mapper for linear Taylor approximations
#' @description
#' Provides a pre-computed affine mapping,
#' internally used to represent and evaluate linearisation information.
#' The `state0` information indicates for which state the `offset` was
#' evaluated;
#' The affine mapper output is defined as
#' `effect(state) = offset + jacobian %*% (state - state0)`
#' @export
#' @rdname bm_taylor
#' @seealso [bru_mapper], [bru_mapper_generics]
#' @family mappers
#'
#' @examples
#' m <- bm_taylor(
#'   offset = rep(2, 3),
#'   jacobian = matrix(1:6, 3, 2),
#'   state0 = c(1, 2)
#' )
#' ibm_eval2(m, state = 2:3)
bm_taylor <- function(offset = NULL, jacobian = NULL, state0 = NULL,
                      values_mapper = NULL) {
  stopifnot(!is.null(offset) || !is.null(jacobian))
  if (is.null(state0)) {
    n_state <- 0
  } else if (is.list(state0)) {
    if (is.null(names(state0))) {
      stop("`bm_taylor(..., state0)` lists must have named elements.")
    }
    n_state <- vapply(
      state0,
      function(x) {
        ifelse(is.null(x),
          0L,
          length(x)
        )
      },
      0L
    )
  } else {
    n_state <- length(state0)
  }

  if (is.null(jacobian)) {
    n_jacobian <- 0
  } else if (is.list(jacobian)) {
    if (is.null(names(jacobian))) {
      stop("`bm_taylor(..., jacobian)` lists must have named elements.")
    }
    n_jacobian <- vapply(
      jacobian,
      function(x) {
        ifelse(is.null(x),
          0L,
          ncol(x)
        )
      },
      0L
    )
  } else {
    n_jacobian <- ncol(jacobian)
  }

  if (!is.null(state0)) {
    if ((sum(n_jacobian) == 0) && (sum(n_state) > 0)) {
      stop(paste0(
        "Jacobian information indicates an empty state, but the state ",
        "information has length > 0"
      ))
    }
    stopifnot(sum(n_jacobian) == sum(n_state))
  }

  n_multi <- n_jacobian
  if (is.null(jacobian)) {
    state0 <- NULL
    n_state <- 0
  } else if (!is.list(jacobian) && is.list(state0)) {
    state0 <- unlist(state0)
    n_state <- length(state0)
  } else if (is.list(jacobian) && !is.null(state0) && !is.list(state0)) {
    # Split vector
    state0 <- lapply(
      seq_along(n_multi),
      function(k) {
        idx <- sum(n_multi[seq_len(k - 1)])
        idx <- idx + seq_len(n_multi[k]) - 1
        state0[idx]
      }
    )
    names(state0) <- names(jacobian)
    n_state <- n_multi
  }
  if (!is.null(state0)) {
    stopifnot(all(n_jacobian == n_state))
  }

  if (is.null(offset)) {
    if (is.list(jacobian)) {
      offset <- numeric(nrow(jacobian[[1]]))
    } else {
      offset <- numeric(nrow(jacobian))
    }
  }
  bru_mapper_define(
    list(
      offset = as.vector(offset),
      jacobian = jacobian,
      state0 = state0,
      n_multi = n_multi,
      n = sum(n_multi),
      n_output = length(offset),
      values_mapper = NULL
    ),
    # TODO: maybe allow values_mapper
    new_class = "bm_taylor"
  )
}

#' @export
#' @rdname bm_taylor
#' @param \dots Arguments passed on to [bm_taylor()]
bru_mapper_taylor <- function(...) {
  bm_taylor(...)
}

#' @export
#' @rdname ibm_n
#' @inheritParams bm_multi
#'
ibm_n.bm_taylor <- function(mapper,
                            inla_f = FALSE,
                            multi = FALSE,
                            ...) {
  if (inla_f && !is.null(mapper[["values_mapper"]])) {
    stop("ibm_n.bm_taylor should not be used with inla_f = TRUE")
  } else if (multi) {
    mapper[["n_multi"]]
  } else {
    mapper[["n"]]
  }
}

#' @export
#' @rdname ibm_n_output
#'
ibm_n_output.bm_taylor <- function(mapper, input, ...) {
  mapper[["n_output"]]
}


#' @export
#' @rdname ibm_values
#'
ibm_values.bm_taylor <- function(mapper,
                                 inla_f = FALSE,
                                 multi = FALSE,
                                 ...) {
  if (inla_f && !is.null(mapper[["values_mapper"]])) {
    stop("ibm_values.bm_taylor should not be used with inla_f = TRUE")
  } else if (multi) {
    lapply(mapper[["n_multi"]], function(k) seq_len(k))
  } else {
    seq_len(mapper[["n"]])
  }
}

#' @export
#' @rdname ibm_jacobian
#'
ibm_jacobian.bm_taylor <- function(mapper, ..., multi = FALSE) {
  if (is.null(mapper[["jacobian"]])) {
    return(Matrix::Matrix(0, mapper[["n_output"]], 0))
  }
  if (multi) {
    mapper[["jacobian"]]
  } else if (is.list(mapper[["jacobian"]])) {
    do.call(cbind, mapper[["jacobian"]])
  } else {
    mapper[["jacobian"]]
  }
}


#' @describeIn ibm_eval
#' Evaluates linearised
#' mapper information at the given `state`. The `input` argument is ignored,
#' so that the usual argument order
#' `ibm_eval(mapper, input, state)` syntax can be used, but also
#' `ibm_eval(mapper, state = state)`.  For a mapper with a named jacobian list,
#' the `state` argument must also be a named list.  If `state` is `NULL`,
#' all-zero is assumed.
#' @export
#'
ibm_eval.bm_taylor <- function(mapper,
                               input = NULL,
                               state = NULL,
                               ...) {
  if (is.null(mapper[["jacobian"]]) ||
    (mapper[["n"]] == 0) ||
    (is.null(state) && is.null(mapper[["state0"]]))) {
    val <- mapper[["offset"]]
  } else if (is.list(mapper[["jacobian"]])) {
    stopifnot(is.null(state) || is.list(state))
    val <- mapper[["offset"]]
    if (is.null(mapper[["state0"]])) {
      missing_names <- setdiff(names(mapper[["jacobian"]]), names(state))
      if (length(missing_names) > 0) {
        stop(
          "Names in jacobian list missing from state list: ",
          paste0(missing_names, collapse = ", ")
        )
      }
      for (nm in names(mapper[["jacobian"]])) {
        val <- val + mapper[["jacobian"]][[nm]] %*% state[[nm]]
      }
    } else if (is.null(state)) {
      for (nm in names(mapper[["jacobian"]])) {
        val <- val - mapper[["jacobian"]][[nm]] %*% mapper[["state0"]][[nm]]
      }
    } else {
      for (nm in names(mapper[["jacobian"]])) {
        val <- val + mapper[["jacobian"]][[nm]] %*%
          (state[[nm]] - mapper[["state0"]][[nm]])
      }
    }
  } else {
    stopifnot(is.null(state) || !is.list(state))
    if (is.null(mapper[["state0"]])) {
      val <- mapper[["offset"]] + mapper[["jacobian"]] %*% state
    } else if (is.null(state)) {
      val <- mapper[["offset"]] - mapper[["jacobian"]] %*% mapper[["state0"]]
    } else {
      val <- mapper[["offset"]] +
        mapper[["jacobian"]] %*% (state - mapper[["state0"]])
    }
  }
  as.vector(val)
}


## _linear ####

#' @title Mapper for a linear effect
#' @export
#' @description Create a mapper for linear effects
#' @rdname bm_linear
#' @seealso [bru_mapper], [bru_mapper_generics]
#' @family mappers
#' @examples
#' m <- bm_linear()
#' ibm_eval(m, input = 1:4, state = 2)
#'
bm_linear <- function() {
  bru_mapper_define(list(), new_class = "bm_linear")
}

#' @export
#' @rdname bm_linear
bru_mapper_linear <- function() {
  bm_linear()
}

#' @export
#' @describeIn ibm_n Returns `1L`
#'
ibm_n.bm_linear <- function(mapper, ...) {
  1L
}
#' @export
#' @describeIn ibm_values Returns `1.0`
#'
ibm_values.bm_linear <- function(mapper, ...) {
  1.0
}

#' @export
#' @rdname ibm_jacobian
#'
ibm_jacobian.bm_linear <- function(mapper, input, ...) {
  if (is.null(input)) {
    return(Matrix::Matrix(0, 0, ibm_n(mapper)))
  }
  if (!is.null(input) &&
    !(is.numeric(input) ||
      is.logical(input) ||
      is(input, "Matrix"))) {
    stop(paste0(
      "The input to a bm_linear evaluation must be numeric or logical."
    ))
  }
  ok <- !is.na(input)
  A <- Matrix::sparseMatrix(
    i = which(ok),
    j = rep(1, sum(ok)),
    x = input[ok],
    dims = c(NROW(input), ibm_n(mapper))
  )
  A
}

## _matrix ####

#' @title Mapper for matrix multiplication
#' @param labels Column labels for matrix mappings; Can be factor, character, or
#'   a single integer specifying the number of columns for integer column
#'   indexing.
#' @export
#' @description Create a matrix mapper, for a given number of columns
#' @rdname bm_matrix
#' @seealso [bru_mapper], [bru_mapper_generics]
#' @family mappers
#'
#' @examples
#' m <- bm_matrix(labels = c("a", "b"))
#' ibm_values(m)
#' ibm_eval2(m, input = matrix(1:6, 3, 2), state = 2:3)
#'
#' m <- bm_matrix(labels = 2L)
#' ibm_values(m)
#' ibm_eval2(m, input = matrix(1:6, 3, 2), state = 2:3)
#'
bm_matrix <- function(labels) {
  if (is.factor(labels)) {
    mapper <- list(
      labels = levels(labels)
    )
  } else if (is.integer(labels) && (length(labels) == 1L)) {
    mapper <- list(
      labels = seq_len(labels)
    )
  } else {
    mapper <- list(
      labels = as.character(labels)
    )
  }
  bru_mapper_define(mapper, new_class = "bm_matrix")
}

#' @export
#' @rdname bm_matrix
#' @param \dots Arguments passed on to [bm_matrix()]
bru_mapper_matrix <- function(...) {
  bm_matrix(...)
}

#' @export
#' @describeIn ibm_n Returns the number of columns in the matrix mapper.
#'
ibm_n.bm_matrix <- function(mapper, ...) {
  length(mapper$labels)
}
#' @export
#' @describeIn ibm_values For integer labels, the vector of labels.
#' For character labels, the labels as a factor variable.
#'
ibm_values.bm_matrix <- function(mapper, ...) {
  if (is.integer(mapper$labels)) {
    mapper$labels
  } else {
    factor(x = mapper$labels, levels = mapper$labels)
  }
}

#' @export
#' @describeIn ibm_jacobian Accepts `input` as a matrix, `Matrix`, `Spatial`,
#' or `sfc_POINT` object.
#'
ibm_jacobian.bm_matrix <- function(mapper, input, state = NULL,
                                   inla_f = FALSE, ...) {
  if (is.null(input)) {
    return(Matrix::Matrix(0, 0, ibm_n(mapper)))
  } else if (is.matrix(input)) {
    A <- as(input, "Matrix")
  } else if (inherits(input, "Matrix")) {
    A <- input
  } else if (inherits(input, "Spatial")) {
    A <- sp::coordinates(input)
    A <- as(A, "Matrix")
  } else if (inherits(input, "sfc_POINT")) {
    A <- sf::st_coordinates(input)
    nm <- intersect(colnames(A), c("X", "Y", "Z", "M"))
    A <- as(A[, nm, drop = FALSE], "Matrix")
  } else {
    A <- as(input, "Matrix")
  }
  if (ncol(A) != ibm_n(mapper)) {
    stop(paste0(
      "Input to matrix mapper has ", ncol(A),
      " columns but should have ", ibm_n(mapper),
      " columns."
    ))
  }
  colnames(A) <- mapper$labels
  A
}


## _factor ####

#' @title Mapper for factor variables
#' @param values Input values calculated by [bru_input.bru_input()]
#' @param factor_mapping character; selects the type of factor mapping.
#' * `'contrast'` for leaving out the first factor level.
#' * `'full'` for keeping all levels.
#' @param indexed logical; if `TRUE`, the `ibm_values()` method
#' will return an integer vector instead of the factor levels.
#' This is needed e.g. for `group` and `replicate` mappers, since
#' `INLA::f()` doesn't accept factor values. Default: `FALSE`, which
#' works for the main input mappers. The default mapper constructions
#' will set it the required setting.
#' @export
#' @description Create a factor mapper
#' @rdname bm_factor
#' @seealso [bru_mapper], [bru_mapper_generics]
#' @family mappers
#' @examples
#' m <- bm_factor(factor(c("a", "b")), "full")
#' ibm_eval2(m, input = c("b", "a", "a", "b"), state = c(1, 3))
#'
#' m <- bm_factor(factor(c("a", "b")), "contrast")
#' ibm_eval2(m, input = factor(c("b", "a", "a", "b")), state = 2)
#'
bm_factor <- function(values, factor_mapping, indexed = FALSE) {
  factor_mapping <- match.arg(factor_mapping, c("full", "contrast"))
  if (is.factor(values)) {
    mapper <- list(
      levels = levels(values),
      factor_mapping = factor_mapping,
      indexed = indexed
    )
  } else if (is.character(values)) {
    mapper <- list(
      levels = unique(values),
      factor_mapping = factor_mapping,
      indexed = indexed
    )
  } else {
    mapper <- list(
      levels = as.character(sort(unique(values))),
      factor_mapping = factor_mapping,
      indexed = indexed
    )
  }
  mapper$n <- length(mapper$levels) - identical(factor_mapping, "contrast")
  if (indexed) {
    bru_mapper_define(mapper,
      new_class = c(
        "bm_factor_index",
        "bm_factor"
      )
    )
  } else {
    bru_mapper_define(mapper, new_class = "bm_factor")
  }
}

#' @export
#' @rdname bm_factor
#' @param \dots Arguments passed on to [bm_factor()]
bru_mapper_factor <- function(...) {
  bm_factor(...)
}

#' @export
#' @describeIn ibm_n Returns the number of levels when `factor_mapping` is
#'   "full", and the number of levels minus one if `factor_mapping` is
#'   "contrast".
#'
ibm_n.bm_factor <- function(mapper, ...) {
  length(mapper[["levels"]]) - identical(mapper[["factor_mapping"]], "contrast")
}
#' @export
#' @describeIn ibm_values Returns the factor levels (minus the first level for
#' `factor_mapping` "contrast"), or an integer vector (if `indexed = TRUE` in
#' [bm_factor()]).
#'
ibm_values.bm_factor <- function(mapper, ...) {
  if (is.null(mapper[["indexed"]]) || !mapper[["indexed"]]) {
    if (identical(mapper[["factor_mapping"]], "contrast")) {
      mapper$levels[-1L]
    } else {
      mapper$levels
    }
  } else {
    seq_len(ibm_n(mapper))
  }
}

#' @export
#' @rdname ibm_jacobian
#'
ibm_jacobian.bm_factor <- function(mapper, input, ...) {
  if (is.null(input)) {
    return(Matrix::Matrix(0, 0, ibm_n(mapper)))
  }
  if (is.factor(input)) {
    if (!identical(levels(input), mapper$levels)) {
      input <- factor(as.character(input), levels = mapper$levels)
    }
  } else {
    input <- factor(input, levels = mapper$levels)
  }
  if (mapper$factor_mapping == "full") {
    input <- as.numeric(input)
  } else if (mapper$factor_mapping == "contrast") {
    input <- as.numeric(input) - 1L
  } else {
    stop("Unknown factor mapping '", mapper$factor_mapping, "'.")
  }
  ok <- !is.na(input)
  ok[which(ok)] <- (input[ok] > 0L)
  A <- Matrix::sparseMatrix(
    i = which(ok),
    j = input[ok],
    x = rep(1.0, sum(ok)),
    dims = c(NROW(input), ibm_n(mapper))
  )
  A
}


## _const ####

#' @title Constant mapper
#' @export
#' @description Create a constant mapper
#' @rdname bm_const
#' @seealso [bru_mapper], [bru_mapper_generics]
#' @family mappers
#' @examples
#' m <- bm_const()
#' ibm_eval2(m, input = 1:4)
#'
bm_const <- function() {
  bru_mapper_define(list(), new_class = "bm_const")
}

#' @export
#' @rdname bm_const
bru_mapper_const <- function() {
  bm_const()
}

#' @export
#' @rdname ibm_n
#'
ibm_n.bm_const <- function(mapper, ...) {
  0L
}
#' @export
#' @rdname ibm_values
#'
ibm_values.bm_const <- function(mapper, ...) {
  NULL
}

#' @export
#' @rdname ibm_jacobian
#'
ibm_jacobian.bm_const <- function(mapper, input, ...) {
  Matrix::Matrix(0, ibm_n_output(mapper, input, ...), 0L)
}

#' @export
#' @describeIn ibm_eval Returns the input values, with NA replaced by 0.
#'
ibm_eval.bm_const <- function(mapper, input, state = NULL, ...) {
  if (is.null(input)) {
    return(numeric(0))
  }
  input <- as.vector(input)
  ok <- !is.na(input)
  input[!ok] <- 0
  input
}


## _shift ####

#' @title Mapper for element-wise shifting
#' @export
#' @description Create a standalone
#' shift mapper that can be used as part of a `bm_pipe`.
#' If `mapper` is non-null, the `bm_shift()` constructor
#' returns
#' `bm_pipe(list(mapper = mapper, shift = bm_shift()))`
#' @param mapper If non-NULL, a `bru_mapper` to be shifted.
#' @rdname bm_shift
#' @seealso [bru_mapper], [bru_mapper_generics]
#' @family mappers
#' @examples
#' m <- bm_shift()
#' ibm_eval2(m, c(1, 2, 1, 2), 1:4)
#'
bm_shift <- function(mapper = NULL) {
  m <- bru_mapper_define(
    list(),
    new_class = "bm_shift"
  )
  if (!is.null(mapper)) {
    m <- bm_pipe(list(mapper = mapper, shift = m))
  }
  m
}

#' @export
#' @rdname bm_shift
#' @param \dots Arguments passed on to [bm_shift()]
bru_mapper_shift <- function(...) {
  bm_shift(...)
}

#' @param n_state integer giving the length of the state vector for mappers
#' that have state dependent output size.
#' @export
#' @rdname ibm_n
#'
ibm_n.bm_shift <- function(mapper, ..., state = NULL, n_state = NULL) {
  # Output size depends on the state size
  if (!is.null(state)) {
    length(state)
  } else if (is.null(n_state)) {
    NA_integer_
  } else {
    n_state
  }
}
#' @export
#' @rdname ibm_n_output
#'
ibm_n_output.bm_shift <- function(mapper, input, state = NULL, ...,
                                  n_state = NULL) {
  if (!is.null(state)) {
    length(state)
  } else if (is.null(n_state)) {
    # To allow scalar input shifts, do not assume NROW(input) size
    NA_integer_
  } else {
    n_state
  }
}
#' @export
#' @rdname ibm_values
#'
ibm_values.bm_shift <- function(mapper, ...,
                                state = NULL, n_state = NULL) {
  n_state <- ibm_n(mapper, state = state, n_state = n_state)
  if (is.na(n_state)) {
    # Don't know how big the mapper will be
    NULL
  } else {
    seq_len(n_state)
  }
}

#' @export
#' @describeIn ibm_jacobian `input` NULL values are interpreted as no shift.
#'
ibm_jacobian.bm_shift <- function(mapper, input, state = NULL, ...) {
  stopifnot(!is.null(state))
  Matrix::Diagonal(n = length(state), 1.0)
}


#' @export
#' @rdname ibm_eval
#' @inheritParams ibm_jacobian
#'
ibm_eval.bm_shift <- function(mapper, input, state = NULL, ...) {
  stopifnot(!is.null(state))
  if (is.null(input)) {
    return(state)
  }
  shift <- as.vector(input)
  ok <- !is.na(shift)
  shift[!ok] <- 0
  state + shift
}


## _scale ####

#' @title Mapper for element-wise scaling
#' @export
#' @description Create a standalone
#' scaling mapper that can be used as part of a `bm_pipe`.
#' If `mapper` is non-null, the `bm_scale()` constructor
#' returns
#' `bm_pipe(list(mapper = mapper, scale = bm_scale()))`
#' @rdname bm_scale
#' @param mapper An optional `bru_mapper` to be scaled.
#' @seealso [bru_mapper], [bru_mapper_generics]
#' @family mappers
#' @examples
#' m <- bm_scale()
#' ibm_eval2(m, c(1, 2, 1, 2), 1:4)
#'
bm_scale <- function(mapper = NULL) {
  m <- bru_mapper_define(
    list(),
    new_class = "bm_scale"
  )
  if (!is.null(mapper)) {
    m <- bm_pipe(list(mapper = mapper, scale = m))
  }
  m
}

#' @export
#' @rdname bm_scale
#' @param \dots Arguments passed on to [bm_scale()]
bru_mapper_scale <- function(...) {
  bm_scale(...)
}

#' @param n_state integer giving the length of the state vector for mappers
#' that have state dependent output size.
#' @export
#' @rdname ibm_n
#'
ibm_n.bm_scale <- function(mapper, ..., state = NULL, n_state = NULL) {
  # Output size depends on the state size
  if (!is.null(state)) {
    length(state)
  } else if (is.null(n_state)) {
    NA_integer_
  } else {
    n_state
  }
}
#' @export
#' @rdname ibm_n_output
#'
ibm_n_output.bm_scale <- function(mapper, input, state = NULL, ...,
                                  n_state = NULL) {
  if (!is.null(state)) {
    length(state)
  } else if (is.null(n_state)) {
    # To allow scalar input weights, do not assume NROW(input) size
    NA_integer_
  } else {
    n_state
  }
}
#' @export
#' @rdname ibm_values
#'
ibm_values.bm_scale <- function(mapper, ...,
                                state = NULL, n_state = NULL) {
  n_state <- ibm_n(mapper, state = state, n_state = n_state)
  if (is.na(n_state)) {
    # Don't know how big the mapper will be
    NULL
  } else {
    seq_len(n_state)
  }
}

#' @export
#' @describeIn ibm_jacobian `input` NULL values
#' are interpreted as no scaling.
#'
ibm_jacobian.bm_scale <- function(mapper, input, state = NULL, ...) {
  stopifnot(!is.null(state))
  if (is.null(input)) {
    # No scaling
    return(Matrix::Diagonal(n = length(state), 1.0))
  }
  if (!(is.numeric(input) ||
    is.logical(input) ||
    is(input, "Matrix"))) {
    stop(
      "The input to a bm_scale evaluation must be numeric or logical."
    )
  }
  scale <- as.vector(input)
  ok <- !is.na(scale)
  scale[!ok] <- 0
  Matrix::Diagonal(n = length(state), scale)
}


#' @export
#' @rdname ibm_eval
#'
ibm_eval.bm_scale <- function(mapper, input, state = NULL, ...) {
  stopifnot(!is.null(state))
  if (is.null(input)) {
    return(state)
  }
  if (!is.null(input) &&
    !(is.numeric(input) ||
      is.logical(input) ||
      is(input, "Matrix"))) {
    stop(
      "The input to a bm_scale evaluation must be numeric or logical."
    )
  }
  scale <- as.vector(input)
  ok <- !is.na(scale)
  scale[!ok] <- 0
  scale * state
}


## _aggregate ####

#' @title Mapper for aggregation
#' @export
#' @param rescale
#' logical; For `bm_aggregate` and `bm_logsumexp`,
#' specifies if the blockwise sums should be normalised by the blockwise weight
#' sums or not:
#' * `FALSE`: (default) Straight weighted sum, no rescaling.
#' * `TRUE`: Divide by the sum of the weight values within each block.
#'   This is useful for integration averages, when the given weights are plain
#'   integration weights. If the weights are `NULL` or all ones, this is
#'   the same as dividing by the number of entries in each block.
#' @param n_block Predetermined number of output blocks. If `NULL`, overrides
#' the maximum block index in the inputs. The priority order is `input$n_block`,
#' the mapper definition `n_block`, then `max(input$block)`.
#' @param type character; if non-NULL, overrides the `rescale` argument, and
#' constructs an aggregation mapper of the given type instead. Supported
#' values are "sum", "average" (for regular [bm_aggregate()]),
#' "logsumexp", "logaverageexp" (for [bm_logsumexp()]), and
#' "logitaverage" (for [bm_logitaverage()]).
#' @description
#' Constructs a mapper
#' that aggregates elements of the input state, so it can be used e.g.
#' for weighted summation or integration over blocks of values.
#' @rdname bm_aggregate
#' @seealso [bru_mapper], [bru_mapper_generics]
#' @family mappers
#' @examples
#' m <- bm_aggregate()
#' ibm_eval2(m, list(block = c(1, 2, 1, 2), weights = 1:4), 11:14)
#' ibm_eval2(m, list(block = c(1, 2, 1, 2), weights = 1:4, n_block = 3), 11:14)
#'
bm_aggregate <- function(rescale = FALSE,
                         n_block = NULL,
                         type = NULL) {
  if (!is.null(type)) {
    type <- match.arg(
      type,
      c("sum", "average", "logsumexp", "logaverageexp", "logitaverage")
    )
    if (type == "sum") {
      rescale <- FALSE
    } else if (type == "average") {
      rescale <- TRUE
    } else if (type == "logsumexp") {
      return(bm_logsumexp(
        rescale = FALSE,
        n_block = n_block
      ))
    } else if (type == "logaverageexp") {
      return(bm_logsumexp(
        rescale = TRUE,
        n_block = n_block
      ))
    } else if (type == "logitaverage") {
      return(bm_logitaverage(
        n_block = n_block
      ))
    }
  }
  bru_mapper_define(
    list(
      rescale = rescale,
      n_block = n_block,
      is_rowwise = FALSE
    ),
    new_class = "bm_aggregate"
  )
}

#' @export
#' @param \dots Arguments passed on to [bm_aggregate()]
#' @rdname bm_aggregate
bru_mapper_aggregate <- function(...) {
  bm_aggregate(...)
}

#' @param n_state integer giving the length of the state vector for mappers
#' that have state dependent output size.
#' @export
#' @rdname ibm_n
#'
ibm_n.bm_aggregate <- function(mapper, ...,
                               input = NULL,
                               state = NULL,
                               n_state = NULL) {
  # Output size depends on the state size
  if (!is.null(state)) {
    length(state)
  } else if (!is.null(n_state)) {
    n_state
  } else if (!is.null(input) &&
    (!is.null(input[["block"]]) ||
      !is.null(input[["weights"]]) ||
      !is.null(input[["log_weights"]]))) {
    max(
      length(input[["block"]]),
      length(input[["weights"]]),
      length(input[["log_weights"]])
    )
    length(input[["block"]])
  } else {
    NA_integer_
  }
}

bm_aggregate_n_block <- function(mapper, input = NULL) {
  if (is.null(input[["n_block"]])) {
    # Note: Allowed to be NULL
    n_block <- mapper[["n_block"]]
  } else {
    n_block <- input[["n_block"]]
  }
  if (is.null(n_block) && !is.null(input[["block"]])) {
    if (is.character(input[["block"]])) {
      lifecycle::deprecate_stop(
        "2.12.0.9001",
        "bm_aggregate(input = 'can not be a `character`')",
        I("`integer` values for the `input` argument.")
      )
    } else {
      n_block <- max(input[["block"]])
    }
  }
  n_block
}

#' @export
#' @rdname ibm_n_output
#'
ibm_n_output.bm_aggregate <- function(mapper,
                                      input = NULL, ...) {
  n_block <- bm_aggregate_n_block(mapper = mapper, input = input)
  if (is.null(n_block)) {
    return(NA_integer_)
  }
  n_block
}
#' @export
#' @rdname ibm_values
#'
ibm_values.bm_aggregate <- function(mapper, ...,
                                    state = NULL,
                                    n_state = NULL) {
  n_state <- ibm_n(mapper, state = state, n_state = n_state)
  if (is.na(n_state)) {
    # Don't know how big the mapper will be
    NULL
  } else {
    seq_len(n_state)
  }
}


#' @export
#' @describeIn ibm_jacobian
#' `input` should be a list with elements `block`
#' and `weights`. `block`
#' should be a vector of the same length as the `state`, or `NULL`, with `NULL`
#' equivalent to all-1.
#' If `weights` is `NULL`, it's interpreted as all-1.
#'
ibm_jacobian.bm_aggregate <- function(mapper,
                                      input,
                                      state = NULL,
                                      ...) {
  n_block <- bm_aggregate_n_block(mapper = mapper, input = input)
  fm_block(
    block = input[["block"]],
    weights = input[["weights"]],
    log_weights = input[["log_weights"]],
    n_block = n_block,
    rescale = mapper[["rescale"]]
  )
}


#' @export
#' @rdname ibm_eval
#'
ibm_eval.bm_aggregate <- function(mapper, input, state = NULL, ...) {
  n_block <- bm_aggregate_n_block(mapper = mapper, input = input)
  val <-
    fm_block_eval(
      block = input[["block"]],
      log_weights = input[["log_weights"]],
      weights = input[["weights"]],
      rescale = mapper[["rescale"]],
      n_block = n_block,
      values = state
    )
  val
}


## _logsumexp ####

#' @title Mapper for log-sum-exp aggregation
#' @export
#' @description
#' Constructs a mapper
#' that aggregates elements of `exp(state)`, with optional non-negative
#' weighting, and then takes the `log()`, so it can be used e.g.
#' for  \eqn{v_k=\log[\sum_{i\in I_k} w_i \exp(u_i)]
#' }{log(blocksum(w*exp(state)))}
#' and \eqn{v_k=\log[\sum_{i\in I_k} w_i \exp(u_i) / \sum_{i\in I_k} w_i]
#' }{log(blocksum(w*exp(state)/blocksum(w)))}
#' calculations.  Relies on the input handling methods for
#' `bm_aggregate`, but also allows the weights to be supplied on a
#' logarithmic scale as `log_weights`. To avoid numerical overflow, it uses the
#' common method of internally shifting the state blockwise;
#' \eqn{v_k=s_k+\log[\sum_{i\in I_k} \exp(u_i + \log(w_i)- s_k)]
#' }{log(blocksum(w*exp(state)))},
#' where \eqn{s_k=\max_{i\in I_k} u_i + \log(w_i)}{s=blockmax(u+log(w))} is the
#' shift for block \eqn{k}{k}.
#' @rdname bm_logsumexp
#' @inheritParams bm_aggregate
#' @seealso [bru_mapper], [bru_mapper_generics]
#' @family mappers
#' @examples
#' m <- bm_logsumexp()
#' ibm_eval2(m, list(block = c(1, 2, 1, 2), weights = 1:4), 11:14)
#' ibm_eval2(m, list(block = c(1, 2, 1, 2), weights = 1:4, n_block = 3), 11:14)
#'
bm_logsumexp <- function(rescale = FALSE,
                         n_block = NULL) {
  # Arguments documented for bm_aggregate
  # Inherit class bm_aggregate to reuse common methods
  bru_mapper_define(
    list(
      rescale = rescale,
      n_block = n_block,
      is_rowwise = FALSE,
      is_linear = FALSE
    ),
    new_class = c("bm_logsumexp", "bm_aggregate")
  )
}

#' @export
#' @rdname bm_logsumexp
#' @param \dots Arguments passed on to [bm_logsumexp()]
bru_mapper_logsumexp <- function(...) {
  bm_logsumexp(...)
}


#' @export
#' @describeIn ibm_jacobian `input` should be a list with elements
#'   `block` and `weights`. `block` should be a vector of the same length as the
#'   `state`, or `NULL`, with `NULL` equivalent to all-1.
#' If `weights` is `NULL`, it's interpreted as all-1.
#'
ibm_jacobian.bm_logsumexp <- function(mapper,
                                      input,
                                      state = NULL,
                                      ...) {
  n_block <- bm_aggregate_n_block(mapper = mapper, input = input)
  input <- fm_block_prep(
    block = input[["block"]],
    log_weights = input[["log_weights"]],
    weights = input[["weights"]],
    force_log = TRUE,
    values = state,
    n_block = n_block
  )

  n_state <- length(input$block)
  n_out <- input$n_block

  log_weights <- fm_block_log_weights(
    block = input[["block"]],
    log_weights = input[["log_weights"]],
    weights = input[["weights"]],
    n_block = n_out,
    rescale = mapper[["rescale"]]
  )

  # Compute shift for stable log-sum-exp
  w_state <- state + log_weights
  shift <- fm_block_log_shift(
    log_weights = w_state,
    block = input[["block"]],
    n_block = n_out
  )

  sum_values <-
    as.vector(
      Matrix::sparseMatrix(
        i = input[["block"]],
        j = rep(1L, n_state),
        x = exp(w_state - shift[input[["block"]]]),
        dims = c(n_out, 1)
      )
    )
  scale <- sum_values[input[["block"]]]

  Matrix::sparseMatrix(
    i = input[["block"]],
    j = seq_len(n_state),
    x = exp(w_state - shift[input[["block"]]]) / scale,
    dims = c(n_out, n_state)
  )
}

#' @export
#' @param log logical; control `log` output. Default `TRUE`, see the
#'   `ibm_eval()` details for `logsumexp` mappers.
#' @describeIn ibm_eval When `log` is `TRUE` (default), `ibm_eval()`
#'   for `logsumexp` returns the log-sum-weight-exp value. If `FALSE`, the
#'   `sum-weight-exp` value is returned.
#'
ibm_eval.bm_logsumexp <- function(mapper, input, state = NULL,
                                  log = TRUE, ...) {
  n_block <- bm_aggregate_n_block(mapper = mapper, input = input)
  val <-
    fm_block_logsumexp_eval(
      block = input[["block"]],
      log_weights = input[["log_weights"]],
      weights = input[["weights"]],
      rescale = mapper[["rescale"]],
      n_block = n_block,
      values = state,
      log = log
    )
  val
}


## _logitaverage ####

#' @title Mapper for logit-sum-inverse-logit aggregation
#' @export
#' @description `r lifecycle::badge("experimental")`
#' Constructs a mapper that averages elements of `plogis(state)`, with optional
#' non-negative weighting, and then takes the `qlogis()`.  Relies on the input
#' handling methods for `bm_aggregate`. To avoid numerical issues, it uses
#' `plogis(x, log.p = TRUE)`, `plogis(-x, log.p = TRUE)`, and the equivalent of
#' two applications of [bm_logsumexp()] to evaluate
#' \eqn{\log(p_k)-\log(1-p_k)}{log(p[k])-log(1-p[k])}, where
#'
#' \eqn{p_k=\sum_{i\in I_k} w_i / (1+e^{-\eta_i}) / \sum_{i\in I_k} w_i
#' }{p[k]=sum(w[i] * plogis(eta[i])) / sum(w[i])}
#' @rdname bm_logitaverage
#' @inheritParams bm_aggregate
#' @seealso [bru_mapper], [bru_mapper_generics]
#' @family mappers
#' @examples
#' m <- bm_logitaverage()
#' ibm_eval2(m, list(block = c(1, 2, 1, 2), weights = 1:4), 11:14)
#' ibm_eval2(m, list(block = c(1, 2, 1, 2), weights = 1:4, n_block = 3), 11:14)
#'
bm_logitaverage <- function(n_block = NULL) {
  # Arguments documented for bm_aggregate
  # Inherit class bm_aggregate to reuse common methods
  bru_mapper_define(
    list(
      rescale = TRUE,
      n_block = n_block,
      is_rowwise = FALSE,
      is_linear = FALSE
    ),
    new_class = c("bm_logitaverage", "bm_aggregate")
  )
}

#' @export
#' @describeIn ibm_jacobian `input` should be a list with elements
#'   `block` and `weights`. `block` should be a vector of the same length as the
#'   `state`, or `NULL`, with `NULL` equivalent to all-1.
#' If `weights` is `NULL`, it's interpreted as all-1.
#'
ibm_jacobian.bm_logitaverage <- function(mapper,
                                         input,
                                         state = NULL,
                                         ...) {
  n_block <- bm_aggregate_n_block(mapper = mapper, input = input)
  input <- fm_block_prep(
    block = input[["block"]],
    log_weights = input[["log_weights"]],
    weights = input[["weights"]],
    force_log = TRUE,
    values = state,
    n_block = n_block
  )

  n_state <- length(input$block)
  n_out <- input$n_block

  log_weights <- fm_block_log_weights(
    block = input[["block"]],
    log_weights = input[["log_weights"]],
    weights = input[["weights"]],
    n_block = n_out,
    rescale = TRUE
  )

  state1 <- plogis(state, log.p = TRUE)
  state2 <- plogis(-state, log.p = TRUE)

  # Compute shift for stable log-sum-exp
  w_state1 <- state1 + log_weights
  w_state2 <- state2 + log_weights
  shift1 <- fm_block_log_shift(
    log_weights = w_state1,
    block = input[["block"]],
    n_block = n_out
  )
  shift2 <- fm_block_log_shift(
    log_weights = w_state2,
    block = input[["block"]],
    n_block = n_out
  )

  sum_values1 <-
    as.vector(
      Matrix::sparseMatrix(
        i = input[["block"]],
        j = rep(1L, n_state),
        x = exp(w_state1 - shift1[input[["block"]]]),
        dims = c(n_out, 1)
      )
    )
  scale1 <- exp(state2[input[["block"]]]) / sum_values1[input[["block"]]]
  sum_values2 <-
    as.vector(
      Matrix::sparseMatrix(
        i = input[["block"]],
        j = rep(1L, n_state),
        x = exp(w_state2 - shift2[input[["block"]]]),
        dims = c(n_out, 1)
      )
    )
  scale2 <- -exp(state1[input[["block"]]]) / sum_values2[input[["block"]]]

  A <- Matrix::sparseMatrix(
    i = input[["block"]],
    j = seq_len(n_state),
    x = exp(w_state1 - shift1[input[["block"]]]) * scale1 -
      exp(w_state2 - shift2[input[["block"]]]) * scale2,
    dims = c(n_out, n_state)
  )
  A
}

#' @export
#' @param logit logical; control `logit` output. Default `TRUE`, see the
#'   `ibm_eval()` details for `logitaverage` mappers.
#' @describeIn ibm_eval When `logit` is `TRUE` (default), `ibm_eval()`
#'   for `logitaverage` returns the logit-sum-weight-inverse-logit value.
#'   If `FALSE`, the `sum-weights=invere-logit` value is returned.
#'
ibm_eval.bm_logitaverage <- function(mapper, input, state = NULL,
                                     logit = TRUE, ...) {
  n_block <- bm_aggregate_n_block(mapper = mapper, input = input)
  state1 <- plogis(state, log.p = TRUE)
  state2 <- plogis(-state, log.p = TRUE)
  val_1 <-
    fm_block_logsumexp_eval(
      block = input[["block"]],
      log_weights = input[["log_weights"]],
      weights = input[["weights"]],
      rescale = TRUE,
      n_block = n_block,
      values = state1,
      log = TRUE
    )
  val_2 <-
    fm_block_logsumexp_eval(
      block = input[["block"]],
      log_weights = input[["log_weights"]],
      weights = input[["weights"]],
      rescale = TRUE,
      n_block = n_block,
      values = state2,
      log = TRUE
    )
  val <- val_1 - val_2
  if (!logit) {
    val <- plogis(val)
  }
  val
}


## _marginal ####

require_args <- function(fun, req) {
  label <- deparse1(substitute(fun))
  if (is.null(fun)) {
    return(invisible())
  }
  fun_args <- formals(args(fun))
  subset <- req %in% names(fun_args)
  if (!all(subset)) {
    subset <- req[!subset]
    stop(paste0(
      "The '", label, "' function is missing the argument",
      if (length(subset) > 1) "s" else "",
      " ",
      paste0("'", subset, "'", collapse = ", "),
      "."
    ))
  }
  invisible()
}

#' @title Mapper for marginal distribution transformation
#' @export
#' @description
#' Constructs a mapper that transforms the marginal distribution `state` from
#' \eqn{\textrm{N}(0,1)}{N(0, 1)} to the distribution of a given (continuous)
#' quantile function. The `...` arguments are used as parameter arguments to
#' `qfun`, `pfun`, `dfun`, and `dqfun`.
#' @param qfun A quantile function, supporting arguments `p`, `lower.tail`, and
#'   `log.p`, like [stats::qnorm()].
#' @param pfun A CDF, supporting arguments `q`, `lower.tail`, and `log.p`,
#' like [stats::pnorm()].  Only needed and used when
#' `xor(mapper[["inverse"]], reverse)` is `TRUE` in a method call.
#' Default `NULL`
#' @param dfun A pdf, supporting arguments `x` and `log`,
#' like [stats::dnorm()]. If `NULL` (default), uses finite
#' differences on `qfun` or `pfun` instead.
#' @param dqfun A function evaluating the reciprocal of the derivative of
#'   `qfun`, i.e. the density at a given CDF value. If `NULL` (default), uses
#'   `dfun(qfun(...),...)` or finite differences on `qfun` or `pfun` instead.
#'   Must support the same arguments as `qfun`.
#' @param inverse logical; If `FALSE` (default), [bm_marginal()]
#' defines a mapping from standard Normal to a specified distribution.
#' If `TRUE`, it defines a mapping from the specified distribution to a standard
#' Normal.
#' @examples
#' m <- bm_marginal(qexp, pexp, rate = 1 / 8)
#' (val <- ibm_eval(m, state = -5:5))
#' ibm_eval(m, state = val, reverse = TRUE)
#' @rdname bm_marginal
#' @seealso [bru_mapper], [bru_mapper_generics]
#' @family mappers
#' @examples
#' m <- bm_marginal(qexp, pexp, dexp, rate = 1 / 8)
#' ibm_eval2(m, state = -3:3)
#'
bm_marginal <- function(qfun,
                        pfun = NULL,
                        dfun = NULL,
                        dqfun = NULL,
                        ...,
                        inverse = FALSE) {
  arg_names <- ...names()
  require_args(qfun, c("p", "lower.tail", "log.p", arg_names))
  require_args(pfun, c("q", "lower.tail", "log.p", arg_names))
  require_args(dfun, c("x", "log", arg_names))
  require_args(dqfun, c("p", "lower.tail", "log.p", "log", arg_names))
  bru_mapper_define(
    list(
      qfun = qfun,
      pfun = pfun,
      dfun = dfun,
      dqfun = dqfun,
      inverse = inverse,
      param = list(...),
      is_linear = FALSE
    ),
    new_class = "bm_marginal"
  )
}

#' @export
#' @rdname bm_marginal
#' @param \dots Arguments passed on to [bm_marginal()]
bru_mapper_marginal <- function(...) {
  bm_marginal(...)
}

#' @export
#' @rdname ibm_n
#'
ibm_n.bm_marginal <- function(mapper,
                              ...,
                              state = NULL,
                              n_state = NULL) {
  # Output size depends on the state size
  if (!is.null(state)) {
    length(state)
  } else if (is.null(n_state)) {
    NA_integer_
  } else {
    n_state
  }
}
#' @export
#' @rdname ibm_n_output
#'
ibm_n_output.bm_marginal <- function(mapper,
                                     input,
                                     state = NULL,
                                     ...,
                                     n_state = NULL) {
  if (!is.null(state)) {
    length(state)
  } else if (is.null(n_state)) {
    # To allow scalar input weights, do not assume NROW(input) size
    NA_integer_
  } else {
    n_state
  }
}
#' @export
#' @rdname ibm_values
#'
ibm_values.bm_marginal <- function(mapper, ...,
                                   state = NULL, n_state = NULL) {
  n_state <- ibm_n(mapper, state = state, n_state = n_state)
  if (is.na(n_state)) {
    # Don't know how big the mapper will be
    NULL
  } else {
    seq_len(n_state)
  }
}

#' @export
#' @param reverse logical; control `bm_marginal` evaluation. Default
#'   `FALSE`. When `TRUE`, reverses the direction of the mapping, see details
#'   for `marginal` mappers.
#' @describeIn ibm_jacobian Non-NULL `input` values are interpreted
#' as a parameter list for `qfun`, overriding that of the mapper itself.
#'
ibm_jacobian.bm_marginal <- function(mapper, input, state = NULL,
                                     ...,
                                     reverse = FALSE) {
  stopifnot(!is.null(state))
  if (!missing(input) && !is.null(input)) {
    mapper$param <- input
  }
  if (!xor(isTRUE(mapper[["inverse"]]), reverse)) {
    if (!is.null(mapper[["dqfun"]])) {
      val <- pnorm(state)
      dq <- do.call(mapper[["dqfun"]], c(list(val), mapper$param))
      der <- dnorm(state) / dq
    } else if (!is.null(mapper[["dfun"]])) {
      mapper[["inverse"]] <- FALSE
      val <- ibm_eval(
        mapper,
        input = NULL,
        state = state,
        reverse = FALSE
      )
      log_dens <-
        do.call(mapper[["dfun"]], c(list(val, log = TRUE), mapper$param))
      der <- exp(dnorm(state, log = TRUE) - log_dens)
    } else {
      der <- NULL
    }
  } else if (!is.null(mapper[["dqfun"]]) ||
    !is.null(mapper[["dfun"]])) {
    mapper[["inverse"]] <- FALSE
    val <- ibm_eval(
      mapper,
      input = NULL,
      state = state,
      reverse = TRUE
    )
    if (!is.null(mapper[["dqfun"]])) {
      dq <- do.call(mapper[["dqfun"]], c(list(val), mapper$param))
      der <- dq / dnorm(val)
    } else { # if (!is.null(mapper[["dfun"]])) {
      log_dens <-
        do.call(mapper[["dfun"]], c(list(state, log = TRUE), mapper$param))
      der <- exp(log_dens - dnorm(val, log = TRUE))
    }
  } else {
    der <- NULL
  }
  if (is.null(der)) {
    eps <- 1e-4 * (1 + abs(state))
    der <-
      (
        ibm_eval(
          mapper,
          input = NULL,
          state = state + eps,
          reverse = reverse
        ) -
          ibm_eval(
            mapper,
            input = NULL,
            state = state - eps,
            reverse = reverse
          )
      ) /
        (2 * eps)
  }
  Matrix::Diagonal(n = length(state), der)
}


#' @export
#' @describeIn ibm_eval When `xor(mapper[["inverse"]], reverse)` is
#' `FALSE`, `ibm_eval()`
#' for `marginal` returns `qfun(pnorm(x), param)`, evaluated in a numerically
#' stable way. Otherwise, evaluates the inverse `qnorm(pfun(x, param))` instead.
#'
ibm_eval.bm_marginal <- function(mapper, input, state = NULL,
                                 ...,
                                 reverse = FALSE) {
  stopifnot(!is.null(state))
  if (!missing(input) && !is.null(input)) {
    mapper$param <- input
  }
  if (!xor(isTRUE(mapper[["inverse"]]), reverse)) {
    val <- do.call(
      bru_forward_transformation,
      c(
        list(
          qfun = mapper[["qfun"]],
          x = state
        ),
        mapper[["param"]]
      )
    )
  } else {
    val <- do.call(
      bru_inverse_transformation,
      c(
        list(
          pfun = mapper[["pfun"]],
          x = state
        ),
        mapper[["param"]]
      )
    )
  }
  as.vector(val)
}


## _pipe ####

#' @title Mapper for linking several mappers in sequence
#' @export
#' @description
#' Create a pipe mapper, where `mappers` is a list of mappers,
#' and the evaluated output of each mapper is handed as the state to the next
#' mapper.
#' The `input` format for the `ibm_eval` and `ibm_jacobian` methods is
#' a list of inputs, one for each mapper.
#' @rdname bm_pipe
#' @inheritParams bm_multi
#' @seealso [bru_mapper], [bru_mapper_generics]
#' @family mappers
#' @examples
#' m <- bm_pipe(list(
#'   scale = bm_scale(),
#'   shift = bm_shift()
#' ))
#' ibm_eval2(m, input = list(scale = 2, shift = 1:4), state = 1:4)
#'
bm_pipe <- function(mappers) {
  mappers <- as_bm_list(mappers)
  is_linear_multi <- vapply(mappers, function(x) ibm_is_linear(x), TRUE)
  is_rowwise_multi <- vapply(mappers, function(x) ibm_is_rowwise(x), TRUE)
  n_multi <- vapply(mappers, function(x) as.integer(ibm_n(x)), 0L)
  n <- ibm_n(mappers[[1]])
  is_linear <- all(is_linear_multi)
  is_rowwise <- all(is_rowwise_multi)
  the_names <- names(mappers)
  if (is.null(the_names)) {
    the_names <- as.character(seq_along(mappers))
    names(mappers) <- the_names
  } else if ("" %in% names(mappers)) {
    stop("Either all or none of the pipe sub-mappers should be named.")
  }
  bru_mapper_define(
    list(
      mappers = mappers,
      is_linear_multi = is_linear_multi,
      is_linear = is_linear,
      is_rowwise = is_rowwise,
      n_multi = n_multi,
      n = n,
      names = the_names
    ),
    new_class = "bm_pipe"
  )
}

#' @export
#' @rdname bm_pipe
#' @param \dots Arguments passed on to [bm_pipe()]
bru_mapper_pipe <- function(...) {
  bm_pipe(...)
}

#' @export
#' @rdname ibm_n
#'
ibm_n.bm_pipe <- function(mapper, ..., input = NULL, state = NULL) {
  if (is.null(mapper[["mappers"]][[1]]) && is.null(state)) {
    return(NA_integer_)
  }
  if (is.null(mapper[["names"]])) {
    mapper[["names"]] <- names(mapper[["mappers"]])
  }
  if (!is.null(input) && is.null(names(input))) {
    names(input) <- mapper[["names"]][seq_along(input)]
  }
  ibm_n(mapper[["mappers"]][[1]], ...,
    input = input[[names(mapper[["mappers"]])[1]]],
    state = state
  )
}

#' @export
#' @rdname ibm_n_output
#'
ibm_n_output.bm_pipe <- function(mapper,
                                 input,
                                 state = NULL,
                                 ...,
                                 n_state = NULL) {
  if (is.null(mapper[["names"]])) {
    mapper[["names"]] <- names(mapper[["mappers"]])
  }
  if (!is.null(input) && is.null(names(input))) {
    names(input) <- mapper[["names"]][seq_along(input)]
  }
  n <- ibm_n(
    mapper[["mappers"]][[1]],
    input = input[[names(mapper[["mappers"]])[1]]],
    state = state,
    ...,
    n_state = n_state
  )
  for (k in names(mapper[["mappers"]])) {
    n <- ibm_n_output(
      mapper[["mappers"]][[k]],
      input = input[[k]],
      ..., n_state = n
    )
  }
  n
}

#' @export
#' @rdname ibm_values
#'
ibm_values.bm_pipe <- function(mapper, ...) {
  ibm_values(mapper[["mappers"]][[1]], ...)
}

#' @export
#' @rdname ibm_jacobian
#'
ibm_jacobian.bm_pipe <- function(mapper, input, state = NULL, ...) {
  ibm_eval2(mapper, input = input, state = state, ...)$jacobian
}


#' @export
#' @rdname ibm_eval
#'
ibm_eval.bm_pipe <- function(mapper, input, state = NULL, ...) {
  if (is.null(mapper[["names"]])) {
    mapper[["names"]] <- names(mapper[["mappers"]])
  }
  if (!is.null(input) && is.null(names(input))) {
    names(input) <- mapper[["names"]][seq_along(input)]
  }
  state_k <- state
  for (k in names(mapper[["mappers"]])) {
    state_k <- ibm_eval(
      mapper[["mappers"]][[k]],
      input = input[[k]],
      state = state_k,
      ...
    )
  }
  state_k
}


#' @export
#' @rdname ibm_eval2
#'
ibm_eval2.bm_pipe <- function(mapper, input, state = NULL, ...) {
  if (is.null(mapper[["names"]])) {
    mapper[["names"]] <- names(mapper[["mappers"]])
  }
  if (!is.null(input) && is.null(names(input))) {
    names(input) <- mapper[["names"]][seq_along(input)]
  }
  state_k <- state
  first <- names(mapper[["mappers"]])[1]
  for (k in names(mapper[["mappers"]])) {
    eval2_k <- ibm_eval2(
      mapper[["mappers"]][[k]],
      input = input[[k]],
      state = state_k,
      ...
    )
    state_k <- eval2_k$offset
    if (k == first) {
      A <- eval2_k$jacobian
    } else {
      A <- eval2_k$jacobian %*% A
    }
  }
  list(offset = state_k, jacobian = A)
}


#' @describeIn ibm_simplify
#' Constructs a simplified `pipe` mapper. For fully linear pipes, calls
#' [ibm_linear()].
#' For partially non-linear pipes, replaces each sequence of linear mappers with
#' a single [bm_taylor()] mapper, while keeping the full list of
#' original mapper names, allowing the original `input` structure to be used
#' also with the simplified mappers, since the `taylor` mappers are not
#' dependent on inputs.
#' @export
#'
ibm_simplify.bm_pipe <- function(mapper,
                                 input = NULL,
                                 state = NULL,
                                 inla_f = FALSE,
                                 ...,
                                 n_state = NULL) {
  if (is.null(mapper[["names"]])) {
    mapper[["names"]] <- names(mapper[["mappers"]])
  }
  if (!is.null(input) && is.null(names(input))) {
    names(input) <- mapper[["names"]][seq_along(input)]
  }
  if (ibm_is_linear(mapper)) {
    if (is.null(state)) {
      n <- ibm_n(mapper[["mappers"]][[1]],
        ...,
        inla_f = FALSE,
        input = input[[names(mapper[["mappers"]])[1]]],
        n_state = n_state
      )
      state <- rep(0, n)
    }
    return(ibm_linear(mapper,
      input = input, state = state, ...,
      inla_f = inla_f, n_state = n_state
    ))
  }

  # Basic version, that just replaces each individual mapper with a
  # simplification.
  n <- ibm_n(
    mapper[["mappers"]][[1]],
    input = input[[names(mapper[["mappers"]])[1]]],
    state = state,
    ...,
    inla_f = FALSE,
    n_state = n_state
  )
  for (k in names(mapper[["mappers"]])) {
    if (ibm_is_linear(mapper[["mappers"]][[k]])) {
      mapper[["mappers"]][[k]] <- ibm_simplify(
        mapper[["mappers"]][[k]],
        input = input[[k]],
        ...,
        n_state = n
      )
    }
    n <- ibm_n_output(
      mapper[["mappers"]][[k]],
      input = input[[k]],
      ...,
      inla_f = FALSE,
      n_state = n
    )
  }

  mapper
}


## _multi ####

#' @title Mapper for tensor product domains
#' @param mappers A list of `bru_mapper` objects
#' @param simplify logical; If `TRUE`, removes trivial submappers.
#' Currently only sub-mappers of class [bm_index()] with `ibm_n() == 1L` are
#' removed, and only if the mappers are named (to avoid ordering mismatches).
#' Default: `FALSE`
#' @description Constructs a row-wise Kronecker product mapping of linear/affine
#'   mappers. Any offset in sub-mappers is added into a combined offset. Only
#'   linear/affine sub-mappers are allowed.
#' @rdname bm_multi
#' @export
#' @seealso [bru_mapper], [bru_mapper_generics]
#' @family mappers
#' @examples
#' (m <- bm_multi(list(a = bm_index(2), b = bm_index(3))))
#' ibm_eval2(m, list(a = c(1, 2, 1), b = c(1, 3, 2)), 1:6)
#'
bm_multi <- function(mappers, simplify = FALSE) {
  if (!is.list(mappers)) {
    stop("bm_multi requires a list of sub-mappers.")
  }
  if (is.null(names(mappers))) {
    names(mappers) <- as.character(seq_along(mappers))
  } else if (any(names(mappers) == "")) {
    stop("Either all or none of the multi sub-mappers should be named.")
  }
  mappers <- as_bm_list(mappers)

  if (simplify) {
    trivial <- vapply(mappers, function(x) {
      inherits(x, "bm_index") && (ibm_n(x) == 1L)
    }, logical(1))
    if (any(trivial)) {
      if (is.null(names(mappers))) {
        warning("Not removing trivial sub-mappers in unnamed bm_multi.")
      } else {
        mappers <- mappers[!trivial]
      }
    }
  }

  mapper <- list(
    mappers = mappers,
    n_multi = lapply(mappers, ibm_n),
    n_inla_multi = lapply(mappers, function(x) {
      ibm_n(x, inla_f = TRUE)
    }),
    values_multi = lapply(mappers, ibm_values),
    values_inla_multi = lapply(mappers, function(x) {
      ibm_values(x, inla_f = TRUE)
    }),
    is_linear_multi = lapply(mappers, ibm_is_linear),
    is_rowwise_multi = lapply(mappers, ibm_is_rowwise)
  )
  mapper[["n"]] <- prod(unlist(mapper[["n_multi"]]))
  mapper[["n_inla"]] <- prod(unlist(mapper[["n_inla_multi"]]))
  mapper[["is_linear"]] <- all(unlist(mapper[["is_linear_multi"]]))
  mapper[["is_rowwise"]] <- all(unlist(mapper[["is_rowwise_multi"]]))

  if (!mapper[["is_linear"]]) {
    stop("bm_multi sub-mappers must be linear mappers")
  }
  bru_mapper_define(mapper, new_class = "bm_multi")
}

#' @export
#' @rdname bm_multi
#' @param \dots Arguments passed on to [bm_multi()]
bru_mapper_multi <- function(...) {
  bm_multi(...)
}

#' @param multi logical;
#' If `TRUE` (or positive), recurse one level into sub-mappers
#' @export
#' @rdname ibm_n
#'
ibm_n.bm_multi <- function(mapper, inla_f = FALSE, multi = FALSE, ...) {
  if (multi) {
    if (inla_f) {
      mapper[["n_inla_multi"]]
    } else {
      mapper[["n_multi"]]
    }
  } else if (inla_f) {
    mapper[["n_inla"]]
  } else {
    mapper[["n"]]
  }
}

#' @export
#' @rdname ibm_n_output
#'
ibm_n_output.bm_multi <- function(mapper, input, ...) {
  input <- bm_multi_prepare_input(mapper, input)
  # Assume that the first mapper fully handles the output size
  nm <- names(mapper[["mappers"]])[1]
  if (is.matrix(input)) {
    ibm_n_output(mapper[["mappers"]][[nm]], input[, nm, drop = TRUE], ...)
  } else {
    stopifnot(is.list(input)) # data.frame or list
    ibm_n_output(mapper[["mappers"]][[nm]], input[[nm]], ...)
  }
}


#' @export
#' @rdname ibm_values
#'
ibm_values.bm_multi <- function(mapper,
                                inla_f = FALSE,
                                multi = FALSE,
                                ...) {
  if (multi) {
    vals <- if (inla_f) {
      mapper[["values_inla_multi"]]
    } else {
      mapper[["values_multi"]]
    }
    if (any(vapply(vals, is.null, logical(1)))) {
      return(NULL)
    }
    # Expand indices/values. First sub-mapper varies fastest
    as.data.frame(
      do.call(
        expand.grid,
        c(
          vals,
          list(
            KEEP.OUT.ATTRS = FALSE,
            stringsAsFactors = FALSE
          )
        )
      )
    )
  } else if (inla_f) {
    seq_len(mapper[["n_inla"]])
  } else {
    seq_len(mapper[["n"]])
  }
}

#' @export
#' @rdname ibm_is_linear
#'
ibm_is_linear.bm_multi <- function(mapper, multi = FALSE, ...) {
  if (multi) {
    mapper[["is_linear_multi"]]
  } else {
    mapper[["is_linear"]]
  }
}

bm_multi_prepare_input <- function(mapper, input) {
  if (is.matrix(input)) {
    input_names <- colnames(input)
    input <- as.data.frame(input)
  } else {
    input_names <- names(input)
  }
  if (is.null(input_names)) {
    if (is.data.frame(input)) {
      names(input) <- names(mapper[["mappers"]])[seq_len(ncol(input))]
    } else {
      if (is.null(input)) {
        input <- rep(list(NULL), length(mapper[["mappers"]]))
      }
      # input should now be a list
      names(input) <- names(mapper[["mappers"]])[seq_along(input)]
    }
  }
  input
}

#' @describeIn ibm_jacobian
#' Accepts a list with
#' named entries, or a list with unnamed but ordered elements.
#' The names must match the sub-mappers, see [ibm_names.bm_multi()].
#' Each list element should take a format accepted by the corresponding
#' sub-mapper. In case each element is a vector, the input can be given as a
#' data.frame with named columns, a matrix with named columns, or a matrix
#' with unnamed but ordered columns.
#' @param sub_A Internal; precomputed Jacobian matrices.
#' @export
#'
ibm_jacobian.bm_multi <- function(mapper,
                                  input,
                                  state = NULL,
                                  inla_f = FALSE,
                                  multi = FALSE,
                                  ...,
                                  sub_A = NULL) {
  input <- bm_multi_prepare_input(mapper, input)

  # Note: the multi-mapper only handles linear/affine sub-mappers.
  if (is.null(sub_A)) {
    sub_A <-
      lapply(
        names(mapper[["mappers"]]),
        function(x) {
          ibm_jacobian(mapper[["mappers"]][[x]],
            input = input[[x]],
            state = NULL,
            inla_f = inla_f,
            multi = FALSE
          )
        }
      )
  }
  if (multi) {
    return(sub_A)
  }
  # Combine the matrices
  # (A1, A2, A3) -> rowkron(A3, rowkron(A2, A1))
  A_ <- sub_A[[1]]
  for (k in seq_len(length(mapper[["mappers"]]) - 1)) {
    A_ <- fm_row_kron(sub_A[[k + 1]], A_)
  }
  A_
}


#' @export
#' @rdname ibm_linear
#'
ibm_linear.bm_multi <- function(mapper, input, state,
                                inla_f = FALSE,
                                ...) {
  eval2 <- ibm_eval2(
    mapper, input,
    state = state,
    inla_f = inla_f,
    multi = FALSE,
    ...
  )
  bru_mapper_taylor(
    offset = eval2$offset,
    jacobian = eval2$jacobian,
    state0 = state,
    values_mapper = mapper
  )
}


#' @param jacobian
#' For `ibm_eval()` methods, an optional pre-computed Jacobian, typically
#' supplied by internal methods that already have the Jacobian.
#' @param pre_A `r lifecycle::badge("deprecated")` in favour of `jacobian`.
#' @export
#' @rdname ibm_eval
#'
ibm_eval.bm_multi <- function(mapper, input, state = NULL,
                              inla_f = FALSE, ...,
                              jacobian = NULL,
                              pre_A = deprecated()) {
  input <- bm_multi_prepare_input(mapper, input)

  if (ibm_n(mapper) == 0L) {
    # TODO: Ensure that this code is only reached when the first sub-mapper
    # is constant (mapper_const or a pipe that is state-invariant)
    val <- 0
    m <- names(mapper[["mappers"]])[1]
    val_add <- ibm_eval(mapper[["mappers"]][[m]],
      input = input[[m]],
      state = NULL,
      inla_f = inla_f,
      multi = FALSE
    )
    val <- val + val_add
    return(val)
  }

  val <- rep(0, ibm_n_output(mapper, input, state = state, inla_f = inla_f))
  for (m in names(mapper[["mappers"]])) {
    val_add <- ibm_eval(mapper[["mappers"]][[m]],
      input = input[[m]],
      state = NULL,
      inla_f = inla_f,
      multi = FALSE
    )
    val <- val + val_add
  }
  if ((ibm_n(mapper) > 0) && !is.null(state)) {
    if (is.null(jacobian)) {
      jacobian <- ibm_jacobian(mapper,
        input = input, state = state,
        inla_f = inla_f, multi = FALSE, ...
      )
    }
    val <- val + jacobian %*% state
  }
  as.vector(val)
}


bm_multi_indexing <- function(mapper, input) {
  if (is.matrix(input)) {
    nms <- colnames(input)
    len <- ncol(input)
  } else {
    stopifnot(is.list(input))
    nms <- names(input)
    len <- length(input)
  }

  nms_mapper <- names(mapper[["mappers"]])
  if (is.null(nms)) {
    indexing <- seq_len(min(length(nms_mapper), length(input)))
  } else {
    indexing <- intersect(nms_mapper, nms)
  }
  names(indexing) <- indexing
  indexing
}


#' @describeIn ibm_invalid_output
#' Accepts a list with
#' named entries, or a list with unnamed but ordered elements.
#' The names must match the sub-mappers, see [ibm_names.bm_multi()].
#' Each list element should take a format accepted by the corresponding
#' sub-mapper. In case each element is a vector, the input can be given as a
#' data.frame with named columns, a matrix with named columns, or a matrix
#' with unnamed but ordered columns.
#' @export
#'
ibm_invalid_output.bm_multi <- function(mapper,
                                        input,
                                        state,
                                        inla_f = FALSE,
                                        multi = FALSE,
                                        ...) {
  indexing <- bm_multi_indexing(mapper, input)
  if (is.matrix(input)) {
    invalid <-
      lapply(
        indexing,
        function(x) {
          ibm_invalid_output(
            mapper[["mappers"]][[x]],
            input = input[, x],
            inla_f = inla_f,
            multi = FALSE
          )
        }
      )
  } else if (is.list(input)) {
    invalid <-
      lapply(
        indexing,
        function(x) {
          ibm_invalid_output(
            mapper[["mappers"]][[x]],
            input = input[[x]],
            inla_f = inla_f,
            multi = FALSE
          )
        }
      )
  } else {
    invalid <- as.list(rep(TRUE, length(mapper[["mappers"]])))
    names(invalid) <- names(mapper[["mappers"]])
  }
  if (!multi) {
    # Combine the vectors (v1, v2, v3) -> v1 | v2 | v3
    invalid_ <- invalid[[1]]
    for (k in seq_len(length(mapper[["mappers"]]) - 1)) {
      invalid_ <- invalid_ | invalid[[k + 1]]
    }
    return(invalid_)
  }
  invalid
}

#' @return
#' * `[`-indexing a `bm_multi` extracts a subset
#'   `bm_multi` object (for drop `FALSE`) or an individual sub-mapper
#'   (for drop `TRUE`, and `i` identifies a single element)
#' @export
#' @param x object from which to extract element(s)
#' @param i indices specifying element(s) to extract
#' @param drop logical;
#' For `[.bm_multi`, whether to extract an individual mapper when
#' `i` identifies a single element. If `FALSE`, a list of sub-mappers is
#' returned (suitable e.g. for creating a new `bm_multi` object).
#' Default: `TRUE`
#' @rdname bm_multi
`[.bm_multi` <- function(x, i, drop = TRUE) {
  if (is.logical(i)) {
    i <- which(i)
  }
  mapper <- x[["mappers"]][i]
  if (drop) {
    if (length(mapper) == 1) {
      mapper <- mapper[[1]]
    } else if (length(mapper) == 0) {
      mapper <- NULL
    }
  }
  mapper
}

#' @rdname bm_multi
#' @export
`[.bru_mapper_multi` <- function(x, i, drop = TRUE) {
  if (is.logical(i)) {
    i <- which(i)
  }
  mapper <- x[["mappers"]][i]
  if (drop) {
    if (length(mapper) == 1) {
      mapper <- mapper[[1]]
    } else if (length(mapper) == 0) {
      mapper <- NULL
    }
  }
  mapper
}

#' @describeIn ibm_names
#' Returns the names from the sub-mappers list
#' @export
#'
`ibm_names.bm_multi` <- function(mapper) {
  names(mapper[["mappers"]])
}

#' @param value a character vector of up to the same length as the number
#' of mappers in the multi-mapper x
#' @export
#' @rdname ibm_names
#'
`ibm_names<-.bm_multi` <- function(mapper, value) {
  names(mapper[["mappers"]]) <- value
  names(mapper[["n_multi"]]) <- value
  names(mapper[["n_inla_multi"]]) <- value
  names(mapper[["values_multi"]]) <- value
  names(mapper[["values_inla_multi"]]) <- value
  mapper
}

#' @export
#' @rdname ibm_names
`ibm_names<-.bru_mapper_multi` <- function(mapper, value) {
  names(mapper[["mappers"]]) <- value
  names(mapper[["n_multi"]]) <- value
  names(mapper[["n_inla_multi"]]) <- value
  names(mapper[["values_multi"]]) <- value
  names(mapper[["values_inla_multi"]]) <- value
  mapper
}

## _harmonics ####

#' @title Mapper for cos/sin functions
#' @param order For `bm_harmonics`, specifies the maximum `cos`/`sin`
#' order. (Default 1)
#' @param scaling For `bm_harmonics`, specifies an optional vector of
#'  scaling factors of length `intercept + order`, or a common single scalar.
#' @param intercept logical; For `bm_harmonics`, if `TRUE`, the first
#'   basis function is a constant. (Default `TRUE`)
#' @param interval numeric length-2 vector specifying a domain interval.
#'   Default `c(0, 1)`.
#' @description Constructs a mapper for `cos`/`sin` functions of orders 1 (if
#'   `intercept` is `TRUE`, otherwise 0) through `order`. The total number of
#'   basis functions is `intercept + 2 * order`.
#'
#'   Optionally, each order can be given a non-unit scaling, via the `scaling`
#'   vector, of length `intercept + order`. This can be used to
#'   give an effective spectral prior. For example, let
#'   ```
#'   scaling = 1 / (1 + (0:4)^2)
#'   x <- seq(0, 1, length.out = 11)
#'   bmh1 = bm_harmonics(order = 4, interval = c(0, 1))
#'   u1 <- ibm_eval(
#'     bmh1,
#'     input = x,
#'     state = rnorm(9, sd = rep(scaling, c(1, 2, 2, 2, 2)))
#'   )
#'   ```
#'   Then, with
#'   ```
#'   bmh2 = bm_harmonics(order = 4, scaling = scaling)
#'   u2 = ibm_eval(bmh2, input = x, state = rnorm(9))
#'   ```
#'   the stochastic properties of `u1` and `u2` will be the same, with
#'   `scaling^2` determining the variance for each frequency contribution.
#'
#'   The period for the first order harmonics is shifted and scaled to match
#'   `interval`.
#' @export
#' @seealso [bru_mapper], [bru_mapper_generics]
#' @family mappers
#' @examples
#' m <- bm_harmonics(2)
#' ibm_eval2(m, input = c(0, pi / 4, pi / 2, 3 * pi / 4), 1:5)
#'
#' @rdname bm_harmonics
#'
bm_harmonics <- function(order = 1,
                         scaling = 1,
                         intercept = TRUE,
                         interval = c(0, 1)) {
  if (length(scaling) == 1) {
    scaling <- rep(scaling, intercept + order)
  } else {
    stopifnot(length(scaling) == intercept + order)
  }
  mapper <- list(
    order = order,
    scaling = scaling,
    intercept = intercept,
    interval = interval
  )
  bru_mapper_define(mapper, new_class = "bm_harmonics")
}

#' @export
#' @rdname bm_harmonics
#' @param \dots Arguments passed on to [bm_harmonics()]
bru_mapper_harmonics <- function(...) {
  bm_harmonics(...)
}

#' @export
#' @rdname ibm_n
#'
ibm_n.bm_harmonics <- function(mapper, inla_f = FALSE, ...) {
  mapper[["order"]] * 2 + mapper[["intercept"]]
}

#' @export
#' @rdname ibm_jacobian
#'
ibm_jacobian.bm_harmonics <- function(mapper,
                                      input,
                                      state = NULL,
                                      inla_f = FALSE,
                                      ...) {
  # Indexing into sparseMatrix is slower than into a dense Matrix,
  # and indexing into a dense Matrix is slower than for a plain matrix(),
  # including conversion to dense Matrix afterwards.
  A <- matrix(0.0, NROW(input), ibm_n(mapper))
  off <- 0
  if (mapper[["intercept"]]) {
    A[, 1] <- 1.0 * mapper[["scaling"]][1]
    off <- off + 1
  }
  if (mapper[["order"]] > 0) {
    input <- (input - mapper[["interval"]][1]) / diff(mapper[["interval"]])
    for (ord in seq_len(mapper[["order"]])) {
      scale <- mapper[["scaling"]][mapper[["intercept"]] + ord]
      angle <- (2 * pi * ord) * input
      A[, off + 1] <- cos(angle) * scale
      A[, off + 2] <- sin(angle) * scale
      off <- off + 2
    }
  }
  as(A, "Matrix")
}


## _reparam ####

#' @title Mapper for reparameterising mapper states
#' @param mapper A `bru_mapper` object
#' @param B a square or rectangular basis conversion matrix
# @param \dots Arguments passed on to submethods
#' @export
#' @description Creates a mapper for handling basis conversions. Functionally
#'   equivalent to `bm_pipe(list(bm_matrix(ncol(B)), mapper))`, but with an
#'   internally stored matrix input `B` for efficiency, and allowing the mapper
#'   `input` format to be identical to that of the original `mapper`.
#' @seealso [bru_mapper], [bru_mapper_generics]
#' @family mappers
#' @rdname bm_reparam
#' @examples
#' # 2->2 reparameterisation; (u1,u2) -> (u1, u1 + u2)
#' (m <- bm_reparam(bm_index(2), B = matrix(c(1, 1, 0, 1), 2, 2)))
#'
#' # 2->3 reparameterisation; (u1,u2) -> (u1, u2, u1+u2)
#' # This is an example of a low-rank representation of a higher-dimensional
#' # state vector.
#' (m <- bm_reparam(bm_index(3), B = cbind(c(1, 0, 1), c(0, 1, 1))))
#'
bm_reparam <- function(mapper, B) {
  map_n <- ibm_n(mapper)
  stopifnot(is.na(map_n) || (map_n == nrow(B)))
  map <- list(
    mapper = mapper,
    B = B,
    n = ncol(B),
    is_linear = ibm_is_linear(mapper),
    is_rowwise = ibm_is_rowwise(mapper)
  )
  bru_mapper_define(map, new_class = "bm_reparam")
}


#' @title Deprecated basis conversion mapper
#' @description `r lifecycle::badge("deprecated")`
#'  Old deprecated name for [bm_reparam()]
#' @param mesh Object supported by a [bru_mapper()] method.
#' @param B a square or tall basis conversion matrix
#' @export
#' @keywords internal
bm_mesh_B <- function(mesh, B) {
  bm_reparam(bru_mapper(mesh), B)
}

#' @export
#' @describeIn bm_mesh_B `r lifecycle::badge("deprecated")`
#'   Deprecated name for `bm_mesh_B`
bru_mapper_mesh_B <- function(...) {
  bm_mesh_B(...)
}

#' @export
#' @rdname ibm_n
#' @inheritParams ibm_n
ibm_n.bm_reparam <- function(mapper, ...) {
  mapper[["n"]]
}
#' @export
#' @rdname ibm_values
ibm_values.bm_reparam <- function(mapper, ...) {
  seq_len(mapper[["n"]])
}
#' @export
#' @rdname ibm_n_output
ibm_n_output.bm_reparam <- function(mapper, ...) {
  ibm_n_output(mapper[["mapper"]], ...)
}
#' @export
#' @rdname ibm_jacobian
ibm_jacobian.bm_reparam <- function(mapper, input, state = NULL, ...) {
  if (is.null(input)) {
    return(Matrix::Matrix(0, 0, ibm_n(mapper)))
  }
  if (!is.null(state)) {
    state <- mapper[["B"]] %*% state
  }
  A <- ibm_jacobian(mapper[["mapper"]], input = input, state = state, ...)
  A %*% mapper[["B"]]
}
#' @export
#' @rdname ibm_eval
ibm_eval.bm_reparam <- function(mapper, input, state = NULL, ...,
                                jacobian = NULL) {
  if (is.null(input)) {
    return(Matrix::Matrix(0, 0, ibm_n(mapper)))
  }
  if (!is.null(state)) {
    state <- mapper[["B"]] %*% state
  }
  # Note: Have to prevent the jacobian to be fed through, as it's the jacobian
  # for the remapped state, not the original state.
  val <- ibm_eval(mapper[["mapper"]], input = input, state = state, ...)
  val
}
