#' Reposition brain slices
#'
#' Function for repositioning
#' pre-joined atlas data (i.e. data and atlas
#' already joined to a single data frame).
#' This makes it possible for users to
#' reposition the geometry data for the atlas
#' for control over final plot layout. For even
#' more detailed control over the positioning,
#' the "hemi" and "view" columns should be
#' converted into factors and ordered by wanted
#' order of appearance.
#'
#' @param data sf-data.frame of joined brain atlas and data
#' @param position Position formula for slices. For cortical atlases, use
#'   formulas like `hemi ~ view`. For subcortical/tract atlases, use
#'   "horizontal", "vertical", or `type ~ .` for type-based layout.
#' @param nrow Number of rows for grid layout (subcortical/tract only)
#' @param ncol Number of columns for grid layout (subcortical/tract only)
#' @param views Character vector specifying view order (subcortical/tract only)
#'
#' @return sf-data.frame with re-positioned slices
#' @export
#'
#' @examples
#' reposition_brain(dk(), hemi ~ view)
#' reposition_brain(dk(), view ~ hemi)
#' reposition_brain(dk(), hemi + view ~ .)
#' reposition_brain(dk(), . ~ hemi + view)
#'
#' \donttest{
#' reposition_brain(aseg(), nrow = 2)
#' reposition_brain(aseg(), views = c("sagittal", "axial_3"))
#' }
reposition_brain <- function(
  data,
  position = "horizontal",
  nrow = NULL,
  ncol = NULL,
  views = NULL
) {
  data <- as.data.frame(
    data,
    stringsAsFactors = FALSE
  ) |>
    frame_2_position(position, nrow = nrow, ncol = ncol, views = views)
}


#' Alter brain atlas position
#'
#' Function to be used in the position argument in geom_brain
#' to alter the position of the brain slice/views.
#'
#' @param position Formula describing the rows ~ columns organisation for
#'   cortical atlases (e.g., `hemi ~ view`). For subcortical/tract atlases,
#'   can be "horizontal", "vertical", or a formula with `type ~ .` where type
#'   is extracted from view names like "axial_1" -> "axial".
#' @param nrow Number of rows for grid layout. If NULL (default), calculated
#'   automatically. Only used for subcortical/tract atlases when position is
#'   not a formula.
#' @param ncol Number of columns for grid layout. If NULL (default), calculated
#'   automatically. Only used for subcortical/tract atlases when position is
#'   not a formula.
#' @param views Character vector specifying which views to include and their
#'   order. If NULL (default), all views are included in their original order.
#'   Only applies to subcortical/tract atlases.
#'
#' @export
#' @return a ggproto object
#' @importFrom ggplot2 ggproto
#' @examples
#' library(ggplot2)
#'
#' # Cortical atlas with formula
#' ggplot() +
#'   geom_brain(
#'     atlas = dk(), aes(fill = region),
#'     position = position_brain(. ~ view + hemi),
#'     show.legend = FALSE
#'   )
#'
#' ggplot() +
#'   geom_brain(
#'     atlas = dk(), aes(fill = region),
#'     position = position_brain(view ~ hemi),
#'     show.legend = FALSE
#'   )
#'
#' \donttest{
#' ggplot() +
#'   geom_brain(
#'     atlas = aseg(), aes(fill = region),
#'     position = position_brain(nrow = 2)
#'   )
#'
#' ggplot() +
#'   geom_brain(
#'     atlas = aseg(), aes(fill = region),
#'     position = position_brain(
#'       views = c("sagittal", "axial_3", "coronal_2"),
#'       nrow = 1
#'     )
#'   )
#'
#' ggplot() +
#'   geom_brain(
#'     atlas = aseg(), aes(fill = region),
#'     position = position_brain(type ~ .)
#'   )
#' }
position_brain <- function(
  position = "horizontal",
  nrow = NULL,
  ncol = NULL,
  views = NULL
) {
  ggproto(
    NULL,
    PositionBrain,
    position = position,
    nrow = nrow,
    ncol = ncol,
    views = views
  )
}

#' ggproto Position class for brain atlas layout
#'
#' Handles coordinate repositioning of brain views/slices during
#' the ggplot2 rendering pipeline. Created by [position_brain()].
#'
#' @keywords internal
#' @importFrom ggplot2 ggproto Position
#' @noRd
PositionBrain <- ggproto(
  "PositionBrain",
  Position,
  position = hemi + view ~ .,
  nrow = NULL,
  ncol = NULL,
  views = NULL,
  setup_params = function(self, data) {
    list(
      position = self$position,
      nrow = self$nrow,
      ncol = self$ncol,
      views = self$views
    )
  },
  compute_layer = function(self, data, params, layout) {
    df3 <- frame_2_position(
      data,
      params$position,
      nrow = params$nrow,
      ncol = params$ncol,
      views = params$views
    )
    bbx <- sf::st_bbox(df3$geometry)

    if (is.null(layout$coord$limits$y)) {
      layout$coord$limits$y <- bbx[c(2, 4)]
    }

    if (is.null(layout$coord$limits$x)) {
      layout$coord$limits$x <- bbx[c(1, 3)]
    }

    df3
  }
)

# geometry movers ----

#' Extract and validate variable names from a position formula
#'
#' @param pos A formula describing the layout.
#' @return Character vector of variable names (excluding `.`).
#' @keywords internal
#' @noRd
parse_formula_vars <- function(pos) {
  chosen <- all.vars(pos, unique = FALSE)
  chosen <- chosen[!grepl(".", chosen, fixed = TRUE)]

  if (anyDuplicated(chosen)) {
    cli::cli_abort(
      "Cannot position brain with the same data as columns and rows"
    )
  }
  chosen
}

#' Detect stacking direction from a formula with `+`
#'
#' For formulas like `hemi + view ~ .` or `. ~ hemi + view`,
#' determines whether the stacked layout is row-based or column-based.
#'
#' @param pos A formula containing `+` on one side and `.` on the other.
#' @return `"rows"` or `"columns"`.
#' @keywords internal
#' @noRd
stacking_direction <- function(pos) {
  if (grepl("~\\s*\\.", deparse(pos))) "rows" else "columns"
}

#' Validate that a single-direction formula includes both `.` and `~`
#'
#' @param pos A formula.
#' @param position The resolved position (`"rows"`, `"columns"`, or vars).
#' @keywords internal
#' @noRd
validate_stacking_formula <- function(pos, position) {
  is_single <- length(position) == 1 && position %in% c("rows", "columns")
  if (!is_single) return(invisible())

  has_both <- sum(grepl("\\.|~", pos)) == 2
  if (!has_both) {
    cli::cli_abort(
      "Formula for a single row or column must contain both a '.' and '~'"
    )
  }
}

#' Parse a position formula into layout instructions
#'
#' Interprets a formula like `hemi ~ view` into row/column
#' variable names and validates against the atlas type.
#'
#' @param pos A formula describing the layout.
#' @param data Data.frame with atlas columns (`type`, `hemi`, `view`).
#'
#' @return A list with `position` (character), `chosen` (variable names),
#'   and `data` (possibly modified data.frame).
#' @keywords internal
#' @noRd
position_formula <- function(pos, data) {
  chosen <- parse_formula_vars(pos)
  atlas_type <- unique(data$type)[1]

  if (atlas_type == "cortical") {
    position <- position_cortical(pos, chosen)
  } else {
    result <- position_subcortical(pos, chosen, data)
    position <- result$position
    chosen <- result$chosen
    data <- result$data
  }

  validate_stacking_formula(pos, position)
  list(position = position, chosen = chosen, data = data)
}

#' Resolve position for a cortical atlas formula
#'
#' @param pos A formula.
#' @param chosen Character vector of variable names.
#' @return Position specification: variable names or `"rows"`/`"columns"`.
#' @keywords internal
#' @noRd
position_cortical <- function(pos, chosen) {
  if (length(chosen) < 2) {
    missing_vars <- c("view", "hemi")[!c("view", "hemi") %in% chosen]
    cli::cli_abort(c(
      "Position formula not correct.",
      "x" = paste("Missing:", paste(missing_vars, collapse = " & "))
    ))
  }
  if (any(grepl("+", pos, fixed = TRUE))) {
    stacking_direction(pos)
  } else {
    chosen
  }
}

#' Resolve position for a subcortical/tract atlas formula
#'
#' @param pos A formula.
#' @param chosen Character vector of variable names.
#' @param data Data.frame with atlas columns.
#' @return A list with `position`, `chosen`, and `data`.
#' @keywords internal
#' @noRd
position_subcortical <- function(pos, chosen, data) {
  if ("type" %in% chosen) {
    data$.view_type <- extract_view_type(data$view)
    chosen[chosen == "type"] <- ".view_type"
  }

  position <- if (length(chosen) == 1) {
    stacking_direction(pos)
  } else {
    chosen
  }

  list(position = position, chosen = chosen, data = data)
}


#' Extract the type prefix from view names
#'
#' Splits view names like `"axial_3"` on underscore and returns
#' the first part (`"axial"`).
#'
#' @param views Character vector of view names.
#'
#' @return Character vector of type prefixes.
#' @keywords internal
#' @noRd
extract_view_type <- function(views) {
  vapply(
    views,
    function(v) {
      parts <- strsplit(v, "_", fixed = TRUE)[[1]]
      if (length(parts) >= 1) parts[1] else v # nocov
    },
    character(1),
    USE.NAMES = FALSE
  )
}

#' Reposition brain views according to layout specification
#'
#' Main dispatcher that splits data by view/hemisphere, gathers
#' geometry, and delegates to the appropriate stacking function.
#'
#' @param data Data.frame with atlas columns and `geometry`.
#' @param pos Position specification: a formula, `"horizontal"`,
#'   or `"vertical"`.
#' @param nrow Number of grid rows (optional).
#' @param ncol Number of grid columns (optional).
#' @param views Character vector of views to include.
#'
#' @return An sf data.frame with repositioned geometry and
#'   adjusted bounding box.
#' @importFrom sf st_as_sf
#' @keywords internal
#' @noRd
frame_2_position <- function(
  data,
  pos,
  nrow = NULL,
  ncol = NULL,
  views = NULL
) {
  if (!is.null(views)) {
    data <- data[data$view %in% views, , drop = FALSE]
    data$view <- factor(data$view, levels = views)
    data <- data[order(data$view), ]
    data$view <- as.character(data$view)
  }

  if (!is.null(nrow) || !is.null(ncol)) {
    dfpos <- split_data_grid(data, nrow, ncol)
  } else {
    dfpos <- split_data(data, pos)
  }

  df2 <- lapply(dfpos$data, gather_geometry)
  posi <- ifelse(length(dfpos$position) > 1, "grid", dfpos$position)

  df3 <- switch(
    posi,
    rows = stack_vertical(df2),
    columns = stack_horizontal(df2),
    grid = stack_grid(df2, dfpos$position[1], dfpos$position[2])
  )

  df4 <- st_as_sf(df3$df)
  attr(sf::st_geometry(df4), "bbox") <- df3$box

  df4
}


#' Split atlas data into a grid of views
#'
#' Assigns grid row/column indices to each view and returns
#' a list of per-view data.frames.
#'
#' @param data Data.frame with a `view` column.
#' @param nrow Number of grid rows (optional, auto-calculated).
#' @param ncol Number of grid columns (optional, auto-calculated).
#'
#' @return A list with `data` (list of data.frames) and
#'   `position` (column names for grid coordinates).
#' @keywords internal
#' @noRd
split_data_grid <- function(data, nrow = NULL, ncol = NULL) {
  view_list <- unique(data$view)
  n_views <- length(view_list)

  if (is.null(nrow) && is.null(ncol)) {
    ncol <- ceiling(sqrt(n_views))
    nrow <- ceiling(n_views / ncol)
  } else if (is.null(nrow)) {
    nrow <- ceiling(n_views / ncol)
  } else if (is.null(ncol)) {
    ncol <- ceiling(n_views / nrow)
  }

  data$.grid_row <- ((seq_along(view_list) - 1) %/% ncol + 1)[
    match(data$view, view_list)
  ]
  data$.grid_col <- ((seq_along(view_list) - 1) %% ncol + 1)[
    match(data$view, view_list)
  ]

  df_list <- lapply(view_list, function(v) {
    data[data$view == v, ]
  })

  list(
    data = df_list,
    position = c(".grid_row", ".grid_col")
  )
}

#' Split atlas data by position specification
#'
#' Routes to formula-based or string-based splitting depending
#' on whether `position` is a formula or character.
#'
#' @param data Data.frame with atlas columns.
#' @param position A formula or character layout specification.
#'
#' @return A list with `data` (list of data.frames) and
#'   `position` (layout direction or variable names).
#' @keywords internal
#' @noRd
split_data <- function(data, position) {
  if (inherits(position, "formula")) {
    pos <- position_formula(position, data)
    if (!is.null(pos$data)) {
      data <- pos$data
    }
    df2 <- dplyr::group_by_at(data, pos$chosen)
    df2 <- dplyr::group_split(df2)
    pos <- pos$position
  } else {
    layout_direction <- "columns"
    if (length(position) == 1) {
      if (position %in% c("horizontal", "vertical")) {
        layout_direction <- if (position == "vertical") "rows" else "columns"
        position <- default_order(data)
      }
    }
    pos <- as.data.frame(
      strsplit(position, " ", fixed = TRUE),
      stringsAsFactors = FALSE
    )
    atlas_type <- unique(data$type)[1]
    if (atlas_type == "cortical") {
      k <- cbind(
        pos[2, ] %in% data$view,
        pos[1, ] %in% data$hemi
      )
      k <- vapply(seq_len(nrow(k)), function(x) sum(k[x, ]), numeric(1))
      pos <- pos[k == 2]

      df2 <- lapply(pos, function(x) {
        data[data$hemi == x[1] & data$view == x[2], ]
      })
    } else {
      df2 <- lapply(pos, function(x) {
        data[data$view == x, ]
      })
    }
    pos <- layout_direction
  }

  list(data = df2, position = pos)
}

#' Zero-origin geometry for a single view
#'
#' Translates geometry so the bounding box starts at (0, 0).
#'
#' @param df Data.frame with a `geometry` column.
#'
#' @return Modified data.frame with shifted geometry.
#' @keywords internal
#' @noRd
gather_geometry <- function(df) {
  bbx <- sf::st_bbox(df$geometry)
  df$geometry <- df$geometry - bbx[c("xmin", "ymin")]
  df
}

#' Center a view within a grid cell
#'
#' Offsets geometry so the view is centered inside a cell of
#' `cell_size` at position `grid_pos`.
#'
#' @param df Data.frame with a `geometry` column.
#' @param cell_size Numeric length-2 vector `c(width, height)`.
#' @param grid_pos Numeric length-2 vector `c(x, y)` offset.
#'
#' @return Modified data.frame with repositioned geometry.
#' @keywords internal
#' @noRd
center_view <- function(df, cell_size, grid_pos) {
  bbox <- sf::st_bbox(df$geometry)
  view_width <- bbox["xmax"] - bbox["xmin"]
  view_height <- bbox["ymax"] - bbox["ymin"]
  view_size <- c(view_width, view_height)
  center_offset <- (cell_size - view_size) / 2
  df$geometry <- df$geometry + grid_pos + center_offset
  df
}

#' Arrange views in a horizontal row
#'
#' @param df List of data.frames, each with a `geometry` column.
#'
#' @return A list with `df` (combined data.frame) and `box` (bbox).
#' @keywords internal
#' @noRd
stack_horizontal <- function(df) {
  sep <- get_sep(df)
  cell_size <- sep / 1.2

  bx <- list()
  for (k in seq_along(df)) {
    df[[k]] <- center_view(df[[k]], cell_size, c((k - 1) * sep[1], 0))
    bx[[k]] <- sf::st_bbox(df[[k]]$geometry)
  }

  list(df = do.call(rbind, df), box = get_box(bx))
}

#' Arrange views in a vertical column
#'
#' @param df List of data.frames, each with a `geometry` column.
#'
#' @return A list with `df` (combined data.frame) and `box` (bbox).
#' @keywords internal
#' @noRd
stack_vertical <- function(df) {
  sep <- get_sep(df)
  cell_size <- sep / 1.2

  bx <- list()
  for (k in seq_along(df)) {
    df[[k]] <- center_view(df[[k]], cell_size, c(0, (k - 1) * sep[2]))
    bx[[k]] <- sf::st_bbox(df[[k]]$geometry)
  }

  list(df = do.call(rbind, df), box = get_box(bx))
}

#' Arrange views in a row-by-column grid
#'
#' @param df List of data.frames, each with a `geometry` column
#'   and grid assignment columns.
#' @param rows Column name identifying the row variable.
#' @param columns Column name identifying the column variable.
#'
#' @return A list with `df` (combined data.frame) and `box` (bbox).
#' @keywords internal
#' @noRd
stack_grid <- function(df, rows, columns) {
  sep <- get_sep(df)
  lookup <- grid_lookup(df, rows, columns)

  grid <- expand.grid(
    col_idx = seq_along(lookup$col_vals),
    row_idx = seq_along(lookup$row_vals)
  )

  cell_size <- sep / 1.2
  df_positioned <- Map(
    function(ri, ci) {
      idx <- which(
        lookup$df_rows == lookup$row_vals[ri] &
          lookup$df_cols == lookup$col_vals[ci]
      )
      if (length(idx) != 1) return(NULL)
      grid_pos <- c((ci - 1) * sep[1], (ri - 1) * sep[2])
      center_view(df[[idx]], cell_size, grid_pos)
    },
    grid$row_idx,
    grid$col_idx
  )

  df_positioned <- Filter(Negate(is.null), df_positioned)
  bx <- lapply(df_positioned, function(x) sf::st_bbox(x$geometry))
  result_df <- drop_temp_columns(do.call(rbind, df_positioned))

  list(df = result_df, box = get_box(bx))
}

#' Build a lookup of row/column values per data.frame element
#'
#' @param df List of data.frames from a split atlas.
#' @param rows Column name for the row variable.
#' @param columns Column name for the column variable.
#'
#' @return A list with `df_rows`, `df_cols` (per-element values),
#'   and `row_vals`, `col_vals` (unique levels).
#' @keywords internal
#' @noRd
grid_lookup <- function(df, rows, columns) {
  as_char <- function(x) {
    if (is.numeric(x)) as.character(x) else x
  }
  df_rows <- vapply(
    df, function(x) as_char(unique(x[[rows]])), character(1)
  )
  df_cols <- vapply(
    df, function(x) as_char(unique(x[[columns]])), character(1)
  )
  list(
    df_rows = df_rows,
    df_cols = df_cols,
    row_vals = unique(df_rows),
    col_vals = unique(df_cols)
  )
}

#' Remove temporary columns added during grid layout
#'
#' @param df A data.frame.
#' @return The data.frame with temp columns removed.
#' @keywords internal
#' @noRd
drop_temp_columns <- function(df) {
  temp <- c(
    "xmin", "xmax", "ymin", "ymax",
    ".grid_row", ".grid_col", ".view_type"
  )
  temp <- temp[temp %in% names(df)]
  if (length(temp) > 0) df[, temp] <- NULL
  df
}

#' Compute a padded bounding box from a list of bboxes
#'
#' @param bx List of sf bbox objects.
#'
#' @return An sf `bbox` object with 1% padding.
#' @keywords internal
#' @noRd
get_box <- function(bx) {
  bx <- do.call(rbind, bx)
  pad <- max(bx) * 0.01
  bx <- c(
    -pad,
    -pad,
    max(bx[, "xmax"]) + pad,
    max(bx[, "ymax"]) + pad
  )
  x <- stats::setNames(bx, c("xmin", "ymin", "xmax", "ymax"))
  class(x) <- "bbox"
  x
}

#' Compute cell separation distances
#'
#' Determines the x and y spacing between grid cells based on
#' the maximum bounding-box dimensions across all views.
#'
#' @param data List of data.frames, each with a `geometry` column.
#'
#' @return Named numeric vector `c(x = ..., y = ...)`.
#' @keywords internal
#' @noRd
get_sep <- function(data) {
  get_bbox <- function(x) sf::st_bbox(x$geometry)
  bboxes <- vapply(data, get_bbox, numeric(4))
  sep <- c(max(bboxes[3, ]), max(bboxes[4, ]))
  c("x" = sep[1] + sep[1] * 0.2, "y" = sep[2] + sep[2] * 0.2)
}

#' Generate the default view ordering for an atlas
#'
#' For cortical atlases, returns `"hemi view"` pairs (left first).
#' For subcortical/tract atlases, returns unique views as-is.
#'
#' @param data Data.frame with `type`, `view`, and `hemi` columns.
#'
#' @return Character vector of ordered view identifiers.
#' @keywords internal
#' @noRd
default_order <- function(data) {
  if (unique(data$type) != "cortical") {
    return(unique(data$view))
  }
  sides <- unique(data$view)
  left_sides <- sides[sides %in% unique(data$view[data$hemi == "left"])]
  right_sides <- sides[sides %in% unique(data$view[data$hemi == "right"])]
  left_views <- paste("left", left_sides)
  right_views <- paste("right", right_sides)
  c(left_views, right_views)
}
