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


from snakemake.shell import shell

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

n = len(snakemake.input.sample)
assert (
    n == 1 or n == 2
), "input->sample must have 1 (single-end) or 2 (paired-end) elements."

if n == 1:
    reads = "-U {}".format(*snakemake.input.sample)
else:
    reads = "-1 {} -2 {}".format(*snakemake.input.sample)

shell(
    "(bowtie2 --threads {snakemake.threads} {extra} "
    "-x {snakemake.params.index} {reads} "
    "| samtools view -Sbh -o {snakemake.output[0]} -) {log}"
)

__author__ = "Daniel Standage"
__copyright__ = "Copyright 2020, Daniel Standage"
__email__ = "[email protected]"
__license__ = "MIT"


from snakemake.shell import shell

extra = snakemake.params.get("extra", "")
log = snakemake.log_fmt_shell(stdout=True, stderr=True)
indexbase = snakemake.output[0].replace(".1.bt2", "")
shell(
    "bowtie2-build --threads {snakemake.threads} {snakemake.params.extra} "
    "{snakemake.input.reference} {indexbase}"
)

__author__ = "Jan Forster"
__copyright__ = "Copyright 2021, Jan Forster"
__email__ = "[email protected]"
__license__ = "MIT"


import os

from snakemake.shell import shell

in_file = snakemake.input[0]
extra = snakemake.params.get("extra", "")
log = snakemake.log_fmt_shell(stdout=False, stderr=True)

if in_file.endswith(".sam") and ("-S" not in extra or "--sam-input" not in extra):
    extra += " --sam-input"

shell(
    "sambamba view {extra} -t {snakemake.threads} "
    "{snakemake.input[0]} > {snakemake.output[0]} "
    "{log}"
)

shell:
	"bamCoverage --bam {input.bam} -o {output} -p {threads} {params.extra}"

shell:
	"""
	seqkit replace -p "(.+)" -r "spikeIn_\$1" -o resources/tmp_spikeIn.fasta.gz {input.spikeIn} 2> {log}
	cat {input.ref} resources/tmp_spikeIn.fasta.gz > {output} 2>> {log}
	rm resources/tmp_spikeIn.fasta.gz
	"""

library(tidyverse)

# import data ------------------------------------------------------------------
# read in all relevant stat files
stat_files <- snakemake@input[["mapping_stats"]]
names(stat_files) <- stat_files

mapping_stats <- stat_files %>%
  map_dfr(read_tsv, .id = "source", col_names = c("chromosome", "chromosome_size", "mapped_reads", "unmapped_reads")) %>%
  mutate(sample_name = gsub("_unireads.idxstats", "", basename(source)))


# add column marking reference or spike-in chromosomes
mapping_stats <- mapping_stats %>%
  mutate(reference_genome = case_when(str_detect(chromosome, "spikeIn")  ~ "spikeIn",
                                      !str_detect(chromosome, "spikeIn")  ~ "reference")) 


# compute % reads mapping to reference and spike-in genomes
mapping_stats <- mapping_stats %>% 
  group_by(sample_name, reference_genome) %>% 
  summarise(total_reads = sum(mapped_reads)) %>%
  pivot_wider(names_from = reference_genome, values_from = total_reads) %>%
  mutate(total_mapped_reads = reference + spikeIn) %>%
  mutate(percent_reference = reference / total_mapped_reads * 100, percent_spikeIn = spikeIn / total_mapped_reads * 100)


# read in unit table
sample_table <- snakemake@config[["units"]] %>% 
  read_tsv() %>% 

  # remove redundant technical replicates from sample_table
  distinct(sample_name, .keep_all = TRUE)


# mark samples as inputs or IPs
sample_table <- sample_table %>% 
  mutate(chip_sample_type = case_when(
    (sample_name %in% .$input) ~ "input",
    !(sample_name %in% .$input) ~ "IP"
  )) %>% 
  mutate(has_input = !is.na(input))


# create new column indicating input/IP pairs
tmp <- sample_table %>% select(input, sample_name) %>% dplyr::rename(IP = sample_name, sample_name = input)
sample_table <- sample_table %>% 
  left_join(tmp, by = "sample_name") %>% 
  mutate(IP_group = case_when(
    is.na(IP) ~ sample_name,
    !is.na(IP) ~ IP
  ))




# prepare sample_table and mapping_stats tables for join
individual_mapping_stats <- mapping_stats %>% 
  select(sample_name, percent_reference, percent_spikeIn)

# join sample_table and mapping stats
individual_scaling_factors <- sample_table %>%
  left_join(individual_mapping_stats, by = "sample_name")


# compute normalization factors for individual samples -------------------------
sample_has_input <- individual_scaling_factors %>% 
  filter(chip_sample_type == "IP") %>% 
  select(has_input) %>% 
  pull()

if (all(sample_has_input)) {

  individual_scaling_factors <- individual_scaling_factors %>% 
    pivot_wider(names_from = chip_sample_type, values_from = percent_spikeIn, id_cols = IP_group) %>% 
    mutate(enrichment = IP / input) %>%
    mutate(scaling_factor = 1 / enrichment) %>%
    mutate(norm_scaling_factor = scaling_factor / max(scaling_factor))

} else {
  warning("not all IPs have corresponding input sample, calculating scaling factors based on %spikeIn for IP samples")
  individual_scaling_factors <- individual_scaling_factors %>% 
    pivot_wider(names_from = chip_sample_type, values_from = percent_spikeIn, id_cols = IP_group) %>% 
    mutate(scaling_factor = 1 / IP) %>%
    mutate(norm_scaling_factor = scaling_factor / max(scaling_factor))
}

individual_scaling_factors %>% 
  write_tsv(snakemake@output[[1]])

# prepare data for merged samples -----------------------------
merged_mapping_stats <- mapping_stats %>% 
  select(sample_name, reference, spikeIn)

# combine replicate reads
merged_mapping_stats <- sample_table %>% 
  left_join(merged_mapping_stats, by = "sample_name") %>% 
  group_by(sample_group) %>% 
  summarise(spikeIn = sum(spikeIn), reference = sum(reference))




# create new column indicating input/IP pairs
tmp <- sample_table %>% select(input, sample_group) %>% dplyr::rename(IP = sample_group, sample_name = input)
sample_table <- sample_table %>% 
  select(-IP, -IP_group) %>% 
  left_join(tmp, by = "sample_name") %>% 
  mutate(IP_group = case_when(
    is.na(IP) ~ sample_group,
    !is.na(IP) ~ IP
  ))



# join sample_table and mapping stats for merged reads
merged_scaling_factors <- merged_mapping_stats %>% 
  left_join(select(sample_table, sample_group, chip_sample_type,IP_group, has_input), by = "sample_group") %>% 
  distinct(.keep_all = TRUE) %>%
  mutate(total_mapped_reads = reference + spikeIn) %>%
  mutate(percent_reference = reference / total_mapped_reads * 100, percent_spikeIn = spikeIn / total_mapped_reads * 100)

# compute normalization factors for merged samples -------------------------
sample_has_input <- merged_scaling_factors %>% 
  filter(chip_sample_type == "IP") %>% 
  select(has_input) %>% 
  pull()
if (all(sample_has_input)) {

  merged_scaling_factors <- merged_scaling_factors %>% 
    pivot_wider(names_from = chip_sample_type, values_from = percent_spikeIn, id_cols = IP_group) %>% 
    mutate(enrichment = IP / input) %>%
    mutate(scaling_factor = 1 / enrichment) %>%
    mutate(norm_scaling_factor = scaling_factor / max(scaling_factor))

} else {
  warning("not all IPs have corresponding input sample, calculating scaling factors based on %spikeIn for IP samples")
  merged_scaling_factors <- merged_scaling_factors %>% 
    pivot_wider(names_from = chip_sample_type, values_from = percent_spikeIn, id_cols = IP_group) %>% 
    mutate(scaling_factor = 1 / IP) %>%
    mutate(norm_scaling_factor = scaling_factor / max(scaling_factor))
}

merged_scaling_factors %>% 
  write_tsv(snakemake@output[[2]])

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

import os
import tempfile
from snakemake.shell import shell
from snakemake_wrapper_utils.snakemake import get_mem


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


# Parse memory
mem_mb = get_mem(snakemake, "MiB")


# Outdir
outdir = os.path.dirname(snakemake.output[0])
if outdir:
    outdir = f"--outdir {outdir}"


# Output compression
compress = ""
mem = f"-m{mem_mb}" if mem_mb else ""

for output in snakemake.output:
    out_name, out_ext = os.path.splitext(output)
    if out_ext == ".gz":
        compress += f"pigz -p {snakemake.threads} {out_name}; "
    elif out_ext == ".bz2":
        compress += f"pbzip2 -p{snakemake.threads} {mem} {out_name}; "


with tempfile.TemporaryDirectory() as tmpdir:
    mem = f"--mem {mem_mb}M" if mem_mb else ""

    shell(
        "(fasterq-dump --temp {tmpdir} --threads {snakemake.threads} {mem} "
        "{extra} {outdir} {snakemake.wildcards.accession}; "
        "{compress}"
        ") {log}"
    )

suppressPackageStartupMessages(library(tidyverse))
suppressPackageStartupMessages(library(rtracklayer))
suppressPackageStartupMessages(library(GenomicRanges))

# define functions -------------------------------------------------------------
# function to merge multiple sets of ChIP peaks into a non-overlapping master set
# supports multiple methods
# peaks should be a GRangesList
# method should be either "overlap" or "merge": 
# overlap method will determine overlaps among all peak sets and combine non-overlapping peaks
# merge method will combine overlapping peaks into a single peak
combine_peaks <- function(peaks, method = "overlap", min_overlap = 1L) {
  require(GenomicRanges)

  # check arguments
  if (length(peaks) < 2) {
    stop("Only one set of peaks provided. For merging peaks, provide multiple peak sets.")
  }

  if (!inherits(peaks, "GRangesList") ) {
    stop("one or more peak set is not a Granges object. peaks must be a GRangesList")
  }

  if (!any(method == c("overlap", "merge"))) {
    stop("Invalid method provided. method should be 'overlap' or 'merge' ")
  }

  if (!is_integer(min_overlap)) {
    stop("min_overlap must be an integer")
  }


  # merge peaks using "overlap" method
  if (method == "overlap") {
    combined_peaks <- peaks[[1]]
    for (i in 2:length(peaks)) {
      i_set <- peaks[[i]]
      combined_peaks <- c(combined_peaks,
                          subsetByOverlaps(i_set, combined_peaks, minoverlap = min_overlap, invert = TRUE))
    }

  }
  # merge peaks using "merge" method
  if (method == "merge") {
    combined_peaks <- GenomicRanges::reduce(unlist(peaks))

  }
  mcols(combined_peaks) <- NULL
  return(combined_peaks)

}

# function to build an overlap table
# takes a Granges list object
# combines all nonoverlapping peaks from all samples, then builds a logical table indicating which peaks from each sample overlap each peak on the master list
peak_overlap_table <- function(peaks, method = "overlap", min_overlap = 1L, combine_peaks = TRUE) {

  # check arguments
  if (length(peaks) < 2) {
    stop("Only one set of peaks provided. For merging peaks, provide multiple peak sets.")
  }

  if (!inherits(peaks, "GRangesList") ) {
    stop("one or more peak set is not a Granges object. peaks must be a GRangesList")
  }

  if (!any(method == c("overlap", "merge"))) {
    stop("Invalid method provided. method should be 'overlap' or 'merge' ")
  }

  if (!is_integer(min_overlap)) {
    stop("min_overlap must be an integer")
  }


  # set names
  if(is.null(names(peaks))) {
    names(peaks) <- paste0("sample_", seq((peaks)))
  }


  if (combine_peaks) {
    # get a master set of all nonoverlapping peaks from all peak sets
    all_peaks_gr <- combine_peaks(peaks, method = method, min_overlap = min_overlap)

  } else {
    all_peaks_gr <- peaks[[1]]
  }
  all_peaks_df <- as.data.frame(all_peaks_gr) %>%
    dplyr::select(1:5)

  # build a logical overlap table indicating which peaks were detected in which samples
  for (i in 1:length(peaks)) {
    sample_name <- names(peaks)[i]
    i_peaks <- peaks[[i]]

    all_peaks_df[,sample_name] <- FALSE
    overlaps <- findOverlaps(all_peaks_gr, i_peaks, minoverlap = min_overlap)@from
    all_peaks_df[overlaps,sample_name] <- TRUE
  }
  return(all_peaks_df)
}

# read peaks into overlap table ------------------------------------------------
merged_peak_file <- snakemake@input[["merged_peaks"]]
names(merged_peak_file) <- "merged_peaks"

replicate_peak_files <- snakemake@input[["ind_peaks"]]
names(replicate_peak_files) <- paste0("rep", seq(replicate_peak_files), "_peaks")

peak_files <- c(merged_peak_file, replicate_peak_files)

overlap_table <- peak_files %>% 
  map(rtracklayer::import) %>% 
  GRangesList() %>% 
  peak_overlap_table(min_overlap = as.integer(snakemake@params[["min_overlap"]]))

# filter to include only peaks detected in multiple replicates -----------------
filtered_peaks <- overlap_table %>% 
  mutate(reps_w_peak = rowSums(.[7:ncol(.)])) %>% 
  filter(merged_peaks, reps_w_peak >= snakemake@params[["replicate_threshold"]]) %>% 
  makeGRangesFromDataFrame()

# export filtered peaks to file  -----------------------------------------------
if (snakemake@wildcards[["peak_type"]] == "narrow") {
  keep_cols <- c(1:3, 6:7, 5, 8:11)
} else if (snakemake@wildcards[["peak_type"]] == "broad") {
  keep_cols <- c(1:3, 6:7, 5, 8:10)
} else {
  stop("Peak format should be 'narrow' or 'broad'. Check input file format")
}

output <- merged_peak_file %>% 
  rtracklayer::import() %>% 
  subsetByOverlaps(filtered_peaks) %>% 
  as.data.frame() %>% 
  select(all_of(keep_cols)) %>% 
  mutate(strand = ".")

write_tsv(output, snakemake@output[[1]], col_names = FALSE)