run:
    import pandas as pd

    # Open all TSV-files of input, transpose them and merge them into one dataframe
    # Keep filename as row name
    df_list = []
    for file in input:
        current_df = pd.read_csv(file, sep=":\s*", engine="python", header=None, index_col=0)
        current_df = current_df.transpose()
        current_df["BAM-file"] = os.path.basename(file)
        df_list.append(current_df)

    # Finalize column order and output merged dataframe
    output_df = pd.concat(df_list, axis=0)
    cols = output_df.columns.tolist()
    final_cols = cols[-1:] + cols[:-1]
    output_df = output_df[final_cols]
    output_df.to_csv(output[0], sep="\t", index=False)

shell:
    "plotPCA -in {input} -o {output[0]} --plotTitle \"{params.plot_title}\" {params.extras} 2> {log}"

wrapper:
    "0.79.0/bio/fastqc"

shell:
    "kraken2 --use-names --threads {threads} --db {params.kraken2_db} "
    "--report {output.kraken2_report} "
    "--paired {input.r1} {input.r2} "
    "> {output.kraken2_kmer_mapping} 2> {log}"

wrapper:
    "v0.86.0/bio/multiqc"

wrapper:
    "v1.12.0/bio/arriba"

shell:
    "samtools view -h {input} | "
    "awk 'BEGIN {{OFS=\"\t\"}} {{"
        "split($6,C,/[0-9]*/); split($6,L,/[SMDIN]/); "
            "if (C[2]==\"S\") {{$10=substr($10,L[1]+1); $11=substr($11,L[1]+1)}}; "
            "if (C[length(C)]==\"S\") {{L1=length($10)-L[length(L)-1]; $10=substr($10,1,L1); $11=substr($11,1,L1); }}; "
        "gsub(/[0-9]*S/,\"\",$6); print}}' - >{output}"

shell:
    'cufflinks '
    '--num-threads {threads}'
    ' --library-type {params.library_type}'
    ' --GTF-guide {params.gtf_file}'                            # use reference transcript annotation to guide assembly, but also includes novel transcripts
    ' --frag-bias-correct {params.ref_seq} '                    # use bias correction - reference fasta required 
    ' --min-isoform-fraction {params.min_isoform_fraction}'     # suppress transcripts below this abundance level (compared with major isoform of the gene)
    ' --min-frags-per-transfrag {params.min_frags_per_transfrag}'   # assembled transfrags supported by fewer than this many aligned RNA-Seq fragments are ignored
    ' --output-dir {params.output_dir} '
    '{params.extra_options}'                                    # additional options
    '{input}'

wrapper:
    "v1.19.2/bio/kallisto/quant"

shell:
    "htseq-count "
    "-n {threads} "
    "--format {params.format} "
    "--order {params.order} "
    "--stranded {params.stranded} "
    "--additional-attr {params.add_attribute} "
    "{input.bam_file} {input.gtf_file} >{output.count_file} 2>{log}"

library("cummeRbund")


global_statistics_and_qc <- function(cuff_obj, output_dir) {
  # ---- Global Statistics and Quality Control -----------

  # Dispersion explained:
  # https://genomebiology.biomedcentral.com/articles/10.1186/s13059-014-0550-8
  # https://www.biostars.org/p/167688/
  # https://support.bioconductor.org/p/75260/

  # ------------- 1. Dispersion ----------------
  # -> visualizes the estimated overdispersion for each sample
  # uses cufflinks emitted data (mean counts, variance, & dispersion)
  # -> http://cole-trapnell-lab.github.io/cufflinks/cuffdiff/
  genes.disp<-dispersionPlot(genes(cuff_obj))
  pdf(file=file.path(output_dir, "cummerbund_figures/dispersion_genes.pdf"))
  plot(genes.disp)				# Plot is displayed
  dev.off()

  cuff_obj.disp<-dispersionPlot(cuff_obj)
  pdf(file=file.path(output_dir, "dispersion_cuff.pdf"))
  plot(cuff_obj.disp)				# Plot is displayed
  dev.off()

  # ------ 2. Distributions of FPKM scores across samples ----------
  # 2.1.) csDensity plots
  dens<-csDensity(genes(cuff_obj))
  pdf(file=file.path(output_dir, "fpkm_density_genes.pdf"))
  plot(dens)
  dev.off()

  dens<-csDensity(isoforms(cuff_obj))
  pdf(file=file.path(output_dir, "fpkm_density_isoforms.pdf"))
  plot(dens)
  dev.off()

  # 2.2.) Boxplots
  b<-csBoxplot(genes(cuff_obj))
  pdf(file=file.path(output_dir, "fpkm_density_boxplot_genes.pdf"))
  plot(b)
  dev.off()

  b<-csBoxplot(isoforms(cuff_obj))
  pdf(file=file.path(output_dir, "fpkm_density_boxplot_isoforms.pdf"))
  plot(b)
  dev.off()

  # 2.3.) Matrix of pairwise scatterplots
  s<-csScatterMatrix(genes(cuff_obj))
  pdf(file=file.path(output_dir, "fpkm_density_matrix_genes.pdf"))
  plot(s)
  dev.off()

  s<-csScatterMatrix(isoforms(cuff_obj))
  pdf(file=file.path(output_dir, "fpkm_density_matrix_isoforms.pdf"))
  plot(s)
  dev.off()

  # 2.4.) Volcano plots -> Explore relationship between fold-change and significance
  v<-csVolcanoMatrix(genes(cuff_obj))
  pdf(file=file.path(output_dir, "cs_volcano_matrix_genes.pdf"))
  plot(v)
  dev.off()

  v<-csVolcanoMatrix(isoforms(cuff_obj))
  pdf(file=file.path(output_dir, "cs_volcano_matrix_isoforms.pdf"))
  plot(v)
  dev.off()
}


analyse_differential_expression <- function(cuff_obj, output_dir) {
  # ------ Differential expression -------------
  # 1.) genes
  # all
  gene.diff <- diffData(genes(cuff_obj))
  write.csv(gene.diff, file.path(output_dir, "gene_diff.csv"), row.names=F)

  # only significant -> with gene names
  sig_gene_ids <- getSig(cuff_obj,level="genes",alpha=0.05)
  if (NROW(sig_gene_ids) > 0) {
    sigFeatures <- getFeatures(cuff_obj,sig_gene_ids,level="genes")
    sigData <- diffData(sigFeatures)
    sigData <- subset(sigData, (significant == 'yes'))
    names <- featureNames(sigFeatures)
    sigOutput <- merge(names, sigData, by.x="tracking_id", by.y="gene_id")

    # Patch the merged table to have the original name for the ID column.
    # This is always the first column for the examples we've seen.
    colnames(sigOutput)[1] <- "gene_id"
    write.table(sigOutput, file.path(output_dir, "gene_diff_only_significant.tsv"), sep='\t', row.names = F,
                col.names = T, quote = F)
  } else {
    sink(file.path(output_dir, "gene_diff_only_significant.tsv"))
    cat("No significantly differently expressed genes detected")
    sink()
  }

  # 2.) isoforms
  # all
  isoform.diff <- diffData(isoforms(cuff_obj))
  write.csv(isoform.diff, file.path(output_dir, "isoform_diff.csv"), row.names=F)

  # only significant -> with gene names
  sig_isoforms_ids <- getSig(cuff_obj, level="isoforms", alpha=0.05)
  if (NROW(sig_isoforms_ids) > 0) {
    sigFeatures <- getFeatures(cuff_obj, sig_isoforms_ids, level="isoforms")
    sigData <- diffData(sigFeatures)
    sigData <- subset(sigData, (significant == 'yes'))
    names <- featureNames(sigFeatures)
    sigOutput <- merge(names, sigData, by.x="tracking_id", by.y="isoform_id")

    # Patch the merged table to have the original name for the ID column.
    # This is always the first column for the examples we've seen.
    colnames(sigOutput)[1] <- "gene_id"
    write.table(sigOutput, file.path(output_dir, "isoform_diff_only_significant.tsv"), sep='\t', row.names = F,
                col.names = T, quote = F)
  } else {
    sink(file.path(output_dir, "isoform_diff_only_significant.tsv"))
    cat("No significantly differently expressed isoforms detected")
    sink()
  }
}


create_count_matrices <- function(cuff_obj, output_dir) {
  " Create count matrices for genes and isoforms "

  # ------------- CSV files ----------------
  # access feature lvl data
  gene.features <- annotation(genes(cuff_obj))
  write.csv(gene.features, file.path(output_dir, "gene_features.csv"), row.names = FALSE)

  gene.fpkm <- fpkm(genes(cuff_obj))
  write.csv(gene.fpkm, file.path(output_dir, "gene_fpkm.csv"), row.names = FALSE)

  # raw and normalized (on sequencing depth?) fragment counts
  gene.counts <- count(genes(cuff_obj))
  write.csv(gene.counts, file.path(output_dir, "gene_counts.csv"), row.names = FALSE)

  # -- isoforms
  isoform.features <- annotation(isoforms(cuff_obj))
  write.csv(isoform.features, file.path(output_dir, "isoform_features.csv"), row.names = FALSE)

  isoform.fpkm <- fpkm(isoforms(cuff_obj))
  write.csv(isoform.fpkm, file.path(output_dir, "isoform_fpkm.csv"), row.names = FALSE)

  isoform.counts <- count(isoforms(cuff_obj))
  write.csv(isoform.counts, file.path(output_dir, "isoform_counts.csv"), row.names = FALSE)


  # ----------- create PDFs -----------
  # FPKM matrices
  gene.fpkm.matrix<-fpkmMatrix(genes(cuff_obj))
  pdf(file=file.path(output_dir, "fpkm_matrix_genes.pdf"))
  plot(gene.fpkm.matrix)
  dev.off()

  isoform.fpkm.matrix<-fpkmMatrix(isoforms(cuff_obj))
  pdf(file=file.path(output_dir, "fpkm_matrix_isoforms.pdf"))
  plot(isoform.fpkm.matrix)
  dev.off()

  # Count matrices
  gene.count.matrix<-countMatrix(genes(cuff_obj))
  pdf(file=file.path(output_dir, "count_matrix_genes.pdf"))
  plot(gene.count.matrix)
  dev.off()

  isoforms.count.matrix<-countMatrix(isoforms(cuff_obj))
  pdf(file=file.path(output_dir, "count_matrix_isoforms.pdf"))
  plot(isoforms.count.matrix)
  dev.off()
}


single_gene_analysis <- function(cuff_obj, gene_of_interest_id, output_dir) {
  " Detailed analysis of a single gene of interest"
  # IV) ---- Single gene -----------
  myGeneId <- gene_of_interest_id
  myGene<-getGene(cuff_obj,myGeneId)
  header <- paste("Single Gene", arg_gene, ":", sep=" ")
  capture.output(myGene, file=file.path(output_dir, "analysis_output.txt"), append = TRUE)

  output_file_genes <- file.path(output_dir, paste(arg_gene, "_gene_fpkm.csv", sep=""))
  write.csv(fpkm(myGene), output_file_genes, row.names = FALSE)

  output_file_isoforms <- file.path(output_dir, paste(arg_gene, "_isoforms_fpkm.csv", sep=""))
  write.csv(fpkm(isoforms(myGene)), output_file_isoforms, row.names = FALSE)

  # Plots
  gl <- expressionPlot(myGene)
  output_file_1 <- file.path(output_dir, paste("expressionPlot_singleGene_", arg_gene, ".pdf", sep=""))
  pdf(file=output_file_1)
  plot(gl)
  dev.off()

  # Expression plot of all isoforms of a single gene with FPKMs exposed
  gl.iso.rep <- expressionPlot(isoforms(myGene))
  output_file_2 <- file.path(output_dir, paste("expressionPlot_isoforms_singleGene_", arg_gene, ".pdf", sep=""))
  pdf(file=output_file_2)
  plot(gl.iso.rep)
  dev.off()

  # Expression plot of all CDS for a single gene with FPKMS exposed
  gl.cds.rep<-expressionPlot(CDS(myGene))
  output_file_3 <- file.path(output_dir, paste("expressionPlot_cds_singleGene_", arg_gene, ".pdf", sep=""))
  pdf(file=output_file_3)
  plot(gl.cds.rep)
  dev.off()

  # Detailed feature graph
  trackList<-list()
  myStart<-min(features(myGene)$start)
  myEnd<-max(features(myGene)$end)
  myChr<-unique(features(myGene)$seqnames)
  genome<-arg_genome
  ideoTrack <- IdeogramTrack(genome = genome, chromosome = myChr)
  trackList<-c(trackList,ideoTrack)   # appending ideoTrack -> chromosome
  axtrack<-GenomeAxisTrack()
  trackList<-c(trackList,axtrack)     # appending axtrack -> genome-Axis
  genetrack<-makeGeneRegionTrack(myGene)

  trackList<-c(trackList,genetrack)   # appending genetrack -> the mapping results
  biomTrack<-BiomartGeneRegionTrack(genome=genome,chromosome=as.character(myChr), start=myStart,end=myEnd,name="ENSEMBL",showId=T)
  trackList<-c(trackList,biomTrack)   # Biomart transcripts
  conservation <- UcscTrack(genome = genome, chromosome = myChr, track = "Conservation", table = "multiz100way",from = myStart-2000, to = myEnd+2000, trackType = "DataTrack",start = "start", end = "end", data = "score",type = "hist", window = "auto", col.histogram = "darkblue",fill.histogram = "darkblue", ylim = c(-3.7, 4),name = "Conservation")
  trackList<-c(trackList,conservation)    # conservation
  # Plot detailed graph: Chromosome on top...
  pdf(file.path(output_dir, "detailed_track.pdf"))
  plotTracks(trackList,from=myStart-2000,to=myEnd+2000)
  dev.off()
}

# ----------- Main -----------
main <- function() {
  # Inputs
  cufflinks_merged_transcriptome_assemblies_gtf <- snakemake@input["merged_cufflinks_transcriptome_assemblies_gtf"]

  # Params
  cufflinks_output_files_dir <- snakemake@params["cufflinks_output_files_dir"]
  genome_build <- snakemake@params["original_genome_build"]
  chosen_genes_of_interest <- snakemake@params[["chosen_genes_of_interest"]]

  # Output directory
  cummerbund_output_dir <- snakemkae@output[["cummerbund_output_dir"]]
  cummerbund_summary_results <- snakemake@output[["cummerbund_summary_results"]]

  # read results -> create SQLite db
  # Rebuild is important: Create always new database
  cuff_obj <- readCufflinks(dir=cufflinks_output_files_dir, gtfFile=cufflinks_merged_transcriptome_assemblies_gtf,
                            genome=genome_build, rebuild=T)
  capture.output(cuff_obj, file=file.path(cummerbund_output_dir, cummerbund_summary_results), append = FALSE)

  # Global statistics and qc
  global_statistics_and_qc(cuff_obj, cummerbund_output_dir)

  # Count matrices
  create_count_matrices(cuff_obj, cummerbund_output_dir)

  # Single gene analysis
  for (gene_of_interest in chosen_genes_of_interest) {
      single_gene_analysis(cuff_obj, gene_of_interest, cummerbund_output_dir)
  }
}

# Execute main
main()

library("BiocParallel")

# Load data
library("tximeta")	# Import transcript quantification data from Salmon
library("tximport")	# Import transcript-level quantification data from Kaleidoscope, Sailfish, Salmon, Kallisto, RSEM, featureCounts, and HTSeq
library("rhdf5")
library("SummarizedExperiment")

library("magrittr")	# Pipe operator
library("DESeq2")		# Differential gene expression analysis

# Plotting libraries
library("pheatmap")
library("RColorBrewer")
library("ggplot2")

# Mixed
library("PoiClaClu")
library("glmpca")
library("apeglm")
library("genefilter")
library("AnnotationDbi")
library("org.Hs.eg.db")


# ---------------- Loading read/fragment quantification data from Salmon output ----------------
load_in_salmon_generated_counts <- function(annotation_table_file) {
	# ----------------- 1. Load annotation -----------------
	# Columns: 1. names, 2. files, 3. condition, 4. additional information
	annotation_table <- read.csv(file=annotation_table_file, sep="\t")

	annotation_data <- data.frame(
									names=annotation_table[,"sample_name"],
									files=file.path(annotation_table[,"salmon_results_file"]),
									condition=annotation_table[,"condition"],
									add_info=annotation_table[,"additional_comment"]
									)
	# Replace None in condition column with "Control"
	annotation_data$condition[annotation_data$condition=="None"] <- "Control"
	annotation_data$condition[annotation_data$condition==""] <- "Control"

	# ----------------- 2. Load into Bioconductor experiment objects -----------------
	# Summarized experiment: Imports quantifications & metadata from all samples -> Each row is a transcript
	se <- tximeta(annotation_data)

	# Summarize transcript-level quantifications to the gene level -> reduces row number: Each row is a gene
	# Includes 3 matrices:
	# 1. counts: Estimated fragment counts per gene & sample
	# 2. abundance: Estimated transcript abundance in TPM
	# 3. length: Effective Length of each gene (including biases as well as transcript usage)
	gse <- summarizeToGene(se)

	# ----------------- 3. Load experiments into DESeq2 object -----------------
	# SummarizedExperiment
	# assayNames(gse)   		# Get all assays -> counts, abundance, length, ...
	# head(assay(gse), 3)     	# Get count results for first 3 genes
	# colSums(assay(gse))     	# Compute sums of mapped fragments
	# rowRanges(gse)          	# Print rowRanges: Ranges of individual genes
	# seqinfo(rowRanges(gse))   # Metadata of sequences (chromosomes in our case)

	gse$condition <- as.factor(gse$condition)
	gse$add_info <- as.factor(gse$add_info)

	# Use relevel to make sure untreated is listed first
	gse$condition %<>% relevel("Control")   # Concise way of saying: gse$condition <- relevel(gse$condition, "Control")

	# Construct DESeqDataSet from gse
	if (gse$add_info %>% unique %>% length >1) {
		# Add info column with more than 1 unique value
		print("More than 1 unique value in add_info column")
		# TODO Need to make sure to avoid:
		# the model matrix is not full rank, so the model cannot be fit as specified.
		#   One or more variables or interaction terms in the design formula are linear
		#   combinations of the others and must be removed.
		# dds <- DESeqDataSet(gse, design = ~condition + add_info)
		print("However, simple DESeq2 analysis will be performed without add_info column")
		dds <- DESeqDataSet(gse, design = ~condition)

	} else {
		print("Only 1 unique value in add_info column")
		dds <- DESeqDataSet(gse, design = ~condition)
	}

	return(dds)
}


# ---------------- Loading read/fragment quantification data from RSEM output ----------------
load_in_rsem_generated_counts <- function(annotation_table_file) {
	# ----------------- 1. Load annotation -----------------
	annotation_table <- read.csv(file=annotation_table_file, sep="\t")
	files <- file.path(annotation_table[,"rsem_results_file"])
	# For sample.genes.results: txIn= FALSE & txOut= FALSE
	# For sample.isoforms.results: txIn= TRUE & txOut= TRUE
	# Check: https://bioconductor.org/packages/devel/bioc/vignettes/tximport/inst/doc/tximport.html
	txi.rsem <- tximport(files, type = "rsem", txIn = FALSE, txOut = FALSE)

	annotation_data <- data.frame(condition=factor(annotation_table[,"condition"]),
									add_info=factor(annotation_table[,"additional_comment"])
						)
	rownames(annotation_data) <- annotation_table[,"sample_name"]

	# Construct DESeqDataSet from tximport
	if (annotation_data$add_info %>% unique %>% length >1) {
		# Add info column with more than 1 unique value
		# dds <- DESeqDataSetFromTximport(txi.rsem, annotation_data, ~condition + add_info)
		dds <- DESeqDataSetFromTximport(txi.rsem, annotation_data, ~condition)
	} else {
		dds <- DESeqDataSetFromTximport(txi.rsem, annotation_data, ~condition)
	}
	return(dds)
}


load_in_kallisto_generated_counts <- function(annotation_table_file) {
	# ----------------- 1. Load annotation -----------------
	annotation_table <- read.csv(file=annotation_table_file, sep="\t")

	files <- file.path(annotation_table[,"kallisto_results_file"])
	txi.kallisto <- tximport(files, type = "kallisto", txOut = TRUE)

	annotation_data <- data.frame(condition=factor(annotation_table[,"condition"]),
									add_info=factor(annotation_table[,"additional_comment"])
						)
	rownames(annotation_data) <- annotation_table[,"sample_name"]

	# Construct DESeqDataSet from tximport
	if (annotation_data$add_info %>% unique %>% length >1) {
		# Add info column with more than 1 unique value
		# dds <- DESeqDataSetFromTximport(txi.kallisto, annotation_data, ~condition + add_info)
		dds <- DESeqDataSetFromTximport(txi.kallisto, annotation_data, ~condition)
	} else {
		dds <- DESeqDataSetFromTximport(txi.kallisto, annotation_data, ~condition)
	}
	return(dds)
}


# ----------------- Main function -----------------
main_function <- function(){
	threads <- snakemake@threads[[1]]
	register(MulticoreParam(workers=threads))

	# Snakemake variables
	annotation_table_file <- snakemake@input[["annotation_table_file"]]
	output_file <- snakemake@output[["deseq_dataset_r_obj"]]
	count_algorithm <- snakemake@params[["count_algorithm"]]

	# Load annotation table & Salmon data into a DESeq2 object
	if (count_algorithm == "salmon") {
		dds <- load_in_salmon_generated_counts(annotation_table_file)
	} else if (count_algorithm == "kallisto") {
		dds <- load_in_kallisto_generated_counts(annotation_table_file)
	} else if (count_algorithm == "rsem") {
		dds <- load_in_rsem_generated_counts(annotation_table_file)
	} else {
		stop("Count algorithm not supported!")
	}

	# Remove rows that have no or nearly no information about the amount of gene expression
	print(paste(c("Number of rows before filtering out counts with values <1", nrow(dds))))
	keep <- rowSums(counts(dds)) > 1 # Counts have to be greater than 1
	dds <- dds[keep,]
	print(paste(c("Number of rows after filtering out counts with values <1", nrow(dds))))

    # Save deseq dataset object
    saveRDS(dds, output_file)
}


# ----------------- Run main function -----------------
main_function()