#!/usr/bin/env Rscript --vanilla

# Script for importing and processing transcript-level quantifications.
# Written by Lorena Pantano, later modified by Jonathan Manning, and released
# under the MIT license.

# tximport::summarizeToGene() reorders gene-level output rows via base R's
# rowsum(reorder=TRUE), which sorts by the process locale's collation order.
# Fix LC_COLLATE so row order (and therefore output content) is reproducible
# across execution environments regardless of their default locale.
Sys.setlocale("LC_COLLATE", "C")

# Loading required libraries
library(SummarizedExperiment)
library(tximport)

################################################
################################################
## Functions                                  ##
################################################
################################################

#' Parse out options from a string without recourse to optparse
#'
#' @param x Long-form argument list like --opt1 val1 --opt2 val2
#'
#' @return named list of options and values similar to optparse

parse_args <- function(x){
    args_list <- unlist(strsplit(x, ' ?--')[[1]])[-1]
    args_vals <- lapply(args_list, function(y) scan(text=y, what='character', quiet = TRUE))

    # Ensure the option vectors are length 2 (key/value) to catch empty ones
    args_vals <- lapply(args_vals, function(z){ length(z) <- 2; z})

    parsed_args <- structure(lapply(args_vals, function(y) y[2]), names = lapply(args_vals, function(y) y[1]))
    parsed_args[! is.na(parsed_args)]
}

#' Build a table from a SummarizedExperiment object
#'
#' This function takes a SummarizedExperiment object and a specific slot name to extract
#' assay data. It then combines the first two columns of the rowData with the specified
#' assay data slot into a new data table.
#'
#' @param se.obj A SummarizedExperiment object from which to build the table.
#' @param slot The name of the slot in the assays list from which to extract data.
#'
#' @return A data frame combining the first two columns of the rowData with the assay data from the specified slot.

build_table <- function(se.obj, slot) {
    data.frame(cbind(rowData(se.obj)[,1:2], assays(se.obj)[[slot]]), check.names = FALSE)
}

#' Write a table to a file from a SummarizedExperiment object with given parameters
#'
#' This function generates a table from a SummarizedExperiment object using specified parameters
#' and writes the resulting table to a file. The file name is constructed using a prefix and a
#' suffix from the parameters, and the table is written with tab separation, without quoting text,
#' and without row names.
#'
#' @param params A list containing the parameters needed for file generation and table writing.
#' The list should include:
#' - `obj`: A SummarizedExperiment object from which to build the table.
#' - `slot`: The name of the slot in the assays list from which to extract data.
#' - `suffix`: Suffix to use for generating the file name.
#'
#' @return NULL The function is called for its side effect of writing a file and does not return anything.

write_se_table <- function(params, prefix) {
    file_name <- paste0(prefix, ".", params\$suffix)
    write.table(build_table(params\$obj, params\$slot), file_name,
                sep="\t", quote=FALSE, row.names = FALSE)
}

#' Read Transcript Metadata from a Given Path
#'
#' This function reads transcript metadata from a specified file path. The file is expected to
#' be a tab-separated values file with a header, containing transcript information. Columns are
#' selected by name using the tx_col, gene_id_col, and gene_name_col parameters. The function
#' checks if the file is empty and stops execution with an error message if so. Additional
#' processing is done to ensure compatibility with a predefined data structure (e.g., `txi[[1]]`),
#' including adding missing entries and reordering based on the transcript IDs found in `txi[[1]]`.
#'
#' @param tinfo_path The file path to the transcript information file.
#' @param tx_col Column name for transcript IDs.
#' @param gene_id_col Column name for gene IDs.
#' @param gene_name_col Column name for gene names.
#'
#' @return A list containing four elements:
#' - `transcript`: A data frame with transcript IDs, gene IDs, and gene names, indexed by transcript IDs.
#' - `gene`: A data frame with unique gene IDs and gene names.
#' - `tx2gene`: A data frame mapping transcript IDs to gene IDs.
#' - `tx2gene_augmented`: The full tx2gene table actually used by tximport (input
#'   mappings plus self-mappings appended for any transcripts present in the
#'   quantification output but missing from the input tx2gene), with the input
#'   column headers preserved.
#' - `extra`: A character vector of transcript IDs found in quantification output but missing from the tx2gene file.

read_transcript_info <- function(tinfo_path, tx_col, gene_id_col, gene_name_col){
    info <- file.info(tinfo_path)
    if (info\$size == 0) {
        stop("tx2gene file is empty")
    }

    # Read file with actual header to handle variable column counts correctly
    raw_info <- read.csv(tinfo_path, sep="\t", header = TRUE, check.names = FALSE)

    # Resolve column references (use named columns if present, otherwise positional)
    resolved_tx_col        <- if (tx_col        %in% colnames(raw_info)) tx_col        else colnames(raw_info)[1]
    resolved_gene_id_col   <- if (gene_id_col   %in% colnames(raw_info)) gene_id_col   else colnames(raw_info)[2]
    resolved_gene_name_col <- if (gene_name_col %in% colnames(raw_info)) gene_name_col else colnames(raw_info)[3]

    transcript_info <- data.frame(
        tx        = raw_info[[resolved_tx_col]],
        gene_id   = raw_info[[resolved_gene_id_col]],
        gene_name = raw_info[[resolved_gene_name_col]],
        check.names = FALSE
    )

    extra <- setdiff(rownames(txi[[1]]), as.character(transcript_info[["tx"]]))
    if (length(extra) > 0) {
        warning(
            length(extra), " transcripts found in quantification output but missing from ",
            "the tx2gene mapping (GTF). These will be included in transcript-level outputs ",
            "but excluded from gene-level summaries. This usually means the transcript FASTA ",
            "and GTF are from different sources or versions. First 5: ",
            paste(head(extra, 5), collapse = ", ")
        )
    }
    transcript_info <- rbind(transcript_info, data.frame(tx=extra, gene_id=extra, gene_name=extra, check.names = FALSE))
    transcript_info <- transcript_info[match(rownames(txi[[1]]), transcript_info[["tx"]]), ]
    rownames(transcript_info) <- transcript_info[["tx"]]

    # Restore input column headers so the augmented file can be fed back to tximport
    tx2gene_augmented <- transcript_info
    colnames(tx2gene_augmented) <- c(resolved_tx_col, resolved_gene_id_col, resolved_gene_name_col)

    list(transcript = transcript_info,
        gene = unique(transcript_info[,2:3]),
        tx2gene = transcript_info[,1:2],
        tx2gene_augmented = tx2gene_augmented,
        extra = extra)
}

#' Create a SummarizedExperiment Object
#'
#' Constructs a SummarizedExperiment object using provided matrices for counts, abundance, and length,
#' along with metadata for columns and rows. This function facilitates the organization of experimental
#' data (e.g., RNA-seq or other high-throughput data) in a structured format that is convenient for
#' further analyses and visualization.
#'
#' @param counts A matrix or DataFrame containing counts data, with rows as features (e.g., genes) and
#' columns as samples.
#' @param abundance A matrix or DataFrame containing abundance data (e.g., TPM or FPKM) with the same
#' dimensions and row/column names as the counts data.
#' @param length A matrix or DataFrame containing feature lengths, matching the dimensions and row/column
#' names of the counts data.
#' @param col_data A DataFrame containing sample-level metadata, with rows corresponding to columns in the
#' counts, abundance, and length matrices.
#' @param row_data A DataFrame containing feature-level metadata, with rows corresponding to features in
#' the counts, abundance, and length matrices.
#'
#' @return A SummarizedExperiment object containing the supplied data and metadata.

create_summarized_experiment <- function(counts, abundance, length, col_data, row_data) {
    SummarizedExperiment(assays = list(counts = counts, abundance = abundance, length = length),
        colData = col_data,
        rowData = row_data)
}

################################################
################################################
## Main script starts here                    ##
################################################
################################################

# Set defaults and classes
opt <- list(
    tx_col = "transcript_id",
    gene_id_col = "gene_id",
    gene_name_col = "gene_name"
)

# Apply parameter overrides from ext.args
args_opt <- parse_args('$task.ext.args')
for (ao in names(args_opt)) {
    if (ao %in% names(opt)) {
        opt[[ao]] <- args_opt[[ao]]
    }
}

# Define pattern for file names based on quantification type
if ('$quant_type' == "rsem") {
    # RSEM: .isoforms.results files are flat in quants/ directory (not in subdirectories)
    fns <- list.files('quants', pattern = 'isoforms\\\\.results\$', recursive = TRUE, full.names = TRUE)
    names <- gsub("\\\\.isoforms\\\\.results\$", "", basename(fns))
    names(fns) <- names

    # Import transcript-level quantifications from RSEM isoform results
    txi <- tximport(fns, type = 'rsem', txIn = TRUE, txOut = TRUE)
} else {
    # Salmon/Kallisto: files are in sample subdirectories
    pattern <- ifelse('$quant_type' == "kallisto", "abundance.tsv", "quant.sf")
    fns <- list.files('quants', pattern = pattern, recursive = TRUE, full.names = TRUE)
    names <- basename(dirname(fns))
    names(fns) <- names
    dropInfReps <- '$quant_type' == "kallisto"

    # Import transcript-level quantifications
    txi <- tximport(fns, type = '$quant_type', txOut = TRUE, dropInfReps = dropInfReps)
}

# Read transcript and sample data
transcript_info <- read_transcript_info('$tx2gene', opt\$tx_col, opt\$gene_id_col, opt\$gene_name_col)

# Make coldata just to appease the summarizedexperiment
coldata <- data.frame(files = fns, names = names, check.names = FALSE)
rownames(coldata) <- coldata[["names"]]

# Create initial SummarizedExperiment object
se <- create_summarized_experiment(txi[["counts"]], txi[["abundance"]], txi[["length"]],
    DataFrame(coldata), transcript_info\$transcript)

# Setting parameters for writing tables
params <- list(
    list(obj = se, slot = "abundance", suffix = "transcript_tpm.tsv"),
    list(obj = se, slot = "counts", suffix = "transcript_counts.tsv"),
    list(obj = se, slot = "length", suffix = "transcript_lengths.tsv")
)

# Process gene-level data if tx2gene mapping is available
if ("tx2gene" %in% names(transcript_info) && !is.null(transcript_info\$tx2gene)) {
    tx2gene <- transcript_info\$tx2gene
    gi <- summarizeToGene(txi, tx2gene = tx2gene)
    gi.ls <- summarizeToGene(txi, tx2gene = tx2gene, countsFromAbundance = "lengthScaledTPM")
    gi.s <- summarizeToGene(txi, tx2gene = tx2gene, countsFromAbundance = "scaledTPM")

    # Remove fake gene entries created by unmapped transcripts (where gene_id
    # was set to the transcript ID). These would break downstream processes
    # like SummarizedExperiment that try to match gene IDs against the tx2gene.
    if (length(transcript_info\$extra) > 0) {
        real_genes <- setdiff(rownames(gi[[1]]), transcript_info\$extra)
        filter_rows <- function(txi_list, genes) {
            lapply(txi_list, function(x) {
                if (is.matrix(x)) x[genes, , drop = FALSE] else x
            })
        }
        gi    <- filter_rows(gi, real_genes)
        gi.ls <- filter_rows(gi.ls, real_genes)
        gi.s  <- filter_rows(gi.s, real_genes)
    }

    gene_info <- transcript_info\$gene[match(rownames(gi[[1]]), transcript_info\$gene[["gene_id"]]),]
    rownames(gene_info) <- NULL
    col_data_frame <- DataFrame(coldata)

    # Create gene-level SummarizedExperiment objects
    gse <- create_summarized_experiment(gi[["counts"]], gi[["abundance"]], gi[["length"]],
        col_data_frame, gene_info)
    gse.ls <- create_summarized_experiment(gi.ls[["counts"]], gi.ls[["abundance"]], gi.ls[["length"]],
        col_data_frame, gene_info)
    gse.s <- create_summarized_experiment(gi.s[["counts"]], gi.s[["abundance"]], gi.s[["length"]],
        col_data_frame, gene_info)

    params <- c(params, list(
        list(obj = gse, slot = "length", suffix = "gene_lengths.tsv"),
        list(obj = gse, slot = "abundance", suffix = "gene_tpm.tsv"),
        list(obj = gse, slot = "counts", suffix = "gene_counts.tsv"),
        list(obj = gse.ls, slot = "counts", suffix = "gene_counts_length_scaled.tsv"),
        list(obj = gse.s, slot = "counts", suffix = "gene_counts_scaled.tsv")
    ))
}

# Writing tables for each set of parameters

prefix <- ''
if ('$task.ext.prefix' != 'null'){
    prefix = '$task.ext.prefix'
} else if ('$meta.id' != 'null'){
    prefix = '$meta.id'
}

done <- lapply(params, write_se_table, prefix)

write.table(transcript_info\$tx2gene_augmented,
    paste0(prefix, ".tx2gene_augmented.tsv"),
    sep = "\t", quote = FALSE, row.names = FALSE)

################################################
################################################
## R SESSION INFO                             ##
################################################
################################################

sink(paste(prefix, "R_sessionInfo.log", sep = '.'))
citation("tximeta")
print(sessionInfo())
sink()

################################################
################################################
## VERSIONS FILE                              ##
################################################
################################################

r.version <- strsplit(version[['version.string']], ' ')[[1]][3]
tximeta.version <- as.character(packageVersion('tximeta'))

writeLines(
    c(
        '"${task.process}":',
        paste('    bioconductor-tximeta:', tximeta.version)
    ),
'versions.yml')

################################################
################################################
################################################
################################################
