#' Fitting "topic models" with PCA+varimax
#'
"_PACKAGE"

#' Fit a varimax-rotated PCA using irlba
#'
#' Extract `n` principal components from the matrix `mx` using `irlba`, then rotate the solution using `varimax`
#' @param mx Matrix of interest
#' @param n Number of principal components / varimax factors to return; can take a vector of values
#' @param prcomp_fn Function to use to extract principal components
#' @param prcomp_opts List of options to pass to `prcomp_fn`
#' @param varimax_fn Function to use for varimax rotation
#' @param varimax_opts List of options to pass to `varimax_fn`
#' @param retx Whether to return the input matrix `mx`
#' @return A list of class `varimaxes`, with elements
#'   - `totalvar`: Total variance, from PCA
#'   - `sdev`:  Standard deviations of the extracted principal components
#'   - `x`: If `retx` is `TRUE`, the input matrix `mx`
#'   - `rotation`:  Rotation matrix (variable loadings) from PCA
#'   - `varimaxes`: A list of class `varimaxes`, containing one fitted varimax model for each value of `n`, with further elements
#'       - `loadings`: Varimax-rotated standardized loadings
#'       - `rotmat`:  Varimax rotation matrix
#'       - `scores`:  Varimax-rotated observation scores
#' @examples
#' \donttest{
#' set.seed(42)
#' theta  = rdirichlet(50, 1, k = 3)
#' phi    = rdirichlet(3, 0.1, k = 20)
#' corpus = draw_corpus(rep(50L, 50), theta, phi)
#' dtm    = tidytext::cast_sparse(corpus, doc, word, n)
#' varimax_irlba(dtm, n = 3)
#' }
#' @export
varimax_irlba = function(
      mx,
      n,
      prcomp_fn = irlba::prcomp_irlba,
      prcomp_opts = NULL,
      varimax_fn = stats::varimax,
      varimax_opts = NULL,
      retx = TRUE
) {
      ## prcomp_irlba loses names
      rows = rownames(mx)
      cols = colnames(mx)
      ## PCA fit
      pca_fit = do.call(prcomp_fn, c(list(x = mx, n = max(n)), prcomp_opts))

      ## Varimax fit
      varimaxes = n |>
            purrr::set_names() |>
            purrr::map(
                  fit_varimax,
                  pca_fit,
                  cols,
                  rows,
                  varimax_fn,
                  varimax_opts
            )

      toreturn = list(
            totalvar = pca_fit$totalvar,
            sdev = pca_fit$sdev,
            rows = rows,
            cols = cols,
            center = pca_fit$center,
            scale = pca_fit$scale,
            x = mx,
            rotation = pca_fit$rotation,
            n = n,
            varimax = varimaxes
      )
      if (!retx) {
            toreturn$x = NULL
      }
      class(toreturn) = c('varimaxes', class(toreturn))
      return(toreturn)
}

#' Given a (rank `n`) PCA fit, return a rank `k < n` varimax fit
#'
#' @param k Desired rank of the fitted varimax model
#' @param pca Fitted PCA model or `varimaxes` object
#' @param feature_names Names of the features (eg, data columns)
#' @param obs_names Names of the observations (eg, data rows)
#' @param varimax_fn Function to use for varimax rotation
#' @param varimax_opts Options passed to `varimax_fn`
#' @param positive_skew Should negative-skewed factors be flipped to have positive skew?
#' @param x PCA scores matrix (n_obs x max_k), as returned by `predict(fitted, newdata)`. Used when the fitted `pca` object does not contain scores.
#' @return List with components
#'     - `loadings`: Rotated feature loadings
#'     - `rotmat`:  Rotation matrix
#'     - `scores`:  Rotated observation scores
#' @details After the initial rotation, factors with negative skew (left tails) are flipped
#'   `pca` must contain `$rotation` (feature loadings matrix) and `$sdev` (standard deviations
#'   per PC); `$x` (PC scores matrix) is also required unless `x` is supplied directly.
#' @examples
#' \donttest{
#' set.seed(42)
#' theta  = rdirichlet(50, 1, k = 3)
#' phi    = rdirichlet(3, 0.1, k = 20)
#' corpus = draw_corpus(rep(50L, 50), theta, phi)
#' dtm    = tidytext::cast_sparse(corpus, doc, word, n)
#' pca    = irlba::prcomp_irlba(dtm, n = 5)
#' fit_varimax(k = 3, pca = pca,
#'             feature_names = colnames(dtm),
#'             obs_names = rownames(dtm))
#' }
#' @export
fit_varimax = function(
      k,
      pca,
      feature_names,
      obs_names,
      varimax_fn = stats::varimax,
      varimax_opts = NULL,
      positive_skew = TRUE,
      x = NULL
) {
      assertthat::assert_that(
            !is.null(pca$rotation),
            msg = '`pca` must have a `$rotation` field (feature loadings matrix)'
      )
      assertthat::assert_that(
            !is.null(pca$sdev),
            msg = '`pca` must have a `$sdev` field (standard deviations per PC)'
      )
      assertthat::assert_that(
            !is.null(pca$x) || !is.null(x),
            msg = 'Data matrix must be passed through either `pca` or `x`'
      )
      raw_loadings = pca$rotation[, 1:k] %*%
            diag(pca$sdev, k, k) |>
            magrittr::set_rownames(feature_names)
      varimax_fit_prelim = do.call(
            varimax_fn,
            c(list(x = raw_loadings), varimax_opts)
      )

      if (positive_skew) {
            ## Reverse factors with negative skew (left tails)
            varimax_fit = purrr::map(
                  varimax_fit_prelim,
                  ~ .x %*%
                        diag(
                              1 -
                                    2 *
                                          (psych::skew(
                                                varimax_fit_prelim$loadings
                                          ) <
                                                0)
                        )
            )
            assertthat::assert_that(
                  all(psych::skew(varimax_fit$loadings) > 0),
                  msg = 'Varimax loadings do not all have positive skew'
            )
      } else {
            varimax_fit = varimax_fit_prelim
      }

      ## Scores
      if (is.null(x)) {
            scores = scale(pca$x[, 1:k]) %*%
                  varimax_fit$rotmat |>
                  magrittr::set_rownames(obs_names)
      } else {
            scores = scale(x[, 1:k]) %*%
                  varimax_fit$rotmat |>
                  magrittr::set_rownames(obs_names)
      }

      toreturn = list(
            loadings = varimax_fit$loadings,
            rotmat = varimax_fit$rotmat,
            scores = scores
      )
      return(toreturn)
}

#' Fit a topic model using PCA+varimax
#'
#' @param dtm Document-term matrix.  Either an object inheriting from `Matrix` or a long dataframe representation with row column `row`, column column `column`, and value column `n`.
#' @param n Number of topics to return
#' @param row In dataframe `dtm`, row column
#' @param column In dataframe `dtm`, column column
#' @param value In dataframe `dtm`, value column
#' @param verbose Should `irlba()` be `verbose`?
#' @param ... Other arguments, passed to `varimax_irlba`
#' @return As per `varimax_irlba`, of class `tmfast`
#' @details If `dtm` is not a matrix, will be cast to a sparse matrix using `tidytext::case_sparse()`
#' @export
tmfast = function(
      dtm,
      n,
      row = 'doc',
      column = 'word',
      value = 'n',
      verbose = FALSE,
      ...
) {
      if (!inherits(dtm, 'Matrix')) {
            if (verbose) {
                  message('Casting dtm to sparse matrix')
            }
            dtm = tidytext::cast_sparse(
                  dtm,
                  {{ row }},
                  {{ column }},
                  {{ value }}
            )
      }
      fitted = varimax_irlba(
            dtm,
            n,
            prcomp_opts = list(.scale = FALSE, verbose = verbose),
            ...
      )
      class(fitted) = c('tmfast', class(fitted))
      return(fitted)
}

#' Insert a topic model into a fitted `tmfast`
#'
#' Apply varimax rotation for a value of k less than the maximum already included in the tmfast.
#' @param fitted Fitted `tmfast` object
#' @param k Desired number of topics for new model
#' @param x Data matrix (document-term matrix), as Matrix object (eg, using `build_matrix()`)
#' @return `tmfast` object, as `fitted`, with additional topic model inserted
#' @examples
#' \donttest{
#' set.seed(42)
#' theta  = rdirichlet(50, 1, k = 4)
#' phi    = rdirichlet(4, 0.1, k = 20)
#' corpus = draw_corpus(rep(50L, 50), theta, phi)
#' model  = tmfast(corpus, n = c(3, 4))
#' insert_topics(model, k = 2)
#' }
#' @export
insert_topics = function(fitted, k, x = NULL) {
      if (k %in% fitted$n) {
            return(fitted)
      }
      if (k > max(fitted$n)) {
            stop(glue::glue(
                  'Can only insert topic models with at most {max(fitted$n)} topics'
            ))
      }

      if (length(k) > 1) {
            stop('Currently can only insert 1 topic model at a time')
      }

      if (is.null(fitted$x) && is.null(x)) {
            stop('Data matrix missing; pass using x')
      }

      raw_data = if (is.null(x)) fitted$x else x
      pc_scores = predict(fitted, raw_data)
      new_varimax = fit_varimax(
            k,
            fitted,
            feature_names = fitted$cols,
            obs_names = fitted$rows,
            x = pc_scores
      )
      fitted$n = c(fitted$n, k)
      fitted$varimax[[as.character(k)]] = new_varimax
      return(fitted)
}
