__author__ = "Johannes Köster"
__copyright__ = "Copyright 2019, Johannes Köster"
__email__ = "[email protected]"
__license__ = "MIT"

import subprocess
import sys
from snakemake.shell import shell

species = snakemake.params.species.lower()
release = int(snakemake.params.release)
fmt = snakemake.params.fmt
build = snakemake.params.build
flavor = snakemake.params.get("flavor", "")

branch = ""
if release >= 81 and build == "GRCh37":
    # use the special grch37 branch for new releases
    branch = "grch37/"

if flavor:
    flavor += "."

log = snakemake.log_fmt_shell(stdout=False, stderr=True)

suffix = ""
if fmt == "gtf":
    suffix = "gtf.gz"
elif fmt == "gff3":
    suffix = "gff3.gz"

url = "ftp://ftp.ensembl.org/pub/{branch}release-{release}/{fmt}/{species}/{species_cap}.{build}.{release}.{flavor}{suffix}".format(
    release=release,
    build=build,
    species=species,
    fmt=fmt,
    species_cap=species.capitalize(),
    suffix=suffix,
    flavor=flavor,
    branch=branch,
)

try:
    shell("(curl -L {url} | gzip -d > {snakemake.output[0]}) {log}")
except subprocess.CalledProcessError as e:
    if snakemake.log:
        sys.stderr = open(snakemake.log[0], "a")
    print(
        "Unable to download annotation data from Ensembl. "
        "Did you check that this combination of species, build, and release is actually provided?",
        file=sys.stderr,
    )
    exit(1)

__author__ = "Julian de Ruiter"
__copyright__ = "Copyright 2017, Julian de Ruiter"
__email__ = "[email protected]"
__license__ = "MIT"


from snakemake.shell import shell


n = len(snakemake.input)
assert n == 1, "Input must contain 1 (single-end) element."

extra = snakemake.params.get("extra", "")
adapters = snakemake.params.get("adapters", "")
log = snakemake.log_fmt_shell(stdout=False, stderr=True)

assert (
    extra != "" or adapters != ""
), "No options provided to cutadapt. Please use 'params: adapters=' or 'params: extra='."

shell(
    "cutadapt"
    " --cores {snakemake.threads}"
    " {adapters}"
    " {extra}"
    " -o {snakemake.output.fastq}"
    " {snakemake.input[0]}"
    " > {snakemake.output.qc} {log}"
)

__author__ = "Patrik Smeds"
__copyright__ = "Copyright 2016, Patrik Smeds"
__email__ = "[email protected]"
__license__ = "MIT"

from os.path import splitext

from snakemake.shell import shell

log = snakemake.log_fmt_shell(stdout=False, stderr=True)

# Check inputs/arguments.
if len(snakemake.input) == 0:
    raise ValueError("A reference genome has to be provided!")
elif len(snakemake.input) > 1:
    raise ValueError("Only one reference genome can be inputed!")

# Prefix that should be used for the database
prefix = snakemake.params.get("prefix", splitext(snakemake.output.idx[0])[0])

if len(prefix) > 0:
    prefix = "-p " + prefix

# Contrunction algorithm that will be used to build the database, default is bwtsw
construction_algorithm = snakemake.params.get("algorithm", "")

if len(construction_algorithm) != 0:
    construction_algorithm = "-a " + construction_algorithm

shell(
    "bwa index" " {prefix}" " {construction_algorithm}" " {snakemake.input[0]}" " {log}"
)

__author__ = "Johannes Köster, Julian de Ruiter"
__copyright__ = "Copyright 2016, Johannes Köster and Julian de Ruiter"
__email__ = "[email protected], [email protected]"
__license__ = "MIT"


import tempfile
from os import path
from snakemake.shell import shell
from snakemake_wrapper_utils.java import get_java_opts
from snakemake_wrapper_utils.samtools import get_samtools_opts


# Extract arguments.
extra = snakemake.params.get("extra", "")
log = snakemake.log_fmt_shell(stdout=False, stderr=True)
sort = snakemake.params.get("sorting", "none")
sort_order = snakemake.params.get("sort_order", "coordinate")
sort_extra = snakemake.params.get("sort_extra", "")
samtools_opts = get_samtools_opts(snakemake)
java_opts = get_java_opts(snakemake)


index = snakemake.input.idx
if isinstance(index, str):
    index = path.splitext(snakemake.input.idx)[0]
else:
    index = path.splitext(snakemake.input.idx[0])[0]


# Check inputs/arguments.
if not isinstance(snakemake.input.reads, str) and len(snakemake.input.reads) not in {
    1,
    2,
}:
    raise ValueError("input must have 1 (single-end) or 2 (paired-end) elements")


if sort_order not in {"coordinate", "queryname"}:
    raise ValueError("Unexpected value for sort_order ({})".format(sort_order))


# Determine which pipe command to use for converting to bam or sorting.
if sort == "none":
    # Simply convert to bam using samtools view.
    pipe_cmd = "samtools view {samtools_opts}"

elif sort == "samtools":
    # Add name flag if needed.
    if sort_order == "queryname":
        sort_extra += " -n"

    # Sort alignments using samtools sort.
    pipe_cmd = "samtools sort {samtools_opts} {sort_extra} -T {tmpdir}"

elif sort == "picard":
    # Sort alignments using picard SortSam.
    pipe_cmd = "picard SortSam {java_opts} {sort_extra} --INPUT /dev/stdin --TMP_DIR {tmpdir} --SORT_ORDER {sort_order} --OUTPUT {snakemake.output[0]}"

else:
    raise ValueError(f"Unexpected value for params.sort ({sort})")

with tempfile.TemporaryDirectory() as tmpdir:
    shell(
        "(bwa mem"
        " -t {snakemake.threads}"
        " {extra}"
        " {index}"
        " {snakemake.input.reads}"
        " | " + pipe_cmd + ") {log}"
    )

__author__ = "Joël Simoneau"
__copyright__ = "Copyright 2019, Joël Simoneau"
__email__ = "[email protected]"
__license__ = "MIT"

from snakemake.shell import shell

# Creating log
log = snakemake.log_fmt_shell(stdout=True, stderr=True)

# Placeholder for optional parameters
extra = snakemake.params.get("extra", "")

# Allowing for multiple FASTA files
fasta = snakemake.input.get("fasta")
assert fasta is not None, "input-> a FASTA-file is required"
fasta = " ".join(fasta) if isinstance(fasta, list) else fasta

shell(
    "kallisto index "  # Tool
    "{extra} "  # Optional parameters
    "--index={snakemake.output.index} "  # Output file
    "{fasta} "  # Input FASTA files
    "{log}"  # Logging
)

__author__ = "Joël Simoneau"
__copyright__ = "Copyright 2019, Joël Simoneau"
__email__ = "[email protected]"
__license__ = "MIT"

from snakemake.shell import shell

# Creating log
log = snakemake.log_fmt_shell(stdout=True, stderr=True)

# Placeholder for optional parameters
extra = snakemake.params.get("extra", "")

# Allowing for multiple FASTQ files
fastq = snakemake.input.get("fastq")
assert fastq is not None, "input-> a FASTQ-file is required"
fastq = " ".join(fastq) if isinstance(fastq, list) else fastq

shell(
    "kallisto quant "  # Tool
    "{extra} "  # Optional parameters
    "--threads={snakemake.threads} "  # Number of threads
    "--index={snakemake.input.index} "  # Input file
    "--output-dir={snakemake.output} "  # Output directory
    "{fastq} "  # Input FASTQ files
    "{log}"  # Logging
)

shell:
    """bioawk -cfastx \
    'BEGIN{{while((getline k <"{input.canonical_ids}")>0)i[k]=1}} \
    {{if(i[$name])print ">"$name"\\n"$seq}}' \
    {input.fasta} > {output}"""

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

library("fgsea")
library(snakemake@params[["bioc_species_pkg"]], character.only = TRUE)

# provides library("tidyverse") and functions load_bioconductor_package() and
# get_prefix_col(), the latter requires snakemake@output[["samples"]] and
# snakemake@params[["covariate"]]
source(snakemake@params[["common_src"]])

gene_sets <- gmtPathways(snakemake@input[["gene_sets"]])
diffexp <- read_tsv(snakemake@input[["diffexp"]]) %>%
                  drop_na(ext_gene) %>%
                  mutate(ext_gene = str_to_upper(ext_gene)) %>%
                  # resolve multiple occurences of the same ext_gene, usually
                  # meaning that several ENSEMBLE genes map to the same gene
                  # symbol -- so we may loose resolution here, as long as gene
                  # symbols are used
                  group_by(ext_gene) %>%
                    filter( qval == min(qval, na.rm = TRUE) ) %>%
                    filter( pval == min(pval, na.rm = TRUE) ) %>%
                    # for the case of min(qval) == 1 and min(pval) == 1, we
                    # need something else to select a unique entry per gene
                    filter( mean_obs == max(mean_obs, na.rm = TRUE) ) %>%
                    mutate(target_id = str_c(target_id, collapse=",")) %>%
                    mutate(ens_gene = str_c(ens_gene, collapse=",")) %>%
                  distinct()

signed_pi <- get_prefix_col("signed_pi_value", colnames(diffexp))

ranked_genes <- diffexp %>%
                  dplyr::select(ext_gene, !!signed_pi) %>%
                  deframe()

# get and write out rank values that are tied -- a way to check up on respective warnings
rank_ties <- enframe(ranked_genes) %>%
               mutate(dup = duplicated(value) | duplicated(value, fromLast = TRUE) ) %>%
               filter(dup == TRUE) %>%
               dplyr::select(ext_gene = name, !!signed_pi := value)
write_tsv(rank_ties, snakemake@output[["rank_ties"]])

print(gene_sets)
print(ranked_genes)

fgsea_res <- fgsea(pathways = gene_sets,
                    stats = ranked_genes,
                    minSize=10,
                    maxSize=700,
                    nproc=snakemake@threads,
                    eps=snakemake@params[["eps"]]
                    ) %>%
                as_tibble()

# Annotation is impossible without any entries, so then just write out empty files
if ( (fgsea_res %>% count() %>% pull(n)) == 0 ) {

    # leadingEdge cannot be a list, and the empty output should at least
    # have the correct columns in the correct order
    unnested <- fgsea_res %>%
                    unnest(leadingEdge) %>%
                    add_column(
                        leading_edge_symbol = NA,
                        leading_edge_ens_gene = NA,
                        leading_edge_entrez_id = NA
                    ) %>%
                    dplyr::relocate(
                        c(
                          leading_edge_symbol,
                          leading_edge_ens_gene,
                          leading_edge_entrez_id,
                          leadingEdge
                        ),
                        .after = last_col()
                    )

    write_tsv(unnested, file = snakemake@output[["enrichment"]])
    write_tsv(unnested, file = snakemake@output[["significant"]])

} else {

    # add further annotation
    annotated <- fgsea_res %>%
                    unnest(leadingEdge) %>%
                    mutate(
                        leading_edge_symbol = str_to_title(leadingEdge),
                        leading_edge_entrez_id = mapIds(leading_edge_symbol, x=get(snakemake@params[["bioc_species_pkg"]]), keytype="SYMBOL", column="ENTREZID"),
                        leading_edge_ens_gene = mapIds(leading_edge_symbol, x=get(snakemake@params[["bioc_species_pkg"]]), keytype="SYMBOL", column="ENSEMBL")
                        ) %>%
                    group_by(pathway) %>%
                    summarise(
                        leadingEdge = str_c(leadingEdge, collapse = ','),
                        leading_edge_symbol = str_c(leading_edge_symbol, collapse = ','),
                        leading_edge_entrez_id = str_c(leading_edge_entrez_id, collapse = ','),
                        leading_edge_ens_gene = str_c(leading_edge_ens_gene, collapse = ',')
                    ) %>%
                    inner_join(fgsea_res %>% dplyr::select(-leadingEdge), by = "pathway") %>%
                    dplyr::relocate(
                        c(
                          leading_edge_symbol,
                          leading_edge_ens_gene,
                          leading_edge_entrez_id,
                          leadingEdge
                        ),
                        .after = last_col()
                    )

    # write out fgsea results for all gene sets
    write_tsv(annotated, file = snakemake@output[["enrichment"]])

    # select significant pathways
    sig_gene_sets <- annotated %>%
                       filter( padj < snakemake@params[["gene_set_fdr"]] )

    # write out fgsea results for gene sets found to be significant
    write_tsv(sig_gene_sets, file = snakemake@output[["significant"]])
}

# select significant pathways
top_pathways <- fgsea_res %>% arrange(padj) %>% head(n=1000) %>% filter(padj <= snakemake@params[["gene_set_fdr"]]) %>% arrange(-NES) %>% pull(pathway)
selected_gene_sets <- gene_sets[top_pathways]

height = .7 * (length(selected_gene_sets) + 2)

# table plot of all gene sets
tg <- plotGseaTable(
            pathways = selected_gene_sets,
            stats = ranked_genes,
            fgseaRes = fgsea_res,
            gseaParam = 1,
            render = FALSE
        )
ggsave(filename = snakemake@output[["plot"]], plot = tg, width = 15, height = height, limitsize=FALSE)

collapsed_pathways <- collapsePathways(fgsea_res %>% arrange(pval) %>% filter(padj <= snakemake@params[["gene_set_fdr"]]), gene_sets, ranked_genes)
main_pathways <- fgsea_res %>% filter(pathway %in% collapsed_pathways$mainPathways) %>% arrange(-NES) %>% pull(pathway)
selected_gene_sets <- gene_sets[main_pathways]
height = .7 * (length(selected_gene_sets) + 2)

# table plot of all gene sets
tg <- plotGseaTable(
            pathways = selected_gene_sets,
            stats = ranked_genes,
            fgseaRes = fgsea_res,
            gseaParam = 1,
            render = FALSE
        )
ggsave(filename = snakemake@output[["plot_collapsed"]], plot = tg, width = 15, height = height, limitsize=FALSE)

import pandas as pd
import numpy as np
import pysam
import sys

sys.stderr = open(snakemake.log[0], "w")

bam_file = snakemake.input["canonical_mapped_bam"]
sample_name = bam_file.split(".canonical.mapped.sorted.bam")[0]
# Bam file reading
bam_file = pysam.AlignmentFile(snakemake.input["canonical_mapped_bam"], "rb")
bam_header = bam_file.header.to_dict()
trans_length_data = pd.DataFrame(bam_header.get("SQ"))
trans_length_data.rename(columns={"SN": "Transcript_ID"}, inplace=True)

# Aligned text file reading
align_bam_txt = pd.read_csv(
    snakemake.input["canonical_mapped_pos"],
    sep="\t",
    names=["read_name", "Transcript_ID", "Start", "read", "Quality"],
)
align_bam_txt["Strand"] = align_bam_txt["Transcript_ID"].str.split("_", 1).str[1]
align_bam_txt["Transcript"] = align_bam_txt["Transcript_ID"].str.split("_", 1).str[0]
merge_data = align_bam_txt.merge(trans_length_data, on="Transcript_ID")

# reads aligned to forward strand
forward_strand = merge_data.loc[merge_data["Strand"] == "1"]
forward_strand[sample_name + "_forward_strand"] = (
    forward_strand["LN"] - forward_strand["Start"]
)
aligned_reads = forward_strand.loc[
    forward_strand.groupby(["read_name", "read"])[
        sample_name + "_forward_strand"
    ].idxmin()
]

# reads aligned to reverse strand
reverse_strand = merge_data.loc[merge_data["Strand"] == "-1"]
read_min = reverse_strand.loc[
    reverse_strand.groupby(["read_name", "read"])["Start"].idxmin()
]
aligned_reads = pd.concat([aligned_reads, read_min])
aligned_reads.to_csv(
    snakemake.output["canonical_mapped_3prime_pos"],
    columns=["read_name"],
    sep="\t",
    index=False,
)

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


library(biomaRt)
library(dplyr)
library(tidyverse)

ensembl <- useEnsembl(
    biomart = "genes", dataset = "hsapiens_gene_ensembl",
    version = snakemake@params["release"]
)
get_transcripts_ids <-
    getBM(
        attributes = c(
            "ensembl_transcript_id_version",
            "transcript_is_canonical", "transcript_mane_select", "chromosome_name"
        ),
        mart = ensembl
    )
# transcript_mane_select- manually curated primary transcript as it has been encoded in NCBI
canonical_ids <- get_transcripts_ids %>%
    select(
        ensembl_transcript_id_version, transcript_is_canonical,
        transcript_mane_select, chromosome_name
    ) %>%
    filter(!str_detect(chromosome_name, "patch|PATCH")) %>%
    filter(str_detect(transcript_mane_select, "")) %>%
    subset(transcript_is_canonical == 1)
write.csv(canonical_ids$ensembl_transcript_id_version,
    file = snakemake@output[[1]], quote = FALSE, row.names = FALSE
)

import json
import pysam

max_read_length = max(
    len(rec.sequence) for f in snakemake.input for rec in pysam.FastxFile(f)
)

with open(snakemake.output[0], "w") as out:
    json.dump(max_read_length, out)

import pandas as pd
import numpy as np
import pysam
from scipy import stats
import sys
import json
import csv


sys.stderr = open(snakemake.log[0], "w")

# Get the read-length
f = open(snakemake.input["read_length"])
read_length = json.load(f)
f.close()

# Get the transcript IDs if executing for individual transcripts or else for all transcripts
transcript_ids = snakemake.params["each_transcript"]

# Define data-frame for each strand
fwrd_allsamp_hist = pd.DataFrame([])
fwrd_allsamp_hist_trim = pd.DataFrame([])
fwrd_allsamp_hist_fil = pd.DataFrame([])
rev_allsamp_hist = pd.DataFrame([])
rev_allsamp_hist_trim = pd.DataFrame([])
rev_allsamp_hist_fil = pd.DataFrame([])

# Get the sample names
sample_name = snakemake.params["samples"]

# Bam file reading
bam_file = pysam.AlignmentFile(snakemake.input["samtools_sort"], "rb")
bam_header = bam_file.header.to_dict()
trans_length_data = pd.DataFrame(bam_header.get("SQ"))
trans_length_data.rename(columns={"SN": "Transcript_ID"}, inplace=True)

# Aligned text file reading
align_bam_txt = pd.read_csv(
    snakemake.input["aligned_file"],
    sep="\t",
    names=["read_Name", "Transcript_ID", "Start", "reads", "Quality"],
)
align_bam_txt["Strand"] = align_bam_txt["Transcript_ID"].str.split("_", 1).str[1]
align_bam_txt["Transcript"] = align_bam_txt["Transcript_ID"].str.split("_", 1).str[0]


# Both transcript len and start postion are merged based on same transcript ID
merge_data = align_bam_txt.merge(trans_length_data, on="Transcript_ID")

# Forward strand
forward_strand = merge_data.loc[merge_data["Strand"] == "1"]

# Each read postion is calcuated
forward_strand[sample_name + "_forward_strand"] = (
    forward_strand["LN"] - forward_strand["Start"]
)
aligned_reads = forward_strand.loc[
    forward_strand.groupby(["read_Name", "reads"])[
        sample_name + "_forward_strand"
    ].idxmin()
]
if aligned_reads["Transcript"].str.contains(transcript_ids).any():
    # Get aligned read postion of the given transcript
    fwrd_filtered_transcript_data = aligned_reads.query("Transcript == @transcript_ids")
    Freq_fwrd_fil, bins_fwrd_fil = np.histogram(
        fwrd_filtered_transcript_data[sample_name + "_forward_strand"],
        bins=read_length,
        range=[0, max(fwrd_filtered_transcript_data["LN"])],
    )
    hist_fwrd_fil = pd.DataFrame(
        {
            "sample_Name": sample_name,
            "Freq_forward": Freq_fwrd_fil,
            "bins_foward": bins_fwrd_fil[:-1],
        }
    )
    fwrd_allsamp_hist_fil = pd.concat([fwrd_allsamp_hist_fil, hist_fwrd_fil])
elif transcript_ids == "all":
    # Values added to corresponding bins
    Freq_fwrd, bins_fwrd = np.histogram(
        aligned_reads[sample_name + "_forward_strand"],
        bins=read_length,
        range=[0, max(aligned_reads["LN"])],
    )
    Freq_fwrd_trim, bins_fwrd_trim = np.histogram(
        forward_strand[sample_name + "_forward_strand"],
        bins=read_length,
        range=[0, 20000],
    )
    # Dataframe created for bins
    hist_fwrd = pd.DataFrame(
        {
            "sample_Name": sample_name,
            "Freq_forward": Freq_fwrd,
            "bins_foward": bins_fwrd[:-1],
        }
    )
    hist_fwrd_trim = pd.DataFrame(
        {
            "sample_Name": sample_name,
            "Freq_forward": Freq_fwrd_trim,
            "bins_foward": bins_fwrd_trim[:-1],
        }
    )
    # Each sample dataframe is concatinated
    fwrd_allsamp_hist = pd.concat([fwrd_allsamp_hist, hist_fwrd])
    fwrd_allsamp_hist_trim = pd.concat([fwrd_allsamp_hist_trim, hist_fwrd_trim])
    fwrd_allsamp_hist.to_csv(
        snakemake.output["fwrd_allsamp_hist"], sep="\t", index=False
    )
    fwrd_allsamp_hist_trim.to_csv(
        snakemake.output["fwrd_allsamp_hist_trim"], sep="\t", index=False
    )

# Reverse strand
reverse_strand = merge_data.loc[merge_data["Strand"] == "-1"]

# Minimum aligned start postion is taken
read_min = reverse_strand.loc[
    reverse_strand.groupby(["read_Name", "reads"])["Start"].idxmin()
]

if read_min["Transcript"].str.contains(transcript_ids).any():
    # Get aligned read postion of the given transcript
    rev_filtered_transcript_data = read_min.query("Transcript == @transcript_ids")
    Freq_rev_fil, bins_rev_fil = np.histogram(
        rev_filtered_transcript_data["Start"],
        bins=read_length,
        range=[0, max(rev_filtered_transcript_data["LN"])],
    )
    hist_rev_fil = pd.DataFrame(
        {
            "sample_Name": sample_name,
            "Freq_rev": Freq_rev_fil,
            "bins_rev": bins_rev_fil[:-1],
        }
    )
    rev_allsamp_hist_fil = pd.concat([rev_allsamp_hist_fil, hist_rev_fil])

elif transcript_ids == "all":
    # Values added to corresponding bins
    Freq_rev, bins_rev = np.histogram(
        read_min["Start"], bins=read_length, range=[0, max(read_min["LN"])]
    )
    Freq_rev_trim, bins_rev_trim = np.histogram(
        read_min["Start"], bins=read_length, range=[0, 20000]
    )

    # Dataframe created for bins
    hist_rev = pd.DataFrame(
        {"sample_Name": sample_name, "Freq_rev": Freq_rev, "bins_rev": bins_rev[:-1]}
    )
    hist_rev_trim = pd.DataFrame(
        {
            "sample_Name": sample_name,
            "Freq_rev": Freq_rev_trim,
            "bins_rev": bins_rev_trim[:-1],
        }
    )

    # Each sample dataframe is concatinated
    rev_allsamp_hist = pd.concat([rev_allsamp_hist, hist_rev])
    rev_allsamp_hist_trim = pd.concat([rev_allsamp_hist_trim, hist_rev_trim])
    rev_allsamp_hist.to_csv(snakemake.output["rev_allsamp_hist"], sep="\t", index=False)
    rev_allsamp_hist_trim.to_csv(
        snakemake.output["rev_allsamp_hist_trim"], sep="\t", index=False
    )


# Write bins to each file
if transcript_ids != "all":
    if fwrd_allsamp_hist_fil.empty:
        with open(snakemake.output["fwrd_allsamp_hist_fil"], "w") as fwrd_csvfile:
            print("no reads aligned", file=fwrd_csvfile)
    else:
        with open(snakemake.output["fwrd_allsamp_hist_fil"], "w") as fwrd_csvfile:
            fwrd_allsamp_hist_fil.to_csv(fwrd_csvfile, sep="\t", index=False)
    if rev_allsamp_hist_fil.empty:
        with open(snakemake.output["rev_allsamp_hist_fil"], "w") as rev_csvfile:
            print("no reads aligned", file=rev_csvfile)
    else:
        with open(snakemake.output["rev_allsamp_hist_fil"], "w") as rev_csvfile:
            rev_allsamp_hist_fil.to_csv(rev_csvfile, sep="\t", index=False)