import pandas as pd


"""Function accepts a STAR output directory and compiles all sample information from ReadsPerGene.out.tab

Args:
    snakemake.input (list): list of globbed wildcards STAR ReadsPerGene.out.tab
    project_title (str): Project title for compiled STAR counts
Returns:
    Compiled STAR counts as tab delimited file.
"""

colnames = snakemake.params.samples
tables = [pd.read_csv(fh, header=None, skiprows=4, usecols=[0,3], index_col=0, sep = '\t',  names = ['Genes',fh.split('/')[-2].split('_')[0]]) for fh in snakemake.input]
joined_table = pd.concat(tables, axis=1)
joined_table.columns = colnames
joined_table_sorted = joined_table.reindex(sorted(joined_table.columns), axis = 1)
joined_table_sorted.to_csv(snakemake.output[0], sep='\t')

sink(file(snakemake@log[[1]], open = "wt"), type = "message")

library(DESeq2)
library(ggplot2)
library(tibble)
library(dplyr)
library(viridisLite)
library(pheatmap)
library(yaml)

parallel <- FALSE
if (snakemake@threads > 1) {
	library("BiocParallel")
	register(MulticoreParam(snakemake@threads))
	parallel <- TRUE
}

if (!dir.exists(snakemake@output$outdir)) {dir.create(snakemake@output$outdir)}

# import data -------------------------------------------------------------------------------------\

# import counts
counts = read.table(snakemake@input$counts, sep = "\t", header = TRUE, stringsAsFactors = FALSE, check.names = FALSE)
gene_names = counts |> pull(Genes)

# import metdata
md = read.table(snakemake@input$md, header = TRUE, stringsAsFactors = FALSE, check.names = FALSE)
rownames(md) = md$SampleID

# subset md by config/GROUPS.txt if provided
# allow user to upload a contrast combination file for easier interpretation of the output
if (file.exists(snakemake@config$GROUPS)) {

	# import contrast combination file
	groups = read.table(snakemake@config$GROUPS) |> pull(1)

	# subset conditions in md with the conditions in config/GROUPS.txt
	md = md |> filter(Condition %in% groups)
}

# make the md reflect the available sample names
md = md[md$SampleID %in% colnames(counts),]

# make sure md rownames and counts colnames in the same order.
rownames(md) = md$SampleID
counts = counts[, rownames(md)]
stopifnot(rownames(md) == colnames(counts))

# DESeq2 ------------------------------------------------------------------------------------------

# differential expression
dds = DESeqDataSetFromMatrix(countData = counts, colData = md, design = as.formula("~Condition"))
dds.lrt <- DESeq(dds, test="LRT", reduced=~1, parallel = parallel)
res.lrt <- results(dds.lrt, cooksCutoff = Inf, independentFiltering=FALSE)

# Obtain normalized counts
# vsd = vst(dds.lrt, blind = FALSE)
rld <- rlog(dds.lrt, blind=FALSE)

# visualization -----------------------------------------------------------------------------------

colors_fwd <- viridis(100)
colors_rev <- viridis(100, direction = -1)

### sample-sample distances
sampleDists = dist(t(assay(rld)), diag = TRUE)

sampleDistsMatrix = sampleDists |> as.matrix()
rownames(sampleDistsMatrix) <- rld$SampleID
colnames(sampleDistsMatrix) <- NULL

distance_plot = paste0(snakemake@output$outdir, "/", snakemake@config$PROJECT_ID, "-group-sample-dists.pdf")
pdf(distance_plot, width = 12, height = 12)
pheatmap(sampleDistsMatrix,
	fontsize = 12,
	clustering_distance_rows = sampleDists,
	clustering_distance_cols = sampleDists,
	cluster_rows = TRUE,
	cluster_cols = TRUE,
	col = colors_rev
)
dev.off()

### PCA plot
pca_plot = paste0(snakemake@output$outdir, "/", snakemake@config$PROJECT_ID, "-group-pca.pdf")
pdf(pca_plot, width = 12, height = 12)
plotPCA(rld, intgroup=c("Condition"))
dev.off()

### Gene expression - top 50 genes post DESeq2 normalization
# top_ge = assay(rld)[top_genes,]
top_gene_indices <- head(order(res.lrt$padj), 50)
top_ge = counts(dds.lrt, normalized=TRUE)[top_gene_indices, ]
rownames(top_ge) = gene_names[top_gene_indices]

top_50 = paste0(snakemake@output$outdir, "/", snakemake@config$PROJECT_ID, "-top-50-genes-heatmap.pdf")
pdf(top_50, width = 12, height = 12)
pheatmap(top_ge,
	main = "Top 50 differential gene expression (DESeq2-normalized counts)",
	scale = "row",
	fontsize = 12,
	cluster_rows = TRUE,
	cluster_cols = TRUE
)
dev.off()

### Gene expression - all genes by individual replicates
all_ge = counts(dds.lrt, normalized=TRUE)
rownames(all_ge) = gene_names
all_ge = all_ge |>
	as.data.frame() |>
	distinct() |>
	as.matrix()
all_ge = all_ge[apply(all_ge, 1, function(x) sd(x) != 0), ]

all_genes_heatmap_individual = paste0(snakemake@output$outdir, "/", snakemake@config$PROJECT_ID, "-all-genes-heatmap-individual.pdf")
pdf(all_genes_heatmap_individual, width = 12, height = 12)
pheatmap(all_ge,
	main = "All gene expression (DESeq2-normalized counts) by individual samples",
	scale = "row",
	fontsize = 12,
	cluster_rows = TRUE,
	cluster_cols = TRUE,
	show_rownames = FALSE
)
dev.off()

### Gene expression - all genes by merged replicates

if (snakemake@config$MERGE_REPLICATES) {

	message("Merging replicates...")

	# read YAML file
	merge_scheme = read_yaml(snakemake@config$REPLICATES)
	for (name in names(merge_scheme)) {

		reps = merge_scheme[[name]]

		if (snakemake@config$MERGE_METHOD == "mean") {
			merged_counts = all_ge[, reps] |> apply(MARGIN = 1, FUN = mean) # same as rowMeans(counts[, reps])
		} else if (snakemake@config$MERGE_METHOD == "median") {
			merged_counts = all_ge[, reps] |> apply(MARGIN = 1, FUN = median)
		} else {
			stop("ERROR: MERGE_METHOD must be either 'mean' or 'median'")
		}

		all_ge = all_ge |>
			as.data.frame() |>
			mutate(!!enexpr(name) := merged_counts) |>
			select(!all_of(reps)) |>
			as.matrix()
		}

	# remove rows with sd=0 one more time.
	all_ge = all_ge[apply(all_ge, 1, function(x) sd(x) != 0), ]

	# define output
	all_genes_heatmap_merged = paste0(snakemake@output$outdir, "/", snakemake@config$PROJECT_ID, "-all-genes-heatmap-merged.pdf")
	pdf(all_genes_heatmap_merged, width = 12, height = 12)
	pheatmap(all_ge,
		main = "All gene expression (DESeq2-normalized counts) by merged samples",
		scale = "row",
		fontsize = 12,
		cluster_rows = TRUE,
		cluster_cols = FALSE, # this respects the order of replicates.yaml
		show_rownames = FALSE
	)
	dev.off()

}

import yaml
import pandas as pd
import os

# identify the annotation provided by the STAR index (user-defined)
# filter_col will be used for filtering columns later on.
if snakemake.config["ANNOTATION_STYLE"] == "Genes":
	filter_col = "Gene name lower"
elif snakemake.config["ANNOTATION_STYLE"] == "UCSC":
	filter_col = "UCSC Stable ID"
elif snakemake.config["ANNOTATION_STYLE"] == "ENSEMBL_ID":
	filter_col = "Gene stable ID"
elif snakemake.config["ANNOTATION_STYLE"] == "REFSEQ":
	filter_col = "RefSeq mRNA ID"
elif snakemake.config["ANNOTATION_STYLE"] == "NCBI":
	filter_col = "NCBI gene (formerly Entrezgene) ID"
else:
	print("ERROR: ANNOTATION_STYLE is undefined. Please select from one of the following options: ['Genes', 'UCSC', 'ENSEMBL_ID', 'REFSEQ', 'NCBI']")
	os.exit()

# annotation file -------------------------------------------------------------

# open annotation file
if snakemake.input["annotation"]:
	annotation = pd.read_csv(snakemake.config["ANNOTATION"], sep = "\t")
else:
	print("ERROR: ANNOTATION key in config.yaml is undefined")
	os.exit()

# edit content of the annotation file

# turn gene type and names to lower-case to ensure string matching downstream.
annotation["Gene type lower"] = annotation["Gene type"].transform(lambda x: str.lower(x))
annotation["Gene name lower"] = annotation["Gene name"].transform(lambda x: str(x).lower())

# import biotypes. depending on their values, get them into a list.

biotypes = snakemake.config["BIOTYPES"]
biotype_type = type(biotypes)

if biotype_type == str:
	biotypes = [ snakemake.config["BIOTYPES"] ]
elif biotype_type == list:
	biotypes = snakemake.config["BIOTYPES"]
else:
	print("ERROR: BIOTYPES must be a single entry <string> or multiple entries <list of strings>")

biotypes = [ i.lower() for i in biotypes ]

# filter annotation file contents by biotype
annotation = annotation[annotation["Gene type lower"].isin(biotypes)]
filter_features = list(annotation[filter_col]) # use these features to filter counts table

# open counts file ------------------------------------------------------------

# import counts
counts = pd.read_csv(snakemake.input["counts"], sep = "\t")

print("PREFILTER: Counts table rows = {}".format(counts.shape[0]))
print("PREFILTER: Counts table columns = {}".format(counts.shape[1]))

# filter features
counts["Gene lower"] = counts["Genes"].transform(lambda x: str.lower(x))
counts = counts[counts["Gene lower"].isin(filter_features)]

print("POSTFILTER: Counts table rows = {}".format(counts.shape[0]))
print("POSTFILTER: Counts table columns = {}".format(counts.shape[1] - 1)) # minus one to adjust for lower-case gene names col

counts = counts.drop("Gene lower", axis = 1)

# export filtered counts
counts.to_csv(str(snakemake.output), sep = "\t", index = False)

shell:
	"fastp "
	"-i {input.r1} "
	"-I {input.r2} "
	"-o {output.r1} "
	"-O {output.r2} "
	"--detect_adapter_for_pe "
	"--thread {threads} "
	"-j {log} "
	"-h /dev/null"

shell:
	"fastqc -t {threads} --outdir data/fastqc {input} > {log} 2>&1"

shell:
	"fastq_screen --aligner bowtie2 --threads {threads} --outdir data/fastq_screen --conf {input.config} --force {input.fastq} > {log} 2>&1"

shell:
	"samtools index -@ {threads} {input}"

shell:
	"preseq c_curve -B {resources.defect_mode} -l 1000000000 -P -o {output} {input} > {log} 2>&1"

shell:
	"bamCoverage -b {input[0]} -o {output} -p {threads} --normalizeUsing CPM --binSize 10 --smoothLength 50"

shell:
	"if [ '{params.singularity}' == 'true' ]; then export LC_ALL=C.UTF-8; export LANG=C.UTF-8; fi && "
	"multiqc data -f --ignore data/tmp -o data/multiqc 2>&1"