project_root <- here::here()
renv::load(project_root)

# import libraries
library(magrittr)
library(optparse)

# source integration functions
utils_path <- file.path(project_root, "scripts", "utils")
suppressPackageStartupMessages({
  source(file.path(utils_path, "integrate-fastMNN.R"))
  source(file.path(utils_path, "integrate-harmony.R"))
  source(file.path(utils_path, "integrate-seurat.R"))
  source(file.path(utils_path, "integration-helpers.R"))
})

# Set up optparse options
option_list <- list(
  make_option(
    opt_str = c("--input_sce_file"),
    type = "character",
    default = NULL,
    help = "Path to RDS file that contains the merged SCE object to integrate"
  ),
  make_option(
    opt_str = c("--output_sce_file"),
    type = "character",
    default = NULL,
    help = "Path to RDS file where the integrated SCE object will be saved"
  ),
  make_option(
    opt_str = c("--method"),
    type = "character",
    default = NULL,
    help = "Integration method to use, either `fastMNN`, `harmony`, or `Seurat` (case-insensitive)."
  ),
  make_option(
    opt_str = c("-b", "--batch_column"),
    type = "character",
    default = "batch",
    help = "Name of the column in the SCE object that indicates batch groupings."
  ),
  make_option(
    opt_str = c("--seed"),
    type = "integer",
    default = NULL,
    help = "random seed to set during integration"
  ),
  make_option(
    opt_str = c("--harmony_covariate_cols"),
    type = "character",
    default = NULL,
    help = "Optional comma-separated list of columns (e.g. patient, sex) to consider as covariates
            during integration with `harmony`."
  ),
  make_option(
    opt_str = c("--harmony_from_expression"),
    action = "store_false",
    default = TRUE,
    help = "Indicate whether to use the gene expression matrix, rather than PCs, during `harmony` integration.
    To use expression instead of PCs (default), use `--harmony_from_expression`"
  ),
  make_option(
    opt_str = c("--fastmnn_no_cosine"),
    action = "store_false",
    default = TRUE,
    help = "Indicate whether to turn off cosine normalization during `fastMNN` integration.
    To turn off cosine normalization, use `--fastmnn_no_cosine`"
  ),
  make_option(
    opt_str = c("--fastmnn_use_all_genes"),
    action = "store_true",
    default = FALSE,
    help = "Indicate whether to use all genes instead of only HVGs during `fastMNN` integration.
    To use all genes, use `--fastmnn_use_all_genes`"
  ),
  make_option(
    opt_str = c("--fastmnn_auto_merge"),
    action = "store_true",
    default = FALSE,
    help = "Indicates whether or not to use the auto.merge option for `fastMNN` integration.
    To perform auto.merge, use `--fastmnn_auto_merge`."
  ),
  make_option(
    opt_str = c("--seurat_reduction_method"),
    type = "character",
    default = NULL,
    help = "Reduction method to use during Seurat integration."
  ),
  make_option(
    opt_str = c("--seurat_num_genes"),
    type = "numeric",
    default = 2000,
    help = "Number of variable genes for Seurat to identify and use."
  ),
  make_option(
    opt_str = c("--seurat_integration_dims"),
    type = "numeric",
    default = 30,
    help = "Number of dimensions for Seurat to use during integration."
  ),
  make_option(
    opt_str = c("--seurat_umap_dims"),
    type = "numeric",
    default = 30,
    help = "Number of dimensions for Seurat to use during UMAP calculation."
  ),
  make_option(
    opt_str = c("--seurat_anchor_threshold"),
    type = "numeric",
    default = 100,
    help = "Minimum threshold for number of neighbors to consider when weighting anchors during integration."
  ),
  make_option(
    opt_str = c("--corrected_only"),
    action = "store_true",
    default = FALSE,
    help = "Indicate whether to only return the corrected gene expression data, and
    not uncorrected expression, in the integrated SCE object. To return only
    corrected expression, use `--corrected_only`."
  )
)


# --corrected_only: Flag to specify that only corrected gene expression values should
#   be returned in the integrated SCE object. Default usage of this script will
#   return all data.

# Setup ------------------------------------------------------------------------
# Parse options
opt <- parse_args(OptionParser(option_list = option_list))



# Check and assign provided method based on available methods
available_methods <- c("fastmnn", "harmony", "seurat")

# helper function for method check fails
stop_method <- function() {
  stop(
    paste("You must specify one of the following (case-insensitive) to --method:",
          paste(available_methods, collapse = ", ")
    )
  )
}

if (is.null(opt$method)) {
  stop_method()
} else {
  integration_method <- tolower(opt$method)
  if (!(integration_method %in% available_methods)) {
    stop_method()
  }
}


# Check that provided input file exists and is an RDS file
if(is.null(opt$input_sce_file)) {
  stop("You must provide the path to the RDS file with merged SCEs to --input_sce_file")
} else {
  if(!file.exists(opt$input_sce_file)) {
    stop("Provided --input_sce_file file does not exist. Make sure that you have
          run `02-prepare-merged-sce.R` or the provided path is correct.")
  }
}

# Check that both input and output files have RDS extensions
if(!(grepl("\\.rds$", opt$input_sce_file, ignore.case = TRUE)) ||
   !(grepl("\\.rds$", opt$output_sce_file, ignore.case = TRUE))) {
  stop("The provided --input_sce_file and --output_sce_file files must be RDS files.")
}


# Check that directory for output file exists and the specified file is an RDS file
integrated_sce_dir <- dirname(opt$output_sce_file)
create_dir(integrated_sce_dir)


# Read in SCE file -----------------------------
merged_sce_obj <- readr::read_rds(opt$input_sce_file)


# Perform integration with fastMNN, if specified -------------------------
if (integration_method == "fastmnn") {

  # Prepare `gene_list` argument
  if (opt$fastmnn_use_all_genes) {
    fastmnn_gene_list <- NULL
  } else {
    fastmnn_gene_list <- metadata(merged_sce_obj)$variable_genes
  }

  # Perform integration
  integrated_sce_obj <- integrate_fastMNN(
    merged_sce_obj,
    batch_column = opt$batch_column,
    cosine_norm  = opt$fastmnn_no_cosine,
    gene_list    = fastmnn_gene_list,
    seed         = opt$seed,
    auto.merge   = opt$fastmnn_auto_merge
  )

  # add note about auto merge to metadata of integrated object
  if(opt$fastmnn_auto_merge){
    auto_merge <- "auto"
  } else {
    auto_merge <- "input"
  }

  metadata(integrated_sce_obj)$fastmnn_auto_merge <- auto_merge

}

# Perform integration with harmony, if specified -------------------------
if (integration_method == "harmony") {

  # Set up `covariate_cols` argument based on user options
  if (is.null(opt$harmony_covariate_cols)) {
    harmony_covariate_cols <- c()
  } else {
    harmony_covariate_cols <- unlist(stringr::str_split(opt$harmony_covariate_cols, "\\s*,\\s*"))
  }

  # Perform integration
  integrated_sce_obj <- integrate_harmony(
    merged_sce_obj,
    batch_column = opt$batch_column,
    covariate_cols = opt$harmony_covariate_cols,
    from_pca = opt$harmony_from_expression,
    seed = opt$seed
  )
}

# Perform integration with seurat, if specified -------------------------
if (integration_method == "seurat") {

  # Convert the merged sce object into a list of seurat objects
  seurat_list <- prepare_seurat_list(merged_sce_obj)

  # Perform integration
  integrated_seurat_obj <- integrate_seurat(
    seurat_list,
    opt$seurat_reduction_method, # integrate_seurat() will check this argument
    batch_column = opt$batch_column,
    num_genes = opt$seurat_num_genes,
    integration_dims = 1:opt$seurat_integration_dims,
    umap_dims = 1:opt$seurat_umap_dims,
    anchor_threshold = opt$seurat_anchor_threshold
  )

  # Converted the integrated seurat object into an SCE object
  integrated_sce_obj <- as.SingleCellExperiment(integrated_seurat_obj)

  # Rename assays appropriately:

  # The converted `logcounts` is the CORRECTED expression
  assay(integrated_sce_obj, paste0(opt$seurat_reduction_method, "_corrected")) <- logcounts(integrated_sce_obj)

  # The logcounts in SCT assay is logcounts
  logcounts(integrated_sce_obj) <- logcounts( altExp(integrated_sce_obj, "SCT") )[rownames(integrated_sce_obj),]

  # The counts in RNA assay is counts  
  counts(integrated_sce_obj) <- counts( altExp(integrated_sce_obj, "RNA") )[rownames(integrated_sce_obj),]

  # Remove altExps
  integrated_sce_obj <- removeAltExps(integrated_sce_obj)

  # Convert reducedDims names back to <lowercase>_<UPPERCASE> because 
  #   `as.SingleCellExperiment` makes them all uppercase
  reducedDimNames(integrated_sce_obj) <- stringr::str_replace_all(
    reducedDimNames(integrated_sce_obj),
    toupper(opt$seurat_reduction_method),
    opt$seurat_reduction_method
  )

}


# Remove uncorrected expression values, if specified ----------
if (opt$corrected_only) {
  integrated_sce_obj <- remove_uncorrected_expression(integrated_sce_obj)
}


# Write integrated SCE object to RDS -------------------
readr::write_rds(integrated_sce_obj, opt$output_sce_file)

import os
import anndata as adata
import argparse
import re
from utils.integrate_scvi import integrate_scvi

# define arguments
parser = argparse.ArgumentParser()
parser.add_argument('-i', '--input_anndata',
                    dest = 'input_anndata',
                    required = True,
                    help = 'Path to HDF5 file with merged AnnData object to integrate')
parser.add_argument('-o', '--output_anndata',
                    dest = 'output_anndata',
                    required = True,
                    help = 'Path to HDF5 file to save the integrated AnnData object')
parser.add_argument('-b', '--batch_column',
                    dest = 'batch_column',
                    default = 'batch',
                    help = 'The name of the column in `anndata.obs` that indicates the batches for each cell, '
                            ' typically this corresponds to the library id.')
parser.add_argument('--categorical_covariates',
                    dest = 'categorical_covariates',
                    default = None,
                    type = str,
                    help = 'A comma-separated list of columns containing additional categorical data to be'
                    ' included as a covariate. Default is None.')
parser.add_argument('--continuous_covariates',
                    dest = 'continuous_covariates',
                    default = "subsets_mito_percent",
                    type = str,
                    help = 'A comma-separated list of columns containing additional continous data to be'
                    ' included as a covariate. Default is "subsets_mito_percent".')
parser.add_argument('--num_latent',
                    dest = 'num_latent',
                    type = int,
                    default = 20,
                    help = 'Number of dimensions to return from the latent space.')
parser.add_argument('--use_hvg',
                    dest = 'use_hvg',
                    action = 'store_true',
                    default = False,
                    help = 'Boolean indicating whether or not to use only highly variable genes for data integration.'
                    'If --use_hvg is used, the returned integrated object will only contain the highly variable genes.')
parser.add_argument('--corrected_only',
                    dest = 'corrected_only',
                    action = 'store_true',
                    help = 'Boolean indicating to return only the corrected data or all raw data.'
                    ' Default will return all data. To return only corrected data, use --corrected_only.')
parser.add_argument('-s', '--seed',
                    dest = 'seed',
                    type=int,
                    default = None,
                    help = 'Random seed to set for scanorama.')

args = parser.parse_args()

# compile extension regex
file_ext = re.compile(r"\.hdf5$|.h5$", re.IGNORECASE)

# check that input file exists, if it does exist, make sure it's an h5 file
if not os.path.exists(args.input_anndata):
    raise FileExistsError("--input_anndata file not found.")
elif not file_ext.search(args.input_anndata):
    raise ValueError("--input_anndata must end in either .hdf5 or .h5 and contain a merged AnnData object.")

# make sure output file path is h5 file
if not file_ext.search(args.output_anndata):
    raise ValueError("--output_anndata must provide a file path ending in either .hdf5 or .h5.")

# check that output file directory exists and create directory if doesn't exist
integrated_adata_dir = os.path.dirname(args.output_anndata)
os.makedirs(integrated_adata_dir, exist_ok = True)

# read in merged anndata object
merged_adata = adata.read_h5ad(args.input_anndata)

if args.use_hvg:
    try:
        var_genes = list(merged_adata.uns['variable_genes'])
    except KeyError:
        print("Variable genes cannot be found in anndata object."
              " If using --use_hvg, make sure HVG are stored in adata.uns['variable_genes'].")
        raise
else:
    var_genes = list(merged_adata.var_names)

# split covariates from comma separated strings into lists
if args.categorical_covariates:
    categorical_covariates = [covariate.strip() for covariate in args.categorical_covariates.split(',')]
else:
    categorical_covariates = None

if args.continuous_covariates:
    if args.continuous_covariates == "None":
        continuous_covariates = None
    else:
        continuous_covariates = [covariate.strip() for covariate in args.continuous_covariates.split(',')]
else:
    continuous_covariates = None

# integrate anndata with scvi
scvi_integrated_adata = integrate_scvi(merged_adata,
                                       integrate_genes = var_genes,
                                       batch_column = args.batch_column,
                                       categorical_covariate_columns = categorical_covariates,
                                       continuous_covariate_columns = continuous_covariates,
                                       num_latent = args.num_latent,
                                       seed = args.seed)

# remove raw data and logcounts to minimize space if corrected_only is true
if args.corrected_only:
    print("Removing raw data and log-normalized data. Only corrected data will be returned.")
    scvi_integrated_adata.X = None
    del scvi_integrated_adata.layers["logcounts"]


# write anndata to h5
scvi_integrated_adata.write(filename = args.output_anndata)

project_root <- here::here()
renv::load(project_root)

# import libraries
suppressPackageStartupMessages({
  library(magrittr)
  library(optparse)
  library(SingleCellExperiment) 
})
source(file.path(project_root, "scripts", "utils", "integration-helpers.R"))

# Load helper functions
source(
  file.path(
    "scripts", 
    "utils", 
    "integration-helpers.R"
  )
)

# Set up optparse options
option_list <- list(
  make_option(
    opt_str = c("--input_anndata_file"),
    type = "character",
    default = NULL,
    help = "Path to H5 file that contains the integrated AnnData object"
  ),
  make_option(
    opt_str = c("--output_sce_file"),
    type = "character",
    default = NULL,
    help = "Path to RDS file where the integrated SCE object will be saved"
  ),
  make_option(
    opt_str = c("--method"),
    type = "character",
    default = NULL,
    help = "Method that was used for integration, either `scanorama` or `scVI` (case-insensitive)."
  ),
  make_option(
    opt_str = c("--seed"),
    type = "integer",
    default = NULL,
    help = "random seed to set"
  ),
  make_option(
    opt_str = c("--corrected_only"),
    action = "store_true",
    default = FALSE,
    help = "Indicate whether the integrated AnnData object contains only 
    corrected gene expression data, and is missing uncorrected expression (no `X` matrix). 
    To indicate the object only contains corrected data, use `--corrected_only`."
  )
)

# Setup ------------------------------------------------------------------------
# Parse options
opt <- parse_args(OptionParser(option_list = option_list))

set.seed(opt$seed)

# Check provided method based on available methods
available_methods <- c("scanorama", "scvi")

# helper function for method check fails
stop_method <- function() {
  stop(
    paste("You must specify one of the following (case-insensitive) to --method:",
          paste(available_methods, collapse = ", ")
    )
  )
}

if (is.null(opt$method)) {
  stop_method()
} else {
  integration_method <- tolower(opt$method)
  if (!(integration_method %in% available_methods)) {
    stop_method()
  }
}

# Check that provided input file exists and is an RDS file
if(is.null(opt$input_anndata_file)) {
  stop("You must provide the path to the H5 file with merged SCEs to --input_anndata_file")
} else {
  if(!file.exists(opt$input_anndata_file)) {
    stop("Provided --input_anndata_file file does not exist.")
  }
}

# Check that both input and output files have correct extensions
if(!(stringr::str_ends(opt$input_anndata_file, ".hdf5|.h5"))){
  stop("`--input_anndata_file` must end in either '.hdf5' or '.h5'")
}
if(!(stringr::str_ends(opt$output_sce_file, ".rds"))){
  stop("`--output_sce_file` must end in '.rds'")
}

# Check that directory for output file exists
integrated_sce_dir <- dirname(opt$output_sce_file)
create_dir(integrated_sce_dir)

# Anndata to SCE post-processing -----------------------------------------------

# read in AnnData object as H5
integrated_sce <- zellkonverter::readH5AD(opt$input_anndata_file)

# Remove uncorrected expression values if corrected_only is set
# We do this after zellkonverter which always retain counts, even if empty.
if (opt$corrected_only) {
  integrated_sce <- remove_uncorrected_expression(integrated_sce)
}


# replace . added in anndata names with - found in other sce objects
names(metadata(integrated_sce)) <- stringr::str_replace_all(names(metadata(integrated_sce)), "\\.", "-")
names(colData(integrated_sce)) <- stringr::str_replace_all(names(colData(integrated_sce)), "\\.", "-")
names(rowData(integrated_sce)) <- stringr::str_replace_all(names(rowData(integrated_sce)), "\\.", "-")

# check for corrected data 
corrected_assay_name <- paste(integration_method, "corrected", sep = "_")
if(!corrected_assay_name %in% assayNames(integrated_sce)){
  stop("integrated SCE missing corrected gene expression data in assays.")
}

# function to check for SVD or latent embeddings
check_pseudo_pca <- function(integrated_sce, 
                             pseudo_pca_name){
  # check corrected data and SVD were converted
  if(!pseudo_pca_name %in% reducedDimNames(integrated_sce)){
    stop(glue::glue("integrated SCE missing {pseudo_pca_name} in reduced dimensions slot."))
  }
}

# function to run UMAP with respective pseudo PCA name 
add_umap <- function(integrated_sce,
                     pseudo_pca_name,
                     umap_name){

  integrated_sce <- integrated_sce %>%
    scater::runUMAP(dimred = pseudo_pca_name,
                    name = umap_name)
}

# calculate UMAP dependent on integration method 
if(integration_method == "scanorama"){

  pseudo_pca_name <- "scanorama_SVD"
  check_pseudo_pca(integrated_sce,
                   pseudo_pca_name)

  umap_name <- "scanorama_UMAP"

  # run UMAP using the SVD as input 
  integrated_sce <- add_umap(integrated_sce,
                             pseudo_pca_name,
                             umap_name)

}

if(integration_method == "scvi"){

  pseudo_pca_name <- "scvi_latent"
  check_pseudo_pca(integrated_sce,
                   pseudo_pca_name)

  umap_name <- "scvi_UMAP"

  # run UMAP using the latent embeddings as input 
  integrated_sce <- add_umap(integrated_sce,
                             pseudo_pca_name,
                             umap_name)

}

# Write RDS --------------------------------------------------------------------

readr::write_rds(integrated_sce, opt$output_sce_file)

shell:
    """
    Rscript -e "renv::restore(lockfile = '{input}')" &> {log}
    date -u -Iseconds  > {output}
    """

shell:
    """
    Rscript scripts/02-prepare-merged-sce.R \
      --library_file "{input.processed_tsv}" \
      --add_celltype {config[add_celltype]} \
      --celltype_info "{config[celltype_file]}" \
      --grouping_var {config[grouping_var]} \
      --groups_to_integrate "{config[groups_to_integrate]}" \
      --merged_sce_dir "{output}" \
      --num_hvg {config[num_hvg]} \
      --subset_hvg \
      --seed {config[seed]} \
      {params.random_merge_flag} \
      &> {log}
    """

shell:
    """
    Rscript scripts/02a-convert-sce-to-anndata.R \
      --library_file "{input.processed_tsv}" \
      --merged_sce_dir "{input.merged_sce_dir}" \
      --grouping_var {config[grouping_var]} \
      --anndata_output_dir "{output}" \
      &> {log}
    """

shell:
    """
    Rscript scripts/03a-integrate-sce.R \
      --input_sce_file "{input.merged_sce_dir}/{params.merged_sce_file}" \
      --output_sce_file "{output}" \
      --method fastMNN \
      --seed {config[seed]} \
      --corrected_only \
      {params.auto_merge_flag} \
      &> {log}
    """

shell:
    """
    Rscript scripts/03a-integrate-sce.R \
      --input_sce_file "{input.merged_sce_dir}/{params.merged_sce_file}" \
      --output_sce_file "{output}" \
      --method harmony \
      --seed {config[seed]} \
      --corrected_only \
      &> {log}
    """

shell:
    """
    Rscript scripts/03a-integrate-sce.R \
      --input_sce_file "{input.merged_sce_dir}/{params.merged_sce_file}" \
      --output_sce_file "{output}" \
      --method seurat \
      --seurat_reduction_method {wildcards.method} \
      --num_genes {params.num_genes} \
      --corrected_only \
      &> {log}
    """

shell:
    """
    python scripts/03b-integrate-scanorama.py \
      --input_anndata "{input.merged_anndata_dir}/{params.merged_anndata_file}" \
      --output_anndata "{output}" \
      --seed {config[seed]} \
      --use_hvg \
      --corrected_only \
      &> {log}
    """

shell:
    """
    python scripts/03c-integrate-scvi.py \
      --input_anndata "{input.merged_anndata_dir}/{params.merged_anndata_file}" \
      --output_anndata "{output}" \
      --continuous_covariates {config[continuous_covariates]} \
      --num_latent {config[num_latent]} \
      --seed {config[seed]} \
      --use_hvg \
      --corrected_only \
      &> {log}
    """

shell:
    """
    Rscript scripts/04-post-process-anndata.R \
      --input_anndata_file "{input}" \
      --output_sce_file "{output}" \
      --method "{wildcards.method}" \
      --seed {config[seed]} \
      --corrected_only \
      &> {log}
    """

shell:
    """
    Rscript -e \
    "rmarkdown::render('analysis_templates/01-single-group-integration-check-template.Rmd', \
                        clean = TRUE, \
                        output_file = '{output}', \
                        output_dir = dirname('{output}'), \
                        params = list(group_name = '{wildcards.project}', \
                                      merged_sce_dir = '{workflow.basedir}/{input.merged_sce_dir}', \
                                      integrated_sce_dir = '{workflow.basedir}/{params.integrated_sce_dir}', \
                                      integration_methods = '{config[integration_methods]}', \
                                      library_metadata = '{workflow.basedir}/{config[processed_tsv]}' , \
                                      max_celltypes = {config[max_celltypes]}))" \
    &> {log}
    """