#' @title DataBackend for Raster Objects
#'
#' @description
#' [mlr3::DataBackend] for [terra::SpatRaster] raster objects.
#'
#' @param rows `integer()`\cr
#'   Row indices. Row indices start with 1 in the upper left corner in the
#'   raster, increase from left to right and then from top to bottom. The last
#'   cell is in the bottom right corner and the row index equals the number of
#'   cells in the raster.
#' @param cols `character()`\cr
#'   Column names.
#'
#' @section Read mode:
#'  There are two different ways the reading of values is performed internally:
#' * "Block mode" reads complete rows of the raster file and subsets the requested cells.
#'   This mode is faster than "cell mode" if the complete raster file is iterated over.
#'
#' * "Cell mode" reads individual cells.
#'   This is faster than "block mode" if only a few cells are sampled.
#'
#' "Block mode" is activated if `$data(rows)` is used with a increasing integer sequence e.g. `200:300`.
#' If only a single cell is requested, "cell mode" is used.
#'
#' @importFrom terra writeRaster writeStart writeStop rast cats sources intersect readStart readStop rowColFromCell readValues head unique ncell nlyr ncol
#' @importFrom utils tail
#' @importFrom methods as
#' @export
DataBackendRaster = R6Class(
  "DataBackendRaster",
  inherit = DataBackend,
  cloneable = FALSE,
  public = list(
    #' @description
    #' Creates a new instance of this [R6][R6::R6Class] class.
    #'
    #' @template param-data
    #'
    initialize = function(data) {
      assert_class(data, "SpatRaster")

      # record the source file and band index within that source for each layer
      # in-memory layers are written to individual files that hold a single band
      layers = map(seq_len(terra::nlyr(data)), function(i) {
        layer = data[[i]]
        if (terra::inMemory(layer)) {
          filename = tempfile(fileext = ".tif")
          terra::writeRaster(layer, filename = filename)
          list(source = filename, band = 1L)
        } else {
          src = terra::sources(layer, bands = TRUE)
          list(source = src$source, band = as.integer(src$bands))
        }
      })

      private$.sources = map_chr(layers, "source")
      private$.bands = map_int(layers, "band")
      private$.categories = terra::cats(data)
      private$.layer_names = names(data)
      private$.crs = terra::crs(data)
    },

    #' @description
    #' Returns a slice of the raster in the specified format.
    #' Currently, the only supported formats is `"data.table"`.
    #'
    #' The rows must be addressed as vector of cells indices, columns must be
    #' referred to via layer names. Queries for rows with no matching row id and
    #' queries for columns with no matching column name are silently ignored.
    #'
    #' Rows are guaranteed to be returned in the same order as `rows`, columns
    #' may be returned in an arbitrary order. Duplicated row ids result in
    #' duplicated rows, duplicated column names lead to an exception.
    data = function(rows, cols) {
      stack = self$stack
      if (is.null(rows)) {
        rows = numeric(0)
      }
      rows = assert_integerish(rows, coerce = TRUE, null.ok = TRUE)
      assert_names(cols, type = "unique")
      cols = intersect(cols, private$.layer_names)

      if (!length(cols)) {
        data.table()
      } else if (
        length(rows) && test_integer(rows, sorted = TRUE, unique = TRUE, len = max(rows[length(rows)] - rows[1] + 1, 0))
      ) {
        # block read (e.g. c(1:10))
        terra::readStart(stack)
        on.exit(terra::readStop(stack))
        # determine rows to read
        cells = terra::rowColFromCell(stack, rows)
        row = cells[1, 1]
        nrows = cells[dim(cells)[1], 1] - cells[1, 1] + 1
        res = as.data.table(terra::readValues(stack, row = row, nrows = nrows, dataframe = TRUE))
        # subset cells and features
        res[cells[1, 2]:(cells[1, 2] + length(rows) - 1), cols, with = FALSE]
      } else {
        # cell read (e.g. c(1, 3, 5, 6, 10))
        as.data.table(terra::extract(stack, rows))[, cols, with = FALSE]
      }
    },

    #' @description
    #' Retrieve the first `n` rows.
    #'
    #' @param n (`integer(1)`)\cr
    #'   Number of rows.
    #'
    #' @return [data.table::data.table()] of the first `n` rows.
    head = function(n = 6L) {
      res = as.data.table(terra::head(self$stack, n))
      if (length(private$.response)) {
        set(res, j = private$.target_names, value = private$.response)
      }
      res
    },

    #' @description
    #' Returns a named list of vectors of distinct values for each column
    #' specified. If `na_rm` is `TRUE`, missing values are removed from the
    #' returned vectors of distinct values. Non-existing rows and columns are
    #' silently ignored.
    #'
    #' @param na_rm `logical(1)`\cr
    #'   Whether to remove NAs or not.
    #'
    #' @return Named `list()` of distinct values.
    distinct = function(rows, cols, na_rm = TRUE) {
      cols = intersect(cols, private$.layer_names)
      rows = rows %??% seq(self$nrow)

      res = if (!length(cols)) {
        named_list()
      } else if (test_integer(rows, sorted = TRUE, unique = TRUE, len = self$nrow)) {
        # fast
        stack = terra::subset(self$stack, cols)
        set_names(
          map(cols, function(layer) {
            if (terra::is.factor(stack[[layer]])) {
              terra::cats(stack[[layer]])[[1]][, 2]
            } else {
              terra::unique(stack[[layer]])[[1]]
            }
          }),
          cols
        )
      } else {
        # slow
        data = self$data(rows, cols)
        res = map(data, unique)
        map_if(res, is.factor, as.character)
      }

      if (na_rm) {
        res = map(res, function(values) values[!is.na(values)])
      }
      res
    },

    #' @description
    #' Returns the number of missing values per column in the specified slice
    #' of data. Non-existing rows and columns are silently ignored.
    #'
    #' @return Total of missing values per column (named `numeric()`).
    missings = function(rows, cols) {
      cols = intersect(cols, private$.layer_names)

      if (!length(cols)) {
        numeric(0)
      } else if (test_integer(rows, sorted = TRUE, unique = TRUE, len = self$nrow)) {
        # fast
        stack = self$stack
        res = terra::freq(stack, value = NA)
        res = set_names(res[, "count"], private$.layer_names)[cols]
        res
      } else {
        # slow
        data = self$data(rows, cols)
        map_int(data, count_missing)
      }
    },

    #' @description
    #' Returns the coordinates of `rows`.
    #' If `rows` is missing, all coordinates are returned.
    #'
    #' @return [data.table::data.table()] of coordinates of `rows`.
    coordinates = function(rows) {
      if (missing(rows)) {
        as.data.table(terra::crds(self$stack, df = TRUE))
      } else {
        as.data.table(terra::xyFromCell(rows))
      }
    }
  ),

  active = list(
    #' @field rownames (`integer()`)\cr
    #' Returns vector of all distinct row identifiers, i.e. the contents of the primary key column.
    rownames = function(rhs) {
      assert_ro_binding(rhs)
      1:terra::ncell(self$stack)
    },

    #' @field colnames (`character()`)\cr
    #' Returns vector of all column names.
    colnames = function(rhs) {
      assert_ro_binding(rhs)
      names(self$stack)
    },

    #' @field nrow (`integer(1)`)\cr
    #' Number of rows (observations).
    nrow = function(rhs) {
      assert_ro_binding(rhs)
      terra::ncell(self$stack)
    },

    #' @field ncol (`integer(1)`)\cr
    #' Number of columns (variables).
    ncol = function(rhs) {
      assert_ro_binding(rhs)
      terra::nlyr(self$stack)
    },

    #' @field stack (`SpatRaster`)\cr
    #' Raster stack.
    stack = function(rhs) {
      assert_ro_binding(rhs)
      # reconstruct the stack from the unique sources and select the recorded
      # bands so that layers reading from a shared multi-band file are not lost
      sources = unique(private$.sources)
      nlyrs = map_dbl(sources, function(s) terra::nlyr(terra::rast(s)))
      offset = c(0, cumsum(nlyrs))[match(private$.sources, sources)]
      stack = terra::rast(sources)[[offset + private$.bands]]
      iwalk(private$.categories, function(category, n) {
        if (!is.null(category)) {
          terra::set.cats(stack, layer = n, value = category)
        }
      })
      terra::set.names(stack, private$.layer_names)
      terra::crs(stack) = private$.crs
      stack
    }
  ),

  private = list(
    .calculate_hash = function() {
      mlr3misc::calculate_hash(self$stack)
    },
    .sources = NULL,
    .bands = NULL,
    .categories = NULL,
    .layer_names = NULL,
    .crs = NULL
  )
)

#' @title Coerce to spatial DataBackend
#'
#' @description
#' Wraps a [mlr3::DataBackend] around spatial objects.
#' Currently these S3 methods are only alternative ways for writing `DataBackendRaster$new()`.
#' They do not support coercing from other backends yet.
#'
#' @template param-data
#' @template param-primary-key
#' @param ... (`any`)\cr
#'   Not used.
#'
#' @return [mlr3::DataBackend].
#' @rdname as_data_backend
#'
#' @exportS3Method
#' @export as_data_backend.stars
#nolint next
as_data_backend.stars = function(data, primary_key = NULL, ...) {
  require_namespaces("stars")
  data = as(data, "SpatRaster")
  DataBackendRaster$new(data)
}

#' @export as_data_backend.SpatRaster
#' @exportS3Method
#' @rdname as_data_backend
#nolint next
as_data_backend.SpatRaster = function(data, primary_key = NULL, ...) {
  DataBackendRaster$new(data)
}

#' @export as_data_backend.RasterBrick
#' @exportS3Method
#' @rdname as_data_backend
#nolint next
as_data_backend.RasterBrick = function(data, primary_key = NULL, ...) {
  data = terra::rast(data)
  DataBackendRaster$new(data)
}

#' @export as_data_backend.RasterStack
#' @exportS3Method
#' @rdname as_data_backend
#nolint next
as_data_backend.RasterStack = function(data, primary_key = NULL, ...) {
  data = terra::rast(data)
  DataBackendRaster$new(data)
}
