####
# Utility functions for imager::cimg objects.
####

# global imports
import::here("dplyr", c("mutate", "sample_n"), .character_only = TRUE)
import::here(
  "imager",
  c(
    "as.cimg", "grayscale", "HSVtoRGB", "imappend", "imsplit",
    "map_il", "RGBtoHSV"
  ),
  .character_only = TRUE
)
import::here("magrittr", "%>%", .character_only = TRUE)
import::here("mclust", "densityMclust", .character_only = TRUE)
import::here("purrr", "modify_at", .character_only = TRUE)

# local imports
import::here("utils.R", "sample_histogram", .character_only = TRUE)


#' DataFrame conversion
#'
#' Convert image to dataframe and expose color channel.
#'
#' @param img an imager::cimg
#'
#' @returns a data.frame
img_to_df <- function(img) {
  out <- img %>%
    as.data.frame() %>%
    mutate(channel = factor(cc, labels = c("R", "G", "B")))
  return(out)
}


#' Histogram equalization
#'
#' Flatten histogram by replacing the pixel value of an image by their rank.
#'
#' @param img an imager::cimg
#'
#' @returns an imager::cimg
hist_eq <- function(img) {
  return(as.cimg(ecdf(img)(img), dim = dim(img)))
}


#' Enhance contrast
#'
#' Enhance the contrasts of an image by running an histogram equalization
#' separately on each channel and combining the results.
#'
#' @param img an imager::cimg
#'
#' @returns an imager::cimg
enhance_contrast <- function(img) {
  out <- img %>%
    imsplit("cc") %>%
    map_il(hist_eq) %>%
    imappend("cc")
  return(out)
}


#' Reduce saturation
#'
#' Reduce the saturation of an image through HSV conversion.
#'
#' @param img an imager::cimg
#' @param ratio an integer, how much to divide the saturation by.
#'
#' @returns an imager::cimg
desaturation <- function(img, ratio = 2L) {
  out <- img %>%
    RGBtoHSV() %>%
    imsplit("cc") %>%
    modify_at(2L, ~ . / ratio) %>%
    imappend("cc") %>%
    HSVtoRGB()
  return(out)
}


#' Correct illumination
#'
#' Correct a gray-scaled image illumination by fitting a linear model and
#' removing the spatial trend.
#'
#' @param img an imager::cimg
#' @param nsamples an integer, pixel subsampling value.
#'
#' @returns an imager::cimg object
correct_illumination <- function(img, nsamples = 1e4L) {
  # convert to grayscale if needed
  if (rev(dim(img))[1L] > 1L) {
    img <- grayscale(img)
  }
  # linear regression trend
  trend <- img %>%
    as.data.frame() %>%
    sample_n(nsamples) %>%
    lm(value ~ x * y, data = .) %>%
    predict(img)
  out <- img - trend
  return(out)
}


#' Invert grayscale image
#'
#' @param img an imager::cimg
#'
#' @returns an imager::cimg
invert_grayscale <- function(img) {
  # convert to grayscale if needed
  if (rev(dim(img))[1L] > 1L) {
    img <- grayscale(img)
  }
  out <- max(img) - img
  return(out)
}


#' Binarize
#'
#' Binarize a grayscale image.
#'
#' @param img an imager::cimg
#' @param quantile a real, the quantile level used for thresholding.
#' @param G an integer, the number of mixture components.
#' @param ... mclust::densityMclust parameters
#'
#' @returns an imager::cimg
binarize <- function(img, quantile = 0.95, G = 1L, ...) {
  # convert to grayscale if needed and invert
  if (rev(dim(img))[1L] > 1L) {
    stop("A grayscale image is expected.")
  }
  # sample 
  sample <- sample_histogram(img)
  # fit Gaussian mixture
  gm <- densityMclust(sample, G = G, plot = FALSE, ...)
  # threshold based on 95% quantile
  threshold <- qnorm(
    quantile, gm$parameters$mean[1L], sqrt(gm$parameters$variance$sigmasq[1L])
  )
  out <- img > threshold
  return(out)
}