#!python

"""
This tool will convert one or more SAM/BAM file/s into a pair of stranded
BigWig files containing the 5' counts. The input files can be either stored
locally or remote URLs. When multiple files are provided, the reads are
concatenated. Only fully mapped reads are kept.
"""

import gzip
import numpy
import argparse
import collections
import multiprocessing

import pysam
import pyBigWig
import pyfaidx

from tqdm import tqdm
from joblib import Parallel
from joblib import delayed

### Load the user-specified arguments

# Base arguments

parser = argparse.ArgumentParser(
    prog='bam2bw',
    description='This tool will convert BAM files to bigwig files without an intermediate.')

parser.add_argument('filename', nargs='+',
	help="""The SAM/BAM or tsv/tsv.gz file to be processed.""")         
parser.add_argument('-s', '--sizes', required=True, 
	help="""A chrom_sizes, .fai, or FASTA file. Only the first two columns of
	a chrom_sizes/.fai file are read. A compressed FASTA must be BGZF, not
	gzip.""")

# Data properties

parser.add_argument('-u', '--unstranded', action='store_true',
    help="Have only one, unstranded, output.")

end_group = parser.add_mutually_exclusive_group()
end_group.add_argument('-f', '--fragments', action='store_true', default=False,
	help="The data is fragments and so both ends should be recorded.")
end_group.add_argument('-3p', '--three_prime', action='store_true', default=False,
	help="Record the 3' end of each read instead of the 5' end.")

parser.add_argument('-ps', '--pos_shift', default=0, type=int,
	help="""A shift to apply to positive strand reads.""")
parser.add_argument('-ns', '--neg_shift', default=0, type=int,
	help="""A shift to apply to negative strand reads.""")

parser.add_argument('-mp', '--mate_pairs', action='store_true', default=False,
	help="""Treat paired-end reads as a single RNA/fragment tag instead of
	counting each mate independently: buffer reads by name, and once both
	mates of a pair are seen, record one jointly-determined position (see
	--rna5 and --opposite_strand). BAM/SAM input only.""")
parser.add_argument('--rna5', choices=('read1', 'read2'), default='read1',
	help="""Which mate carries the 5' end of the RNA/fragment. The other
	mate's own 5' end is used as the RNA's 3' end. Only used with
	--mate_pairs.""")
parser.add_argument('--opposite_strand', action='store_true', default=False,
	help="""Report the strand of the mate opposite the one chosen by
	--rna5, instead of that mate's own strand. Only used with
	--mate_pairs.""")

parser.add_argument('-sf', '--scale_factor', default=1, type=float,
	help="""A scaling factor to multiply each position by.""")
parser.add_argument('-r', '--read_depth', default=False, action='store_true',
	help="""Whether to divide through by total (pre-scaled) read depth.""")

# Misc arguments

parser.add_argument('-p', '--parallel', default=1, type=int,
	help="The number of jobs to use, max of one per input file.")
parser.add_argument('-n', '--name', required=True)
parser.add_argument('-z', '--zooms', default=0, type=int,
    help="""The number of zooms to store in the bigwig.""")
parser.add_argument('-v', '--verbose', action='store_true')
args = parser.parse_args()


###

# First check that the files are the right format

acceptable_formats = '.bam', '.sam', '.tsv', '.tsv.gz', '.bed', '.bed.gz'

for filename in args.filename:
	if not filename.endswith(acceptable_formats):
		raise ValueError("Filenames must end in one of {}.".format(', '.join(acceptable_formats)))

if args.mate_pairs:
	for filename in args.filename:
		if not filename.endswith(('.bam', '.sam')):
			raise ValueError("--mate_pairs only supports BAM/SAM input files.")

# Here, we are determining the chromosomes and their sizes. This is necessary
# for creating the bigWig header(s) and for figuring out which reads to filter
# out (those that do not map to the provided chromosomes).
#
# Because we allow you to provide either a two-column chrom_sizes file or a
# FASTA file, we need code to handle the two situations. We use pyfaidx to
# quickly process the FASTA file so that we do not have to scan through the
# entire thing just to get the sizes.

def is_fasta(path):
    base = path.rstrip(".gz")
    return base.endswith((".fa", ".fasta", ".fna", ".fas"))

chrom_sizes = []

# If provided a FASTA file, read the lengths of the sequences. pyfaidx reads
# BGZF but not plain gzip, and its own error does not say what to do about it.
if is_fasta(args.sizes):
	try:
		fa = pyfaidx.Fasta(args.sizes)
	except pyfaidx.UnsupportedCompressionFormat:
		raise ValueError(("{} is compressed with gzip. Only BGZF is supported "
			"for compressed FASTA files -- recompress it with `bgzip`, or pass "
			"a chrom_sizes file instead.").format(args.sizes))

	for chrom, seq in fa.items():
		chrom_sizes.append((chrom, len(seq)))

# If provided a chrom_sizes file, just use the provided lengths. Only the first
# two columns are read, so a samtools .fai index works here as well as a
# two-column chrom_sizes file, and blank lines and # comments are skipped.
else:
	with open(args.sizes, "r") as size_file:
		for i, line in enumerate(size_file):
			line = line.strip()
			if len(line) == 0 or line.startswith("#"):
				continue

			fields = line.split()
			if len(fields) < 2:
				raise ValueError(("{}, line {}: expected a chromosome name and "
					"a length, got {}.").format(args.sizes, i + 1, repr(line)))

			try:
				size = int(fields[1])
			except ValueError:
				raise ValueError(("{}, line {}: expected an integer length for "
					"{}, got {}.").format(args.sizes, i + 1, fields[0],
						repr(fields[1])))

			chrom_sizes.append((fields[0], size))

###

# This is the main loop that goes through the reads and records them in
# one or two dictionaries (depending on if the data is stranded). The
# processing of BAM/SAM files relies on pysam, whereas tsv/tsv.gz files use
# basic file iteration. The processing of reads in both cases is largely the
# same except that, usually, the -f flag will be passed in for .tsv/.tsv.gz
# files because those come from fragments from ATAC-seq-like experiments,
# whereas BAM/SAM files are usually just reads.

# First, a function that reads a single file and returns one or two
# dictionaries of reads. This can be parallelized across files. 

MateInfo = collections.namedtuple('MateInfo', ['start', 'end', 'is_forward'])

def record_positions(pos_reads, neg_reads, chrom, five_prime, three_prime, is_forward):
	reads = pos_reads if is_forward else neg_reads

	reads[chrom][three_prime if args.three_prime else five_prime] += 1

	if args.fragments:
		reads[chrom][five_prime if args.three_prime else three_prime] += 1

def record_mate_pair(pos_reads, neg_reads, chrom, anchor, other):
	# The RNA's 5' end is the anchor mate's own 5' end (per --rna5); the RNA's
	# 3' end is the other mate's own 5' end, since that mate started
	# sequencing from the opposite end of the fragment.
	five_prime = anchor.start if anchor.is_forward else anchor.end - 1
	three_prime = other.start if other.is_forward else other.end - 1
	is_forward = other.is_forward if args.opposite_strand else anchor.is_forward

	record_positions(pos_reads, neg_reads, chrom, five_prime, three_prime, is_forward)

def extract_reads(args, chrom_sizes, idx):
	missing_chroms = set()
	pos_reads = {}
	neg_reads = pos_reads if args.unstranded else {}

	# Create dictionaries for each chrom, regardless
	for chrom, _ in chrom_sizes:
		pos_reads[chrom] = collections.defaultdict(int)
		neg_reads[chrom] = collections.defaultdict(int) # Redundant if unstranded

	# Extract the file being considered	
	filename = args.filename[idx]
	name = filename.split("/")[-1]
	if filename.endswith((".bam", ".sam")):
		# pysam has to be told which of the two it is -- "rb" is binary BAM
		# and "r" is text SAM. Neither needs an index because the loop below
		# fetches until_eof.
		mode = "rb" if filename.endswith(".bam") else "r"
		bam = pysam.AlignmentFile(filename, mode)
		pending_mates = {}

		# These are read once per read in the loop below, which is hot enough
		# that the attribute lookups show up in a profile.
		mate_pairs = args.mate_pairs
		fragments = args.fragments
		use_three_prime = args.three_prime

		for read in tqdm(bam.fetch(until_eof=True), disable=not args.verbose, position=idx, desc=name):
			if read.is_unmapped:
				continue

			# Check whether the chrom is in the allowable chroms. Otherwise, discard.

			chrom = read.reference_name
			if chrom not in pos_reads:
				if chrom not in missing_chroms:
					missing_chroms.add(chrom)
					if args.verbose:
						tqdm.write("{} encountered in input but not in FASTA/chrom sizes.".format(
							chrom))

				continue

			# A mapped record whose CIGAR is '*' has no reference end, so
			# there is no position to record for the far end of the read.
			# reference_end is computed from the CIGAR on every access, so it
			# is read once here and reused below.
			reference_end = read.reference_end
			if reference_end is None:
				raise ValueError(("{}: read {} is mapped to {} but has no "
					"CIGAR, so its reference end is unknown.").format(name,
						read.query_name, chrom))

			start = read.reference_start + args.pos_shift
			end = reference_end + args.neg_shift

			if mate_pairs:
				# Only consider one alignment per mate so that the two halves
				# of a pair can be matched up unambiguously by read name.
				if not read.is_proper_pair or read.is_secondary or read.is_supplementary:
					continue

				partner = pending_mates.pop(read.query_name, None)
				mate_info = MateInfo(start, end, read.is_forward)

				if partner is None or partner[0] == read.is_read1:
					# Either the first mate seen for this name, or an
					# unexpected duplicate record for the same mate -- in the
					# latter case the stale entry is dropped in favor of this one.
					pending_mates[read.query_name] = (read.is_read1, mate_info)
					continue

				_, partner_info = partner
				read1_info = mate_info if read.is_read1 else partner_info
				read2_info = partner_info if read.is_read1 else mate_info

				if args.rna5 == 'read1':
					anchor, other = read1_info, read2_info
				else:
					anchor, other = read2_info, read1_info

				record_mate_pair(pos_reads, neg_reads, chrom, anchor, other)
			else:
				# Here, we need to deal with two related issues.
				#
				#    (1) Does the read map to the fwd or rev strand?
				#    (2) Are we mapping the start or the strand and the end (fragments)?
				#
				# Accordingly, we first check to see the strand the read is on and take the
				# start of the read (start for fwd, end-1 for bwd reads). Then, we need to
				# check whether we want both starts and ends and record both if so. This
				# strategy works even if the underlying data is not stranded because
				# pos_reads and neg_reads are the same dictionary in that case.

				if read.is_forward:
					five_prime, three_prime = start, end - 1
					reads = pos_reads
				else:
					five_prime, three_prime = end - 1, start
					reads = neg_reads

				reads[chrom][three_prime if use_three_prime else five_prime] += 1
				if fragments:
					reads[chrom][five_prime if use_three_prime else three_prime] += 1

		bam.close()
	
	elif filename[-4:] in ('.tsv', '.bed') or filename[-7:] in ('.tsv.gz', '.bed.gz'):
		# Open the file using the correct opener -- the standard one if the file is
		# not compressed, otherwise the gzip opener if gzipped.
		
		if filename[-4:] in ('.tsv', '.bed'):
			f = open(filename, "r")
		elif filename[-7:] in ('.tsv.gz', '.bed.gz'):
			f = gzip.open(filename, "rt")
		
		# Here, we process the entries in a similar manner to using pysam except that
		# we assume the coordinates are all fwd strand. We do not explicitly assume
		# that we want both the start and the end of the entry, which is controlled
		# using the -f flag, but we do not handle strandedness here.

		for i, line in enumerate(tqdm(f, disable=not args.verbose, position=idx,
			desc=name)):
			fields = line.split()

			# Blank lines and # comments are skipped, matching the chrom_sizes
			# reader above. A 10x CellRanger fragments file opens with a #
			# header, so rejecting it would fail on the most common input.
			# Testing fields[0] rather than the raw line costs no allocation on
			# this per-read path and still catches a # behind leading
			# whitespace; split() never yields an empty field, so [0][0] is safe.

			if len(fields) == 0 or fields[0][0] == '#':
				continue

			if len(fields) < 3:
				raise ValueError(("{}, line {}: expected at least three columns "
					"holding a chromosome, a start and an end, got {}.").format(
						filename, i + 1, repr(line.strip())))

			# Check whether the chrom is in the allowable chroms. Otherwise, discard.

			chrom, start, end = fields[:3]
			if chrom not in pos_reads:
				if chrom not in missing_chroms:
					missing_chroms.add(chrom)
					if args.verbose:
						tqdm.write("{} encountered in input but not in FASTA/chrom sizes.".format(
							chrom))

				continue

			try:
				start = int(float(start)) + args.pos_shift
				end = int(float(end)) + args.neg_shift
			except ValueError:
				raise ValueError(("{}, line {}: expected numeric coordinates, "
					"got {} and {}.").format(filename, i + 1, repr(start),
						repr(end)))

			pos_reads[chrom][end-1 if args.three_prime else start] += 1
			if args.fragments:
				pos_reads[chrom][start if args.three_prime else end-1] += 1

	return pos_reads, neg_reads

# Share a single lock across the worker processes so that each file's tqdm
# progress bar renders on its own line (via position=idx) instead of the
# processes clobbering each other's cursor movements on stdout. The lock is
# created before Parallel so it is inherited by the forked workers.
tqdm.set_lock(multiprocessing.RLock())

f = delayed(extract_reads)
reads = Parallel(n_jobs=args.parallel, backend='multiprocessing')(
	f(args, chrom_sizes, i) for i in range(len(args.filename))
)

# The persistent per-file bars leave the cursor part-way up the screen, so
# drop below them before printing anything else.
if args.verbose:
	print("\n" * len(args.filename))

### Collect the reads into a single object.

pos_reads, neg_reads = reads[0]
if len(reads) > 1:
	for pos_reads_, neg_reads_ in reads[1:]:
		for chrom, reads_ in pos_reads_.items():
			for idx, count in reads_.items():
				if idx in pos_reads[chrom]:
					pos_reads[chrom][idx] += count
				else:
					pos_reads[chrom][idx] = count

		if not args.unstranded:
			for chrom, reads_ in neg_reads_.items():
				for idx, count in reads_.items():
					if idx in neg_reads[chrom]:
						neg_reads[chrom][idx] += count
					else:
						neg_reads[chrom][idx] = count

###

# A bigWig can only hold positions that fall inside the chromosome it declares,
# and pyBigWig discards anything else without raising or returning an error.
# That silently removes reads from the output -- either because the sizes file
# disagrees with the BAM header, or because a shift pushed an end off the edge
# of a chromosome -- and, worse, read depth would be summed over reads that
# never reach the file, leaving the track normalized to less than it claims.
# Discarding them here instead means it can be counted, reported, and excluded
# from the read depth below.

discarded = collections.defaultdict(int)

for chrom, size in chrom_sizes:
	strands = [pos_reads[chrom]]
	if not args.unstranded:
		strands.append(neg_reads[chrom])

	for counts in strands:
		idxs = numpy.fromiter(counts.keys(), dtype='int64', count=len(counts))

		for idx in idxs[(idxs < 0) | (idxs >= size)]:
			discarded[chrom] += counts.pop(idx)

if len(discarded) > 0:
	print("{} entries fell outside the provided chromosome sizes and were discarded ({}).".format(
		sum(discarded.values()), ", ".join("{}: {}".format(chrom, count)
			for chrom, count in discarded.items())))

###

# Now that we have our dictionary(ies) of reads, we need to create bigWig
# objects and store them. Because the entries need to be sorted along the
# length of each chromosome we have to convert the dictionaries to numpy arrays
# and then sort them, but this usually is not that big of a hassle. It is much
# faster to sort the arrays at the end like this that it is to try to keep
# everything in order as you see the reads.

# Here, we open the bigWigs that we will be saving data into. If the data is
# stranded, we are saving two bigWigs. If the data is not stranded, we are only
# saving one bigWig.

if args.unstranded:
    bw_pos = pyBigWig.open(args.name + ".bw", "w")
    bw_pos.addHeader(chrom_sizes, maxZooms=args.zooms)

else:
    bw_pos = pyBigWig.open(args.name + ".+.bw", "w")
    bw_neg = pyBigWig.open(args.name + ".-.bw", "w")

    bw_pos.addHeader(chrom_sizes, maxZooms=args.zooms)
    bw_neg.addHeader(chrom_sizes, maxZooms=args.zooms)


# We use pyBigWig for our tool to create bigWig files. We choose to save
# entries using only the coordinate and the value, so two numbers per non-zero
# position, rather than as spans of size 1 which would be three numbers per
# non-zero position. This reduces file size and also I/O time.

if args.read_depth:
	read_depth = sum([sum(pos_reads[chrom].values()) for chrom, _ in chrom_sizes])
	if not args.unstranded:
		read_depth += sum([sum(neg_reads[chrom].values()) for chrom, _ in chrom_sizes])

	if args.verbose:
		print("Dividing through by a read depth of {}.".format(read_depth))


for chrom, _ in chrom_sizes:
	reads = pos_reads[chrom]
	if len(reads) > 0:
		pos_starts = numpy.array(list(reads.keys()), dtype='int64')
		pos_values = numpy.array(list(reads.values()), dtype='float64')
		pos_values *= args.scale_factor

		if args.read_depth:
			pos_values /= read_depth

		idxs = numpy.argsort(pos_starts)
		bw_pos.addEntries(chrom, pos_starts[idxs], values=pos_values[idxs], span=1)

	###
	
	reads = neg_reads[chrom]
	if len(reads) > 0 and not args.unstranded:
		neg_starts = numpy.array(list(reads.keys()), dtype='int64')
		neg_values = numpy.array(list(reads.values()), dtype='float64')
		neg_values *= args.scale_factor

		if args.read_depth:
			neg_values /= read_depth

		idxs = numpy.argsort(neg_starts)
		bw_neg.addEntries(chrom, neg_starts[idxs], values=neg_values[idxs], span=1)

bw_pos.close()
if not args.unstranded:
	bw_neg.close()
